module Control.Monad.Trans.Push where
import Control.Monad.ST.Trans (runST, STRef, readSTRef, writeSTRef, newSTRef)
import Control.Monad.ST.Trans.Internal (STT(STT), STTRet(STTRet))
import GHC.ST (ST(ST))
import qualified Data.Vector as V
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable as VGM
import Control.Monad.Push.Class (MonadPush, push)
data Res a = Res !Int
!a deriving Functor
newtype PushT v p m a = PushT (forall s . Int -> (STRef s (v s p)) -> (STT s m (Res a)))
deriving instance Functor m => Functor (PushT v p m)
instance Monad m => Applicative (PushT v p m) where
pure a = PushT $ \u _ -> (return (Res u a))
(PushT f) <*> (PushT g) = PushT $ \u v -> f u v >>= (\(Res u' o1) -> g u' v >>= (\(Res u'' o2) -> return (Res u'' (o1 o2))))
instance Monad m => Monad (PushT v p m) where
return = pure
PushT g >>= f = PushT $ \u v -> ((g u v) >>= (\(Res u' x) -> let PushT h = f x in h u' v))
PushT g >> PushT h = PushT $ \u v -> g u v >>= (\(Res u' _) -> h u' v)
liftST :: Applicative m => ST s a -> STT s m a
liftST (ST f) = STT (\s -> let (# s', a #) = f s in pure (STTRet s' a))
instance (Monad m, VGM.MVector v p) => MonadPush p (PushT v p m) where
push a = PushT $ \used vec' -> do
vec <- readSTRef vec'
if (VGM.length vec == used)
then do
bigger <- liftST $ VGM.grow vec used
liftST $ VGM.write bigger used a
writeSTRef vec' bigger
else do
liftST $ VGM.write vec used a
return $ Res (used+1) ()
runPushT :: (Monad m, VG.Vector v p) => PushT (VG.Mutable v) p m a -> m (a, v p)
runPushT (PushT action) = runST $ do
initial <- liftST $ VGM.new 1
vecRef <- newSTRef initial
(Res used out) <- action 0 vecRef
final <- readSTRef vecRef
vec <- liftST $ VG.freeze (VGM.slice 0 used final)
return (out, vec)
runPushTU :: forall p a m . (VU.Unbox p, Monad m) => PushT (VU.MVector) p m a -> m (a, VU.Vector p)
runPushTU = runPushT
runPushTB :: forall p a m . Monad m => PushT (V.MVector) p m a -> m (a, V.Vector p)
runPushTB = runPushT
runPushTS :: forall p a m . (VS.Storable p, Monad m) => PushT (VS.MVector) p m a -> m (a, VS.Vector p)
runPushTS = runPushT