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 :: 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))
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)
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 (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