Scala Memoization Patterns, in depth: pure-function caches, recursive memoization, key design, thread safety and bounded caches

When memoization is valid

A function can be memoized safely only if it is referentially transparent: called with equal arguments it always returns equal results, and it has no side effects you rely on happening every time. That is the property described in referential transparency in Scala. A function that reads the clock, a mutable field, a database or a random generator is not a candidate, because the cached answer goes stale. A function that logs, sends metrics or writes a file will silently stop doing those things after the first call.

The second condition is economic. A cache lookup costs a hash computation, an equality check and memory for every stored entry. Memoizing a function that adds two integers makes it slower. The wins come from functions whose cost is high relative to hashing the key, whose argument space repeats, and whose result is small enough to keep.

A generic memoize helper

For a single-threaded function of one argument, a mutable hash map and getOrElseUpdate are enough. The helper below returns a new function that closes over its own private cache, so the cache lives exactly as long as the function value does.

import scala.collection.mutable

def memoize[A, B](f: A => B): A => B = {
  val cache = mutable.HashMap.empty[A, B]
  a => cache.getOrElseUpdate(a, f(a))
}

def slowSquare(n: Int): Int = { Thread.sleep(500); n * n }

val fastSquare = memoize(slowSquare)
fastSquare(12)   // 500 ms
fastSquare(12)   // microseconds: cache hit

Functions of several arguments are memoized by turning the arguments into one key. Scala does this with tupled: memoize((g _).tupled) gives a function from a tuple, whose equality and hash code are derived from its elements. A case class works the same way and gives the key a name, which helps when the cache is exposed in metrics or logs.

Recursive functions: why the obvious version does not work

The classic demonstration is Fibonacci, and it hides the most common memoization bug. Wrapping a recursive function with the helper above caches only the outermost call. The recursion inside fib calls fib, the original method, not the memoized value, so fib(40) still makes hundreds of millions of calls the first time.

The fix is open recursion: write the function so that it receives the function to use for its recursive calls as a parameter, then tie the knot so that the parameter is the memoized version.

def memoRec[A, B](f: (A => B, A) => B): A => B = {
  val cache = mutable.HashMap.empty[A, B]
  lazy val self: A => B = a =>
    cache.get(a) match {
      case Some(b) => b
      case None =>
        val b = f(self, a)      // recursive calls go back through self
        cache.update(a, b)      // insert after computing, not inside getOrElseUpdate
        b
    }
  self
}

val fib: Int => BigInt = memoRec[Int, BigInt] { (rec, n) =>
  if (n < 2) BigInt(n) else rec(n - 1) + rec(n - 2)
}
fib(90)   // 91 distinct computations instead of about 10^19 calls

Two choices in that code are deliberate. It uses get followed by update rather than getOrElseUpdate, because the recursive calls insert other keys into the same map while the outer computation is still running. Modifying a map from inside the update function of its own getOrElseUpdate is outside that method's contract, so do not depend on whatever a particular version happens to do. And the depth of the recursion is unchanged by memoization: fib(100000) on a cold cache still recurses a hundred thousand frames deep and overflows the stack. Warm the cache bottom-up, or rewrite as a loop, as discussed in Scala tail recursion, in depth.

Key design

A cache is only as correct as its keys. The rules follow from how hash maps work.

• Keys must have value equality. Tuples, case classes, strings and boxed numbers do. Arrays do not: two arrays with identical contents are different keys, so convert to Vector or ArraySeq first.
• Keys must be immutable. A mutable object whose hash changes after insertion becomes unreachable in the map but still occupies memory.
• Normalise before lookup. If "Alice" and "alice " mean the same thing, lower-case and trim before using the string as a key, or you store the same result twice.
• Be careful with floating-point keys. 0.1 + 0.2 is not 0.3, and NaN-related equality is surprising. Round or quantise to the precision that actually matters.
• Keep keys small. A key that contains a large object pins that object in memory for as long as the entry lives. Key by an ID, not by the whole record.

Thread safety

A mutable.HashMap shared between threads can lose entries or be corrupted, and a memoized function stored in an object is shared by every thread that calls it. There are three standard options, and they differ in a property people rarely check: whether the underlying function may run more than once for the same key.

If the function is cheap and idempotent, TrieMap is fine: an occasional duplicate computation costs a little CPU. If it is an expensive remote call, use single-flight caching, where the cache stores the in-progress Future rather than the result, so concurrent callers all attach to the same computation. Remove failed futures, or one transient error is cached forever. The execution context behaviour that this relies on is covered in Scala Futures and execution contexts.

