ghc/lib/std/PrelArr.lhs

   1 %
   2 % (c) The AQUA Project, Glasgow University, 1994-1996
   3 %
   4 \section[PrelArr]{Module @PrelArr@}
   5
   6 Array implementation, @PrelArr@ exports the basic array
   7 types and operations.
   8
   9 For byte-arrays see @PrelByteArr@.
  10
  11 \begin{code}
  12 {-# OPTIONS -fno-implicit-prelude #-}
  13
  14 module PrelArr where
  15
  16 import {-# SOURCE #-} PrelErr ( error )
  17 import Ix
  18 import PrelList (foldl)
  19 import PrelST
  20 import PrelBase
  21 import PrelAddr
  22 import PrelGHC
  23 import PrelShow
  24
  25 infixl 9  !, //
  26
  27 default ()
  28 \end{code}
  29
  30 \begin{code}
  31 {-# SPECIALISE array :: (Int,Int) -> [(Int,b)] -> Array Int b #-}
  32 array                 :: (Ix a) => (a,a) -> [(a,b)] -> Array a b
  33
  34 {-# SPECIALISE (!) :: Array Int b -> Int -> b #-}
  35 (!)                   :: (Ix a) => Array a b -> a -> b
  36
  37 {-# SPECIALISE (//) :: Array Int b -> [(Int,b)] -> Array Int b #-}
  38 (//)                  :: (Ix a) => Array a b -> [(a,b)] -> Array a b
  39
  40 {-# SPECIALISE accum  :: (b -> c -> b) -> Array Int b -> [(Int,c)] -> Array Int b #-}
  41 accum                 :: (Ix a) => (b -> c -> b) -> Array a b -> [(a,c)] -> Array a b
  42
  43 {-# SPECIALISE accumArray :: (b -> c -> b) -> b -> (Int,Int) -> [(Int,c)] -> Array Int b #-}
  44 accumArray            :: (Ix a) => (b -> c -> b) -> b -> (a,a) -> [(a,c)] -> Array a b
  45
  46 bounds                :: (Ix a) => Array a b -> (a,a)
  47 assocs                :: (Ix a) => Array a b -> [(a,b)]
  48 indices               :: (Ix a) => Array a b -> [a]
  49 \end{code}
  50
  51
  52 %*********************************************************
  53 %*                                                      *
  54 \subsection{The @Array@ types}
  55 %*                                                      *
  56 %*********************************************************
  57
  58 \begin{code}
  59 type IPr = (Int, Int)
  60
  61 data Ix ix => Array     ix elt = Array   ix ix (Array# elt)
  62 data Ix ix => STArray s ix elt = STArray ix ix (MutableArray# s elt)
  63
  64
  65 data STRef s a = STRef (MutVar# s a)
  66
  67 instance Eq (STRef s a) where
  68         STRef v1# == STRef v2#
  69                 = sameMutVar# v1# v2#
  70
  71 -- just pointer equality on arrays:
  72 instance Eq (STArray s ix elt) where
  73         STArray _ _ arr1# == STArray _ _ arr2#
  74                 = sameMutableArray# arr1# arr2#
  75 \end{code}
  76
  77 %*********************************************************
  78 %*                                                      *
  79 \subsection{Operations on mutable variables}
  80 %*                                                      *
  81 %*********************************************************
  82
  83 \begin{code}
  84 newSTRef   :: a -> ST s (STRef s a)
  85 readSTRef  :: STRef s a -> ST s a
  86 writeSTRef :: STRef s a -> a -> ST s ()
  87
  88 newSTRef init = ST $ \ s# ->
  89     case (newMutVar# init s#)     of { (# s2#, var# #) ->
  90     (# s2#, STRef var# #) }
  91
  92 readSTRef (STRef var#) = ST $ \ s# -> readMutVar# var# s#
  93
  94 writeSTRef (STRef var#) val = ST $ \ s# ->
  95     case writeMutVar# var# val s# of { s2# ->
  96     (# s2#, () #) }
  97 \end{code}
  98
  99 %*********************************************************
 100 %*                                                      *
 101 \subsection{Operations on immutable arrays}
 102 %*                                                      *
 103 %*********************************************************
 104
 105 "array", "!" and "bounds" are basic; the rest can be defined in terms of them
 106
 107 \begin{code}
 108 {-# INLINE bounds #-}
 109 bounds (Array l u _)  = (l,u)
 110
 111 {-# INLINE assocs #-}   -- Want to fuse the list comprehension
 112 assocs a              =  [(i, a!i) | i <- indices a]
 113
 114 {-# INLINE indices #-}
 115 indices               =  range . bounds
 116
 117 {-# SPECIALISE amap :: (b -> c) -> Array Int b -> Array Int c #-}
 118 amap                  :: (Ix a) => (b -> c) -> Array a b -> Array a c
 119 amap f a              =  array b [(i, f (a!i)) | i <- range b]
 120                          where b = bounds a
 121
 122 (Array l u arr#) ! i
 123   = let n# = case (index (l,u) i) of { I# x -> x } -- index fails if out of range
 124     in
 125     case (indexArray# arr# n#) of
 126       (# v #) -> v
 127
 128 {-# INLINE array #-}
 129 array ixs ivs
 130   = case rangeSize ixs                          of { I# n ->
 131     runST ( ST $ \ s1 ->
 132         case newArray# n arrEleBottom s1        of { (# s2, marr #) ->
 133         foldr (fill ixs marr) (done ixs marr) ivs s2
 134     })}
 135
 136 fill :: Ix ix => (ix,ix)  -> MutableArray# s elt
 137               -> (ix,elt) -> STRep s a -> STRep s a
 138 {-# INLINE fill #-}
 139 fill ixs marr (i,v) next = \s1 -> case index ixs i      of { I# n ->
 140                                   case writeArray# marr n v s1  of { s2 ->
 141                                   next s2 }}
 142
 143 done :: Ix ix => (ix,ix) -> MutableArray# s elt
 144               -> STRep s (Array ix elt)
 145 {-# INLINE done #-}
 146 done (l,u) marr = \s1 ->
 147    case unsafeFreezeArray# marr s1 of { (# s2, arr #) ->
 148    (# s2, Array l u arr #) }
 149
 150 arrEleBottom :: a
 151 arrEleBottom = error "(Array.!): undefined array element"
 152
 153
 154 -----------------------------------------------------------------------
 155 -- These also go better with magic: (//), accum, accumArray
 156 -- *** NB *** We INLINE them all so that their foldr's get to the call site
 157
 158 {-# INLINE (//) #-}
 159 old_array // ivs
 160   = runST (do
 161         -- copy the old array:
 162         arr <- thawSTArray old_array
 163         -- now write the new elements into the new array:
 164         fill_it_in arr ivs
 165         freezeSTArray arr
 166     )
 167
 168 fill_it_in :: Ix ix => STArray s ix elt -> [(ix, elt)] -> ST s ()
 169 {-# INLINE fill_it_in #-}
 170 fill_it_in arr lst = foldr (fill_one_in arr) (return ()) lst
 171          -- **** STRICT **** (but that's OK...)
 172
 173 fill_one_in arr (i, v) rst = writeSTArray arr i v >> rst
 174
 175 zap_with_f :: Ix ix => (elt -> elt2 -> elt) -> STArray s ix elt -> [(ix,elt2)] -> ST s ()
 176 -- zap_with_f: reads an elem out first, then uses "f" on that and the new value
 177 {-# INLINE zap_with_f #-}
 178
 179 zap_with_f f arr lst
 180   = foldr (zap_one f arr) (return ()) lst
 181
 182 zap_one f arr (i, new_v) rst = do
 183         old_v <- readSTArray arr i
 184         writeSTArray arr i (f old_v new_v)
 185         rst
 186
 187 {-# INLINE accum #-}
 188 accum f old_array ivs
 189   = runST (do
 190         -- copy the old array:
 191         arr <- thawSTArray old_array
 192         -- now zap the elements in question with "f":
 193         zap_with_f f arr ivs
 194         freezeSTArray arr
 195     )
 196
 197 {-# INLINE accumArray #-}
 198 accumArray f zero ixs ivs
 199   = runST (do
 200         arr <- newSTArray ixs zero
 201         zap_with_f f arr ivs
 202         freezeSTArray arr
 203     )
 204 \end{code}
 205
 206
 207 %*********************************************************
 208 %*                                                      *
 209 \subsection{Array instances}
 210 %*                                                      *
 211 %*********************************************************
 212
 213
 214 \begin{code}
 215 instance Ix a => Functor (Array a) where
 216   fmap = amap
 217
 218 instance  (Ix a, Eq b)  => Eq (Array a b)  where
 219     a == a'             =  assocs a == assocs a'
 220     a /= a'             =  assocs a /= assocs a'
 221
 222 instance  (Ix a, Ord b) => Ord (Array a b)  where
 223     compare a b = compare (assocs a) (assocs b)
 224
 225 instance  (Ix a, Show a, Show b) => Show (Array a b)  where
 226     showsPrec p a = showParen (p > 9) (
 227                     showString "array " .
 228                     shows (bounds a) . showChar ' ' .
 229                     shows (assocs a)                  )
 230     showList = showList__ (showsPrec 0)
 231
 232 {-
 233 instance  (Ix a, Read a, Read b) => Read (Array a b)  where
 234     readsPrec p = readParen (p > 9)
 235            (\r -> [(array b as, u) | ("array",s) <- lex r,
 236                                      (b,t)       <- reads s,
 237                                      (as,u)      <- reads t   ])
 238     readList = readList__ (readsPrec 0)
 239 -}
 240 \end{code}
 241
 242
 243 %*********************************************************
 244 %*                                                      *
 245 \subsection{Operations on mutable arrays}
 246 %*                                                      *
 247 %*********************************************************
 248
 249 Idle ADR question: What's the tradeoff here between flattening these
 250 datatypes into @STArray ix ix (MutableArray# s elt)@ and using
 251 it as is?  As I see it, the former uses slightly less heap and
 252 provides faster access to the individual parts of the bounds while the
 253 code used has the benefit of providing a ready-made @(lo, hi)@ pair as
 254 required by many array-related functions.  Which wins? Is the
 255 difference significant (probably not).
 256
 257 Idle AJG answer: When I looked at the outputted code (though it was 2
 258 years ago) it seems like you often needed the tuple, and we build
 259 it frequently. Now we've got the overloading specialiser things
 260 might be different, though.
 261
 262 \begin{code}
 263 newSTArray :: Ix ix => (ix,ix) -> elt -> ST s (STArray s ix elt)
 264
 265 {-# SPECIALIZE newSTArray :: IPr       -> elt -> ST s (STArray s Int elt),
 266                              (IPr,IPr) -> elt -> ST s (STArray s IPr elt)
 267   #-}
 268 newSTArray (l,u) init = ST $ \ s# ->
 269     case rangeSize (l,u)          of { I# n# ->
 270     case (newArray# n# init s#)   of { (# s2#, arr# #) ->
 271     (# s2#, STArray l u arr# #) }}
 272
 273
 274
 275 boundsSTArray     :: Ix ix => STArray s ix elt -> (ix, ix)
 276 {-# SPECIALIZE boundsSTArray :: STArray s Int elt -> IPr #-}
 277 boundsSTArray     (STArray     l u _) = (l,u)
 278
 279 readSTArray     :: Ix ix => STArray s ix elt -> ix -> ST s elt
 280 {-# SPECIALIZE readSTArray :: STArray s Int elt -> Int -> ST s elt,
 281                               STArray s IPr elt -> IPr -> ST s elt
 282   #-}
 283
 284 readSTArray (STArray l u arr#) n = ST $ \ s# ->
 285     case (index (l,u) n)                of { I# n# ->
 286     case readArray# arr# n# s#          of { (# s2#, r #) ->
 287     (# s2#, r #) }}
 288
 289 writeSTArray     :: Ix ix => STArray s ix elt -> ix -> elt -> ST s ()
 290 {-# SPECIALIZE writeSTArray :: STArray s Int elt -> Int -> elt -> ST s (),
 291                                STArray s IPr elt -> IPr -> elt -> ST s ()
 292   #-}
 293
 294 writeSTArray (STArray l u arr#) n ele = ST $ \ s# ->
 295     case index (l,u) n                      of { I# n# ->
 296     case writeArray# arr# n# ele s#         of { s2# ->
 297     (# s2#, () #) }}
 298 \end{code}
 299
 300
 301 %*********************************************************
 302 %*                                                      *
 303 \subsection{Moving between mutable and immutable}
 304 %*                                                      *
 305 %*********************************************************
 306
 307 \begin{code}
 308 freezeSTArray     :: Ix ix => STArray s ix elt -> ST s (Array ix elt)
 309 {-# SPECIALISE freezeSTArray :: STArray s Int elt -> ST s (Array Int elt),
 310                               STArray s IPr elt -> ST s (Array IPr elt)
 311   #-}
 312
 313 freezeSTArray (STArray l u arr#) = ST $ \ s# ->
 314     case rangeSize (l,u)     of { I# n# ->
 315     case freeze arr# n# s# of { (# s2#, frozen# #) ->
 316     (# s2#, Array l u frozen# #) }}
 317   where
 318     freeze  :: MutableArray# s ele      -- the thing
 319             -> Int#                     -- size of thing to be frozen
 320             -> State# s                 -- the Universe and everything
 321             -> (# State# s, Array# ele #)
 322     freeze m_arr# n# s#
 323       = case newArray# n# init s#             of { (# s2#, newarr1# #) ->
 324         case copy 0# n# m_arr# newarr1# s2#   of { (# s3#, newarr2# #) ->
 325         unsafeFreezeArray# newarr2# s3#
 326         }}
 327       where
 328         init = error "freezeArray: element not copied"
 329
 330         copy :: Int# -> Int#
 331              -> MutableArray# s ele
 332              -> MutableArray# s ele
 333              -> State# s
 334              -> (# State# s, MutableArray# s ele #)
 335
 336         copy cur# end# from# to# st#
 337           | cur# ==# end#
 338             = (# st#, to# #)
 339           | otherwise
 340             = case readArray#  from# cur#     st#  of { (# s1#, ele #) ->
 341               case writeArray# to#   cur# ele s1# of { s2# ->
 342               copy (cur# +# 1#) end# from# to# s2#
 343               }}
 344
 345 unsafeFreezeSTArray     :: Ix ix => STArray s ix elt -> ST s (Array ix elt)
 346 unsafeFreezeSTArray (STArray l u arr#) = ST $ \ s# ->
 347     case unsafeFreezeArray# arr# s# of { (# s2#, frozen# #) ->
 348     (# s2#, Array l u frozen# #) }
 349
 350 --This takes a immutable array, and copies it into a mutable array, in a
 351 --hurry.
 352
 353 thawSTArray :: Ix ix => Array ix elt -> ST s (STArray s ix elt)
 354 {-# SPECIALISE thawSTArray :: Array Int elt -> ST s (STArray s Int elt),
 355                               Array IPr elt -> ST s (STArray s IPr elt)
 356   #-}
 357
 358 thawSTArray (Array l u arr#) = ST $ \ s# ->
 359     case rangeSize (l,u) of { I# n# ->
 360     case thaw arr# n# s# of { (# s2#, thawed# #) ->
 361     (# s2#, STArray l u thawed# #)}}
 362   where
 363     thaw  :: Array# ele                 -- the thing
 364             -> Int#                     -- size of thing to be thawed
 365             -> State# s                 -- the Universe and everything
 366             -> (# State# s, MutableArray# s ele #)
 367
 368     thaw arr1# n# s#
 369       = case newArray# n# init s#             of { (# s2#, newarr1# #) ->
 370         copy 0# n# arr1# newarr1# s2# }
 371       where
 372         init = error "thawSTArray: element not copied"
 373
 374         copy :: Int# -> Int#
 375              -> Array# ele
 376              -> MutableArray# s ele
 377              -> State# s
 378              -> (# State# s, MutableArray# s ele #)
 379
 380         copy cur# end# from# to# st#
 381           | cur# ==# end#
 382             = (# st#, to# #)
 383           | otherwise
 384             = case indexArray#  from# cur#        of { (# ele #) ->
 385               case writeArray# to#   cur# ele st# of { s1# ->
 386               copy (cur# +# 1#) end# from# to# s1#
 387               }}
 388
 389 -- this is a quicker version of the above, just flipping the type
 390 -- (& representation) of an immutable array. And placing a
 391 -- proof obligation on the programmer.
 392 unsafeThawSTArray :: Ix ix => Array ix elt -> ST s (STArray s ix elt)
 393 unsafeThawSTArray (Array l u arr#) = ST $ \ s# ->
 394    case unsafeThawArray# arr# s# of
 395       (# s2#, marr# #) -> (# s2#, STArray l u marr# #)
 396 \end{code}