Arr.hs 32.7 KB
Newer Older
1
{-# LANGUAGE Unsafe #-}
2
{-# LANGUAGE NoImplicitPrelude, MagicHash, UnboxedTuples, RoleAnnotations #-}
3
{-# OPTIONS_HADDOCK hide #-}
4

5 6 7 8 9
-----------------------------------------------------------------------------
-- |
-- Module      :  GHC.Arr
-- Copyright   :  (c) The University of Glasgow, 1994-2000
-- License     :  see libraries/base/LICENSE
Jan Stolarek's avatar
Jan Stolarek committed
10
--
11 12 13 14
-- Maintainer  :  cvs-ghc@haskell.org
-- Stability   :  internal
-- Portability :  non-portable (GHC extensions)
--
15
-- GHC\'s array implementation.
Jan Stolarek's avatar
Jan Stolarek committed
16
--
17
-----------------------------------------------------------------------------
18

19 20 21 22 23 24 25 26 27 28 29 30 31
module GHC.Arr (
        Ix(..), Array(..), STArray(..),

        indexError, hopelessIndexError,
        arrEleBottom, array, listArray,
        (!), safeRangeSize, negRange, safeIndex, badSafeIndex,
        bounds, numElements, numElementsSTArray, indices, elems,
        assocs, accumArray, adjust, (//), accum,
        amap, ixmap,
        eqArray, cmpArray, cmpIntArray,
        newSTArray, boundsSTArray,
        readSTArray, writeSTArray,
        freezeSTArray, thawSTArray,
32 33
        foldlElems, foldlElems', foldl1Elems,
        foldrElems, foldrElems', foldr1Elems,
34 35 36 37 38 39 40 41 42

        -- * Unsafe operations
        fill, done,
        unsafeArray, unsafeArray',
        lessSafeIndex, unsafeAt, unsafeReplace,
        unsafeAccumArray, unsafeAccumArray', unsafeAccum,
        unsafeReadSTArray, unsafeWriteSTArray,
        unsafeFreezeSTArray, unsafeThawSTArray,
    ) where
43 44 45 46 47 48

import GHC.Enum
import GHC.Num
import GHC.ST
import GHC.Base
import GHC.List
Simon Peyton Jones's avatar
Simon Peyton Jones committed
49
import GHC.Real( fromIntegral )
50 51 52 53 54 55
import GHC.Show

infixl 9  !, //

default ()

ross's avatar
ross committed
56 57
-- | The 'Ix' class is used to map a contiguous subrange of values in
-- a type onto integers.  It is used primarily for array indexing
Ian Lynagh's avatar
Ian Lynagh committed
58
-- (see the array package).
ross's avatar
ross committed
59 60 61 62 63 64 65
--
-- The first argument @(l,u)@ of each of these operations is a pair
-- specifying the lower and upper bounds of a contiguous subrange of values.
--
-- An implementation is entitled to assume the following laws about these
-- operations:
--
66
-- * @'inRange' (l,u) i == 'elem' i ('range' (l,u))@ @ @
ross's avatar
ross committed
67 68 69
--
-- * @'range' (l,u) '!!' 'index' (l,u) i == i@, when @'inRange' (l,u) i@
--
70
-- * @'map' ('index' (l,u)) ('range' (l,u))) == [0..'rangeSize' (l,u)-1]@ @ @
ross's avatar
ross committed
71
--
72
-- * @'rangeSize' (l,u) == 'length' ('range' (l,u))@ @ @
ross's avatar
ross committed
73
--
74
class (Ord a) => Ix a where
75 76
    {-# MINIMAL range, (index | unsafeIndex), inRange #-}

ross's avatar
ross committed
77
    -- | The list of values in the subrange defined by a bounding pair.
Don Stewart's avatar
Don Stewart committed
78
    range               :: (a,a) -> [a]
ross's avatar
ross committed
79
    -- | The position of a subscript in the subrange.
Don Stewart's avatar
Don Stewart committed
80
    index               :: (a,a) -> a -> Int
ross's avatar
ross committed
81
    -- | Like 'index', but without checking that the value is in range.
Don Stewart's avatar
Don Stewart committed
82
    unsafeIndex         :: (a,a) -> a -> Int
ross's avatar
ross committed
83 84
    -- | Returns 'True' the given subscript lies in the range defined
    -- the bounding pair.
Don Stewart's avatar
Don Stewart committed
85
    inRange             :: (a,a) -> a -> Bool
ross's avatar
ross committed
86
    -- | The size of the subrange defined by a bounding pair.
Don Stewart's avatar
Don Stewart committed
87
    rangeSize           :: (a,a) -> Int
ross's avatar
ross committed
88 89
    -- | like 'rangeSize', but without checking that the upper bound is
    -- in range.
90
    unsafeRangeSize     :: (a,a) -> Int
91

Don Stewart's avatar
Don Stewart committed
92
        -- Must specify one of index, unsafeIndex
93

94 95 96
        -- 'index' is typically over-ridden in instances, with essentially
        -- the same code, but using indexError instead of hopelessIndexError
        -- Reason: we have 'Show' at the instances
97
    {-# INLINE index #-}  -- See Note [Inlining index]
Jan Stolarek's avatar
Jan Stolarek committed
98
    index b i | inRange b i = unsafeIndex b i
99 100
              | otherwise   = hopelessIndexError

101
    unsafeIndex b i = index b i
102 103

    rangeSize b@(_l,h) | inRange b h = unsafeIndex b h + 1
Don Stewart's avatar
Don Stewart committed
104 105 106 107
                       | otherwise   = 0        -- This case is only here to
                                                -- check for an empty range
        -- NB: replacing (inRange b h) by (l <= h) fails for
        --     tuples.  E.g.  (1,2) <= (2,1) but the range is empty
108

109
    unsafeRangeSize b@(_l,h) = unsafeIndex b h + 1
110

111
{-
112
Note that the following is NOT right
Don Stewart's avatar
Don Stewart committed
113 114
        rangeSize (l,h) | l <= h    = index b h + 1
                        | otherwise = 0
115 116 117

Because it might be the case that l<h, but the range
is nevertheless empty.  Consider
Don Stewart's avatar
Don Stewart committed
118
        ((1,2),(2,1))
119 120
Here l<h, but the second index ranges from 2..1 and
hence is empty
121 122


123 124
Note [Inlining index]
~~~~~~~~~~~~~~~~~~~~~
Jan Stolarek's avatar
Jan Stolarek committed
125
We inline the 'index' operation,
126

Jan Stolarek's avatar
Jan Stolarek committed
127
 * Partly because it generates much faster code
128 129 130 131 132 133 134
   (although bigger); see Trac #1216

 * Partly because it exposes the bounds checks to the simplifier which
   might help a big.

If you make a per-instance index method, you may consider inlining it.

135 136 137 138
Note [Double bounds-checking of index values]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
When you index an array, a!x, there are two possible bounds checks we might make:

Jan Stolarek's avatar
Jan Stolarek committed
139
  (A) Check that (inRange (bounds a) x) holds.
140 141 142

      (A) is checked in the method for 'index'

Jan Stolarek's avatar
Jan Stolarek committed
143
  (B) Check that (index (bounds a) x) lies in the range 0..n,
144 145 146 147
      where n is the size of the underlying array

      (B) is checked in the top-level function (!), in safeIndex.

Jan Stolarek's avatar
Jan Stolarek committed
148
Of course it *should* be the case that (A) holds iff (B) holds, but that
149 150 151
is a property of the particular instances of index, bounds, and inRange,
so GHC cannot guarantee it.

Jan Stolarek's avatar
Jan Stolarek committed
152
 * If you do (A) and not (B), then you might get a seg-fault,
153 154 155 156 157 158 159 160 161 162
   by indexing at some bizarre location.  Trac #1610

 * If you do (B) but not (A), you may get no complaint when you index
   an array out of its semantic bounds.  Trac #2120

At various times we have had (A) and not (B), or (B) and not (A); both
led to complaints.  So now we implement *both* checks (Trac #2669).

For 1-d, 2-d, and 3-d arrays of Int we have specialised instances to avoid this.

163 164 165 166 167
Note [Out-of-bounds error messages]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The default method for 'index' generates hoplelessIndexError, because
Ix doesn't have Show as a superclass.  For particular base types we
can do better, so we override the default method for index.
168
-}
169 170

-- Abstract these errors from the relevant index functions so that
171 172
-- the guts of the function will be small enough to inline.

Simon Peyton Jones's avatar
Simon Peyton Jones committed
173
{-# NOINLINE indexError #-}
174 175 176 177
indexError :: Show a => (a,a) -> a -> String -> b
indexError rng i tp
  = error (showString "Ix{" . showString tp . showString "}.index: Index " .
           showParen True (showsPrec 0 i) .
Don Stewart's avatar
Don Stewart committed
178 179
           showString " out of range " $
           showParen True (showsPrec 0 rng) "")
180

181 182 183
hopelessIndexError :: Int -- Try to use 'indexError' instead!
hopelessIndexError = error "Error in array index"

184 185 186 187 188 189 190 191
----------------------------------------------------------------------
instance  Ix Char  where
    {-# INLINE range #-}
    range (m,n) = [m..n]

    {-# INLINE unsafeIndex #-}
    unsafeIndex (m,_n) i = fromEnum i - fromEnum m

192 193
    {-# INLINE index #-}  -- See Note [Out-of-bounds error messages]
                          -- and Note [Inlining index]
194
    index b i | inRange b i =  unsafeIndex b i
Don Stewart's avatar
Don Stewart committed
195
              | otherwise   =  indexError b i "Char"
196

Don Stewart's avatar
Don Stewart committed
197
    inRange (m,n) i     =  m <= i && i <= n
198 199 200 201

----------------------------------------------------------------------
instance  Ix Int  where
    {-# INLINE range #-}
Don Stewart's avatar
Don Stewart committed
202 203
        -- The INLINE stops the build in the RHS from getting inlined,
        -- so that callers can fuse with the result of range
204 205 206 207 208
    range (m,n) = [m..n]

    {-# INLINE unsafeIndex #-}
    unsafeIndex (m,_n) i = i - m

209 210
    {-# INLINE index #-}  -- See Note [Out-of-bounds error messages]
                          -- and Note [Inlining index]
211
    index b i | inRange b i =  unsafeIndex b i
Don Stewart's avatar
Don Stewart committed
212
              | otherwise   =  indexError b i "Int"
213 214

    {-# INLINE inRange #-}
215
    inRange (I# m,I# n) (I# i) =  isTrue# (m <=# i) && isTrue# (i <=# n)
216

217 218 219 220 221
instance Ix Word where
    range (m,n)         = [m..n]
    unsafeIndex (m,_) i = fromIntegral (i - m)
    inRange (m,n) i     = m <= i && i <= n

222 223 224 225 226 227 228 229
----------------------------------------------------------------------
instance  Ix Integer  where
    {-# INLINE range #-}
    range (m,n) = [m..n]

    {-# INLINE unsafeIndex #-}
    unsafeIndex (m,_n) i   = fromInteger (i - m)

230 231
    {-# INLINE index #-}  -- See Note [Out-of-bounds error messages]
                          -- and Note [Inlining index]
232
    index b i | inRange b i =  unsafeIndex b i
Don Stewart's avatar
Don Stewart committed
233
              | otherwise   =  indexError b i "Integer"
234

Don Stewart's avatar
Don Stewart committed
235
    inRange (m,n) i     =  m <= i && i <= n
236 237 238 239 240 241 242 243 244

----------------------------------------------------------------------
instance Ix Bool where -- as derived
    {-# INLINE range #-}
    range (m,n) = [m..n]

    {-# INLINE unsafeIndex #-}
    unsafeIndex (l,_) i = fromEnum i - fromEnum l

245 246
    {-# INLINE index #-}  -- See Note [Out-of-bounds error messages]
                          -- and Note [Inlining index]
247
    index b i | inRange b i =  unsafeIndex b i
Don Stewart's avatar
Don Stewart committed
248
              | otherwise   =  indexError b i "Bool"
249 250 251 252 253 254 255 256 257 258 259

    inRange (l,u) i = fromEnum i >= fromEnum l && fromEnum i <= fromEnum u

----------------------------------------------------------------------
instance Ix Ordering where -- as derived
    {-# INLINE range #-}
    range (m,n) = [m..n]

    {-# INLINE unsafeIndex #-}
    unsafeIndex (l,_) i = fromEnum i - fromEnum l

260 261
    {-# INLINE index #-}  -- See Note [Out-of-bounds error messages]
                          -- and Note [Inlining index]
262
    index b i | inRange b i =  unsafeIndex b i
Don Stewart's avatar
Don Stewart committed
263
              | otherwise   =  indexError b i "Ordering"
264 265 266 267 268 269 270 271 272 273 274

    inRange (l,u) i = fromEnum i >= fromEnum l && fromEnum i <= fromEnum u

----------------------------------------------------------------------
instance Ix () where
    {-# INLINE range #-}
    range   ((), ())    = [()]
    {-# INLINE unsafeIndex #-}
    unsafeIndex   ((), ()) () = 0
    {-# INLINE inRange #-}
    inRange ((), ()) () = True
275 276

    {-# INLINE index #-}  -- See Note [Inlining index]
277 278 279 280 281 282
    index b i = unsafeIndex b i

----------------------------------------------------------------------
instance (Ix a, Ix b) => Ix (a, b) where -- as derived
    {-# SPECIALISE instance Ix (Int,Int) #-}

283
    {-# INLINE range #-}
284 285 286
    range ((l1,l2),(u1,u2)) =
      [ (i1,i2) | i1 <- range (l1,u1), i2 <- range (l2,u2) ]

287
    {-# INLINE unsafeIndex #-}
288 289 290
    unsafeIndex ((l1,l2),(u1,u2)) (i1,i2) =
      unsafeIndex (l1,u1) i1 * unsafeRangeSize (l2,u2) + unsafeIndex (l2,u2) i2

291
    {-# INLINE inRange #-}
292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353
    inRange ((l1,l2),(u1,u2)) (i1,i2) =
      inRange (l1,u1) i1 && inRange (l2,u2) i2

    -- Default method for index

----------------------------------------------------------------------
instance  (Ix a1, Ix a2, Ix a3) => Ix (a1,a2,a3)  where
    {-# SPECIALISE instance Ix (Int,Int,Int) #-}

    range ((l1,l2,l3),(u1,u2,u3)) =
        [(i1,i2,i3) | i1 <- range (l1,u1),
                      i2 <- range (l2,u2),
                      i3 <- range (l3,u3)]

    unsafeIndex ((l1,l2,l3),(u1,u2,u3)) (i1,i2,i3) =
      unsafeIndex (l3,u3) i3 + unsafeRangeSize (l3,u3) * (
      unsafeIndex (l2,u2) i2 + unsafeRangeSize (l2,u2) * (
      unsafeIndex (l1,u1) i1))

    inRange ((l1,l2,l3),(u1,u2,u3)) (i1,i2,i3) =
      inRange (l1,u1) i1 && inRange (l2,u2) i2 &&
      inRange (l3,u3) i3

    -- Default method for index

----------------------------------------------------------------------
instance  (Ix a1, Ix a2, Ix a3, Ix a4) => Ix (a1,a2,a3,a4)  where
    range ((l1,l2,l3,l4),(u1,u2,u3,u4)) =
      [(i1,i2,i3,i4) | i1 <- range (l1,u1),
                       i2 <- range (l2,u2),
                       i3 <- range (l3,u3),
                       i4 <- range (l4,u4)]

    unsafeIndex ((l1,l2,l3,l4),(u1,u2,u3,u4)) (i1,i2,i3,i4) =
      unsafeIndex (l4,u4) i4 + unsafeRangeSize (l4,u4) * (
      unsafeIndex (l3,u3) i3 + unsafeRangeSize (l3,u3) * (
      unsafeIndex (l2,u2) i2 + unsafeRangeSize (l2,u2) * (
      unsafeIndex (l1,u1) i1)))

    inRange ((l1,l2,l3,l4),(u1,u2,u3,u4)) (i1,i2,i3,i4) =
      inRange (l1,u1) i1 && inRange (l2,u2) i2 &&
      inRange (l3,u3) i3 && inRange (l4,u4) i4

    -- Default method for index

instance  (Ix a1, Ix a2, Ix a3, Ix a4, Ix a5) => Ix (a1,a2,a3,a4,a5)  where
    range ((l1,l2,l3,l4,l5),(u1,u2,u3,u4,u5)) =
      [(i1,i2,i3,i4,i5) | i1 <- range (l1,u1),
                          i2 <- range (l2,u2),
                          i3 <- range (l3,u3),
                          i4 <- range (l4,u4),
                          i5 <- range (l5,u5)]

    unsafeIndex ((l1,l2,l3,l4,l5),(u1,u2,u3,u4,u5)) (i1,i2,i3,i4,i5) =
      unsafeIndex (l5,u5) i5 + unsafeRangeSize (l5,u5) * (
      unsafeIndex (l4,u4) i4 + unsafeRangeSize (l4,u4) * (
      unsafeIndex (l3,u3) i3 + unsafeRangeSize (l3,u3) * (
      unsafeIndex (l2,u2) i2 + unsafeRangeSize (l2,u2) * (
      unsafeIndex (l1,u1) i1))))

    inRange ((l1,l2,l3,l4,l5),(u1,u2,u3,u4,u5)) (i1,i2,i3,i4,i5) =
      inRange (l1,u1) i1 && inRange (l2,u2) i2 &&
Jan Stolarek's avatar
Jan Stolarek committed
354
      inRange (l3,u3) i3 && inRange (l4,u4) i4 &&
355 356
      inRange (l5,u5) i5

357
    -- Default method for index
358

ross's avatar
ross committed
359 360
-- | The type of immutable non-strict (boxed) arrays
-- with indices in @i@ and elements in @e@.
361
data Array i e
362 363 364 365 366 367
   = Array            !i         -- the lower bound, l
                      !i         -- the upper bound, u
       {-# UNPACK #-} !Int       -- A cache of (rangeSize (l,u))
                                 -- used to make sure an index is
                                 -- really in range
                      (Array# e) -- The actual elements
368 369 370 371 372 373

-- | Mutable, boxed, non-strict arrays in the 'ST' monad.  The type
-- arguments are as follows:
--
--  * @s@: the state variable argument for the 'ST' type
--
ross's avatar
ross committed
374
--  * @i@: the index type of the array (should be an instance of 'Ix')
375 376 377
--
--  * @e@: the element type of the array.
--
378
data STArray s i e
379 380 381
  = STArray           !i               -- the lower bound, l
                      !i               -- the upper bound, u
      {-# UNPACK #-}  !Int             -- A cache of (rangeSize (l,u))
382 383 384
                                       -- used to make sure an index is
                                       -- really in range
                   (MutableArray# s e) -- The actual elements
Don Stewart's avatar
Don Stewart committed
385 386
        -- No Ix context for STArray.  They are stupid,
        -- and force an Ix context on the equality instance.
387

388 389 390 391
-- Index types should have nominal role, because of Ix class. See also #9220.
type role Array nominal representational
type role STArray nominal nominal representational

392 393
-- Just pointer equality on mutable arrays:
instance Eq (STArray s i e) where
394
    STArray _ _ _ arr1# == STArray _ _ _ arr2# =
395
        isTrue# (sameMutableArray# arr1# arr2#)
396

397 398
----------------------------------------------------------------------
-- Operations on immutable arrays
399 400 401 402 403

{-# NOINLINE arrEleBottom #-}
arrEleBottom :: a
arrEleBottom = error "(Array.!): undefined array element"

ross's avatar
ross committed
404 405 406 407
-- | Construct an array with the specified bounds and containing values
-- for given indices within these bounds.
--
-- The array is undefined (i.e. bottom) if any index in the list is
408
-- out of bounds.  The Haskell 2010 Report further specifies that if any
ross's avatar
ross committed
409 410 411 412 413 414 415
-- two associations in the list have the same index, the value at that
-- index is undefined (i.e. bottom).  However in GHC's implementation,
-- the value at such an index is the value part of the last association
-- with that index in the list.
--
-- Because the indices must be checked for these errors, 'array' is
-- strict in the bounds argument and in the indices of the association
416
-- list, but non-strict in the values.  Thus, recurrences such as the
ross's avatar
ross committed
417 418 419 420 421 422 423 424 425 426 427 428
-- following are possible:
--
-- > a = array (1,100) ((1,1) : [(i, i * a!(i-1)) | i <- [2..100]])
--
-- Not every index within the bounds of the array need appear in the
-- association list, but the values associated with indices that do not
-- appear will be undefined (i.e. bottom).
--
-- If, in any dimension, the lower bound is greater than the upper bound,
-- then the array is legal, but empty.  Indexing an empty array always
-- gives an array-bounds error, but 'bounds' still yields the bounds
-- with which the array was constructed.
429
{-# INLINE array #-}
ross's avatar
ross committed
430
array :: Ix i
Don Stewart's avatar
Don Stewart committed
431 432 433 434 435 436 437 438 439 440 441 442
        => (i,i)        -- ^ a pair of /bounds/, each of the index type
                        -- of the array.  These bounds are the lowest and
                        -- highest indices in the array, in that order.
                        -- For example, a one-origin vector of length
                        -- '10' has bounds '(1,10)', and a one-origin '10'
                        -- by '10' matrix has bounds '((1,1),(10,10))'.
        -> [(i, e)]     -- ^ a list of /associations/ of the form
                        -- (/index/, /value/).  Typically, this list will
                        -- be expressed as a comprehension.  An
                        -- association '(i, x)' defines the value of
                        -- the array at index 'i' to be 'x'.
        -> Array i e
443 444 445 446
array (l,u) ies
    = let n = safeRangeSize (l,u)
      in unsafeArray' (l,u) n
                      [(safeIndex (l,u) n i, e) | (i, e) <- ies]
447 448 449

{-# INLINE unsafeArray #-}
unsafeArray :: Ix i => (i,i) -> [(Int, e)] -> Array i e
450 451 452
unsafeArray b ies = unsafeArray' b (rangeSize b) ies

{-# INLINE unsafeArray' #-}
453
unsafeArray' :: (i,i) -> Int -> [(Int, e)] -> Array i e
454 455 456 457
unsafeArray' (l,u) n@(I# n#) ies = runST (ST $ \s1# ->
    case newArray# n# arrEleBottom s1# of
        (# s2#, marr# #) ->
            foldr (fill marr#) (done l u n marr#) ies s2#)
458 459 460

{-# INLINE fill #-}
fill :: MutableArray# s e -> (Int, e) -> STRep s a -> STRep s a
Jan Stolarek's avatar
Jan Stolarek committed
461 462 463 464 465
-- NB: put the \s after the "=" so that 'fill'
--     inlines when applied to three args
fill marr# (I# i#, e) next
 = \s1# -> case writeArray# marr# i# e s1# of
             s2# -> next s2#
466 467

{-# INLINE done #-}
468
done :: i -> i -> Int -> MutableArray# s e -> STRep s (Array i e)
469
-- See NB on 'fill'
470 471
-- Make sure it is strict in 'n'
done l u n@(I# _) marr#
472 473
  = \s1# -> case unsafeFreezeArray# marr# s1# of
              (# s2#, arr# #) -> (# s2#, Array l u n arr# #)
474

ross's avatar
ross committed
475 476
-- | Construct an array from a pair of bounds and a list of values in
-- index order.
477 478 479
{-# INLINE listArray #-}
listArray :: Ix i => (i,i) -> [e] -> Array i e
listArray (l,u) es = runST (ST $ \s1# ->
480
    case safeRangeSize (l,u)            of { n@(I# n#) ->
481
    case newArray# n# arrEleBottom s1#  of { (# s2#, marr# #) ->
David Feuer's avatar
David Feuer committed
482 483 484 485 486 487 488 489 490 491 492
      let
        go y r = \ i# s3# ->
            case writeArray# marr# i# y s3# of
              s4# -> if (isTrue# (i# ==# n# -# 1#))
                     then s4#
                     else r (i# +# 1#) s4#
      in
        done l u n marr# (
          if n == 0
          then s2#
          else foldr go (\_ s# -> s#) es 0# s2#)}})
493

ross's avatar
ross committed
494
-- | The value at the given index in an array.
495 496
{-# INLINE (!) #-}
(!) :: Ix i => Array i e -> i -> e
497 498 499 500 501
arr@(Array l u n _) ! i = unsafeAt arr $ safeIndex (l,u) n i

{-# INLINE safeRangeSize #-}
safeRangeSize :: Ix i => (i, i) -> Int
safeRangeSize (l,u) = let r = rangeSize (l, u)
502
                      in if r < 0 then negRange
503 504
                                  else r

505
-- Don't inline this error message everywhere!!
506
negRange :: Int   -- Uninformative, but Ix does not provide Show
507 508
negRange = error "Negative range size"

509 510 511
{-# INLINE[1] safeIndex #-}
-- See Note [Double bounds-checking of index values]
-- Inline *after* (!) so the rules can fire
512
-- Make sure it is strict in n
513
safeIndex :: Ix i => (i, i) -> Int -> i -> Int
514
safeIndex (l,u) n@(I# _) i
515 516 517 518
  | (0 <= i') && (i' < n) = i'
  | otherwise             = badSafeIndex i' n
  where
    i' = index (l,u) i
519

520 521 522 523 524 525 526 527 528 529
-- See Note [Double bounds-checking of index values]
{-# RULES
"safeIndex/I"       safeIndex = lessSafeIndex :: (Int,Int) -> Int -> Int -> Int
"safeIndex/(I,I)"   safeIndex = lessSafeIndex :: ((Int,Int),(Int,Int)) -> Int -> (Int,Int) -> Int
"safeIndex/(I,I,I)" safeIndex = lessSafeIndex :: ((Int,Int,Int),(Int,Int,Int)) -> Int -> (Int,Int,Int) -> Int
  #-}

lessSafeIndex :: Ix i => (i, i) -> Int -> i -> Int
-- See Note [Double bounds-checking of index values]
-- Do only (A), the semantic check
Jan Stolarek's avatar
Jan Stolarek committed
530
lessSafeIndex (l,u) _ i = index (l,u) i
531

532 533 534 535
-- Don't inline this long error message everywhere!!
badSafeIndex :: Int -> Int -> Int
badSafeIndex i' n = error ("Error in array index; " ++ show i' ++
                        " not in range [0.." ++ show n ++ ")")
536 537

{-# INLINE unsafeAt #-}
538
unsafeAt :: Array i e -> Int -> e
539
unsafeAt (Array _ _ _ arr#) (I# i#) =
540 541
    case indexArray# arr# i# of (# e #) -> e

ross's avatar
ross committed
542
-- | The bounds with which an array was constructed.
543
{-# INLINE bounds #-}
544
bounds :: Array i e -> (i,i)
545 546 547 548
bounds (Array l u _ _) = (l,u)

-- | The number of elements in the array.
{-# INLINE numElements #-}
549
numElements :: Array i e -> Int
550
numElements (Array _ _ n _) = n
551

ross's avatar
ross committed
552
-- | The list of indices of an array in ascending order.
553 554
{-# INLINE indices #-}
indices :: Ix i => Array i e -> [i]
555
indices (Array l u _ _) = range (l,u)
556

ross's avatar
ross committed
557
-- | The list of elements of an array in index order.
558
{-# INLINE elems #-}
559
elems :: Array i e -> [e]
Ian Lynagh's avatar
Ian Lynagh committed
560
elems arr@(Array _ _ n _) =
561
    [unsafeAt arr i | i <- [0 .. n - 1]]
562

563 564
-- | A right fold over the elements
{-# INLINABLE foldrElems #-}
565
foldrElems :: (a -> b -> b) -> b -> Array i a -> b
566 567 568 569 570 571 572 573
foldrElems f b0 = \ arr@(Array _ _ n _) ->
  let
    go i | i == n    = b0
         | otherwise = f (unsafeAt arr i) (go (i+1))
  in go 0

-- | A left fold over the elements
{-# INLINABLE foldlElems #-}
574
foldlElems :: (b -> a -> b) -> b -> Array i a -> b
575 576 577 578 579 580 581 582
foldlElems f b0 = \ arr@(Array _ _ n _) ->
  let
    go i | i == (-1) = b0
         | otherwise = f (go (i-1)) (unsafeAt arr i)
  in go (n-1)

-- | A strict right fold over the elements
{-# INLINABLE foldrElems' #-}
583
foldrElems' :: (a -> b -> b) -> b -> Array i a -> b
584 585 586 587 588 589 590 591
foldrElems' f b0 = \ arr@(Array _ _ n _) ->
  let
    go i a | i == (-1) = a
           | otherwise = go (i-1) (f (unsafeAt arr i) $! a)
  in go (n-1) b0

-- | A strict left fold over the elements
{-# INLINABLE foldlElems' #-}
592
foldlElems' :: (b -> a -> b) -> b -> Array i a -> b
593 594 595 596 597 598 599 600
foldlElems' f b0 = \ arr@(Array _ _ n _) ->
  let
    go i a | i == n    = a
           | otherwise = go (i+1) (a `seq` f a (unsafeAt arr i))
  in go 0 b0

-- | A left fold over the elements with no starting value
{-# INLINABLE foldl1Elems #-}
601
foldl1Elems :: (a -> a -> a) -> Array i a -> a
602 603 604 605 606 607 608 609 610
foldl1Elems f = \ arr@(Array _ _ n _) ->
  let
    go i | i == 0    = unsafeAt arr 0
         | otherwise = f (go (i-1)) (unsafeAt arr i)
  in
    if n == 0 then error "foldl1: empty Array" else go (n-1)

-- | A right fold over the elements with no starting value
{-# INLINABLE foldr1Elems #-}
611
foldr1Elems :: (a -> a -> a) -> Array i a -> a
612 613 614 615 616 617 618
foldr1Elems f = \ arr@(Array _ _ n _) ->
  let
    go i | i == n-1  = unsafeAt arr i
         | otherwise = f (unsafeAt arr i) (go (i + 1))
  in
    if n == 0 then error "foldr1: empty Array" else go 0

ross's avatar
ross committed
619
-- | The list of associations of an array in index order.
620 621
{-# INLINE assocs #-}
assocs :: Ix i => Array i e -> [(i, e)]
622 623
assocs arr@(Array l u _ _) =
    [(i, arr ! i) | i <- range (l,u)]
624

Simon Marlow's avatar
Simon Marlow committed
625
-- | The 'accumArray' function deals with repeated indices in the association
ross's avatar
ross committed
626 627 628 629 630 631 632 633 634 635 636 637 638
-- list using an /accumulating function/ which combines the values of
-- associations with the same index.
-- For example, given a list of values of some index type, @hist@
-- produces a histogram of the number of occurrences of each index within
-- a specified range:
--
-- > hist :: (Ix a, Num b) => (a,a) -> [a] -> Array a b
-- > hist bnds is = accumArray (+) 0 bnds [(i, 1) | i<-is, inRange bnds i]
--
-- If the accumulating function is strict, then 'accumArray' is strict in
-- the values, as well as the indices, in the association list.  Thus,
-- unlike ordinary arrays built with 'array', accumulated arrays should
-- not in general be recursive.
639
{-# INLINE accumArray #-}
ross's avatar
ross committed
640
accumArray :: Ix i
Don Stewart's avatar
Don Stewart committed
641 642 643 644 645
        => (e -> a -> e)        -- ^ accumulating function
        -> e                    -- ^ initial value
        -> (i,i)                -- ^ bounds of the array
        -> [(i, a)]             -- ^ association list
        -> Array i e
Ian Lynagh's avatar
Ian Lynagh committed
646
accumArray f initial (l,u) ies =
647
    let n = safeRangeSize (l,u)
Ian Lynagh's avatar
Ian Lynagh committed
648
    in unsafeAccumArray' f initial (l,u) n
649
                         [(safeIndex (l,u) n i, e) | (i, e) <- ies]
650 651 652

{-# INLINE unsafeAccumArray #-}
unsafeAccumArray :: Ix i => (e -> a -> e) -> e -> (i,i) -> [(Int, a)] -> Array i e
Ian Lynagh's avatar
Ian Lynagh committed
653
unsafeAccumArray f initial b ies = unsafeAccumArray' f initial b (rangeSize b) ies
654 655

{-# INLINE unsafeAccumArray' #-}
656
unsafeAccumArray' :: (e -> a -> e) -> e -> (i,i) -> Int -> [(Int, a)] -> Array i e
Ian Lynagh's avatar
Ian Lynagh committed
657 658
unsafeAccumArray' f initial (l,u) n@(I# n#) ies = runST (ST $ \s1# ->
    case newArray# n# initial s1#          of { (# s2#, marr# #) ->
659
    foldr (adjust f marr#) (done l u n marr#) ies s2# })
660 661 662

{-# INLINE adjust #-}
adjust :: (e -> a -> e) -> MutableArray# s e -> (Int, a) -> STRep s b -> STRep s b
663 664 665
-- See NB on 'fill'
adjust f marr# (I# i#, new) next
  = \s1# -> case readArray# marr# i# s1# of
666 667 668
                (# s2#, old #) ->
                    case writeArray# marr# i# (f old new) s2# of
                        s3# -> next s3#
669

ross's avatar
ross committed
670 671 672 673 674 675 676 677 678
-- | Constructs an array identical to the first argument except that it has
-- been updated by the associations in the right argument.
-- For example, if @m@ is a 1-origin, @n@ by @n@ matrix, then
--
-- > m//[((i,i), 0) | i <- [1..n]]
--
-- is the same matrix, except with the diagonal zeroed.
--
-- Repeated indices in the association list are handled as for 'array':
679
-- Haskell 2010 specifies that the resulting array is undefined (i.e. bottom),
ross's avatar
ross committed
680
-- but GHC's implementation uses the last association for each index.
681 682
{-# INLINE (//) #-}
(//) :: Ix i => Array i e -> [(i, e)] -> Array i e
683 684
arr@(Array l u n _) // ies =
    unsafeReplace arr [(safeIndex (l,u) n i, e) | (i, e) <- ies]
685 686

{-# INLINE unsafeReplace #-}
687
unsafeReplace :: Array i e -> [(Int, e)] -> Array i e
688 689 690
unsafeReplace arr ies = runST (do
    STArray l u n marr# <- thawSTArray arr
    ST (foldr (fill marr#) (done l u n marr#) ies))
691

ross's avatar
ross committed
692 693 694 695 696 697
-- | @'accum' f@ takes an array and an association list and accumulates
-- pairs from the list into the array with the accumulating function @f@.
-- Thus 'accumArray' can be defined using 'accum':
--
-- > accumArray f z b = accum f (array b [(i, z) | i <- range b])
--
698 699
{-# INLINE accum #-}
accum :: Ix i => (e -> a -> e) -> Array i e -> [(i, a)] -> Array i e
700 701
accum f arr@(Array l u n _) ies =
    unsafeAccum f arr [(safeIndex (l,u) n i, e) | (i, e) <- ies]
702 703

{-# INLINE unsafeAccum #-}
704
unsafeAccum :: (e -> a -> e) -> Array i e -> [(Int, a)] -> Array i e
705 706 707
unsafeAccum f arr ies = runST (do
    STArray l u n marr# <- thawSTArray arr
    ST (foldr (adjust f marr#) (done l u n marr#) ies))
708

Simon Peyton Jones's avatar
Simon Peyton Jones committed
709
{-# INLINE [1] amap #-}  -- See Note [amap]
710
amap :: (a -> b) -> Array i a -> Array i b
711 712 713 714 715 716 717 718
amap f arr@(Array l u n@(I# n#) _) = runST (ST $ \s1# ->
    case newArray# n# arrEleBottom s1# of
        (# s2#, marr# #) ->
          let go i s#
                | i == n    = done l u n marr# s#
                | otherwise = fill marr# (i, f (unsafeAt arr i)) (go (i+1)) s#
          in go 0 s2# )

Simon Peyton Jones's avatar
Simon Peyton Jones committed
719 720
{- Note [amap]
~~~~~~~~~~~~~~
721 722 723 724 725 726 727 728
amap was originally defined like this:

 amap f arr@(Array l u n _) =
     unsafeArray' (l,u) n [(i, f (unsafeAt arr i)) | i <- [0 .. n - 1]]

There are two problems:

1. The enumFromTo implementation produces (spurious) code for the impossible
Simon Peyton Jones's avatar
Simon Peyton Jones committed
729
   case of n<0 that ends up duplicating the array freezing code.
730

Simon Peyton Jones's avatar
Simon Peyton Jones committed
731 732 733 734
2. This implementation relies on list fusion for efficiency. In order
   to implement the "amap/coerce" rule, we need to delay inlining amap
   until simplifier phase 1, which is when the eftIntList rule kicks
   in and makes that impossible.  (c.f. Trac #8767)
735 736 737 738 739 740 741
-}


-- See Breitner, Eisenberg, Peyton Jones, and Weirich, "Safe Zero-cost
-- Coercions for Haskell", section 6.5:
--   http://research.microsoft.com/en-us/um/people/simonpj/papers/ext-f/coercible.pdf
{-# RULES
Simon Peyton Jones's avatar
Simon Peyton Jones committed
742
"amap/coerce" amap coerce = coerce  -- See Note [amap]
743 744 745 746 747 748
 #-}

-- Second functor law:
{-# RULES
"amap/amap" forall f g a . amap f (amap g a) = amap (f . g) a
 #-}
749

ross's avatar
ross committed
750 751 752 753 754 755
-- | 'ixmap' allows for transformations on array indices.
-- It may be thought of as providing function composition on the right
-- with the mapping that the original array embodies.
--
-- A similar transformation of array values may be achieved using 'fmap'
-- from the 'Array' instance of the 'Functor' class.
756 757 758
{-# INLINE ixmap #-}
ixmap :: (Ix i, Ix j) => (i,i) -> (i -> j) -> Array j e -> Array i e
ixmap (l,u) f arr =
759
    array (l,u) [(i, arr ! f i) | i <- range (l,u)]
760 761 762

{-# INLINE eqArray #-}
eqArray :: (Ix i, Eq e) => Array i e -> Array i e -> Bool
763 764
eqArray arr1@(Array l1 u1 n1 _) arr2@(Array l2 u2 n2 _) =
    if n1 == 0 then n2 == 0 else
765
    l1 == l2 && u1 == u2 &&
766
    and [unsafeAt arr1 i == unsafeAt arr2 i | i <- [0 .. n1 - 1]]
767

768
{-# INLINE [1] cmpArray #-}
769 770 771 772 773
cmpArray :: (Ix i, Ord e) => Array i e -> Array i e -> Ordering
cmpArray arr1 arr2 = compare (assocs arr1) (assocs arr2)

{-# INLINE cmpIntArray #-}
cmpIntArray :: Ord e => Array Int e -> Array Int e -> Ordering
774 775 776 777 778 779 780 781
cmpIntArray arr1@(Array l1 u1 n1 _) arr2@(Array l2 u2 n2 _) =
    if n1 == 0 then
        if n2 == 0 then EQ else LT
    else if n2 == 0 then GT
    else case compare l1 l2 of
             EQ    -> foldr cmp (compare u1 u2) [0 .. (n1 `min` n2) - 1]
             other -> other
  where
782 783 784 785 786 787
    cmp i rest = case compare (unsafeAt arr1 i) (unsafeAt arr2 i) of
        EQ    -> rest
        other -> other

{-# RULES "cmpArray/Int" cmpArray = cmpIntArray #-}

788 789
----------------------------------------------------------------------
-- Array instances
790

791
instance Functor (Array i) where
792 793 794 795 796 797 798 799 800 801
    fmap = amap

instance (Ix i, Eq e) => Eq (Array i e) where
    (==) = eqArray

instance (Ix i, Ord e) => Ord (Array i e) where
    compare = cmpArray

instance (Ix a, Show a, Show b) => Show (Array a b) where
    showsPrec p a =
802
        showParen (p > appPrec) $
803
        showString "array " .
804
        showsPrec appPrec1 (bounds a) .
805
        showChar ' ' .
806
        showsPrec appPrec1 (assocs a)
Don Stewart's avatar
Don Stewart committed
807
        -- Precedence of 'array' is the precedence of application
808

809
-- The Read instance is in GHC.Read
810

811 812
----------------------------------------------------------------------
-- Operations on mutable arrays
813

814
{-
815 816 817 818 819 820 821 822 823 824 825 826
Idle ADR question: What's the tradeoff here between flattening these
datatypes into @STArray ix ix (MutableArray# s elt)@ and using
it as is?  As I see it, the former uses slightly less heap and
provides faster access to the individual parts of the bounds while the
code used has the benefit of providing a ready-made @(lo, hi)@ pair as
required by many array-related functions.  Which wins? Is the
difference significant (probably not).

Idle AJG answer: When I looked at the outputted code (though it was 2
years ago) it seems like you often needed the tuple, and we build
it frequently. Now we've got the overloading specialiser things
might be different, though.
827
-}
828 829 830

{-# INLINE newSTArray #-}
newSTArray :: Ix i => (i,i) -> e -> ST s (STArray s i e)
Ian Lynagh's avatar
Ian Lynagh committed
831
newSTArray (l,u) initial = ST $ \s1# ->
832
    case safeRangeSize (l,u)            of { n@(I# n#) ->
Ian Lynagh's avatar
Ian Lynagh committed
833
    case newArray# n# initial s1#       of { (# s2#, marr# #) ->
834
    (# s2#, STArray l u n marr# #) }}
835 836

{-# INLINE boundsSTArray #-}
Jan Stolarek's avatar
Jan Stolarek committed
837
boundsSTArray :: STArray s i e -> (i,i)
838 839 840 841 842
boundsSTArray (STArray l u _ _) = (l,u)

{-# INLINE numElementsSTArray #-}
numElementsSTArray :: STArray s i e -> Int
numElementsSTArray (STArray _ _ n _) = n
843 844 845

{-# INLINE readSTArray #-}
readSTArray :: Ix i => STArray s i e -> i -> ST s e
846 847
readSTArray marr@(STArray l u n _) i =
    unsafeReadSTArray marr (safeIndex (l,u) n i)
848 849

{-# INLINE unsafeReadSTArray #-}
850
unsafeReadSTArray :: STArray s i e -> Int -> ST s e
851 852
unsafeReadSTArray (STArray _ _ _ marr#) (I# i#)
    = ST $ \s1# -> readArray# marr# i# s1#
853 854

{-# INLINE writeSTArray #-}
Jan Stolarek's avatar
Jan Stolarek committed
855
writeSTArray :: Ix i => STArray s i e -> i -> e -> ST s ()
856 857
writeSTArray marr@(STArray l u n _) i e =
    unsafeWriteSTArray marr (safeIndex (l,u) n i) e
858 859

{-# INLINE unsafeWriteSTArray #-}
860
unsafeWriteSTArray :: STArray s i e -> Int -> e -> ST s ()
861 862 863
unsafeWriteSTArray (STArray _ _ _ marr#) (I# i#) e = ST $ \s1# ->
    case writeArray# marr# i# e s1# of
        s2# -> (# s2#, () #)
864

865 866
----------------------------------------------------------------------
-- Moving between mutable and immutable
867

868
freezeSTArray :: STArray s i e -> ST s (Array i e)
869
freezeSTArray (STArray l u n@(I# n#) marr#) = ST $ \s1# ->
870
    case newArray# n# arrEleBottom s1#  of { (# s2#, marr'# #) ->
871
    let copy i# s3# | isTrue# (i# ==# n#) = s3#
872 873 874 875 876 877
                    | otherwise =
            case readArray# marr# i# s3# of { (# s4#, e #) ->
            case writeArray# marr'# i# e s4# of { s5# ->
            copy (i# +# 1#) s5# }} in
    case copy 0# s2#                    of { s3# ->
    case unsafeFreezeArray# marr'# s3#  of { (# s4#, arr# #) ->
878
    (# s4#, Array l u n arr# #) }}}
879 880

{-# INLINE unsafeFreezeSTArray #-}
881
unsafeFreezeSTArray :: STArray s i e -> ST s (Array i e)
882
unsafeFreezeSTArray (STArray l u n marr#) = ST $ \s1# ->
883
    case unsafeFreezeArray# marr# s1#   of { (# s2#, arr# #) ->
884
    (# s2#, Array l u n arr# #) }
885

886
thawSTArray :: Array i e -> ST s (STArray s i e)
887
thawSTArray (Array l u n@(I# n#) arr#) = ST $ \s1# ->
888
    case newArray# n# arrEleBottom s1#  of { (# s2#, marr# #) ->
889
    let copy i# s3# | isTrue# (i# ==# n#) = s3#
890 891 892 893 894
                    | otherwise =
            case indexArray# arr# i#    of { (# e #) ->
            case writeArray# marr# i# e s3# of { s4# ->
            copy (i# +# 1#) s4# }} in
    case copy 0# s2#                    of { s3# ->
895
    (# s3#, STArray l u n marr# #) }}
896 897

{-# INLINE unsafeThawSTArray #-}
898
unsafeThawSTArray :: Array i e -> ST s (STArray s i e)
899
unsafeThawSTArray (Array l u n arr#) = ST $ \s1# ->
900
    case unsafeThawArray# arr# s1#      of { (# s2#, marr# #) ->
901
    (# s2#, STArray l u n marr# #) }