Improves cache utilities (microsoft#210471)

Author: hediet

src/vs/base/common/cache.ts

@@ -36,47 +36,85 @@ export class Cache<T> {
 	}
 }
 
+export function identity<T>(t: T): T {
+	return t;
+}
+
+interface ICacheOptions<TArg> {
+	/**
+	 * The cache key is used to identify the cache entry.
+	 * Strict equality is used to compare cache keys.
+	*/
+	getCacheKey: (arg: TArg) => unknown;
+}
+
 /**
  * Uses a LRU cache to make a given parametrized function cached.
- * Caches just the last value.
+ * Caches just the last key/value.
 */
 export class LRUCachedFunction<TArg, TComputed> {
 	private lastCache: TComputed | undefined = undefined;
 	private lastArgKey: unknown | undefined = undefined;
 
-	constructor(
-		private readonly fn: (arg: TArg) => TComputed,
-		private readonly _computeKey: (arg: TArg) => unknown = JSON.stringify,
-	) {
+	private readonly _fn: (arg: TArg) => TComputed;
+	private readonly _computeKey: (arg: TArg) => unknown;
+
+	constructor(fn: (arg: TArg) => TComputed);
+	constructor(options: ICacheOptions<TArg>, fn: (arg: TArg) => TComputed);
+	constructor(arg1: ICacheOptions<TArg> | ((arg: TArg) => TComputed), arg2?: (arg: TArg) => TComputed) {
+		if (typeof arg1 === 'function') {
+			this._fn = arg1;
+			this._computeKey = identity;
+		} else {
+			this._fn = arg2!;
+			this._computeKey = arg1.getCacheKey;
+		}
 	}
 
 	public get(arg: TArg): TComputed {
 		const key = this._computeKey(arg);
 		if (this.lastArgKey !== key) {
 			this.lastArgKey = key;
-			this.lastCache = this.fn(arg);
+			this.lastCache = this._fn(arg);
 		}
 		return this.lastCache!;
 	}
 }
 
 /**
- * Uses an unbounded cache (referential equality) to memoize the results of the given function.
+ * Uses an unbounded cache to memoize the results of the given function.
 */
-export class CachedFunction<TArg, TValue> {
-	private readonly _map = new Map<TArg, TValue>();
-	public get cachedValues(): ReadonlyMap<TArg, TValue> {
+export class CachedFunction<TArg, TComputed> {
+	private readonly _map = new Map<TArg, TComputed>();
+	private readonly _map2 = new Map<unknown, TComputed>();
+	public get cachedValues(): ReadonlyMap<TArg, TComputed> {
 		return this._map;
 	}
 
-	constructor(private readonly fn: (arg: TArg) => TValue) { }
+	private readonly _fn: (arg: TArg) => TComputed;
+	private readonly _computeKey: (arg: TArg) => unknown;
 
-	public get(arg: TArg): TValue {
-		if (this._map.has(arg)) {
-			return this._map.get(arg)!;
+	constructor(fn: (arg: TArg) => TComputed);
+	constructor(options: ICacheOptions<TArg>, fn: (arg: TArg) => TComputed);
+	constructor(arg1: ICacheOptions<TArg> | ((arg: TArg) => TComputed), arg2?: (arg: TArg) => TComputed) {
+		if (typeof arg1 === 'function') {
+			this._fn = arg1;
+			this._computeKey = identity;
+		} else {
+			this._fn = arg2!;
+			this._computeKey = arg1.getCacheKey;
 		}
-		const value = this.fn(arg);
+	}
+
+	public get(arg: TArg): TComputed {
+		const key = this._computeKey(arg);
+		if (this._map2.has(key)) {
+			return this._map2.get(key)!;
+		}
+
+		const value = this._fn(arg);
 		this._map.set(arg, value);
+		this._map2.set(key, value);
 		return value;
 	}
 }

src/vs/base/common/strings.ts

@@ -1149,7 +1149,7 @@ export class AmbiguousCharacters {
 	private static readonly cache = new LRUCachedFunction<
 		string[],
 		AmbiguousCharacters
-	>((locales) => {
+	>({ getCacheKey: JSON.stringify }, (locales) => {
 		function arrayToMap(arr: number[]): Map<number, number> {
 			const result = new Map<number, number>();
 			for (let i = 0; i < arr.length; i += 2) {

src/vs/workbench/contrib/bulkEdit/browser/preview/bulkEditPane.ts

@@ -410,7 +410,7 @@ export class BulkEditPane extends ViewPane {
 			resources,
 			getResourceDiffEditorInputIdOfOperation
 		};
-	}, key => key);
+	});
 
 	private _onContextMenu(e: ITreeContextMenuEvent<any>): void {