@@ -0,0 +1,966 @@
1+{-# LANGUAGE BangPatterns #-}
2+{-# LANGUAGE CPP #-}
3+{-# LANGUAGE MagicHash #-}
4+{-# LANGUAGE RankNTypes #-}
5+{-# LANGUAGE ScopedTypeVariables #-}
6+{-# LANGUAGE TypeFamilies #-}
7+{-# LANGUAGE UnboxedTuples #-}
8+
9+
10+-- |
11+-- Module : Data.Primitive.PrimArray
12+-- Copyright : (c) Roman Leshchinskiy 2009-2012
13+-- License : BSD-style
14+--
15+-- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au>
16+-- Portability : non-portable
17+--
18+-- Arrays of unboxed primitive types. The function provided by this module
19+-- match the behavior of those provided by @Data.Primitive.ByteArray@, and
20+-- the underlying types and primops that back them are the same.
21+-- However, the type constructors 'PrimArray' and 'MutablePrimArray' take one additional
22+-- argument than their respective counterparts 'ByteArray' and 'MutableByteArray'.
23+-- This argument is used to designate the type of element in the array.
24+-- Consequently, all function this modules accepts length and incides in
25+-- terms of elements, not bytes.
26+--
27+-- @since 0.6.4.0
28+module Data.Primitive.PrimArray
29+ ( -- * Types
30+ PrimArray(..)
31+ , MutablePrimArray(..)
32+ -- * Allocation
33+ , newPrimArray
34+ , resizeMutablePrimArray
35+#if __GLASGOW_HASKELL__ >= 710
36+ , shrinkMutablePrimArray
37+#endif
38+ -- * Element Access
39+ , readPrimArray
40+ , writePrimArray
41+ , indexPrimArray
42+ -- * Freezing and Thawing
43+ , unsafeFreezePrimArray
44+ , unsafeThawPrimArray
45+ -- * Block Operations
46+ , copyPrimArray
47+ , copyMutablePrimArray
48+#if __GLASGOW_HASKELL__ >= 708
49+ , copyPrimArrayToPtr
50+ , copyMutablePrimArrayToPtr
51+#endif
52+ , setPrimArray
53+ -- * Information
54+ , sameMutablePrimArray
55+ , getSizeofMutablePrimArray
56+ , sizeofMutablePrimArray
57+ , sizeofPrimArray
58+ -- * List Conversion
59+ , primArrayToList
60+ , primArrayFromList
61+ , primArrayFromListN
62+ -- * Folding
63+ , foldrPrimArray
64+ , foldrPrimArray'
65+ , foldlPrimArray
66+ , foldlPrimArray'
67+ , foldlPrimArrayM'
68+ -- * Effectful Folding
69+ , traversePrimArray_
70+ , itraversePrimArray_
71+ -- * Map/Create
72+ , mapPrimArray
73+ , imapPrimArray
74+ , generatePrimArray
75+ , replicatePrimArray
76+ , filterPrimArray
77+ , mapMaybePrimArray
78+ -- * Effectful Map/Create
79+ -- $effectfulMapCreate
80+ -- ** Lazy Applicative
81+ , traversePrimArray
82+ , itraversePrimArray
83+ , generatePrimArrayA
84+ , replicatePrimArrayA
85+ , filterPrimArrayA
86+ , mapMaybePrimArrayA
87+ -- ** Strict Primitive Monadic
88+ , traversePrimArrayP
89+ , itraversePrimArrayP
90+ , generatePrimArrayP
91+ , replicatePrimArrayP
92+ , filterPrimArrayP
93+ , mapMaybePrimArrayP
94+ ) where
95+
96+import GHC.Exts
97+import GHC.Base ( Int(..) )
98+import Data.Primitive.Internal.Compat (isTrue#)
99+import Data.Primitive.Types
100+import Data.Primitive.ByteArray (ByteArray(..))
101+import Data.Monoid (Monoid(..),(<>))
102+import Control.Applicative
103+import Control.Monad.Primitive
104+import Control.Monad.ST
105+import qualified Data.List as L
106+import qualified Data.Primitive.ByteArray as PB
107+import qualified Data.Primitive.Types as PT
108+
109+#if MIN_VERSION_base(4,7,0)
110+import GHC.Exts (IsList(..))
111+#endif
112+
113+#if MIN_VERSION_base(4,9,0)
114+import Data.Semigroup (Semigroup)
115+import qualified Data.Semigroup as SG
116+#endif
117+
118+-- | Arrays of unboxed elements. This accepts types like 'Double', 'Char',
119+-- 'Int', and 'Word', as well as their fixed-length variants ('Word8',
120+-- 'Word16', etc.). Since the elements are unboxed, a 'PrimArray' is strict
121+-- in its elements. This differs from the behavior of 'Array', which is lazy
122+-- in its elements.
123+data PrimArray a = PrimArray ByteArray#
124+
125+-- | Mutable primitive arrays associated with a primitive state token.
126+-- These can be written to and read from in a monadic context that supports
127+-- sequencing such as 'IO' or 'ST'. Typically, a mutable primitive array will
128+-- be built and then convert to an immutable primitive array using
129+-- 'unsafeFreezePrimArray'. However, it is also acceptable to simply discard
130+-- a mutable primitive array since it lives in managed memory and will be
131+-- garbage collected when no longer referenced.
132+data MutablePrimArray s a = MutablePrimArray (MutableByteArray# s)
133+
134+sameByteArray :: ByteArray# -> ByteArray# -> Bool
135+sameByteArray ba1 ba2 =
136+ case reallyUnsafePtrEquality# (unsafeCoerce# ba1 :: ()) (unsafeCoerce# ba2 :: ()) of
137+#if __GLASGOW_HASKELL__ >= 708
138+ r -> isTrue# r
139+#else
140+ 1# -> True
141+ _ -> False
142+#endif
143+
144+-- | @since 0.6.4.0
145+instance (Eq a, Prim a) => Eq (PrimArray a) where
146+ a1@(PrimArray ba1#) == a2@(PrimArray ba2#)
147+ | sameByteArray ba1# ba2# = True
148+ | sz1 /= sz2 = False
149+ | otherwise = loop (quot sz1 (sizeOf (undefined :: a)) - 1)
150+ where
151+ -- Here, we take the size in bytes, not in elements. We do this
152+ -- since it allows us to defer performing the division to
153+ -- calculate the size in elements.
154+ sz1 = PB.sizeofByteArray (ByteArray ba1#)
155+ sz2 = PB.sizeofByteArray (ByteArray ba2#)
156+ loop !i
157+ | i < 0 = True
158+ | otherwise = indexPrimArray a1 i == indexPrimArray a2 i && loop (i-1)
159+ {-# INLINE (==) #-}
160+
161+-- | Lexicographic ordering. Subject to change between major versions.
162+--
163+-- @since 0.6.4.0
164+instance (Ord a, Prim a) => Ord (PrimArray a) where
165+ compare a1@(PrimArray ba1#) a2@(PrimArray ba2#)
166+ | sameByteArray ba1# ba2# = EQ
167+ | otherwise = loop 0
168+ where
169+ sz1 = PB.sizeofByteArray (ByteArray ba1#)
170+ sz2 = PB.sizeofByteArray (ByteArray ba2#)
171+ sz = quot (min sz1 sz2) (sizeOf (undefined :: a))
172+ loop !i
173+ | i < sz = compare (indexPrimArray a1 i) (indexPrimArray a2 i) <> loop (i+1)
174+ | otherwise = compare sz1 sz2
175+ {-# INLINE compare #-}
176+
177+#if MIN_VERSION_base(4,7,0)
178+-- | @since 0.6.4.0
179+instance Prim a => IsList (PrimArray a) where
180+ type Item (PrimArray a) = a
181+ fromList = primArrayFromList
182+ fromListN = primArrayFromListN
183+ toList = primArrayToList
184+#endif
185+
186+-- | @since 0.6.4.0
187+instance (Show a, Prim a) => Show (PrimArray a) where
188+ showsPrec p a = showParen (p > 10) $
189+ showString "fromListN " . shows (sizeofPrimArray a) . showString " "
190+ . shows (primArrayToList a)
191+
192+die :: String -> String -> a
193+die fun problem = error $ "Data.Primitive.PrimArray." ++ fun ++ ": " ++ problem
194+
195+primArrayFromList :: Prim a => [a] -> PrimArray a
196+primArrayFromList vs = primArrayFromListN (L.length vs) vs
197+
198+primArrayFromListN :: forall a. Prim a => Int -> [a] -> PrimArray a
199+primArrayFromListN len vs = runST run where
200+ run :: forall s. ST s (PrimArray a)
201+ run = do
202+ arr <- newPrimArray len
203+ let go :: [a] -> Int -> ST s ()
204+ go [] !ix = if ix == len
205+ then return ()
206+ else die "fromListN" "list length less than specified size"
207+ go (a : as) !ix = if ix < len
208+ then do
209+ writePrimArray arr ix a
210+ go as (ix + 1)
211+ else die "fromListN" "list length greater than specified size"
212+ go vs 0
213+ unsafeFreezePrimArray arr
214+
215+-- | Convert the primitive array to a list.
216+{-# INLINE primArrayToList #-}
217+primArrayToList :: forall a. Prim a => PrimArray a -> [a]
218+primArrayToList xs = build (\c n -> foldrPrimArray c n xs)
219+
220+primArrayToByteArray :: PrimArray a -> PB.ByteArray
221+primArrayToByteArray (PrimArray x) = PB.ByteArray x
222+
223+byteArrayToPrimArray :: ByteArray -> PrimArray a
224+byteArrayToPrimArray (PB.ByteArray x) = PrimArray x
225+
226+#if MIN_VERSION_base(4,9,0)
227+-- | @since 0.6.4.0
228+instance Semigroup (PrimArray a) where
229+ x <> y = byteArrayToPrimArray (primArrayToByteArray x SG.<> primArrayToByteArray y)
230+ sconcat = byteArrayToPrimArray . SG.sconcat . fmap primArrayToByteArray
231+ stimes i arr = byteArrayToPrimArray (SG.stimes i (primArrayToByteArray arr))
232+#endif
233+
234+-- | @since 0.6.4.0
235+instance Monoid (PrimArray a) where
236+ mempty = emptyPrimArray
237+#if !(MIN_VERSION_base(4,11,0))
238+ mappend x y = byteArrayToPrimArray (mappend (primArrayToByteArray x) (primArrayToByteArray y))
239+#endif
240+ mconcat = byteArrayToPrimArray . mconcat . map primArrayToByteArray
241+
242+-- | The empty primitive array.
243+emptyPrimArray :: PrimArray a
244+{-# NOINLINE emptyPrimArray #-}
245+emptyPrimArray = runST $ primitive $ \s0# -> case newByteArray# 0# s0# of
246+ (# s1#, arr# #) -> case unsafeFreezeByteArray# arr# s1# of
247+ (# s2#, arr'# #) -> (# s2#, PrimArray arr'# #)
248+
249+-- | Create a new mutable primitive array of the given length. The
250+-- underlying memory is left uninitialized.
251+newPrimArray :: forall m a. (PrimMonad m, Prim a) => Int -> m (MutablePrimArray (PrimState m) a)
252+{-# INLINE newPrimArray #-}
253+newPrimArray (I# n#)
254+ = primitive (\s# ->
255+ case newByteArray# (n# *# sizeOf# (undefined :: a)) s# of
256+ (# s'#, arr# #) -> (# s'#, MutablePrimArray arr# #)
257+ )
258+
259+-- | Resize a mutable primitive array. The new size is given in elements.
260+--
261+-- This will either resize the array in-place or, if not possible, allocate the
262+-- contents into a new, unpinned array and copy the original array\'s contents.
263+--
264+-- To avoid undefined behaviour, the original 'MutablePrimArray' shall not be
265+-- accessed anymore after a 'resizeMutablePrimArray' has been performed.
266+-- Moreover, no reference to the old one should be kept in order to allow
267+-- garbage collection of the original 'MutablePrimArray' in case a new
268+-- 'MutablePrimArray' had to be allocated.
269+resizeMutablePrimArray :: forall m a. (PrimMonad m, Prim a)
270+ => MutablePrimArray (PrimState m) a
271+ -> Int -- ^ new size
272+ -> m (MutablePrimArray (PrimState m) a)
273+{-# INLINE resizeMutablePrimArray #-}
274+#if __GLASGOW_HASKELL__ >= 710
275+resizeMutablePrimArray (MutablePrimArray arr#) (I# n#)
276+ = primitive (\s# -> case resizeMutableByteArray# arr# (n# *# sizeOf# (undefined :: a)) s# of
277+ (# s'#, arr'# #) -> (# s'#, MutablePrimArray arr'# #))
278+#else
279+resizeMutablePrimArray arr n
280+ = do arr' <- newPrimArray n
281+ copyMutablePrimArray arr' 0 arr 0 (min (sizeofMutablePrimArray arr) n)
282+ return arr'
283+#endif
284+
285+-- Although it is possible to shim resizeMutableByteArray for old GHCs, this
286+-- is not the case with shrinkMutablePrimArray.
287+#if __GLASGOW_HASKELL__ >= 710
288+-- | Shrink a mutable primitive array. The new size is given in elements.
289+-- It must be smaller than the old size. The array will be resized in place.
290+-- This function is only available when compiling with GHC 7.10 or newer.
291+shrinkMutablePrimArray :: forall m a. (PrimMonad m, Prim a)
292+ => MutablePrimArray (PrimState m) a
293+ -> Int -- ^ new size
294+ -> m ()
295+{-# INLINE shrinkMutablePrimArray #-}
296+shrinkMutablePrimArray (MutablePrimArray arr#) (I# n#)
297+ = primitive_ (shrinkMutableByteArray# arr# (n# *# sizeOf# (undefined :: a)))
298+#endif
299+
300+readPrimArray :: (Prim a, PrimMonad m) => MutablePrimArray (PrimState m) a -> Int -> m a
301+{-# INLINE readPrimArray #-}
302+readPrimArray (MutablePrimArray arr#) (I# i#)
303+ = primitive (readByteArray# arr# i#)
304+
305+-- | Write an element to the given index.
306+writePrimArray ::
307+ (Prim a, PrimMonad m)
308+ => MutablePrimArray (PrimState m) a -- ^ array
309+ -> Int -- ^ index
310+ -> a -- ^ element
311+ -> m ()
312+{-# INLINE writePrimArray #-}
313+writePrimArray (MutablePrimArray arr#) (I# i#) x
314+ = primitive_ (writeByteArray# arr# i# x)
315+
316+-- | Copy part of a mutable array into another mutable array.
317+-- In the case that the destination and
318+-- source arrays are the same, the regions may overlap.
319+copyMutablePrimArray :: forall m a.
320+ (PrimMonad m, Prim a)
321+ => MutablePrimArray (PrimState m) a -- ^ destination array
322+ -> Int -- ^ offset into destination array
323+ -> MutablePrimArray (PrimState m) a -- ^ source array
324+ -> Int -- ^ offset into source array
325+ -> Int -- ^ number of elements to copy
326+ -> m ()
327+{-# INLINE copyMutablePrimArray #-}
328+copyMutablePrimArray (MutablePrimArray dst#) (I# doff#) (MutablePrimArray src#) (I# soff#) (I# n#)
329+ = primitive_ (copyMutableByteArray#
330+ src#
331+ (soff# *# (sizeOf# (undefined :: a)))
332+ dst#
333+ (doff# *# (sizeOf# (undefined :: a)))
334+ (n# *# (sizeOf# (undefined :: a)))
335+ )
336+
337+-- | Copy part of an array into another mutable array.
338+copyPrimArray :: forall m a.
339+ (PrimMonad m, Prim a)
340+ => MutablePrimArray (PrimState m) a -- ^ destination array
341+ -> Int -- ^ offset into destination array
342+ -> PrimArray a -- ^ source array
343+ -> Int -- ^ offset into source array
344+ -> Int -- ^ number of elements to copy
345+ -> m ()
346+{-# INLINE copyPrimArray #-}
347+copyPrimArray (MutablePrimArray dst#) (I# doff#) (PrimArray src#) (I# soff#) (I# n#)
348+ = primitive_ (copyByteArray#
349+ src#
350+ (soff# *# (sizeOf# (undefined :: a)))
351+ dst#
352+ (doff# *# (sizeOf# (undefined :: a)))
353+ (n# *# (sizeOf# (undefined :: a)))
354+ )
355+
356+#if __GLASGOW_HASKELL__ >= 708
357+-- | Copy a slice of an immutable primitive array to an address.
358+-- The offset and length are given in elements of type @a@.
359+-- This function assumes that the 'Prim' instance of @a@
360+-- agrees with the 'Storable' instance. This function is only
361+-- available when building with GHC 7.8 or newer.
362+copyPrimArrayToPtr :: forall m a. (PrimMonad m, Prim a)
363+ => Ptr a -- ^ destination pointer
364+ -> PrimArray a -- ^ source array
365+ -> Int -- ^ offset into source array
366+ -> Int -- ^ number of prims to copy
367+ -> m ()
368+{-# INLINE copyPrimArrayToPtr #-}
369+copyPrimArrayToPtr (Ptr addr#) (PrimArray ba#) (I# soff#) (I# n#) =
370+ primitive (\ s# ->
371+ let s'# = copyByteArrayToAddr# ba# (soff# *# siz#) addr# (n# *# siz#) s#
372+ in (# s'#, () #))
373+ where siz# = sizeOf# (undefined :: a)
374+
375+-- | Copy a slice of an immutable primitive array to an address.
376+-- The offset and length are given in elements of type @a@.
377+-- This function assumes that the 'Prim' instance of @a@
378+-- agrees with the 'Storable' instance. This function is only
379+-- available when building with GHC 7.8 or newer.
380+copyMutablePrimArrayToPtr :: forall m a. (PrimMonad m, Prim a)
381+ => Ptr a -- ^ destination pointer
382+ -> MutablePrimArray (PrimState m) a -- ^ source array
383+ -> Int -- ^ offset into source array
384+ -> Int -- ^ number of prims to copy
385+ -> m ()
386+{-# INLINE copyMutablePrimArrayToPtr #-}
387+copyMutablePrimArrayToPtr (Ptr addr#) (MutablePrimArray mba#) (I# soff#) (I# n#) =
388+ primitive (\ s# ->
389+ let s'# = copyMutableByteArrayToAddr# mba# (soff# *# siz#) addr# (n# *# siz#) s#
390+ in (# s'#, () #))
391+ where siz# = sizeOf# (undefined :: a)
392+#endif
393+
394+-- | Fill a slice of a mutable primitive array with a value.
395+setPrimArray
396+ :: (Prim a, PrimMonad m)
397+ => MutablePrimArray (PrimState m) a -- ^ array to fill
398+ -> Int -- ^ offset into array
399+ -> Int -- ^ number of values to fill
400+ -> a -- ^ value to fill with
401+ -> m ()
402+{-# INLINE setPrimArray #-}
403+setPrimArray (MutablePrimArray dst#) (I# doff#) (I# sz#) x
404+ = primitive_ (PT.setByteArray# dst# doff# sz# x)
405+
406+-- | Get the size of a mutable primitive array in elements. Unlike 'sizeofMutablePrimArray',
407+-- this function ensures sequencing in the presence of resizing.
408+getSizeofMutablePrimArray :: forall m a. (PrimMonad m, Prim a)
409+ => MutablePrimArray (PrimState m) a -- ^ array
410+ -> m Int
411+{-# INLINE getSizeofMutablePrimArray #-}
412+#if __GLASGOW_HASKELL__ >= 801
413+getSizeofMutablePrimArray (MutablePrimArray arr#)
414+ = primitive (\s# ->
415+ case getSizeofMutableByteArray# arr# s# of
416+ (# s'#, sz# #) -> (# s'#, I# (quotInt# sz# (sizeOf# (undefined :: a))) #)
417+ )
418+#else
419+-- On older GHCs, it is not possible to resize a byte array, so
420+-- this provides behavior consistent with the implementation for
421+-- newer GHCs.
422+getSizeofMutablePrimArray arr
423+ = return (sizeofMutablePrimArray arr)
424+#endif
425+
426+-- | Size of the mutable primitive array in elements. This function shall not
427+-- be used on primitive arrays that are an argument to or a result of
428+-- 'resizeMutablePrimArray' or 'shrinkMutablePrimArray'.
429+sizeofMutablePrimArray :: forall s a. Prim a => MutablePrimArray s a -> Int
430+{-# INLINE sizeofMutablePrimArray #-}
431+sizeofMutablePrimArray (MutablePrimArray arr#) =
432+ I# (quotInt# (sizeofMutableByteArray# arr#) (sizeOf# (undefined :: a)))
433+
434+-- | Check if the two arrays refer to the same memory block.
435+sameMutablePrimArray :: MutablePrimArray s a -> MutablePrimArray s a -> Bool
436+{-# INLINE sameMutablePrimArray #-}
437+sameMutablePrimArray (MutablePrimArray arr#) (MutablePrimArray brr#)
438+ = isTrue# (sameMutableByteArray# arr# brr#)
439+
440+-- | Convert a mutable byte array to an immutable one without copying. The
441+-- array should not be modified after the conversion.
442+unsafeFreezePrimArray
443+ :: PrimMonad m => MutablePrimArray (PrimState m) a -> m (PrimArray a)
444+{-# INLINE unsafeFreezePrimArray #-}
445+unsafeFreezePrimArray (MutablePrimArray arr#)
446+ = primitive (\s# -> case unsafeFreezeByteArray# arr# s# of
447+ (# s'#, arr'# #) -> (# s'#, PrimArray arr'# #))
448+
449+-- | Convert an immutable array to a mutable one without copying. The
450+-- original array should not be used after the conversion.
451+unsafeThawPrimArray
452+ :: PrimMonad m => PrimArray a -> m (MutablePrimArray (PrimState m) a)
453+{-# INLINE unsafeThawPrimArray #-}
454+unsafeThawPrimArray (PrimArray arr#)
455+ = primitive (\s# -> (# s#, MutablePrimArray (unsafeCoerce# arr#) #))
456+
457+-- | Read a primitive value from the primitive array.
458+indexPrimArray :: forall a. Prim a => PrimArray a -> Int -> a
459+{-# INLINE indexPrimArray #-}
460+indexPrimArray (PrimArray arr#) (I# i#) = indexByteArray# arr# i#
461+
462+-- | Get the size, in elements, of the primitive array.
463+sizeofPrimArray :: forall a. Prim a => PrimArray a -> Int
464+{-# INLINE sizeofPrimArray #-}
465+sizeofPrimArray (PrimArray arr#) = I# (quotInt# (sizeofByteArray# arr#) (sizeOf# (undefined :: a)))
466+
467+-- | Lazy right-associated fold over the elements of a 'PrimArray'.
468+{-# INLINE foldrPrimArray #-}
469+foldrPrimArray :: forall a b. Prim a => (a -> b -> b) -> b -> PrimArray a -> b
470+foldrPrimArray f z arr = go 0
471+ where
472+ !sz = sizeofPrimArray arr
473+ go !i
474+ | sz > i = f (indexPrimArray arr i) (go (i+1))
475+ | otherwise = z
476+
477+-- | Strict right-associated fold over the elements of a 'PrimArray'.
478+{-# INLINE foldrPrimArray' #-}
479+foldrPrimArray' :: forall a b. Prim a => (a -> b -> b) -> b -> PrimArray a -> b
480+foldrPrimArray' f z0 arr = go (sizeofPrimArray arr - 1) z0
481+ where
482+ go !i !acc
483+ | i < 0 = acc
484+ | otherwise = go (i - 1) (f (indexPrimArray arr i) acc)
485+
486+-- | Lazy left-associated fold over the elements of a 'PrimArray'.
487+{-# INLINE foldlPrimArray #-}
488+foldlPrimArray :: forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
489+foldlPrimArray f z arr = go (sizeofPrimArray arr - 1)
490+ where
491+ go !i
492+ | i < 0 = z
493+ | otherwise = f (go (i - 1)) (indexPrimArray arr i)
494+
495+-- | Strict left-associated fold over the elements of a 'PrimArray'.
496+{-# INLINE foldlPrimArray' #-}
497+foldlPrimArray' :: forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
498+foldlPrimArray' f z0 arr = go 0 z0
499+ where
500+ !sz = sizeofPrimArray arr
501+ go !i !acc
502+ | i < sz = go (i + 1) (f acc (indexPrimArray arr i))
503+ | otherwise = acc
504+
505+-- | Strict left-associated fold over the elements of a 'PrimArray'.
506+{-# INLINE foldlPrimArrayM' #-}
507+foldlPrimArrayM' :: (Prim a, Monad m) => (b -> a -> m b) -> b -> PrimArray a -> m b
508+foldlPrimArrayM' f z0 arr = go 0 z0
509+ where
510+ !sz = sizeofPrimArray arr
511+ go !i !acc1
512+ | i < sz = do
513+ acc2 <- f acc1 (indexPrimArray arr i)
514+ go (i + 1) acc2
515+ | otherwise = return acc1
516+
517+-- | Traverse a primitive array. The traversal forces the resulting values and
518+-- writes them to the new primitive array as it performs the monadic effects.
519+-- Consequently:
520+--
521+-- >>> traversePrimArrayP (\x -> print x $> bool x undefined (x == 2)) (fromList [1, 2, 3 :: Int])
522+-- 1
523+-- 2
524+-- *** Exception: Prelude.undefined
525+--
526+-- In many situations, 'traversePrimArrayP' can replace 'traversePrimArray',
527+-- changing the strictness characteristics of the traversal but typically improving
528+-- the performance. Consider the following short-circuiting traversal:
529+--
530+-- > incrPositiveA :: PrimArray Int -> Maybe (PrimArray Int)
531+-- > incrPositiveA xs = traversePrimArray (\x -> bool Nothing (Just (x + 1)) (x > 0)) xs
532+--
533+-- This can be rewritten using 'traversePrimArrayP'. To do this, we must
534+-- change the traversal context to @MaybeT (ST s)@, which has a 'PrimMonad'
535+-- instance:
536+--
537+-- > incrPositiveB :: PrimArray Int -> Maybe (PrimArray Int)
538+-- > incrPositiveB xs = runST $ runMaybeT $ traversePrimArrayP
539+-- > (\x -> bool (MaybeT (return Nothing)) (MaybeT (return (Just (x + 1)))) (x > 0))
540+-- > xs
541+--
542+-- Benchmarks demonstrate that the second implementation runs 150 times
543+-- faster than the first. It also results in fewer allocations.
544+{-# INLINE traversePrimArrayP #-}
545+traversePrimArrayP :: (PrimMonad m, Prim a, Prim b)
546+ => (a -> m b)
547+ -> PrimArray a
548+ -> m (PrimArray b)
549+traversePrimArrayP f arr = do
550+ let !sz = sizeofPrimArray arr
551+ marr <- newPrimArray sz
552+ let go !ix = if ix < sz
553+ then do
554+ b <- f (indexPrimArray arr ix)
555+ writePrimArray marr ix b
556+ go (ix + 1)
557+ else return ()
558+ go 0
559+ unsafeFreezePrimArray marr
560+
561+-- | Filter the primitive array, keeping the elements for which the monadic
562+-- predicate evaluates true.
563+{-# INLINE filterPrimArrayP #-}
564+filterPrimArrayP :: (PrimMonad m, Prim a)
565+ => (a -> m Bool)
566+ -> PrimArray a
567+ -> m (PrimArray a)
568+filterPrimArrayP f arr = do
569+ let !sz = sizeofPrimArray arr
570+ marr <- newPrimArray sz
571+ let go !ixSrc !ixDst = if ixSrc < sz
572+ then do
573+ let a = indexPrimArray arr ixSrc
574+ b <- f a
575+ if b
576+ then do
577+ writePrimArray marr ixDst a
578+ go (ixSrc + 1) (ixDst + 1)
579+ else go (ixSrc + 1) ixDst
580+ else return ixDst
581+ lenDst <- go 0 0
582+ marr' <- resizeMutablePrimArray marr lenDst
583+ unsafeFreezePrimArray marr'
584+
585+-- | Map over the primitive array, keeping the elements for which the monadic
586+-- predicate provides a 'Just'.
587+{-# INLINE mapMaybePrimArrayP #-}
588+mapMaybePrimArrayP :: (PrimMonad m, Prim a, Prim b)
589+ => (a -> m (Maybe b))
590+ -> PrimArray a
591+ -> m (PrimArray b)
592+mapMaybePrimArrayP f arr = do
593+ let !sz = sizeofPrimArray arr
594+ marr <- newPrimArray sz
595+ let go !ixSrc !ixDst = if ixSrc < sz
596+ then do
597+ let a = indexPrimArray arr ixSrc
598+ mb <- f a
599+ case mb of
600+ Just b -> do
601+ writePrimArray marr ixDst b
602+ go (ixSrc + 1) (ixDst + 1)
603+ Nothing -> go (ixSrc + 1) ixDst
604+ else return ixDst
605+ lenDst <- go 0 0
606+ marr' <- resizeMutablePrimArray marr lenDst
607+ unsafeFreezePrimArray marr'
608+
609+-- | Generate a primitive array by evaluating the monadic generator function
610+-- at each index.
611+{-# INLINE generatePrimArrayP #-}
612+generatePrimArrayP :: (PrimMonad m, Prim a)
613+ => Int -- ^ length
614+ -> (Int -> m a) -- ^ generator
615+ -> m (PrimArray a)
616+generatePrimArrayP sz f = do
617+ marr <- newPrimArray sz
618+ let go !ix = if ix < sz
619+ then do
620+ b <- f ix
621+ writePrimArray marr ix b
622+ go (ix + 1)
623+ else return ()
624+ go 0
625+ unsafeFreezePrimArray marr
626+
627+-- | Execute the monadic action the given number of times and store the
628+-- results in a primitive array.
629+{-# INLINE replicatePrimArrayP #-}
630+replicatePrimArrayP :: (PrimMonad m, Prim a)
631+ => Int
632+ -> m a
633+ -> m (PrimArray a)
634+replicatePrimArrayP sz f = do
635+ marr <- newPrimArray sz
636+ let go !ix = if ix < sz
637+ then do
638+ b <- f
639+ writePrimArray marr ix b
640+ go (ix + 1)
641+ else return ()
642+ go 0
643+ unsafeFreezePrimArray marr
644+
645+
646+-- | Map over the elements of a primitive array.
647+{-# INLINE mapPrimArray #-}
648+mapPrimArray :: (Prim a, Prim b)
649+ => (a -> b)
650+ -> PrimArray a
651+ -> PrimArray b
652+mapPrimArray f arr = runST $ do
653+ let !sz = sizeofPrimArray arr
654+ marr <- newPrimArray sz
655+ let go !ix = if ix < sz
656+ then do
657+ let b = f (indexPrimArray arr ix)
658+ writePrimArray marr ix b
659+ go (ix + 1)
660+ else return ()
661+ go 0
662+ unsafeFreezePrimArray marr
663+
664+-- | Indexed map over the elements of a primitive array.
665+{-# INLINE imapPrimArray #-}
666+imapPrimArray :: (Prim a, Prim b)
667+ => (Int -> a -> b)
668+ -> PrimArray a
669+ -> PrimArray b
670+imapPrimArray f arr = runST $ do
671+ let !sz = sizeofPrimArray arr
672+ marr <- newPrimArray sz
673+ let go !ix = if ix < sz
674+ then do
675+ let b = f ix (indexPrimArray arr ix)
676+ writePrimArray marr ix b
677+ go (ix + 1)
678+ else return ()
679+ go 0
680+ unsafeFreezePrimArray marr
681+
682+-- | Filter elements of a primitive array according to a predicate.
683+{-# INLINE filterPrimArray #-}
684+filterPrimArray :: Prim a
685+ => (a -> Bool)
686+ -> PrimArray a
687+ -> PrimArray a
688+filterPrimArray p arr = runST $ do
689+ let !sz = sizeofPrimArray arr
690+ marr <- newPrimArray sz
691+ let go !ixSrc !ixDst = if ixSrc < sz
692+ then do
693+ let !a = indexPrimArray arr ixSrc
694+ if p a
695+ then do
696+ writePrimArray marr ixDst a
697+ go (ixSrc + 1) (ixDst + 1)
698+ else go (ixSrc + 1) ixDst
699+ else return ixDst
700+ dstLen <- go 0 0
701+ marr' <- resizeMutablePrimArray marr dstLen
702+ unsafeFreezePrimArray marr'
703+
704+-- | Filter the primitive array, keeping the elements for which the monadic
705+-- predicate evaluates true.
706+filterPrimArrayA ::
707+ (Applicative f, Prim a)
708+ => (a -> f Bool) -- ^ mapping function
709+ -> PrimArray a -- ^ primitive array
710+ -> f (PrimArray a)
711+filterPrimArrayA f = \ !ary ->
712+ let
713+ !len = sizeofPrimArray ary
714+ go !ixSrc
715+ | ixSrc == len = pure $ IxSTA $ \ixDst _ -> return ixDst
716+ | otherwise = let x = indexPrimArray ary ixSrc in
717+ liftA2
718+ (\keep (IxSTA m) -> IxSTA $ \ixDst mary -> if keep
719+ then writePrimArray (MutablePrimArray mary) ixDst x >> m (ixDst + 1) mary
720+ else m ixDst mary
721+ )
722+ (f x)
723+ (go (ixSrc + 1))
724+ in if len == 0
725+ then pure emptyPrimArray
726+ else runIxSTA len <$> go 0
727+
728+-- | Map over the primitive array, keeping the elements for which the applicative
729+-- predicate provides a 'Just'.
730+mapMaybePrimArrayA ::
731+ (Applicative f, Prim a, Prim b)
732+ => (a -> f (Maybe b)) -- ^ mapping function
733+ -> PrimArray a -- ^ primitive array
734+ -> f (PrimArray b)
735+mapMaybePrimArrayA f = \ !ary ->
736+ let
737+ !len = sizeofPrimArray ary
738+ go !ixSrc
739+ | ixSrc == len = pure $ IxSTA $ \ixDst _ -> return ixDst
740+ | otherwise = let x = indexPrimArray ary ixSrc in
741+ liftA2
742+ (\mb (IxSTA m) -> IxSTA $ \ixDst mary -> case mb of
743+ Just b -> writePrimArray (MutablePrimArray mary) ixDst b >> m (ixDst + 1) mary
744+ Nothing -> m ixDst mary
745+ )
746+ (f x)
747+ (go (ixSrc + 1))
748+ in if len == 0
749+ then pure emptyPrimArray
750+ else runIxSTA len <$> go 0
751+
752+-- | Map over a primitive array, optionally discarding some elements. This
753+-- has the same behavior as @Data.Maybe.mapMaybe@.
754+{-# INLINE mapMaybePrimArray #-}
755+mapMaybePrimArray :: (Prim a, Prim b)
756+ => (a -> Maybe b)
757+ -> PrimArray a
758+ -> PrimArray b
759+mapMaybePrimArray p arr = runST $ do
760+ let !sz = sizeofPrimArray arr
761+ marr <- newPrimArray sz
762+ let go !ixSrc !ixDst = if ixSrc < sz
763+ then do
764+ let !a = indexPrimArray arr ixSrc
765+ case p a of
766+ Just b -> do
767+ writePrimArray marr ixDst b
768+ go (ixSrc + 1) (ixDst + 1)
769+ Nothing -> go (ixSrc + 1) ixDst
770+ else return ixDst
771+ dstLen <- go 0 0
772+ marr' <- resizeMutablePrimArray marr dstLen
773+ unsafeFreezePrimArray marr'
774+
775+
776+-- | Traverse a primitive array. The traversal performs all of the applicative
777+-- effects /before/ forcing the resulting values and writing them to the new
778+-- primitive array. Consequently:
779+--
780+-- >>> traversePrimArray (\x -> print x $> bool x undefined (x == 2)) (fromList [1, 2, 3 :: Int])
781+-- 1
782+-- 2
783+-- 3
784+-- *** Exception: Prelude.undefined
785+--
786+-- The function 'traversePrimArrayP' always outperforms this function, but it
787+-- requires a 'PrimMonad' constraint, and it forces the values as
788+-- it performs the effects.
789+traversePrimArray ::
790+ (Applicative f, Prim a, Prim b)
791+ => (a -> f b) -- ^ mapping function
792+ -> PrimArray a -- ^ primitive array
793+ -> f (PrimArray b)
794+traversePrimArray f = \ !ary ->
795+ let
796+ !len = sizeofPrimArray ary
797+ go !i
798+ | i == len = pure $ STA $ \mary -> unsafeFreezePrimArray (MutablePrimArray mary)
799+ | x <- indexPrimArray ary i
800+ = liftA2 (\b (STA m) -> STA $ \mary ->
801+ writePrimArray (MutablePrimArray mary) i b >> m mary)
802+ (f x) (go (i + 1))
803+ in if len == 0
804+ then pure emptyPrimArray
805+ else runSTA len <$> go 0
806+
807+-- | Traverse a primitive array with the index of each element.
808+itraversePrimArray ::
809+ (Applicative f, Prim a, Prim b)
810+ => (Int -> a -> f b)
811+ -> PrimArray a
812+ -> f (PrimArray b)
813+itraversePrimArray f = \ !ary ->
814+ let
815+ !len = sizeofPrimArray ary
816+ go !i
817+ | i == len = pure $ STA $ \mary -> unsafeFreezePrimArray (MutablePrimArray mary)
818+ | x <- indexPrimArray ary i
819+ = liftA2 (\b (STA m) -> STA $ \mary ->
820+ writePrimArray (MutablePrimArray mary) i b >> m mary)
821+ (f i x) (go (i + 1))
822+ in if len == 0
823+ then pure emptyPrimArray
824+ else runSTA len <$> go 0
825+
826+-- | Traverse a primitive array with the indices. The traversal forces the
827+-- resulting values and writes them to the new primitive array as it performs
828+-- the monadic effects.
829+{-# INLINE itraversePrimArrayP #-}
830+itraversePrimArrayP :: (Prim a, Prim b, PrimMonad m)
831+ => (Int -> a -> m b)
832+ -> PrimArray a
833+ -> m (PrimArray b)
834+itraversePrimArrayP f arr = do
835+ let !sz = sizeofPrimArray arr
836+ marr <- newPrimArray sz
837+ let go !ix
838+ | ix < sz = do
839+ writePrimArray marr ix =<< f ix (indexPrimArray arr ix)
840+ go (ix + 1)
841+ | otherwise = return ()
842+ go 0
843+ unsafeFreezePrimArray marr
844+
845+-- | Generate a primitive array.
846+{-# INLINE generatePrimArray #-}
847+generatePrimArray :: Prim a
848+ => Int -- ^ length
849+ -> (Int -> a) -- ^ element from index
850+ -> PrimArray a
851+generatePrimArray len f = runST $ do
852+ marr <- newPrimArray len
853+ let go !ix = if ix < len
854+ then do
855+ writePrimArray marr ix (f ix)
856+ go (ix + 1)
857+ else return ()
858+ go 0
859+ unsafeFreezePrimArray marr
860+
861+-- | Create a primitive array by copying the element the given
862+-- number of times.
863+{-# INLINE replicatePrimArray #-}
864+replicatePrimArray :: Prim a
865+ => Int -- ^ length
866+ -> a -- ^ element
867+ -> PrimArray a
868+replicatePrimArray len a = runST $ do
869+ marr <- newPrimArray len
870+ setPrimArray marr 0 len a
871+ unsafeFreezePrimArray marr
872+
873+-- | Generate a primitive array by evaluating the applicative generator
874+-- function at each index.
875+{-# INLINE generatePrimArrayA #-}
876+generatePrimArrayA ::
877+ (Applicative f, Prim a)
878+ => Int -- ^ length
879+ -> (Int -> f a) -- ^ element from index
880+ -> f (PrimArray a)
881+generatePrimArrayA len f =
882+ let
883+ go !i
884+ | i == len = pure $ STA $ \mary -> unsafeFreezePrimArray (MutablePrimArray mary)
885+ | otherwise
886+ = liftA2 (\b (STA m) -> STA $ \mary ->
887+ writePrimArray (MutablePrimArray mary) i b >> m mary)
888+ (f i) (go (i + 1))
889+ in if len == 0
890+ then pure emptyPrimArray
891+ else runSTA len <$> go 0
892+
893+-- | Execute the applicative action the given number of times and store the
894+-- results in a vector.
895+{-# INLINE replicatePrimArrayA #-}
896+replicatePrimArrayA ::
897+ (Applicative f, Prim a)
898+ => Int -- ^ length
899+ -> f a -- ^ applicative element producer
900+ -> f (PrimArray a)
901+replicatePrimArrayA len f =
902+ let
903+ go !i
904+ | i == len = pure $ STA $ \mary -> unsafeFreezePrimArray (MutablePrimArray mary)
905+ | otherwise
906+ = liftA2 (\b (STA m) -> STA $ \mary ->
907+ writePrimArray (MutablePrimArray mary) i b >> m mary)
908+ f (go (i + 1))
909+ in if len == 0
910+ then pure emptyPrimArray
911+ else runSTA len <$> go 0
912+
913+-- | Traverse the primitive array, discarding the results. There
914+-- is no 'PrimMonad' variant of this function since it would not provide
915+-- any performance benefit.
916+traversePrimArray_ ::
917+ (Applicative f, Prim a)
918+ => (a -> f b)
919+ -> PrimArray a
920+ -> f ()
921+traversePrimArray_ f a = go 0 where
922+ !sz = sizeofPrimArray a
923+ go !ix = if ix < sz
924+ then f (indexPrimArray a ix) *> go (ix + 1)
925+ else pure ()
926+
927+-- | Traverse the primitive array with the indices, discarding the results.
928+-- There is no 'PrimMonad' variant of this function since it would not
929+-- provide any performance benefit.
930+itraversePrimArray_ ::
931+ (Applicative f, Prim a)
932+ => (Int -> a -> f b)
933+ -> PrimArray a
934+ -> f ()
935+itraversePrimArray_ f a = go 0 where
936+ !sz = sizeofPrimArray a
937+ go !ix = if ix < sz
938+ then f ix (indexPrimArray a ix) *> go (ix + 1)
939+ else pure ()
940+
941+newtype IxSTA a = IxSTA {_runIxSTA :: forall s. Int -> MutableByteArray# s -> ST s Int}
942+
943+runIxSTA :: forall a. Prim a
944+ => Int -- maximum possible size
945+ -> IxSTA a
946+ -> PrimArray a
947+runIxSTA !szUpper = \ (IxSTA m) -> runST $ do
948+ ar :: MutablePrimArray s a <- newPrimArray szUpper
949+ sz <- m 0 (unMutablePrimArray ar)
950+ ar' <- resizeMutablePrimArray ar sz
951+ unsafeFreezePrimArray ar'
952+{-# INLINE runIxSTA #-}
953+
954+newtype STA a = STA {_runSTA :: forall s. MutableByteArray# s -> ST s (PrimArray a)}
955+
956+runSTA :: forall a. Prim a => Int -> STA a -> PrimArray a
957+runSTA !sz = \ (STA m) -> runST $ newPrimArray sz >>= \ (ar :: MutablePrimArray s a) -> m (unMutablePrimArray ar)
958+{-# INLINE runSTA #-}
959+
960+unMutablePrimArray :: MutablePrimArray s a -> MutableByteArray# s
961+unMutablePrimArray (MutablePrimArray m) = m
962+
963+{- $effectfulMapCreate
964+The naming conventions adopted in this section are explained in the
965+documentation of the @Data.Primitive@ module.
966+-}