monadic memoization

gistfile1.scala

import scalaz._
import Scalaz._

object MonadicMemoization extends App {
  type Memo[Arg, Ret] = State[Map[Arg, Ret], Ret]

  def ret[Arg, Ret](r: Ret): Memo[Arg, Ret] =
    state[Map[Arg, Ret], Ret](x => (x, r))

  def memo[Arg, Ret](s: Arg => Memo[Arg, Ret])(k: Arg): Memo[Arg, Ret] =
    for {
      cache <- init[Map[Arg, Ret]]
      value <- cache.get(k) match {
	case Some(r) => ret[Arg, Ret](r)
	case None => for {
          r <- s(k)
          _ <- modify[Map[Arg, Ret]](_ + (k -> r))
        } yield r
      }
    } yield value

  def fib(n: BigInt): Memo[BigInt, BigInt] =
    if (n == 0) ret(1)
    else if (n == 1) ret(1)
    else for {
      a <- memo(fib)(n - 1)
      b <- memo(fib)(n - 2)
    } yield a + b

  def run[Arg, Ret](s: Memo[Arg, Ret]): Ret =
    s ! Map()

  println(run(fib(100)))
}