module PrimitiveExtras.FoldMs
where

import PrimitiveExtras.Prelude hiding (fold, foldM)
import PrimitiveExtras.Types
import Control.Foldl
import qualified PrimitiveExtras.UnliftedArray as UA


{-|
Given a size of the array,
construct a fold, which produces an array of elements.
-}
primArray :: Prim a => Int {-^ Array size -} -> FoldM IO a (PrimArray a)
primArray :: forall a. Prim a => Int -> FoldM IO a (PrimArray a)
primArray Int
size = forall (m :: * -> *) a b x.
(x -> a -> m x) -> m x -> (x -> m b) -> FoldM m a b
FoldM forall {m :: * -> *} {a}.
(Prim a, PrimMonad m) =>
Product2 Int (MutablePrimArray (PrimState m) a)
-> a -> m (Product2 Int (MutablePrimArray (PrimState m) a))
step IO (Product2 Int (MutablePrimArray RealWorld a))
init forall {m :: * -> *} {a} {a}.
PrimMonad m =>
Product2 a (MutablePrimArray (PrimState m) a) -> m (PrimArray a)
extract where
  init :: IO (Product2 Int (MutablePrimArray RealWorld a))
init = forall a b. a -> b -> Product2 a b
Product2 Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
size
  step :: Product2 Int (MutablePrimArray (PrimState m) a)
-> a -> m (Product2 Int (MutablePrimArray (PrimState m) a))
step (Product2 Int
index MutablePrimArray (PrimState m) a
mutable) a
a = do
    forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray (PrimState m) a
mutable Int
index a
a
    forall (m :: * -> *) a. Monad m => a -> m a
return (forall a b. a -> b -> Product2 a b
Product2 (forall a. Enum a => a -> a
succ Int
index) MutablePrimArray (PrimState m) a
mutable)
  extract :: Product2 a (MutablePrimArray (PrimState m) a) -> m (PrimArray a)
extract (Product2 a
_ MutablePrimArray (PrimState m) a
mutable) = forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray (PrimState m) a
mutable