1{-# LANGUAGE BangPatterns, CPP, MagicHash, Rank2Types, UnboxedTuples, ScopedTypeVariables #-}
2{-# OPTIONS_GHC -fno-full-laziness -funbox-strict-fields #-}
3{-# OPTIONS_HADDOCK not-home #-}
4
5-- | = WARNING
6--
7-- This module is considered __internal__.
8--
9-- The Package Versioning Policy __does not apply__.
10--
11-- The contents of this module may change __in any way whatsoever__
12-- and __without any warning__ between minor versions of this package.
13--
14-- Authors importing this module are expected to track development
15-- closely.
16--
17-- = Description
18--
19-- Zero based arrays.
20--
21-- Note that no bounds checking are performed.
22module Data.HashMap.Internal.Array
23    ( Array
24    , MArray
25
26      -- * Creation
27    , new
28    , new_
29    , singleton
30    , singletonM
31    , pair
32
33      -- * Basic interface
34    , length
35    , lengthM
36    , read
37    , write
38    , index
39    , indexM
40    , index#
41    , update
42    , updateWith'
43    , unsafeUpdateM
44    , insert
45    , insertM
46    , delete
47    , sameArray1
48    , trim
49
50    , unsafeFreeze
51    , unsafeThaw
52    , unsafeSameArray
53    , run
54    , copy
55    , copyM
56
57      -- * Folds
58    , foldl
59    , foldl'
60    , foldr
61    , foldr'
62    , foldMap
63    , all
64
65    , thaw
66    , map
67    , map'
68    , traverse
69    , traverse'
70    , toList
71    , fromList
72    ) where
73
74#if !MIN_VERSION_base(4,8,0)
75import Control.Applicative (Applicative (..), (<$>))
76#endif
77import Control.Applicative (liftA2)
78import Control.DeepSeq
79import GHC.Exts(Int(..), Int#, reallyUnsafePtrEquality#, tagToEnum#, unsafeCoerce#, State#)
80import GHC.ST (ST(..))
81import Control.Monad.ST (stToIO)
82
83#if __GLASGOW_HASKELL__ >= 709
84import Prelude hiding (filter, foldMap, foldr, foldl, length, map, read, traverse, all)
85#else
86import Prelude hiding (filter, foldr, foldl, length, map, read, all)
87#endif
88
89#if __GLASGOW_HASKELL__ >= 710
90import GHC.Exts (SmallArray#, newSmallArray#, readSmallArray#, writeSmallArray#,
91                 indexSmallArray#, unsafeFreezeSmallArray#, unsafeThawSmallArray#,
92                 SmallMutableArray#, sizeofSmallArray#, copySmallArray#, thawSmallArray#,
93                 sizeofSmallMutableArray#, copySmallMutableArray#, cloneSmallMutableArray#)
94
95#else
96import GHC.Exts (Array#, newArray#, readArray#, writeArray#,
97                 indexArray#, unsafeFreezeArray#, unsafeThawArray#,
98                 MutableArray#, sizeofArray#, copyArray#, thawArray#,
99                 sizeofMutableArray#, copyMutableArray#, cloneMutableArray#)
100import Data.Monoid (Monoid (..))
101#endif
102
103#if defined(ASSERTS)
104import qualified Prelude
105#endif
106
107import Data.HashMap.Internal.Unsafe (runST)
108import Control.Monad ((>=>))
109
110
111#if __GLASGOW_HASKELL__ >= 710
112type Array# a = SmallArray# a
113type MutableArray# a = SmallMutableArray# a
114
115newArray# :: Int# -> a -> State# d -> (# State# d, SmallMutableArray# d a #)
116newArray# = newSmallArray#
117
118unsafeFreezeArray# :: SmallMutableArray# d a
119                   -> State# d -> (# State# d, SmallArray# a #)
120unsafeFreezeArray# = unsafeFreezeSmallArray#
121
122readArray# :: SmallMutableArray# d a
123           -> Int# -> State# d -> (# State# d, a #)
124readArray# = readSmallArray#
125
126writeArray# :: SmallMutableArray# d a
127            -> Int# -> a -> State# d -> State# d
128writeArray# = writeSmallArray#
129
130indexArray# :: SmallArray# a -> Int# -> (# a #)
131indexArray# = indexSmallArray#
132
133unsafeThawArray# :: SmallArray# a
134                 -> State# d -> (# State# d, SmallMutableArray# d a #)
135unsafeThawArray# = unsafeThawSmallArray#
136
137sizeofArray# :: SmallArray# a -> Int#
138sizeofArray# = sizeofSmallArray#
139
140copyArray# :: SmallArray# a
141           -> Int#
142           -> SmallMutableArray# d a
143           -> Int#
144           -> Int#
145           -> State# d
146           -> State# d
147copyArray# = copySmallArray#
148
149cloneMutableArray# :: SmallMutableArray# s a
150                   -> Int#
151                   -> Int#
152                   -> State# s
153                   -> (# State# s, SmallMutableArray# s a #)
154cloneMutableArray# = cloneSmallMutableArray#
155
156thawArray# :: SmallArray# a
157           -> Int#
158           -> Int#
159           -> State# d
160           -> (# State# d, SmallMutableArray# d a #)
161thawArray# = thawSmallArray#
162
163sizeofMutableArray# :: SmallMutableArray# s a -> Int#
164sizeofMutableArray# = sizeofSmallMutableArray#
165
166copyMutableArray# :: SmallMutableArray# d a
167                  -> Int#
168                  -> SmallMutableArray# d a
169                  -> Int#
170                  -> Int#
171                  -> State# d
172                  -> State# d
173copyMutableArray# = copySmallMutableArray#
174#endif
175
176------------------------------------------------------------------------
177
178#if defined(ASSERTS)
179-- This fugly hack is brought by GHC's apparent reluctance to deal
180-- with MagicHash and UnboxedTuples when inferring types. Eek!
181# define CHECK_BOUNDS(_func_,_len_,_k_) \
182if (_k_) < 0 || (_k_) >= (_len_) then error ("Data.HashMap.Internal.Array." ++ (_func_) ++ ": bounds error, offset " ++ show (_k_) ++ ", length " ++ show (_len_)) else
183# define CHECK_OP(_func_,_op_,_lhs_,_rhs_) \
184if not ((_lhs_) _op_ (_rhs_)) then error ("Data.HashMap.Internal.Array." ++ (_func_) ++ ": Check failed: _lhs_ _op_ _rhs_ (" ++ show (_lhs_) ++ " vs. " ++ show (_rhs_) ++ ")") else
185# define CHECK_GT(_func_,_lhs_,_rhs_) CHECK_OP(_func_,>,_lhs_,_rhs_)
186# define CHECK_LE(_func_,_lhs_,_rhs_) CHECK_OP(_func_,<=,_lhs_,_rhs_)
187# define CHECK_EQ(_func_,_lhs_,_rhs_) CHECK_OP(_func_,==,_lhs_,_rhs_)
188#else
189# define CHECK_BOUNDS(_func_,_len_,_k_)
190# define CHECK_OP(_func_,_op_,_lhs_,_rhs_)
191# define CHECK_GT(_func_,_lhs_,_rhs_)
192# define CHECK_LE(_func_,_lhs_,_rhs_)
193# define CHECK_EQ(_func_,_lhs_,_rhs_)
194#endif
195
196data Array a = Array {
197      unArray :: !(Array# a)
198    }
199
200instance Show a => Show (Array a) where
201    show = show . toList
202
203-- Determines whether two arrays have the same memory address.
204-- This is more reliable than testing pointer equality on the
205-- Array wrappers, but it's still slightly bogus.
206unsafeSameArray :: Array a -> Array b -> Bool
207unsafeSameArray (Array xs) (Array ys) =
208  tagToEnum# (unsafeCoerce# reallyUnsafePtrEquality# xs ys)
209
210sameArray1 :: (a -> b -> Bool) -> Array a -> Array b -> Bool
211sameArray1 eq !xs0 !ys0
212  | lenxs /= lenys = False
213  | otherwise = go 0 xs0 ys0
214  where
215    go !k !xs !ys
216      | k == lenxs = True
217      | (# x #) <- index# xs k
218      , (# y #) <- index# ys k
219      = eq x y && go (k + 1) xs ys
220
221    !lenxs = length xs0
222    !lenys = length ys0
223
224length :: Array a -> Int
225length ary = I# (sizeofArray# (unArray ary))
226{-# INLINE length #-}
227
228data MArray s a = MArray {
229      unMArray :: !(MutableArray# s a)
230    }
231
232lengthM :: MArray s a -> Int
233lengthM mary = I# (sizeofMutableArray# (unMArray mary))
234{-# INLINE lengthM #-}
235
236------------------------------------------------------------------------
237
238instance NFData a => NFData (Array a) where
239    rnf = rnfArray
240
241rnfArray :: NFData a => Array a -> ()
242rnfArray ary0 = go ary0 n0 0
243  where
244    n0 = length ary0
245    go !ary !n !i
246        | i >= n = ()
247        | (# x #) <- index# ary i
248        = rnf x `seq` go ary n (i+1)
249-- We use index# just in case GHC can't see that the
250-- relevant rnf is strict, or in case it actually isn't.
251{-# INLINE rnfArray #-}
252
253-- | Create a new mutable array of specified size, in the specified
254-- state thread, with each element containing the specified initial
255-- value.
256new :: Int -> a -> ST s (MArray s a)
257new (I# n#) b =
258    CHECK_GT("new",n,(0 :: Int))
259    ST $ \s ->
260        case newArray# n# b s of
261            (# s', ary #) -> (# s', MArray ary #)
262{-# INLINE new #-}
263
264new_ :: Int -> ST s (MArray s a)
265new_ n = new n undefinedElem
266
267singleton :: a -> Array a
268singleton x = runST (singletonM x)
269{-# INLINE singleton #-}
270
271singletonM :: a -> ST s (Array a)
272singletonM x = new 1 x >>= unsafeFreeze
273{-# INLINE singletonM #-}
274
275pair :: a -> a -> Array a
276pair x y = run $ do
277    ary <- new 2 x
278    write ary 1 y
279    return ary
280{-# INLINE pair #-}
281
282read :: MArray s a -> Int -> ST s a
283read ary _i@(I# i#) = ST $ \ s ->
284    CHECK_BOUNDS("read", lengthM ary, _i)
285        readArray# (unMArray ary) i# s
286{-# INLINE read #-}
287
288write :: MArray s a -> Int -> a -> ST s ()
289write ary _i@(I# i#) b = ST $ \ s ->
290    CHECK_BOUNDS("write", lengthM ary, _i)
291        case writeArray# (unMArray ary) i# b s of
292            s' -> (# s' , () #)
293{-# INLINE write #-}
294
295index :: Array a -> Int -> a
296index ary _i@(I# i#) =
297    CHECK_BOUNDS("index", length ary, _i)
298        case indexArray# (unArray ary) i# of (# b #) -> b
299{-# INLINE index #-}
300
301index# :: Array a -> Int -> (# a #)
302index# ary _i@(I# i#) =
303    CHECK_BOUNDS("index#", length ary, _i)
304        indexArray# (unArray ary) i#
305{-# INLINE index# #-}
306
307indexM :: Array a -> Int -> ST s a
308indexM ary _i@(I# i#) =
309    CHECK_BOUNDS("indexM", length ary, _i)
310        case indexArray# (unArray ary) i# of (# b #) -> return b
311{-# INLINE indexM #-}
312
313unsafeFreeze :: MArray s a -> ST s (Array a)
314unsafeFreeze mary
315    = ST $ \s -> case unsafeFreezeArray# (unMArray mary) s of
316                   (# s', ary #) -> (# s', Array ary #)
317{-# INLINE unsafeFreeze #-}
318
319unsafeThaw :: Array a -> ST s (MArray s a)
320unsafeThaw ary
321    = ST $ \s -> case unsafeThawArray# (unArray ary) s of
322                   (# s', mary #) -> (# s', MArray mary #)
323{-# INLINE unsafeThaw #-}
324
325run :: (forall s . ST s (MArray s e)) -> Array e
326run act = runST $ act >>= unsafeFreeze
327{-# INLINE run #-}
328
329-- | Unsafely copy the elements of an array. Array bounds are not checked.
330copy :: Array e -> Int -> MArray s e -> Int -> Int -> ST s ()
331copy !src !_sidx@(I# sidx#) !dst !_didx@(I# didx#) _n@(I# n#) =
332    CHECK_LE("copy", _sidx + _n, length src)
333    CHECK_LE("copy", _didx + _n, lengthM dst)
334        ST $ \ s# ->
335        case copyArray# (unArray src) sidx# (unMArray dst) didx# n# s# of
336            s2 -> (# s2, () #)
337
338-- | Unsafely copy the elements of an array. Array bounds are not checked.
339copyM :: MArray s e -> Int -> MArray s e -> Int -> Int -> ST s ()
340copyM !src !_sidx@(I# sidx#) !dst !_didx@(I# didx#) _n@(I# n#) =
341    CHECK_BOUNDS("copyM: src", lengthM src, _sidx + _n - 1)
342    CHECK_BOUNDS("copyM: dst", lengthM dst, _didx + _n - 1)
343    ST $ \ s# ->
344    case copyMutableArray# (unMArray src) sidx# (unMArray dst) didx# n# s# of
345        s2 -> (# s2, () #)
346
347cloneM :: MArray s a -> Int -> Int -> ST s (MArray s a)
348cloneM _mary@(MArray mary#) _off@(I# off#) _len@(I# len#) =
349    CHECK_BOUNDS("cloneM_off", lengthM _mary, _off - 1)
350    CHECK_BOUNDS("cloneM_end", lengthM _mary, _off + _len - 1)
351    ST $ \ s ->
352    case cloneMutableArray# mary# off# len# s of
353      (# s', mary'# #) -> (# s', MArray mary'# #)
354
355-- | Create a new array of the @n@ first elements of @mary@.
356trim :: MArray s a -> Int -> ST s (Array a)
357trim mary n = cloneM mary 0 n >>= unsafeFreeze
358{-# INLINE trim #-}
359
360-- | /O(n)/ Insert an element at the given position in this array,
361-- increasing its size by one.
362insert :: Array e -> Int -> e -> Array e
363insert ary idx b = runST (insertM ary idx b)
364{-# INLINE insert #-}
365
366-- | /O(n)/ Insert an element at the given position in this array,
367-- increasing its size by one.
368insertM :: Array e -> Int -> e -> ST s (Array e)
369insertM ary idx b =
370    CHECK_BOUNDS("insertM", count + 1, idx)
371        do mary <- new_ (count+1)
372           copy ary 0 mary 0 idx
373           write mary idx b
374           copy ary idx mary (idx+1) (count-idx)
375           unsafeFreeze mary
376  where !count = length ary
377{-# INLINE insertM #-}
378
379-- | /O(n)/ Update the element at the given position in this array.
380update :: Array e -> Int -> e -> Array e
381update ary idx b = runST (updateM ary idx b)
382{-# INLINE update #-}
383
384-- | /O(n)/ Update the element at the given position in this array.
385updateM :: Array e -> Int -> e -> ST s (Array e)
386updateM ary idx b =
387    CHECK_BOUNDS("updateM", count, idx)
388        do mary <- thaw ary 0 count
389           write mary idx b
390           unsafeFreeze mary
391  where !count = length ary
392{-# INLINE updateM #-}
393
394-- | /O(n)/ Update the element at the given positio in this array, by
395-- applying a function to it.  Evaluates the element to WHNF before
396-- inserting it into the array.
397updateWith' :: Array e -> Int -> (e -> e) -> Array e
398updateWith' ary idx f
399  | (# x #) <- index# ary idx
400  = update ary idx $! f x
401{-# INLINE updateWith' #-}
402
403-- | /O(1)/ Update the element at the given position in this array,
404-- without copying.
405unsafeUpdateM :: Array e -> Int -> e -> ST s ()
406unsafeUpdateM ary idx b =
407    CHECK_BOUNDS("unsafeUpdateM", length ary, idx)
408        do mary <- unsafeThaw ary
409           write mary idx b
410           _ <- unsafeFreeze mary
411           return ()
412{-# INLINE unsafeUpdateM #-}
413
414foldl' :: (b -> a -> b) -> b -> Array a -> b
415foldl' f = \ z0 ary0 -> go ary0 (length ary0) 0 z0
416  where
417    go ary n i !z
418        | i >= n = z
419        | otherwise
420        = case index# ary i of
421            (# x #) -> go ary n (i+1) (f z x)
422{-# INLINE foldl' #-}
423
424foldr' :: (a -> b -> b) -> b -> Array a -> b
425foldr' f = \ z0 ary0 -> go ary0 (length ary0 - 1) z0
426  where
427    go !_ary (-1) z = z
428    go !ary i !z
429      | (# x #) <- index# ary i
430      = go ary (i - 1) (f x z)
431{-# INLINE foldr' #-}
432
433foldr :: (a -> b -> b) -> b -> Array a -> b
434foldr f = \ z0 ary0 -> go ary0 (length ary0) 0 z0
435  where
436    go ary n i z
437        | i >= n = z
438        | otherwise
439        = case index# ary i of
440            (# x #) -> f x (go ary n (i+1) z)
441{-# INLINE foldr #-}
442
443foldl :: (b -> a -> b) -> b -> Array a -> b
444foldl f = \ z0 ary0 -> go ary0 (length ary0 - 1) z0
445  where
446    go _ary (-1) z = z
447    go ary i z
448      | (# x #) <- index# ary i
449      = f (go ary (i - 1) z) x
450{-# INLINE foldl #-}
451
452-- We go to a bit of trouble here to avoid appending an extra mempty.
453-- The below implementation is by Mateusz Kowalczyk, who indicates that
454-- benchmarks show it to be faster than one that avoids lifting out
455-- lst.
456foldMap :: Monoid m => (a -> m) -> Array a -> m
457foldMap f = \ary0 -> case length ary0 of
458  0 -> mempty
459  len ->
460    let !lst = len - 1
461        go i | (# x #) <- index# ary0 i, let fx = f x =
462          if i == lst then fx else fx `mappend` go (i + 1)
463    in go 0
464{-# INLINE foldMap #-}
465
466-- | Verifies that a predicate holds for all elements of an array.
467all :: (a -> Bool) -> Array a -> Bool
468all p = foldr (\a acc -> p a && acc) True
469{-# INLINE all #-}
470
471undefinedElem :: a
472undefinedElem = error "Data.HashMap.Internal.Array: Undefined element"
473{-# NOINLINE undefinedElem #-}
474
475thaw :: Array e -> Int -> Int -> ST s (MArray s e)
476thaw !ary !_o@(I# o#) (I# n#) =
477    CHECK_LE("thaw", _o + n, length ary)
478        ST $ \ s -> case thawArray# (unArray ary) o# n# s of
479            (# s2, mary# #) -> (# s2, MArray mary# #)
480{-# INLINE thaw #-}
481
482-- | /O(n)/ Delete an element at the given position in this array,
483-- decreasing its size by one.
484delete :: Array e -> Int -> Array e
485delete ary idx = runST (deleteM ary idx)
486{-# INLINE delete #-}
487
488-- | /O(n)/ Delete an element at the given position in this array,
489-- decreasing its size by one.
490deleteM :: Array e -> Int -> ST s (Array e)
491deleteM ary idx = do
492    CHECK_BOUNDS("deleteM", count, idx)
493        do mary <- new_ (count-1)
494           copy ary 0 mary 0 idx
495           copy ary (idx+1) mary idx (count-(idx+1))
496           unsafeFreeze mary
497  where !count = length ary
498{-# INLINE deleteM #-}
499
500map :: (a -> b) -> Array a -> Array b
501map f = \ ary ->
502    let !n = length ary
503    in run $ do
504        mary <- new_ n
505        go ary mary 0 n
506  where
507    go ary mary i n
508        | i >= n    = return mary
509        | otherwise = do
510             x <- indexM ary i
511             write mary i $ f x
512             go ary mary (i+1) n
513{-# INLINE map #-}
514
515-- | Strict version of 'map'.
516map' :: (a -> b) -> Array a -> Array b
517map' f = \ ary ->
518    let !n = length ary
519    in run $ do
520        mary <- new_ n
521        go ary mary 0 n
522  where
523    go ary mary i n
524        | i >= n    = return mary
525        | otherwise = do
526             x <- indexM ary i
527             write mary i $! f x
528             go ary mary (i+1) n
529{-# INLINE map' #-}
530
531fromList :: Int -> [a] -> Array a
532fromList n xs0 =
533    CHECK_EQ("fromList", n, Prelude.length xs0)
534        run $ do
535            mary <- new_ n
536            go xs0 mary 0
537  where
538    go [] !mary !_   = return mary
539    go (x:xs) mary i = do write mary i x
540                          go xs mary (i+1)
541
542toList :: Array a -> [a]
543toList = foldr (:) []
544
545newtype STA a = STA {_runSTA :: forall s. MutableArray# s a -> ST s (Array a)}
546
547runSTA :: Int -> STA a -> Array a
548runSTA !n (STA m) = runST $ new_ n >>= \ (MArray ar) -> m ar
549
550traverse :: Applicative f => (a -> f b) -> Array a -> f (Array b)
551traverse f = \ !ary ->
552  let
553    !len = length ary
554    go !i
555      | i == len = pure $ STA $ \mary -> unsafeFreeze (MArray mary)
556      | (# x #) <- index# ary i
557      = liftA2 (\b (STA m) -> STA $ \mary ->
558                  write (MArray mary) i b >> m mary)
559               (f x) (go (i + 1))
560  in runSTA len <$> go 0
561{-# INLINE [1] traverse #-}
562
563-- TODO: Would it be better to just use a lazy traversal
564-- and then force the elements of the result? My guess is
565-- yes.
566traverse' :: Applicative f => (a -> f b) -> Array a -> f (Array b)
567traverse' f = \ !ary ->
568  let
569    !len = length ary
570    go !i
571      | i == len = pure $ STA $ \mary -> unsafeFreeze (MArray mary)
572      | (# x #) <- index# ary i
573      = liftA2 (\ !b (STA m) -> STA $ \mary ->
574                    write (MArray mary) i b >> m mary)
575               (f x) (go (i + 1))
576  in runSTA len <$> go 0
577{-# INLINE [1] traverse' #-}
578
579-- Traversing in ST, we don't need to get fancy; we
580-- can just do it directly.
581traverseST :: (a -> ST s b) -> Array a -> ST s (Array b)
582traverseST f = \ ary0 ->
583  let
584    !len = length ary0
585    go k !mary
586      | k == len = return mary
587      | otherwise = do
588          x <- indexM ary0 k
589          y <- f x
590          write mary k y
591          go (k + 1) mary
592  in new_ len >>= (go 0 >=> unsafeFreeze)
593{-# INLINE traverseST #-}
594
595traverseIO :: (a -> IO b) -> Array a -> IO (Array b)
596traverseIO f = \ ary0 ->
597  let
598    !len = length ary0
599    go k !mary
600      | k == len = return mary
601      | otherwise = do
602          x <- stToIO $ indexM ary0 k
603          y <- f x
604          stToIO $ write mary k y
605          go (k + 1) mary
606  in stToIO (new_ len) >>= (go 0 >=> stToIO . unsafeFreeze)
607{-# INLINE traverseIO #-}
608
609
610-- Why don't we have similar RULES for traverse'? The efficient
611-- way to traverse strictly in IO or ST is to force results as
612-- they come in, which leads to different semantics. In particular,
613-- we need to ensure that
614--
615--  traverse' (\x -> print x *> pure undefined) xs
616--
617-- will actually print all the values and then return undefined.
618-- We could add a strict mapMWithIndex, operating in an arbitrary
619-- Monad, that supported such rules, but we don't have that right now.
620{-# RULES
621"traverse/ST" forall f. traverse f = traverseST f
622"traverse/IO" forall f. traverse f = traverseIO f
623 #-}
624