Memoizing (safe use of unsafePerformIO)

Memoize.hs

module Memoize (memoize) where

import Data.Map as M
import Data.IORef
import System.IO.Unsafe

-- | Memoize a function. The use of unsafePerformIO is fine since referential transparency is maintained.
-- `memoize f` will return a new function which behaves the same as `f` but memoizes its results everytime
-- it is called with an argument it hasn't seen yet.
memoize :: Ord a => (a -> b) -> (a -> b)
memoize f = \x -> unsafePerformIO $ do
      table <- readIORef dict
      case M.lookup x table of
        (Just y) -> return y
        Nothing  -> do
          let y = f x
          modifyIORef dict $ insert x y
          return y
  where
    dict :: IORef (Map a b)
    dict = unsafePerformIO $ newIORef empty