[Haskell-cafe] Memoisation + unsafePerformIO

[Ryan Ingram]

Funny that you mention using unsafePerformIO to do memoizing; I just
implemented it a couple of days ago to familiarize myself with
techniques that use global state.  Here's my implementation which uses
trees (Data.Map).

module Memoize (memoize, memoizefix) where
import System.IO.Unsafe
import qualified Data.Map as Map
import Data.IORef

memoized :: Ord a => IORef (Map.Map a b) -> (a -> b) -> a -> b
memoized memTbl f a = unsafePerformIO $ do
   memo <- readIORef memTbl
   catch (Map.lookup a memo) $ \_ -> do
      let ans = f a
      let memo' = Map.insert a ans memo
      writeIORef memTbl memo'
      return ans

memoize :: Ord a => (a -> b) -> (a -> b)
memoize f = unsafePerformIO $ do
    memTbl <- newIORef Map.empty
    return $ memoized memTbl f

memFix :: Ord a => IORef (Map.Map a b) -> ((a -> b) -> (a -> b)) -> (a -> b)
memFix memTbl f = let x = f (memoized memTbl x) in x

memoizefix :: Ord a => ((a -> b) -> (a -> b)) -> (a -> b)
memoizefix f = unsafePerformIO $ do
    memTbl <- newIORef Map.empty
    return $ memFix memTbl f

A test case:
module Main where
import Memoize

fix f = let x = f x in x
nfibr f x = if x <= 1 then (1::Integer) else f (x-1) + f (x-2)
nfib = fix nfibr
mfib = memoizefix nfibr

main = sequence_ $ map (print . mfib) [1..]

(compare replacing mfib with nfib in main).

I think you'll have problems with overlapping instances with your
Ix/Ord instance declarations; better to not use typeclasses there and
just have memoIx and memoTree, I think.

I am a bit nervous about what happens if you turn optimizations on;
you might need to sprinkle a few {-# NOINLINE function_name #-}
pragmas into Memoize.hs to make sure it works properly in a real
codebase, but it's a good start.

  -- ryan