{-# 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
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)
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