module Control.Monad.Sharing.Implementation.CPS (
collect,
Store, emptyStore, freshLabel, lookupValue, storeValue,
Untyped(..), typed
) where
import Control.Monad ( MonadPlus(..) )
import Control.Monad.State ( MonadState(..), gets, modify )
import Control.Monad.Sharing.Classes
import Unsafe.Coerce
import qualified Data.IntMap as M
newtype Lazy n a = Lazy {
fromLazy :: (a -> Store -> n) -> Store -> n
}
collect :: Nondet n => (forall s. Sharing s => s n) -> n
collect a = runLazy a
runLazy :: Nondet n => Lazy n n -> n
runLazy m = fromLazy m (\a _ -> 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 Nondet n => Monad (Lazy n)
where
return x = Lazy (\c -> c x)
a >>= k = Lazy (\c s -> fromLazy a (\x -> fromLazy (k x) c) s)
fail _ = Lazy (\_ _ -> failure)
instance Nondet n => MonadPlus (Lazy n)
where
mzero = Lazy (\_ _ -> failure)
a `mplus` b = Lazy (\c s -> fromLazy a c s ? fromLazy b c s)
instance Nondet n => MonadState Store (Lazy n)
where
get = Lazy (\c s -> c s s)
put s = Lazy (\c _ -> c () s)
instance Nondet n => Sharing (Lazy n)
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