import scala.collection.concurrent.TrieMap
import scala.concurrent.{ExecutionContext, Future}
import scala.util.Failure

final class SingleFlight[K, V](load: K => Future[V])(implicit ec: ExecutionContext) {
  private val inFlight = TrieMap.empty[K, Future[V]]

  def get(k: K): Future[V] =
    inFlight.get(k).getOrElse {
      val p = scala.concurrent.Promise[V]()
      inFlight.putIfAbsent(k, p.future) match {
        case Some(existing) => existing           // another caller won the race
        case None =>
          p.completeWith(Future.delegate(load(k)))   // a synchronous throw fails the future too
          p.future.onComplete {
            case Failure(_) => inFlight.remove(k, p.future)   // do not cache failures
            case _          => ()
          }
          p.future
      }
    }
}

The Promise is created before the insert and only started by the winner, so the loader runs at most once per key even though putIfAbsent may lose a race. Successful entries stay in the map, so give it a bound as described next.

Bounded caches

An unbounded memo table is a memory leak with good intentions. A long-running service that memoizes by user ID, URL or query string will eventually hold every one it has ever seen. Before shipping a memoized function, decide what bounds it: a maximum entry count, a time to live, or a lifetime tied to a request or batch so that the whole cache is discarded together.

For a small, single-threaded bound, the JDK's LinkedHashMap in access order gives a least-recently-used cache in a few lines. For concurrent services, use a purpose-built library such as Caffeine, which provides size and time bounds, statistics and per-key atomic loading.

final class Lru[K, V](max: Int) {
  private val m = new java.util.LinkedHashMap[K, V](16, 0.75f, true) {   // true = access order
    override def removeEldestEntry(e: java.util.Map.Entry[K, V]): Boolean = size() > max
  }
  def getOrCompute(k: K)(f: K => V): V = m.synchronized {
    val hit = m.get(k)
    if (hit != null || m.containsKey(k)) hit
    else { val v = f(k); m.put(k, v); v }
  }
}

This version holds the lock while computing, which serialises all callers; that is acceptable for fast functions and wrong for slow ones. Whatever the implementation, export the hit rate. A memo cache with a 3 percent hit rate is costing memory and buying nothing.

Worked example: top-down memo versus a bottom-up table

Consider counting the ways to make change for an amount from a list of coin values. The recursive definition is natural: the ways to make n with coins from index i onward are the ways that use coin i at least once plus the ways that skip it. Memoized with memoRec on the key (n, i), it runs in time proportional to amount times coins.

def waysMemo(coins: Vector[Int], amount: Int): BigInt = {
  val go = memoRec[(Int, Int), BigInt] { case (rec, (n, i)) =>
    if (n == 0) 1
    else if (n < 0 || i == coins.length) 0
    else rec((n - coins(i), i)) + rec((n, i + 1))
  }
  go((amount, 0))
}

def waysTable(coins: Vector[Int], amount: Int): BigInt = {
  val t = Array.fill[BigInt](amount + 1)(0)
  t(0) = 1
  for (c <- coins; n <- c to amount) t(n) += t(n - c)
  t(amount)
}

Both return 292 for waysTable(Vector(1, 5, 10, 25, 50), 100). The memoized version reads like the definition and only computes the states it actually needs, which matters when most states are unreachable. But for amount = 1,000,000 it recurses a million frames deep and stores millions of boxed tuple keys. The table uses one array, has no recursion and no hashing, and is typically far faster. A good workflow is to write the memoized version first to get the recurrence right, test it on small inputs, and convert to a table once the evaluation order is obvious and the input sizes are large.

Effects and memoization

In code built on Cats Effect or ZIO, a function returning IO[B] is a description of work, not the work. Memoizing that function caches the description, so running the cached value still performs the effect every time. To share the result of an effect, use the effect system's own memoization operators, covered in the lazy evaluation article linked above, or a cache type built for that runtime. The same reasoning about failures applies: decide explicitly whether an error should be cached or retried.

Failure modes

What to do next

1. List the functions you want to memoize and confirm each is pure; drop any that read mutable state.
2. Add the generic memoize helper and use it for one expensive single-threaded function.
3. Rewrite any memoized recursion with open recursion and test that the inner calls hit the cache.
4. Check every cache key for value equality, immutability and normalisation.
5. For shared caches, pick TrieMap, computeIfAbsent or single-flight futures based on whether duplicate work is acceptable.
6. Put a size or time bound on every cache that outlives a request, and export its hit rate.
7. For large dynamic programming inputs, convert the memo to a bottom-up table.