    {-# OPTIONS_GHC -Wunused-imports #-}

    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' s (Maybe a) -> m a -> m a
    memo :: forall s (m :: * -> *) a.
    MonadState s m =>
    Lens' s (Maybe a) -> m a -> m a
    memo Lens' s (Maybe a)
    tbl m a
    compute = do
      mv <- Lens' s (Maybe a) -> m (Maybe a)
    forall o (m :: * -> *) i. MonadState o m => Lens' o i -> m i
    use (Maybe a -> f (Maybe a)) -> s -> f s
    Lens' s (Maybe a)
    tbl
      case mv of
        Just a
    x  -> a -> m a
    forall a. a -> m a
    forall (m :: * -> *) a. Monad m => a -> m a
    return a
    x
        Maybe a
    Nothing -> do
          x <- m a
    compute
          x <$ (tbl .= Just x)

    -- | Recursive memoisation, second argument is the value you get
    --   on recursive calls.
    memoRec :: MonadState s m => Lens' s (Maybe a) -> a -> m a -> m a
    memoRec :: forall s (m :: * -> *) a.
    MonadState s m =>
    Lens' s (Maybe a) -> a -> m a -> m a
    memoRec Lens' s (Maybe a)
    tbl a
    ih m a
    compute = do
      mv <- Lens' s (Maybe a) -> m (Maybe a)
    forall o (m :: * -> *) i. MonadState o m => Lens' o i -> m i
    use (Maybe a -> f (Maybe a)) -> s -> f s
    Lens' s (Maybe a)
    tbl
      case mv of
        Just a
    x  -> a -> m a
    forall a. a -> m a
    forall (m :: * -> *) a. Monad m => a -> m a
    return a
    x
        Maybe a
    Nothing -> do
          (Maybe a -> f (Maybe a)) -> s -> f s
    Lens' s (Maybe a)
    tbl Lens' s (Maybe a) -> Maybe a -> m ()
    forall o (m :: * -> *) i. MonadState o m => Lens' o i -> i -> m ()
    .= a -> Maybe a
    forall a. a -> Maybe a
    Just a
    ih
          x <- m a
    compute
          x <$ (tbl .= Just 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
      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
      return (unsafePerformIO . f' tbl)
      where
        f' :: IORef (Map a b) -> a -> IO b
    f' IORef (Map a b)
    tbl a
    x = do
          m <- IORef (Map a b) -> IO (Map a b)
    forall a. IORef a -> IO a
    readIORef IORef (Map a b)
    tbl
          case Map.lookup x m of
            Just b
    y  -> b -> IO b
    forall a. a -> IO a
    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 a. a -> IO a
    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
      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
      return (unsafePerformIO . f' tbl)
      where
        f' :: IORef (HashMap a b) -> a -> IO b
    f' IORef (HashMap a b)
    tbl a
    x = do
          m <- IORef (HashMap a b) -> IO (HashMap a b)
    forall a. IORef a -> IO a
    readIORef IORef (HashMap a b)
    tbl
          case HMap.lookup x m of
            Just b
    y  -> b -> IO b
    forall a. a -> IO a
    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 a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return b
    y