10.1 KB344 lines
Blame
1/**
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 *
4 * This source code is licensed under the MIT license found in the
5 * LICENSE file in the root directory of this source tree.
6 */
7
8import type {ValueObject} from 'immutable';
9
10import deepEqual from 'fast-deep-equal';
11import {is, isValueObject, List} from 'immutable';
12import {isPromise} from './utils';
13
14type LRUKey = LRUHashKey | ValueObject;
15type LRUHashKey = string | number | boolean | null | undefined | object;
16
17/**
18 * Simple least-recently-used cache which holds at most `max` entries.
19 *
20 * The map uses `Immutable.is` [1] instead of `Object.is` to compare keys.
21 * So it can support complex key types (ex. a tuple of immutable
22 * collections and other types). Right now this is done by using data
23 * structures from `immutable`. If the native `Map` supports customizing
24 * the compare function we can replace `Object.is` with `Immutable.is`
25 * and use native `Map` for performance.
26 *
27 * [1]: https://immutable-js.com/docs/v4.3.0/is()/.
28 */
29export class LRU<K extends LRUKey, V> {
30 // Implementation is based on Map having stable insertion order and O(1).
31 // To support immutable objects, the cache map uses hashCode of "key" as
32 // the first key, then the actual key in the nested map.
33 private cache = new Map<LRUHashKey, Map<K, V>>();
34
35 constructor(
36 private maxItems: number,
37 private maxHashCollision = 3,
38 ) {}
39
40 get(key: K): V | undefined {
41 let result = undefined;
42 const hashKey = getHashKey(key);
43 const valueMap = this.cache.get(hashKey);
44 if (valueMap !== undefined) {
45 // Fast path: by object reference.
46 const maybeValue = valueMap.get(key);
47 if (maybeValue !== undefined) {
48 result = maybeValue;
49 } else {
50 // Slower path: immutable.is
51 for (const [k, v] of valueMap) {
52 // The order matters. `is(key, k)` updates `key` (user-provided) to
53 // the `k` (cache) reference. See `immutableExt.withSelfUpdateEquals`.
54 if (is(key, k)) {
55 result = v;
56 break;
57 }
58 }
59 }
60 this.cache.delete(hashKey);
61 this.cache.set(hashKey, valueMap);
62 }
63 return result;
64 }
65
66 set(key: K, value: V) {
67 const hashKey = getHashKey(key);
68 let valueMap = this.cache.get(hashKey);
69 if (valueMap === undefined || valueMap.size >= this.maxHashCollision) {
70 valueMap = new Map([[key, value]]);
71 } else {
72 valueMap.set(key, value);
73 }
74 // ensure refresh by deleting before setting
75 this.cache.delete(hashKey);
76
77 if (value !== undefined) {
78 this.cache.set(hashKey, valueMap);
79 if (this.cache.size > this.maxItems) {
80 // evict oldest
81 // iteration order guarantees oldest item is first
82 const next = this.cache.keys().next();
83 // undefined is a valid value, so use iterator `done` to know whether to delete or not.
84 if (!next.done) {
85 this.cache.delete(next.value);
86 }
87 }
88 }
89 }
90
91 delete(key: K) {
92 const hashKey = getHashKey(key);
93 this.cache.delete(hashKey);
94 }
95
96 clear() {
97 this.cache.clear();
98 }
99}
100
101function getHashKey<K extends LRUKey>(key: K): LRUHashKey {
102 // @ts-expect-error (string)?.hashCode is valid JavaScript.
103 const hashCodeFunc = key?.hashCode;
104 if (hashCodeFunc !== undefined) {
105 return hashCodeFunc.apply(key);
106 }
107 return key;
108}
109
110// Neither `unknown` nor `never` works here.
111// eslint-disable-next-line @typescript-eslint/no-explicit-any
112type AnyFunction<T> = (this: T, ...args: any[]) => any;
113
114type CacheStats = {
115 hit?: number;
116 miss?: number;
117 skip?: number;
118};
119
120export interface LRUWithStats extends LRU<LRUKey, unknown> {
121 stats?: CacheStats;
122}
123
124export interface WithCache {
125 cache: LRUWithStats;
126}
127
128/** Cache options. */
129type CacheOpts<This> = {
130 /**
131 * If set, use the specified cache.
132 *
133 * Callsite can assign `cache.stats = {hit: 0, miss: 0, skip: 0}`
134 * to collect statistics.
135 */
136 cache?: LRUWithStats;
137
138 /**
139 * If set, and cache is not set, create cache of the given size.
140 * Default value: 200.
141 */
142 cacheSize?: number;
143
144 /** If set, use the returned values as extra cache keys. */
145 getExtraKeys?: (this: This) => unknown[];
146
147 /**
148 * Cached functions are expected to be "pure" - give same output
149 * for the same inputs. If set to true, compare the cached value
150 * with a fresh recalculation and throws on mismatch.
151 */
152 audit?: boolean;
153
154 /**
155 * Track the `cache` so it can be cleared by `clearTrackedCache`.
156 * Default: true.
157 */
158 track?: boolean;
159};
160
161type DecoratorFunc = (target: unknown, propertyKey: string, descriptor: PropertyDescriptor) => void;
162
163/**
164 * Decorator to make a class method cached.
165 *
166 * This is similar to calling `cached` on the function, with
167 * an auto-generated `opts.getExtraKeys` function that turns
168 * `this` into extra cache keys. Immutable `this` is used
169 * as the extra cache key directly. Otherwise, cacheable
170 * properties of `this` are used as extra cache keys.
171 */
172export function cached<T>(opts?: CacheOpts<T>): DecoratorFunc;
173
174/**
175 * Wraps the given function with a LRU cache.
176 * Returns the wrapped function.
177 *
178 * If the function depends on extra inputs outside the
179 * parameters, use `opts.getExtraKeys` to provide them.
180 *
181 * The cache can be accessed via `returnedFunction.cache`.
182 *
183 * Cache is used only when all parameters are cacheable [1].
184 * For example, if a parameter is a function or `null`,
185 * then cache is only used when that parameter is `null`,
186 * since functions are not cacheable.
187 *
188 * [1]: See `isCachable` for cacheable types.
189 */
190export function cached<T, F extends AnyFunction<T>>(func: F, opts?: CacheOpts<T>): F & WithCache;
191
192// union of the above
193export function cached<T, F extends AnyFunction<T>>(
194 arg1?: F | CacheOpts<T>,
195 arg2?: CacheOpts<T>,
196): (F & WithCache) | DecoratorFunc {
197 if (typeof arg1 === 'function') {
198 // cached(func)
199 return cachedFunction(arg1, arg2);
200 } else {
201 // @cached(opts)
202 return cacheDecorator(arg1);
203 }
204}
205
206const trackedCaches = new Set<WeakRef<LRUWithStats>>();
207
208/** Clear tracked LRU caches. By default, `@cached` */
209export function clearTrackedCache() {
210 for (const weakRef of trackedCaches) {
211 const cache = weakRef.deref();
212 if (cache === undefined) {
213 trackedCaches.delete(weakRef);
214 } else {
215 cache.clear();
216 }
217 }
218}
219
220function cachedFunction<T, F extends AnyFunction<T>>(func: F, opts?: CacheOpts<T>): F & WithCache {
221 const cache: LRUWithStats = opts?.cache ?? new LRU(opts?.cacheSize ?? 200);
222 const audit = opts?.audit ?? false;
223 const getExtraKeys = opts?.getExtraKeys;
224 const track = opts?.track ?? true;
225 if (track) {
226 trackedCaches.add(new WeakRef(cache));
227 }
228 const cachedFunc = function (this: T, ...args: Parameters<F>): ReturnType<F> {
229 const stats = cache.stats;
230 if (!args.every(isCachable)) {
231 if (stats != null) {
232 stats.skip = (stats.skip ?? 0) + 1;
233 }
234 return func.apply(this, args) as ReturnType<F>;
235 }
236 const cacheKey = List(getExtraKeys ? [...getExtraKeys.apply(this), ...args] : args);
237 const cachedValue = cache.get(cacheKey);
238 if (cachedValue !== undefined) {
239 if (stats != null) {
240 stats.hit = (stats.hit ?? 0) + 1;
241 }
242 if (audit) {
243 const result = func.apply(this, args) as ReturnType<F>;
244 const equal = isValueObject(result)
245 ? is(result, cachedValue)
246 : deepEqual(result, cachedValue);
247 if (!equal) {
248 const argsStr = args.map(a => a.toString()).join(', ');
249 throw new Error(
250 `cached value mismatch on ${func.name}(${argsStr}): cached ${cachedValue}, actual ${result}`,
251 );
252 }
253 }
254 return cachedValue as ReturnType<F>;
255 }
256 if (stats != null) {
257 stats.miss = (stats.miss ?? 0) + 1;
258 }
259 const result = func.apply(this, args) as ReturnType<F>;
260 if (isPromise(result)) {
261 return result.then((value: unknown) => {
262 cache.set(cacheKey, value);
263 return value;
264 });
265 }
266 cache.set(cacheKey, result);
267 return result;
268 };
269 cachedFunc.cache = cache;
270 return cachedFunc as WithCache & F;
271}
272
273// See https://www.typescriptlang.org/docs/handbook/decorators.html.
274function cacheDecorator<T>(opts?: CacheOpts<T>) {
275 const getExtraKeys =
276 opts?.getExtraKeys ??
277 function (this: T): unknown[] {
278 // Use `this` as extra key if it's a value object (hash + eq).
279 if (isValueObject(this)) {
280 return [this];
281 }
282 // Scan through cacheable properties.
283 if (this != null && typeof this === 'object') {
284 return Object.values(this).filter(isCachable);
285 }
286 // Give up - do not add extra cache keys.
287 return [];
288 };
289 return function (_target: unknown, _propertyKey: string, descriptor: PropertyDescriptor) {
290 const originalFunc = descriptor.value;
291 descriptor.value = cachedFunction(originalFunc, {...opts, getExtraKeys});
292 };
293}
294
295/** Cache an "instance" function like `this.foo`. */
296export function cachedMethod<T, F extends AnyFunction<T>>(originalFunc: F, opts?: CacheOpts<T>): F {
297 const getExtraKeys =
298 opts?.getExtraKeys ??
299 function (this: T): unknown[] {
300 // Use `this` as extra key if it's a value object (hash + eq).
301 if (isValueObject(this)) {
302 return [this];
303 }
304 // Scan through cacheable properties.
305 if (this != null && typeof this === 'object') {
306 return Object.values(this).filter(isCachable);
307 }
308 // Give up - do not add extra cache keys.
309 return [];
310 };
311 return cachedFunction(originalFunc, {...opts, getExtraKeys});
312}
313
314const cachableTypeNames = new Set([
315 'number',
316 'string',
317 'boolean',
318 'symbol',
319 'bigint',
320 'undefined',
321 'null',
322]);
323
324/**
325 * Returns true if `arg` can be used as cache keys.
326 * Primitive types (string, number, boolean, null, undefined)
327 * can be used as cache keys.
328 * Objects can be used as cache keys if they are immutable.
329 */
330function isCachable(arg: unknown): boolean {
331 // null is a special case, since typeof(null) returns 'object'.
332 if (arg == null) {
333 return true;
334 }
335 const typeName = typeof arg;
336 if (cachableTypeNames.has(typeName)) {
337 return true;
338 }
339 if (typeName === 'object' && isValueObject(arg)) {
340 return true;
341 }
342 return false;
343}
344