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