module Data.MemoUgly(memoIO, memo) where
import Control.Concurrent.MVar
import qualified Data.Map as M
import System.IO.Unsafe(unsafePerformIO)

-- | Memoize the given function by allocating a memo table,
-- and then updating the memo table on each function call.
memoIO :: (Ord a)
       => (a -> b)           -- ^Function to memoize
       -> IO (a -> IO b)
memoIO :: forall a b. Ord a => (a -> b) -> IO (a -> IO b)
memoIO a -> b
f = do
    MVar (Map a b)
v <- Map a b -> IO (MVar (Map a b))
forall a. a -> IO (MVar a)
newMVar Map a b
forall k a. Map k a
M.empty
    let f' :: a -> IO b
f' a
x = do
            Map a b
m <- MVar (Map a b) -> IO (Map a b)
forall a. MVar a -> IO a
readMVar MVar (Map a b)
v
            case a -> Map a b -> Maybe b
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
x Map a b
m of
                Maybe b
Nothing -> do let { r :: b
r = a -> b
f a
x }; MVar (Map a b) -> (Map a b -> IO (Map a b)) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar (Map a b)
v (Map a b -> IO (Map a b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Map a b -> IO (Map a b))
-> (Map a b -> Map a b) -> Map a b -> IO (Map a b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b -> Map a b -> Map a b
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
x b
r); b -> IO b
forall (m :: * -> *) a. Monad m => a -> m a
return b
r
                Just b
r  -> b -> IO b
forall (m :: * -> *) a. Monad m => a -> m a
return b
r
    (a -> IO b) -> IO (a -> IO b)
forall (m :: * -> *) a. Monad m => a -> m a
return a -> IO b
f'

-- | The pure version of 'memoIO'.
memo :: (Ord a)
     => (a -> b)           -- ^Function to memoize
     -> (a -> b)
memo :: forall a b. Ord a => (a -> b) -> a -> b
memo a -> b
f = let f' :: a -> IO b
f' = IO (a -> IO b) -> a -> IO b
forall a. IO a -> a
unsafePerformIO ((a -> b) -> IO (a -> IO b)
forall a b. Ord a => (a -> b) -> IO (a -> IO b)
memoIO a -> b
f) in \ a
x -> IO b -> b
forall a. IO a -> a
unsafePerformIO (a -> IO b
f' a
x)