memorization via ST hashtable

memorize.hs

{-# LANGUAGE RankNTypes #-}

import Data.Int
import Control.Monad.ST
import Control.Monad
import Data.Function
import Data.Hashable
import qualified Data.HashTable.ST.Cuckoo as H

type Func a b = (a -> b) -> (a -> b)
type Func2 a b = forall m. Monad m => (a -> m b) -> (a -> m b)

memorize :: (Eq a, Hashable a) => Func2 a b -> (a -> b)
memorize func a = runST $ do
    h <- H.new
    let loop a = do
        res <- H.lookup h a
        case res of
            Just b  -> return b
            Nothing -> do
                b <- func loop a
                H.insert h a b
                return b
    loop a

fib1 :: Func2 Integer Integer
fib1 rec n | n <= 2    = return 1
           | otherwise = liftM2 (+) (rec (n - 1)) (rec (n - 2))

fib2 :: Func Integer Integer
fib2 rec n | n <= 2    = 1
           | otherwise = rec (n - 1) + rec (n - 2)

main = do
    let fib1' = memorize fib1
        fib2' = fix fib2
    print (map fib1' [1..100])
    print (map fib2' [1..10])