addons/shared/LRU.tsblame
View source
b69ab311/**
b69ab312 * Copyright (c) Meta Platforms, Inc. and affiliates.
b69ab313 *
b69ab314 * This source code is licensed under the MIT license found in the
b69ab315 * LICENSE file in the root directory of this source tree.
b69ab316 */
b69ab317
b69ab318import type {ValueObject} from 'immutable';
b69ab319
b69ab3110import deepEqual from 'fast-deep-equal';
b69ab3111import {is, isValueObject, List} from 'immutable';
b69ab3112import {isPromise} from './utils';
b69ab3113
b69ab3114type LRUKey = LRUHashKey | ValueObject;
b69ab3115type LRUHashKey = string | number | boolean | null | undefined | object;
b69ab3116
b69ab3117/**
b69ab3118 * Simple least-recently-used cache which holds at most `max` entries.
b69ab3119 *
b69ab3120 * The map uses `Immutable.is` [1] instead of `Object.is` to compare keys.
b69ab3121 * So it can support complex key types (ex. a tuple of immutable
b69ab3122 * collections and other types). Right now this is done by using data
b69ab3123 * structures from `immutable`. If the native `Map` supports customizing
b69ab3124 * the compare function we can replace `Object.is` with `Immutable.is`
b69ab3125 * and use native `Map` for performance.
b69ab3126 *
b69ab3127 * [1]: https://immutable-js.com/docs/v4.3.0/is()/.
b69ab3128 */
b69ab3129export class LRU<K extends LRUKey, V> {
b69ab3130 // Implementation is based on Map having stable insertion order and O(1).
b69ab3131 // To support immutable objects, the cache map uses hashCode of "key" as
b69ab3132 // the first key, then the actual key in the nested map.
b69ab3133 private cache = new Map<LRUHashKey, Map<K, V>>();
b69ab3134
b69ab3135 constructor(
b69ab3136 private maxItems: number,
b69ab3137 private maxHashCollision = 3,
b69ab3138 ) {}
b69ab3139
b69ab3140 get(key: K): V | undefined {
b69ab3141 let result = undefined;
b69ab3142 const hashKey = getHashKey(key);
b69ab3143 const valueMap = this.cache.get(hashKey);
b69ab3144 if (valueMap !== undefined) {
b69ab3145 // Fast path: by object reference.
b69ab3146 const maybeValue = valueMap.get(key);
b69ab3147 if (maybeValue !== undefined) {
b69ab3148 result = maybeValue;
b69ab3149 } else {
b69ab3150 // Slower path: immutable.is
b69ab3151 for (const [k, v] of valueMap) {
b69ab3152 // The order matters. `is(key, k)` updates `key` (user-provided) to
b69ab3153 // the `k` (cache) reference. See `immutableExt.withSelfUpdateEquals`.
b69ab3154 if (is(key, k)) {
b69ab3155 result = v;
b69ab3156 break;
b69ab3157 }
b69ab3158 }
b69ab3159 }
b69ab3160 this.cache.delete(hashKey);
b69ab3161 this.cache.set(hashKey, valueMap);
b69ab3162 }
b69ab3163 return result;
b69ab3164 }
b69ab3165
b69ab3166 set(key: K, value: V) {
b69ab3167 const hashKey = getHashKey(key);
b69ab3168 let valueMap = this.cache.get(hashKey);
b69ab3169 if (valueMap === undefined || valueMap.size >= this.maxHashCollision) {
b69ab3170 valueMap = new Map([[key, value]]);
b69ab3171 } else {
b69ab3172 valueMap.set(key, value);
b69ab3173 }
b69ab3174 // ensure refresh by deleting before setting
b69ab3175 this.cache.delete(hashKey);
b69ab3176
b69ab3177 if (value !== undefined) {
b69ab3178 this.cache.set(hashKey, valueMap);
b69ab3179 if (this.cache.size > this.maxItems) {
b69ab3180 // evict oldest
b69ab3181 // iteration order guarantees oldest item is first
b69ab3182 const next = this.cache.keys().next();
b69ab3183 // undefined is a valid value, so use iterator `done` to know whether to delete or not.
b69ab3184 if (!next.done) {
b69ab3185 this.cache.delete(next.value);
b69ab3186 }
b69ab3187 }
b69ab3188 }
b69ab3189 }
b69ab3190
b69ab3191 delete(key: K) {
b69ab3192 const hashKey = getHashKey(key);
b69ab3193 this.cache.delete(hashKey);
b69ab3194 }
b69ab3195
b69ab3196 clear() {
b69ab3197 this.cache.clear();
b69ab3198 }
b69ab3199}
b69ab31100
b69ab31101function getHashKey<K extends LRUKey>(key: K): LRUHashKey {
b69ab31102 // @ts-expect-error (string)?.hashCode is valid JavaScript.
b69ab31103 const hashCodeFunc = key?.hashCode;
b69ab31104 if (hashCodeFunc !== undefined) {
b69ab31105 return hashCodeFunc.apply(key);
b69ab31106 }
b69ab31107 return key;
b69ab31108}
b69ab31109
b69ab31110// Neither `unknown` nor `never` works here.
b69ab31111// eslint-disable-next-line @typescript-eslint/no-explicit-any
b69ab31112type AnyFunction<T> = (this: T, ...args: any[]) => any;
b69ab31113
b69ab31114type CacheStats = {
b69ab31115 hit?: number;
b69ab31116 miss?: number;
b69ab31117 skip?: number;
b69ab31118};
b69ab31119
b69ab31120export interface LRUWithStats extends LRU<LRUKey, unknown> {
b69ab31121 stats?: CacheStats;
b69ab31122}
b69ab31123
b69ab31124export interface WithCache {
b69ab31125 cache: LRUWithStats;
b69ab31126}
b69ab31127
b69ab31128/** Cache options. */
b69ab31129type CacheOpts<This> = {
b69ab31130 /**
b69ab31131 * If set, use the specified cache.
b69ab31132 *
b69ab31133 * Callsite can assign `cache.stats = {hit: 0, miss: 0, skip: 0}`
b69ab31134 * to collect statistics.
b69ab31135 */
b69ab31136 cache?: LRUWithStats;
b69ab31137
b69ab31138 /**
b69ab31139 * If set, and cache is not set, create cache of the given size.
b69ab31140 * Default value: 200.
b69ab31141 */
b69ab31142 cacheSize?: number;
b69ab31143
b69ab31144 /** If set, use the returned values as extra cache keys. */
b69ab31145 getExtraKeys?: (this: This) => unknown[];
b69ab31146
b69ab31147 /**
b69ab31148 * Cached functions are expected to be "pure" - give same output
b69ab31149 * for the same inputs. If set to true, compare the cached value
b69ab31150 * with a fresh recalculation and throws on mismatch.
b69ab31151 */
b69ab31152 audit?: boolean;
b69ab31153
b69ab31154 /**
b69ab31155 * Track the `cache` so it can be cleared by `clearTrackedCache`.
b69ab31156 * Default: true.
b69ab31157 */
b69ab31158 track?: boolean;
b69ab31159};
b69ab31160
b69ab31161type DecoratorFunc = (target: unknown, propertyKey: string, descriptor: PropertyDescriptor) => void;
b69ab31162
b69ab31163/**
b69ab31164 * Decorator to make a class method cached.
b69ab31165 *
b69ab31166 * This is similar to calling `cached` on the function, with
b69ab31167 * an auto-generated `opts.getExtraKeys` function that turns
b69ab31168 * `this` into extra cache keys. Immutable `this` is used
b69ab31169 * as the extra cache key directly. Otherwise, cacheable
b69ab31170 * properties of `this` are used as extra cache keys.
b69ab31171 */
b69ab31172export function cached<T>(opts?: CacheOpts<T>): DecoratorFunc;
b69ab31173
b69ab31174/**
b69ab31175 * Wraps the given function with a LRU cache.
b69ab31176 * Returns the wrapped function.
b69ab31177 *
b69ab31178 * If the function depends on extra inputs outside the
b69ab31179 * parameters, use `opts.getExtraKeys` to provide them.
b69ab31180 *
b69ab31181 * The cache can be accessed via `returnedFunction.cache`.
b69ab31182 *
b69ab31183 * Cache is used only when all parameters are cacheable [1].
b69ab31184 * For example, if a parameter is a function or `null`,
b69ab31185 * then cache is only used when that parameter is `null`,
b69ab31186 * since functions are not cacheable.
b69ab31187 *
b69ab31188 * [1]: See `isCachable` for cacheable types.
b69ab31189 */
b69ab31190export function cached<T, F extends AnyFunction<T>>(func: F, opts?: CacheOpts<T>): F & WithCache;
b69ab31191
b69ab31192// union of the above
b69ab31193export function cached<T, F extends AnyFunction<T>>(
b69ab31194 arg1?: F | CacheOpts<T>,
b69ab31195 arg2?: CacheOpts<T>,
b69ab31196): (F & WithCache) | DecoratorFunc {
b69ab31197 if (typeof arg1 === 'function') {
b69ab31198 // cached(func)
b69ab31199 return cachedFunction(arg1, arg2);
b69ab31200 } else {
b69ab31201 // @cached(opts)
b69ab31202 return cacheDecorator(arg1);
b69ab31203 }
b69ab31204}
b69ab31205
b69ab31206const trackedCaches = new Set<WeakRef<LRUWithStats>>();
b69ab31207
b69ab31208/** Clear tracked LRU caches. By default, `@cached` */
b69ab31209export function clearTrackedCache() {
b69ab31210 for (const weakRef of trackedCaches) {
b69ab31211 const cache = weakRef.deref();
b69ab31212 if (cache === undefined) {
b69ab31213 trackedCaches.delete(weakRef);
b69ab31214 } else {
b69ab31215 cache.clear();
b69ab31216 }
b69ab31217 }
b69ab31218}
b69ab31219
b69ab31220function cachedFunction<T, F extends AnyFunction<T>>(func: F, opts?: CacheOpts<T>): F & WithCache {
b69ab31221 const cache: LRUWithStats = opts?.cache ?? new LRU(opts?.cacheSize ?? 200);
b69ab31222 const audit = opts?.audit ?? false;
b69ab31223 const getExtraKeys = opts?.getExtraKeys;
b69ab31224 const track = opts?.track ?? true;
b69ab31225 if (track) {
b69ab31226 trackedCaches.add(new WeakRef(cache));
b69ab31227 }
b69ab31228 const cachedFunc = function (this: T, ...args: Parameters<F>): ReturnType<F> {
b69ab31229 const stats = cache.stats;
b69ab31230 if (!args.every(isCachable)) {
b69ab31231 if (stats != null) {
b69ab31232 stats.skip = (stats.skip ?? 0) + 1;
b69ab31233 }
b69ab31234 return func.apply(this, args) as ReturnType<F>;
b69ab31235 }
b69ab31236 const cacheKey = List(getExtraKeys ? [...getExtraKeys.apply(this), ...args] : args);
b69ab31237 const cachedValue = cache.get(cacheKey);
b69ab31238 if (cachedValue !== undefined) {
b69ab31239 if (stats != null) {
b69ab31240 stats.hit = (stats.hit ?? 0) + 1;
b69ab31241 }
b69ab31242 if (audit) {
b69ab31243 const result = func.apply(this, args) as ReturnType<F>;
b69ab31244 const equal = isValueObject(result)
b69ab31245 ? is(result, cachedValue)
b69ab31246 : deepEqual(result, cachedValue);
b69ab31247 if (!equal) {
b69ab31248 const argsStr = args.map(a => a.toString()).join(', ');
b69ab31249 throw new Error(
b69ab31250 `cached value mismatch on ${func.name}(${argsStr}): cached ${cachedValue}, actual ${result}`,
b69ab31251 );
b69ab31252 }
b69ab31253 }
b69ab31254 return cachedValue as ReturnType<F>;
b69ab31255 }
b69ab31256 if (stats != null) {
b69ab31257 stats.miss = (stats.miss ?? 0) + 1;
b69ab31258 }
b69ab31259 const result = func.apply(this, args) as ReturnType<F>;
b69ab31260 if (isPromise(result)) {
b69ab31261 return result.then((value: unknown) => {
b69ab31262 cache.set(cacheKey, value);
b69ab31263 return value;
b69ab31264 });
b69ab31265 }
b69ab31266 cache.set(cacheKey, result);
b69ab31267 return result;
b69ab31268 };
b69ab31269 cachedFunc.cache = cache;
b69ab31270 return cachedFunc as WithCache & F;
b69ab31271}
b69ab31272
b69ab31273// See https://www.typescriptlang.org/docs/handbook/decorators.html.
b69ab31274function cacheDecorator<T>(opts?: CacheOpts<T>) {
b69ab31275 const getExtraKeys =
b69ab31276 opts?.getExtraKeys ??
b69ab31277 function (this: T): unknown[] {
b69ab31278 // Use `this` as extra key if it's a value object (hash + eq).
b69ab31279 if (isValueObject(this)) {
b69ab31280 return [this];
b69ab31281 }
b69ab31282 // Scan through cacheable properties.
b69ab31283 if (this != null && typeof this === 'object') {
b69ab31284 return Object.values(this).filter(isCachable);
b69ab31285 }
b69ab31286 // Give up - do not add extra cache keys.
b69ab31287 return [];
b69ab31288 };
b69ab31289 return function (_target: unknown, _propertyKey: string, descriptor: PropertyDescriptor) {
b69ab31290 const originalFunc = descriptor.value;
b69ab31291 descriptor.value = cachedFunction(originalFunc, {...opts, getExtraKeys});
b69ab31292 };
b69ab31293}
b69ab31294
b69ab31295/** Cache an "instance" function like `this.foo`. */
b69ab31296export function cachedMethod<T, F extends AnyFunction<T>>(originalFunc: F, opts?: CacheOpts<T>): F {
b69ab31297 const getExtraKeys =
b69ab31298 opts?.getExtraKeys ??
b69ab31299 function (this: T): unknown[] {
b69ab31300 // Use `this` as extra key if it's a value object (hash + eq).
b69ab31301 if (isValueObject(this)) {
b69ab31302 return [this];
b69ab31303 }
b69ab31304 // Scan through cacheable properties.
b69ab31305 if (this != null && typeof this === 'object') {
b69ab31306 return Object.values(this).filter(isCachable);
b69ab31307 }
b69ab31308 // Give up - do not add extra cache keys.
b69ab31309 return [];
b69ab31310 };
b69ab31311 return cachedFunction(originalFunc, {...opts, getExtraKeys});
b69ab31312}
b69ab31313
b69ab31314const cachableTypeNames = new Set([
b69ab31315 'number',
b69ab31316 'string',
b69ab31317 'boolean',
b69ab31318 'symbol',
b69ab31319 'bigint',
b69ab31320 'undefined',
b69ab31321 'null',
b69ab31322]);
b69ab31323
b69ab31324/**
b69ab31325 * Returns true if `arg` can be used as cache keys.
b69ab31326 * Primitive types (string, number, boolean, null, undefined)
b69ab31327 * can be used as cache keys.
b69ab31328 * Objects can be used as cache keys if they are immutable.
b69ab31329 */
b69ab31330function isCachable(arg: unknown): boolean {
b69ab31331 // null is a special case, since typeof(null) returns 'object'.
b69ab31332 if (arg == null) {
b69ab31333 return true;
b69ab31334 }
b69ab31335 const typeName = typeof arg;
b69ab31336 if (cachableTypeNames.has(typeName)) {
b69ab31337 return true;
b69ab31338 }
b69ab31339 if (typeName === 'object' && isValueObject(arg)) {
b69ab31340 return true;
b69ab31341 }
b69ab31342 return false;
b69ab31343}