    module Agda.Utils.Memo where

    import Control.Monad.State
    import System.IO.Unsafe
    import Data.IORef
    import qualified Data.Map as Map
    import qualified Data.HashMap.Strict as HMap
    import Data.Hashable

    import Agda.Utils.Lens

    -- Simple memoisation in a state monad

    -- | Simple, non-reentrant memoisation.
    memo :: MonadState s m => Lens' (Maybe a) s -> m a -> m a
    memo :: forall s (m :: * -> *) a.
    MonadState s m =>
    Lens' (Maybe a) s -> m a -> m a
    memo Lens' (Maybe a) s
    tbl m a
    compute = do
      Maybe a
    mv <- Lens' (Maybe a) s -> m (Maybe a)
    forall o (m :: * -> *) i. MonadState o m => Lens' i o -> m i
    use Lens' (Maybe a) s
    tbl
      case Maybe a
    mv of
        Just a
    x  -> a -> m a
    forall (m :: * -> *) a. Monad m => a -> m a
    return a
    x
        Maybe a
    Nothing -> do
          a
    x <- m a
    compute
          a
    x a -> m () -> m a
    forall (f :: * -> *) a b. Functor f => a -> f b -> f a
    <$ (Lens' (Maybe a) s
    tbl Lens' (Maybe a) s -> Maybe a -> m ()
    forall o (m :: * -> *) i. MonadState o m => Lens' i o -> i -> m ()
    .= a -> Maybe a
    forall a. a -> Maybe a
    Just a
    x)

    -- | Recursive memoisation, second argument is the value you get
    --   on recursive calls.
    memoRec :: MonadState s m => Lens' (Maybe a) s -> a -> m a -> m a
    memoRec :: forall s (m :: * -> *) a.
    MonadState s m =>
    Lens' (Maybe a) s -> a -> m a -> m a
    memoRec Lens' (Maybe a) s
    tbl a
    ih m a
    compute = do
      Maybe a
    mv <- Lens' (Maybe a) s -> m (Maybe a)
    forall o (m :: * -> *) i. MonadState o m => Lens' i o -> m i
    use Lens' (Maybe a) s
    tbl
      case Maybe a
    mv of
        Just a
    x  -> a -> m a
    forall (m :: * -> *) a. Monad m => a -> m a
    return a
    x
        Maybe a
    Nothing -> do
          Lens' (Maybe a) s
    tbl Lens' (Maybe a) s -> Maybe a -> m ()
    forall o (m :: * -> *) i. MonadState o m => Lens' i o -> i -> m ()
    .= a -> Maybe a
    forall a. a -> Maybe a
    Just a
    ih
          a
    x <- m a
    compute
          a
    x a -> m () -> m a
    forall (f :: * -> *) a b. Functor f => a -> f b -> f a
    <$ (Lens' (Maybe a) s
    tbl Lens' (Maybe a) s -> Maybe a -> m ()
    forall o (m :: * -> *) i. MonadState o m => Lens' i o -> i -> m ()
    .= a -> Maybe a
    forall a. a -> Maybe a
    Just a
    x)

    {-# NOINLINE memoUnsafe #-}
    memoUnsafe :: Ord a => (a -> b) -> (a -> b)
    memoUnsafe :: forall a b. Ord a => (a -> b) -> a -> b
    memoUnsafe a -> b
    f = IO (a -> b) -> a -> b
    forall a. IO a -> a
    unsafePerformIO (IO (a -> b) -> a -> b) -> IO (a -> b) -> a -> b
    forall a b. (a -> b) -> a -> b
    $ do
      IORef (Map a b)
    tbl <- Map a b -> IO (IORef (Map a b))
    forall a. a -> IO (IORef a)
    newIORef Map a b
    forall k a. Map k a
    Map.empty
      (a -> b) -> IO (a -> b)
    forall (m :: * -> *) a. Monad m => a -> m a
    return (IO b -> b
    forall a. IO a -> a
    unsafePerformIO (IO b -> b) -> (a -> IO b) -> a -> b
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . IORef (Map a b) -> a -> IO b
    f' IORef (Map a b)
    tbl)
      where
        f' :: IORef (Map a b) -> a -> IO b
    f' IORef (Map a b)
    tbl a
    x = do
          Map a b
    m <- IORef (Map a b) -> IO (Map a b)
    forall a. IORef a -> IO a
    readIORef IORef (Map a b)
    tbl
          case a -> Map a b -> Maybe b
    forall k a. Ord k => k -> Map k a -> Maybe a
    Map.lookup a
    x Map a b
    m of
            Just b
    y  -> b -> IO b
    forall (m :: * -> *) a. Monad m => a -> m a
    return b
    y
            Maybe b
    Nothing -> do
              let y :: b
    y = a -> b
    f a
    x
              IORef (Map a b) -> Map a b -> IO ()
    forall a. IORef a -> a -> IO ()
    writeIORef IORef (Map a b)
    tbl (a -> b -> Map a b -> Map a b
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    Map.insert a
    x b
    y Map a b
    m)
              b -> IO b
    forall (m :: * -> *) a. Monad m => a -> m a
    return b
    y

    {-# NOINLINE memoUnsafeH #-}
    memoUnsafeH :: (Eq a, Hashable a) => (a -> b) -> (a -> b)
    memoUnsafeH :: forall a b. (Eq a, Hashable a) => (a -> b) -> a -> b
    memoUnsafeH a -> b
    f = IO (a -> b) -> a -> b
    forall a. IO a -> a
    unsafePerformIO (IO (a -> b) -> a -> b) -> IO (a -> b) -> a -> b
    forall a b. (a -> b) -> a -> b
    $ do
      IORef (HashMap a b)
    tbl <- HashMap a b -> IO (IORef (HashMap a b))
    forall a. a -> IO (IORef a)
    newIORef HashMap a b
    forall k v. HashMap k v
    HMap.empty
      (a -> b) -> IO (a -> b)
    forall (m :: * -> *) a. Monad m => a -> m a
    return (IO b -> b
    forall a. IO a -> a
    unsafePerformIO (IO b -> b) -> (a -> IO b) -> a -> b
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . IORef (HashMap a b) -> a -> IO b
    f' IORef (HashMap a b)
    tbl)
      where
        f' :: IORef (HashMap a b) -> a -> IO b
    f' IORef (HashMap a b)
    tbl a
    x = do
          HashMap a b
    m <- IORef (HashMap a b) -> IO (HashMap a b)
    forall a. IORef a -> IO a
    readIORef IORef (HashMap a b)
    tbl
          case a -> HashMap a b -> Maybe b
    forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
    HMap.lookup a
    x HashMap a b
    m of
            Just b
    y  -> b -> IO b
    forall (m :: * -> *) a. Monad m => a -> m a
    return b
    y
            Maybe b
    Nothing -> do
              let y :: b
    y = a -> b
    f a
    x
              IORef (HashMap a b) -> HashMap a b -> IO ()
    forall a. IORef a -> a -> IO ()
    writeIORef IORef (HashMap a b)
    tbl (a -> b -> HashMap a b -> HashMap a b
    forall k v.
    (Eq k, Hashable k) =>
    k -> v -> HashMap k v -> HashMap k v
    HMap.insert a
    x b
    y HashMap a b
    m)
              b -> IO b
    forall (m :: * -> *) a. Monad m => a -> m a
    return b
    y