1 files changed, 98 insertions, 0 deletions
diff --git a/users/grfn/xanthous/src/Xanthous/Data/Memo.hs b/users/grfn/xanthous/src/Xanthous/Data/Memo.hs
new file mode 100644
index 000000000000..2b2ee0f96028
--- /dev/null
+++ b/users/grfn/xanthous/src/Xanthous/Data/Memo.hs
@@ -0,0 +1,98 @@
+--------------------------------------------------------------------------------
+-- | Memoized values
+--------------------------------------------------------------------------------
+module Xanthous.Data.Memo
+  ( Memoized(UnMemoized)
+  , memoizeWith
+  , getMemoized
+  , runMemoized
+  , fillWith
+  , fillWithM
+  ) where
+--------------------------------------------------------------------------------
+import Xanthous.Prelude
+import Data.Aeson (FromJSON, ToJSON)
+import Test.QuickCheck (Arbitrary (arbitrary), oneof, CoArbitrary, Function)
+import Test.QuickCheck.Checkers (EqProp)
+import Xanthous.Util (EqEqProp(EqEqProp))
+import Control.Monad.State.Class (MonadState)
+--------------------------------------------------------------------------------
+
+-- | A memoized value, keyed by a key
+--
+-- If key is different than what is stored here, then val is invalid
+data Memoized key val = Memoized key val | UnMemoized
+  deriving stock (Show, Eq, Generic)
+  deriving anyclass (Hashable, FromJSON, ToJSON, NFData, CoArbitrary, Function)
+  deriving EqProp via EqEqProp (Memoized key val)
+
+instance (Arbitrary k, Arbitrary v) => Arbitrary (Memoized k v) where
+  arbitrary = oneof [ pure UnMemoized
+                    , Memoized <$> arbitrary <*> arbitrary
+                    ]
+
+-- | Construct a memoized value with the given key
+memoizeWith :: forall key val. key -> val -> Memoized key val
+memoizeWith = Memoized
+{-# INLINE memoizeWith #-}
+
+-- | Retrieve a memoized value providing the key. If the value is unmemoized or
+-- the keys do not match, returns Nothing.
+--
+-- >>> getMemoized 1 (memoizeWith @Int @Int 1 2)
+-- Just 2
+--
+-- >>> getMemoized 2 (memoizeWith @Int @Int 1 2)
+-- Nothing
+--
+-- >>> getMemoized 1 (UnMemoized :: Memoized Int Int)
+-- Nothing
+getMemoized :: Eq key => key -> Memoized key val -> Maybe val
+getMemoized key (Memoized key' v)
+  | key == key' = Just v
+  | otherwise = Nothing
+getMemoized _ UnMemoized = Nothing
+{-# INLINE getMemoized #-}
+
+-- | Get a memoized value using an applicative action to obtain the key
+runMemoized
+  :: (Eq key, Applicative m)
+  => Memoized key val
+  -> m key
+  -> m (Maybe val)
+runMemoized m mk = getMemoized <$> mk <*> pure m
+
+-- | In a monadic state containing a 'MemoState', look up the current memoized
+-- target of some lens keyed by k, filling it with v if not present and
+-- returning either the new or old value
+fillWith
+  :: forall m s k v.
+    (MonadState s m, Eq k)
+  => Lens' s (Memoized k v)
+  -> k
+  -> v
+  -> m v
+fillWith l k v' = do
+  uses l (getMemoized k) >>= \case
+    Just v -> pure v
+    Nothing -> do
+      l .= memoizeWith k v'
+      pure v'
+
+-- | In a monadic state, look up the current memoized target of some lens keyed
+-- by k, filling it with the result of some monadic action v if not present and
+-- returning either the new or old value
+fillWithM
+  :: forall m s k v.
+    (MonadState s m, Eq k)
+  => Lens' s (Memoized k v)
+  -> k
+  -> m v
+  -> m v
+fillWithM l k mv = do
+  uses l (getMemoized k) >>= \case
+    Just v -> pure v
+    Nothing -> do
+      v' <- mv
+      l .= memoizeWith k v'
+      pure v'