{-# OPTIONS_GHC -Wunused-imports #-}

module Agda.Utils.Memo where

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

import Agda.Utils.Lens
import qualified Data.IntMap as IntMap

-- 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 <- Getting (Maybe a) s (Maybe a) -> m (Maybe a)
forall s (m :: * -> *) a. MonadState s m => Getting a s a -> m a
use Getting (Maybe a) s (Maybe a)
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 <- Getting (Maybe a) s (Maybe a) -> m (Maybe a)
forall s (m :: * -> *) a. MonadState s m => Getting a s a -> m a
use Getting (Maybe a) s (Maybe a)
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 -> Identity (Maybe a)) -> s -> Identity s
Lens' s (Maybe a)
tbl ((Maybe a -> Identity (Maybe a)) -> s -> Identity s)
-> Maybe a -> m ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> 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 :: (Hashable a) => (a -> b) -> (a -> b)
memoUnsafeH :: forall a b. 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. 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

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