{-# LANGUAGE UndecidableInstances #-}
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 :: forall a p.
(Ord a, Typeable a, Typeable p) =>
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 a p.
(Ord a, Typeable a, Typeable p) =>
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 :: forall (m :: * -> *) a p.
(MonadMemoize m, Ord a, Typeable a, Typeable p) =>
Name -> (a -> m p) -> a -> m p
memoize Name
name a -> m p
f a
a = Name -> a -> m p -> m p
forall a p.
(Ord a, Typeable a, Typeable p) =>
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
  { forall (m :: * -> *) a.
MemoizeT m a -> StateT (DMap MemoizationKey Identity) m a
unMemoizeT :: StateT (DMap MemoizationKey Identity) m a
  }
  deriving ((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
$cfmap :: forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> MemoizeT m a -> MemoizeT m b
fmap :: forall a b. (a -> b) -> MemoizeT m a -> MemoizeT m b
$c<$ :: forall (m :: * -> *) a b.
Functor m =>
a -> MemoizeT m b -> MemoizeT m a
<$ :: forall a b. a -> MemoizeT m b -> MemoizeT m a
Functor, Functor (MemoizeT m)
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)
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
$cpure :: forall (m :: * -> *) a. Monad m => a -> MemoizeT m a
pure :: forall a. a -> MemoizeT m a
$c<*> :: forall (m :: * -> *) a b.
Monad m =>
MemoizeT m (a -> b) -> MemoizeT m a -> MemoizeT m b
<*> :: forall a b. MemoizeT m (a -> b) -> MemoizeT m a -> MemoizeT m b
$cliftA2 :: forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> c) -> MemoizeT m a -> MemoizeT m b -> MemoizeT m c
liftA2 :: forall a b c.
(a -> b -> c) -> MemoizeT m a -> MemoizeT m b -> MemoizeT m c
$c*> :: forall (m :: * -> *) a b.
Monad m =>
MemoizeT m a -> MemoizeT m b -> MemoizeT m b
*> :: forall a b. MemoizeT m a -> MemoizeT m b -> MemoizeT m b
$c<* :: forall (m :: * -> *) a b.
Monad m =>
MemoizeT m a -> MemoizeT m b -> MemoizeT m a
<* :: forall a b. MemoizeT m a -> MemoizeT m b -> MemoizeT m a
Applicative, Applicative (MemoizeT m)
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)
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
$c>>= :: forall (m :: * -> *) a b.
Monad m =>
MemoizeT m a -> (a -> MemoizeT m b) -> MemoizeT m b
>>= :: forall a b. MemoizeT m a -> (a -> MemoizeT m b) -> MemoizeT m b
$c>> :: forall (m :: * -> *) a b.
Monad m =>
MemoizeT m a -> MemoizeT m b -> MemoizeT m b
>> :: forall a b. MemoizeT m a -> MemoizeT m b -> MemoizeT m b
$creturn :: forall (m :: * -> *) a. Monad m => a -> MemoizeT m a
return :: forall a. a -> MemoizeT m a
Monad, MonadError e, MonadReader r, (forall (m :: * -> *) a. Monad m => m a -> MemoizeT m a)
-> MonadTrans MemoizeT
forall (m :: * -> *) a. Monad m => m a -> MemoizeT m a
forall (t :: (* -> *) -> * -> *).
(forall (m :: * -> *) a. Monad m => m a -> t m a) -> MonadTrans t
$clift :: forall (m :: * -> *) a. Monad m => m a -> MemoizeT m a
lift :: forall (m :: * -> *) a. Monad m => m a -> MemoizeT m a
MonadTrans)
instance (MonadState s m) => MonadState s (MemoizeT m) where
  get :: MemoizeT m s
get = m s -> MemoizeT m s
forall (m :: * -> *) a. Monad m => m a -> MemoizeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m s
forall s (m :: * -> *). MonadState s m => m s
get
  put :: s -> MemoizeT m ()
put = m () -> MemoizeT m ()
forall (m :: * -> *) a. Monad m => m a -> MemoizeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> MemoizeT m ()) -> (s -> m ()) -> s -> MemoizeT m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put
runMemoizeT :: forall m a. (Monad m) => MemoizeT m a -> m a
runMemoizeT :: forall (m :: * -> *) a. Monad m => 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 :: forall a p.
(Ord a, Typeable a, Typeable p) =>
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 a. a -> StateT (DMap MemoizationKey Identity) m a
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 a. IO a -> StateT (DMap MemoizationKey Identity) m a
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 a. IO a -> StateT (DMap MemoizationKey Identity) m a
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 a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                Just p
parser -> Identity p -> IO (Identity p)
forall a. a -> IO a
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 a. IO a -> StateT (DMap MemoizationKey Identity) m a
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 a. a -> StateT (DMap MemoizationKey Identity) m a
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 :: forall a b. 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 <- forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a1 TypeRep a -> TypeRep a -> Maybe (a :~: a)
forall a b. 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` forall a. 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 <- forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @t1 TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall a b. 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` forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @t2 =
          (a :~: b) -> Maybe (a :~: b)
forall a. a -> Maybe a
Just a :~: a
a :~: b
forall {k} (a :: k). a :~: a
Refl
      | Bool
otherwise = Maybe (a :~: b)
forall a. Maybe a
Nothing
instance GCompare MemoizationKey where
  gcompare :: forall a b. 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 a b. 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 (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a1) (forall a. 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 a b. 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 (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @t1) (forall a. Typeable a => TypeRep a
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` GOrdering a a
GOrdering a b
(a ~ b) => GOrdering a b
forall {k} (a :: k). GOrdering a a
GEQ