関数の自動メモ化 (>scala 2.9.0)

Memoize.scala

object Memoize {

  // 1変数関数をメモ化する(2変数以上は tupled/untupledで対応)
  def memoize[ArgT, RetT](f: ArgT => RetT): ArgT => RetT = {
    val memo = scala.collection.mutable.Map.empty[ArgT, RetT]
    arg => memo.getOrElseUpdate(arg, f(arg))
  }
  
  // scala 2.9.0.RC1 or later
  def memoize[T1, T2, R](f: Function2[T1, T2, R]): Function2[T1, T2, R] =
    Function.untupled(memoize(f.tupled))

  def memoize[T1, T2, T3, R](f: Function3[T1, T2, T3, R]): Function3[T1, T2, T3, R] =
    Function.untupled(memoize(f.tupled))

  def memoize[T1, T2, T3, T4, R](f: Function4[T1, T2, T3, T4, R]): Function4[T1, T2, T3, T4, R] =
    Function.untupled(memoize(f.tupled))

  def memoize[T1, T2, T3, T4, T5, R](f: Function5[T1, T2, T3, T4, T5, R]): Function5[T1, T2, T3, T4, T5, R] =
    Function.untupled(memoize(f.tupled))
  
  
  // Functionのメソッドに memoized が付け足されたかのようにする
  implicit def equipMemoizedWithFunction[ArgT, RetT](f: ArgT => RetT) =
    new { def memoized = memoize(f) }
  
  // scala 2.9.0.RC1 or later
  implicit def equipMemoizedWithFunction2[T1, T2, R](f: Function2[T1, T2, R]) =
    new { def memoized = Function.untupled(memoize(f.tupled)) }
    
  implicit def equipMemoizedWithFunction3[T1, T2, T3, R](f: Function3[T1, T2, T3, R]) =
    new { def memoized = Function.untupled(memoize(f.tupled)) }
    
  implicit def equipMemoizedWithFunction4[T1, T2, T3, T4, R](f: Function4[T1, T2, T3, T4, R]) =
    new { def memoized = Function.untupled(memoize(f.tupled)) }
    
  implicit def equipMemoizedWithFunction5[T1, T2, T3, T4, T5, R](f: Function5[T1, T2, T3, T4, T5, R]) =
    new { def memoized = Function.untupled(memoize(f.tupled)) }

  
  def memoizeWith[ArgT, RetT](
      initialMemo: (ArgT, RetT)*)(f: ArgT => RetT): ArgT => RetT = {
    val memo = scala.collection.mutable.Map(initialMemo: _*)
    arg => memo.getOrElseUpdate(arg, f(arg))
  }
}

zmain.scala

object Main extends App {
  import Memoize._
  
  val fib1: Int => Long = memoize {
    case 0 => 0
    case 1 => 1
    case n => fib1(n - 1) + fib1(n - 2)
  }
  
  val fib2: Int => Long = 
    memoizeWith(0 -> 0L, 1 -> 1L){ n => fib2(n - 1) + fib2(n - 2) }
  
  var fib3: Int => Long = {
    case 0 => 0
    case 1 => 1
    case n => fib3(n - 1) + fib3(n - 2)
  }
  fib3 = fib3.memoized
  
  println("fib1(40) == "+fib1(40))
  println("fib2(40) == "+fib2(40))
  println("fib3(40) == "+fib3(40))
}

zzGenerator.scala

// 自動生成用(Function2 ~ Function5)
object MetaGenerator extends App {
  val templateFunction = """
    |  def memoize[%s](f: Function%d[%s]): Function%d[%s] =
    |    Function.untupled(memoize(f.tupled))
""".stripMargin

  for (arity <- 2 to 5) {
    val args = (1 to arity map { "T" +}):+ "R" mkString ", "
    print(templateFunction.format(args, arity, args, arity, args))
  }
  
  val templateImplicit = """    
    |  implicit def equipMemoizedWith%s(f: %s) =
    |    new { def memoized = Function.untupled(memoize(f.tupled)) }
""".stripMargin

  for (arity <- 2 to 5) {
    val args = 1 to arity map { "T"+ } mkString ", "
    val functionN = "Function%d[%s, R]".format(arity, args) 
    print(templateImplicit.format(Seq.fill(10)(functionN):_*))
  }
}

2.9.0.RC0って・・・？

RC1 のミスです。ご指摘ありがとう。