A production-ready memoization utility that intelligently caches function results using WeakMap for object arguments (preventing memory leaks) and Map for primitive values, with support for custom cache key serialization.
type Primitive = string | number | boolean | bigint | symbol | null | undefined;
type CacheKeySerializer = (args: unknown[]) => string;
interface MemoizeOptions {
serializer?: CacheKeySerializer;
maxSize?: number;
}
interface CacheNode {
primitiveCache: Map<string, unknown>;
objectCache: WeakMap<object, CacheNode>;
result?: { value: unknown };
}
function createCacheNode(): CacheNode {
return {
primitiveCache: new Map(),
objectCache: new WeakMap(),
};
}
function isPrimitive(value: unknown): value is Primitive {
return value === null || (typeof value !== 'object' && typeof value !== 'function');
}
const defaultSerializer: CacheKeySerializer = (args) => {
return args.map((arg) => {
if (arg === null) return 'null';
if (arg === undefined) return 'undefined';
if (typeof arg === 'symbol') return arg.toString();
if (typeof arg === 'bigint') return `${arg}n`;
return String(arg);
}).join('|');
};
function memoize<TArgs extends unknown[], TResult>(
fn: (...args: TArgs) => TResult,
options: MemoizeOptions = {}
): ((...args: TArgs) => TResult) & { cache: { clear: () => void; size: () => number } } {
const { serializer = defaultSerializer, maxSize } = options;
const rootCache = createCacheNode();
const primitiveKeys: string[] = [];
function getOrCreateCacheNode(args: TArgs): CacheNode {
let currentNode = rootCache;
for (const arg of args) {
if (isPrimitive(arg)) {
const key = serializer([arg]);
let nextNode = currentNode.primitiveCache.get(key) as CacheNode | undefined;
if (!nextNode) {
nextNode = createCacheNode();
currentNode.primitiveCache.set(key, nextNode);
if (maxSize && primitiveKeys.length >= maxSize) {
const oldestKey = primitiveKeys.shift();
if (oldestKey) {
rootCache.primitiveCache.delete(oldestKey);
}
}
primitiveKeys.push(key);
}
currentNode = nextNode;
} else {
const objArg = arg as object;
let nextNode = currentNode.objectCache.get(objArg);
if (!nextNode) {
nextNode = createCacheNode();
currentNode.objectCache.set(objArg, nextNode);
}
currentNode = nextNode;
}
}
return currentNode;
}
function memoized(...args: TArgs): TResult {
const cacheNode = getOrCreateCacheNode(args);
if (cacheNode.result !== undefined) {
return cacheNode.result.value as TResult;
}
const result = fn(...args);
cacheNode.result = { value: result };
return result;
}
memoized.cache = {
clear: () => {
rootCache.primitiveCache.clear();
rootCache.objectCache = new WeakMap();
rootCache.result = undefined;
primitiveKeys.length = 0;
},
size: () => primitiveKeys.length,
};
return memoized;
}
// Example usage and demonstration
// Basic usage with primitives
const fibonacci = memoize((n: number): number => {
console.log(`Computing fibonacci(${n})`);
if (n <= 1) return n;
return fibonacci(n - 1) + fibonacci(n - 2);
});
console.log('Fibonacci examples:');
console.log(fibonacci(10));
console.log(fibonacci(10)); // Cached - no log
// Usage with object arguments (WeakMap prevents memory leaks)
interface User {
id: number;
name: string;
}
const getUserPermissions = memoize((user: User, resource: string): string[] => {
console.log(`Fetching permissions for ${user.name} on ${resource}`);
return ['read', 'write'];
});
console.log('\nObject argument examples:');
const user1: User = { id: 1, name: 'Alice' };
console.log(getUserPermissions(user1, 'documents'));
console.log(getUserPermissions(user1, 'documents')); // Cached
// Custom serializer example
const customMemoize = memoize(
(date: Date, format: string): string => {
console.log(`Formatting date with ${format}`);
return date.toISOString();
},
{
serializer: (args) => {
const [date, format] = args as [Date, string];
return `${date.getTime()}-${format}`;
},
}
);
console.log('\nCustom serializer examples:');
const date = new Date('2024-01-01');
console.log(customMemoize(date, 'ISO'));
console.log(customMemoize(date, 'ISO')); // Cached
// With max cache size (LRU-like behavior for primitives)
const limitedCache = memoize(
(x: number) => {
console.log(`Computing for ${x}`);
return x * 2;
},
{ maxSize: 3 }
);
console.log('\nLimited cache examples:');
console.log(limitedCache(1));
console.log(limitedCache(2));
console.log(limitedCache(3));
console.log(limitedCache(4)); // Evicts cache for 1
console.log(limitedCache(1)); // Recomputed
// Cache management
console.log('\nCache size:', limitedCache.cache.size());
limitedCache.cache.clear();
console.log('Cache cleared, size:', limitedCache.cache.size());This memoization implementation solves a fundamental challenge in JavaScript caching: handling both primitive and object arguments without causing memory leaks. The key insight is that WeakMap allows garbage collection of cached entries when object keys are no longer referenced elsewhere, while Map is needed for primitives since WeakMap only accepts objects as keys.
The cache structure uses a tree-like approach where each argument in the function call traverses deeper into nested cache nodes. For primitive arguments, we serialize them to strings and store the next cache node in a Map. For object arguments, we use the object reference directly as a WeakMap key. This means calling fn(user, 'read') creates a path: rootCache → objectCache[user] → primitiveCache['read'] → result. The final node stores the computed result wrapped in an object to distinguish between cached undefined values and cache misses.
The custom serializer option addresses cases where the default primitive serialization isn't sufficient. For example, two different Date objects with the same timestamp would be treated as different cache keys by default (since they're different object references). By providing a serializer that converts dates to timestamps, you can achieve semantic equality. The serializer is only used for primitive values; objects always use reference equality via WeakMap.
The optional maxSize parameter implements basic cache eviction for primitive keys using a FIFO strategy. This prevents unbounded memory growth when memoizing functions called with many different primitive values. Note that object-based cache entries are automatically managed by the garbage collector through WeakMap, so maxSize only affects primitive entries. The size() method only reports primitive cache entries since WeakMap doesn't expose its size.
Use this pattern for expensive computations, API response caching, or recursive algorithms like dynamic programming solutions. Avoid it for functions with side effects, functions where you need time-based expiration, or when dealing with very large numbers of unique argument combinations that might exceed memory limits. The WeakMap approach is particularly valuable in long-running applications where temporary objects (like request contexts or DOM elements) are frequently passed to memoized functions.