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
28module 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
96import GHC.Exts
97import GHC.Base ( Int(..) )
98import Data.Primitive.Internal.Compat (isTrue#)
99import Data.Primitive.Types
100import Data.Primitive.ByteArray (ByteArray(..))
101import Data.Monoid (Monoid(..),(<>))
102import Control.Applicative
103import Control.Monad.Primitive
104import Control.Monad.ST
105import qualified Data.List as L
106import qualified Data.Primitive.ByteArray as PB
107import qualified Data.Primitive.Types as PT
108
109#if MIN_VERSION_base(4,7,0)
110import GHC.Exts (IsList(..))
111#endif
112
113#if MIN_VERSION_base(4,9,0)
114import Data.Semigroup (Semigroup)
115import 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.
123data 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.
132data MutablePrimArray s a = MutablePrimArray (MutableByteArray# s)
133
134sameByteArray :: ByteArray# -> ByteArray# -> Bool
135sameByteArray 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
145instance (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
164instance (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
179instance 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
187instance (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
192die :: String -> String -> a
193die fun problem = error $ "Data.Primitive.PrimArray." ++ fun ++ ": " ++ problem
194
195primArrayFromList :: Prim a => [a] -> PrimArray a
196primArrayFromList vs = primArrayFromListN (L.length vs) vs
197
198primArrayFromListN :: forall a. Prim a => Int -> [a] -> PrimArray a
199primArrayFromListN 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 #-}
217primArrayToList :: forall a. Prim a => PrimArray a -> [a]
218primArrayToList xs = build (\c n -> foldrPrimArray c n xs)
219
220primArrayToByteArray :: PrimArray a -> PB.ByteArray
221primArrayToByteArray (PrimArray x) = PB.ByteArray x
222
223byteArrayToPrimArray :: ByteArray -> PrimArray a
224byteArrayToPrimArray (PB.ByteArray x) = PrimArray x
225
226#if MIN_VERSION_base(4,9,0)
227-- | @since 0.6.4.0
228instance 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
235instance 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.
243emptyPrimArray :: PrimArray a
244{-# NOINLINE emptyPrimArray #-}
245emptyPrimArray = 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.
251newPrimArray :: forall m a. (PrimMonad m, Prim a) => Int -> m (MutablePrimArray (PrimState m) a)
252{-# INLINE newPrimArray #-}
253newPrimArray (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.
269resizeMutablePrimArray :: 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
275resizeMutablePrimArray (MutablePrimArray arr#) (I# n#)
276  = primitive (\s# -> case resizeMutableByteArray# arr# (n# *# sizeOf# (undefined :: a)) s# of
277                        (# s'#, arr'# #) -> (# s'#, MutablePrimArray arr'# #))
278#else
279resizeMutablePrimArray 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.
291shrinkMutablePrimArray :: forall m a. (PrimMonad m, Prim a)
292  => MutablePrimArray (PrimState m) a
293  -> Int -- ^ new size
294  -> m ()
295{-# INLINE shrinkMutablePrimArray #-}
296shrinkMutablePrimArray (MutablePrimArray arr#) (I# n#)
297  = primitive_ (shrinkMutableByteArray# arr# (n# *# sizeOf# (undefined :: a)))
298#endif
299
300readPrimArray :: (Prim a, PrimMonad m) => MutablePrimArray (PrimState m) a -> Int -> m a
301{-# INLINE readPrimArray #-}
302readPrimArray (MutablePrimArray arr#) (I# i#)
303  = primitive (readByteArray# arr# i#)
304
305-- | Write an element to the given index.
306writePrimArray ::
307     (Prim a, PrimMonad m)
308  => MutablePrimArray (PrimState m) a -- ^ array
309  -> Int -- ^ index
310  -> a -- ^ element
311  -> m ()
312{-# INLINE writePrimArray #-}
313writePrimArray (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.
319copyMutablePrimArray :: 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 #-}
328copyMutablePrimArray (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.
338copyPrimArray :: 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 #-}
347copyPrimArray (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.
362copyPrimArrayToPtr :: 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 #-}
369copyPrimArrayToPtr (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.
380copyMutablePrimArrayToPtr :: 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 #-}
387copyMutablePrimArrayToPtr (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.
395setPrimArray
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 #-}
403setPrimArray (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.
408getSizeofMutablePrimArray :: forall m a. (PrimMonad m, Prim a)
409  => MutablePrimArray (PrimState m) a -- ^ array
410  -> m Int
411{-# INLINE getSizeofMutablePrimArray #-}
412#if __GLASGOW_HASKELL__ >= 801
413getSizeofMutablePrimArray (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.
422getSizeofMutablePrimArray 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'.
429sizeofMutablePrimArray :: forall s a. Prim a => MutablePrimArray s a -> Int
430{-# INLINE sizeofMutablePrimArray #-}
431sizeofMutablePrimArray (MutablePrimArray arr#) =
432  I# (quotInt# (sizeofMutableByteArray# arr#) (sizeOf# (undefined :: a)))
433
434-- | Check if the two arrays refer to the same memory block.
435sameMutablePrimArray :: MutablePrimArray s a -> MutablePrimArray s a -> Bool
436{-# INLINE sameMutablePrimArray #-}
437sameMutablePrimArray (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.
442unsafeFreezePrimArray
443  :: PrimMonad m => MutablePrimArray (PrimState m) a -> m (PrimArray a)
444{-# INLINE unsafeFreezePrimArray #-}
445unsafeFreezePrimArray (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.
451unsafeThawPrimArray
452  :: PrimMonad m => PrimArray a -> m (MutablePrimArray (PrimState m) a)
453{-# INLINE unsafeThawPrimArray #-}
454unsafeThawPrimArray (PrimArray arr#)
455  = primitive (\s# -> (# s#, MutablePrimArray (unsafeCoerce# arr#) #))
456
457-- | Read a primitive value from the primitive array.
458indexPrimArray :: forall a. Prim a => PrimArray a -> Int -> a
459{-# INLINE indexPrimArray #-}
460indexPrimArray (PrimArray arr#) (I# i#) = indexByteArray# arr# i#
461
462-- | Get the size, in elements, of the primitive array.
463sizeofPrimArray :: forall a. Prim a => PrimArray a -> Int
464{-# INLINE sizeofPrimArray #-}
465sizeofPrimArray (PrimArray arr#) = I# (quotInt# (sizeofByteArray# arr#) (sizeOf# (undefined :: a)))
466
467-- | Lazy right-associated fold over the elements of a 'PrimArray'.
468{-# INLINE foldrPrimArray #-}
469foldrPrimArray :: forall a b. Prim a => (a -> b -> b) -> b -> PrimArray a -> b
470foldrPrimArray 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' #-}
479foldrPrimArray' :: forall a b. Prim a => (a -> b -> b) -> b -> PrimArray a -> b
480foldrPrimArray' 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 #-}
488foldlPrimArray :: forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
489foldlPrimArray 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' #-}
497foldlPrimArray' :: forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
498foldlPrimArray' 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' #-}
507foldlPrimArrayM' :: (Prim a, Monad m) => (b -> a -> m b) -> b -> PrimArray a -> m b
508foldlPrimArrayM' 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 #-}
545traversePrimArrayP :: (PrimMonad m, Prim a, Prim b)
546  => (a -> m b)
547  -> PrimArray a
548  -> m (PrimArray b)
549traversePrimArrayP 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 #-}
564filterPrimArrayP :: (PrimMonad m, Prim a)
565  => (a -> m Bool)
566  -> PrimArray a
567  -> m (PrimArray a)
568filterPrimArrayP 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 #-}
588mapMaybePrimArrayP :: (PrimMonad m, Prim a, Prim b)
589  => (a -> m (Maybe b))
590  -> PrimArray a
591  -> m (PrimArray b)
592mapMaybePrimArrayP 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 #-}
612generatePrimArrayP :: (PrimMonad m, Prim a)
613  => Int -- ^ length
614  -> (Int -> m a) -- ^ generator
615  -> m (PrimArray a)
616generatePrimArrayP 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 #-}
630replicatePrimArrayP :: (PrimMonad m, Prim a)
631  => Int
632  -> m a
633  -> m (PrimArray a)
634replicatePrimArrayP 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 #-}
648mapPrimArray :: (Prim a, Prim b)
649  => (a -> b)
650  -> PrimArray a
651  -> PrimArray b
652mapPrimArray 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 #-}
666imapPrimArray :: (Prim a, Prim b)
667  => (Int -> a -> b)
668  -> PrimArray a
669  -> PrimArray b
670imapPrimArray 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 #-}
684filterPrimArray :: Prim a
685  => (a -> Bool)
686  -> PrimArray a
687  -> PrimArray a
688filterPrimArray 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.
706filterPrimArrayA ::
707     (Applicative f, Prim a)
708  => (a -> f Bool) -- ^ mapping function
709  -> PrimArray a -- ^ primitive array
710  -> f (PrimArray a)
711filterPrimArrayA 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'.
730mapMaybePrimArrayA ::
731     (Applicative f, Prim a, Prim b)
732  => (a -> f (Maybe b)) -- ^ mapping function
733  -> PrimArray a -- ^ primitive array
734  -> f (PrimArray b)
735mapMaybePrimArrayA 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 #-}
755mapMaybePrimArray :: (Prim a, Prim b)
756  => (a -> Maybe b)
757  -> PrimArray a
758  -> PrimArray b
759mapMaybePrimArray 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.
789traversePrimArray ::
790     (Applicative f, Prim a, Prim b)
791  => (a -> f b) -- ^ mapping function
792  -> PrimArray a -- ^ primitive array
793  -> f (PrimArray b)
794traversePrimArray 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.
808itraversePrimArray ::
809     (Applicative f, Prim a, Prim b)
810  => (Int -> a -> f b)
811  -> PrimArray a
812  -> f (PrimArray b)
813itraversePrimArray 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 #-}
830itraversePrimArrayP :: (Prim a, Prim b, PrimMonad m)
831  => (Int -> a -> m b)
832  -> PrimArray a
833  -> m (PrimArray b)
834itraversePrimArrayP 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 #-}
847generatePrimArray :: Prim a
848  => Int -- ^ length
849  -> (Int -> a) -- ^ element from index
850  -> PrimArray a
851generatePrimArray 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 #-}
864replicatePrimArray :: Prim a
865  => Int -- ^ length
866  -> a -- ^ element
867  -> PrimArray a
868replicatePrimArray 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 #-}
876generatePrimArrayA ::
877     (Applicative f, Prim a)
878  => Int -- ^ length
879  -> (Int -> f a) -- ^ element from index
880  -> f (PrimArray a)
881generatePrimArrayA 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 #-}
896replicatePrimArrayA ::
897     (Applicative f, Prim a)
898  => Int -- ^ length
899  -> f a -- ^ applicative element producer
900  -> f (PrimArray a)
901replicatePrimArrayA 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.
916traversePrimArray_ ::
917     (Applicative f, Prim a)
918  => (a -> f b)
919  -> PrimArray a
920  -> f ()
921traversePrimArray_ 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.
930itraversePrimArray_ ::
931     (Applicative f, Prim a)
932  => (Int -> a -> f b)
933  -> PrimArray a
934  -> f ()
935itraversePrimArray_ 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
941newtype IxSTA a = IxSTA {_runIxSTA :: forall s. Int -> MutableByteArray# s -> ST s Int}
942
943runIxSTA :: forall a. Prim a
944  => Int -- maximum possible size
945  -> IxSTA a
946  -> PrimArray a
947runIxSTA !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
954newtype STA a = STA {_runSTA :: forall s. MutableByteArray# s -> ST s (PrimArray a)}
955
956runSTA :: forall a. Prim a => Int -> STA a -> PrimArray a
957runSTA !sz = \ (STA m) -> runST $ newPrimArray sz >>= \ (ar :: MutablePrimArray s a) -> m (unMutablePrimArray ar)
958{-# INLINE runSTA #-}
959
960unMutablePrimArray :: MutablePrimArray s a -> MutableByteArray# s
961unMutablePrimArray (MutablePrimArray m) = m
962
963{- $effectfulMapCreate
964The naming conventions adopted in this section are explained in the
965documentation of the @Data.Primitive@ module.
966-}
967
968
969