@@ -1,70 +1,125 @@
11{-# LANGUAGE AllowAmbiguousTypes #-}
22-- | Utilities for effectfully memoizing other, more effectful functions.
33module OpenAPI.Checker.Memo
4− ( Progress(..)
5− , MonadMemo
4+ ( MonadMemo
65 , MemoState
76 , runMemo
8− , tryMemo
9− , tryMemoWithTag
7+ , modifyMemoNonce
8+ , KnotTier(..)
9+ , unknot
10+ , memoWithKnot
11+ , memoTaggedWithKnot
1012 ) where
1113
1214import Control.Monad.State
15+import Data.Dynamic
1316import qualified Data.Map as M
1417import Data.Tagged
18+import Data.Void
1519import qualified Data.TypeRepMap as T
1620import Type.Reflection
1721
18−-- | A computation has either 'Finished', or it is 'InProgress'.
19−data Progress a = Finished a | InProgress
20− deriving (Eq, Ord, Show)
22+data Progress a = Finished a | Started | TyingKnot Dynamic
2123
2224data MemoMap a where
2325 MemoMap :: !(M.Map k (Progress v)) -> MemoMap (k, v)
2426
25−newtype MemoState = MemoState (T.TypeRepMap MemoMap)
27+data MemoState s = MemoState s (T.TypeRepMap MemoMap)
2628
2729-- | An effectful memoization monad.
28−type MonadMemo m = MonadState MemoState m
30+type MonadMemo s m = MonadState (MemoState s) m
2931
3032memoStateLookup
31− :: forall k v. (Typeable k, Typeable v, Ord k)
32− => k -> MemoState -> Maybe (Progress v)
33−memoStateLookup k (MemoState t) = case T.lookup @(k, v) t of
33+ :: forall k v s. (Typeable k, Typeable v, Ord k)
34+ => k -> MemoState s -> Maybe (Progress v)
35+memoStateLookup k (MemoState _ t) = case T.lookup @(k, v) t of
3436 Just (MemoMap m) -> M.lookup k m
3537 Nothing -> Nothing
3638
3739memoStateInsert
38− :: forall k v. (Typeable k, Typeable v, Ord k)
39− => k -> Progress v -> MemoState -> MemoState
40−memoStateInsert k x (MemoState t) = MemoState $ T.insert (MemoMap m'') t
40+ :: forall k v s. (Typeable k, Typeable v, Ord k)
41+ => k -> Progress v -> MemoState s -> MemoState s
42+memoStateInsert k x (MemoState s t) = MemoState s $ T.insert (MemoMap m'') t
4143 where
4244 m'' = M.insert k x m'
4345 m' = case T.lookup @(k, v) t of
4446 Just (MemoMap m) -> m
4547 Nothing -> M.empty
4648
49+modifyMemoNonce :: MonadMemo s m => (s -> s) -> m s
50+modifyMemoNonce f = do
51+ MemoState s t <- get
52+ put $ MemoState (f s) t
53+ pure s
54+
4755-- | Run a memoized computation.
48−runMemo :: Monad m => StateT MemoState m a -> m a
49−runMemo = (`evalStateT` MemoState T.empty)
56+runMemo :: Monad m => s -> StateT (MemoState s) m a -> m a
57+runMemo s = (`evalStateT` MemoState s T.empty)
58+
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
67+ }
68+
69+unknot :: KnotTier v Void m
70+unknot = KnotTier
71+ { onKnotFound = error "Recursion detected"
72+ , onKnotUsed = absurd
73+ , tieKnot = absurd
74+ }
5075
51−-- | Try to call a function whilst memoizing its result. If we are already
52−-- inside an evaluation of some function with these arguments, this returns
53−-- @InProgress@. Beware of different functions that have the same type.
54−tryMemo
55− :: forall k v m. (Typeable k, Typeable v, Ord k, MonadMemo m)
56− => (k -> m v) -> k -> m (Progress v)
57−tryMemo f k = memoStateLookup k <$> get >>= \case
58− Just x -> pure x
76+-- | Run a potentially recursive computation. The provided key will be used to
77+-- refer to the result of this computation. If during the computation, another
78+-- attempt to run the computation with the same key is made, we run a
79+-- tying-the-knot procedure.
80+--
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.
83+memoWithKnot
84+ :: forall k v d m s.
85+ (Typeable k, Typeable v, Typeable d, Ord k, MonadMemo s m)
86+ => KnotTier v d m
87+ -> m v -- ^ the computation to memoize
88+ -> k -- ^ key for memoization
89+ -> m v
90+memoWithKnot tier f k = memoStateLookup @k @v k <$> get >>= \case
91+ Just (Finished v) -> pure v
92+ Just Started -> do
93+ d <- onKnotFound tier
94+ modify $ memoStateInsert @k @v k (TyingKnot $ toDyn d)
95+ onKnotUsed tier 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)
59101 Nothing -> do
60− modify $ memoStateInsert k (InProgress @v)
61− result <- Finished <$> f k
62− modify $ memoStateInsert k result
63− pure result
102+ modify $ memoStateInsert @k @v k Started
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))
115+ modify $ memoStateInsert @k @v k (Finished v')
116+ pure v'
64117
65118-- | Disambiguate memoized computations with an arbitrary tag.
66−tryMemoWithTag
67− :: forall t k v m. (Typeable t, Typeable k, Typeable v, Ord k, MonadMemo m)
68− => (k -> m v) -> k -> m (Progress v)
69−tryMemoWithTag f k = withTypeable (typeRepKind $ typeRep @t) $
70− tryMemo (f . unTagged) (Tagged @t k)
119+memoTaggedWithKnot
120+ :: forall t k v d m s.
121+ ( Typeable t, Typeable k, Typeable v, Typeable d
122+ , Ord k, MonadMemo s m )
123+ => KnotTier v d m -> m v -> k -> m v
124+memoTaggedWithKnot tier f k = withTypeable (typeRepKind $ typeRep @t) $
125+ memoWithKnot tier f (Tagged @t k)