module Control.Monad.Trans.StateBag.Primitive (
StateBaggerT,
runBagger,
addItem,
topItem,
stackItem,
StateBagT,
makeBag,
getItem,
putItem,
modifyItemM,
ElementCount(),
ElementIndex(),
) where
import Control.Applicative
import Control.Monad
import Control.Monad.Trans.Class
import Control.Monad.Trans.Reader
import Control.Monad.Trans.State
import Control.Monad.Trans.StateBag.Internal
import Control.Monad.IO.Class
import Control.Monad.Primitive
import Data.Proxy
import GHC.Prim (Any, unsafeCoerce#)
import qualified Data.Vector.Mutable as V
newtype BagImpl s (bag :: [*]) = BagImpl (V.MVector s Any)
newtype StateBaggerT full (bag :: [*]) m a
= StateBaggerT (ReaderT (BagImpl (PrimState m) full) m a)
deriving (Functor, Applicative, Monad, MonadIO)
instance MonadTrans (StateBaggerT full bag) where
lift = StateBaggerT . lift
instance (PrimMonad m) => PrimMonad (StateBaggerT full bag m) where
type PrimState (StateBaggerT full bag m) = PrimState m
primitive = lift . primitive
runBagger :: forall full m a. (PrimMonad m, ElementCount full) =>
StateBaggerT full '[] m a -> m a
runBagger (StateBaggerT r) = do
vec <- V.new $ elemCount (Proxy :: Proxy full)
runReaderT r $ BagImpl vec
addItem :: forall item full bag m a. (PrimMonad m, ElementIndex item full) =>
item -> StateBaggerT full (item ': bag) m a -> StateBaggerT full bag m a
addItem item (StateBaggerT chain) = StateBaggerT $ do
(BagImpl vec) <- ask
V.write vec (elemIndex (Proxy :: Proxy item) (Proxy :: Proxy full)) $
unsafeCoerce# item
lift $ runReaderT chain $ BagImpl vec
topItem :: forall item full bag m. (PrimMonad m, ElementIndex item full) =>
StateBaggerT full (item ': bag) m item
topItem = StateBaggerT $ do
(BagImpl vec) <- ask
fmap unsafeCoerce# $ V.read vec $
elemIndex (Proxy :: Proxy item) (Proxy :: Proxy full)
stackItem :: forall item full bag m a. (PrimMonad m, ElementIndex item full) =>
item -> StateBaggerT full (item ': bag) m a ->
StateBaggerT full bag m (a, item)
stackItem item chain =
addItem item $ liftM2 (,) chain topItem
newtype StateBagT bag m a = StateBagT (ReaderT (BagImpl (PrimState m) bag) m a)
deriving (Functor, Applicative, Monad, MonadIO)
instance MonadTrans (StateBagT bag) where
lift = StateBagT . lift
instance (PrimMonad m) => PrimMonad (StateBagT bag m) where
type PrimState (StateBagT bag m) = PrimState m
primitive = lift . primitive
makeBag ::
forall bag m a. (PrimMonad m, ElementCount bag) =>
StateBagT bag m a -> StateBaggerT bag bag m a
makeBag (StateBagT r) = StateBaggerT r
itemImpl :: forall m item bag.
(PrimMonad m, ElementIndex item bag) =>
StateBagT bag m (StateBagT bag m item, item -> StateBagT bag m ())
itemImpl = do
let i = elemIndex (Proxy :: Proxy item) (Proxy :: Proxy bag)
(BagImpl vec) <- StateBagT ask
let geti = fmap unsafeCoerce# $ V.read vec i
let puti item = V.write vec i $ unsafeCoerce# item
return (geti, puti)
getItem :: forall m item bag.
(PrimMonad m, ElementIndex item bag) =>
StateBagT bag m item
getItem = itemImpl >>= fst
putItem :: forall m item bag.
(PrimMonad m, ElementIndex item bag) =>
item -> StateBagT bag m ()
putItem item = itemImpl >>= flip snd item
modifyItemM :: forall m item bag.
(PrimMonad m, ElementIndex item bag) =>
(item -> StateBagT bag m item) -> StateBagT bag m ()
modifyItemM f = do
(get, put) <- itemImpl
item <- get
item' <- f item
put item'