Memoizing an effectful function

Question

I started working on a project defining a cellular automaton as a local transition function:

newtype Cellular g a = Cellular { delta :: (g -> a) -> a }

Whenever g is a Monoid, it is possible to define a global transition by shifting the focus before applying the local transition. This gives us the following step function:

step :: Monoid g => Cellular g a -> (g -> a) -> (g -> a)
step cell init g = delta cell $ init . (g <>)

Now, we can simply run the automaton by using iterate. And we can save a lot (and I do mean a lot: it literally saves hours) of re-computations by memoizing each one of the steps:

run :: (Monoid g, Memoizable g) => Cellular g a -> (g -> a) -> [g -> a]
run cell = iterate (memo . step cell)

My problem is that I generalised Cellular to CelluarT so that I would be able to use side effects in the local rules (e.g. copying a random neighbour):

newtype CellularT m g a = Cellular { delta :: (g -> m a) -> m a }

However, I only want the effects to be run once so that if you ask a cell multiple times what its value is, the answers are all consistent. memo fails us here because it saves the effectful computation rather than its result.

I don't expect this to be achievable without using unsafe features. I've tried to have a go at it using unsafePerformIO, an IORef and a Map g a to store the values already computed:

memoM :: (Ord k, Monad m) => (k -> m v) -> (k -> m v)
memoM =
  let ref = unsafePerformIO (newIORef empty) in
  ref `seq` loopM ref

loopM :: (Monad m, Ord k) => IORef (Map k v) -> (k -> m v) -> (k -> m v)
loopM ref f k =
  let m = unsafePerformIO (readIORef ref) in
  case Map.lookup k m of
    Just v  -> return v
    Nothing -> do
      v <- f k
      let upd = unsafePerformIO (writeIORef ref $ insert k v m)
      upd `seq` return v

But it behaves in unpredictable ways: memoM putStrLn is correctly memoized whilst memoM (\ str -> getLine) keeps fetching lines despite the same argument being passed to it.

Answer 1

This can be achieved safely if you give yourself an opportunity to allocate the reference to hold the map.

import Control.Monad.IO.Class

memoM :: (Ord k, MonadIO m) => (k -> m v) -> m (k -> m v)
                 |                           |
                 |        opportunity to allocate the map
                 get to IO correctly

I'm going to use an MVar instead of an IORef to get most of concurrency correct. This is for correctness, In case it's used concurrently, not for performance. For performance we might be fancier than this and use double-check locks or a concurrent map with finer lock granularity.

import Control.Concurrent    
import Control.Monad.IO.Class    
import qualified Data.Map as Map

memoM :: (Ord k, Monad m, MonadIO m) => (k -> m v) -> m (k -> m v)
memoM once = do 
    mapVar <- liftIO $ newMVar Map.empty    
    return (\k -> inMVar mapVar (lookupInsertM once k))

-- like withMVar, but isn't exception safe   
inMVar :: (MonadIO m) => MVar a -> (a -> m (a, b)) -> m b
inMVar mvar step = do
    (a, b) <- liftIO (takeMVar mvar) >>= step
    liftIO $ putMVar mvar a
    return b

lookupInsertM :: (Ord k, Monad m) => (k -> m v) -> k -> Map.Map k v -> m (Map.Map k v, v)
lookupInsertM once k map = 
    case Map.lookup k map of
        Just v -> return (map, v)
        Nothing -> do
            v <- once k
            return (Map.insert k v map, v)

We aren't really using IO, we're just passing around state. Any monad should be able to do so with a transformer applied to it, so why are we mucking about in IO? It's because we want to be able to allocate those maps so that memoM can be used for more than one different function. If we only ever care about a single memoized effectful function, we can just use a state transformer instead.

{-# LANGUAGE GeneralizedNewtypeDeriving #-}

import Control.Applicative
import Control.Monad.Trans.Class
import Control.Monad.Trans.State

newtype MemoT k v m a = MemoT {getMemoT :: StateT (k -> m v, Map.Map k v) m a}
    deriving (Functor, Applicative, Monad, MonadIO)

instance MonadTrans (MemoT k v) where
    lift = MemoT . lift

This transformer adds the ability to lookup a value from the memoized effectful function

lookupMemoT :: (Ord k, Monad m) => k -> MemoT k v m v
lookupMemoT k = MemoT . StateT $ \(once, map) -> do
                                                    (map', v) <- lookupInsertM once k map
                                                    return (v, (once, map'))

To run it and get at the underlying monad, we need to provide the effectful function we want to memoize.

runMemoT :: (Monad m) => MemoT k v m a -> (k -> m v) -> m a
runMemoT memo once = evalStateT (getMemoT memo) (once, Map.empty)

Our MemoT uses a Map for every function. Some functions might be memoized some other way. The monad-memo package has an mtl-style class for monads that provide memoization for a specific function, and a more elaborate mechanism for building them that doesn't necessarily use Maps.