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 = (Product2 Int (MutablePrimArray RealWorld a)
 -> a -> IO (Product2 Int (MutablePrimArray RealWorld a)))
-> IO (Product2 Int (MutablePrimArray RealWorld a))
-> (Product2 Int (MutablePrimArray RealWorld a)
    -> IO (PrimArray a))
-> FoldM IO a (PrimArray a)
forall (m :: * -> *) a b x.
(x -> a -> m x) -> m x -> (x -> m b) -> FoldM m a b
FoldM Product2 Int (MutablePrimArray RealWorld a)
-> a -> IO (Product2 Int (MutablePrimArray RealWorld a))
Product2 Int (MutablePrimArray (PrimState IO) a)
-> a -> IO (Product2 Int (MutablePrimArray (PrimState IO) a))
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 Product2 Int (MutablePrimArray RealWorld a) -> IO (PrimArray a)
Product2 Int (MutablePrimArray (PrimState IO) a)
-> IO (PrimArray a)
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 = Int
-> MutablePrimArray RealWorld a
-> Product2 Int (MutablePrimArray RealWorld a)
forall a b. a -> b -> Product2 a b
Product2 Int
0 (MutablePrimArray RealWorld a
 -> Product2 Int (MutablePrimArray RealWorld a))
-> IO (MutablePrimArray RealWorld a)
-> IO (Product2 Int (MutablePrimArray RealWorld a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO (MutablePrimArray (PrimState IO) a)
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
      MutablePrimArray (PrimState m) a -> Int -> a -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray (PrimState m) a
mutable Int
index a
a
      Product2 Int (MutablePrimArray (PrimState m) a)
-> m (Product2 Int (MutablePrimArray (PrimState m) a))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
-> MutablePrimArray (PrimState m) a
-> Product2 Int (MutablePrimArray (PrimState m) a)
forall a b. a -> b -> Product2 a b
Product2 (Int -> Int
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) = MutablePrimArray (PrimState m) a -> m (PrimArray a)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray (PrimState m) a
mutable