Search code examples
javascriptarraysobjectmemoization

Memoization: can the arguments be used as key in the cache object?


I have this solution for a memoization function.

const slice = Array.prototype.slice
function memoize(fn){
    const cache = {}
    return (...args) => {
        const params = slice.call(args)
        console.log(params)
        if(cache[params]){
            console.log('cached')
            return cache[params]
        } else{
            let result = fn(...args)
            cache[params] = result
            console.log('not cached')
            return result
        }
    }
}

cache[params] is the point. cache is an object and params is an array. Will this always work? After some researching I am still not confident this code is correct.


Solution

  • Such memoization can work for very simple cases, but it does not work in many other cases, for instance when:

    • the arguments are objects. They will often stringify to "[object Object]", and so different objects are treated as if they are the same.
    • the arguments are strings but cannot be distinguished because of the poor way they are separated when the arguments array is stringified (comma delimiter)

    Some demos of failures:

    1. Arguments are Objects

    const slice = Array.prototype.slice
    
    function memoize(fn){
        const cache = {}
        return (...args) => {
            const params = slice.call(args)
            if(cache[params]){
                return cache[params]
            } else{
                let result = fn(...args)
                cache[params] = result
                return result
            }
        }
    }
    
    // The function we will test with:
    function sum(a) {
        let total = 0;
        for (let value of a) total += value;
        return total;
    }
    
    function test() {
        console.log(sum(new Set([1,2,3]))); // should be 6
        console.log(sum(new Set([2,4,6]))); // should be 12
    }
    
    console.log("Run test without memoization...");
    test(); // Perform test without memoization
    sum = memoize(sum); // Memoize the function
    console.log("Run test WITH memoization...");
    test(); // Test again, but now with memoization

    1. Arguments are strings that have commas:

    const slice = Array.prototype.slice
    
    function memoize(fn){
        const cache = {}
        return (...args) => {
            const params = slice.call(args)
            if(cache[params]){
                return cache[params]
            } else{
                let result = fn(...args)
                cache[params] = result
                return result
            }
        }
    }
    
    // The function we will test with:
    function compareString(a, b) {
        return a.localeCompare(b); // returns -1, 0 or 1.
    }
    
    function test() {
        console.log(compareString("b,a", "c")); // should be -1
        // Change the arguments such that the concatenation with comma remains the same
        console.log(compareString("b", "a,c")); // should be  1
    }
    
    console.log("Run test without memoization...");
    test(); // Perform test without memoization
    compareString = memoize(compareString); // Memoize the function
    console.log("Run test WITH memoization...");
    test(); // Test again, but now with memoization

    Other remarks on the code

    • Calling slice is useless, as args is already a new array.
    • if(cache[params]) will evaluate to false when the already cached value is a falsy value, like 0, "", false. Doing if (params in cache) would avoid that issue
    • params will be stringified, which (as shown above) is not guaranteed to uniquely identify an array of arguments.

    Improvement

    If we can require that the arguments passed to our function are immutable, then we can use these immutable values or references as keys in a Map. This Map would become a tree of Maps when there are multiple arguments, so that when a lookup is made for the first argument in the main Map, it returns a nested Map, and then in that Map the next argument would be used as key, ...etc, until the deep Map is found for the last argument. In that final Map, a reserved key is used to retrieve the cached value. This reserved key can be Symbol that is only known by the memoize function, so that it can never collide with an argument value.

    Disclaimer: This will not work when objects can mutate between calls.

    Here is how that looks:

    function memoize(fn){
        const end = Symbol("end"); // A unique reference, only known here
        const cache = new Map;
        return (...args) => {
            let node = args.reduce((node, arg) => {
                 if (!node.has(arg)) node.set(arg, new Map);
                 return node.get(arg);
            }, cache);
            if (!node.has(end)) node.set(end, fn(...args));
            return node.get(end);
        }
    }
    
    // The function we will test with:
    let numCalls = 0;
    function length(a) { // Length of a linked list
        numCalls++; // Keep track of the number of calls made
        return a ? length(a.next) + 1 : 0;
    }
    
    function test() {
        numCalls = 0; // Reset the number of calls
        let head = { next: { next: { next: {}}}}; // Linked list with 4 nodes
        console.log(length(head)); // should be 4
        // Now exclude the head node:
        console.log(length(head.next)); // should be 3
        console.log("number of calls: ", numCalls); 
    }
    
    console.log("Run test without memoization...");
    test(); // Perform test without memoization
    length = memoize(length); // Memoize the function
    console.log("Run test WITH memoization...");
    test(); // Test again, but now with memoization

    Again, this cannot be used when objects mutate between calls. But for all other scenarios it should work fine.