I wrote up a polyfill of Function.memoize using WeakRefs (https://caniuse.com/?search=weakref) and a cache trie. The main point of this implementation is not the speed of memoization or retrieval of memoized values, but eliminating memory leaks where memoization caches hold on to otherwise unreachable objects, preventing them from being garbage collected.
I didn't do anything special with the hash function (like interpreting iterable return values as separate keys as @claudiameadows discussed) but that is something that could be further explored.
After implementing this, I realized that the consumer could also pass in a comparison function (currently this is done with === in the get() function) but that may be a footgun since doing deep or complex comparisons on so many values may be slower than the computation that was memoized away.
I also had difficulty in writing tests to determine if the cache had been modified by pruneTree() after a value had been deleted, so I relied on breakpoints in the browser DevTools to confirm that it worked as expected.
// Polyfill
Function.memoize = (function() {
const NO_VALUE = Symbol("NO_VALUE");
const VALUE_AT = Symbol("VALUE_AT");
const registry = new FinalizationRegistry(function schedulePruning(
cacheInstance
) {
// An object key in the cacheInstance has been Garbage Collected.
// Schedule a cleanup of the cache at the next idle period.
window.requestIdleCallback(() => pruneTree(cacheInstance));
});
function isWeakRef(wr) {
return Object.prototype.toString.call(wr).slice(8, -1) === "WeakRef";
}
function createCacheNode() {
const cache = new Map();
cache.set(VALUE_AT, NO_VALUE);
return cache;
}
// remove all branches that are keyed off of empty (GC'd) WeakRefs
function pruneTree(cache) {
for (const [k, v] of cache.entries()) {
if (isWeakRef(v) && !v.deref()) {
cache.delete(k);
} else if (k !== VALUE_AT) {
pruneTree(v);
}
}
}
// essentially Map.get where cache is a map with some WeakRef keys that
// need to be dereferenced to be checked, and the value is a cache node
function get(cache, key) {
for (const [k, v] of cache.entries()) {
const _key = isWeakRef(k) ? k.deref() : k;
if (_key === key) {
return v === NO_VALUE ? null : v;
}
}
return null;
}
function findValueByKeys(cache, ...keys) {
let currentCache = cache;
while (keys.length) {
currentCache = get(currentCache, keys.shift());
if (currentCache === null) {
return NO_VALUE;
}
}
return currentCache.get(VALUE_AT);
}
function setValueAt(cache, value, ...keys) {
let currentCache = cache;
// traverse to the leaf node
while (keys.length) {
let key = keys.shift();
let nextCache = get(currentCache, key);
if (nextCache === null) {
nextCache = createCacheNode();
if (key !== null && typeof key === "object") {
registry.register(key, currentCache);
key = new WeakRef(key);
}
currentCache.set(key, nextCache);
}
currentCache = nextCache;
}
currentCache.set(VALUE_AT, value);
}
return function(fn, hashFn) {
const rootCacheNode = createCacheNode();
const shouldHash = typeof hashFn === "function";
return function() {
let args = shouldHash ? hashFn(...arguments) : arguments;
let val = findValueByKeys(rootCacheNode, ...args);
if (val === NO_VALUE) {
val = fn.apply(null, arguments);
setValueAt(rootCacheNode, val, ...args);
}
return val;
};
};
})();
/**********************************/
/* Tests */
/**********************************/
const ctx = { maybeExists: { val: "foo" }, alwaysExists: { val: "bar" } };
let unmemoizedRunCount = 0;
function combineValues(a, b) {
unmemoizedRunCount++;
return (a && a.val) + (b && b.val);
}
const memoCombo = Function.memoize((a, b) => combineValues(a, b));
memoCombo(ctx.maybeExists, ctx.alwaysExists);
console.assert(
unmemoizedRunCount === 1,
"combineValues should have run unmemoized"
);
memoCombo(ctx.maybeExists, ctx.alwaysExists);
console.assert(
unmemoizedRunCount === 1,
"combineValues should have not have run"
);
memoCombo(ctx.maybeExists, ctx.alwaysExists);
console.assert(
unmemoizedRunCount === 1,
"combineValues should have not have run"
);
delete ctx.maybeExists;
memoCombo(ctx.maybeExists, ctx.alwaysExists);
console.assert(
unmemoizedRunCount === 2,
"combineValues should have run unmemoized"
);
EDIT: Better polyfill
Function.memoize = (function() {
const NO_VALUE = Symbol("NO_VALUE");
const registry = new FinalizationRegistry(function schedulePruning(
cacheInstance
) {
// An object key in the cacheInstance has been Garbage Collected.
// Schedule a cleanup of the cache at the next idle period.
window.requestIdleCallback(() => pruneTree(cacheInstance));
});
function isWeakRef(wr) {
return Object.prototype.toString.call(wr).slice(8, -1) === "WeakRef";
}
function createNode() {
const node = {
cache: new Map(),
value: NO_VALUE,
}
return node;
}
// remove all branches that are keyed off of empty (GC'd) WeakRefs
function pruneTree(cache) {
for (const [k, v] of cache.entries()) {
if (isWeakRef(k) && !k.deref()) {
cache.delete(k);
} else {
pruneTree(v.cache);
}
}
}
// essentially Map.get where cache is a map with some WeakRef keys that
// need to be dereferenced to be checked, and the value is a cache node
function getNode(node, key, comparisonFn) {
for (const [k, nextNode] of node.cache.entries()) {
const _key = isWeakRef(k) ? k.deref() : k;
if (comparisonFn(_key, key)) {
return nextNode;
}
}
return null;
}
function findValueByKeys(node, comparisonFn, ...keys) {
let currentNode = node;
while (keys.length) {
currentNode = getNode(currentNode, keys.shift(), comparisonFn);
if (currentNode === null) {
return NO_VALUE;
}
}
return currentNode.value;
}
function setValueAt(node, comparisonFn, value, ...keys) {
let currentNode = node;
// traverse to the leaf node
while (keys.length) {
let key = keys.shift();
let nextNode = getNode(currentNode, key, comparisonFn);
if (nextNode === null) {
nextNode = createNode();
if (key !== null && typeof key === "object") {
registry.register(key, currentNode.cache);
key = new WeakRef(key);
}
currentNode.cache.set(key, nextNode);
}
currentNode = nextNode;
}
currentNode.value = value;
}
function defaultCompare(a, b) {
return a === b;
}
return function(fn, hashFn, comparisonFn) {
const rootCacheNode = createNode();
hashFn = typeof hashFn === "function" ? hashFn : null;
comparisonFn = typeof comparisonFn === 'function' ? comparisonFn : defaultCompare;
return function() {
let args = hashFn ? hashFn(...arguments) : arguments;
let val = findValueByKeys(rootCacheNode, comparisonFn, ...args);
if (val === NO_VALUE) {
val = fn.apply(null, arguments);
setValueAt(rootCacheNode, comparisonFn, val, ...args);
}
return val;
};
};
})();
/**********************************/
/* Tests */
/**********************************/
const ctx = { maybeExists: { val: "foo" }, alwaysExists: { val: "bar" } };
let unmemoizedRunCount = 0;
function combineValues(a, b) {
unmemoizedRunCount++;
return (a && a.val) + (b && b.val);
}
const memoCombo = Function.memoize((a, b) => combineValues(a, b));
memoCombo(ctx.maybeExists, ctx.alwaysExists);
console.assert(
unmemoizedRunCount === 1,
"combineValues should have run unmemoized"
);
memoCombo(ctx.maybeExists, ctx.alwaysExists);
console.assert(
unmemoizedRunCount === 1,
"combineValues should have not have run"
);
memoCombo(ctx.maybeExists, ctx.alwaysExists);
console.assert(
unmemoizedRunCount === 1,
"combineValues should have not have run"
);
delete ctx.maybeExists;
memoCombo(ctx.maybeExists, ctx.alwaysExists);
console.assert(
unmemoizedRunCount === 2,
"combineValues should have run unmemoized"
);
.
