Haskell memoization

memoize.hs

-- from http://www.nadineloveshenry.com/haskell/memofib.html#memoize
-- had some problems when function had more arguments though, so saving this for future me :)

import Debug.Trace
import Data.Map as Map
import Control.Monad.State.Lazy as State

--------------------------------------------

type StateMap a b = State (Map a b) b

memoizeM :: (Show a, Show b, Ord a) => 
            ((a -> StateMap a b) -> (a -> StateMap a b)) -> (a -> b)
memoizeM t x = evalState (f x) Map.empty where
  g x = do
    y <- t f x  
    m <- get
    put $ Map.insert x y m
    newM <- get
    return y
    --return $ trace ("Map now contains\n" ++ Map.showTree newM) y
  f x = get >>= \m -> maybe (g x) return (Map.lookup x m)


--------------------------------------------
{-

how to do it:

 - type:
     :: Type1 -> Type2 -> Endtype
   rewrite into
     :: Monad m =>
        ((Type1, Type2) -> m Endtype)
      -> (Type1, Type2) -> m Endtype

 - definition:
     myFun x y
   rewrite into
     myFun f (x,y)

 - call:
     myFun x y
   rewrite into
     f (x,y)
   and probably inside return or its own var (cause it's monad)

-}
--------------------------------------------

-- one argument

normalFib :: Integer -> Integer
normalFib 0 = 1
normalFib 1 = 1
normalFib n = normalFib (n-1) + normalFib (n-2)

fibM :: Monad m =>
        (Integer -> m Integer)
      -> Integer -> m Integer
fibM f 0 = return 1
fibM f 1 = return 1
fibM f n = do
  a <- f (n-1)
  b <- f (n-2)
  return (a+b)

fib n = memoizeM fibM n

--------------------------------------------

-- more arguments

normalAdd :: Int -> Int -> Int
normalAdd 0 y = y
normalAdd x y = normalAdd (x-1) (y+1)

addM :: Monad m =>
        ((Int, Int) -> m Int)
      -> (Int, Int) -> m Int
addM f (0,y) = return y
addM f (x,y) = do
  answer <- f ((x-1), (y+1))
  return answer

add x y = memoizeM addM (x,y)

--------------------------------------------

main = do
        print $ normalFib 10
        print $ fib 10
        print $ normalAdd 5 6
        print $ add 5 6