{-# language CPP #-}

#define LAZY 1
#if(LAZY)
import Data.STRef.Lazy
import Control.Monad.ST.Lazy
#else
import Data.STRef.Strict
import Control.Monad.ST.Strict
#endif

import Data.Map (Map)
import qualified Data.Map as M

main = print $ runST $ do
    r <- newSTRef ( M.empty :: Map Int Int )
    fib r 2

fib :: STRef s ( Map Int Int )
    -> Int
    -> ST s Int
fib r x = cached r x $ case x of
    0 -> return 0
    1 -> return 1
    _ -> do 
        a <- fib r (x-1) 
        b <- fib r (x-2) 
        return $ a + b

cached :: Ord a 
       => STRef s ( Map a b )
       -> a
       -> ( ST s b )
       -> ST s b
cached r x computation = do
    m <- readSTRef r
    case M.lookup x m of
        Just y -> return y
        Nothing -> do
            y <- computation
            m <- readSTRef r
            writeSTRef r $ M.insert x y m
            return y
