{-# language BangPatterns #-} {-# language FlexibleInstances #-} {-# language MagicHash #-} {-# language RankNTypes #-} {-# language ScopedTypeVariables #-} {-# language TypeFamilies #-} {-# language TypeFamilyDependencies #-} {-# language UnboxedTuples #-} -- | The contiguous typeclass parameterises over a contiguous array type. -- This allows us to have a common API to a number of contiguous -- array types and their mutable counterparts. module Data.Primitive.Contiguous ( -- * Accessors -- ** Length Information size , sizeMutable , null -- ** Indexing , index , index# , read -- ** Monadic indexing , indexM -- * Construction -- ** Initialisation , empty , new , singleton , doubleton , tripleton , replicate , replicateMutable , generate , generateMutable , iterateN , iterateMutableN , write -- ** Monadic initialisation , replicateMutableM , generateMutableM , iterateMutableNM , create , createT -- ** Unfolding , unfoldr , unfoldrN , unfoldrMutable -- ** Enumeration , enumFromN , enumFromMutableN -- ** Concatenation , append -- * Modifying arrays -- ** Permutations , reverse , reverseMutable -- ** Resizing , resize -- * Elementwise operations -- ** Mapping , map , map' , mapMutable , mapMutable' , imap , imap' , imapMutable , imapMutable' , modify , modify' , mapMaybe -- * Working with predicates -- ** Filtering , filter , ifilter -- ** Comparing for equality , equals , equalsMutable , same -- * Folds , foldl , foldl' , foldr , foldr' , foldMap , foldMap' , foldlMap' , ifoldl' , ifoldr' , ifoldlMap' , ifoldlMap1' , foldlM' -- * Traversals , traverse , traverse_ , itraverse , itraverse_ , traverseP -- * Conversions -- ** Lists , fromList , fromListN , fromListMutable , fromListMutableN , unsafeFromListN , unsafeFromListReverseN , unsafeFromListReverseMutableN , toList , toListMutable -- ** Other array types , convert , lift , unlift -- ** Between mutable and immutable variants , clone , cloneMutable , copy , copyMutable , freeze , thaw , unsafeFreeze -- * Hashing , liftHashWithSalt -- * Forcing an array and its contents , rnf -- * Classes , Contiguous(Mutable,Element) , Always ) where import Prelude hiding (map,foldr,foldMap,traverse,read,filter,replicate,null,reverse,foldl,foldr) import Control.Applicative (liftA2) import Control.DeepSeq (NFData) import Control.Monad.Primitive import Control.Monad.ST (runST,ST) import Data.Bits (xor) import Data.Kind (Type) import Data.Primitive hiding (fromList,fromListN) import Data.Semigroup (Semigroup,(<>)) import Data.Word (Word8) import GHC.Base (build) import GHC.Exts (MutableArrayArray#,ArrayArray#,Constraint,sizeofByteArray#,sizeofArray#,sizeofArrayArray#,unsafeCoerce#,sameMutableArrayArray#,isTrue#,dataToTag#,Int(..)) import qualified Control.DeepSeq as DS -- | A typeclass that is satisfied by all types. This is used -- used to provide a fake constraint for 'Array' and 'SmallArray'. class Always a instance Always a -- | The 'Contiguous' typeclass as an interface to a multitude of -- contiguous structures. class Contiguous (arr :: Type -> Type) where -- | The Mutable counterpart to the array. type family Mutable arr = (r :: Type -> Type -> Type) | r -> arr -- | The constraint needed to store elements in the array. type family Element arr :: Type -> Constraint -- | The empty array. empty :: arr a -- | Test whether the array is empty. null :: arr b -> Bool -- | Allocate a new mutable array of the given size. new :: (PrimMonad m, Element arr b) => Int -> m (Mutable arr (PrimState m) b) -- | @'replicateMutable' n x@ is a mutable array of length @n@ with @x@ the value of every element. replicateMutable :: (PrimMonad m, Element arr b) => Int -> b -> m (Mutable arr (PrimState m) b) -- | Index into an array at the given index. index :: Element arr b => arr b -> Int -> b -- | Index into an array at the given index, yielding an unboxed one-tuple of the element. index# :: Element arr b => arr b -> Int -> (# b #) -- | Indexing in a monad. -- -- The monad allows operations to be strict in the array -- when necessary. Suppose array copying is implemented like this: -- -- > copy mv v = ... write mv i (v ! i) ... -- -- For lazy arrays, @v ! i@ would not be not be evaluated, -- which means that @mv@ would unnecessarily retain a reference -- to @v@ in each element written. -- -- With 'indexM', copying can be implemented like this instead: -- -- > copy mv v = ... do -- > x <- indexM v i -- > write mv i x -- -- Here, no references to @v@ are retained because indexing -- (but /not/ the elements) is evaluated eagerly. indexM :: (Element arr b, Monad m) => arr b -> Int -> m b -- | Read a mutable array at the given index. read :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -> Int -> m b -- | Write to a mutable array at the given index. write :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -> Int -> b -> m () -- | Resize an array into one with the given size. resize :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -> Int -> m (Mutable arr (PrimState m) b) -- | The size of the array size :: Element arr b => arr b -> Int -- | The size of the mutable array sizeMutable :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -> m Int -- | Turn a mutable array into an immutable one without copying. -- The mutable array should not be used after this conversion. unsafeFreeze :: PrimMonad m => Mutable arr (PrimState m) b -> m (arr b) -- | Turn a mutable array into an immutable one with copying, using a slice of the mutable array. freeze :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -> Int -> Int -> m (arr b) -- | Copy a slice of an immutable array into a new mutable array. thaw :: (PrimMonad m, Element arr b) => arr b -> Int -> Int -> m (Mutable arr (PrimState m) b) -- | Copy a slice of an array into a mutable array. copy :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -- ^ destination array -> Int -- ^ offset into destination array -> arr b -- ^ source array -> Int -- ^ offset into source array -> Int -- ^ number of elements to copy -> m () -- | Copy a slice of a mutable array into another mutable array. -- In the case that the destination and source arrays are the -- same, the regions may overlap. copyMutable :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -- ^ destination array -> Int -- ^ offset into destination array -> Mutable arr (PrimState m) b -- ^ source array -> Int -- ^ offset into source array -> Int -- ^ number of elements to copy -> m () -- | Clone a slice of an array. clone :: Element arr b => arr b -> Int -> Int -> arr b -- | Clone a slice of a mutable array. cloneMutable :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -> Int -> Int -> m (Mutable arr (PrimState m) b) -- | Test the two arrays for equality. equals :: (Element arr b, Eq b) => arr b -> arr b -> Bool -- | Test the two mutable arrays for pointer equality. -- Does not check equality of elements. equalsMutable :: Mutable arr s a -> Mutable arr s a -> Bool -- | Unlift an array into an 'ArrayArray#'. unlift :: arr b -> ArrayArray# -- | Lift an 'ArrayArray#' into an array. lift :: ArrayArray# -> arr b -- | Create a singleton array. singleton :: Element arr a => a -> arr a -- | Create a doubleton array. doubleton :: Element arr a => a -> a -> arr a -- | Create a tripleton array. tripleton :: Element arr a => a -> a -> a -> arr a -- | Reduce the array and all of its elements to WHNF. rnf :: (NFData a, Element arr a) => arr a -> () instance Contiguous SmallArray where type Mutable SmallArray = SmallMutableArray type Element SmallArray = Always empty = mempty new n = newSmallArray n errorThunk index = indexSmallArray indexM = indexSmallArrayM index# = indexSmallArray## read = readSmallArray write = writeSmallArray null a = case sizeofSmallArray a of 0 -> True _ -> False freeze = freezeSmallArray size = sizeofSmallArray sizeMutable = pure . sizeofSmallMutableArray unsafeFreeze = unsafeFreezeSmallArray thaw = thawSmallArray equals = (==) equalsMutable = (==) singleton a = runST $ do marr <- newSmallArray 1 errorThunk writeSmallArray marr 0 a unsafeFreezeSmallArray marr doubleton a b = runST $ do m <- newSmallArray 2 errorThunk writeSmallArray m 0 a writeSmallArray m 1 b unsafeFreezeSmallArray m tripleton a b c = runST $ do m <- newSmallArray 3 errorThunk writeSmallArray m 0 a writeSmallArray m 1 b writeSmallArray m 2 c unsafeFreezeSmallArray m rnf !ary = let !sz = sizeofSmallArray ary go !ix = if ix < sz then let !(# x #) = indexSmallArray## ary ix in DS.rnf x `seq` go (ix + 1) else () in go 0 clone = cloneSmallArray cloneMutable = cloneSmallMutableArray lift = fromArrayArray# unlift = toArrayArray# copy = copySmallArray copyMutable = copySmallMutableArray replicateMutable = replicateSmallMutableArray resize = resizeSmallArray {-# inline empty #-} {-# inline null #-} {-# inline new #-} {-# inline replicateMutable #-} {-# inline index #-} {-# inline index# #-} {-# inline indexM #-} {-# inline read #-} {-# inline write #-} {-# inline resize #-} {-# inline size #-} {-# inline sizeMutable #-} {-# inline unsafeFreeze #-} {-# inline freeze #-} {-# inline thaw #-} {-# inline copy #-} {-# inline copyMutable #-} {-# inline clone #-} {-# inline cloneMutable #-} {-# inline equals #-} {-# inline equalsMutable #-} {-# inline unlift #-} {-# inline lift #-} {-# inline singleton #-} {-# inline doubleton #-} {-# inline tripleton #-} {-# inline rnf #-} instance Contiguous PrimArray where type Mutable PrimArray = MutablePrimArray type Element PrimArray = Prim empty = mempty new = newPrimArray replicateMutable = replicateMutablePrimArray index = indexPrimArray index# arr ix = (# indexPrimArray arr ix #) indexM arr ix = pure (indexPrimArray arr ix) read = readPrimArray write = writePrimArray resize = resizeMutablePrimArray size = sizeofPrimArray sizeMutable = getSizeofMutablePrimArray freeze = freezePrimArray unsafeFreeze = unsafeFreezePrimArray thaw = thawPrimArray copy = copyPrimArray copyMutable = copyMutablePrimArray clone = clonePrimArray cloneMutable = cloneMutablePrimArray equals = (==) unlift = toArrayArray# lift = fromArrayArray# null (PrimArray a) = case sizeofByteArray# a of 0# -> True _ -> False equalsMutable = sameMutablePrimArray rnf (PrimArray !_) = () singleton a = runST $ do marr <- newPrimArray 1 writePrimArray marr 0 a unsafeFreezePrimArray marr doubleton a b = runST $ do m <- newPrimArray 2 writePrimArray m 0 a writePrimArray m 1 b unsafeFreezePrimArray m tripleton a b c = runST $ do m <- newPrimArray 3 writePrimArray m 0 a writePrimArray m 1 b writePrimArray m 2 c unsafeFreezePrimArray m {-# inline empty #-} {-# inline null #-} {-# inline new #-} {-# inline replicateMutable #-} {-# inline index #-} {-# inline index# #-} {-# inline indexM #-} {-# inline read #-} {-# inline write #-} {-# inline resize #-} {-# inline size #-} {-# inline sizeMutable #-} {-# inline unsafeFreeze #-} {-# inline freeze #-} {-# inline thaw #-} {-# inline copy #-} {-# inline copyMutable #-} {-# inline clone #-} {-# inline cloneMutable #-} {-# inline equals #-} {-# inline equalsMutable #-} {-# inline unlift #-} {-# inline lift #-} {-# inline singleton #-} {-# inline doubleton #-} {-# inline tripleton #-} {-# inline rnf #-} instance Contiguous Array where type Mutable Array = MutableArray type Element Array = Always empty = mempty new n = newArray n errorThunk replicateMutable = newArray index = indexArray index# = indexArray## indexM = indexArrayM read = readArray write = writeArray resize = resizeArray size = sizeofArray sizeMutable = pure . sizeofMutableArray freeze = freezeArray unsafeFreeze = unsafeFreezeArray thaw = thawArray copy = copyArray copyMutable = copyMutableArray clone = cloneArray cloneMutable = cloneMutableArray equals = (==) unlift = toArrayArray# lift = fromArrayArray# null (Array a) = case sizeofArray# a of 0# -> True _ -> False equalsMutable = sameMutableArray rnf !ary = let !sz = sizeofArray ary go !i | i == sz = () | otherwise = let !(# x #) = indexArray## ary i in DS.rnf x `seq` go (i+1) in go 0 singleton a = runST (newArray 1 a >>= unsafeFreezeArray) doubleton a b = runST $ do m <- newArray 2 a writeArray m 1 b unsafeFreezeArray m tripleton a b c = runST $ do m <- newArray 3 a writeArray m 1 b writeArray m 2 c unsafeFreezeArray m {-# inline empty #-} {-# inline null #-} {-# inline new #-} {-# inline replicateMutable #-} {-# inline index #-} {-# inline index# #-} {-# inline indexM #-} {-# inline read #-} {-# inline write #-} {-# inline resize #-} {-# inline size #-} {-# inline sizeMutable #-} {-# inline unsafeFreeze #-} {-# inline freeze #-} {-# inline thaw #-} {-# inline copy #-} {-# inline copyMutable #-} {-# inline clone #-} {-# inline cloneMutable #-} {-# inline equals #-} {-# inline equalsMutable #-} {-# inline unlift #-} {-# inline lift #-} {-# inline singleton #-} {-# inline doubleton #-} {-# inline tripleton #-} {-# inline rnf #-} instance Contiguous UnliftedArray where type Mutable UnliftedArray = MutableUnliftedArray type Element UnliftedArray = PrimUnlifted empty = emptyUnliftedArray new = unsafeNewUnliftedArray replicateMutable = newUnliftedArray index = indexUnliftedArray index# arr ix = (# indexUnliftedArray arr ix #) indexM arr ix = pure (indexUnliftedArray arr ix) read = readUnliftedArray write = writeUnliftedArray resize = resizeUnliftedArray size = sizeofUnliftedArray sizeMutable = pure . sizeofMutableUnliftedArray freeze = freezeUnliftedArray unsafeFreeze = unsafeFreezeUnliftedArray thaw = thawUnliftedArray copy = copyUnliftedArray copyMutable = copyMutableUnliftedArray clone = cloneUnliftedArray cloneMutable = cloneMutableUnliftedArray equals = (==) unlift = toArrayArray# lift = fromArrayArray# null (UnliftedArray a) = case sizeofArrayArray# a of 0# -> True _ -> False equalsMutable = sameMutableUnliftedArray rnf !ary = let !sz = sizeofUnliftedArray ary go !i | i == sz = () | otherwise = let x = indexUnliftedArray ary i in DS.rnf x `seq` go (i+1) in go 0 singleton a = runST (newUnliftedArray 1 a >>= unsafeFreezeUnliftedArray) doubleton a b = runST $ do m <- newUnliftedArray 2 a writeUnliftedArray m 1 b unsafeFreezeUnliftedArray m tripleton a b c = runST $ do m <- newUnliftedArray 3 a writeUnliftedArray m 1 b writeUnliftedArray m 2 c unsafeFreezeUnliftedArray m {-# inline empty #-} {-# inline null #-} {-# inline new #-} {-# inline replicateMutable #-} {-# inline index #-} {-# inline index# #-} {-# inline indexM #-} {-# inline read #-} {-# inline write #-} {-# inline resize #-} {-# inline size #-} {-# inline sizeMutable #-} {-# inline unsafeFreeze #-} {-# inline freeze #-} {-# inline thaw #-} {-# inline copy #-} {-# inline copyMutable #-} {-# inline clone #-} {-# inline cloneMutable #-} {-# inline equals #-} {-# inline equalsMutable #-} {-# inline unlift #-} {-# inline lift #-} {-# inline singleton #-} {-# inline doubleton #-} {-# inline tripleton #-} {-# inline rnf #-} errorThunk :: a errorThunk = error "Contiguous typeclass: unitialized element" {-# noinline errorThunk #-} freezePrimArray :: (PrimMonad m, Prim a) => MutablePrimArray (PrimState m) a -> Int -> Int -> m (PrimArray a) freezePrimArray !src !off !len = do dst <- newPrimArray len copyMutablePrimArray dst 0 src off len unsafeFreezePrimArray dst {-# inline freezePrimArray #-} resizeArray :: PrimMonad m => MutableArray (PrimState m) a -> Int -> m (MutableArray (PrimState m) a) resizeArray !src !sz = do dst <- newArray sz errorThunk copyMutableArray dst 0 src 0 (min sz (sizeofMutableArray src)) pure dst {-# inline resizeArray #-} resizeSmallArray :: PrimMonad m => SmallMutableArray (PrimState m) a -> Int -> m (SmallMutableArray (PrimState m) a) resizeSmallArray !src !sz = do dst <- newSmallArray sz errorThunk copySmallMutableArray dst 0 src 0 (min sz (sizeofSmallMutableArray src)) pure dst {-# inline resizeSmallArray #-} resizeUnliftedArray :: (PrimMonad m, PrimUnlifted a) => MutableUnliftedArray (PrimState m) a -> Int -> m (MutableUnliftedArray (PrimState m) a) resizeUnliftedArray !src !sz = do dst <- unsafeNewUnliftedArray sz copyMutableUnliftedArray dst 0 src 0 (min sz (sizeofMutableUnliftedArray src)) pure dst {-# inline resizeUnliftedArray #-} emptyUnliftedArray :: UnliftedArray a emptyUnliftedArray = runST (unsafeNewUnliftedArray 0 >>= unsafeFreezeUnliftedArray) {-# noinline emptyUnliftedArray #-} -- | Append two arrays. append :: (Contiguous arr, Element arr a) => arr a -> arr a -> arr a append !a !b = runST $ do let !szA = size a let !szB = size b m <- new (szA + szB) copy m 0 a 0 szA copy m szA b 0 szB unsafeFreeze m {-# inline append #-} -- | Map over the elements of an array with the index. imap :: (Contiguous arr1, Element arr1 b, Contiguous arr2, Element arr2 c) => (Int -> b -> c) -> arr1 b -> arr2 c imap f a = runST $ do mb <- new (size a) let go !i | i == size a = pure () | otherwise = do x <- indexM a i write mb i (f i x) go (i+1) go 0 unsafeFreeze mb {-# inline imap #-} -- | Map strictly over the elements of an array with the index. -- -- Note that because a new array must be created, the resulting -- array type can be /different/ than the original. imap' :: (Contiguous arr1, Element arr1 b, Contiguous arr2, Element arr2 c) => (Int -> b -> c) -> arr1 b -> arr2 c imap' f a = runST $ do mb <- new (size a) let go !i | i == size a = pure () | otherwise = do x <- indexM a i let !b = f i x write mb i b go (i + 1) go 0 unsafeFreeze mb {-# INLINABLE imap' #-} -- | Map over the elements of an array. -- -- Note that because a new array must be created, the resulting -- array type can be /different/ than the original. map :: (Contiguous arr1, Element arr1 b, Contiguous arr2, Element arr2 c) => (b -> c) -> arr1 b -> arr2 c map f a = runST $ do mb <- new (size a) let go !i | i == size a = pure () | otherwise = do x <- indexM a i write mb i (f x) go (i+1) go 0 unsafeFreeze mb {-# inline map #-} -- | Map strictly over the elements of an array. -- -- Note that because a new array must be created, the resulting -- array type can be /different/ than the original. map' :: (Contiguous arr1, Element arr1 b, Contiguous arr2, Element arr2 c) => (b -> c) -> arr1 b -> arr2 c map' f a = runST $ do mb <- new (size a) let go !i | i == size a = pure () | otherwise = do x <- indexM a i let !b = f x write mb i b go (i+1) go 0 unsafeFreeze mb {-# inline map' #-} -- | Convert one type of array into another. convert :: (Contiguous arr1, Element arr1 b, Contiguous arr2, Element arr2 b) => arr1 b -> arr2 b convert a = map id a {-# inline convert #-} -- | Right fold over the element of an array. foldr :: (Contiguous arr, Element arr a) => (a -> b -> b) -> b -> arr a -> b {-# inline foldr #-} foldr f z arr = go 0 where !sz = size arr go !i | sz > i = case index# arr i of (# x #) -> f x (go (i+1)) | otherwise = z -- | Strict right fold over the elements of an array. foldr' :: (Contiguous arr, Element arr a) => (a -> b -> b) -> b -> arr a -> b foldr' f !z !ary = let go i !acc | i == -1 = acc | !(# x #) <- index# ary i = go (i-1) (f x acc) in go (size ary - 1) z {-# inline foldr' #-} -- | Left fold over the elements of an array. foldl :: (Contiguous arr, Element arr a) => (b -> a -> b) -> b -> arr a -> b foldl f z ary = go 0 z where !sz = size ary go !i acc | i == sz = acc | otherwise = let (# x #) = index# ary i in go (i+1) (f acc x) -- | Strict left fold over the elements of an array. foldl' :: (Contiguous arr, Element arr a) => (b -> a -> b) -> b -> arr a -> b foldl' f !z !ary = let !sz = size ary go !i !acc | i == sz = acc | !(# x #) <- index# ary i = go (i+1) (f acc x) in go 0 z {-# inline foldl' #-} -- | Strict left fold over the elements of an array, where the accumulating -- function cares about the index of the element. ifoldl' :: (Contiguous arr, Element arr a) => (b -> Int -> a -> b) -> b -> arr a -> b ifoldl' f !z !ary = let !sz = size ary go !i !acc | i == sz = acc | (# x #) <- index# ary i = go (i+1) (f acc i x) in go 0 z {-# inline ifoldl' #-} -- | Strict right fold over the elements of an array, where the accumulating -- function cares about the index of the element. ifoldr' :: (Contiguous arr, Element arr a) => (Int -> a -> b -> b) -> b -> arr a -> b ifoldr' f !z !arr = let !sz = size arr go !i !acc = if i == (-1) then acc else let !(# x #) = index# arr i in go (i-1) (f i x acc) in go (sz-1) z {-# inline ifoldr' #-} -- | Monoidal fold over the element of an array. foldMap :: (Contiguous arr, Element arr a, Monoid m) => (a -> m) -> arr a -> m foldMap f arr = go 0 where !sz = size arr go !i | sz > i = case index# arr i of (# x #) -> mappend (f x) (go (i+1)) | otherwise = mempty {-# inline foldMap #-} -- | Strict monoidal fold over the elements of an array. foldMap' :: (Contiguous arr, Element arr a, Monoid m) => (a -> m) -> arr a -> m foldMap' f !ary = let !sz = size ary go !i !acc | i == sz = acc | (# x #) <- index# ary i = go (i+1) (mappend acc (f x)) in go 0 mempty {-# inline foldMap' #-} -- | Strict left monoidal fold over the elements of an array. foldlMap' :: (Contiguous arr, Element arr a, Monoid m) => (a -> m) -> arr a -> m foldlMap' f !ary = let !sz = size ary go !i !acc | i == sz = acc | (# x #) <- index# ary i = go (i+1) (mappend acc (f x)) in go 0 mempty {-# inline foldlMap' #-} -- | Strict monoidal fold over the elements of an array. ifoldlMap' :: (Contiguous arr, Element arr a, Monoid m) => (Int -> a -> m) -> arr a -> m ifoldlMap' f !ary = let !sz = size ary go !i !acc | i == sz = acc | (# x #) <- index# ary i = go (i+1) (mappend acc (f i x)) in go 0 mempty {-# inline ifoldlMap' #-} -- | Strict monoidal fold over the elements of an array. ifoldlMap1' :: (Contiguous arr, Element arr a, Semigroup m) => (Int -> a -> m) -> arr a -> m ifoldlMap1' f !ary = let !sz = size ary go !i !acc | i == sz = acc | (# x #) <- index# ary i = go (i+1) (acc <> f i x) !(# e0 #) = index# ary 0 in go 1 (f 0 e0) {-# inline ifoldlMap1' #-} -- | Strict left monadic fold over the elements of an array. foldlM' :: (Contiguous arr, Element arr a, Monad m) => (b -> a -> m b) -> b -> arr a -> m b foldlM' f z0 arr = go 0 z0 where !sz = size arr go !i !acc1 | i < sz = do let (# x #) = index# arr i acc2 <- f acc1 x go (i + 1) acc2 | otherwise = pure acc1 {-# inline foldlM' #-} -- | Drop elements that do not satisfy the predicate. filter :: (Contiguous arr, Element arr a) => (a -> Bool) -> arr a -> arr a filter p arr = ifilter (\_ a -> p a) arr {-# inline filter #-} -- | Drop elements that do not satisfy the predicate which -- is applied to values and their indices. ifilter :: (Contiguous arr, Element arr a) => (Int -> a -> Bool) -> arr a -> arr a ifilter p arr = runST $ do marr :: MutablePrimArray s Word8 <- newPrimArray sz let go1 :: Int -> Int -> ST s Int go1 !ix !numTrue = if ix < sz then do atIx <- indexM arr ix let !keep = p ix atIx let !keepTag = I# (dataToTag# keep) writePrimArray marr ix (fromIntegral keepTag) go1 (ix + 1) (numTrue + keepTag) else pure numTrue numTrue <- go1 0 0 if numTrue == sz then pure arr else do marrTrues <- new numTrue let go2 !ixSrc !ixDst = if ixDst < numTrue then do atIxKeep <- readPrimArray marr ixSrc if isTrue atIxKeep then do atIxVal <- indexM arr ixSrc write marrTrues ixDst atIxVal go2 (ixSrc + 1) (ixDst + 1) else go2 (ixSrc + 1) ixDst else pure () go2 0 0 unsafeFreeze marrTrues where !sz = size arr {-# inline ifilter #-} -- | The 'mapMaybe' function is a version of 'map' which can throw out elements. -- In particular, the functional arguments returns something of type @'Maybe' b@. -- If this is 'Nothing', no element is added on to the result array. If it is -- @'Just' b@, then @b@ is included in the result array. mapMaybe :: forall arr1 arr2 a b. (Contiguous arr1, Element arr1 a, Contiguous arr2, Element arr2 b) => (a -> Maybe b) -> arr1 a -> arr2 b mapMaybe f arr = runST $ do let !sz = size arr let go :: Int -> Int -> [b] -> ST s ([b],Int) go !ix !numJusts justs = if ix < sz then do atIx <- indexM arr ix case f atIx of Nothing -> go (ix+1) numJusts justs Just x -> go (ix+1) (numJusts+1) (x:justs) else pure (justs,numJusts) !(bs,!numJusts) <- go 0 0 [] !marr <- unsafeFromListReverseMutableN numJusts bs unsafeFreeze marr {-# inline mapMaybe #-} {-# inline isTrue #-} isTrue :: Word8 -> Bool isTrue 0 = False isTrue _ = True thawPrimArray :: (PrimMonad m, Prim a) => PrimArray a -> Int -> Int -> m (MutablePrimArray (PrimState m) a) thawPrimArray !arr !off !len = do marr <- newPrimArray len copyPrimArray marr 0 arr off len pure marr {-# inline thawPrimArray #-} clonePrimArray :: Prim a => PrimArray a -> Int -> Int -> PrimArray a clonePrimArray !arr !off !len = runST $ do marr <- newPrimArray len copyPrimArray marr 0 arr off len unsafeFreezePrimArray marr {-# inline clonePrimArray #-} cloneMutablePrimArray :: (PrimMonad m, Prim a) => MutablePrimArray (PrimState m) a -> Int -> Int -> m (MutablePrimArray (PrimState m) a) cloneMutablePrimArray !arr !off !len = do marr <- newPrimArray len copyMutablePrimArray marr 0 arr off len pure marr {-# inline cloneMutablePrimArray #-} -- | @'replicate' n x@ is an array of length @n@ with @x@ the value of every element. replicate :: (Contiguous arr, Element arr a) => Int -> a -> arr a replicate n x = create (replicateMutable n x) {-# inline replicate #-} -- | @'replicateMutableM' n act@ performs the action n times, gathering the results. replicateMutableM :: (PrimMonad m, Contiguous arr, Element arr a) => Int -> m a -> m (Mutable arr (PrimState m) a) replicateMutableM len act = do marr <- new len let go !ix = if ix < len then do x <- act write marr ix x go (ix + 1) else pure () go 0 pure marr {-# inline replicateMutableM #-} replicateMutablePrimArray :: (PrimMonad m, Prim a) => Int -- ^ length -> a -- ^ element -> m (MutablePrimArray (PrimState m) a) replicateMutablePrimArray len a = do marr <- newPrimArray len setPrimArray marr 0 len a pure marr {-# inline replicateMutablePrimArray #-} replicateSmallMutableArray :: (PrimMonad m) => Int -> a -> m (SmallMutableArray (PrimState m) a) replicateSmallMutableArray len a = do marr <- newSmallArray len errorThunk let go !ix = if ix < len then writeSmallArray marr ix a >> go (ix + 1) else pure () go 0 pure marr {-# inline replicateSmallMutableArray #-} -- | Create an array from a list. If the given length does -- not match the actual length, this function has undefined -- behavior. unsafeFromListN :: (Contiguous arr, Element arr a) => Int -- ^ length of list -> [a] -- ^ list -> arr a unsafeFromListN n l = create (unsafeFromListMutableN n l) {-# inline unsafeFromListN #-} unsafeFromListMutableN :: (Contiguous arr, Element arr a, PrimMonad m) => Int -> [a] -> m (Mutable arr (PrimState m) a) unsafeFromListMutableN n l = do m <- new n let go !_ [] = pure m go !ix (x : xs) = do write m ix x go (ix+1) xs go 0 l {-# inline unsafeFromListMutableN #-} -- | Create a mutable array from a list, reversing the order of -- the elements. If the given length does not match the actual length, -- this function has undefined behavior. unsafeFromListReverseMutableN :: (Contiguous arr, Element arr a, PrimMonad m) => Int -> [a] -> m (Mutable arr (PrimState m) a) unsafeFromListReverseMutableN n l = do m <- new n let go !_ [] = pure m go !ix (x : xs) = do write m ix x go (ix-1) xs go (n - 1) l {-# inline unsafeFromListReverseMutableN #-} -- | Create an array from a list, reversing the order of the -- elements. If the given length does not match the actual length, -- this function has undefined behavior. unsafeFromListReverseN :: (Contiguous arr, Element arr a) => Int -> [a] -> arr a unsafeFromListReverseN n l = create (unsafeFromListReverseMutableN n l) {-# inline unsafeFromListReverseN #-} -- | Map over a mutable array, modifying the elements in place. mapMutable :: (Contiguous arr, Element arr a, PrimMonad m) => (a -> a) -> Mutable arr (PrimState m) a -> m () mapMutable f = \ !mary -> do !sz <- sizeMutable mary let go !ix = if ix < sz then do a <- read mary ix write mary ix (f a) go (ix + 1) else pure () go 0 {-# inline mapMutable #-} -- | Strictly map over a mutable array, modifying the elements in place. mapMutable' :: (PrimMonad m, Contiguous arr, Element arr a) => (a -> a) -> Mutable arr (PrimState m) a -> m () mapMutable' f = \ !mary -> do !sz <- sizeMutable mary let go !i | i == sz = pure () | otherwise = do a <- read mary i let !b = f a write mary i b go (i + 1) go 0 {-# inline mapMutable' #-} -- | Map over a mutable array with indices, modifying the elements in place. imapMutable :: (Contiguous arr, Element arr a, PrimMonad m) => (Int -> a -> a) -> Mutable arr (PrimState m) a -> m () imapMutable f = \ !mary -> do !sz <- sizeMutable mary let go !ix = if ix < sz then do a <- read mary ix write mary ix (f ix a) go (ix + 1) else pure () go 0 {-# inline imapMutable #-} -- | Strictly map over a mutable array with indices, modifying the elements in place. imapMutable' :: (PrimMonad m, Contiguous arr, Element arr a) => (Int -> a -> a) -> Mutable arr (PrimState m) a -> m () imapMutable' f = \ !mary -> do !sz <- sizeMutable mary let go !i | i == sz = pure () | otherwise = do a <- read mary i let !b = f i a write mary i b go (i + 1) go 0 {-# inline imapMutable' #-} -- | Map each element of the array to an action, evaluate these -- actions from left to right, and collect the results in a -- new array. traverseP :: (PrimMonad m, Contiguous arr1, Contiguous arr2, Element arr1 a, Element arr2 b) => (a -> m b) -> arr1 a -> m (arr2 b) traverseP f = \ !ary -> let !sz = size ary go !i !mary | i == sz = unsafeFreeze mary | otherwise = do a <- indexM ary i b <- f a write mary i b go (i + 1) mary in do mary <- new sz go 0 mary {-# inline traverseP #-} newtype STA v a = STA {_runSTA :: forall s. Mutable v s a -> ST s (v a)} runSTA :: (Contiguous v, Element v a) => Int -> STA v a -> v a runSTA !sz (STA m) = runST $ new sz >>= \ ar -> m ar {-# inline runSTA #-} -- | Map each element of the array to an action, evaluate these -- actions from left to right, and collect the results. -- For a version that ignores the results, see 'traverse_'. traverse :: (Contiguous arr, Element arr a, Element arr b, Applicative f) => (a -> f b) -> arr a -> f (arr b) traverse f = \ !ary -> let !len = size ary go !i | i == len = pure $ STA $ \mary -> unsafeFreeze mary | (# x #) <- index# ary i = liftA2 (\b (STA m) -> STA $ \mary -> write mary i b >> m mary) (f x) (go (i + 1)) in if len == 0 then pure empty else runSTA len <$> go 0 {-# inline traverse #-} -- | Map each element of the array to an action, evaluate these -- actions from left to right, and ignore the results. -- For a version that doesn't ignore the results, see 'traverse'. traverse_ :: (Contiguous arr, Element arr a, Applicative f) => (a -> f b) -> arr a -> f () traverse_ f a = go 0 where !sz = size a go !ix = if ix < sz then f (index a ix) *> go (ix + 1) else pure () {-# inline traverse_ #-} -- | Map each element of the array and its index to an action, -- evaluating these actions from left to right. itraverse :: (Contiguous arr, Element arr a, Element arr b, Applicative f) => (Int -> a -> f b) -> arr a -> f (arr b) itraverse f ary = let !len = size ary go !ix | ix == len = pure $ STA $ \mary -> unsafeFreeze mary | (# x #) <- index# ary ix = liftA2 (\b (STA m) -> STA $ \mary -> write mary ix b >> m mary) (f ix x) (go (ix + 1)) in if len == 0 then pure empty else runSTA len <$> go 0 {-# inline itraverse #-} -- | Map each element of the array and its index to an action, -- evaluate these actions from left to right, and ignore the results. -- For a version that doesn't ignore the results, see 'itraverse'. itraverse_ :: (Contiguous arr, Element arr a, Applicative f) => (Int -> a -> f b) -> arr a -> f () itraverse_ f a = go 0 where !sz = size a go !ix = if ix < sz then f ix (index a ix) *> go (ix + 1) else pure () {-# inline itraverse_ #-} -- | Construct an array of the given length by applying -- the function to each index. generate :: (Contiguous arr, Element arr a) => Int -> (Int -> a) -> arr a generate len f = runST (generateMutable len f >>= unsafeFreeze) {-# inline generate #-} -- | Construct a mutable array of the given length by applying -- the function to each index. generateMutable :: (Contiguous arr, Element arr a, PrimMonad m) => Int -> (Int -> a) -> m (Mutable arr (PrimState m) a) generateMutable len f = generateMutableM len (pure . f) {-# inline generateMutable #-} -- | Construct a mutable array of the given length by applying -- the monadic action to each index. generateMutableM :: (Contiguous arr, Element arr a, PrimMonad m) => Int -> (Int -> m a) -> m (Mutable arr (PrimState m) a) generateMutableM !len f = do marr <- new len let go !ix = if ix < len then do x <- f ix write marr ix x go (ix + 1) else pure () go 0 pure marr {-# inline generateMutableM #-} -- | Apply a function @n@ times to a value and construct an array -- where each consecutive element is the result of an additional -- application of this function. The zeroth element is the original value. -- -- @'iterateN' 5 ('+' 1) 0 = 'fromListN' 5 [0,1,2,3,4]@ iterateN :: (Contiguous arr, Element arr a) => Int -> (a -> a) -> a -> arr a iterateN len f z0 = runST (iterateMutableN len f z0 >>= unsafeFreeze) {-# inline iterateN #-} -- | Apply a function @n@ times to a value and construct a mutable array -- where each consecutive element is the result of an additional -- application of this function. The zeroth element is the original value. iterateMutableN :: (Contiguous arr, Element arr a, PrimMonad m) => Int -> (a -> a) -> a -> m (Mutable arr (PrimState m) a) iterateMutableN len f z0 = iterateMutableNM len (pure . f) z0 {-# inline iterateMutableN #-} -- | Apply a monadic function @n@ times to a value and construct a mutable array -- where each consecutive element is the result of an additional -- application of this function. The zeroth element is the original value. iterateMutableNM :: (Contiguous arr, Element arr a, PrimMonad m) => Int -> (a -> m a) -> a -> m (Mutable arr (PrimState m) a) iterateMutableNM !len f z0 = do marr <- new len -- we are strict in the accumulator because -- otherwise we could build up a ton of `f (f (f (f .. (f a))))` -- thunks for no reason. let go !ix !acc | ix <= 0 = write marr ix z0 >> go (ix + 1) z0 | ix == len = pure () | otherwise = do a <- f acc write marr ix a go (ix + 1) a go 0 z0 pure marr {-# inline iterateMutableNM #-} -- | Execute the monad action and freeze the resulting array. create :: (Contiguous arr, Element arr a) => (forall s. ST s (Mutable arr s a)) -> arr a create x = runST (unsafeFreeze =<< x) {-# inline create #-} -- | Execute the monadic action and freeze the resulting array. createT :: (Contiguous arr, Element arr a, Traversable f) => (forall s. ST s (f (Mutable arr s a))) -> f (arr a) createT p = runST (mapM unsafeFreeze =<< p) {-# inline createT #-} -- | Construct an array by repeatedly applying a generator -- function to a seed. The generator function yields 'Just' the -- next element and the new seed or 'Nothing' if there are no more -- elements. -- -- >>> unfoldr (\n -> if n == 0 then Nothing else Just (n,n-1) 10 -- <10,9,8,7,6,5,4,3,2,1> -- Unfortunately, because we don't know ahead of time when to stop, -- we need to construct a list and then turn it into an array. unfoldr :: (Contiguous arr, Element arr a) => (b -> Maybe (a,b)) -> b -> arr a unfoldr f z0 = create (unfoldrMutable f z0) {-# inline unfoldr #-} -- | Construct a mutable array by repeatedly applying a generator -- function to a seed. The generator function yields 'Just' the -- next element and the new seed or 'Nothing' if there are no more -- elements. -- -- >>> unfoldrMutable (\n -> if n == 0 then Nothing else Just (n,n-1) 10 -- <10,9,8,7,6,5,4,3,2,1> -- Unfortunately, because we don't know ahead of time when to stop, -- we need to construct a list and then turn it into an array. unfoldrMutable :: (Contiguous arr, Element arr a, PrimMonad m) => (b -> Maybe (a,b)) -> b -> m (Mutable arr (PrimState m) a) unfoldrMutable f z0 = do let go !sz s !xs = case f s of Nothing -> pure (sz,xs) Just (x,s') -> go (sz + 1) s' (x : xs) (sz,xs) <- go 0 z0 [] unsafeFromListReverseMutableN sz xs {-# inline unfoldrMutable #-} -- | Construct an array with at most n elements by repeatedly -- applying the generator function to a seed. The generator function -- yields 'Just' the next element and the new seed or 'Nothing' if -- there are no more elements. unfoldrN :: (Contiguous arr, Element arr a) => Int -> (b -> Maybe (a, b)) -> b -> arr a unfoldrN maxSz f z0 = create (unfoldrMutableN maxSz f z0) {-# inline unfoldrN #-} -- | Construct a mutable array with at most n elements by repeatedly -- applying the generator function to a seed. The generator function -- yields 'Just' the next element and the new seed or 'Nothing' if -- there are no more elements. unfoldrMutableN :: (Contiguous arr, Element arr a, PrimMonad m) => Int -> (b -> Maybe (a, b)) -> b -> m (Mutable arr (PrimState m) a) unfoldrMutableN !maxSz f z0 = do m <- new maxSz let go !ix s = if ix < maxSz then case f s of Nothing -> pure ix Just (x,s') -> do write m ix x go (ix + 1) s' else pure ix sz <- go 0 z0 case compare maxSz sz of EQ -> pure m GT -> resize m sz LT -> error "Data.Primitive.Contiguous.unfoldrMutableN: internal error" {-# inline unfoldrMutableN #-} -- | Convert an array to a list. toList :: (Contiguous arr, Element arr a) => arr a -> [a] toList arr = build (\c n -> foldr c n arr) {-# inline toList #-} -- | Convert a mutable array to a list. -- I don't think this can be expressed in terms of foldr/build, -- so we just loop through the array. toListMutable :: (Contiguous arr, Element arr a, PrimMonad m) => Mutable arr (PrimState m) a -> m [a] toListMutable marr = do sz <- sizeMutable marr let go !ix !acc = if ix >= 0 then do x <- read marr ix go (ix - 1) (x : acc) else pure acc go (sz - 1) [] {-# inline toListMutable #-} -- | Given an 'Int' that is representative of the length of -- the list, convert the list into a mutable array of the -- given length. -- -- /Note/: calls 'error' if the given length is incorrect. fromListMutableN :: (Contiguous arr, Element arr a, PrimMonad m) => Int -> [a] -> m (Mutable arr (PrimState m) a) fromListMutableN len vs = do marr <- new len let go [] !ix = if ix == len then pure () else error "Data.Primitive.Contiguous.fromListN: list length less than specified size." go (a:as) !ix = if ix < len then do write marr ix a go as (ix + 1) else error "Data.Primitive.Contiguous.fromListN: list length greater than specified size." go vs 0 pure marr {-# inline fromListMutableN #-} -- | Convert a list into a mutable array of the given length. fromListMutable :: (Contiguous arr, Element arr a, PrimMonad m) => [a] -> m (Mutable arr (PrimState m) a) fromListMutable xs = fromListMutableN (length xs) xs {-# inline fromListMutable #-} -- | Given an 'Int' that is representative of the length of -- the list, convert the list into a mutable array of the -- given length. -- -- /Note/: calls 'error' if the given length is incorrect. fromListN :: (Contiguous arr, Element arr a) => Int -> [a] -> arr a fromListN len vs = create (fromListMutableN len vs) {-# inline fromListN #-} -- | Convert a list into an array. fromList :: (Contiguous arr, Element arr a) => [a] -> arr a fromList vs = create (fromListMutable vs) {-# inline fromList #-} -- | Modify the elements of a mutable array in-place. modify :: (Contiguous arr, Element arr a, PrimMonad m) => (a -> a) -> Mutable arr (PrimState m) a -> m () modify f marr = do !sz <- sizeMutable marr let go !ix = if ix < sz then do x <- read marr ix write marr ix (f x) go (ix + 1) else pure () go 0 {-# inline modify #-} -- | Strictly modify the elements of a mutable array in-place. modify' :: (Contiguous arr, Element arr a, PrimMonad m) => (a -> a) -> Mutable arr (PrimState m) a -> m () modify' f marr = do !sz <- sizeMutable marr let go !ix = if ix < sz then do x <- read marr ix let !y = f x write marr ix y go (ix + 1) else pure () go 0 {-# inline modify' #-} -- | Yield an array of the given length containing the values -- @x, 'succ' x, 'succ' ('succ' x)@ etc. enumFromN :: (Contiguous arr, Element arr a, Enum a) => a -> Int -> arr a enumFromN z0 sz = create (enumFromMutableN z0 sz) {-# inline enumFromN #-} -- | Yield a mutable array of the given length containing the values -- @x, 'succ' x, 'succ' ('succ' x)@ etc. enumFromMutableN :: (Contiguous arr, Element arr a, PrimMonad m, Enum a) => a -> Int -> m (Mutable arr (PrimState m) a) enumFromMutableN z0 !sz = do m <- new sz let go !ix z = if ix < sz then do write m ix z go (ix + 1) (succ z) else pure m go 0 z0 {-# inline enumFromMutableN #-} -- | Lift an accumulating hash function over the elements of the array, -- returning the final accumulated hash. liftHashWithSalt :: (Contiguous arr, Element arr a) => (Int -> a -> Int) -> Int -> arr a -> Int liftHashWithSalt f s0 arr = go 0 s0 where sz = size arr go !ix !s = if ix < sz then let !(# x #) = index# arr ix in go (ix + 1) (f s x) else hashIntWithSalt s ix {-# inline liftHashWithSalt #-} -- | Reverse the elements of an array. reverse :: (Contiguous arr, Element arr a) => arr a -> arr a reverse arr = runST $ do marr <- new sz copy marr 0 arr 0 sz reverseMutable marr unsafeFreeze marr where !sz = size arr {-# inline reverse #-} -- | Reverse the elements of a mutable array, in-place. reverseMutable :: (Contiguous arr, Element arr a, PrimMonad m) => Mutable arr (PrimState m) a -> m () reverseMutable marr = do !sz <- sizeMutable marr let go !start !end = if start >= end then pure () else do tmp <- read marr start write marr start =<< read marr end write marr end tmp go (start+1) (end-1) go 0 (sz-1) {-# inline reverseMutable #-} -- | This function does not behave deterministically. Optimization level and -- inlining can affect its results. However, the one thing that can be counted -- on is that if it returns 'True', the two immutable arrays are definitely the -- same. This is useful as shortcut for equality tests. However, keep in mind -- that a result of 'False' tells us nothing about the arguments. same :: Contiguous arr => arr a -> arr a -> Bool same a b = isTrue# (sameMutableArrayArray# (unsafeCoerce# (unlift a) :: MutableArrayArray# s) (unsafeCoerce# (unlift b) :: MutableArrayArray# s)) hashIntWithSalt :: Int -> Int -> Int hashIntWithSalt salt x = salt `combine` x {-# inline hashIntWithSalt #-} combine :: Int -> Int -> Int combine h1 h2 = (h1 * 16777619) `xor` h2 {-# inline combine #-}