{-# language BangPatterns #-}
{-# language FlexibleInstances #-}
{-# language MagicHash #-}
{-# language RankNTypes #-}
{-# language ScopedTypeVariables #-}
{-# language TypeFamilies #-}
{-# language TypeFamilyDependencies #-}
{-# language UnboxedTuples #-}
module Data.Primitive.Contiguous
(
size
, sizeMutable
, null
, index
, index#
, read
, indexM
, empty
, new
, singleton
, doubleton
, tripleton
, replicate
, replicateMutable
, generate
, generateMutable
, iterateN
, iterateMutableN
, write
, replicateMutableM
, generateMutableM
, iterateMutableNM
, create
, createT
, unfoldr
, unfoldrN
, unfoldrMutable
, enumFromN
, enumFromMutableN
, append
, reverse
, reverseMutable
, resize
, map
, map'
, mapMutable
, mapMutable'
, imap
, imap'
, imapMutable
, imapMutable'
, modify
, modify'
, mapMaybe
, filter
, ifilter
, equals
, equalsMutable
, same
, foldl
, foldl'
, foldr
, foldr'
, foldMap
, foldMap'
, foldlMap'
, ifoldl'
, ifoldr'
, ifoldlMap'
, ifoldlMap1'
, foldlM'
, traverse
, traverse_
, itraverse
, itraverse_
, traverseP
, fromList
, fromListN
, fromListMutable
, fromListMutableN
, unsafeFromListN
, unsafeFromListReverseN
, unsafeFromListReverseMutableN
, toList
, toListMutable
, convert
, lift
, unlift
, clone
, cloneMutable
, copy
, copyMutable
, freeze
, thaw
, unsafeFreeze
, liftHashWithSalt
, rnf
, 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
class Always a
instance Always a
class Contiguous (arr :: Type -> Type) where
type family Mutable arr = (r :: Type -> Type -> Type) | r -> arr
type family Element arr :: Type -> Constraint
empty :: arr a
null :: arr b -> Bool
new :: (PrimMonad m, Element arr b) => Int -> m (Mutable arr (PrimState m) b)
replicateMutable :: (PrimMonad m, Element arr b) => Int -> b -> m (Mutable arr (PrimState m) b)
index :: Element arr b => arr b -> Int -> b
index# :: Element arr b => arr b -> Int -> (# b #)
indexM :: (Element arr b, Monad m) => arr b -> Int -> m b
read :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -> Int -> m b
write :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -> Int -> b -> m ()
resize :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -> Int -> m (Mutable arr (PrimState m) b)
size :: Element arr b => arr b -> Int
sizeMutable :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -> m Int
unsafeFreeze :: PrimMonad m => Mutable arr (PrimState m) b -> m (arr b)
freeze :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -> Int -> Int -> m (arr b)
thaw :: (PrimMonad m, Element arr b) => arr b -> Int -> Int -> m (Mutable arr (PrimState m) b)
copy :: (PrimMonad m, Element arr b)
=> Mutable arr (PrimState m) b
-> Int
-> arr b
-> Int
-> Int
-> m ()
copyMutable :: (PrimMonad m, Element arr b)
=> Mutable arr (PrimState m) b
-> Int
-> Mutable arr (PrimState m) b
-> Int
-> Int
-> m ()
clone :: Element arr b
=> arr b
-> Int
-> Int
-> arr b
cloneMutable :: (PrimMonad m, Element arr b)
=> Mutable arr (PrimState m) b
-> Int
-> Int
-> m (Mutable arr (PrimState m) b)
equals :: (Element arr b, Eq b) => arr b -> arr b -> Bool
equalsMutable :: Mutable arr s a -> Mutable arr s a -> Bool
unlift :: arr b -> ArrayArray#
lift :: ArrayArray# -> arr b
singleton :: Element arr a => a -> arr a
doubleton :: Element arr a => a -> a -> arr a
tripleton :: Element arr a => a -> a -> a -> arr a
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 :: (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 #-}
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 #-}
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 :: (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' :: (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 :: (Contiguous arr1, Element arr1 b, Contiguous arr2, Element arr2 b) => arr1 b -> arr2 b
convert a = map id a
{-# inline convert #-}
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
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' #-}
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)
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' #-}
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' #-}
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' #-}
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 #-}
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' #-}
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' #-}
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' #-}
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' #-}
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' #-}
filter :: (Contiguous arr, Element arr a)
=> (a -> Bool)
-> arr a
-> arr a
filter p arr = ifilter (\_ a -> p a) arr
{-# inline filter #-}
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 #-}
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 :: (Contiguous arr, Element arr a) => Int -> a -> arr a
replicate n x = create (replicateMutable n x)
{-# inline replicate #-}
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
-> a
-> 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 #-}
unsafeFromListN :: (Contiguous arr, Element arr a)
=> Int
-> [a]
-> 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 #-}
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 #-}
unsafeFromListReverseN :: (Contiguous arr, Element arr a)
=> Int
-> [a]
-> arr a
unsafeFromListReverseN n l = create (unsafeFromListReverseMutableN n l)
{-# inline unsafeFromListReverseN #-}
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 #-}
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' #-}
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 #-}
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' #-}
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 #-}
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 #-}
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_ #-}
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 #-}
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_ #-}
generate :: (Contiguous arr, Element arr a)
=> Int
-> (Int -> a)
-> arr a
generate len f = runST (generateMutable len f >>= unsafeFreeze)
{-# inline generate #-}
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 #-}
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 #-}
iterateN :: (Contiguous arr, Element arr a)
=> Int
-> (a -> a)
-> a
-> arr a
iterateN len f z0 = runST (iterateMutableN len f z0 >>= unsafeFreeze)
{-# inline iterateN #-}
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 #-}
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
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 #-}
create :: (Contiguous arr, Element arr a)
=> (forall s. ST s (Mutable arr s a))
-> arr a
create x = runST (unsafeFreeze =<< x)
{-# inline create #-}
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 #-}
unfoldr :: (Contiguous arr, Element arr a)
=> (b -> Maybe (a,b))
-> b
-> arr a
unfoldr f z0 = create (unfoldrMutable f z0)
{-# inline unfoldr #-}
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 #-}
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 #-}
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 #-}
toList :: (Contiguous arr, Element arr a)
=> arr a
-> [a]
toList arr = build (\c n -> foldr c n arr)
{-# inline toList #-}
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 #-}
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 #-}
fromListMutable :: (Contiguous arr, Element arr a, PrimMonad m)
=> [a]
-> m (Mutable arr (PrimState m) a)
fromListMutable xs = fromListMutableN (length xs) xs
{-# inline fromListMutable #-}
fromListN :: (Contiguous arr, Element arr a)
=> Int
-> [a]
-> arr a
fromListN len vs = create (fromListMutableN len vs)
{-# inline fromListN #-}
fromList :: (Contiguous arr, Element arr a)
=> [a]
-> arr a
fromList vs = create (fromListMutable vs)
{-# inline fromList #-}
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 #-}
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' #-}
enumFromN :: (Contiguous arr, Element arr a, Enum a)
=> a
-> Int
-> arr a
enumFromN z0 sz = create (enumFromMutableN z0 sz)
{-# inline enumFromN #-}
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 #-}
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 :: (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 #-}
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 #-}
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 #-}