MemoFix Example

memoFix.hs

import Control.Concurrent.MVar
import Control.Monad.Fix (mfix)
import qualified Data.Map as M

fix :: (t -> t) -> t
fix f = f (fix f)

fib :: (Eq a, Num a) => (a -> a) -> a -> a
fib _ 0 = 0
fib _ 1 = 1
fib r n = r (n-1) + r (n-2)

{-
fix fib : a-> a, using t |-> a -> a
-}

memo :: (a -> b) -> a -> b
memo f = f -- TODO a pure memo!

memoFix :: ((a -> b) -> a -> b) -> a -> b
memoFix f = let mf = memo (f mf) in mf

----------------------------------------
-- Monadic memoization

fibM :: (Eq a, Num a) => (a -> IO a) -> a -> IO a
fibM _ 0 = return 0
fibM _ 1 = return 1
fibM r n = do
  a <- r (n-1)
  b <- r (n-2)
  return $ a + b

{-
fix fib : a-> a, using t |-> a -> a
-}

memoIO :: (Ord a)
       => (a -> IO b)
       -> IO (a -> IO b)
memoIO f = do
  v <- newMVar M.empty
  return $ \x -> do
    m <- readMVar v
    case M.lookup x m of
      Nothing -> do
        r <- f x
        modifyMVar_ v (return . M.insert x r)
        return r
      Just r  -> return r

memoFixIO :: Ord a => ((a -> IO b) -> a -> IO b) -> IO (a -> IO b)
memoFixIO f = mfix $ \mf -> memoIO (f mf)

test = do
  fibMemo <- memoFixIO fibM
  fibMemo 10000