Memo.hs

{-# LANGUAGE FlexibleInstances #-}

-- ==================================================================
-- Memo module
--   the code from http://www.sampou.org/cgi-bin/haskell.cgi?Memoise
--   with non-essential patches by @nfunato on 2013-09-07
--

module Memo
    (
     Table(..),
     Memo,
     memoize, evalMemo,
     fix2
     ) where

import Control.Monad.State
import Data.Function (fix)

----------------------------------------------------------------------
-- Fixpoint utilities
--

-- "fix" is now provided by Data.Function
-- fix :: (a -> a) -> a
-- fix f = let x = f x in x

fix2 :: (a -> b) -> (b -> a) -> b
fix2 = (fix .) . (.)

----------------------------------------------------------------------
-- Memo
--

class Table t where
  emptyTable  :: Ord a => t a b
  lookupTable :: Ord a => a -> t a b -> Maybe b
  insertTable :: Ord a => a -> b -> t a b -> t a b

instance (Table t, Ord a, Eq b, Num b) => Eq (State (t a b) b) where
  sx == sy  = evalState sx emptyTable == evalState sy emptyTable

instance (Table t, Ord a, Num b, Show b) => Show (State (t a b) b) where
  show sx   = show (evalState sx emptyTable)

instance (Table t, Ord a, Num b) => Num (State (t a b) b) where
  (+)    = liftM2 (+)
  (-)    = liftM2 (-)
  (*)    = liftM2 (*)
  negate = liftM negate
  abs    = liftM abs
  signum = liftM signum
  fromInteger = return . fromInteger

type Memo t a b = a -> State (t a b) b

memoize :: (Table t, Ord a) => Memo t a b -> Memo t a b
memoize mf x = do
    prev <- find x
    case prev of
      Just y  -> return y
      Nothing -> do
        y <- mf x
        ins x y
        return y
  where find x  = get >>= return . lookupTable x
        ins x y = get >>= put . insertTable x y

evalMemo :: (Table t, Ord a) => (Memo t) a b -> (->) a b
evalMemo m v = evalState (m v) emptyTable

----------------------------------------------------------------------
-- an example: fibonacchi
--

{-
import Memo
import qualified Data.Map as M

instance Table M.Map where
  emptyTable  = M.empty
  lookupTable = M.lookup
  insertTable = M.insert

fibF f 0 = 0
fibF f 1 = 1
fibF f n = f (n-2) + f (n-1)

ifib :: (->) Integer Integer
ifib = fix2 ($)     fibF    -- same with "ifib = fix fibF".

mfib :: Table t => (Memo t) Integer Integer
mfib = fix2 memoize fibF

memofib :: Integer -> Integer
memofib = evalMemo (mfib :: (Memo M.Map) Integer Integer)

*Main> memofib 100
354224848179261915075
-}