module Control.Monad.Memoize
( MonadMemoize (..),
memoize,
MemoizeT,
runMemoizeT,
)
where
import Control.Monad.Except
import Control.Monad.Reader (MonadReader, ReaderT, mapReaderT)
import Control.Monad.State.Strict (MonadState (..), StateT, evalStateT)
import Data.Dependent.Map (DMap)
import Data.Dependent.Map qualified as DM
import Data.Functor.Identity
import Data.GADT.Compare.Extended
import Data.IORef
import Data.Kind qualified as K
import Language.Haskell.TH qualified as TH
import System.IO.Unsafe (unsafeInterleaveIO)
import Type.Reflection (Typeable, typeRep, (:~:) (..))
import Prelude
class Monad m => MonadMemoize m where
memoizeOn ::
forall a p.
(Ord a, Typeable a, Typeable p) =>
TH.Name ->
a ->
m p ->
m p
instance
MonadMemoize m =>
MonadMemoize (ReaderT a m)
where
memoizeOn :: Name -> a -> ReaderT a m p -> ReaderT a m p
memoizeOn Name
name a
key = (m p -> m p) -> ReaderT a m p -> ReaderT a m p
forall (m :: * -> *) a (n :: * -> *) b r.
(m a -> n b) -> ReaderT r m a -> ReaderT r n b
mapReaderT (Name -> a -> m p -> m p
forall (m :: * -> *) a p.
(MonadMemoize m, Ord a, Typeable a, Typeable p) =>
Name -> a -> m p -> m p
memoizeOn Name
name a
key)
memoize ::
(MonadMemoize m, Ord a, Typeable a, Typeable p) =>
TH.Name ->
(a -> m p) ->
(a -> m p)
memoize :: Name -> (a -> m p) -> a -> m p
memoize Name
name a -> m p
f a
a = Name -> a -> m p -> m p
forall (m :: * -> *) a p.
(MonadMemoize m, Ord a, Typeable a, Typeable p) =>
Name -> a -> m p -> m p
memoizeOn Name
name a
a (a -> m p
f a
a)
newtype MemoizeT m a = MemoizeT
{ MemoizeT m a -> StateT (DMap MemoizationKey Identity) m a
unMemoizeT :: StateT (DMap MemoizationKey Identity) m a
}
deriving (a -> MemoizeT m b -> MemoizeT m a
(a -> b) -> MemoizeT m a -> MemoizeT m b
(forall a b. (a -> b) -> MemoizeT m a -> MemoizeT m b)
-> (forall a b. a -> MemoizeT m b -> MemoizeT m a)
-> Functor (MemoizeT m)
forall a b. a -> MemoizeT m b -> MemoizeT m a
forall a b. (a -> b) -> MemoizeT m a -> MemoizeT m b
forall (m :: * -> *) a b.
Functor m =>
a -> MemoizeT m b -> MemoizeT m a
forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> MemoizeT m a -> MemoizeT m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> MemoizeT m b -> MemoizeT m a
$c<$ :: forall (m :: * -> *) a b.
Functor m =>
a -> MemoizeT m b -> MemoizeT m a
fmap :: (a -> b) -> MemoizeT m a -> MemoizeT m b
$cfmap :: forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> MemoizeT m a -> MemoizeT m b
Functor, Functor (MemoizeT m)
a -> MemoizeT m a
Functor (MemoizeT m)
-> (forall a. a -> MemoizeT m a)
-> (forall a b.
MemoizeT m (a -> b) -> MemoizeT m a -> MemoizeT m b)
-> (forall a b c.
(a -> b -> c) -> MemoizeT m a -> MemoizeT m b -> MemoizeT m c)
-> (forall a b. MemoizeT m a -> MemoizeT m b -> MemoizeT m b)
-> (forall a b. MemoizeT m a -> MemoizeT m b -> MemoizeT m a)
-> Applicative (MemoizeT m)
MemoizeT m a -> MemoizeT m b -> MemoizeT m b
MemoizeT m a -> MemoizeT m b -> MemoizeT m a
MemoizeT m (a -> b) -> MemoizeT m a -> MemoizeT m b
(a -> b -> c) -> MemoizeT m a -> MemoizeT m b -> MemoizeT m c
forall a. a -> MemoizeT m a
forall a b. MemoizeT m a -> MemoizeT m b -> MemoizeT m a
forall a b. MemoizeT m a -> MemoizeT m b -> MemoizeT m b
forall a b. MemoizeT m (a -> b) -> MemoizeT m a -> MemoizeT m b
forall a b c.
(a -> b -> c) -> MemoizeT m a -> MemoizeT m b -> MemoizeT m c
forall (m :: * -> *). Monad m => Functor (MemoizeT m)
forall (m :: * -> *) a. Monad m => a -> MemoizeT m a
forall (m :: * -> *) a b.
Monad m =>
MemoizeT m a -> MemoizeT m b -> MemoizeT m a
forall (m :: * -> *) a b.
Monad m =>
MemoizeT m a -> MemoizeT m b -> MemoizeT m b
forall (m :: * -> *) a b.
Monad m =>
MemoizeT m (a -> b) -> MemoizeT m a -> MemoizeT m b
forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> c) -> MemoizeT m a -> MemoizeT m b -> MemoizeT m c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: MemoizeT m a -> MemoizeT m b -> MemoizeT m a
$c<* :: forall (m :: * -> *) a b.
Monad m =>
MemoizeT m a -> MemoizeT m b -> MemoizeT m a
*> :: MemoizeT m a -> MemoizeT m b -> MemoizeT m b
$c*> :: forall (m :: * -> *) a b.
Monad m =>
MemoizeT m a -> MemoizeT m b -> MemoizeT m b
liftA2 :: (a -> b -> c) -> MemoizeT m a -> MemoizeT m b -> MemoizeT m c
$cliftA2 :: forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> c) -> MemoizeT m a -> MemoizeT m b -> MemoizeT m c
<*> :: MemoizeT m (a -> b) -> MemoizeT m a -> MemoizeT m b
$c<*> :: forall (m :: * -> *) a b.
Monad m =>
MemoizeT m (a -> b) -> MemoizeT m a -> MemoizeT m b
pure :: a -> MemoizeT m a
$cpure :: forall (m :: * -> *) a. Monad m => a -> MemoizeT m a
$cp1Applicative :: forall (m :: * -> *). Monad m => Functor (MemoizeT m)
Applicative, Applicative (MemoizeT m)
a -> MemoizeT m a
Applicative (MemoizeT m)
-> (forall a b.
MemoizeT m a -> (a -> MemoizeT m b) -> MemoizeT m b)
-> (forall a b. MemoizeT m a -> MemoizeT m b -> MemoizeT m b)
-> (forall a. a -> MemoizeT m a)
-> Monad (MemoizeT m)
MemoizeT m a -> (a -> MemoizeT m b) -> MemoizeT m b
MemoizeT m a -> MemoizeT m b -> MemoizeT m b
forall a. a -> MemoizeT m a
forall a b. MemoizeT m a -> MemoizeT m b -> MemoizeT m b
forall a b. MemoizeT m a -> (a -> MemoizeT m b) -> MemoizeT m b
forall (m :: * -> *). Monad m => Applicative (MemoizeT m)
forall (m :: * -> *) a. Monad m => a -> MemoizeT m a
forall (m :: * -> *) a b.
Monad m =>
MemoizeT m a -> MemoizeT m b -> MemoizeT m b
forall (m :: * -> *) a b.
Monad m =>
MemoizeT m a -> (a -> MemoizeT m b) -> MemoizeT m b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: a -> MemoizeT m a
$creturn :: forall (m :: * -> *) a. Monad m => a -> MemoizeT m a
>> :: MemoizeT m a -> MemoizeT m b -> MemoizeT m b
$c>> :: forall (m :: * -> *) a b.
Monad m =>
MemoizeT m a -> MemoizeT m b -> MemoizeT m b
>>= :: MemoizeT m a -> (a -> MemoizeT m b) -> MemoizeT m b
$c>>= :: forall (m :: * -> *) a b.
Monad m =>
MemoizeT m a -> (a -> MemoizeT m b) -> MemoizeT m b
$cp1Monad :: forall (m :: * -> *). Monad m => Applicative (MemoizeT m)
Monad, MonadError e, MonadReader r)
runMemoizeT :: forall m a. Monad m => MemoizeT m a -> m a
runMemoizeT :: MemoizeT m a -> m a
runMemoizeT = (StateT (DMap MemoizationKey Identity) m a
-> DMap MemoizationKey Identity -> m a)
-> DMap MemoizationKey Identity
-> StateT (DMap MemoizationKey Identity) m a
-> m a
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT (DMap MemoizationKey Identity) m a
-> DMap MemoizationKey Identity -> m a
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT DMap MemoizationKey Identity
forall a. Monoid a => a
mempty (StateT (DMap MemoizationKey Identity) m a -> m a)
-> (MemoizeT m a -> StateT (DMap MemoizationKey Identity) m a)
-> MemoizeT m a
-> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MemoizeT m a -> StateT (DMap MemoizationKey Identity) m a
forall (m :: * -> *) a.
MemoizeT m a -> StateT (DMap MemoizationKey Identity) m a
unMemoizeT
instance
MonadIO m =>
MonadMemoize (MemoizeT m)
where
memoizeOn :: Name -> a -> MemoizeT m p -> MemoizeT m p
memoizeOn Name
name a
key MemoizeT m p
buildParser = StateT (DMap MemoizationKey Identity) m p -> MemoizeT m p
forall (m :: * -> *) a.
StateT (DMap MemoizationKey Identity) m a -> MemoizeT m a
MemoizeT do
let parserId :: MemoizationKey p
parserId = Name -> a -> MemoizationKey p
forall a p.
(Ord a, Typeable a, Typeable p) =>
Name -> a -> MemoizationKey p
MemoizationKey Name
name a
key
DMap MemoizationKey Identity
parsersById <- StateT
(DMap MemoizationKey Identity) m (DMap MemoizationKey Identity)
forall s (m :: * -> *). MonadState s m => m s
get
case MemoizationKey p
-> DMap MemoizationKey Identity -> Maybe (Identity p)
forall k1 (k2 :: k1 -> *) (f :: k1 -> *) (v :: k1).
GCompare k2 =>
k2 v -> DMap k2 f -> Maybe (f v)
DM.lookup MemoizationKey p
parserId DMap MemoizationKey Identity
parsersById of
Just (Identity p
parser) -> p -> StateT (DMap MemoizationKey Identity) m p
forall (f :: * -> *) a. Applicative f => a -> f a
pure p
parser
Maybe (Identity p)
Nothing -> do
IORef (Maybe p)
cell <- IO (IORef (Maybe p))
-> StateT (DMap MemoizationKey Identity) m (IORef (Maybe p))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IORef (Maybe p))
-> StateT (DMap MemoizationKey Identity) m (IORef (Maybe p)))
-> IO (IORef (Maybe p))
-> StateT (DMap MemoizationKey Identity) m (IORef (Maybe p))
forall a b. (a -> b) -> a -> b
$ Maybe p -> IO (IORef (Maybe p))
forall a. a -> IO (IORef a)
newIORef Maybe p
forall a. Maybe a
Nothing
Identity p
parserById <-
IO (Identity p)
-> StateT (DMap MemoizationKey Identity) m (Identity p)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Identity p)
-> StateT (DMap MemoizationKey Identity) m (Identity p))
-> IO (Identity p)
-> StateT (DMap MemoizationKey Identity) m (Identity p)
forall a b. (a -> b) -> a -> b
$
IO (Identity p) -> IO (Identity p)
forall a. IO a -> IO a
unsafeInterleaveIO (IO (Identity p) -> IO (Identity p))
-> IO (Identity p) -> IO (Identity p)
forall a b. (a -> b) -> a -> b
$
IORef (Maybe p) -> IO (Maybe p)
forall a. IORef a -> IO a
readIORef IORef (Maybe p)
cell IO (Maybe p) -> (Maybe p -> IO (Identity p)) -> IO (Identity p)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
Just p
parser -> Identity p -> IO (Identity p)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Identity p -> IO (Identity p)) -> Identity p -> IO (Identity p)
forall a b. (a -> b) -> a -> b
$ p -> Identity p
forall a. a -> Identity a
Identity p
parser
Maybe p
Nothing ->
[Char] -> IO (Identity p)
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO (Identity p)) -> [Char] -> IO (Identity p)
forall a b. (a -> b) -> a -> b
$
[[Char]] -> [Char]
unlines
[ [Char]
"memoize: parser was forced before being fully constructed",
[Char]
" parser constructor: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Name -> [Char]
forall a. Ppr a => a -> [Char]
TH.pprint Name
name
]
DMap MemoizationKey Identity
-> StateT (DMap MemoizationKey Identity) m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (DMap MemoizationKey Identity
-> StateT (DMap MemoizationKey Identity) m ())
-> DMap MemoizationKey Identity
-> StateT (DMap MemoizationKey Identity) m ()
forall a b. (a -> b) -> a -> b
$! MemoizationKey p
-> Identity p
-> DMap MemoizationKey Identity
-> DMap MemoizationKey Identity
forall k1 (k2 :: k1 -> *) (f :: k1 -> *) (v :: k1).
GCompare k2 =>
k2 v -> f v -> DMap k2 f -> DMap k2 f
DM.insert MemoizationKey p
parserId Identity p
parserById DMap MemoizationKey Identity
parsersById
p
parser <- MemoizeT m p -> StateT (DMap MemoizationKey Identity) m p
forall (m :: * -> *) a.
MemoizeT m a -> StateT (DMap MemoizationKey Identity) m a
unMemoizeT MemoizeT m p
buildParser
IO () -> StateT (DMap MemoizationKey Identity) m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> StateT (DMap MemoizationKey Identity) m ())
-> IO () -> StateT (DMap MemoizationKey Identity) m ()
forall a b. (a -> b) -> a -> b
$ IORef (Maybe p) -> Maybe p -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe p)
cell (p -> Maybe p
forall a. a -> Maybe a
Just p
parser)
p -> StateT (DMap MemoizationKey Identity) m p
forall (f :: * -> *) a. Applicative f => a -> f a
pure p
parser
data MemoizationKey (t :: K.Type) where
MemoizationKey :: (Ord a, Typeable a, Typeable p) => TH.Name -> a -> MemoizationKey p
instance GEq MemoizationKey where
geq :: MemoizationKey a -> MemoizationKey b -> Maybe (a :~: b)
geq
(MemoizationKey Name
name1 (a
arg1 :: a1) :: MemoizationKey t1)
(MemoizationKey Name
name2 (a
arg2 :: a2) :: MemoizationKey t2)
| Name
name1 Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
name2,
Just a :~: a
Refl <- Typeable a => TypeRep a
forall k (a :: k). Typeable a => TypeRep a
typeRep @a1 TypeRep a -> TypeRep a -> Maybe (a :~: a)
forall k (f :: k -> *) (a :: k) (b :: k).
GEq f =>
f a -> f b -> Maybe (a :~: b)
`geq` Typeable a => TypeRep a
forall k (a :: k). Typeable a => TypeRep a
typeRep @a2,
a
arg1 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
a
arg2,
Just a :~: b
Refl <- Typeable a => TypeRep a
forall k (a :: k). Typeable a => TypeRep a
typeRep @t1 TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall k (f :: k -> *) (a :: k) (b :: k).
GEq f =>
f a -> f b -> Maybe (a :~: b)
`geq` Typeable b => TypeRep b
forall k (a :: k). Typeable a => TypeRep a
typeRep @t2 =
(a :~: a) -> Maybe (a :~: a)
forall a. a -> Maybe a
Just a :~: a
forall k (a :: k). a :~: a
Refl
| Bool
otherwise = Maybe (a :~: b)
forall a. Maybe a
Nothing
instance GCompare MemoizationKey where
gcompare :: MemoizationKey a -> MemoizationKey b -> GOrdering a b
gcompare
(MemoizationKey Name
name1 (a
arg1 :: a1) :: MemoizationKey t1)
(MemoizationKey Name
name2 (a
arg2 :: a2) :: MemoizationKey t2) =
Ordering -> GOrdering Any Any
forall k (a :: k). Ordering -> GOrdering a a
strengthenOrdering (Name -> Name -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Name
name1 Name
name2)
GOrdering Any Any
-> ((Any ~ Any) => GOrdering a b) -> GOrdering a b
forall k1 k2 (a :: k1) (b :: k1) (c :: k2) (d :: k2).
GOrdering a b -> ((a ~ b) => GOrdering c d) -> GOrdering c d
`extendGOrdering` TypeRep a -> TypeRep a -> GOrdering a a
forall k (f :: k -> *) (a :: k) (b :: k).
GCompare f =>
f a -> f b -> GOrdering a b
gcompare (Typeable a => TypeRep a
forall k (a :: k). Typeable a => TypeRep a
typeRep @a1) (Typeable a => TypeRep a
forall k (a :: k). Typeable a => TypeRep a
typeRep @a2)
GOrdering a a -> ((a ~ a) => GOrdering a b) -> GOrdering a b
forall k1 k2 (a :: k1) (b :: k1) (c :: k2) (d :: k2).
GOrdering a b -> ((a ~ b) => GOrdering c d) -> GOrdering c d
`extendGOrdering` Ordering -> GOrdering Any Any
forall k (a :: k). Ordering -> GOrdering a a
strengthenOrdering (a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a
arg1 a
a
arg2)
GOrdering Any Any
-> ((Any ~ Any) => GOrdering a b) -> GOrdering a b
forall k1 k2 (a :: k1) (b :: k1) (c :: k2) (d :: k2).
GOrdering a b -> ((a ~ b) => GOrdering c d) -> GOrdering c d
`extendGOrdering` TypeRep a -> TypeRep b -> GOrdering a b
forall k (f :: k -> *) (a :: k) (b :: k).
GCompare f =>
f a -> f b -> GOrdering a b
gcompare (Typeable a => TypeRep a
forall k (a :: k). Typeable a => TypeRep a
typeRep @t1) (Typeable b => TypeRep b
forall k (a :: k). Typeable a => TypeRep a
typeRep @t2)
GOrdering a b -> ((a ~ b) => GOrdering a b) -> GOrdering a b
forall k1 k2 (a :: k1) (b :: k1) (c :: k2) (d :: k2).
GOrdering a b -> ((a ~ b) => GOrdering c d) -> GOrdering c d
`extendGOrdering` (a ~ b) => GOrdering a b
forall k (a :: k). GOrdering a a
GEQ