From 45bdd6c76c7029c7b06a9715f4030f322d319437 Mon Sep 17 00:00:00 2001
From: vstefanovic97 <stefanoviv@ae.com>
Date: Wed, 12 Mar 2025 13:16:59 +0100
Subject: [PATCH] Fix memory leak

---
 src/memoize-decorator.ts             | 56 +++++++++++++++++-----------
 test/specs/memoize-decorator.spec.ts |  8 ++--
 2 files changed, 40 insertions(+), 24 deletions(-)

diff --git a/src/memoize-decorator.ts b/src/memoize-decorator.ts
index f2f060e..6ca7702 100644
--- a/src/memoize-decorator.ts
+++ b/src/memoize-decorator.ts
@@ -4,6 +4,11 @@ interface MemoizeArgs {
 	tags?: string[];
 }
 
+interface CacheValue {
+	resultsMap: Map<any,any>;
+	tagVersions?: Record<string, symbol>;
+}
+
 export function Memoize(args?: MemoizeArgs | MemoizeArgs['hashFunction']) {
 	let hashFunction: MemoizeArgs['hashFunction'];
 	let duration: MemoizeArgs['expiring'];
@@ -35,26 +40,29 @@ export function MemoizeExpiring(expiring: number, hashFunction?: MemoizeArgs['ha
 	});
 }
 
-const clearCacheTagsMap: Map<string, Map<any, any>[]> = new Map();
+const latestTagVersions: Map<string, symbol> = new Map();
 
-export function clear (tags: string[]): number {
-	const cleared: Set<Map<any, any>> = new Set();
+export function clear (tags: string[]): void {
 	for (const tag of tags) {
-		const maps = clearCacheTagsMap.get(tag);
-		if (maps) {
-			for (const mp of maps) {
-				if (!cleared.has(mp)) {
-					mp.clear();
-					cleared.add(mp);
-				}
-			}
+		if (latestTagVersions.has(tag)) {
+			latestTagVersions.set(tag, Symbol());
 		}
 	}
-	return cleared.size;
+}
+
+function getLatestTagVersionsForTags (tags: string[]) {
+	return tags.reduce<Record<string, symbol>>((acc, tag) => {
+		if (!latestTagVersions.has(tag)) {
+			const symbolForTag = Symbol();
+			latestTagVersions.set(tag, symbolForTag);
+			return Object.assign(acc, { [tag]: symbolForTag });
+		}
+		return Object.assign(acc, { [tag]: latestTagVersions.get(tag) });
+	}, {});
 }
 
 function getNewFunction(originalMethod: () => void, hashFunction?: MemoizeArgs['hashFunction'], duration: number = 0, tags?: MemoizeArgs['tags']) {
-	const propMapName = Symbol(`__memoized_map__`);
+	const propMapName = Symbol(`__cache__`);
 
 	// The function returned here gets called instead of originalMethod.
 	return function (...args: any[]) {
@@ -62,23 +70,29 @@ function getNewFunction(originalMethod: () => void, hashFunction?: MemoizeArgs['
 
 		// Get or create map
 		if (!this.hasOwnProperty(propMapName)) {
+			const value: CacheValue = { resultsMap: new Map<any, any>() };
+			if (Array.isArray(tags)) {
+				value.tagVersions = getLatestTagVersionsForTags(tags);
+			}
 			Object.defineProperty(this, propMapName, {
 				configurable: false,
 				enumerable: false,
 				writable: false,
-				value: new Map<any, any>()
+				value,
 			});
 		}
-		let myMap: Map<any, any> = this[propMapName];
+
+		const cache = this[propMapName] as CacheValue;
+		let myMap: Map<any, any> = cache.resultsMap;
 
 		if (Array.isArray(tags)) {
-			for (const tag of tags) {
-				if (clearCacheTagsMap.has(tag)) {
-					clearCacheTagsMap.get(tag).push(myMap);
-				} else {
-					clearCacheTagsMap.set(tag, [myMap]);
-				}
+			const tagVersions = this[propMapName].tagVersions;
+			const isAtLeastOneTagStale = tags.some((tag) => tagVersions[tag] !== latestTagVersions.get(tag));
+			if (isAtLeastOneTagStale) {
+				myMap.clear();
+				cache.tagVersions = getLatestTagVersionsForTags(tags);
 			}
+
 		}
 
 		if (hashFunction || args.length > 0 || duration > 0) {
diff --git a/test/specs/memoize-decorator.spec.ts b/test/specs/memoize-decorator.spec.ts
index f710c9f..ba8924a 100644
--- a/test/specs/memoize-decorator.spec.ts
+++ b/test/specs/memoize-decorator.spec.ts
@@ -252,9 +252,10 @@ describe('Memoize()', () => {
 			let val3 = a.getGreeting4('Hello', 'World');
 			clear(["foo"]);
 			let val4 = a.getGreeting4('Hello', 'Moon');
-			let val5 = a.getGreeting4('Hello', 'World');
-			clear(["bar"]);
+			let val5 = a.getGreeting4('Hello', 'Moon');
 			let val6 = a.getGreeting4('Hello', 'World');
+			clear(["bar"]);
+			let val7 = a.getGreeting4('Hello', 'World');
 
 			clear(["unknown"]);
 
@@ -262,8 +263,9 @@ describe('Memoize()', () => {
 			expect(val2).toEqual('Hello, Moon');
 			expect(val3).toEqual('Hello, World');
 			expect(val4).toEqual('Hello, Moon');
-			expect(val5).toEqual('Hello, World');
+			expect(val5).toEqual('Hello, Moon');
 			expect(val6).toEqual('Hello, World');
+			expect(val7).toEqual('Hello, World');
 
 			expect(getGreetingSpy).toHaveBeenCalledTimes(5);
 		});