module PrimitiveExtras.FoldMs where

import Control.Foldl
import PrimitiveExtras.Prelude hiding (fold, foldM)

-- |
-- Given a size of the array,
-- construct a fold, which produces an array of elements.
primArray ::
  (Prim a) =>
  -- | Array size
  Int ->
  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