Index arrays more eagerly
[ghc.git] / libraries / base / GHC / Arr.hs
1 {-# LANGUAGE Unsafe #-}
2 {-# LANGUAGE NoImplicitPrelude, MagicHash, UnboxedTuples, RoleAnnotations #-}
3 {-# LANGUAGE BangPatterns #-}
4 {-# OPTIONS_HADDOCK hide #-}
5
6 -----------------------------------------------------------------------------
7 -- |
8 -- Module : GHC.Arr
9 -- Copyright : (c) The University of Glasgow, 1994-2000
10 -- License : see libraries/base/LICENSE
11 --
12 -- Maintainer : cvs-ghc@haskell.org
13 -- Stability : internal
14 -- Portability : non-portable (GHC extensions)
15 --
16 -- GHC\'s array implementation.
17 --
18 -----------------------------------------------------------------------------
19
20 module GHC.Arr (
21 Ix(..), Array(..), STArray(..),
22
23 indexError, hopelessIndexError,
24 arrEleBottom, array, listArray,
25 (!), safeRangeSize, negRange, safeIndex, badSafeIndex,
26 bounds, numElements, numElementsSTArray, indices, elems,
27 assocs, accumArray, adjust, (//), accum,
28 amap, ixmap,
29 eqArray, cmpArray, cmpIntArray,
30 newSTArray, boundsSTArray,
31 readSTArray, writeSTArray,
32 freezeSTArray, thawSTArray,
33 foldlElems, foldlElems', foldl1Elems,
34 foldrElems, foldrElems', foldr1Elems,
35
36 -- * Unsafe operations
37 fill, done,
38 unsafeArray, unsafeArray',
39 lessSafeIndex, unsafeAt, unsafeReplace,
40 unsafeAccumArray, unsafeAccumArray', unsafeAccum,
41 unsafeReadSTArray, unsafeWriteSTArray,
42 unsafeFreezeSTArray, unsafeThawSTArray,
43 ) where
44
45 import GHC.Enum
46 import GHC.Num
47 import GHC.ST
48 import GHC.Base
49 import GHC.List
50 import GHC.Real( fromIntegral )
51 import GHC.Show
52
53 infixl 9 !, //
54
55 default ()
56
57 -- | The 'Ix' class is used to map a contiguous subrange of values in
58 -- a type onto integers. It is used primarily for array indexing
59 -- (see the array package).
60 --
61 -- The first argument @(l,u)@ of each of these operations is a pair
62 -- specifying the lower and upper bounds of a contiguous subrange of values.
63 --
64 -- An implementation is entitled to assume the following laws about these
65 -- operations:
66 --
67 -- * @'inRange' (l,u) i == 'elem' i ('range' (l,u))@ @ @
68 --
69 -- * @'range' (l,u) '!!' 'index' (l,u) i == i@, when @'inRange' (l,u) i@
70 --
71 -- * @'map' ('index' (l,u)) ('range' (l,u))) == [0..'rangeSize' (l,u)-1]@ @ @
72 --
73 -- * @'rangeSize' (l,u) == 'length' ('range' (l,u))@ @ @
74 --
75 class (Ord a) => Ix a where
76 {-# MINIMAL range, (index | unsafeIndex), inRange #-}
77
78 -- | The list of values in the subrange defined by a bounding pair.
79 range :: (a,a) -> [a]
80 -- | The position of a subscript in the subrange.
81 index :: (a,a) -> a -> Int
82 -- | Like 'index', but without checking that the value is in range.
83 unsafeIndex :: (a,a) -> a -> Int
84 -- | Returns 'True' the given subscript lies in the range defined
85 -- the bounding pair.
86 inRange :: (a,a) -> a -> Bool
87 -- | The size of the subrange defined by a bounding pair.
88 rangeSize :: (a,a) -> Int
89 -- | like 'rangeSize', but without checking that the upper bound is
90 -- in range.
91 unsafeRangeSize :: (a,a) -> Int
92
93 -- Must specify one of index, unsafeIndex
94
95 -- 'index' is typically over-ridden in instances, with essentially
96 -- the same code, but using indexError instead of hopelessIndexError
97 -- Reason: we have 'Show' at the instances
98 {-# INLINE index #-} -- See Note [Inlining index]
99 index b i | inRange b i = unsafeIndex b i
100 | otherwise = hopelessIndexError
101
102 unsafeIndex b i = index b i
103
104 rangeSize b@(_l,h) | inRange b h = unsafeIndex b h + 1
105 | otherwise = 0 -- This case is only here to
106 -- check for an empty range
107 -- NB: replacing (inRange b h) by (l <= h) fails for
108 -- tuples. E.g. (1,2) <= (2,1) but the range is empty
109
110 unsafeRangeSize b@(_l,h) = unsafeIndex b h + 1
111
112 {-
113 Note that the following is NOT right
114 rangeSize (l,h) | l <= h = index b h + 1
115 | otherwise = 0
116
117 Because it might be the case that l<h, but the range
118 is nevertheless empty. Consider
119 ((1,2),(2,1))
120 Here l<h, but the second index ranges from 2..1 and
121 hence is empty
122
123
124 Note [Inlining index]
125 ~~~~~~~~~~~~~~~~~~~~~
126 We inline the 'index' operation,
127
128 * Partly because it generates much faster code
129 (although bigger); see Trac #1216
130
131 * Partly because it exposes the bounds checks to the simplifier which
132 might help a big.
133
134 If you make a per-instance index method, you may consider inlining it.
135
136 Note [Double bounds-checking of index values]
137 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
138 When you index an array, a!x, there are two possible bounds checks we might make:
139
140 (A) Check that (inRange (bounds a) x) holds.
141
142 (A) is checked in the method for 'index'
143
144 (B) Check that (index (bounds a) x) lies in the range 0..n,
145 where n is the size of the underlying array
146
147 (B) is checked in the top-level function (!), in safeIndex.
148
149 Of course it *should* be the case that (A) holds iff (B) holds, but that
150 is a property of the particular instances of index, bounds, and inRange,
151 so GHC cannot guarantee it.
152
153 * If you do (A) and not (B), then you might get a seg-fault,
154 by indexing at some bizarre location. Trac #1610
155
156 * If you do (B) but not (A), you may get no complaint when you index
157 an array out of its semantic bounds. Trac #2120
158
159 At various times we have had (A) and not (B), or (B) and not (A); both
160 led to complaints. So now we implement *both* checks (Trac #2669).
161
162 For 1-d, 2-d, and 3-d arrays of Int we have specialised instances to avoid this.
163
164 Note [Out-of-bounds error messages]
165 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
166 The default method for 'index' generates hoplelessIndexError, because
167 Ix doesn't have Show as a superclass. For particular base types we
168 can do better, so we override the default method for index.
169 -}
170
171 -- Abstract these errors from the relevant index functions so that
172 -- the guts of the function will be small enough to inline.
173
174 {-# NOINLINE indexError #-}
175 indexError :: Show a => (a,a) -> a -> String -> b
176 indexError rng i tp
177 = errorWithoutStackTrace (showString "Ix{" . showString tp . showString "}.index: Index " .
178 showParen True (showsPrec 0 i) .
179 showString " out of range " $
180 showParen True (showsPrec 0 rng) "")
181
182 hopelessIndexError :: Int -- Try to use 'indexError' instead!
183 hopelessIndexError = errorWithoutStackTrace "Error in array index"
184
185 ----------------------------------------------------------------------
186 -- | @since 2.01
187 instance Ix Char where
188 {-# INLINE range #-}
189 range (m,n) = [m..n]
190
191 {-# INLINE unsafeIndex #-}
192 unsafeIndex (m,_n) i = fromEnum i - fromEnum m
193
194 {-# INLINE index #-} -- See Note [Out-of-bounds error messages]
195 -- and Note [Inlining index]
196 index b i | inRange b i = unsafeIndex b i
197 | otherwise = indexError b i "Char"
198
199 inRange (m,n) i = m <= i && i <= n
200
201 ----------------------------------------------------------------------
202 -- | @since 2.01
203 instance Ix Int where
204 {-# INLINE range #-}
205 -- The INLINE stops the build in the RHS from getting inlined,
206 -- so that callers can fuse with the result of range
207 range (m,n) = [m..n]
208
209 {-# INLINE unsafeIndex #-}
210 unsafeIndex (m,_n) i = i - m
211
212 {-# INLINE index #-} -- See Note [Out-of-bounds error messages]
213 -- and Note [Inlining index]
214 index b i | inRange b i = unsafeIndex b i
215 | otherwise = indexError b i "Int"
216
217 {-# INLINE inRange #-}
218 inRange (I# m,I# n) (I# i) = isTrue# (m <=# i) && isTrue# (i <=# n)
219
220 -- | @since 4.6.0.0
221 instance Ix Word where
222 range (m,n) = [m..n]
223 unsafeIndex (m,_) i = fromIntegral (i - m)
224 inRange (m,n) i = m <= i && i <= n
225
226 ----------------------------------------------------------------------
227 -- | @since 2.01
228 instance Ix Integer where
229 {-# INLINE range #-}
230 range (m,n) = [m..n]
231
232 {-# INLINE unsafeIndex #-}
233 unsafeIndex (m,_n) i = fromInteger (i - m)
234
235 {-# INLINE index #-} -- See Note [Out-of-bounds error messages]
236 -- and Note [Inlining index]
237 index b i | inRange b i = unsafeIndex b i
238 | otherwise = indexError b i "Integer"
239
240 inRange (m,n) i = m <= i && i <= n
241
242 ----------------------------------------------------------------------
243 -- | @since 2.01
244 instance Ix Bool where -- as derived
245 {-# INLINE range #-}
246 range (m,n) = [m..n]
247
248 {-# INLINE unsafeIndex #-}
249 unsafeIndex (l,_) i = fromEnum i - fromEnum l
250
251 {-# INLINE index #-} -- See Note [Out-of-bounds error messages]
252 -- and Note [Inlining index]
253 index b i | inRange b i = unsafeIndex b i
254 | otherwise = indexError b i "Bool"
255
256 inRange (l,u) i = fromEnum i >= fromEnum l && fromEnum i <= fromEnum u
257
258 ----------------------------------------------------------------------
259 -- | @since 2.01
260 instance Ix Ordering where -- as derived
261 {-# INLINE range #-}
262 range (m,n) = [m..n]
263
264 {-# INLINE unsafeIndex #-}
265 unsafeIndex (l,_) i = fromEnum i - fromEnum l
266
267 {-# INLINE index #-} -- See Note [Out-of-bounds error messages]
268 -- and Note [Inlining index]
269 index b i | inRange b i = unsafeIndex b i
270 | otherwise = indexError b i "Ordering"
271
272 inRange (l,u) i = fromEnum i >= fromEnum l && fromEnum i <= fromEnum u
273
274 ----------------------------------------------------------------------
275 -- | @since 2.01
276 instance Ix () where
277 {-# INLINE range #-}
278 range ((), ()) = [()]
279 {-# INLINE unsafeIndex #-}
280 unsafeIndex ((), ()) () = 0
281 {-# INLINE inRange #-}
282 inRange ((), ()) () = True
283
284 {-# INLINE index #-} -- See Note [Inlining index]
285 index b i = unsafeIndex b i
286
287 ----------------------------------------------------------------------
288 -- | @since 2.01
289 instance (Ix a, Ix b) => Ix (a, b) where -- as derived
290 {-# SPECIALISE instance Ix (Int,Int) #-}
291
292 {-# INLINE range #-}
293 range ((l1,l2),(u1,u2)) =
294 [ (i1,i2) | i1 <- range (l1,u1), i2 <- range (l2,u2) ]
295
296 {-# INLINE unsafeIndex #-}
297 unsafeIndex ((l1,l2),(u1,u2)) (i1,i2) =
298 unsafeIndex (l1,u1) i1 * unsafeRangeSize (l2,u2) + unsafeIndex (l2,u2) i2
299
300 {-# INLINE inRange #-}
301 inRange ((l1,l2),(u1,u2)) (i1,i2) =
302 inRange (l1,u1) i1 && inRange (l2,u2) i2
303
304 -- Default method for index
305
306 ----------------------------------------------------------------------
307 -- | @since 2.01
308 instance (Ix a1, Ix a2, Ix a3) => Ix (a1,a2,a3) where
309 {-# SPECIALISE instance Ix (Int,Int,Int) #-}
310
311 range ((l1,l2,l3),(u1,u2,u3)) =
312 [(i1,i2,i3) | i1 <- range (l1,u1),
313 i2 <- range (l2,u2),
314 i3 <- range (l3,u3)]
315
316 unsafeIndex ((l1,l2,l3),(u1,u2,u3)) (i1,i2,i3) =
317 unsafeIndex (l3,u3) i3 + unsafeRangeSize (l3,u3) * (
318 unsafeIndex (l2,u2) i2 + unsafeRangeSize (l2,u2) * (
319 unsafeIndex (l1,u1) i1))
320
321 inRange ((l1,l2,l3),(u1,u2,u3)) (i1,i2,i3) =
322 inRange (l1,u1) i1 && inRange (l2,u2) i2 &&
323 inRange (l3,u3) i3
324
325 -- Default method for index
326
327 ----------------------------------------------------------------------
328 -- | @since 2.01
329 instance (Ix a1, Ix a2, Ix a3, Ix a4) => Ix (a1,a2,a3,a4) where
330 range ((l1,l2,l3,l4),(u1,u2,u3,u4)) =
331 [(i1,i2,i3,i4) | i1 <- range (l1,u1),
332 i2 <- range (l2,u2),
333 i3 <- range (l3,u3),
334 i4 <- range (l4,u4)]
335
336 unsafeIndex ((l1,l2,l3,l4),(u1,u2,u3,u4)) (i1,i2,i3,i4) =
337 unsafeIndex (l4,u4) i4 + unsafeRangeSize (l4,u4) * (
338 unsafeIndex (l3,u3) i3 + unsafeRangeSize (l3,u3) * (
339 unsafeIndex (l2,u2) i2 + unsafeRangeSize (l2,u2) * (
340 unsafeIndex (l1,u1) i1)))
341
342 inRange ((l1,l2,l3,l4),(u1,u2,u3,u4)) (i1,i2,i3,i4) =
343 inRange (l1,u1) i1 && inRange (l2,u2) i2 &&
344 inRange (l3,u3) i3 && inRange (l4,u4) i4
345
346 -- Default method for index
347 -- | @since 2.01
348 instance (Ix a1, Ix a2, Ix a3, Ix a4, Ix a5) => Ix (a1,a2,a3,a4,a5) where
349 range ((l1,l2,l3,l4,l5),(u1,u2,u3,u4,u5)) =
350 [(i1,i2,i3,i4,i5) | i1 <- range (l1,u1),
351 i2 <- range (l2,u2),
352 i3 <- range (l3,u3),
353 i4 <- range (l4,u4),
354 i5 <- range (l5,u5)]
355
356 unsafeIndex ((l1,l2,l3,l4,l5),(u1,u2,u3,u4,u5)) (i1,i2,i3,i4,i5) =
357 unsafeIndex (l5,u5) i5 + unsafeRangeSize (l5,u5) * (
358 unsafeIndex (l4,u4) i4 + unsafeRangeSize (l4,u4) * (
359 unsafeIndex (l3,u3) i3 + unsafeRangeSize (l3,u3) * (
360 unsafeIndex (l2,u2) i2 + unsafeRangeSize (l2,u2) * (
361 unsafeIndex (l1,u1) i1))))
362
363 inRange ((l1,l2,l3,l4,l5),(u1,u2,u3,u4,u5)) (i1,i2,i3,i4,i5) =
364 inRange (l1,u1) i1 && inRange (l2,u2) i2 &&
365 inRange (l3,u3) i3 && inRange (l4,u4) i4 &&
366 inRange (l5,u5) i5
367
368 -- Default method for index
369
370 -- | The type of immutable non-strict (boxed) arrays
371 -- with indices in @i@ and elements in @e@.
372 data Array i e
373 = Array !i -- the lower bound, l
374 !i -- the upper bound, u
375 {-# UNPACK #-} !Int -- A cache of (rangeSize (l,u))
376 -- used to make sure an index is
377 -- really in range
378 (Array# e) -- The actual elements
379
380 -- | Mutable, boxed, non-strict arrays in the 'ST' monad. The type
381 -- arguments are as follows:
382 --
383 -- * @s@: the state variable argument for the 'ST' type
384 --
385 -- * @i@: the index type of the array (should be an instance of 'Ix')
386 --
387 -- * @e@: the element type of the array.
388 --
389 data STArray s i e
390 = STArray !i -- the lower bound, l
391 !i -- the upper bound, u
392 {-# UNPACK #-} !Int -- A cache of (rangeSize (l,u))
393 -- used to make sure an index is
394 -- really in range
395 (MutableArray# s e) -- The actual elements
396 -- No Ix context for STArray. They are stupid,
397 -- and force an Ix context on the equality instance.
398
399 -- Index types should have nominal role, because of Ix class. See also #9220.
400 type role Array nominal representational
401 type role STArray nominal nominal representational
402
403 -- Just pointer equality on mutable arrays:
404 -- | @since 2.01
405 instance Eq (STArray s i e) where
406 STArray _ _ _ arr1# == STArray _ _ _ arr2# =
407 isTrue# (sameMutableArray# arr1# arr2#)
408
409 ----------------------------------------------------------------------
410 -- Operations on immutable arrays
411
412 {-# NOINLINE arrEleBottom #-}
413 arrEleBottom :: a
414 arrEleBottom = errorWithoutStackTrace "(Array.!): undefined array element"
415
416 -- | Construct an array with the specified bounds and containing values
417 -- for given indices within these bounds.
418 --
419 -- The array is undefined (i.e. bottom) if any index in the list is
420 -- out of bounds. The Haskell 2010 Report further specifies that if any
421 -- two associations in the list have the same index, the value at that
422 -- index is undefined (i.e. bottom). However in GHC's implementation,
423 -- the value at such an index is the value part of the last association
424 -- with that index in the list.
425 --
426 -- Because the indices must be checked for these errors, 'array' is
427 -- strict in the bounds argument and in the indices of the association
428 -- list, but non-strict in the values. Thus, recurrences such as the
429 -- following are possible:
430 --
431 -- > a = array (1,100) ((1,1) : [(i, i * a!(i-1)) | i <- [2..100]])
432 --
433 -- Not every index within the bounds of the array need appear in the
434 -- association list, but the values associated with indices that do not
435 -- appear will be undefined (i.e. bottom).
436 --
437 -- If, in any dimension, the lower bound is greater than the upper bound,
438 -- then the array is legal, but empty. Indexing an empty array always
439 -- gives an array-bounds error, but 'bounds' still yields the bounds
440 -- with which the array was constructed.
441 {-# INLINE array #-}
442 array :: Ix i
443 => (i,i) -- ^ a pair of /bounds/, each of the index type
444 -- of the array. These bounds are the lowest and
445 -- highest indices in the array, in that order.
446 -- For example, a one-origin vector of length
447 -- '10' has bounds '(1,10)', and a one-origin '10'
448 -- by '10' matrix has bounds '((1,1),(10,10))'.
449 -> [(i, e)] -- ^ a list of /associations/ of the form
450 -- (/index/, /value/). Typically, this list will
451 -- be expressed as a comprehension. An
452 -- association '(i, x)' defines the value of
453 -- the array at index 'i' to be 'x'.
454 -> Array i e
455 array (l,u) ies
456 = let n = safeRangeSize (l,u)
457 in unsafeArray' (l,u) n
458 [(safeIndex (l,u) n i, e) | (i, e) <- ies]
459
460 {-# INLINE unsafeArray #-}
461 unsafeArray :: Ix i => (i,i) -> [(Int, e)] -> Array i e
462 unsafeArray b ies = unsafeArray' b (rangeSize b) ies
463
464 {-# INLINE unsafeArray' #-}
465 unsafeArray' :: (i,i) -> Int -> [(Int, e)] -> Array i e
466 unsafeArray' (l,u) n@(I# n#) ies = runST (ST $ \s1# ->
467 case newArray# n# arrEleBottom s1# of
468 (# s2#, marr# #) ->
469 foldr (fill marr#) (done l u n marr#) ies s2#)
470
471 {-# INLINE fill #-}
472 fill :: MutableArray# s e -> (Int, e) -> STRep s a -> STRep s a
473 -- NB: put the \s after the "=" so that 'fill'
474 -- inlines when applied to three args
475 fill marr# (I# i#, e) next
476 = \s1# -> case writeArray# marr# i# e s1# of
477 s2# -> next s2#
478
479 {-# INLINE done #-}
480 done :: i -> i -> Int -> MutableArray# s e -> STRep s (Array i e)
481 -- See NB on 'fill'
482 -- Make sure it is strict in 'n'
483 done l u n@(I# _) marr#
484 = \s1# -> case unsafeFreezeArray# marr# s1# of
485 (# s2#, arr# #) -> (# s2#, Array l u n arr# #)
486
487 -- | Construct an array from a pair of bounds and a list of values in
488 -- index order.
489 {-# INLINE listArray #-}
490 listArray :: Ix i => (i,i) -> [e] -> Array i e
491 listArray (l,u) es = runST (ST $ \s1# ->
492 case safeRangeSize (l,u) of { n@(I# n#) ->
493 case newArray# n# arrEleBottom s1# of { (# s2#, marr# #) ->
494 let
495 go y r = \ i# s3# ->
496 case writeArray# marr# i# y s3# of
497 s4# -> if (isTrue# (i# ==# n# -# 1#))
498 then s4#
499 else r (i# +# 1#) s4#
500 in
501 done l u n marr# (
502 if n == 0
503 then s2#
504 else foldr go (\_ s# -> s#) es 0# s2#)}})
505
506 -- | The value at the given index in an array.
507 {-# INLINE (!) #-}
508 (!) :: Ix i => Array i e -> i -> e
509 (!) arr@(Array l u n _) i = unsafeAt arr $ safeIndex (l,u) n i
510
511 {-# INLINE (!#) #-}
512 (!#) :: Ix i => Array i e -> i -> (# e #)
513 (!#) arr@(Array l u n _) i = unsafeAt# arr $ safeIndex (l,u) n i
514
515 {-# INLINE safeRangeSize #-}
516 safeRangeSize :: Ix i => (i, i) -> Int
517 safeRangeSize (l,u) = let r = rangeSize (l, u)
518 in if r < 0 then negRange
519 else r
520
521 -- Don't inline this error message everywhere!!
522 negRange :: Int -- Uninformative, but Ix does not provide Show
523 negRange = errorWithoutStackTrace "Negative range size"
524
525 {-# INLINE[1] safeIndex #-}
526 -- See Note [Double bounds-checking of index values]
527 -- Inline *after* (!) so the rules can fire
528 -- Make sure it is strict in n
529 safeIndex :: Ix i => (i, i) -> Int -> i -> Int
530 safeIndex (l,u) n@(I# _) i
531 | (0 <= i') && (i' < n) = i'
532 | otherwise = badSafeIndex i' n
533 where
534 i' = index (l,u) i
535
536 -- See Note [Double bounds-checking of index values]
537 {-# RULES
538 "safeIndex/I" safeIndex = lessSafeIndex :: (Int,Int) -> Int -> Int -> Int
539 "safeIndex/(I,I)" safeIndex = lessSafeIndex :: ((Int,Int),(Int,Int)) -> Int -> (Int,Int) -> Int
540 "safeIndex/(I,I,I)" safeIndex = lessSafeIndex :: ((Int,Int,Int),(Int,Int,Int)) -> Int -> (Int,Int,Int) -> Int
541 #-}
542
543 lessSafeIndex :: Ix i => (i, i) -> Int -> i -> Int
544 -- See Note [Double bounds-checking of index values]
545 -- Do only (A), the semantic check
546 lessSafeIndex (l,u) _ i = index (l,u) i
547
548 -- Don't inline this long error message everywhere!!
549 badSafeIndex :: Int -> Int -> Int
550 badSafeIndex i' n = errorWithoutStackTrace ("Error in array index; " ++ show i' ++
551 " not in range [0.." ++ show n ++ ")")
552
553 {-# INLINE unsafeAt #-}
554 unsafeAt :: Array i e -> Int -> e
555 unsafeAt (Array _ _ _ arr#) (I# i#) =
556 case indexArray# arr# i# of (# e #) -> e
557
558 -- | Look up an element in an array without forcing it
559 unsafeAt# :: Array i e -> Int -> (# e #)
560 unsafeAt# (Array _ _ _ arr#) (I# i#) = indexArray# arr# i#
561
562 -- | A convenient version of unsafeAt#
563 unsafeAtA :: Applicative f
564 => Array i e -> Int -> f e
565 unsafeAtA ary i = case unsafeAt# ary i of (# e #) -> pure e
566
567 -- | The bounds with which an array was constructed.
568 {-# INLINE bounds #-}
569 bounds :: Array i e -> (i,i)
570 bounds (Array l u _ _) = (l,u)
571
572 -- | The number of elements in the array.
573 {-# INLINE numElements #-}
574 numElements :: Array i e -> Int
575 numElements (Array _ _ n _) = n
576
577 -- | The list of indices of an array in ascending order.
578 {-# INLINE indices #-}
579 indices :: Ix i => Array i e -> [i]
580 indices (Array l u _ _) = range (l,u)
581
582 -- | The list of elements of an array in index order.
583 {-# INLINE elems #-}
584 elems :: Array i e -> [e]
585 elems arr@(Array _ _ n _) =
586 [e | i <- [0 .. n - 1], e <- unsafeAtA arr i]
587
588 -- | A right fold over the elements
589 {-# INLINABLE foldrElems #-}
590 foldrElems :: (a -> b -> b) -> b -> Array i a -> b
591 foldrElems f b0 = \ arr@(Array _ _ n _) ->
592 let
593 go i | i == n = b0
594 | (# e #) <- unsafeAt# arr i
595 = f e (go (i+1))
596 in go 0
597
598 -- | A left fold over the elements
599 {-# INLINABLE foldlElems #-}
600 foldlElems :: (b -> a -> b) -> b -> Array i a -> b
601 foldlElems f b0 = \ arr@(Array _ _ n _) ->
602 let
603 go i | i == (-1) = b0
604 | (# e #) <- unsafeAt# arr i
605 = f (go (i-1)) e
606 in go (n-1)
607
608 -- | A strict right fold over the elements
609 {-# INLINABLE foldrElems' #-}
610 foldrElems' :: (a -> b -> b) -> b -> Array i a -> b
611 foldrElems' f b0 = \ arr@(Array _ _ n _) ->
612 let
613 go i a | i == (-1) = a
614 | (# e #) <- unsafeAt# arr i
615 = go (i-1) (f e $! a)
616 in go (n-1) b0
617
618 -- | A strict left fold over the elements
619 {-# INLINABLE foldlElems' #-}
620 foldlElems' :: (b -> a -> b) -> b -> Array i a -> b
621 foldlElems' f b0 = \ arr@(Array _ _ n _) ->
622 let
623 go i a | i == n = a
624 | (# e #) <- unsafeAt# arr i
625 = go (i+1) (a `seq` f a e)
626 in go 0 b0
627
628 -- | A left fold over the elements with no starting value
629 {-# INLINABLE foldl1Elems #-}
630 foldl1Elems :: (a -> a -> a) -> Array i a -> a
631 foldl1Elems f = \ arr@(Array _ _ n _) ->
632 let
633 go i | i == 0 = unsafeAt arr 0
634 | (# e #) <- unsafeAt# arr i
635 = f (go (i-1)) e
636 in
637 if n == 0 then errorWithoutStackTrace "foldl1: empty Array" else go (n-1)
638
639 -- | A right fold over the elements with no starting value
640 {-# INLINABLE foldr1Elems #-}
641 foldr1Elems :: (a -> a -> a) -> Array i a -> a
642 foldr1Elems f = \ arr@(Array _ _ n _) ->
643 let
644 go i | i == n-1 = unsafeAt arr i
645 | (# e #) <- unsafeAt# arr i
646 = f e (go (i + 1))
647 in
648 if n == 0 then errorWithoutStackTrace "foldr1: empty Array" else go 0
649
650 -- | The list of associations of an array in index order.
651 {-# INLINE assocs #-}
652 assocs :: Ix i => Array i e -> [(i, e)]
653 assocs arr@(Array l u _ _) =
654 [(i, e) | i <- range (l,u), let !(# e #) = arr !# i]
655
656 -- | The 'accumArray' function deals with repeated indices in the association
657 -- list using an /accumulating function/ which combines the values of
658 -- associations with the same index.
659 --
660 -- For example, given a list of values of some index type, @hist@
661 -- produces a histogram of the number of occurrences of each index within
662 -- a specified range:
663 --
664 -- > hist :: (Ix a, Num b) => (a,a) -> [a] -> Array a b
665 -- > hist bnds is = accumArray (+) 0 bnds [(i, 1) | i<-is, inRange bnds i]
666 --
667 -- @accumArray@ is strict in each result of applying the accumulating
668 -- function, although it is lazy in the initial value. Thus, unlike
669 -- arrays built with 'array', accumulated arrays should not in general
670 -- be recursive.
671 {-# INLINE accumArray #-}
672 accumArray :: Ix i
673 => (e -> a -> e) -- ^ accumulating function
674 -> e -- ^ initial value
675 -> (i,i) -- ^ bounds of the array
676 -> [(i, a)] -- ^ association list
677 -> Array i e
678 accumArray f initial (l,u) ies =
679 let n = safeRangeSize (l,u)
680 in unsafeAccumArray' f initial (l,u) n
681 [(safeIndex (l,u) n i, e) | (i, e) <- ies]
682
683 {-# INLINE unsafeAccumArray #-}
684 unsafeAccumArray :: Ix i => (e -> a -> e) -> e -> (i,i) -> [(Int, a)] -> Array i e
685 unsafeAccumArray f initial b ies = unsafeAccumArray' f initial b (rangeSize b) ies
686
687 {-# INLINE unsafeAccumArray' #-}
688 unsafeAccumArray' :: (e -> a -> e) -> e -> (i,i) -> Int -> [(Int, a)] -> Array i e
689 unsafeAccumArray' f initial (l,u) n@(I# n#) ies = runST (ST $ \s1# ->
690 case newArray# n# initial s1# of { (# s2#, marr# #) ->
691 foldr (adjust' f marr#) (done l u n marr#) ies s2# })
692
693 {-# INLINE adjust #-}
694 adjust :: (e -> a -> e) -> MutableArray# s e -> (Int, a) -> STRep s b -> STRep s b
695 -- See NB on 'fill'
696 adjust f marr# (I# i#, new) next
697 = \s1# -> case readArray# marr# i# s1# of
698 (# s2#, old #) ->
699 case writeArray# marr# i# (f old new) s2# of
700 s3# -> next s3#
701
702 {-# INLINE adjust' #-}
703 adjust' :: (e -> a -> e)
704 -> MutableArray# s e
705 -> (Int, a)
706 -> STRep s b -> STRep s b
707 adjust' f marr# (I# i#, new) next
708 = \s1# -> case readArray# marr# i# s1# of
709 (# s2#, old #) ->
710 let !combined = f old new
711 in next (writeArray# marr# i# combined s2#)
712
713
714 -- | Constructs an array identical to the first argument except that it has
715 -- been updated by the associations in the right argument.
716 -- For example, if @m@ is a 1-origin, @n@ by @n@ matrix, then
717 --
718 -- > m//[((i,i), 0) | i <- [1..n]]
719 --
720 -- is the same matrix, except with the diagonal zeroed.
721 --
722 -- Repeated indices in the association list are handled as for 'array':
723 -- Haskell 2010 specifies that the resulting array is undefined (i.e. bottom),
724 -- but GHC's implementation uses the last association for each index.
725 {-# INLINE (//) #-}
726 (//) :: Ix i => Array i e -> [(i, e)] -> Array i e
727 arr@(Array l u n _) // ies =
728 unsafeReplace arr [(safeIndex (l,u) n i, e) | (i, e) <- ies]
729
730 {-# INLINE unsafeReplace #-}
731 unsafeReplace :: Array i e -> [(Int, e)] -> Array i e
732 unsafeReplace arr ies = runST (do
733 STArray l u n marr# <- thawSTArray arr
734 ST (foldr (fill marr#) (done l u n marr#) ies))
735
736 -- | @'accum' f@ takes an array and an association list and accumulates
737 -- pairs from the list into the array with the accumulating function @f@.
738 -- Thus 'accumArray' can be defined using 'accum':
739 --
740 -- > accumArray f z b = accum f (array b [(i, z) | i <- range b])
741 --
742 -- @accum@ is strict in all the results of applying the accumulation.
743 -- However, it is lazy in the initial values of the array.
744 {-# INLINE accum #-}
745 accum :: Ix i => (e -> a -> e) -> Array i e -> [(i, a)] -> Array i e
746 accum f arr@(Array l u n _) ies =
747 unsafeAccum f arr [(safeIndex (l,u) n i, e) | (i, e) <- ies]
748
749 {-# INLINE unsafeAccum #-}
750 unsafeAccum :: (e -> a -> e) -> Array i e -> [(Int, a)] -> Array i e
751 unsafeAccum f arr ies = runST (do
752 STArray l u n marr# <- thawSTArray arr
753 ST (foldr (adjust' f marr#) (done l u n marr#) ies))
754
755 {-# INLINE [1] amap #-} -- See Note [amap]
756 amap :: (a -> b) -> Array i a -> Array i b
757 amap f arr@(Array l u n@(I# n#) _) = runST (ST $ \s1# ->
758 case newArray# n# arrEleBottom s1# of
759 (# s2#, marr# #) ->
760 let go i s#
761 | i == n = done l u n marr# s#
762 | (# e #) <- unsafeAt# arr i
763 = fill marr# (i, f e) (go (i+1)) s#
764 in go 0 s2# )
765
766 {- Note [amap]
767 ~~~~~~~~~~~~~~
768 amap was originally defined like this:
769
770 amap f arr@(Array l u n _) =
771 unsafeArray' (l,u) n [(i, f (unsafeAt arr i)) | i <- [0 .. n - 1]]
772
773 There are two problems:
774
775 1. The enumFromTo implementation produces (spurious) code for the impossible
776 case of n<0 that ends up duplicating the array freezing code.
777
778 2. This implementation relies on list fusion for efficiency. In order
779 to implement the "amap/coerce" rule, we need to delay inlining amap
780 until simplifier phase 1, which is when the eftIntList rule kicks
781 in and makes that impossible. (c.f. Trac #8767)
782 -}
783
784
785 -- See Breitner, Eisenberg, Peyton Jones, and Weirich, "Safe Zero-cost
786 -- Coercions for Haskell", section 6.5:
787 -- http://research.microsoft.com/en-us/um/people/simonpj/papers/ext-f/coercible.pdf
788 {-# RULES
789 "amap/coerce" amap coerce = coerce -- See Note [amap]
790 #-}
791
792 -- Second functor law:
793 {-# RULES
794 "amap/amap" forall f g a . amap f (amap g a) = amap (f . g) a
795 #-}
796
797 -- | 'ixmap' allows for transformations on array indices.
798 -- It may be thought of as providing function composition on the right
799 -- with the mapping that the original array embodies.
800 --
801 -- A similar transformation of array values may be achieved using 'fmap'
802 -- from the 'Array' instance of the 'Functor' class.
803 {-# INLINE ixmap #-}
804 ixmap :: (Ix i, Ix j) => (i,i) -> (i -> j) -> Array j e -> Array i e
805 ixmap (l,u) f arr =
806 array (l,u) [(i, arr ! f i) | i <- range (l,u)]
807
808 {-# INLINE eqArray #-}
809 eqArray :: (Ix i, Eq e) => Array i e -> Array i e -> Bool
810 eqArray arr1@(Array l1 u1 n1 _) arr2@(Array l2 u2 n2 _) =
811 if n1 == 0 then n2 == 0 else
812 l1 == l2 && u1 == u2 &&
813 and [unsafeAt arr1 i == unsafeAt arr2 i | i <- [0 .. n1 - 1]]
814
815 {-# INLINE [1] cmpArray #-}
816 cmpArray :: (Ix i, Ord e) => Array i e -> Array i e -> Ordering
817 cmpArray arr1 arr2 = compare (assocs arr1) (assocs arr2)
818
819 {-# INLINE cmpIntArray #-}
820 cmpIntArray :: Ord e => Array Int e -> Array Int e -> Ordering
821 cmpIntArray arr1@(Array l1 u1 n1 _) arr2@(Array l2 u2 n2 _) =
822 if n1 == 0 then
823 if n2 == 0 then EQ else LT
824 else if n2 == 0 then GT
825 else case compare l1 l2 of
826 EQ -> foldr cmp (compare u1 u2) [0 .. (n1 `min` n2) - 1]
827 other -> other
828 where
829 cmp i rest = case compare (unsafeAt arr1 i) (unsafeAt arr2 i) of
830 EQ -> rest
831 other -> other
832
833 {-# RULES "cmpArray/Int" cmpArray = cmpIntArray #-}
834
835 ----------------------------------------------------------------------
836 -- Array instances
837
838 -- | @since 2.01
839 instance Functor (Array i) where
840 fmap = amap
841
842 -- | @since 2.01
843 instance (Ix i, Eq e) => Eq (Array i e) where
844 (==) = eqArray
845
846 -- | @since 2.01
847 instance (Ix i, Ord e) => Ord (Array i e) where
848 compare = cmpArray
849
850 -- | @since 2.01
851 instance (Ix a, Show a, Show b) => Show (Array a b) where
852 showsPrec p a =
853 showParen (p > appPrec) $
854 showString "array " .
855 showsPrec appPrec1 (bounds a) .
856 showChar ' ' .
857 showsPrec appPrec1 (assocs a)
858 -- Precedence of 'array' is the precedence of application
859
860 -- The Read instance is in GHC.Read
861
862 ----------------------------------------------------------------------
863 -- Operations on mutable arrays
864
865 {-
866 Idle ADR question: What's the tradeoff here between flattening these
867 datatypes into @STArray ix ix (MutableArray# s elt)@ and using
868 it as is? As I see it, the former uses slightly less heap and
869 provides faster access to the individual parts of the bounds while the
870 code used has the benefit of providing a ready-made @(lo, hi)@ pair as
871 required by many array-related functions. Which wins? Is the
872 difference significant (probably not).
873
874 Idle AJG answer: When I looked at the outputted code (though it was 2
875 years ago) it seems like you often needed the tuple, and we build
876 it frequently. Now we've got the overloading specialiser things
877 might be different, though.
878 -}
879
880 {-# INLINE newSTArray #-}
881 newSTArray :: Ix i => (i,i) -> e -> ST s (STArray s i e)
882 newSTArray (l,u) initial = ST $ \s1# ->
883 case safeRangeSize (l,u) of { n@(I# n#) ->
884 case newArray# n# initial s1# of { (# s2#, marr# #) ->
885 (# s2#, STArray l u n marr# #) }}
886
887 {-# INLINE boundsSTArray #-}
888 boundsSTArray :: STArray s i e -> (i,i)
889 boundsSTArray (STArray l u _ _) = (l,u)
890
891 {-# INLINE numElementsSTArray #-}
892 numElementsSTArray :: STArray s i e -> Int
893 numElementsSTArray (STArray _ _ n _) = n
894
895 {-# INLINE readSTArray #-}
896 readSTArray :: Ix i => STArray s i e -> i -> ST s e
897 readSTArray marr@(STArray l u n _) i =
898 unsafeReadSTArray marr (safeIndex (l,u) n i)
899
900 {-# INLINE unsafeReadSTArray #-}
901 unsafeReadSTArray :: STArray s i e -> Int -> ST s e
902 unsafeReadSTArray (STArray _ _ _ marr#) (I# i#)
903 = ST $ \s1# -> readArray# marr# i# s1#
904
905 {-# INLINE writeSTArray #-}
906 writeSTArray :: Ix i => STArray s i e -> i -> e -> ST s ()
907 writeSTArray marr@(STArray l u n _) i e =
908 unsafeWriteSTArray marr (safeIndex (l,u) n i) e
909
910 {-# INLINE unsafeWriteSTArray #-}
911 unsafeWriteSTArray :: STArray s i e -> Int -> e -> ST s ()
912 unsafeWriteSTArray (STArray _ _ _ marr#) (I# i#) e = ST $ \s1# ->
913 case writeArray# marr# i# e s1# of
914 s2# -> (# s2#, () #)
915
916 ----------------------------------------------------------------------
917 -- Moving between mutable and immutable
918
919 freezeSTArray :: STArray s i e -> ST s (Array i e)
920 freezeSTArray (STArray l u n@(I# n#) marr#) = ST $ \s1# ->
921 case newArray# n# arrEleBottom s1# of { (# s2#, marr'# #) ->
922 let copy i# s3# | isTrue# (i# ==# n#) = s3#
923 | otherwise =
924 case readArray# marr# i# s3# of { (# s4#, e #) ->
925 case writeArray# marr'# i# e s4# of { s5# ->
926 copy (i# +# 1#) s5# }} in
927 case copy 0# s2# of { s3# ->
928 case unsafeFreezeArray# marr'# s3# of { (# s4#, arr# #) ->
929 (# s4#, Array l u n arr# #) }}}
930
931 {-# INLINE unsafeFreezeSTArray #-}
932 unsafeFreezeSTArray :: STArray s i e -> ST s (Array i e)
933 unsafeFreezeSTArray (STArray l u n marr#) = ST $ \s1# ->
934 case unsafeFreezeArray# marr# s1# of { (# s2#, arr# #) ->
935 (# s2#, Array l u n arr# #) }
936
937 thawSTArray :: Array i e -> ST s (STArray s i e)
938 thawSTArray (Array l u n@(I# n#) arr#) = ST $ \s1# ->
939 case newArray# n# arrEleBottom s1# of { (# s2#, marr# #) ->
940 let copy i# s3# | isTrue# (i# ==# n#) = s3#
941 | otherwise =
942 case indexArray# arr# i# of { (# e #) ->
943 case writeArray# marr# i# e s3# of { s4# ->
944 copy (i# +# 1#) s4# }} in
945 case copy 0# s2# of { s3# ->
946 (# s3#, STArray l u n marr# #) }}
947
948 {-# INLINE unsafeThawSTArray #-}
949 unsafeThawSTArray :: Array i e -> ST s (STArray s i e)
950 unsafeThawSTArray (Array l u n arr#) = ST $ \s1# ->
951 case unsafeThawArray# arr# s1# of { (# s2#, marr# #) ->
952 (# s2#, STArray l u n marr# #) }