module Control.Monad.Sharing.Implementation.CPS (
Lazy, evalLazy, runLazy,
Store, emptyStore, freshLabel, lookupValue, storeValue
) where
import Control.Monad ( MonadPlus(..) )
import Control.Monad.State ( MonadState(..), gets, modify )
import Control.Monad.Trans ( MonadTrans(..), MonadIO(..) )
import Control.Monad.Sharing.Classes
import Unsafe.Coerce
import qualified Data.IntMap as M
newtype Lazy m a = Lazy {
fromLazy :: forall w . (a -> Store -> m w) -> Store -> m w
}
evalLazy :: (Monad m, Convertible (Lazy m) a b) => Lazy m a -> m b
evalLazy m = runLazy (m >>= convert)
runLazy :: Monad m => Lazy m a -> m a
runLazy m = fromLazy m (\a _ -> return a) emptyStore
data Store = Store { nextLabel :: Int, heap :: M.IntMap Untyped }
emptyStore :: Store
emptyStore = Store 1 M.empty
freshLabel :: MonadState Store m => m Int
freshLabel = do s <- get
put (s { nextLabel = nextLabel s + 1 })
return (nextLabel s)
lookupValue :: MonadState Store m => Int -> m (Maybe a)
lookupValue k = gets (fmap typed . M.lookup k . heap)
storeValue :: MonadState Store m => Int -> a -> m ()
storeValue k v = modify (\s -> s { heap = M.insert k (Untyped v) (heap s) })
instance Monad m => Monad (Lazy m)
where
return x = Lazy (\c -> c x)
a >>= k = Lazy (\c s -> fromLazy a (\x -> fromLazy (k x) c) s)
fail err = Lazy (\_ _ -> fail err)
instance MonadPlus m => MonadPlus (Lazy m)
where
mzero = Lazy (\_ _ -> mzero)
a `mplus` b = Lazy (\c s -> fromLazy a c s `mplus` fromLazy b c s)
instance Monad m => MonadState Store (Lazy m)
where
get = Lazy (\c s -> c s s)
put s = Lazy (\c _ -> c () s)
instance MonadTrans Lazy
where
lift a = Lazy (\c s -> a >>= \x -> c x s)
instance MonadIO m => MonadIO (Lazy m)
where
liftIO = lift . liftIO
instance Monad m => Sharing (Lazy m)
where
share a = memo (a >>= shareArgs share)
memo :: Lazy m a -> Lazy m (Lazy m a)
memo a = Lazy (\c (Store key heap) ->
c (Lazy (\c s@(Store _ heap) ->
case M.lookup key heap of
Just x -> c (typed x) s
Nothing -> fromLazy a
(\x (Store other heap) ->
c x (Store other (M.insert key (Untyped x) heap))) s))
(Store (succ key) heap))
data Untyped = forall a . Untyped a
typed :: Untyped -> a
typed (Untyped x) = unsafeCoerce x