Sunday, January 11, 2009

Haskell memoization with unsafePerformIO

Searching for Haskell memoization gives a lot of interesting material.. It is summarized at the Haskell wiki (Memoization). But unsafe/impure memoize functions are usually burried deep enough to make them hard to find. By impure I mean functions with this signature: memoize :: (a -> b) -> (a -> b) Or at least: memoize :: Ord a => (a -> b) -> (a -> b) where Ord allows to use log-n collections for storage, such as maps. Like this:

module Memoize (memoize, memofix) where

import qualified Data.Map as Data.Map
import System.IO.Unsafe
import Control.Concurrent.MVar

memoize :: Ord a => (a -> b) -> (a -> b) 
memoize f =
    unsafePerformIO $
    do cacheRef <- newMVar Data.Map.empty
       return (\x -> unsafePerformIO (g cacheRef x))
      g cacheRef x = 
          do cache <- readMVar cacheRef
             case Data.Map.lookup x cache of
               Just y  -> return y
               Nothing -> 
                   do let y     = f x 
                      let cache = Data.Map.insert x y cache
                      swapMVar cacheRef cache
                      return y

memofix f = let g = memoize (f g) in g
I think the reason such memoize functions are shunned upon is because they run against the spirit of Haskell (use global state), and are usually unsafe in consuming unbounded memory (unless somehow cut off in the implementation). But they are still useful and much easier to use than pure versions.

import Memoize
let fib = memofix f where
  f rec 0 = 1
  f rec 1 = 1
  f rec n = rec (n - 1) + rec (n - 2)