@@ -1,5 +1,4 @@
11{-# LANGUAGE BangPatterns #-}
2−{-# LANGUAGE CPP #-}
32{-# LANGUAGE MagicHash #-}
43{-# LANGUAGE RankNTypes #-}
54{-# LANGUAGE ScopedTypeVariables #-}
@@ -33,83 +32,23 @@ module Streamly.Internal.Data.Prim.Array.Types
3332 , resizeMutablePrimArray
3433 , shrinkMutablePrimArray
3534 -- * Element Access
36− , readPrimArray
3735 , writePrimArray
3836 , indexPrimArray
3937 -- * Freezing and Thawing
4038 , unsafeFreezePrimArray
41− , unsafeThawPrimArray
42− -- * Block Operations
43− , copyPrimArray
44− , copyMutablePrimArray
45− , copyPrimArrayToPtr
46− , copyMutablePrimArrayToPtr
47− , setPrimArray
4839 -- * Information
49− , sameMutablePrimArray
50− , getSizeofMutablePrimArray
51− , sizeofMutablePrimArray
5240 , sizeofPrimArray
53− -- * List Conversion
54− , primArrayToList
55− , primArrayFromList
56− , primArrayFromListN
5741 -- * Folding
5842 , foldrPrimArray
59− , foldrPrimArray'
60− , foldlPrimArray
6143 , foldlPrimArray'
62− , foldlPrimArrayM'
63− -- * Effectful Folding
64− , traversePrimArray_
65− , itraversePrimArray_
66− -- * Map/Create
67− , mapPrimArray
68− , imapPrimArray
69− , generatePrimArray
70− , replicatePrimArray
71− , filterPrimArray
72− , mapMaybePrimArray
73− -- * Effectful Map/Create
74− -- $effectfulMapCreate
75− -- ** Lazy Applicative
76− , traversePrimArray
77− , itraversePrimArray
78− , generatePrimArrayA
79− , replicatePrimArrayA
80− , filterPrimArrayA
81− , mapMaybePrimArrayA
82− -- ** Strict Primitive Monadic
83− , traversePrimArrayP
84− , itraversePrimArrayP
85− , generatePrimArrayP
86− , replicatePrimArrayP
87− , filterPrimArrayP
88− , mapMaybePrimArrayP
8944 ) where
9045
9146import GHC.Exts
9247
9348import Data.Primitive.Types
9449import Data.Primitive.ByteArray (ByteArray(..))
95−#if !MIN_VERSION_base(4,11,0)
96−import Data.Monoid (Monoid(..),(<>))
97−#endif
98−import Control.Applicative
9950import Control.Monad.Primitive
100−import Control.Monad.ST
101−import qualified Data.List as L
10251import qualified Data.Primitive.ByteArray as PB
103−import qualified Data.Primitive.Types as PT
104−
105−#if MIN_VERSION_base(4,9,0) && !MIN_VERSION_base(4,11,0)
106−import Data.Semigroup (Semigroup)
107−#endif
108−
109−#if MIN_VERSION_base(4,9,0) && MIN_VERSION_primitive(0,6,4)
110−import qualified Data.Semigroup as SG
111−#endif
112−
11352
11453-- | Arrays of unboxed elements. This accepts types like 'Double', 'Char',
11554-- 'Int', and 'Word', as well as their fixed-length variants ('Word8',
@@ -157,88 +96,28 @@ instance (Ord a, Prim a) => Ord (PrimArray a) where
15796 | sameByteArray ba1# ba2# = EQ
15897 | otherwise = loop 0
15998 where
99+ cmp LT _ = LT
100+ cmp EQ y = y
101+ cmp GT _ = GT
160102 sz1 = PB.sizeofByteArray (ByteArray ba1#)
161103 sz2 = PB.sizeofByteArray (ByteArray ba2#)
162104 sz = quot (min sz1 sz2) (sizeOf (undefined :: a))
163105 loop !i
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)
165107 | otherwise = compare sz1 sz2
166108 {-# INLINE compare #-}
167109
168−-- | @since 0.6.4.0
169−instance Prim a => IsList (PrimArray a) where
170− type Item (PrimArray a) = a
171− fromList = primArrayFromList
172− fromListN = primArrayFromListN
173− toList = primArrayToList
174−
175110-- | @since 0.6.4.0
176111instance (Show a, Prim a) => Show (PrimArray a) where
177112 showsPrec p a = showParen (p > 10) $
178113 showString "fromListN " . shows (sizeofPrimArray a) . showString " "
179114 . shows (primArrayToList a)
180115
181−die :: String -> String -> a
182−die fun problem = error $ "Data.Primitive.PrimArray." ++ fun ++ ": " ++ problem
183−
184−primArrayFromList :: Prim a => [a] -> PrimArray a
185−primArrayFromList vs = primArrayFromListN (L.length vs) vs
186−
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)
190− run = do
191− arr <- newPrimArray len
192− let go :: [a] -> Int -> ST s ()
193− go [] !ix = if ix == len
194− then return ()
195− else die "fromListN" "list length less than specified size"
196− go (a : as) !ix = if ix < len
197− then do
198− writePrimArray arr ix a
199− go as (ix + 1)
200− else die "fromListN" "list length greater than specified size"
201− go vs 0
202− unsafeFreezePrimArray arr
203−
204116-- | Convert the primitive array to a list.
205117{-# INLINE primArrayToList #-}
206118primArrayToList :: forall a. Prim a => PrimArray a -> [a]
207119primArrayToList xs = build (\c n -> foldrPrimArray c n xs)
208120
209−#if MIN_VERSION_primitive(0,6,4)
210−
211−primArrayToByteArray :: PrimArray a -> PB.ByteArray
212−primArrayToByteArray (PrimArray x) = PB.ByteArray x
213−
214−byteArrayToPrimArray :: ByteArray -> PrimArray a
215−byteArrayToPrimArray (PB.ByteArray x) = PrimArray x
216−
217−#if MIN_VERSION_base(4,9,0)
218−-- | @since 0.6.4.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))
223−#endif
224−
225−-- | @since 0.6.4.0
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))
230−#endif
231− mconcat = byteArrayToPrimArray . mconcat . map primArrayToByteArray
232−
233−#endif
234−
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'# #)
241−
242121-- | Create a new mutable primitive array of the given length. The
243122-- underlying memory is left uninitialized.
244123newPrimArray :: 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)
282161shrinkMutablePrimArray (MutablePrimArray arr#) (I# n#)
283162 = primitive_ (shrinkMutableByteArray# arr# (n# *# sizeOf# (undefined :: a)))
284163
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#)
289−
290164-- | Write an element to the given index.
291165writePrimArray ::
292166 (Prim a, PrimMonad m)
@@ -298,128 +172,6 @@ writePrimArray ::
298172writePrimArray (MutablePrimArray arr#) (I# i#) x
299173 = primitive_ (writeByteArray# arr# i# x)
300174
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
311− -> m ()
312−{-# INLINE copyMutablePrimArray #-}
313−copyMutablePrimArray (MutablePrimArray dst#) (I# doff#) (MutablePrimArray src#) (I# soff#) (I# n#)
314− = primitive_ (copyMutableByteArray#
315− src#
316− (soff# *# (sizeOf# (undefined :: a)))
317− dst#
318− (doff# *# (sizeOf# (undefined :: a)))
319− (n# *# (sizeOf# (undefined :: a)))
320− )
321−
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
330− -> m ()
331−{-# INLINE copyPrimArray #-}
332−copyPrimArray (MutablePrimArray dst#) (I# doff#) (PrimArray src#) (I# soff#) (I# n#)
333− = primitive_ (copyByteArray#
334− src#
335− (soff# *# (sizeOf# (undefined :: a)))
336− dst#
337− (doff# *# (sizeOf# (undefined :: a)))
338− (n# *# (sizeOf# (undefined :: a)))
339− )
340−
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
351− -> m ()
352−{-# INLINE copyPrimArrayToPtr #-}
353−copyPrimArrayToPtr (Ptr addr#) (PrimArray ba#) (I# soff#) (I# n#) =
354− primitive (\ s# ->
355− let s'# = copyByteArrayToAddr# ba# (soff# *# siz#) addr# (n# *# siz#) s#
356− in (# s'#, () #))
357− where siz# = sizeOf# (undefined :: a)
358−
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
369− -> m ()
370−{-# INLINE copyMutablePrimArrayToPtr #-}
371−copyMutablePrimArrayToPtr (Ptr addr#) (MutablePrimArray mba#) (I# soff#) (I# n#) =
372− primitive (\ s# ->
373− let s'# = copyMutableByteArrayToAddr# mba# (soff# *# siz#) addr# (n# *# siz#) s#
374− in (# s'#, () #))
375− where siz# = sizeOf# (undefined :: a)
376−
377−-- | Fill a slice of a mutable primitive array with a value.
378−setPrimArray
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
384− -> m ()
385−{-# INLINE setPrimArray #-}
386−setPrimArray (MutablePrimArray dst#) (I# doff#) (I# sz#) x
387− = primitive_ (PT.setByteArray# dst# doff# sz# x)
388−
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
393− -> m Int
394−{-# INLINE getSizeofMutablePrimArray #-}
395−#if __GLASGOW_HASKELL__ >= 801
396−getSizeofMutablePrimArray (MutablePrimArray arr#)
397− = primitive (\s# ->
398− case getSizeofMutableByteArray# arr# s# of
399− (# s'#, sz# #) -> (# s'#, I# (quotInt# sz# (sizeOf# (undefined :: a))) #)
400− )
401−#else
402−-- On older GHCs, it is not possible to resize a byte array, so
403−-- this provides behavior consistent with the implementation for
404−-- newer GHCs.
405−getSizeofMutablePrimArray arr
406− = return (sizeofMutablePrimArray arr)
407−#endif
408−
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)))
416−
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#)
422−
423175-- | Convert a mutable byte array to an immutable one without copying. The
424176-- array should not be modified after the conversion.
425177unsafeFreezePrimArray
@@ -429,14 +181,6 @@ unsafeFreezePrimArray (MutablePrimArray arr#)
429181 = primitive (\s# -> case unsafeFreezeByteArray# arr# s# of
430182 (# s'#, arr'# #) -> (# s'#, PrimArray arr'# #))
431183
432−-- | Convert an immutable array to a mutable one without copying. The
433−-- original array should not be used after the conversion.
434−unsafeThawPrimArray
435− :: PrimMonad m => PrimArray a -> m (MutablePrimArray (PrimState m) a)
436−{-# INLINE unsafeThawPrimArray #-}
437−unsafeThawPrimArray (PrimArray arr#)
438− = primitive (\s# -> (# s#, MutablePrimArray (unsafeCoerce# arr#) #))
439−
440184-- | Read a primitive value from the primitive array.
441185indexPrimArray :: forall a. Prim a => PrimArray a -> Int -> a
442186{-# INLINE indexPrimArray #-}
@@ -457,24 +201,6 @@ foldrPrimArray f z arr = go 0
457201 | sz > i = f (indexPrimArray arr i) (go (i+1))
458202 | otherwise = z
459203
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
464− where
465− go !i !acc
466− | i < 0 = acc
467− | otherwise = go (i - 1) (f (indexPrimArray arr i) acc)
468−
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)
473− where
474− go !i
475− | i < 0 = z
476− | otherwise = f (go (i - 1)) (indexPrimArray arr i)
477−
478204-- | Strict left-associated fold over the elements of a 'PrimArray'.
479205{-# INLINE foldlPrimArray' #-}
480206foldlPrimArray' :: forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
@@ -484,466 +210,3 @@ foldlPrimArray' f z0 arr = go 0 z0
484210 go !i !acc
485211 | i < sz = go (i + 1) (f acc (indexPrimArray arr i))
486212 | otherwise = acc
487−
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
492− where
493− !sz = sizeofPrimArray arr
494− go !i !acc1
495− | i < sz = do
496− acc2 <- f acc1 (indexPrimArray arr i)
497− go (i + 1) acc2
498− | otherwise = return acc1
499−
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.
502−-- Consequently:
503−--
504−-- >>> traversePrimArrayP (\x -> print x $> bool x undefined (x == 2)) (fromList [1, 2, 3 :: Int])
505−-- 1
506−-- 2
507−-- *** Exception: Prelude.undefined
508−--
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:
512−--
513−-- > incrPositiveA :: PrimArray Int -> Maybe (PrimArray Int)
514−-- > incrPositiveA xs = traversePrimArray (\x -> bool Nothing (Just (x + 1)) (x > 0)) xs
515−--
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'
518−-- instance:
519−--
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))
523−-- > xs
524−--
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)
529− => (a -> m b)
530− -> PrimArray a
531− -> m (PrimArray b)
532−traversePrimArrayP f arr = do
533− let !sz = sizeofPrimArray arr
534− marr <- newPrimArray sz
535− let go !ix = if ix < sz
536− then do
537− b <- f (indexPrimArray arr ix)
538− writePrimArray marr ix b
539− go (ix + 1)
540− else return ()
541− go 0
542− unsafeFreezePrimArray marr
543−
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)
548− => (a -> m Bool)
549− -> PrimArray a
550− -> m (PrimArray a)
551−filterPrimArrayP f arr = do
552− let !sz = sizeofPrimArray arr
553− marr <- newPrimArray sz
554− let go !ixSrc !ixDst = if ixSrc < sz
555− then do
556− let a = indexPrimArray arr ixSrc
557− b <- f a
558− if b
559− then do
560− writePrimArray marr ixDst a
561− go (ixSrc + 1) (ixDst + 1)
562− else go (ixSrc + 1) ixDst
563− else return ixDst
564− lenDst <- go 0 0
565− marr' <- resizeMutablePrimArray marr lenDst
566− unsafeFreezePrimArray marr'
567−
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))
573− -> PrimArray a
574− -> m (PrimArray b)
575−mapMaybePrimArrayP f arr = do
576− let !sz = sizeofPrimArray arr
577− marr <- newPrimArray sz
578− let go !ixSrc !ixDst = if ixSrc < sz
579− then do
580− let a = indexPrimArray arr ixSrc
581− mb <- f a
582− case mb of
583− Just b -> do
584− writePrimArray marr ixDst b
585− go (ixSrc + 1) (ixDst + 1)
586− Nothing -> go (ixSrc + 1) ixDst
587− else return ixDst
588− lenDst <- go 0 0
589− marr' <- resizeMutablePrimArray marr lenDst
590− unsafeFreezePrimArray marr'
591−
592−-- | Generate a primitive array by evaluating the monadic generator function
593−-- at each index.
594−{-# INLINE generatePrimArrayP #-}
595−generatePrimArrayP :: (PrimMonad m, Prim a)
596− => Int -- ^ length
597− -> (Int -> m a) -- ^ generator
598− -> m (PrimArray a)
599−generatePrimArrayP sz f = do
600− marr <- newPrimArray sz
601− let go !ix = if ix < sz
602− then do
603− b <- f ix
604− writePrimArray marr ix b
605− go (ix + 1)
606− else return ()
607− go 0
608− unsafeFreezePrimArray marr
609−
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)
614− => Int
615− -> m a
616− -> m (PrimArray a)
617−replicatePrimArrayP sz f = do
618− marr <- newPrimArray sz
619− let go !ix = if ix < sz
620− then do
621− b <- f
622− writePrimArray marr ix b
623− go (ix + 1)
624− else return ()
625− go 0
626− unsafeFreezePrimArray marr
627−
628−
629−-- | Map over the elements of a primitive array.
630−{-# INLINE mapPrimArray #-}
631−mapPrimArray :: (Prim a, Prim b)
632− => (a -> b)
633− -> PrimArray a
634− -> PrimArray b
635−mapPrimArray f arr = runST $ do
636− let !sz = sizeofPrimArray arr
637− marr <- newPrimArray sz
638− let go !ix = if ix < sz
639− then do
640− let b = f (indexPrimArray arr ix)
641− writePrimArray marr ix b
642− go (ix + 1)
643− else return ()
644− go 0
645− unsafeFreezePrimArray marr
646−
647−-- | Indexed map over the elements of a primitive array.
648−{-# INLINE imapPrimArray #-}
649−imapPrimArray :: (Prim a, Prim b)
650− => (Int -> a -> b)
651− -> PrimArray a
652− -> PrimArray b
653−imapPrimArray f arr = runST $ do
654− let !sz = sizeofPrimArray arr
655− marr <- newPrimArray sz
656− let go !ix = if ix < sz
657− then do
658− let b = f ix (indexPrimArray arr ix)
659− writePrimArray marr ix b
660− go (ix + 1)
661− else return ()
662− go 0
663− unsafeFreezePrimArray marr
664−
665−-- | Filter elements of a primitive array according to a predicate.
666−{-# INLINE filterPrimArray #-}
667−filterPrimArray :: Prim a
668− => (a -> Bool)
669− -> PrimArray a
670− -> PrimArray a
671−filterPrimArray p arr = runST $ do
672− let !sz = sizeofPrimArray arr
673− marr <- newPrimArray sz
674− let go !ixSrc !ixDst = if ixSrc < sz
675− then do
676− let !a = indexPrimArray arr ixSrc
677− if p a
678− then do
679− writePrimArray marr ixDst a
680− go (ixSrc + 1) (ixDst + 1)
681− else go (ixSrc + 1) ixDst
682− else return ixDst
683− dstLen <- go 0 0
684− marr' <- resizeMutablePrimArray marr dstLen
685− unsafeFreezePrimArray marr'
686−
687−-- | Filter the primitive array, keeping the elements for which the monadic
688−-- predicate evaluates true.
689−filterPrimArrayA ::
690− (Applicative f, Prim a)
691− => (a -> f Bool) -- ^ mapping function
692− -> PrimArray a -- ^ primitive array
693− -> f (PrimArray a)
694−filterPrimArrayA f = \ !ary ->
695− let
696− !len = sizeofPrimArray ary
697− go !ixSrc
698− | ixSrc == len = pure $ IxSTA $ \ixDst _ -> return ixDst
699− | otherwise = let x = indexPrimArray ary ixSrc in
700− liftA2
701− (\keep (IxSTA m) -> IxSTA $ \ixDst mary -> if keep
702− then writePrimArray (MutablePrimArray mary) ixDst x >> m (ixDst + 1) mary
703− else m ixDst mary
704− )
705− (f x)
706− (go (ixSrc + 1))
707− in if len == 0
708− then pure emptyPrimArray
709− else runIxSTA len <$> go 0
710−
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
717− -> f (PrimArray b)
718−mapMaybePrimArrayA f = \ !ary ->
719− let
720− !len = sizeofPrimArray ary
721− go !ixSrc
722− | ixSrc == len = pure $ IxSTA $ \ixDst _ -> return ixDst
723− | otherwise = let x = indexPrimArray ary ixSrc in
724− liftA2
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
728− )
729− (f x)
730− (go (ixSrc + 1))
731− in if len == 0
732− then pure emptyPrimArray
733− else runIxSTA len <$> go 0
734−
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)
739− => (a -> Maybe b)
740− -> PrimArray a
741− -> PrimArray b
742−mapMaybePrimArray p arr = runST $ do
743− let !sz = sizeofPrimArray arr
744− marr <- newPrimArray sz
745− let go !ixSrc !ixDst = if ixSrc < sz
746− then do
747− let !a = indexPrimArray arr ixSrc
748− case p a of
749− Just b -> do
750− writePrimArray marr ixDst b
751− go (ixSrc + 1) (ixDst + 1)
752− Nothing -> go (ixSrc + 1) ixDst
753− else return ixDst
754− dstLen <- go 0 0
755− marr' <- resizeMutablePrimArray marr dstLen
756− unsafeFreezePrimArray marr'
757−
758−
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:
762−--
763−-- >>> traversePrimArray (\x -> print x $> bool x undefined (x == 2)) (fromList [1, 2, 3 :: Int])
764−-- 1
765−-- 2
766−-- 3
767−-- *** Exception: Prelude.undefined
768−--
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.
772−traversePrimArray ::
773− (Applicative f, Prim a, Prim b)
774− => (a -> f b) -- ^ mapping function
775− -> PrimArray a -- ^ primitive array
776− -> f (PrimArray b)
777−traversePrimArray f = \ !ary ->
778− let
779− !len = sizeofPrimArray ary
780− go !i
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)
785− (f x) (go (i + 1))
786− in if len == 0
787− then pure emptyPrimArray
788− else runSTA len <$> go 0
789−
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)
794− -> PrimArray a
795− -> f (PrimArray b)
796−itraversePrimArray f = \ !ary ->
797− let
798− !len = sizeofPrimArray ary
799− go !i
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))
805− in if len == 0
806− then pure emptyPrimArray
807− else runSTA len <$> go 0
808−
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)
815− -> PrimArray a
816− -> m (PrimArray b)
817−itraversePrimArrayP f arr = do
818− let !sz = sizeofPrimArray arr
819− marr <- newPrimArray sz
820− let go !ix
821− | ix < sz = do
822− writePrimArray marr ix =<< f ix (indexPrimArray arr ix)
823− go (ix + 1)
824− | otherwise = return ()
825− go 0
826− unsafeFreezePrimArray marr
827−
828−-- | Generate a primitive array.
829−{-# INLINE generatePrimArray #-}
830−generatePrimArray :: Prim a
831− => Int -- ^ length
832− -> (Int -> a) -- ^ element from index
833− -> PrimArray a
834−generatePrimArray len f = runST $ do
835− marr <- newPrimArray len
836− let go !ix = if ix < len
837− then do
838− writePrimArray marr ix (f ix)
839− go (ix + 1)
840− else return ()
841− go 0
842− unsafeFreezePrimArray marr
843−
844−-- | Create a primitive array by copying the element the given
845−-- number of times.
846−{-# INLINE replicatePrimArray #-}
847−replicatePrimArray :: Prim a
848− => Int -- ^ length
849− -> a -- ^ element
850− -> PrimArray a
851−replicatePrimArray len a = runST $ do
852− marr <- newPrimArray len
853− setPrimArray marr 0 len a
854− unsafeFreezePrimArray marr
855−
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)
861− => Int -- ^ length
862− -> (Int -> f a) -- ^ element from index
863− -> f (PrimArray a)
864−generatePrimArrayA len f =
865− let
866− go !i
867− | i == len = pure $ STA $ \mary -> unsafeFreezePrimArray (MutablePrimArray mary)
868− | otherwise
869− = liftA2 (\b (STA m) -> STA $ \mary ->
870− writePrimArray (MutablePrimArray mary) i b >> m mary)
871− (f i) (go (i + 1))
872− in if len == 0
873− then pure emptyPrimArray
874− else runSTA len <$> go 0
875−
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)
881− => Int -- ^ length
882− -> f a -- ^ applicative element producer
883− -> f (PrimArray a)
884−replicatePrimArrayA len f =
885− let
886− go !i
887− | i == len = pure $ STA $ \mary -> unsafeFreezePrimArray (MutablePrimArray mary)
888− | otherwise
889− = liftA2 (\b (STA m) -> STA $ \mary ->
890− writePrimArray (MutablePrimArray mary) i b >> m mary)
891− f (go (i + 1))
892− in if len == 0
893− then pure emptyPrimArray
894− else runSTA len <$> go 0
895−
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)
901− => (a -> f b)
902− -> PrimArray a
903− -> f ()
904−traversePrimArray_ f a = go 0 where
905− !sz = sizeofPrimArray a
906− go !ix = if ix < sz
907− then f (indexPrimArray a ix) *> go (ix + 1)
908− else pure ()
909−
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)
916− -> PrimArray a
917− -> f ()
918−itraversePrimArray_ f a = go 0 where
919− !sz = sizeofPrimArray a
920− go !ix = if ix < sz
921− then f ix (indexPrimArray a ix) *> go (ix + 1)
922− else pure ()
923−
924−newtype IxSTA a = IxSTA {_runIxSTA :: forall s. Int -> MutableByteArray# s -> ST s Int}
925−
926−runIxSTA :: forall a. Prim a
927− => Int -- maximum possible size
928− -> IxSTA a
929− -> PrimArray a
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 #-}
936−
937−newtype STA a = STA {_runSTA :: forall s. MutableByteArray# s -> ST s (PrimArray a)}
938−
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 #-}
942−
943−unMutablePrimArray :: MutablePrimArray s a -> MutableByteArray# s
944−unMutablePrimArray (MutablePrimArray m) = m
945−
946−{- $effectfulMapCreate
947−The naming conventions adopted in this section are explained in the
948−documentation of the @Data.Primitive@ module.
949−-}