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