How do I make this higher-order memoization function work for recursive functions?

Question

I have a basic memoization function written as

function memo(func) {
  const cache = new Map()

  return function (...args) {
    const cacheKey = args.join('-')
    if (!cache.has(cacheKey)) {
      const value = func(...args)
      cache.set(cacheKey, value)
      return value
    }

    return cache.get(cacheKey)
  }
}

It doesn't work with functions that recursively calls itself. For example:

const fibonacci = (n) => {
  if (n <= 1) return 1
  return fibonacci(n - 1) + fibonacci(n - 2)
}

const memoizedFib = memo(fibonacci)
memoizedFib(20)

Here inside fibonacci, it still does a lot of duplicate calculations.

I guess a way to avoid that is to insert the memoization into the implementation for the function.

const cache = new Map();
const memoFibonacci = (n) => {
  if (memory.has(n)) return memory.get(n);
  if (n <= 1) return 1;
  const result = memoFibonacci(n - 1) + memoFibonacci(n - 2);
  memory.set(n, result);
  return result;
};

I wonder if there is a way to make the higher order util function work with recursive functions like fibonacci?

Answer 1

Here is a non-memoised recursive function as a benchmark:

const fibonacci = (n) => {
  console.log(`fibonacci(${n})...`)
  if (n <= 1) return 1
  return fibonacci(n - 1) + fibonacci(n - 2)
}

const result = fibonacci(5);

console.log("result:", result);

.as-console-wrapper { max-height: 100% !important }

Directly define a memoised function

You can define the function and memoise which would make all recursive calls use the memoised version:

How it looks

const fibonacci = memo((n) => {
  console.log(`fibonacci(${n})...`)
  if (n <= 1) return 1
  return fibonacci(n - 1) + fibonacci(n - 2)
});

Demo

function memo(func) {
  const cache = new Map()

  return function (...args) {
    const cacheKey = args.join('-')
    if (!cache.has(cacheKey)) {
      const value = func(...args)
      cache.set(cacheKey, value)
      return value
    }

    return cache.get(cacheKey)
  }
}

const fibonacci = memo((n) => {
  console.log(`fibonacci(${n})...`)
  if (n <= 1) return 1
  return fibonacci(n - 1) + fibonacci(n - 2)
});

const result = fibonacci(5);

console.log("result:", result);

.as-console-wrapper { max-height: 100% !important }

Replace original with memoised

You can also memoise it later by replacing the original binding for it. For this to work, the function should not be defined as const. This would still make future calls use the memoised version:

How it looks

let fibonacci = (n) => {
  console.log(`fibonacci(${n})...`)
  if (n <= 1) return 1
  return fibonacci(n - 1) + fibonacci(n - 2)
};

fibonacci = memo(fibonacci);

OR

function fibonacci(n) {
  console.log(`fibonacci(${n})...`)
  if (n <= 1) return 1
  return fibonacci(n - 1) + fibonacci(n - 2)
};

fibonacci = memo(fibonacci);

Demo

function memo(func) {
  const cache = new Map()

  return function (...args) {
    const cacheKey = args.join('-')
    if (!cache.has(cacheKey)) {
      const value = func(...args)
      cache.set(cacheKey, value)
      return value
    }

    return cache.get(cacheKey)
  }
}

let fibonacci = (n) => {
  console.log(`fibonacci(${n})...`)
  if (n <= 1) return 1
  return fibonacci(n - 1) + fibonacci(n - 2)
};

fibonacci = memo(fibonacci);

const result = fibonacci(5);

console.log("result:", result);

.as-console-wrapper { max-height: 100% !important }

---

Warnings

Be aware that this will only work for recursive functions that refer to a binding that you can control. Not all recursive functions are like that. Few examples:

Named function expressions

For example if the function is defined with a local name which only it can use to refer to itself:

How it looks

let fibonacci = function fibonacci(n) {
  console.log(`fibonacci(${n})...`)
  if (n <= 1) return 1
  return fibonacci(n - 1) + fibonacci(n - 2)
};

fibonacci = memo(fibonacci);

Demo

function memo(func) {
  const cache = new Map()

  return function (...args) {
    const cacheKey = args.join('-')
    if (!cache.has(cacheKey)) {
      const value = func(...args)
      cache.set(cacheKey, value)
      return value
    }

    return cache.get(cacheKey)
  }
}

let fibonacci = function fibonacci(n) {
  console.log(`fibonacci(${n})...`)
  if (n <= 1) return 1
  return fibonacci(n - 1) + fibonacci(n - 2)
};

fibonacci = memo(fibonacci);

const result = fibonacci(5);

console.log("result:", result);

.as-console-wrapper { max-height: 100% !important }

This is because the name of the function expression is usable in the body of the function and cannot be overwritten from the outside. Effectively, it is the same as making it:

let fibonacci = function recursive(n) {
  console.log(`fibonacci(${n})...`)
  if (n <= 1) return 1
  return recursive(n - 1) + recursive(n - 2)
};

fibonacci = memo(fibonacci);

Inner functions

It would also not work for functions that use an inner recursion helper:

let fibonacci = (n) => {
  const fibonacciHelper = (n) => {
    console.log(`fibonacci(${n})...`);
    if (n <= 1) return 1;
    return fibonacciHelper(n - 1) + fibonacciHelper(n - 2)
  }
  
  return fibonacciHelper(n);
};

fibonacci = memo(fibonacci);

Demo

function memo(func) {
  const cache = new Map()

  return function (...args) {
    const cacheKey = args.join('-')
    if (!cache.has(cacheKey)) {
      const value = func(...args)
      cache.set(cacheKey, value)
      return value
    }

    return cache.get(cacheKey)
  }
}

let fibonacci = (n) => {
  const fibonacciHelper = (n) => {
    console.log(`fibonacci(${n})...`);
    if (n <= 1) return 1;
    return fibonacciHelper(n - 1) + fibonacciHelper(n - 2)
  }
  
  return fibonacciHelper(n);
};

fibonacci = memo(fibonacci);

const result = fibonacci(5);

console.log("result:", result);

.as-console-wrapper { max-height: 100% !important }

Modules / other scopes

If you do not have access to the definition of the function, then you cannot really control the binding it uses to call itself. Easiest to see when using modules:

How it looks

fibonacci.js

export let fibonacci = (n) => {
  console.log(`fibonacci(${n})...`)
  if (n <= 1) return 1
  return fibonacci(n - 1) + fibonacci(n - 2)
}

index.js

import { fibonacci as fib } from "./fibonacci.js"

//assume memo() exists here

let fibonacci = memo(fib);
fibonacci(5);

This will not affect the recursive function since it still refers to itself from the module scope.