@@ -1,70 +1,125 @@
1{-# LANGUAGE AllowAmbiguousTypes #-}
1{-# LANGUAGE AllowAmbiguousTypes #-}
2-- | Utilities for effectfully memoizing other, more effectful functions.
2-- | Utilities for effectfully memoizing other, more effectful functions.
3module OpenAPI.Checker.Memo
3module OpenAPI.Checker.Memo
5− , MonadMemo
4+ ( MonadMemo
6 , MemoState
5 , MemoState
8− , tryMemo
7+ , modifyMemoNonce
9− , tryMemoWithTag
8+ , KnotTier(..)
12import Control.Monad.State
14import Control.Monad.State
13import qualified Data.Map as M
16import qualified Data.Map as M
14import Data.Tagged
17import Data.Tagged
15import qualified Data.TypeRepMap as T
19import qualified Data.TypeRepMap as T
16import Type.Reflection
20import Type.Reflection
18−-- | A computation has either 'Finished', or it is 'InProgress'.
19−data Progress a = Finished a | InProgress
22+data Progress a = Finished a | Started | TyingKnot Dynamic
20− deriving (Eq, Ord, Show)
22data MemoMap a where
24data MemoMap a where
23 MemoMap :: !(M.Map k (Progress v)) -> MemoMap (k, v)
25 MemoMap :: !(M.Map k (Progress v)) -> MemoMap (k, v)
25−newtype MemoState = MemoState (T.TypeRepMap MemoMap)
27+data MemoState s = MemoState s (T.TypeRepMap MemoMap)
27-- | An effectful memoization monad.
29-- | An effectful memoization monad.
28−type MonadMemo m = MonadState MemoState m
30+type MonadMemo s m = MonadState (MemoState s) m
30memoStateLookup
32memoStateLookup
31− :: forall k v. (Typeable k, Typeable v, Ord k)
33+ :: forall k v s. (Typeable k, Typeable v, Ord k)
32− => k -> MemoState -> Maybe (Progress v)
34+ => k -> MemoState s -> Maybe (Progress v)
33−memoStateLookup k (MemoState t) = case T.lookup @(k, v) t of
35+memoStateLookup k (MemoState _ t) = case T.lookup @(k, v) t of
34 Just (MemoMap m) -> M.lookup k m
36 Just (MemoMap m) -> M.lookup k m
35 Nothing -> Nothing
37 Nothing -> Nothing
37memoStateInsert
39memoStateInsert
38− :: forall k v. (Typeable k, Typeable v, Ord k)
40+ :: forall k v s. (Typeable k, Typeable v, Ord k)
39− => k -> Progress v -> MemoState -> MemoState
41+ => k -> Progress v -> MemoState s -> MemoState s
40−memoStateInsert k x (MemoState t) = MemoState $ T.insert (MemoMap m'') t
42+memoStateInsert k x (MemoState s t) = MemoState s $ T.insert (MemoMap m'') t
42 m'' = M.insert k x m'
44 m'' = M.insert k x m'
43 m' = case T.lookup @(k, v) t of
45 m' = case T.lookup @(k, v) t of
44 Just (MemoMap m) -> m
46 Just (MemoMap m) -> m
45 Nothing -> M.empty
47 Nothing -> M.empty
49+modifyMemoNonce :: MonadMemo s m => (s -> s) -> m s
50+modifyMemoNonce f = do
52+ put $ MemoState (f s) t
47-- | Run a memoized computation.
55-- | Run a memoized computation.
48−runMemo :: Monad m => StateT MemoState m a -> m a
56+runMemo :: Monad m => s -> StateT (MemoState s) m a -> m a
49−runMemo = (`evalStateT` MemoState T.empty)
57+runMemo s = (`evalStateT` MemoState s T.empty)
59+-- | A description of how to effectfully tie knots in type @v@, using the @m@
60+-- monad, and by sharing some @d@ data among the recursive instances.
61+data KnotTier v d m = KnotTier
62+ { onKnotFound :: m d -- ^ Create some data that will be connected to this knot
63+ , onKnotUsed :: d -> m v -- ^ This is what the knot will look like as a value
64+ -- to the inner computations
65+ , tieKnot :: d -> v -> m v -- ^ Once we're done and we're outside, tie the
66+ -- knot using the datum
69+unknot :: KnotTier v Void m
71+ { onKnotFound = error "Recursion detected"
72+ , onKnotUsed = absurd
51−-- | Try to call a function whilst memoizing its result. If we are already
76+-- | Run a potentially recursive computation. The provided key will be used to
52−-- inside an evaluation of some function with these arguments, this returns
77+-- refer to the result of this computation. If during the computation, another
53−-- @InProgress@. Beware of different functions that have the same type.
78+-- attempt to run the computation with the same key is made, we run a
54−tryMemo
79+-- tying-the-knot procedure.
81+-- If another attempt to run the computation with the same key is made
82+-- *after we're done*, we will return the memoized value.
55− :: forall k v m. (Typeable k, Typeable v, Ord k, MonadMemo m)
85+ (Typeable k, Typeable v, Typeable d, Ord k, MonadMemo s m)
56− => (k -> m v) -> k -> m (Progress v)
86+ => KnotTier v d m
87+ -> m v -- ^ the computation to memoize
88+ -> k -- ^ key for memoization
57−tryMemo f k = memoStateLookup k <$> get >>= \case
90+memoWithKnot tier f k = memoStateLookup @k @v k <$> get >>= \case
58− Just x -> pure x
91+ Just (Finished v) -> pure v
93+ d <- onKnotFound tier
94+ modify $ memoStateInsert @k @v k (TyingKnot $ toDyn d)
96+ Just (TyingKnot dyn) -> case fromDynamic dyn of
97+ Just d -> onKnotUsed tier d
98+ Nothing -> error $ "Type mismatch when examining the knot of "
99+ <> show (typeRep @(k -> v)) <> ": expected " <> show (typeRep @d)
100+ <> ", got " <> show (dynTypeRep dyn)
59 Nothing -> do
101 Nothing -> do
60− modify $ memoStateInsert k (InProgress @v)
102+ modify $ memoStateInsert @k @v k Started
61− result <- Finished <$> f k
103+ v <- f
104+ v' <- memoStateLookup @k @v k <$> get >>= \case
105+ Just Started -> pure v
106+ Just (TyingKnot dyn) -> case fromDynamic dyn of
107+ Just d -> tieKnot tier d v
108+ Nothing -> error $ "Type mismatch when tying the knot of "
109+ <> show (typeRep @(k -> v)) <> ": expected " <> show (typeRep @d)
110+ <> ", got " <> show (dynTypeRep dyn)
111+ Just (Finished _) -> error $ "Unexpected Finished when memoizing "
112+ <> show (typeRep @(k -> v))
113+ Nothing -> error $ "No key found when memoizing "
114+ <> show (typeRep @(k -> v))
62− modify $ memoStateInsert k result
115+ modify $ memoStateInsert @k @v k (Finished v')
63− pure result
116+ pure v'
65-- | Disambiguate memoized computations with an arbitrary tag.
118-- | Disambiguate memoized computations with an arbitrary tag.
66−tryMemoWithTag
119+memoTaggedWithKnot
120+ :: forall t k v d m s.
67− :: forall t k v m. (Typeable t, Typeable k, Typeable v, Ord k, MonadMemo m)
121+ ( Typeable t, Typeable k, Typeable v, Typeable d
68− => (k -> m v) -> k -> m (Progress v)
122+ , Ord k, MonadMemo s m )
123+ => KnotTier v d m -> m v -> k -> m v
69−tryMemoWithTag f k = withTypeable (typeRepKind $ typeRep @t) $
124+memoTaggedWithKnot tier f k = withTypeable (typeRepKind $ typeRep @t) $
70− tryMemo (f . unTagged) (Tagged @t k)
125+ memoWithKnot tier f (Tagged @t k)