scala: memoize a function no matter how many arguments the function takes?

Question

i want to write a memoize function in scala that can be applied to any function object no matter what that function object is. i want to do so in a way that lets me use a single implementation of memoize. i'm flexible about the syntax, but ideally the memoize appears somewhere very close to the declaration of the function as opposed to after the function. i'd also like to avoid first declaring the original function and then a second declaration for the memoized version.

so some ideal syntax might be this:

def slowFunction(<some args left intentionally vague>) = memoize {
  // the original implementation of slow function
}

or even this would be acceptable:

def slowFUnction = memoize { <some args left intentionally vague> => {
  // the original implementation of slow function
}}

i've seen ways to do this where memoize must be redefined for each arity function, but i want to avoid this approach. the reason is that i will need to implement dozens of functions similar to memoize (i.e. other decorators) and it's too much to ask to have to copy each one for each arity function.

one way to do memoize that does require you to repeat memoize declarations (so it's no good) is at What type to use to store an in-memory mutable data table in Scala?.

Answer 1

You can use a type-class approach to deal with the arity issue. You will still need to deal with each function arity you want to support, but not for every arity/decorator combination:

/**
 * A type class that can tuple and untuple function types.
 * @param [U] an untupled function type
 * @param [T] a tupled function type
 */
sealed class Tupler[U, T](val tupled: U => T, 
                          val untupled: T => U)

object Tupler {
   implicit def function0[R]: Tupler[() => R, Unit => R] =
      new Tupler((f: () => R) => (_: Unit) => f(),
                 (f: Unit => R) => () => f(()))
   implicit def function1[T, R]: Tupler[T => R, T => R] = 
      new Tupler(identity, identity)
   implicit def function2[T1, T2, R]: Tupler[(T1, T2) => R, ((T1, T2)) => R] = 
      new Tupler(_.tupled, Function.untupled[T1, T2, R]) 
   // ... more tuplers
}

You can then implement the decorator as follows:

/**
 * A memoized unary function.
 *
 * @param f A unary function to memoize
 * @param [T] the argument type
 * @param [R] the return type
 */
class Memoize1[-T, +R](f: T => R) extends (T => R) {
   // memoization implementation
}

object Memoize {
   /**
    * Memoize a function.
    *
    * @param f the function to memoize
    */
   def memoize[T, R, F](f: F)(implicit e: Tupler[F, T => R]): F = 
      e.untupled(new Memoize1(e.tupled(f)))
}

Your "ideal" syntax won't work because the compiler would assume that the block passed into memoize is a 0-argument lexical closure. You can, however, use your latter syntax:

// edit: this was originally (and incorrectly) a def
lazy val slowFn = memoize { (n: Int) => 
   // compute the prime decomposition of n
}

Edit:

To eliminate a lot of the boilerplate for defining new decorators, you can create a trait:

trait FunctionDecorator {
   final def apply[T, R, F](f: F)(implicit e: Tupler[F, T => R]): F = 
      e.untupled(decorate(e.tupled(f)))

   protected def decorate[T, R](f: T => R): T => R
}

This allows you to redefine the Memoize decorator as

object Memoize extends FunctionDecorator {
   /**
    * Memoize a function.
    *
    * @param f the function to memoize
    */
   protected def decorate[T, R](f: T => R) = new Memoize1(f)
}

Rather than invoking a memoize method on the Memoize object, you apply the Memoize object directly:

// edit: this was originally (and incorrectly) a def
lazy val slowFn = Memoize(primeDecomposition _)

or

lazy val slowFn = Memoize { (n: Int) =>
   // compute the prime decomposition of n
}