@@ -1,5 +1,4 @@
1{-# LANGUAGE BangPatterns #-}
1{-# LANGUAGE BangPatterns #-}
3{-# LANGUAGE MagicHash #-}
2{-# LANGUAGE MagicHash #-}
4{-# LANGUAGE RankNTypes #-}
3{-# LANGUAGE RankNTypes #-}
5{-# LANGUAGE ScopedTypeVariables #-}
4{-# LANGUAGE ScopedTypeVariables #-}
@@ -33,83 +32,23 @@ module Streamly.Internal.Data.Prim.Array.Types
33 , resizeMutablePrimArray
32 , resizeMutablePrimArray
34 , shrinkMutablePrimArray
33 , shrinkMutablePrimArray
35 -- * Element Access
34 -- * Element Access
37 , writePrimArray
35 , writePrimArray
38 , indexPrimArray
36 , indexPrimArray
39 -- * Freezing and Thawing
37 -- * Freezing and Thawing
40 , unsafeFreezePrimArray
38 , unsafeFreezePrimArray
41− , unsafeThawPrimArray
42− -- * Block Operations
44− , copyMutablePrimArray
46− , copyMutablePrimArrayToPtr
48 -- * Information
39 -- * Information
49− , sameMutablePrimArray
50− , getSizeofMutablePrimArray
51− , sizeofMutablePrimArray
52 , sizeofPrimArray
40 , sizeofPrimArray
57 -- * Folding
41 -- * Folding
58 , foldrPrimArray
42 , foldrPrimArray
61 , foldlPrimArray'
43 , foldlPrimArray'
63− -- * Effectful Folding
65− , itraversePrimArray_
73− -- * Effectful Map/Create
74− -- $effectfulMapCreate
75− -- ** Lazy Applicative
79− , replicatePrimArrayA
82− -- ** Strict Primitive Monadic
84− , itraversePrimArrayP
86− , replicatePrimArrayP
91import GHC.Exts
46import GHC.Exts
93import Data.Primitive.Types
48import Data.Primitive.Types
94import Data.Primitive.ByteArray (ByteArray(..))
49import Data.Primitive.ByteArray (ByteArray(..))
95−#if !MIN_VERSION_base(4,11,0)
96−import Data.Monoid (Monoid(..),(<>))
98−import Control.Applicative
99import Control.Monad.Primitive
50import Control.Monad.Primitive
100−import Control.Monad.ST
101−import qualified Data.List as L
102import qualified Data.Primitive.ByteArray as PB
51import qualified Data.Primitive.ByteArray as PB
103−import qualified Data.Primitive.Types as PT
105−#if MIN_VERSION_base(4,9,0) && !MIN_VERSION_base(4,11,0)
106−import Data.Semigroup (Semigroup)
109−#if MIN_VERSION_base(4,9,0) && MIN_VERSION_primitive(0,6,4)
110−import qualified Data.Semigroup as SG
114-- | Arrays of unboxed elements. This accepts types like 'Double', 'Char',
53-- | Arrays of unboxed elements. This accepts types like 'Double', 'Char',
115-- 'Int', and 'Word', as well as their fixed-length variants ('Word8',
54-- 'Int', and 'Word', as well as their fixed-length variants ('Word8',
@@ -157,88 +96,28 @@ instance (Ord a, Prim a) => Ord (PrimArray a) where
157 | sameByteArray ba1# ba2# = EQ
96 | sameByteArray ba1# ba2# = EQ
158 | otherwise = loop 0
97 | otherwise = loop 0
160 sz1 = PB.sizeofByteArray (ByteArray ba1#)
102 sz1 = PB.sizeofByteArray (ByteArray ba1#)
161 sz2 = PB.sizeofByteArray (ByteArray ba2#)
103 sz2 = PB.sizeofByteArray (ByteArray ba2#)
162 sz = quot (min sz1 sz2) (sizeOf (undefined :: a))
104 sz = quot (min sz1 sz2) (sizeOf (undefined :: a))
164− | i < sz = compare (indexPrimArray a1 i) (indexPrimArray a2 i) <> loop (i+1)
106+ | i < sz = compare (indexPrimArray a1 i) (indexPrimArray a2 i) `cmp` loop (i+1)
165 | otherwise = compare sz1 sz2
107 | otherwise = compare sz1 sz2
166 {-# INLINE compare #-}
108 {-# INLINE compare #-}
169−instance Prim a => IsList (PrimArray a) where
170− type Item (PrimArray a) = a
171− fromList = primArrayFromList
172− fromListN = primArrayFromListN
173− toList = primArrayToList
175-- | @since 0.6.4.0
110-- | @since 0.6.4.0
176instance (Show a, Prim a) => Show (PrimArray a) where
111instance (Show a, Prim a) => Show (PrimArray a) where
177 showsPrec p a = showParen (p > 10) $
112 showsPrec p a = showParen (p > 10) $
178 showString "fromListN " . shows (sizeofPrimArray a) . showString " "
113 showString "fromListN " . shows (sizeofPrimArray a) . showString " "
179 . shows (primArrayToList a)
114 . shows (primArrayToList a)
181−die :: String -> String -> a
182−die fun problem = error $ "Data.Primitive.PrimArray." ++ fun ++ ": " ++ problem
184−primArrayFromList :: Prim a => [a] -> PrimArray a
185−primArrayFromList vs = primArrayFromListN (L.length vs) vs
187−primArrayFromListN :: forall a. Prim a => Int -> [a] -> PrimArray a
188−primArrayFromListN len vs = runST run where
189− run :: forall s. ST s (PrimArray a)
191− arr <- newPrimArray len
192− let go :: [a] -> Int -> ST s ()
193− go [] !ix = if ix == len
195− else die "fromListN" "list length less than specified size"
196− go (a : as) !ix = if ix < len
198− writePrimArray arr ix a
200− else die "fromListN" "list length greater than specified size"
202− unsafeFreezePrimArray arr
204-- | Convert the primitive array to a list.
116-- | Convert the primitive array to a list.
205{-# INLINE primArrayToList #-}
117{-# INLINE primArrayToList #-}
206primArrayToList :: forall a. Prim a => PrimArray a -> [a]
118primArrayToList :: forall a. Prim a => PrimArray a -> [a]
207primArrayToList xs = build (\c n -> foldrPrimArray c n xs)
119primArrayToList xs = build (\c n -> foldrPrimArray c n xs)
209−#if MIN_VERSION_primitive(0,6,4)
211−primArrayToByteArray :: PrimArray a -> PB.ByteArray
212−primArrayToByteArray (PrimArray x) = PB.ByteArray x
214−byteArrayToPrimArray :: ByteArray -> PrimArray a
215−byteArrayToPrimArray (PB.ByteArray x) = PrimArray x
217−#if MIN_VERSION_base(4,9,0)
219−instance Semigroup (PrimArray a) where
220− x <> y = byteArrayToPrimArray (primArrayToByteArray x SG.<> primArrayToByteArray y)
221− sconcat = byteArrayToPrimArray . SG.sconcat . fmap primArrayToByteArray
222− stimes i arr = byteArrayToPrimArray (SG.stimes i (primArrayToByteArray arr))
226−instance Monoid (PrimArray a) where
227− mempty = emptyPrimArray
228−#if !(MIN_VERSION_base(4,11,0))
229− mappend x y = byteArrayToPrimArray (mappend (primArrayToByteArray x) (primArrayToByteArray y))
231− mconcat = byteArrayToPrimArray . mconcat . map primArrayToByteArray
235−-- | The empty primitive array.
236−emptyPrimArray :: PrimArray a
237−{-# NOINLINE emptyPrimArray #-}
238−emptyPrimArray = runST $ primitive $ \s0# -> case newByteArray# 0# s0# of
239− (# s1#, arr# #) -> case unsafeFreezeByteArray# arr# s1# of
240− (# s2#, arr'# #) -> (# s2#, PrimArray arr'# #)
242-- | Create a new mutable primitive array of the given length. The
121-- | Create a new mutable primitive array of the given length. The
243-- underlying memory is left uninitialized.
122-- underlying memory is left uninitialized.
244newPrimArray :: forall m a. (PrimMonad m, Prim a) => Int -> m (MutablePrimArray (PrimState m) a)
123newPrimArray :: forall m a. (PrimMonad m, Prim a) => Int -> m (MutablePrimArray (PrimState m) a)
@@ -282,11 +161,6 @@ shrinkMutablePrimArray :: forall m a. (PrimMonad m, Prim a)
282shrinkMutablePrimArray (MutablePrimArray arr#) (I# n#)
161shrinkMutablePrimArray (MutablePrimArray arr#) (I# n#)
283 = primitive_ (shrinkMutableByteArray# arr# (n# *# sizeOf# (undefined :: a)))
162 = primitive_ (shrinkMutableByteArray# arr# (n# *# sizeOf# (undefined :: a)))
285−readPrimArray :: (Prim a, PrimMonad m) => MutablePrimArray (PrimState m) a -> Int -> m a
286−{-# INLINE readPrimArray #-}
287−readPrimArray (MutablePrimArray arr#) (I# i#)
288− = primitive (readByteArray# arr# i#)
290-- | Write an element to the given index.
164-- | Write an element to the given index.
291writePrimArray ::
165writePrimArray ::
292 (Prim a, PrimMonad m)
166 (Prim a, PrimMonad m)
@@ -298,128 +172,6 @@ writePrimArray ::
298writePrimArray (MutablePrimArray arr#) (I# i#) x
172writePrimArray (MutablePrimArray arr#) (I# i#) x
299 = primitive_ (writeByteArray# arr# i# x)
173 = primitive_ (writeByteArray# arr# i# x)
301−-- | Copy part of a mutable array into another mutable array.
302−-- In the case that the destination and
303−-- source arrays are the same, the regions may overlap.
304−copyMutablePrimArray :: forall m a.
305− (PrimMonad m, Prim a)
306− => MutablePrimArray (PrimState m) a -- ^ destination array
307− -> Int -- ^ offset into destination array
308− -> MutablePrimArray (PrimState m) a -- ^ source array
309− -> Int -- ^ offset into source array
310− -> Int -- ^ number of elements to copy
312−{-# INLINE copyMutablePrimArray #-}
313−copyMutablePrimArray (MutablePrimArray dst#) (I# doff#) (MutablePrimArray src#) (I# soff#) (I# n#)
314− = primitive_ (copyMutableByteArray#
316− (soff# *# (sizeOf# (undefined :: a)))
318− (doff# *# (sizeOf# (undefined :: a)))
319− (n# *# (sizeOf# (undefined :: a)))
322−-- | Copy part of an array into another mutable array.
323−copyPrimArray :: forall m a.
324− (PrimMonad m, Prim a)
325− => MutablePrimArray (PrimState m) a -- ^ destination array
326− -> Int -- ^ offset into destination array
327− -> PrimArray a -- ^ source array
328− -> Int -- ^ offset into source array
329− -> Int -- ^ number of elements to copy
331−{-# INLINE copyPrimArray #-}
332−copyPrimArray (MutablePrimArray dst#) (I# doff#) (PrimArray src#) (I# soff#) (I# n#)
333− = primitive_ (copyByteArray#
335− (soff# *# (sizeOf# (undefined :: a)))
337− (doff# *# (sizeOf# (undefined :: a)))
338− (n# *# (sizeOf# (undefined :: a)))
341−-- | Copy a slice of an immutable primitive array to an address.
342−-- The offset and length are given in elements of type @a@.
343−-- This function assumes that the 'Prim' instance of @a@
344−-- agrees with the 'Storable' instance. This function is only
345−-- available when building with GHC 7.8 or newer.
346−copyPrimArrayToPtr :: forall m a. (PrimMonad m, Prim a)
347− => Ptr a -- ^ destination pointer
348− -> PrimArray a -- ^ source array
349− -> Int -- ^ offset into source array
350− -> Int -- ^ number of prims to copy
352−{-# INLINE copyPrimArrayToPtr #-}
353−copyPrimArrayToPtr (Ptr addr#) (PrimArray ba#) (I# soff#) (I# n#) =
355− let s'# = copyByteArrayToAddr# ba# (soff# *# siz#) addr# (n# *# siz#) s#
357− where siz# = sizeOf# (undefined :: a)
359−-- | Copy a slice of an immutable primitive array to an address.
360−-- The offset and length are given in elements of type @a@.
361−-- This function assumes that the 'Prim' instance of @a@
362−-- agrees with the 'Storable' instance. This function is only
363−-- available when building with GHC 7.8 or newer.
364−copyMutablePrimArrayToPtr :: forall m a. (PrimMonad m, Prim a)
365− => Ptr a -- ^ destination pointer
366− -> MutablePrimArray (PrimState m) a -- ^ source array
367− -> Int -- ^ offset into source array
368− -> Int -- ^ number of prims to copy
370−{-# INLINE copyMutablePrimArrayToPtr #-}
371−copyMutablePrimArrayToPtr (Ptr addr#) (MutablePrimArray mba#) (I# soff#) (I# n#) =
373− let s'# = copyMutableByteArrayToAddr# mba# (soff# *# siz#) addr# (n# *# siz#) s#
375− where siz# = sizeOf# (undefined :: a)
377−-- | Fill a slice of a mutable primitive array with a value.
379− :: (Prim a, PrimMonad m)
380− => MutablePrimArray (PrimState m) a -- ^ array to fill
381− -> Int -- ^ offset into array
382− -> Int -- ^ number of values to fill
383− -> a -- ^ value to fill with
385−{-# INLINE setPrimArray #-}
386−setPrimArray (MutablePrimArray dst#) (I# doff#) (I# sz#) x
387− = primitive_ (PT.setByteArray# dst# doff# sz# x)
389−-- | Get the size of a mutable primitive array in elements. Unlike 'sizeofMutablePrimArray',
390−-- this function ensures sequencing in the presence of resizing.
391−getSizeofMutablePrimArray :: forall m a. (PrimMonad m, Prim a)
392− => MutablePrimArray (PrimState m) a -- ^ array
394−{-# INLINE getSizeofMutablePrimArray #-}
395−#if __GLASGOW_HASKELL__ >= 801
396−getSizeofMutablePrimArray (MutablePrimArray arr#)
398− case getSizeofMutableByteArray# arr# s# of
399− (# s'#, sz# #) -> (# s'#, I# (quotInt# sz# (sizeOf# (undefined :: a))) #)
402−-- On older GHCs, it is not possible to resize a byte array, so
403−-- this provides behavior consistent with the implementation for
405−getSizeofMutablePrimArray arr
406− = return (sizeofMutablePrimArray arr)
409−-- | Size of the mutable primitive array in elements. This function shall not
410−-- be used on primitive arrays that are an argument to or a result of
411−-- 'resizeMutablePrimArray' or 'shrinkMutablePrimArray'.
412−sizeofMutablePrimArray :: forall s a. Prim a => MutablePrimArray s a -> Int
413−{-# INLINE sizeofMutablePrimArray #-}
414−sizeofMutablePrimArray (MutablePrimArray arr#) =
415− I# (quotInt# (sizeofMutableByteArray# arr#) (sizeOf# (undefined :: a)))
417−-- | Check if the two arrays refer to the same memory block.
418−sameMutablePrimArray :: MutablePrimArray s a -> MutablePrimArray s a -> Bool
419−{-# INLINE sameMutablePrimArray #-}
420−sameMutablePrimArray (MutablePrimArray arr#) (MutablePrimArray brr#)
421− = isTrue# (sameMutableByteArray# arr# brr#)
423-- | Convert a mutable byte array to an immutable one without copying. The
175-- | Convert a mutable byte array to an immutable one without copying. The
424-- array should not be modified after the conversion.
176-- array should not be modified after the conversion.
425unsafeFreezePrimArray
177unsafeFreezePrimArray
@@ -429,14 +181,6 @@ unsafeFreezePrimArray (MutablePrimArray arr#)
429 = primitive (\s# -> case unsafeFreezeByteArray# arr# s# of
181 = primitive (\s# -> case unsafeFreezeByteArray# arr# s# of
430 (# s'#, arr'# #) -> (# s'#, PrimArray arr'# #))
182 (# s'#, arr'# #) -> (# s'#, PrimArray arr'# #))
432−-- | Convert an immutable array to a mutable one without copying. The
433−-- original array should not be used after the conversion.
435− :: PrimMonad m => PrimArray a -> m (MutablePrimArray (PrimState m) a)
436−{-# INLINE unsafeThawPrimArray #-}
437−unsafeThawPrimArray (PrimArray arr#)
438− = primitive (\s# -> (# s#, MutablePrimArray (unsafeCoerce# arr#) #))
440-- | Read a primitive value from the primitive array.
184-- | Read a primitive value from the primitive array.
441indexPrimArray :: forall a. Prim a => PrimArray a -> Int -> a
185indexPrimArray :: forall a. Prim a => PrimArray a -> Int -> a
442{-# INLINE indexPrimArray #-}
186{-# INLINE indexPrimArray #-}
@@ -457,24 +201,6 @@ foldrPrimArray f z arr = go 0
457 | sz > i = f (indexPrimArray arr i) (go (i+1))
201 | sz > i = f (indexPrimArray arr i) (go (i+1))
458 | otherwise = z
202 | otherwise = z
460−-- | Strict right-associated fold over the elements of a 'PrimArray'.
461−{-# INLINE foldrPrimArray' #-}
462−foldrPrimArray' :: forall a b. Prim a => (a -> b -> b) -> b -> PrimArray a -> b
463−foldrPrimArray' f z0 arr = go (sizeofPrimArray arr - 1) z0
467− | otherwise = go (i - 1) (f (indexPrimArray arr i) acc)
469−-- | Lazy left-associated fold over the elements of a 'PrimArray'.
470−{-# INLINE foldlPrimArray #-}
471−foldlPrimArray :: forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
472−foldlPrimArray f z arr = go (sizeofPrimArray arr - 1)
476− | otherwise = f (go (i - 1)) (indexPrimArray arr i)
478-- | Strict left-associated fold over the elements of a 'PrimArray'.
204-- | Strict left-associated fold over the elements of a 'PrimArray'.
479{-# INLINE foldlPrimArray' #-}
205{-# INLINE foldlPrimArray' #-}
480foldlPrimArray' :: forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
206foldlPrimArray' :: forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
@@ -484,466 +210,3 @@ foldlPrimArray' f z0 arr = go 0 z0
484 go !i !acc
210 go !i !acc
485 | i < sz = go (i + 1) (f acc (indexPrimArray arr i))
211 | i < sz = go (i + 1) (f acc (indexPrimArray arr i))
486 | otherwise = acc
212 | otherwise = acc
488−-- | Strict left-associated fold over the elements of a 'PrimArray'.
489−{-# INLINE foldlPrimArrayM' #-}
490−foldlPrimArrayM' :: (Prim a, Monad m) => (b -> a -> m b) -> b -> PrimArray a -> m b
491−foldlPrimArrayM' f z0 arr = go 0 z0
493− !sz = sizeofPrimArray arr
496− acc2 <- f acc1 (indexPrimArray arr i)
498− | otherwise = return acc1
500−-- | Traverse a primitive array. The traversal forces the resulting values and
501−-- writes them to the new primitive array as it performs the monadic effects.
504−-- >>> traversePrimArrayP (\x -> print x $> bool x undefined (x == 2)) (fromList [1, 2, 3 :: Int])
507−-- *** Exception: Prelude.undefined
509−-- In many situations, 'traversePrimArrayP' can replace 'traversePrimArray',
510−-- changing the strictness characteristics of the traversal but typically improving
511−-- the performance. Consider the following short-circuiting traversal:
513−-- > incrPositiveA :: PrimArray Int -> Maybe (PrimArray Int)
514−-- > incrPositiveA xs = traversePrimArray (\x -> bool Nothing (Just (x + 1)) (x > 0)) xs
516−-- This can be rewritten using 'traversePrimArrayP'. To do this, we must
517−-- change the traversal context to @MaybeT (ST s)@, which has a 'PrimMonad'
520−-- > incrPositiveB :: PrimArray Int -> Maybe (PrimArray Int)
521−-- > incrPositiveB xs = runST $ runMaybeT $ traversePrimArrayP
522−-- > (\x -> bool (MaybeT (return Nothing)) (MaybeT (return (Just (x + 1)))) (x > 0))
525−-- Benchmarks demonstrate that the second implementation runs 150 times
526−-- faster than the first. It also results in fewer allocations.
527−{-# INLINE traversePrimArrayP #-}
528−traversePrimArrayP :: (PrimMonad m, Prim a, Prim b)
532−traversePrimArrayP f arr = do
533− let !sz = sizeofPrimArray arr
534− marr <- newPrimArray sz
535− let go !ix = if ix < sz
537− b <- f (indexPrimArray arr ix)
538− writePrimArray marr ix b
542− unsafeFreezePrimArray marr
544−-- | Filter the primitive array, keeping the elements for which the monadic
545−-- predicate evaluates true.
546−{-# INLINE filterPrimArrayP #-}
547−filterPrimArrayP :: (PrimMonad m, Prim a)
551−filterPrimArrayP f arr = do
552− let !sz = sizeofPrimArray arr
553− marr <- newPrimArray sz
554− let go !ixSrc !ixDst = if ixSrc < sz
556− let a = indexPrimArray arr ixSrc
560− writePrimArray marr ixDst a
561− go (ixSrc + 1) (ixDst + 1)
562− else go (ixSrc + 1) ixDst
565− marr' <- resizeMutablePrimArray marr lenDst
566− unsafeFreezePrimArray marr'
568−-- | Map over the primitive array, keeping the elements for which the monadic
569−-- predicate provides a 'Just'.
570−{-# INLINE mapMaybePrimArrayP #-}
571−mapMaybePrimArrayP :: (PrimMonad m, Prim a, Prim b)
572− => (a -> m (Maybe b))
575−mapMaybePrimArrayP f arr = do
576− let !sz = sizeofPrimArray arr
577− marr <- newPrimArray sz
578− let go !ixSrc !ixDst = if ixSrc < sz
580− let a = indexPrimArray arr ixSrc
584− writePrimArray marr ixDst b
585− go (ixSrc + 1) (ixDst + 1)
586− Nothing -> go (ixSrc + 1) ixDst
589− marr' <- resizeMutablePrimArray marr lenDst
590− unsafeFreezePrimArray marr'
592−-- | Generate a primitive array by evaluating the monadic generator function
594−{-# INLINE generatePrimArrayP #-}
595−generatePrimArrayP :: (PrimMonad m, Prim a)
597− -> (Int -> m a) -- ^ generator
599−generatePrimArrayP sz f = do
600− marr <- newPrimArray sz
601− let go !ix = if ix < sz
604− writePrimArray marr ix b
608− unsafeFreezePrimArray marr
610−-- | Execute the monadic action the given number of times and store the
611−-- results in a primitive array.
612−{-# INLINE replicatePrimArrayP #-}
613−replicatePrimArrayP :: (PrimMonad m, Prim a)
617−replicatePrimArrayP sz f = do
618− marr <- newPrimArray sz
619− let go !ix = if ix < sz
622− writePrimArray marr ix b
626− unsafeFreezePrimArray marr
629−-- | Map over the elements of a primitive array.
630−{-# INLINE mapPrimArray #-}
631−mapPrimArray :: (Prim a, Prim b)
635−mapPrimArray f arr = runST $ do
636− let !sz = sizeofPrimArray arr
637− marr <- newPrimArray sz
638− let go !ix = if ix < sz
640− let b = f (indexPrimArray arr ix)
641− writePrimArray marr ix b
645− unsafeFreezePrimArray marr
647−-- | Indexed map over the elements of a primitive array.
648−{-# INLINE imapPrimArray #-}
649−imapPrimArray :: (Prim a, Prim b)
653−imapPrimArray f arr = runST $ do
654− let !sz = sizeofPrimArray arr
655− marr <- newPrimArray sz
656− let go !ix = if ix < sz
658− let b = f ix (indexPrimArray arr ix)
659− writePrimArray marr ix b
663− unsafeFreezePrimArray marr
665−-- | Filter elements of a primitive array according to a predicate.
666−{-# INLINE filterPrimArray #-}
667−filterPrimArray :: Prim a
671−filterPrimArray p arr = runST $ do
672− let !sz = sizeofPrimArray arr
673− marr <- newPrimArray sz
674− let go !ixSrc !ixDst = if ixSrc < sz
676− let !a = indexPrimArray arr ixSrc
679− writePrimArray marr ixDst a
680− go (ixSrc + 1) (ixDst + 1)
681− else go (ixSrc + 1) ixDst
684− marr' <- resizeMutablePrimArray marr dstLen
685− unsafeFreezePrimArray marr'
687−-- | Filter the primitive array, keeping the elements for which the monadic
688−-- predicate evaluates true.
690− (Applicative f, Prim a)
691− => (a -> f Bool) -- ^ mapping function
692− -> PrimArray a -- ^ primitive array
694−filterPrimArrayA f = \ !ary ->
696− !len = sizeofPrimArray ary
698− | ixSrc == len = pure $ IxSTA $ \ixDst _ -> return ixDst
699− | otherwise = let x = indexPrimArray ary ixSrc in
701− (\keep (IxSTA m) -> IxSTA $ \ixDst mary -> if keep
702− then writePrimArray (MutablePrimArray mary) ixDst x >> m (ixDst + 1) mary
708− then pure emptyPrimArray
709− else runIxSTA len <$> go 0
711−-- | Map over the primitive array, keeping the elements for which the applicative
712−-- predicate provides a 'Just'.
713−mapMaybePrimArrayA ::
714− (Applicative f, Prim a, Prim b)
715− => (a -> f (Maybe b)) -- ^ mapping function
716− -> PrimArray a -- ^ primitive array
718−mapMaybePrimArrayA f = \ !ary ->
720− !len = sizeofPrimArray ary
722− | ixSrc == len = pure $ IxSTA $ \ixDst _ -> return ixDst
723− | otherwise = let x = indexPrimArray ary ixSrc in
725− (\mb (IxSTA m) -> IxSTA $ \ixDst mary -> case mb of
726− Just b -> writePrimArray (MutablePrimArray mary) ixDst b >> m (ixDst + 1) mary
727− Nothing -> m ixDst mary
732− then pure emptyPrimArray
733− else runIxSTA len <$> go 0
735−-- | Map over a primitive array, optionally discarding some elements. This
736−-- has the same behavior as @Data.Maybe.mapMaybe@.
737−{-# INLINE mapMaybePrimArray #-}
738−mapMaybePrimArray :: (Prim a, Prim b)
742−mapMaybePrimArray p arr = runST $ do
743− let !sz = sizeofPrimArray arr
744− marr <- newPrimArray sz
745− let go !ixSrc !ixDst = if ixSrc < sz
747− let !a = indexPrimArray arr ixSrc
750− writePrimArray marr ixDst b
751− go (ixSrc + 1) (ixDst + 1)
752− Nothing -> go (ixSrc + 1) ixDst
755− marr' <- resizeMutablePrimArray marr dstLen
756− unsafeFreezePrimArray marr'
759−-- | Traverse a primitive array. The traversal performs all of the applicative
760−-- effects /before/ forcing the resulting values and writing them to the new
761−-- primitive array. Consequently:
763−-- >>> traversePrimArray (\x -> print x $> bool x undefined (x == 2)) (fromList [1, 2, 3 :: Int])
767−-- *** Exception: Prelude.undefined
769−-- The function 'traversePrimArrayP' always outperforms this function, but it
770−-- requires a 'PrimMonad' constraint, and it forces the values as
771−-- it performs the effects.
773− (Applicative f, Prim a, Prim b)
774− => (a -> f b) -- ^ mapping function
775− -> PrimArray a -- ^ primitive array
777−traversePrimArray f = \ !ary ->
779− !len = sizeofPrimArray ary
781− | i == len = pure $ STA $ \mary -> unsafeFreezePrimArray (MutablePrimArray mary)
782− | x <- indexPrimArray ary i
783− = liftA2 (\b (STA m) -> STA $ \mary ->
784− writePrimArray (MutablePrimArray mary) i b >> m mary)
787− then pure emptyPrimArray
788− else runSTA len <$> go 0
790−-- | Traverse a primitive array with the index of each element.
791−itraversePrimArray ::
792− (Applicative f, Prim a, Prim b)
793− => (Int -> a -> f b)
796−itraversePrimArray f = \ !ary ->
798− !len = sizeofPrimArray ary
800− | i == len = pure $ STA $ \mary -> unsafeFreezePrimArray (MutablePrimArray mary)
801− | x <- indexPrimArray ary i
802− = liftA2 (\b (STA m) -> STA $ \mary ->
803− writePrimArray (MutablePrimArray mary) i b >> m mary)
804− (f i x) (go (i + 1))
806− then pure emptyPrimArray
807− else runSTA len <$> go 0
809−-- | Traverse a primitive array with the indices. The traversal forces the
810−-- resulting values and writes them to the new primitive array as it performs
811−-- the monadic effects.
812−{-# INLINE itraversePrimArrayP #-}
813−itraversePrimArrayP :: (Prim a, Prim b, PrimMonad m)
814− => (Int -> a -> m b)
817−itraversePrimArrayP f arr = do
818− let !sz = sizeofPrimArray arr
819− marr <- newPrimArray sz
822− writePrimArray marr ix =<< f ix (indexPrimArray arr ix)
824− | otherwise = return ()
826− unsafeFreezePrimArray marr
828−-- | Generate a primitive array.
829−{-# INLINE generatePrimArray #-}
830−generatePrimArray :: Prim a
832− -> (Int -> a) -- ^ element from index
834−generatePrimArray len f = runST $ do
835− marr <- newPrimArray len
836− let go !ix = if ix < len
838− writePrimArray marr ix (f ix)
842− unsafeFreezePrimArray marr
844−-- | Create a primitive array by copying the element the given
846−{-# INLINE replicatePrimArray #-}
847−replicatePrimArray :: Prim a
851−replicatePrimArray len a = runST $ do
852− marr <- newPrimArray len
853− setPrimArray marr 0 len a
854− unsafeFreezePrimArray marr
856−-- | Generate a primitive array by evaluating the applicative generator
857−-- function at each index.
858−{-# INLINE generatePrimArrayA #-}
859−generatePrimArrayA ::
860− (Applicative f, Prim a)
862− -> (Int -> f a) -- ^ element from index
864−generatePrimArrayA len f =
867− | i == len = pure $ STA $ \mary -> unsafeFreezePrimArray (MutablePrimArray mary)
869− = liftA2 (\b (STA m) -> STA $ \mary ->
870− writePrimArray (MutablePrimArray mary) i b >> m mary)
873− then pure emptyPrimArray
874− else runSTA len <$> go 0
876−-- | Execute the applicative action the given number of times and store the
877−-- results in a vector.
878−{-# INLINE replicatePrimArrayA #-}
879−replicatePrimArrayA ::
880− (Applicative f, Prim a)
882− -> f a -- ^ applicative element producer
884−replicatePrimArrayA len f =
887− | i == len = pure $ STA $ \mary -> unsafeFreezePrimArray (MutablePrimArray mary)
889− = liftA2 (\b (STA m) -> STA $ \mary ->
890− writePrimArray (MutablePrimArray mary) i b >> m mary)
893− then pure emptyPrimArray
894− else runSTA len <$> go 0
896−-- | Traverse the primitive array, discarding the results. There
897−-- is no 'PrimMonad' variant of this function since it would not provide
898−-- any performance benefit.
899−traversePrimArray_ ::
900− (Applicative f, Prim a)
904−traversePrimArray_ f a = go 0 where
905− !sz = sizeofPrimArray a
907− then f (indexPrimArray a ix) *> go (ix + 1)
910−-- | Traverse the primitive array with the indices, discarding the results.
911−-- There is no 'PrimMonad' variant of this function since it would not
912−-- provide any performance benefit.
913−itraversePrimArray_ ::
914− (Applicative f, Prim a)
915− => (Int -> a -> f b)
918−itraversePrimArray_ f a = go 0 where
919− !sz = sizeofPrimArray a
921− then f ix (indexPrimArray a ix) *> go (ix + 1)
924−newtype IxSTA a = IxSTA {_runIxSTA :: forall s. Int -> MutableByteArray# s -> ST s Int}
926−runIxSTA :: forall a. Prim a
927− => Int -- maximum possible size
930−runIxSTA !szUpper = \ (IxSTA m) -> runST $ do
931− ar :: MutablePrimArray s a <- newPrimArray szUpper
932− sz <- m 0 (unMutablePrimArray ar)
933− ar' <- resizeMutablePrimArray ar sz
934− unsafeFreezePrimArray ar'
935−{-# INLINE runIxSTA #-}
937−newtype STA a = STA {_runSTA :: forall s. MutableByteArray# s -> ST s (PrimArray a)}
939−runSTA :: forall a. Prim a => Int -> STA a -> PrimArray a
940−runSTA !sz = \ (STA m) -> runST $ newPrimArray sz >>= \ (ar :: MutablePrimArray s a) -> m (unMutablePrimArray ar)
941−{-# INLINE runSTA #-}
943−unMutablePrimArray :: MutablePrimArray s a -> MutableByteArray# s
944−unMutablePrimArray (MutablePrimArray m) = m
946−{- $effectfulMapCreate
947−The naming conventions adopted in this section are explained in the
948−documentation of the @Data.Primitive@ module.