Pure functional memoization based on FunWithTypeFuns.

1.md

Port of memoization example from FunWithTypeFuns.

scala> memo.MemoExamples.b2i(true)
b=true
res1: Int = 1

scala> memo.MemoExamples.b2i(false)
b=false
res2: Int = 0

scala> memo.MemoExamples.b2i(false)
res3: Int = 0

scala> memo.MemoExamples.fib(5)
n=5
n=4
n=3
n=2
n=1
n=0
res0: Int = 8

scala> memo.MemoExamples.fib(3)
res4: Int = 3

scala> memo.MemoExamples.fib(6)
n=6
res5: Int = 13

Examples.scala

package memo

object MemoExamples {
  import Memo._

  val b2i = memo { b: Boolean =>
    if (b) {
      println("b=true")
      1
    } else {
      println("b=false")
      0
    }
  }

  @tailrec
  val fib: Int => Int = memo[Int, Int] {
    case 0 =>
      println("n=0")
      1
    case 1 =>
      println("n=1")
      1
    case n =>
      println("n=" + n)
      fib(n - 1) + fib(n - 2)
  }
}

Memo.scala

package memo

import scala.annotation.tailrec

trait Base[U, V]
trait Memo[K] {
  type Table[R] <: Base[K, R]

  def toTable[R](f: K => R): Table[R]
  def fromTable[R](t: Table[R]): K => R
}

object Memo {
  implicit def boolMemo = new Memo[Boolean] {
    class Table[W](u: => W, v: => W) extends Base[Boolean, W] {
      lazy val _1 = u
      lazy val _2 = v
    }
    def toTable[R](f: Boolean => R) = new Table(f(true), f(false))
    def fromTable[R](t: Table[R]) = x => if (x) t._1 else t._2
  }

  class TList[A, W, Q](w: => W, t: => Q) extends Base[Seq[A], W] {
    lazy val _1 = w
    lazy val _2 = t
  }

  implicit def listMemo[A](implicit M: Memo[A]) = new Memo[Seq[A]] { self =>
    type Table[W] = TList[A, W, M.Table[Base[Seq[A], W]]]

    def toTable[R](f: Seq[A] => R) = new Table[R](
      f(Nil),
      M.toTable { x => self.toTable { xs: Seq[A] => f(x +: xs) } }
    )

    @tailrec
    def fromTableRec[R](t: Table[R], s: Seq[A]): R = s match {
      case Nil => t._1
      case xs =>
        val table = M.fromTable(t._2)(xs.head)
        fromTableRec(table.asInstanceOf[Table[R]], xs.tail)
    }

    def fromTable[R](t: Table[R]): Seq[A] => R = s => fromTableRec(t, s)
  }

  implicit def intMemo = new Memo[Int] {
    val lmb = listMemo[Boolean]

    class Table[R](w: => lmb.Table[R]) extends Base[Int, R] {
      lazy val _1 = w
    }

    def toTable[R](f: Int => R) =
      new Table(lmb.toTable { bs => f(bitsToInt(bs)) })

    def fromTable[R](t: Table[R]): Int => R =
      n => lmb.fromTable(t._1)(intToBits(n))

    private[this] def bitsToInt(bits: Seq[Boolean]): Int = {
      var n = 0
      for (b <- bits) n = (n << 1) | (if (b) 1 else 0)
      n
    }

    private[this] def intToBits(n: Int): Seq[Boolean] = 
      (0 to 7).reverse map isBitSet(n)

    private[this] def isBitSet(n: Int)(bit: Int) = ((n >> bit) & 1) == 1

  }

  def memo[K, R](f: K => R)(implicit M: Memo[K]): K => R =
    M.fromTable(M.toTable(f))
}