module PrimitiveExtras.Folds
(
indexCounts,
unliftedArray,
primMultiArray,
)
where
import PrimitiveExtras.Prelude hiding (fold, foldM)
import PrimitiveExtras.Types
import Control.Foldl
import qualified PrimitiveExtras.UnliftedArray as UA
unsafeIO :: (state -> input -> IO state) -> IO state -> (state -> IO output) -> Fold input output
unsafeIO :: forall state input output.
(state -> input -> IO state)
-> IO state -> (state -> IO output) -> Fold input output
unsafeIO state -> input -> IO state
stepInIO IO state
initInIO state -> IO output
extractInIO =
forall a b x. (x -> a -> x) -> x -> (x -> b) -> Fold a b
Fold
(\ !state
state input
input -> forall a. IO a -> a
unsafeDupablePerformIO (state -> input -> IO state
stepInIO state
state input
input))
(forall a. IO a -> a
unsafeDupablePerformIO IO state
initInIO)
(\ state
state -> let !output :: output
output = forall a. IO a -> a
unsafePerformIO (state -> IO output
extractInIO state
state) in output
output)
foldMInUnsafeDupableIO :: FoldM IO input output -> Fold input output
foldMInUnsafeDupableIO :: forall input output. FoldM IO input output -> Fold input output
foldMInUnsafeDupableIO (FoldM x -> input -> IO x
step IO x
init x -> IO output
extract) = forall state input output.
(state -> input -> IO state)
-> IO state -> (state -> IO output) -> Fold input output
unsafeIO x -> input -> IO x
step IO x
init x -> IO output
extract
indexCounts :: (Integral count, Prim count) => Int -> Fold Int (PrimArray count)
indexCounts :: forall count.
(Integral count, Prim count) =>
Int -> Fold Int (PrimArray count)
indexCounts Int
size = forall state input output.
(state -> input -> IO state)
-> IO state -> (state -> IO output) -> Fold input output
unsafeIO forall {m :: * -> *} {a}.
(Prim a, PrimMonad m, Enum a) =>
MutablePrimArray (PrimState m) a
-> Int -> m (MutablePrimArray (PrimState m) a)
step IO (MutablePrimArray (PrimState IO) count)
init forall {m :: * -> *} {a}.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
extract where
init :: IO (MutablePrimArray (PrimState IO) count)
init = forall (m :: * -> *) a.
PrimMonad m =>
PrimArray a -> m (MutablePrimArray (PrimState m) a)
unsafeThawPrimArray (forall a. Prim a => Int -> a -> PrimArray a
replicatePrimArray Int
size count
0)
step :: MutablePrimArray (PrimState m) a
-> Int -> m (MutablePrimArray (PrimState m) a)
step MutablePrimArray (PrimState m) a
mutable Int
i = do
a
count <- forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray (PrimState m) a
mutable Int
i
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray (PrimState m) a
mutable Int
i (forall a. Enum a => a -> a
succ a
count)
forall (m :: * -> *) a. Monad m => a -> m a
return MutablePrimArray (PrimState m) a
mutable
extract :: MutablePrimArray (PrimState m) a -> m (PrimArray a)
extract = forall {m :: * -> *} {a}.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray
unliftedArray :: PrimUnlifted element => Int -> Fold (Int, element) (UnliftedArray element)
unliftedArray :: forall element.
PrimUnlifted element =>
Int -> Fold (Int, element) (UnliftedArray element)
unliftedArray Int
size =
forall state input output.
(state -> input -> IO state)
-> IO state -> (state -> IO output) -> Fold input output
unsafeIO forall {f :: * -> *} {a}.
(PrimMonad f, PrimUnlifted a) =>
MutableUnliftedArray (PrimState f) a
-> (Int, a) -> f (MutableUnliftedArray (PrimState f) a)
step IO (MutableUnliftedArray (PrimState IO) element)
init forall {m :: * -> *} {a}.
PrimMonad m =>
MutableUnliftedArray (PrimState m) a -> m (UnliftedArray a)
extract
where
step :: MutableUnliftedArray (PrimState f) a
-> (Int, a) -> f (MutableUnliftedArray (PrimState f) a)
step MutableUnliftedArray (PrimState f) a
mutable (Int
index, a
element) =
forall (m :: * -> *) a.
(PrimMonad m, PrimUnlifted a) =>
MutableUnliftedArray (PrimState m) a -> Int -> a -> m ()
writeUnliftedArray MutableUnliftedArray (PrimState f) a
mutable Int
index a
element forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> MutableUnliftedArray (PrimState f) a
mutable
init :: IO (MutableUnliftedArray (PrimState IO) element)
init =
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MutableUnliftedArray (PrimState m) a)
unsafeNewUnliftedArray Int
size
extract :: MutableUnliftedArray (PrimState m) a -> m (UnliftedArray a)
extract =
forall {m :: * -> *} {a}.
PrimMonad m =>
MutableUnliftedArray (PrimState m) a -> m (UnliftedArray a)
unsafeFreezeUnliftedArray
primMultiArray :: forall size element. (Integral size, Prim size, Prim element) => PrimArray size -> Fold (Int, element) (PrimMultiArray element)
primMultiArray :: forall size element.
(Integral size, Prim size, Prim element) =>
PrimArray size -> Fold (Int, element) (PrimMultiArray element)
primMultiArray PrimArray size
sizeArray =
forall state input output.
(state -> input -> IO state)
-> IO state -> (state -> IO output) -> Fold input output
unsafeIO forall {m :: * -> *} {a} {a}.
(PrimMonad m, Prim a, Prim a, Integral a) =>
Product2
(MutablePrimArray (PrimState m) a)
(UnliftedArray (MutablePrimArray (PrimState m) a))
-> (Int, a)
-> m (Product2
(MutablePrimArray (PrimState m) a)
(UnliftedArray (MutablePrimArray (PrimState m) a)))
step IO
(Product2
(MutablePrimArray RealWorld size)
(UnliftedArray (MutablePrimArray RealWorld element)))
init Product2
(MutablePrimArray RealWorld size)
(UnliftedArray (MutablePrimArray RealWorld element))
-> IO (PrimMultiArray element)
extract
where
outerLength :: Int
outerLength = forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray size
sizeArray
init :: IO
(Product2
(MutablePrimArray RealWorld size)
(UnliftedArray (MutablePrimArray RealWorld element)))
init =
forall a b. a -> b -> Product2 a b
Product2 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (MutablePrimArray RealWorld size)
initIndexArray forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IO (UnliftedArray (MutablePrimArray RealWorld element))
initMultiArray
where
initIndexArray :: IO (MutablePrimArray RealWorld size)
initIndexArray :: IO (MutablePrimArray RealWorld size)
initIndexArray =
forall (m :: * -> *) a.
PrimMonad m =>
PrimArray a -> m (MutablePrimArray (PrimState m) a)
unsafeThawPrimArray (forall a. Prim a => Int -> a -> PrimArray a
replicatePrimArray Int
outerLength size
0)
initMultiArray :: IO (UnliftedArray (MutablePrimArray RealWorld element))
initMultiArray :: IO (UnliftedArray (MutablePrimArray RealWorld element))
initMultiArray =
forall a.
PrimUnlifted a =>
Int -> (Int -> IO a) -> IO (UnliftedArray a)
UA.generate Int
outerLength forall a b. (a -> b) -> a -> b
$ \ Int
index -> do
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray (forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray size
sizeArray Int
index))
step :: Product2
(MutablePrimArray (PrimState m) a)
(UnliftedArray (MutablePrimArray (PrimState m) a))
-> (Int, a)
-> m (Product2
(MutablePrimArray (PrimState m) a)
(UnliftedArray (MutablePrimArray (PrimState m) a)))
step (Product2 MutablePrimArray (PrimState m) a
indexArray UnliftedArray (MutablePrimArray (PrimState m) a)
multiArray) (Int
outerIndex, a
element) = do
let innerArray :: MutablePrimArray (PrimState m) a
innerArray = forall a. PrimUnlifted a => UnliftedArray a -> Int -> a
indexUnliftedArray UnliftedArray (MutablePrimArray (PrimState m) a)
multiArray Int
outerIndex
a
innerIndex <- forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray (PrimState m) a
indexArray Int
outerIndex
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray (PrimState m) a
indexArray Int
outerIndex (forall a. Enum a => a -> a
succ a
innerIndex)
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray (PrimState m) a
innerArray (forall a b. (Integral a, Num b) => a -> b
fromIntegral a
innerIndex) a
element
forall (m :: * -> *) a. Monad m => a -> m a
return (forall a b. a -> b -> Product2 a b
Product2 MutablePrimArray (PrimState m) a
indexArray UnliftedArray (MutablePrimArray (PrimState m) a)
multiArray)
extract :: Product2
(MutablePrimArray RealWorld size)
(UnliftedArray (MutablePrimArray RealWorld element))
-> IO (PrimMultiArray element)
extract (Product2 MutablePrimArray RealWorld size
_ UnliftedArray (MutablePrimArray RealWorld element)
multiArray) = do
MutableUnliftedArray RealWorld (PrimArray element)
copied <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MutableUnliftedArray (PrimState m) a)
unsafeNewUnliftedArray Int
outerLength
forall (m :: * -> *) a.
Applicative m =>
Int -> (Int -> m a) -> m ()
forMFromZero_ Int
outerLength forall a b. (a -> b) -> a -> b
$ \ Int
outerIndex -> do
let mutableInnerArray :: MutablePrimArray RealWorld element
mutableInnerArray = forall a. PrimUnlifted a => UnliftedArray a -> Int -> a
indexUnliftedArray UnliftedArray (MutablePrimArray RealWorld element)
multiArray Int
outerIndex
PrimArray element
frozenInnerArray <- forall {m :: * -> *} {a}.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray RealWorld element
mutableInnerArray
forall (m :: * -> *) a.
(PrimMonad m, PrimUnlifted a) =>
MutableUnliftedArray (PrimState m) a -> Int -> a -> m ()
writeUnliftedArray MutableUnliftedArray RealWorld (PrimArray element)
copied Int
outerIndex PrimArray element
frozenInnerArray
UnliftedArray (PrimArray element)
result <- forall {m :: * -> *} {a}.
PrimMonad m =>
MutableUnliftedArray (PrimState m) a -> m (UnliftedArray a)
unsafeFreezeUnliftedArray MutableUnliftedArray RealWorld (PrimArray element)
copied
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. UnliftedArray (PrimArray a) -> PrimMultiArray a
PrimMultiArray forall a b. (a -> b) -> a -> b
$ UnliftedArray (PrimArray element)
result