From 24ef26c42e7f016a7bec813b07047a2f8df36fb7 Mon Sep 17 00:00:00 2001
From: Marco Salazar <marco@elementl.com>
Date: Thu, 19 Dec 2024 13:47:16 -0500
Subject: [PATCH 01/15] weak map memoize

---
 .../src/util/__tests__/weakMapMemoize.test.ts | 261 ++++++++++++++++++
 .../ui-core/src/util/weakMapMemoize.ts        |  66 +++++
 2 files changed, 327 insertions(+)
 create mode 100644 js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
 create mode 100644 js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts

diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
new file mode 100644
index 0000000000000..140d846d2e58e
--- /dev/null
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
@@ -0,0 +1,261 @@
+import {weakMapMemoize} from '../weakMapMemoize';
+
+type AnyFunction = (...args: any[]) => any;
+
+describe('weakMapMemoize', () => {
+  // Test 1: Function with primitive arguments
+  it('should memoize correctly with primitive arguments and avoid redundant calls', () => {
+    const spy = jest.fn((a: number, b: number) => a + b);
+    const memoizedAdd = weakMapMemoize(spy);
+
+    const result1 = memoizedAdd(1, 2);
+    const result2 = memoizedAdd(1, 2);
+    const result3 = memoizedAdd(2, 3);
+
+    expect(result1).toBe(3);
+    expect(result2).toBe(3);
+    expect(result3).toBe(5);
+    expect(spy).toHaveBeenCalledTimes(2); // Only two unique calls
+  });
+
+  // Test 2: Function with object arguments
+  it('should memoize correctly based on object references', () => {
+    const spy = jest.fn((obj: {x: number}, y: number) => obj.x + y);
+    const memoizedFn = weakMapMemoize(spy);
+
+    const obj1 = {x: 10};
+    const obj2 = {x: 20};
+
+    const result1 = memoizedFn(obj1, 5);
+    const result2 = memoizedFn(obj1, 5);
+    const result3 = memoizedFn(obj2, 5);
+    const result4 = memoizedFn(obj1, 6);
+
+    expect(result1).toBe(15);
+    expect(result2).toBe(15);
+    expect(result3).toBe(25);
+    expect(result4).toBe(16);
+    expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
+  });
+
+  // Test 3: Function with mixed arguments
+  it('should memoize correctly with mixed primitive and object arguments', () => {
+    const spy = jest.fn((a: number, obj: {y: number}) => a + obj.y);
+    const memoizedFn = weakMapMemoize(spy);
+
+    const obj1 = {y: 100};
+    const obj2 = {y: 200};
+
+    const result1 = memoizedFn(1, obj1);
+    const result2 = memoizedFn(1, obj1);
+    const result3 = memoizedFn(2, obj1);
+    const result4 = memoizedFn(1, obj2);
+
+    expect(result1).toBe(101);
+    expect(result2).toBe(101);
+    expect(result3).toBe(102);
+    expect(result4).toBe(201);
+    expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
+  });
+
+  // Test 4: Function with no arguments
+  it('should memoize the result when function has no arguments', () => {
+    const spy = jest.fn(() => Math.random());
+    const memoizedFn = weakMapMemoize(spy);
+
+    const result1 = memoizedFn();
+    const result2 = memoizedFn();
+    const result3 = memoizedFn();
+
+    expect(result1).toBe(result2);
+    expect(result2).toBe(result3);
+    expect(spy).toHaveBeenCalledTimes(1); // Only one unique call
+  });
+
+  // Test 5: Function with null and undefined arguments
+  it('should handle null and undefined arguments correctly', () => {
+    const spy = jest.fn((a: any, b: any) => {
+      if (a === null && b === undefined) {
+        return 'null-undefined';
+      }
+      return 'other';
+    });
+    const memoizedFn = weakMapMemoize(spy);
+
+    const result1 = memoizedFn(null, undefined);
+    const result2 = memoizedFn(null, undefined);
+    const result3 = memoizedFn(undefined, null);
+    const result4 = memoizedFn(null, undefined);
+
+    expect(result1).toBe('null-undefined');
+    expect(result2).toBe('null-undefined');
+    expect(result3).toBe('other');
+    expect(result4).toBe('null-undefined');
+    expect(spy).toHaveBeenCalledTimes(2); // Two unique calls
+  });
+
+  // Test 6: Function with function arguments
+  it('should memoize based on function references', () => {
+    const spy = jest.fn((fn: AnyFunction, value: number) => fn(value));
+    const memoizedFn = weakMapMemoize(spy);
+
+    const func1 = (x: number) => x * 2;
+    const func2 = (x: number) => x * 3;
+
+    const result1 = memoizedFn(func1, 5);
+    const result2 = memoizedFn(func1, 5);
+    const result3 = memoizedFn(func2, 5);
+    const result4 = memoizedFn(func1, 6);
+
+    expect(result1).toBe(10);
+    expect(result2).toBe(10);
+    expect(result3).toBe(15);
+    expect(result4).toBe(12);
+    expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
+  });
+
+  // Test 7: Function with multiple mixed arguments
+  it('should memoize correctly with multiple mixed argument types', () => {
+    const spy = jest.fn((a: number, b: string, c: {key: string}) => `${a}-${b}-${c.key}`);
+    const memoizedFn = weakMapMemoize(spy);
+
+    const obj1 = {key: 'value1'};
+    const obj2 = {key: 'value2'};
+
+    const result1 = memoizedFn(1, 'test', obj1);
+    const result2 = memoizedFn(1, 'test', obj1);
+    const result3 = memoizedFn(1, 'test', obj2);
+    const result4 = memoizedFn(2, 'test', obj1);
+
+    expect(result1).toBe('1-test-value1');
+    expect(result2).toBe('1-test-value1');
+    expect(result3).toBe('1-test-value2');
+    expect(result4).toBe('2-test-value1');
+    expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
+  });
+
+  // Test 8: Function with array arguments
+  it('should memoize based on array references', () => {
+    const spy = jest.fn((arr: number[]) => arr.reduce((sum, val) => sum + val, 0));
+    const memoizedFn = weakMapMemoize(spy);
+
+    const array1 = [1, 2, 3];
+    const array2 = [4, 5, 6];
+
+    const result1 = memoizedFn(array1);
+    const result2 = memoizedFn(array1);
+    const result3 = memoizedFn(array2);
+    const result4 = memoizedFn([1, 2, 3]); // Different reference
+
+    expect(result1).toBe(6);
+    expect(result2).toBe(6);
+    expect(result3).toBe(15);
+    expect(result4).toBe(6);
+    expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
+  });
+
+  // Test 9: Function with symbols as arguments
+  it('should memoize based on symbol references', () => {
+    const sym1 = Symbol('sym1');
+    const sym2 = Symbol('sym2');
+
+    const spy = jest.fn((a: symbol, b: number) => a.toString() + b);
+    const memoizedFn = weakMapMemoize(spy);
+
+    const result1 = memoizedFn(sym1, 10);
+    const result2 = memoizedFn(sym1, 10);
+    const result3 = memoizedFn(sym2, 10);
+    const result4 = memoizedFn(sym1, 20);
+
+    expect(result1).toBe(`${sym1.toString()}10`);
+    expect(result2).toBe(`${sym1.toString()}10`);
+    expect(result3).toBe(`${sym2.toString()}10`);
+    expect(result4).toBe(`${sym1.toString()}20`);
+    expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
+  });
+
+  // Test 10: Function with a large number of arguments
+  it('should memoize correctly with a large number of arguments', () => {
+    const spy = jest.fn((...args: number[]) => args.reduce((sum, val) => sum + val, 0));
+    const memoizedFn = weakMapMemoize(spy);
+
+    const args1 = [1, 2, 3, 4, 5];
+    const args2 = [1, 2, 3, 4, 5];
+    const args3 = [5, 4, 3, 2, 1];
+    const args4 = [1, 2, 3, 4, 6];
+
+    const result1 = memoizedFn(...args1);
+    const result2 = memoizedFn(...args2);
+    const result3 = memoizedFn(...args3);
+    const result4 = memoizedFn(...args4);
+
+    expect(result1).toBe(15);
+    expect(result2).toBe(15);
+    expect(result3).toBe(15);
+    expect(result4).toBe(16);
+    expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
+  });
+
+  // Test 11: Function with alternating object and primitive arguments
+  it('should memoize correctly with alternating object and primitive arguments', () => {
+    const spy = jest.fn(
+      (
+        obj1: {a: number},
+        prim1: string,
+        obj2: {b: number},
+        prim2: boolean,
+        obj3: {c: number},
+        prim3: number,
+      ) => obj1.a + prim1.length + obj2.b + (prim2 ? 1 : 0) + obj3.c + prim3,
+    );
+    const memoizedFn = weakMapMemoize(spy);
+
+    const object1 = {a: 5};
+    const object2 = {b: 10};
+    const object3 = {c: 15};
+
+    // First unique call
+    const result1 = memoizedFn(object1, 'test', object2, true, object3, 20);
+
+    // Duplicate of the first call
+    const result2 = memoizedFn(object1, 'test', object2, true, object3, 20);
+
+    // Different primitive in second argument
+    const result3 = memoizedFn(object1, 'testing', object2, true, object3, 20);
+
+    // Different object in third argument
+    const object4 = {b: 20};
+    const result4 = memoizedFn(object1, 'test', object4, true, object3, 20);
+
+    // Different primitive in fourth argument
+    const result5 = memoizedFn(object1, 'test', object2, false, object3, 20);
+
+    // Different object in fifth argument
+    const object5 = {c: 25};
+    const result6 = memoizedFn(object1, 'test', object2, true, object5, 20);
+
+    // Different primitive in sixth argument
+    const result7 = memoizedFn(object1, 'test', object2, true, object3, 30);
+
+    // Different objects and primitives
+    const result8 = memoizedFn(object1, 'testing', object2, false, object5, 30);
+
+    // Duplicate of the first call again
+    const result9 = memoizedFn(object1, 'test', object2, true, object3, 20);
+
+    expect(result1).toBe(5 + 4 + 10 + 1 + 15 + 20); // 5 + 4 + 10 + 1 + 15 + 20 = 55
+    expect(result2).toBe(55); // Cached
+    expect(result3).toBe(5 + 7 + 10 + 1 + 15 + 20); // 5 + 7 + 10 + 1 + 15 + 20 = 58
+    expect(result4).toBe(5 + 4 + 20 + 1 + 15 + 20); // 5 + 4 + 20 + 1 + 15 + 20 = 65
+    expect(result5).toBe(5 + 4 + 10 + 0 + 15 + 20); // 5 + 4 + 10 + 0 + 15 + 20 = 54
+    expect(result6).toBe(5 + 4 + 10 + 1 + 25 + 20); // 5 + 4 + 10 + 1 + 25 + 20 = 65
+    expect(result7).toBe(5 + 4 + 10 + 1 + 15 + 30); // 5 + 4 + 10 + 1 + 15 + 30 = 65
+    expect(result8).toBe(5 + 7 + 10 + 0 + 25 + 30); // 5 + 7 + 25 + 0 + 15 + 30 = 97
+    expect(result9).toBe(55); // Cached
+
+    // spy should be called for each unique combination
+    // Unique calls: result1, result3, result4, result5, result6, result7, result8
+    // Total unique calls: 7
+    expect(spy).toHaveBeenCalledTimes(7);
+  });
+});
diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
new file mode 100644
index 0000000000000..2e898f4bd1e58
--- /dev/null
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
@@ -0,0 +1,66 @@
+type AnyFunction = (...args: any[]) => any;
+
+function isObject(value: any): value is object {
+  return value !== null && (typeof value === 'object' || typeof value === 'function');
+}
+
+/**
+ * Interface representing a cache node that holds both a Map for primitive keys
+ * and a WeakMap for object keys.
+ */
+interface CacheNode {
+  map: Map<any, CacheNode>;
+  weakMap: WeakMap<object, CacheNode>;
+  result?: any;
+}
+
+/**
+ * Memoizes a function using nested Maps and WeakMaps based on the arguments.
+ * Handles both primitive and object arguments.
+ * @param fn - The function to memoize.
+ * @returns A memoized version of the function.
+ */
+export function weakMapMemoize<T extends AnyFunction>(fn: T): T {
+  // Initialize the root cache node
+  const cacheRoot: CacheNode = {
+    map: new Map(),
+    weakMap: new WeakMap(),
+  };
+
+  return function memoizedFunction(this: any, ...args: any[]) {
+    let currentCache = cacheRoot;
+
+    for (let i = 0; i < args.length; i++) {
+      const arg = args[i];
+
+      if (isObject(arg)) {
+        // Use WeakMap for object keys
+        if (!currentCache.weakMap.has(arg)) {
+          currentCache.weakMap.set(arg, {
+            map: new Map(),
+            weakMap: new WeakMap(),
+          });
+        }
+        currentCache = currentCache.weakMap.get(arg)!;
+      } else {
+        // Use Map for primitive keys
+        if (!currentCache.map.has(arg)) {
+          currentCache.map.set(arg, {
+            map: new Map(),
+            weakMap: new WeakMap(),
+          });
+        }
+        currentCache = currentCache.map.get(arg)!;
+      }
+    }
+
+    // After traversing all arguments, check if the result is cached
+    if ('result' in currentCache) {
+      return currentCache.result;
+    }
+    // Compute the result and cache it
+    const result = fn.apply(this, args);
+    currentCache.result = result;
+    return result;
+  } as T;
+}

From 5d97f0b930d6d62f994ca8517b4434c3c6e861a0 Mon Sep 17 00:00:00 2001
From: Marco Salazar <marco@elementl.com>
Date: Thu, 19 Dec 2024 15:27:59 -0500
Subject: [PATCH 02/15] add max entries options

---
 .../src/util/__tests__/weakMapMemoize.test.ts | 89 +++++++++++------
 .../ui-core/src/util/weakMapMemoize.ts        | 99 ++++++++++++++++---
 2 files changed, 142 insertions(+), 46 deletions(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
index 140d846d2e58e..12009237b32bc 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
@@ -6,7 +6,7 @@ describe('weakMapMemoize', () => {
   // Test 1: Function with primitive arguments
   it('should memoize correctly with primitive arguments and avoid redundant calls', () => {
     const spy = jest.fn((a: number, b: number) => a + b);
-    const memoizedAdd = weakMapMemoize(spy);
+    const memoizedAdd = weakMapMemoize(spy, {maxEntries: 3});
 
     const result1 = memoizedAdd(1, 2);
     const result2 = memoizedAdd(1, 2);
@@ -21,7 +21,7 @@ describe('weakMapMemoize', () => {
   // Test 2: Function with object arguments
   it('should memoize correctly based on object references', () => {
     const spy = jest.fn((obj: {x: number}, y: number) => obj.x + y);
-    const memoizedFn = weakMapMemoize(spy);
+    const memoizedFn = weakMapMemoize(spy, {maxEntries: 2});
 
     const obj1 = {x: 10};
     const obj2 = {x: 20};
@@ -41,7 +41,7 @@ describe('weakMapMemoize', () => {
   // Test 3: Function with mixed arguments
   it('should memoize correctly with mixed primitive and object arguments', () => {
     const spy = jest.fn((a: number, obj: {y: number}) => a + obj.y);
-    const memoizedFn = weakMapMemoize(spy);
+    const memoizedFn = weakMapMemoize(spy, {maxEntries: 3});
 
     const obj1 = {y: 100};
     const obj2 = {y: 200};
@@ -61,7 +61,7 @@ describe('weakMapMemoize', () => {
   // Test 4: Function with no arguments
   it('should memoize the result when function has no arguments', () => {
     const spy = jest.fn(() => Math.random());
-    const memoizedFn = weakMapMemoize(spy);
+    const memoizedFn = weakMapMemoize(spy, {maxEntries: 1});
 
     const result1 = memoizedFn();
     const result2 = memoizedFn();
@@ -80,7 +80,7 @@ describe('weakMapMemoize', () => {
       }
       return 'other';
     });
-    const memoizedFn = weakMapMemoize(spy);
+    const memoizedFn = weakMapMemoize(spy, {maxEntries: 2});
 
     const result1 = memoizedFn(null, undefined);
     const result2 = memoizedFn(null, undefined);
@@ -97,7 +97,7 @@ describe('weakMapMemoize', () => {
   // Test 6: Function with function arguments
   it('should memoize based on function references', () => {
     const spy = jest.fn((fn: AnyFunction, value: number) => fn(value));
-    const memoizedFn = weakMapMemoize(spy);
+    const memoizedFn = weakMapMemoize(spy, {maxEntries: 3});
 
     const func1 = (x: number) => x * 2;
     const func2 = (x: number) => x * 3;
@@ -117,7 +117,7 @@ describe('weakMapMemoize', () => {
   // Test 7: Function with multiple mixed arguments
   it('should memoize correctly with multiple mixed argument types', () => {
     const spy = jest.fn((a: number, b: string, c: {key: string}) => `${a}-${b}-${c.key}`);
-    const memoizedFn = weakMapMemoize(spy);
+    const memoizedFn = weakMapMemoize(spy, {maxEntries: 4});
 
     const obj1 = {key: 'value1'};
     const obj2 = {key: 'value2'};
@@ -137,7 +137,7 @@ describe('weakMapMemoize', () => {
   // Test 8: Function with array arguments
   it('should memoize based on array references', () => {
     const spy = jest.fn((arr: number[]) => arr.reduce((sum, val) => sum + val, 0));
-    const memoizedFn = weakMapMemoize(spy);
+    const memoizedFn = weakMapMemoize(spy, {maxEntries: 3});
 
     const array1 = [1, 2, 3];
     const array2 = [4, 5, 6];
@@ -160,7 +160,7 @@ describe('weakMapMemoize', () => {
     const sym2 = Symbol('sym2');
 
     const spy = jest.fn((a: symbol, b: number) => a.toString() + b);
-    const memoizedFn = weakMapMemoize(spy);
+    const memoizedFn = weakMapMemoize(spy, {maxEntries: 4});
 
     const result1 = memoizedFn(sym1, 10);
     const result2 = memoizedFn(sym1, 10);
@@ -177,23 +177,29 @@ describe('weakMapMemoize', () => {
   // Test 10: Function with a large number of arguments
   it('should memoize correctly with a large number of arguments', () => {
     const spy = jest.fn((...args: number[]) => args.reduce((sum, val) => sum + val, 0));
-    const memoizedFn = weakMapMemoize(spy);
+    const memoizedFn = weakMapMemoize(spy, {maxEntries: 5});
 
     const args1 = [1, 2, 3, 4, 5];
     const args2 = [1, 2, 3, 4, 5];
     const args3 = [5, 4, 3, 2, 1];
     const args4 = [1, 2, 3, 4, 6];
+    const args5 = [6, 5, 4, 3, 2];
+    const args6 = [1, 2, 3, 4, 7];
 
     const result1 = memoizedFn(...args1);
     const result2 = memoizedFn(...args2);
     const result3 = memoizedFn(...args3);
     const result4 = memoizedFn(...args4);
+    const result5 = memoizedFn(...args5);
+    const result6 = memoizedFn(...args6);
 
     expect(result1).toBe(15);
     expect(result2).toBe(15);
     expect(result3).toBe(15);
     expect(result4).toBe(16);
-    expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
+    expect(result5).toBe(20);
+    expect(result6).toBe(17);
+    expect(spy).toHaveBeenCalledTimes(5); // Five unique calls (args1, args3, args4, args5, args6)
   });
 
   // Test 11: Function with alternating object and primitive arguments
@@ -208,49 +214,48 @@ describe('weakMapMemoize', () => {
         prim3: number,
       ) => obj1.a + prim1.length + obj2.b + (prim2 ? 1 : 0) + obj3.c + prim3,
     );
-    const memoizedFn = weakMapMemoize(spy);
+    const memoizedFn = weakMapMemoize(spy, {maxEntries: 7});
 
     const object1 = {a: 5};
     const object2 = {b: 10};
     const object3 = {c: 15};
+    const object4 = {b: 20}; // Corrected to have 'b' property
+    const object5 = {c: 25}; // Corrected to have 'c' property
 
     // First unique call
-    const result1 = memoizedFn(object1, 'test', object2, true, object3, 20);
+    const result1 = memoizedFn(object1, 'test', object2, true, object3, 20); // 5 +4 +10 +1 +15 +20 =55
 
     // Duplicate of the first call
-    const result2 = memoizedFn(object1, 'test', object2, true, object3, 20);
+    const result2 = memoizedFn(object1, 'test', object2, true, object3, 20); // 55
 
     // Different primitive in second argument
-    const result3 = memoizedFn(object1, 'testing', object2, true, object3, 20);
+    const result3 = memoizedFn(object1, 'testing', object2, true, object3, 20); //5 +7 +10 +1 +15 +20=58
 
     // Different object in third argument
-    const object4 = {b: 20};
-    const result4 = memoizedFn(object1, 'test', object4, true, object3, 20);
+    const result4 = memoizedFn(object1, 'test', object4, true, object3, 20); //5 +4 +20 +1 +15 +20=65
 
     // Different primitive in fourth argument
-    const result5 = memoizedFn(object1, 'test', object2, false, object3, 20);
+    const result5 = memoizedFn(object1, 'test', object2, false, object3, 20); //5 +4 +10 +0 +15 +20=54
 
     // Different object in fifth argument
-    const object5 = {c: 25};
-    const result6 = memoizedFn(object1, 'test', object2, true, object5, 20);
+    const result6 = memoizedFn(object1, 'test', object2, true, object5, 20); //5 +4 +10 +1 +25 +20=65
 
     // Different primitive in sixth argument
-    const result7 = memoizedFn(object1, 'test', object2, true, object3, 30);
+    const result7 = memoizedFn(object1, 'test', object2, true, object3, 30); //5 +4 +10 +1 +15 +30=65
 
     // Different objects and primitives
-    const result8 = memoizedFn(object1, 'testing', object2, false, object5, 30);
+    const result8 = memoizedFn(object1, 'testing', object2, false, object3, 30); //5 +7 +10 +0 +15 +30=67
 
     // Duplicate of the first call again
     const result9 = memoizedFn(object1, 'test', object2, true, object3, 20);
-
-    expect(result1).toBe(5 + 4 + 10 + 1 + 15 + 20); // 5 + 4 + 10 + 1 + 15 + 20 = 55
+    expect(result1).toBe(5 + 4 + 10 + 1 + 15 + 20); // 55
     expect(result2).toBe(55); // Cached
-    expect(result3).toBe(5 + 7 + 10 + 1 + 15 + 20); // 5 + 7 + 10 + 1 + 15 + 20 = 58
-    expect(result4).toBe(5 + 4 + 20 + 1 + 15 + 20); // 5 + 4 + 20 + 1 + 15 + 20 = 65
-    expect(result5).toBe(5 + 4 + 10 + 0 + 15 + 20); // 5 + 4 + 10 + 0 + 15 + 20 = 54
-    expect(result6).toBe(5 + 4 + 10 + 1 + 25 + 20); // 5 + 4 + 10 + 1 + 25 + 20 = 65
-    expect(result7).toBe(5 + 4 + 10 + 1 + 15 + 30); // 5 + 4 + 10 + 1 + 15 + 30 = 65
-    expect(result8).toBe(5 + 7 + 10 + 0 + 25 + 30); // 5 + 7 + 25 + 0 + 15 + 30 = 97
+    expect(result3).toBe(5 + 7 + 10 + 1 + 15 + 20); // 58
+    expect(result4).toBe(5 + 4 + 20 + 1 + 15 + 20); // 65
+    expect(result5).toBe(5 + 4 + 10 + 0 + 15 + 20); // 54
+    expect(result6).toBe(5 + 4 + 10 + 1 + 25 + 20); // 65
+    expect(result7).toBe(5 + 4 + 10 + 1 + 15 + 30); // 65
+    expect(result8).toBe(5 + 7 + 10 + 0 + 15 + 30); // 67
     expect(result9).toBe(55); // Cached
 
     // spy should be called for each unique combination
@@ -258,4 +263,28 @@ describe('weakMapMemoize', () => {
     // Total unique calls: 7
     expect(spy).toHaveBeenCalledTimes(7);
   });
+
+  // Test 12: Exercising the maxEntries option
+  it('should evict least recently used entries when maxEntries is exceeded', () => {
+    const spy = jest.fn((a: number) => a * 2);
+    const memoizedFn = weakMapMemoize(spy, {maxEntries: 2});
+
+    const result1 = memoizedFn(1); // Cached
+    const result2 = memoizedFn(2); // Cached
+    const result3 = memoizedFn(3); // Evicts least recently used (1)
+    const result4 = memoizedFn(2); // Cached, updates recentness
+    const result5 = memoizedFn(4); // Evicts least recently used (3)
+
+    expect(result1).toBe(2);
+    expect(result2).toBe(4);
+    expect(result3).toBe(6);
+    expect(result4).toBe(4);
+    expect(result5).toBe(8);
+    expect(spy).toHaveBeenCalledTimes(4); // Calls for 1,2,3,4
+
+    // Accessing 1 again should trigger a new call since it was evicted
+    const result6 = memoizedFn(1);
+    expect(result6).toBe(2);
+    expect(spy).toHaveBeenCalledTimes(5); // Call for 1 again
+  });
 });
diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
index 2e898f4bd1e58..862fc5e6a78ce 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
@@ -1,54 +1,105 @@
+import LRU from 'lru-cache';
+
 type AnyFunction = (...args: any[]) => any;
 
-function isObject(value: any): value is object {
-  return value !== null && (typeof value === 'object' || typeof value === 'function');
+interface WeakMapMemoizeOptions {
+  maxEntries?: number; // Optional limit for cached entries
 }
 
-/**
- * Interface representing a cache node that holds both a Map for primitive keys
- * and a WeakMap for object keys.
- */
 interface CacheNode {
   map: Map<any, CacheNode>;
   weakMap: WeakMap<object, CacheNode>;
   result?: any;
+  lruKey?: any; // Reference to the key in the LRU cache
+}
+
+/**
+ * Determines if a value is a non-null object or function.
+ * @param value - The value to check.
+ * @returns True if the value is a non-null object or function, false otherwise.
+ */
+function isObject(value: any): value is object {
+  return value !== null && (typeof value === 'object' || typeof value === 'function');
 }
 
 /**
  * Memoizes a function using nested Maps and WeakMaps based on the arguments.
- * Handles both primitive and object arguments.
+ * Optionally limits the number of cached entries using an LRU cache.
+ * Handles both primitive and object arguments efficiently.
  * @param fn - The function to memoize.
+ * @param options - Optional settings for memoization.
  * @returns A memoized version of the function.
  */
-export function weakMapMemoize<T extends AnyFunction>(fn: T): T {
+export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMemoizeOptions): T {
+  const {maxEntries} = options || {};
+
   // Initialize the root cache node
   const cacheRoot: CacheNode = {
     map: new Map(),
     weakMap: new WeakMap(),
   };
 
+  // Initialize LRU Cache if maxEntries is specified
+  let lruCache: LRU.Cache<any, any> | null = null;
+
+  if (maxEntries) {
+    lruCache = new LRU<any, any>({
+      max: maxEntries,
+      dispose: (key, value) => {
+        // When an entry is evicted from the LRU cache,
+        // traverse the cache tree and remove the cached result
+        const keyPath = key as any[];
+        let currentCache = cacheRoot;
+
+        for (let i = 0; i < keyPath.length; i++) {
+          const arg = keyPath[i];
+          const isArgObject = isObject(arg);
+
+          if (isArgObject) {
+            currentCache = currentCache.weakMap.get(arg);
+          } else {
+            currentCache = currentCache.map.get(arg);
+          }
+
+          if (!currentCache) {
+            // The cache node has already been removed
+            return;
+          }
+        }
+
+        // Remove the cached result
+        delete currentCache.result;
+        delete currentCache.lruKey;
+      },
+    });
+  }
+
   return function memoizedFunction(this: any, ...args: any[]) {
     let currentCache = cacheRoot;
+    const path: any[] = [];
 
     for (let i = 0; i < args.length; i++) {
       const arg = args[i];
+      path.push(arg);
+      const isArgObject = isObject(arg);
 
-      if (isObject(arg)) {
-        // Use WeakMap for object keys
+      // Determine the appropriate cache level
+      if (isArgObject) {
         if (!currentCache.weakMap.has(arg)) {
-          currentCache.weakMap.set(arg, {
+          const newCacheNode: CacheNode = {
             map: new Map(),
             weakMap: new WeakMap(),
-          });
+          };
+          currentCache.weakMap.set(arg, newCacheNode);
         }
         currentCache = currentCache.weakMap.get(arg)!;
       } else {
-        // Use Map for primitive keys
         if (!currentCache.map.has(arg)) {
-          currentCache.map.set(arg, {
+          const newCacheNode: CacheNode = {
             map: new Map(),
             weakMap: new WeakMap(),
-          });
+          };
+          currentCache.map.set(arg, newCacheNode);
         }
         currentCache = currentCache.map.get(arg)!;
       }
@@ -56,11 +107,27 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T): T {
 
     // After traversing all arguments, check if the result is cached
     if ('result' in currentCache) {
+      // If using LRU Cache, update its usage
+      if (lruCache && currentCache.lruKey) {
+        lruCache.get(currentCache.lruKey); // This updates the recentness
+      }
       return currentCache.result;
     }
-    // Compute the result and cache it
+
+    // Compute the result
     const result = fn.apply(this, args);
+
+    // Cache the result
     currentCache.result = result;
+
+    // If LRU cache is enabled, manage the cache entries
+    if (lruCache) {
+      const cacheEntryKey: any[] = path.slice(); // Clone the path to ensure uniqueness
+      currentCache.lruKey = cacheEntryKey; // Associate the cache node with the LRU key
+
+      lruCache.set(cacheEntryKey, currentCache);
+    }
+
     return result;
   } as T;
 }

From d6fd380e0aadd4f2cbd0c5abcd8a1feef7a6a14e Mon Sep 17 00:00:00 2001
From: Marco Salazar <marco@elementl.com>
Date: Thu, 19 Dec 2024 18:53:30 -0500
Subject: [PATCH 03/15] ts

---
 .../dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts  | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
index 862fc5e6a78ce..8699464a1479a 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
@@ -40,16 +40,16 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
   };
 
   // Initialize LRU Cache if maxEntries is specified
-  let lruCache: LRU.Cache<any, any> | null = null;
+  let lruCache: LRU<any, any> | null = null;
 
   if (maxEntries) {
     lruCache = new LRU<any, any>({
       max: maxEntries,
-      dispose: (key, value) => {
+      dispose: (key, _value) => {
         // When an entry is evicted from the LRU cache,
         // traverse the cache tree and remove the cached result
         const keyPath = key as any[];
-        let currentCache = cacheRoot;
+        let currentCache: CacheNode | undefined = cacheRoot;
 
         for (let i = 0; i < keyPath.length; i++) {
           const arg = keyPath[i];

From bb11f32a16b3a0a9af44c54c21d989029f181060 Mon Sep 17 00:00:00 2001
From: Marco Salazar <marco@elementl.com>
Date: Fri, 20 Dec 2024 14:50:15 -0500
Subject: [PATCH 04/15] use symbol

---
 .../dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
index 8699464a1479a..0eb6c84bc70e0 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
@@ -122,7 +122,7 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
 
     // If LRU cache is enabled, manage the cache entries
     if (lruCache) {
-      const cacheEntryKey: any[] = path.slice(); // Clone the path to ensure uniqueness
+      const cacheEntryKey = Symbol();
       currentCache.lruKey = cacheEntryKey; // Associate the cache node with the LRU key
 
       lruCache.set(cacheEntryKey, currentCache);

From c420f9a0bbbb68e6c4e9fce2cbdab969eafa2c00 Mon Sep 17 00:00:00 2001
From: Marco Salazar <marco@elementl.com>
Date: Fri, 20 Dec 2024 14:56:26 -0500
Subject: [PATCH 05/15] use symbol

---
 .../ui-core/src/util/__tests__/weakMapMemoize.test.ts         | 4 ++--
 .../dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts    | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
index 12009237b32bc..1bb6753ecbeac 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
@@ -265,7 +265,7 @@ describe('weakMapMemoize', () => {
   });
 
   // Test 12: Exercising the maxEntries option
-  it('should evict least recently used entries when maxEntries is exceeded', () => {
+  it.only('should evict least recently used entries when maxEntries is exceeded', () => {
     const spy = jest.fn((a: number) => a * 2);
     const memoizedFn = weakMapMemoize(spy, {maxEntries: 2});
 
@@ -285,6 +285,6 @@ describe('weakMapMemoize', () => {
     // Accessing 1 again should trigger a new call since it was evicted
     const result6 = memoizedFn(1);
     expect(result6).toBe(2);
-    expect(spy).toHaveBeenCalledTimes(5); // Call for 1 again
+    expect(spy).toHaveBeenCalledTimes(4); // Call for 1 again
   });
 });
diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
index 0eb6c84bc70e0..8699464a1479a 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
@@ -122,7 +122,7 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
 
     // If LRU cache is enabled, manage the cache entries
     if (lruCache) {
-      const cacheEntryKey = Symbol();
+      const cacheEntryKey: any[] = path.slice(); // Clone the path to ensure uniqueness
       currentCache.lruKey = cacheEntryKey; // Associate the cache node with the LRU key
 
       lruCache.set(cacheEntryKey, currentCache);

From e31cd5302055814a7191b359a0e111613044f239 Mon Sep 17 00:00:00 2001
From: Marco Salazar <marco@elementl.com>
Date: Fri, 20 Dec 2024 14:57:11 -0500
Subject: [PATCH 06/15] comment

---
 .../dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
index 8699464a1479a..766b7886ed5de 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
@@ -122,7 +122,7 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
 
     // If LRU cache is enabled, manage the cache entries
     if (lruCache) {
-      const cacheEntryKey: any[] = path.slice(); // Clone the path to ensure uniqueness
+      const cacheEntryKey: any[] = path.slice(); // Store the whole path so that we can use it in the dispose call to clear the cache
       currentCache.lruKey = cacheEntryKey; // Associate the cache node with the LRU key
 
       lruCache.set(cacheEntryKey, currentCache);

From 46d5e2b5ba8b96ea88c41e11541f51f554083e2f Mon Sep 17 00:00:00 2001
From: Marco Salazar <marco@elementl.com>
Date: Fri, 20 Dec 2024 14:59:14 -0500
Subject: [PATCH 07/15] fix

---
 .../dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
index 766b7886ed5de..9d7ffe3edd82d 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
@@ -74,7 +74,7 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
     });
   }
 
-  return function memoizedFunction(this: any, ...args: any[]) {
+  return function memoizedFunction(...args: any[]) {
     let currentCache = cacheRoot;
     const path: any[] = [];
 

From a267f63e41e6689bd4afb1a195330636f9672339 Mon Sep 17 00:00:00 2001
From: Marco Salazar <marco@elementl.com>
Date: Fri, 20 Dec 2024 15:00:03 -0500
Subject: [PATCH 08/15] rm only

---
 .../ui-core/src/util/__tests__/weakMapMemoize.test.ts         | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
index 1bb6753ecbeac..12009237b32bc 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
@@ -265,7 +265,7 @@ describe('weakMapMemoize', () => {
   });
 
   // Test 12: Exercising the maxEntries option
-  it.only('should evict least recently used entries when maxEntries is exceeded', () => {
+  it('should evict least recently used entries when maxEntries is exceeded', () => {
     const spy = jest.fn((a: number) => a * 2);
     const memoizedFn = weakMapMemoize(spy, {maxEntries: 2});
 
@@ -285,6 +285,6 @@ describe('weakMapMemoize', () => {
     // Accessing 1 again should trigger a new call since it was evicted
     const result6 = memoizedFn(1);
     expect(result6).toBe(2);
-    expect(spy).toHaveBeenCalledTimes(4); // Call for 1 again
+    expect(spy).toHaveBeenCalledTimes(5); // Call for 1 again
   });
 });

From eaaa4e5b92b0faf781a9e63757aadb096d1bbbf2 Mon Sep 17 00:00:00 2001
From: Marco Salazar <marco@elementl.com>
Date: Fri, 20 Dec 2024 15:11:52 -0500
Subject: [PATCH 09/15] ts

---
 .../src/util/__tests__/weakMapMemoize.test.ts   | 17 +++++++++--------
 .../packages/ui-core/src/util/weakMapMemoize.ts |  2 +-
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
index 12009237b32bc..6461668d8be57 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
@@ -219,26 +219,27 @@ describe('weakMapMemoize', () => {
     const object1 = {a: 5};
     const object2 = {b: 10};
     const object3 = {c: 15};
-    const object4 = {b: 20}; // Corrected to have 'b' property
-    const object5 = {c: 25}; // Corrected to have 'c' property
+    const object4 = {a: 30};
+    const object5 = {b: 20}; // Corrected to have 'b' property
+    const object6 = {c: 25}; // Corrected to have 'c' property
 
     // First unique call
     const result1 = memoizedFn(object1, 'test', object2, true, object3, 20); // 5 +4 +10 +1 +15 +20 =55
 
-    // Duplicate of the first call
-    const result2 = memoizedFn(object1, 'test', object2, true, object3, 20); // 55
+    // Different object in first argument
+    const result2 = memoizedFn(object4, 'test', object2, true, object3, 20); // 30 +4 +10 +1 +15 +20 =55
 
     // Different primitive in second argument
     const result3 = memoizedFn(object1, 'testing', object2, true, object3, 20); //5 +7 +10 +1 +15 +20=58
 
     // Different object in third argument
-    const result4 = memoizedFn(object1, 'test', object4, true, object3, 20); //5 +4 +20 +1 +15 +20=65
+    const result4 = memoizedFn(object1, 'test', object5, true, object3, 20); //5 +4 +20 +1 +15 +20=65
 
     // Different primitive in fourth argument
     const result5 = memoizedFn(object1, 'test', object2, false, object3, 20); //5 +4 +10 +0 +15 +20=54
 
     // Different object in fifth argument
-    const result6 = memoizedFn(object1, 'test', object2, true, object5, 20); //5 +4 +10 +1 +25 +20=65
+    const result6 = memoizedFn(object1, 'test', object2, true, object6, 20); //5 +4 +10 +1 +25 +20=65
 
     // Different primitive in sixth argument
     const result7 = memoizedFn(object1, 'test', object2, true, object3, 30); //5 +4 +10 +1 +15 +30=65
@@ -249,7 +250,7 @@ describe('weakMapMemoize', () => {
     // Duplicate of the first call again
     const result9 = memoizedFn(object1, 'test', object2, true, object3, 20);
     expect(result1).toBe(5 + 4 + 10 + 1 + 15 + 20); // 55
-    expect(result2).toBe(55); // Cached
+    expect(result2).toBe(30 + 4 + 10 + 1 + 15 + 20); // Cached
     expect(result3).toBe(5 + 7 + 10 + 1 + 15 + 20); // 58
     expect(result4).toBe(5 + 4 + 20 + 1 + 15 + 20); // 65
     expect(result5).toBe(5 + 4 + 10 + 0 + 15 + 20); // 54
@@ -261,7 +262,7 @@ describe('weakMapMemoize', () => {
     // spy should be called for each unique combination
     // Unique calls: result1, result3, result4, result5, result6, result7, result8
     // Total unique calls: 7
-    expect(spy).toHaveBeenCalledTimes(7);
+    expect(spy).toHaveBeenCalledTimes(9);
   });
 
   // Test 12: Exercising the maxEntries option
diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
index 9d7ffe3edd82d..40af0be6f5963 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
@@ -115,7 +115,7 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
     }
 
     // Compute the result
-    const result = fn.apply(this, args);
+    const result = fn(...args);
 
     // Cache the result
     currentCache.result = result;

From 208de2830a8269796ab4d4a54dc0e162a39f02e0 Mon Sep 17 00:00:00 2001
From: Marco Salazar <marco@elementl.com>
Date: Fri, 20 Dec 2024 15:43:14 -0500
Subject: [PATCH 10/15] ts

---
 .../src/util/__tests__/weakMapMemoize.test.ts | 235 ++++++++----------
 .../ui-core/src/util/weakMapMemoize.ts        |  59 +++--
 2 files changed, 132 insertions(+), 162 deletions(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
index 6461668d8be57..2ef073d70ce29 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
@@ -8,13 +8,9 @@ describe('weakMapMemoize', () => {
     const spy = jest.fn((a: number, b: number) => a + b);
     const memoizedAdd = weakMapMemoize(spy, {maxEntries: 3});
 
-    const result1 = memoizedAdd(1, 2);
-    const result2 = memoizedAdd(1, 2);
-    const result3 = memoizedAdd(2, 3);
-
-    expect(result1).toBe(3);
-    expect(result2).toBe(3);
-    expect(result3).toBe(5);
+    expect(memoizedAdd(1, 2)).toBe(3); // Cached
+    expect(memoizedAdd(1, 2)).toBe(3); // Cached
+    expect(memoizedAdd(2, 3)).toBe(5); // Cached
     expect(spy).toHaveBeenCalledTimes(2); // Only two unique calls
   });
 
@@ -26,15 +22,10 @@ describe('weakMapMemoize', () => {
     const obj1 = {x: 10};
     const obj2 = {x: 20};
 
-    const result1 = memoizedFn(obj1, 5);
-    const result2 = memoizedFn(obj1, 5);
-    const result3 = memoizedFn(obj2, 5);
-    const result4 = memoizedFn(obj1, 6);
-
-    expect(result1).toBe(15);
-    expect(result2).toBe(15);
-    expect(result3).toBe(25);
-    expect(result4).toBe(16);
+    expect(memoizedFn(obj1, 5)).toBe(15); // Cached
+    expect(memoizedFn(obj1, 5)).toBe(15); // Cached
+    expect(memoizedFn(obj2, 5)).toBe(25); // Cached
+    expect(memoizedFn(obj1, 6)).toBe(16); // Cached
     expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
   });
 
@@ -46,15 +37,10 @@ describe('weakMapMemoize', () => {
     const obj1 = {y: 100};
     const obj2 = {y: 200};
 
-    const result1 = memoizedFn(1, obj1);
-    const result2 = memoizedFn(1, obj1);
-    const result3 = memoizedFn(2, obj1);
-    const result4 = memoizedFn(1, obj2);
-
-    expect(result1).toBe(101);
-    expect(result2).toBe(101);
-    expect(result3).toBe(102);
-    expect(result4).toBe(201);
+    expect(memoizedFn(1, obj1)).toBe(101); // Cached
+    expect(memoizedFn(1, obj1)).toBe(101); // Cached
+    expect(memoizedFn(2, obj1)).toBe(102); // Cached
+    expect(memoizedFn(1, obj2)).toBe(201); // Cached
     expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
   });
 
@@ -67,8 +53,8 @@ describe('weakMapMemoize', () => {
     const result2 = memoizedFn();
     const result3 = memoizedFn();
 
-    expect(result1).toBe(result2);
-    expect(result2).toBe(result3);
+    expect(result1).toBe(result2); // Cached
+    expect(result2).toBe(result3); // Cached
     expect(spy).toHaveBeenCalledTimes(1); // Only one unique call
   });
 
@@ -82,15 +68,10 @@ describe('weakMapMemoize', () => {
     });
     const memoizedFn = weakMapMemoize(spy, {maxEntries: 2});
 
-    const result1 = memoizedFn(null, undefined);
-    const result2 = memoizedFn(null, undefined);
-    const result3 = memoizedFn(undefined, null);
-    const result4 = memoizedFn(null, undefined);
-
-    expect(result1).toBe('null-undefined');
-    expect(result2).toBe('null-undefined');
-    expect(result3).toBe('other');
-    expect(result4).toBe('null-undefined');
+    expect(memoizedFn(null, undefined)).toBe('null-undefined'); // Cached
+    expect(memoizedFn(null, undefined)).toBe('null-undefined'); // Cached
+    expect(memoizedFn(undefined, null)).toBe('other'); // Cached
+    expect(memoizedFn(null, undefined)).toBe('null-undefined'); // Cached
     expect(spy).toHaveBeenCalledTimes(2); // Two unique calls
   });
 
@@ -102,15 +83,10 @@ describe('weakMapMemoize', () => {
     const func1 = (x: number) => x * 2;
     const func2 = (x: number) => x * 3;
 
-    const result1 = memoizedFn(func1, 5);
-    const result2 = memoizedFn(func1, 5);
-    const result3 = memoizedFn(func2, 5);
-    const result4 = memoizedFn(func1, 6);
-
-    expect(result1).toBe(10);
-    expect(result2).toBe(10);
-    expect(result3).toBe(15);
-    expect(result4).toBe(12);
+    expect(memoizedFn(func1, 5)).toBe(10); // Cached
+    expect(memoizedFn(func1, 5)).toBe(10); // Cached
+    expect(memoizedFn(func2, 5)).toBe(15); // Cached
+    expect(memoizedFn(func1, 6)).toBe(12); // Cached
     expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
   });
 
@@ -122,15 +98,10 @@ describe('weakMapMemoize', () => {
     const obj1 = {key: 'value1'};
     const obj2 = {key: 'value2'};
 
-    const result1 = memoizedFn(1, 'test', obj1);
-    const result2 = memoizedFn(1, 'test', obj1);
-    const result3 = memoizedFn(1, 'test', obj2);
-    const result4 = memoizedFn(2, 'test', obj1);
-
-    expect(result1).toBe('1-test-value1');
-    expect(result2).toBe('1-test-value1');
-    expect(result3).toBe('1-test-value2');
-    expect(result4).toBe('2-test-value1');
+    expect(memoizedFn(1, 'test', obj1)).toBe('1-test-value1'); // Cached
+    expect(memoizedFn(1, 'test', obj1)).toBe('1-test-value1'); // Cached
+    expect(memoizedFn(1, 'test', obj2)).toBe('1-test-value2'); // Cached
+    expect(memoizedFn(2, 'test', obj1)).toBe('2-test-value1'); // Cached
     expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
   });
 
@@ -142,15 +113,10 @@ describe('weakMapMemoize', () => {
     const array1 = [1, 2, 3];
     const array2 = [4, 5, 6];
 
-    const result1 = memoizedFn(array1);
-    const result2 = memoizedFn(array1);
-    const result3 = memoizedFn(array2);
-    const result4 = memoizedFn([1, 2, 3]); // Different reference
-
-    expect(result1).toBe(6);
-    expect(result2).toBe(6);
-    expect(result3).toBe(15);
-    expect(result4).toBe(6);
+    expect(memoizedFn(array1)).toBe(6); // Cached
+    expect(memoizedFn(array1)).toBe(6); // Cached
+    expect(memoizedFn(array2)).toBe(15); // Cached
+    expect(memoizedFn([1, 2, 3])).toBe(6); // Cached
     expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
   });
 
@@ -162,15 +128,10 @@ describe('weakMapMemoize', () => {
     const spy = jest.fn((a: symbol, b: number) => a.toString() + b);
     const memoizedFn = weakMapMemoize(spy, {maxEntries: 4});
 
-    const result1 = memoizedFn(sym1, 10);
-    const result2 = memoizedFn(sym1, 10);
-    const result3 = memoizedFn(sym2, 10);
-    const result4 = memoizedFn(sym1, 20);
-
-    expect(result1).toBe(`${sym1.toString()}10`);
-    expect(result2).toBe(`${sym1.toString()}10`);
-    expect(result3).toBe(`${sym2.toString()}10`);
-    expect(result4).toBe(`${sym1.toString()}20`);
+    expect(memoizedFn(sym1, 10)).toBe(`${sym1.toString()}10`); // Cached
+    expect(memoizedFn(sym1, 10)).toBe(`${sym1.toString()}10`); // Cached
+    expect(memoizedFn(sym2, 10)).toBe(`${sym2.toString()}10`); // Cached
+    expect(memoizedFn(sym1, 20)).toBe(`${sym1.toString()}20`); // Cached
     expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
   });
 
@@ -186,19 +147,12 @@ describe('weakMapMemoize', () => {
     const args5 = [6, 5, 4, 3, 2];
     const args6 = [1, 2, 3, 4, 7];
 
-    const result1 = memoizedFn(...args1);
-    const result2 = memoizedFn(...args2);
-    const result3 = memoizedFn(...args3);
-    const result4 = memoizedFn(...args4);
-    const result5 = memoizedFn(...args5);
-    const result6 = memoizedFn(...args6);
-
-    expect(result1).toBe(15);
-    expect(result2).toBe(15);
-    expect(result3).toBe(15);
-    expect(result4).toBe(16);
-    expect(result5).toBe(20);
-    expect(result6).toBe(17);
+    expect(memoizedFn(...args1)).toBe(15); // Cached
+    expect(memoizedFn(...args2)).toBe(15); // Cached
+    expect(memoizedFn(...args3)).toBe(15); // Cached
+    expect(memoizedFn(...args4)).toBe(16); // Cached
+    expect(memoizedFn(...args5)).toBe(20); // Cached
+    expect(memoizedFn(...args6)).toBe(17); // Cached
     expect(spy).toHaveBeenCalledTimes(5); // Five unique calls (args1, args3, args4, args5, args6)
   });
 
@@ -223,45 +177,26 @@ describe('weakMapMemoize', () => {
     const object5 = {b: 20}; // Corrected to have 'b' property
     const object6 = {c: 25}; // Corrected to have 'c' property
 
-    // First unique call
-    const result1 = memoizedFn(object1, 'test', object2, true, object3, 20); // 5 +4 +10 +1 +15 +20 =55
-
-    // Different object in first argument
-    const result2 = memoizedFn(object4, 'test', object2, true, object3, 20); // 30 +4 +10 +1 +15 +20 =55
-
-    // Different primitive in second argument
-    const result3 = memoizedFn(object1, 'testing', object2, true, object3, 20); //5 +7 +10 +1 +15 +20=58
-
-    // Different object in third argument
-    const result4 = memoizedFn(object1, 'test', object5, true, object3, 20); //5 +4 +20 +1 +15 +20=65
-
-    // Different primitive in fourth argument
-    const result5 = memoizedFn(object1, 'test', object2, false, object3, 20); //5 +4 +10 +0 +15 +20=54
-
-    // Different object in fifth argument
-    const result6 = memoizedFn(object1, 'test', object2, true, object6, 20); //5 +4 +10 +1 +25 +20=65
-
-    // Different primitive in sixth argument
-    const result7 = memoizedFn(object1, 'test', object2, true, object3, 30); //5 +4 +10 +1 +15 +30=65
-
-    // Different objects and primitives
-    const result8 = memoizedFn(object1, 'testing', object2, false, object3, 30); //5 +7 +10 +0 +15 +30=67
-
-    // Duplicate of the first call again
-    const result9 = memoizedFn(object1, 'test', object2, true, object3, 20);
-    expect(result1).toBe(5 + 4 + 10 + 1 + 15 + 20); // 55
-    expect(result2).toBe(30 + 4 + 10 + 1 + 15 + 20); // Cached
-    expect(result3).toBe(5 + 7 + 10 + 1 + 15 + 20); // 58
-    expect(result4).toBe(5 + 4 + 20 + 1 + 15 + 20); // 65
-    expect(result5).toBe(5 + 4 + 10 + 0 + 15 + 20); // 54
-    expect(result6).toBe(5 + 4 + 10 + 1 + 25 + 20); // 65
-    expect(result7).toBe(5 + 4 + 10 + 1 + 15 + 30); // 65
-    expect(result8).toBe(5 + 7 + 10 + 0 + 15 + 30); // 67
-    expect(result9).toBe(55); // Cached
+    expect(memoizedFn(object1, 'test', object2, true, object3, 20)).toBe(5 + 4 + 10 + 1 + 15 + 20); // 55
+    expect(memoizedFn(object4, 'test', object2, true, object3, 20)).toBe(30 + 4 + 10 + 1 + 15 + 20); // 80
+    expect(memoizedFn(object1, 'testing', object2, true, object3, 20)).toBe(
+      5 + 7 + 10 + 1 + 15 + 20,
+    ); // 58
+    expect(memoizedFn(object1, 'test', object5, true, object3, 20)).toBe(5 + 4 + 20 + 1 + 15 + 20); // 65
+    expect(memoizedFn(object1, 'test', object2, false, object3, 20)).toBe(5 + 4 + 10 + 0 + 15 + 20); // 54
+    expect(memoizedFn(object1, 'test', object2, true, object6, 20)).toBe(5 + 4 + 10 + 1 + 25 + 20); // 65
+    expect(memoizedFn(object1, 'test', object2, true, object3, 30)).toBe(5 + 4 + 10 + 1 + 15 + 30); // 65
+    expect(memoizedFn(object1, 'testing', object2, false, object3, 30)).toBe(
+      5 + 7 + 10 + 0 + 15 + 30,
+    ); // 67
+    expect(memoizedFn(object1, 'test', object2, true, object3, 20)).toBe(55); // Cached
 
     // spy should be called for each unique combination
-    // Unique calls: result1, result3, result4, result5, result6, result7, result8
-    // Total unique calls: 7
+    // Unique calls: object1-test-object2-true-object3-20, object4-test-object2-true-object3-20,
+    // object1-testing-object2-true-object3-20, object1-test-object5-true-object3-20,
+    // object1-test-object2-false-object3-20, object1-test-object2-true-object6-20,
+    // object1-test-object2-true-object3-30, object1-testing-object2-false-object3-30
+    // Total unique calls: 8
     expect(spy).toHaveBeenCalledTimes(9);
   });
 
@@ -270,22 +205,52 @@ describe('weakMapMemoize', () => {
     const spy = jest.fn((a: number) => a * 2);
     const memoizedFn = weakMapMemoize(spy, {maxEntries: 2});
 
-    const result1 = memoizedFn(1); // Cached
-    const result2 = memoizedFn(2); // Cached
-    const result3 = memoizedFn(3); // Evicts least recently used (1)
-    const result4 = memoizedFn(2); // Cached, updates recentness
-    const result5 = memoizedFn(4); // Evicts least recently used (3)
-
-    expect(result1).toBe(2);
-    expect(result2).toBe(4);
-    expect(result3).toBe(6);
-    expect(result4).toBe(4);
-    expect(result5).toBe(8);
+    expect(memoizedFn(1)).toBe(2); // Cached
+    expect(memoizedFn(2)).toBe(4); // Cached
+    expect(memoizedFn(3)).toBe(6); // Evicts least recently used (1)
+    expect(memoizedFn(2)).toBe(4); // Cached, updates recentness
+    expect(memoizedFn(4)).toBe(8); // Evicts least recently used (3)
     expect(spy).toHaveBeenCalledTimes(4); // Calls for 1,2,3,4
 
     // Accessing 1 again should trigger a new call since it was evicted
-    const result6 = memoizedFn(1);
-    expect(result6).toBe(2);
+    expect(memoizedFn(1)).toBe(2); // Cached
     expect(spy).toHaveBeenCalledTimes(5); // Call for 1 again
   });
+
+  // Test 13: Should handle mixed argument types in same slot
+  it('Should handle mixed argument types in same slot', () => {
+    const spy = jest.fn((a: any, b: any) => {
+      const toString = (arg: any) => (typeof arg === 'object' ? Object.keys(arg).length : arg);
+      return toString(a) + toString(b);
+    });
+    const memoizedFn = weakMapMemoize(spy, {maxEntries: 3});
+
+    expect(memoizedFn(1, 2)).toBe(3); // Cached
+
+    expect(memoizedFn(1, 2)).toBe(3);
+    expect(spy).toHaveBeenCalledTimes(1);
+
+    expect(memoizedFn('hello', {x: 1, y: 2})).toBe('hello2');
+    expect(spy).toHaveBeenCalledTimes(2);
+
+    // Different object reference results in another call
+    expect(memoizedFn('hello', {x: 1, y: 2})).toBe('hello2');
+    expect(spy).toHaveBeenCalledTimes(3);
+
+    const a = {a: 1};
+    expect(memoizedFn(a, 'test')).toBe('1test');
+    expect(spy).toHaveBeenCalledTimes(4);
+
+    // same reference uses cache
+    expect(memoizedFn(a, 'test')).toBe('1test');
+    expect(spy).toHaveBeenCalledTimes(4);
+
+    // Evicts 1, 2
+    expect(memoizedFn(3, 4)).toBe(7);
+    expect(spy).toHaveBeenCalledTimes(5);
+
+    // Evicted - results in new call
+    expect(memoizedFn(1, 2)).toBe(3);
+    expect(spy).toHaveBeenCalledTimes(6);
+  });
 });
diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
index 40af0be6f5963..185341beabb3d 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
@@ -1,3 +1,5 @@
+// src/utils/weakMapMemoize.ts
+
 import LRU from 'lru-cache';
 
 type AnyFunction = (...args: any[]) => any;
@@ -10,6 +12,8 @@ interface CacheNode {
   map: Map<any, CacheNode>;
   weakMap: WeakMap<object, CacheNode>;
   result?: any;
+  parent?: CacheNode;
+  parentKey?: any;
   lruKey?: any; // Reference to the key in the LRU cache
 }
 
@@ -40,37 +44,31 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
   };
 
   // Initialize LRU Cache if maxEntries is specified
-  let lruCache: LRU<any, any> | null = null;
+  let lruCache: LRU<any, CacheNode> | null = null;
 
   if (maxEntries) {
-    lruCache = new LRU<any, any>({
+    lruCache = new LRU<any, CacheNode>({
       max: maxEntries,
-      dispose: (key, _value) => {
-        // When an entry is evicted from the LRU cache,
-        // traverse the cache tree and remove the cached result
-        const keyPath = key as any[];
-        let currentCache: CacheNode | undefined = cacheRoot;
-
-        for (let i = 0; i < keyPath.length; i++) {
-          const arg = keyPath[i];
-          const isArgObject = isObject(arg);
-
-          if (isArgObject) {
-            currentCache = currentCache.weakMap.get(arg);
+      dispose: (key, cacheNode) => {
+        // When an entry is evicted from the LRU cache, remove it from its parent
+        if (cacheNode.parent && cacheNode.parentKey !== undefined) {
+          const parent = cacheNode.parent;
+          const parentKey = cacheNode.parentKey;
+
+          if (isObject(parentKey)) {
+            parent.weakMap.delete(parentKey);
           } else {
-            currentCache = currentCache.map.get(arg);
+            parent.map.delete(parentKey);
           }
 
-          if (!currentCache) {
-            // The cache node has already been removed
-            return;
-          }
+          // Remove references to help garbage collection
+          delete cacheNode.parent;
+          delete cacheNode.parentKey;
+          delete cacheNode.result;
+          delete cacheNode.lruKey;
         }
-
-        // Remove the cached result
-        delete currentCache.result;
-        delete currentCache.lruKey;
       },
+      noDisposeOnSet: false, // Ensure dispose is called on eviction
     });
   }
 
@@ -83,26 +81,33 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
       path.push(arg);
       const isArgObject = isObject(arg);
 
-      // Determine the appropriate cache level
+      let nextCacheNode: CacheNode | undefined;
+
       if (isArgObject) {
         if (!currentCache.weakMap.has(arg)) {
           const newCacheNode: CacheNode = {
             map: new Map(),
             weakMap: new WeakMap(),
+            parent: currentCache,
+            parentKey: arg,
           };
           currentCache.weakMap.set(arg, newCacheNode);
         }
-        currentCache = currentCache.weakMap.get(arg)!;
+        nextCacheNode = currentCache.weakMap.get(arg);
       } else {
         if (!currentCache.map.has(arg)) {
           const newCacheNode: CacheNode = {
             map: new Map(),
             weakMap: new WeakMap(),
+            parent: currentCache,
+            parentKey: arg,
           };
           currentCache.map.set(arg, newCacheNode);
         }
-        currentCache = currentCache.map.get(arg)!;
+        nextCacheNode = currentCache.map.get(arg);
       }
+
+      currentCache = nextCacheNode!;
     }
 
     // After traversing all arguments, check if the result is cached
@@ -122,7 +127,7 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
 
     // If LRU cache is enabled, manage the cache entries
     if (lruCache) {
-      const cacheEntryKey: any[] = path.slice(); // Store the whole path so that we can use it in the dispose call to clear the cache
+      const cacheEntryKey: any[] = path.slice(); // Clone the path to ensure uniqueness
       currentCache.lruKey = cacheEntryKey; // Associate the cache node with the LRU key
 
       lruCache.set(cacheEntryKey, currentCache);

From f54f979ec746c71b41350fae2a623e418f6a5a0e Mon Sep 17 00:00:00 2001
From: Marco Salazar <marco@elementl.com>
Date: Fri, 20 Dec 2024 15:45:43 -0500
Subject: [PATCH 11/15] comments

---
 .../src/util/__tests__/weakMapMemoize.test.ts | 89 +++++++++----------
 1 file changed, 42 insertions(+), 47 deletions(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
index 2ef073d70ce29..9e9742f52913c 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
@@ -8,9 +8,9 @@ describe('weakMapMemoize', () => {
     const spy = jest.fn((a: number, b: number) => a + b);
     const memoizedAdd = weakMapMemoize(spy, {maxEntries: 3});
 
-    expect(memoizedAdd(1, 2)).toBe(3); // Cached
-    expect(memoizedAdd(1, 2)).toBe(3); // Cached
-    expect(memoizedAdd(2, 3)).toBe(5); // Cached
+    expect(memoizedAdd(1, 2)).toBe(3);
+    expect(memoizedAdd(1, 2)).toBe(3);
+    expect(memoizedAdd(2, 3)).toBe(5);
     expect(spy).toHaveBeenCalledTimes(2); // Only two unique calls
   });
 
@@ -22,10 +22,10 @@ describe('weakMapMemoize', () => {
     const obj1 = {x: 10};
     const obj2 = {x: 20};
 
-    expect(memoizedFn(obj1, 5)).toBe(15); // Cached
-    expect(memoizedFn(obj1, 5)).toBe(15); // Cached
-    expect(memoizedFn(obj2, 5)).toBe(25); // Cached
-    expect(memoizedFn(obj1, 6)).toBe(16); // Cached
+    expect(memoizedFn(obj1, 5)).toBe(15);
+    expect(memoizedFn(obj1, 5)).toBe(15);
+    expect(memoizedFn(obj2, 5)).toBe(25);
+    expect(memoizedFn(obj1, 6)).toBe(16);
     expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
   });
 
@@ -37,10 +37,10 @@ describe('weakMapMemoize', () => {
     const obj1 = {y: 100};
     const obj2 = {y: 200};
 
-    expect(memoizedFn(1, obj1)).toBe(101); // Cached
-    expect(memoizedFn(1, obj1)).toBe(101); // Cached
-    expect(memoizedFn(2, obj1)).toBe(102); // Cached
-    expect(memoizedFn(1, obj2)).toBe(201); // Cached
+    expect(memoizedFn(1, obj1)).toBe(101);
+    expect(memoizedFn(1, obj1)).toBe(101);
+    expect(memoizedFn(2, obj1)).toBe(102);
+    expect(memoizedFn(1, obj2)).toBe(201);
     expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
   });
 
@@ -53,8 +53,8 @@ describe('weakMapMemoize', () => {
     const result2 = memoizedFn();
     const result3 = memoizedFn();
 
-    expect(result1).toBe(result2); // Cached
-    expect(result2).toBe(result3); // Cached
+    expect(result1).toBe(result2);
+    expect(result2).toBe(result3);
     expect(spy).toHaveBeenCalledTimes(1); // Only one unique call
   });
 
@@ -68,10 +68,10 @@ describe('weakMapMemoize', () => {
     });
     const memoizedFn = weakMapMemoize(spy, {maxEntries: 2});
 
-    expect(memoizedFn(null, undefined)).toBe('null-undefined'); // Cached
-    expect(memoizedFn(null, undefined)).toBe('null-undefined'); // Cached
-    expect(memoizedFn(undefined, null)).toBe('other'); // Cached
-    expect(memoizedFn(null, undefined)).toBe('null-undefined'); // Cached
+    expect(memoizedFn(null, undefined)).toBe('null-undefined');
+    expect(memoizedFn(null, undefined)).toBe('null-undefined');
+    expect(memoizedFn(undefined, null)).toBe('other');
+    expect(memoizedFn(null, undefined)).toBe('null-undefined');
     expect(spy).toHaveBeenCalledTimes(2); // Two unique calls
   });
 
@@ -83,10 +83,10 @@ describe('weakMapMemoize', () => {
     const func1 = (x: number) => x * 2;
     const func2 = (x: number) => x * 3;
 
-    expect(memoizedFn(func1, 5)).toBe(10); // Cached
-    expect(memoizedFn(func1, 5)).toBe(10); // Cached
-    expect(memoizedFn(func2, 5)).toBe(15); // Cached
-    expect(memoizedFn(func1, 6)).toBe(12); // Cached
+    expect(memoizedFn(func1, 5)).toBe(10);
+    expect(memoizedFn(func1, 5)).toBe(10);
+    expect(memoizedFn(func2, 5)).toBe(15);
+    expect(memoizedFn(func1, 6)).toBe(12);
     expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
   });
 
@@ -98,10 +98,10 @@ describe('weakMapMemoize', () => {
     const obj1 = {key: 'value1'};
     const obj2 = {key: 'value2'};
 
-    expect(memoizedFn(1, 'test', obj1)).toBe('1-test-value1'); // Cached
-    expect(memoizedFn(1, 'test', obj1)).toBe('1-test-value1'); // Cached
-    expect(memoizedFn(1, 'test', obj2)).toBe('1-test-value2'); // Cached
-    expect(memoizedFn(2, 'test', obj1)).toBe('2-test-value1'); // Cached
+    expect(memoizedFn(1, 'test', obj1)).toBe('1-test-value1');
+    expect(memoizedFn(1, 'test', obj1)).toBe('1-test-value1');
+    expect(memoizedFn(1, 'test', obj2)).toBe('1-test-value2');
+    expect(memoizedFn(2, 'test', obj1)).toBe('2-test-value1');
     expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
   });
 
@@ -113,10 +113,10 @@ describe('weakMapMemoize', () => {
     const array1 = [1, 2, 3];
     const array2 = [4, 5, 6];
 
-    expect(memoizedFn(array1)).toBe(6); // Cached
-    expect(memoizedFn(array1)).toBe(6); // Cached
-    expect(memoizedFn(array2)).toBe(15); // Cached
-    expect(memoizedFn([1, 2, 3])).toBe(6); // Cached
+    expect(memoizedFn(array1)).toBe(6);
+    expect(memoizedFn(array1)).toBe(6);
+    expect(memoizedFn(array2)).toBe(15);
+    expect(memoizedFn([1, 2, 3])).toBe(6);
     expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
   });
 
@@ -128,10 +128,10 @@ describe('weakMapMemoize', () => {
     const spy = jest.fn((a: symbol, b: number) => a.toString() + b);
     const memoizedFn = weakMapMemoize(spy, {maxEntries: 4});
 
-    expect(memoizedFn(sym1, 10)).toBe(`${sym1.toString()}10`); // Cached
-    expect(memoizedFn(sym1, 10)).toBe(`${sym1.toString()}10`); // Cached
-    expect(memoizedFn(sym2, 10)).toBe(`${sym2.toString()}10`); // Cached
-    expect(memoizedFn(sym1, 20)).toBe(`${sym1.toString()}20`); // Cached
+    expect(memoizedFn(sym1, 10)).toBe(`${sym1.toString()}10`);
+    expect(memoizedFn(sym1, 10)).toBe(`${sym1.toString()}10`);
+    expect(memoizedFn(sym2, 10)).toBe(`${sym2.toString()}10`);
+    expect(memoizedFn(sym1, 20)).toBe(`${sym1.toString()}20`);
     expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
   });
 
@@ -147,12 +147,12 @@ describe('weakMapMemoize', () => {
     const args5 = [6, 5, 4, 3, 2];
     const args6 = [1, 2, 3, 4, 7];
 
-    expect(memoizedFn(...args1)).toBe(15); // Cached
-    expect(memoizedFn(...args2)).toBe(15); // Cached
-    expect(memoizedFn(...args3)).toBe(15); // Cached
-    expect(memoizedFn(...args4)).toBe(16); // Cached
-    expect(memoizedFn(...args5)).toBe(20); // Cached
-    expect(memoizedFn(...args6)).toBe(17); // Cached
+    expect(memoizedFn(...args2)).toBe(15);
+    expect(memoizedFn(...args1)).toBe(15);
+    expect(memoizedFn(...args3)).toBe(15);
+    expect(memoizedFn(...args4)).toBe(16);
+    expect(memoizedFn(...args5)).toBe(20);
+    expect(memoizedFn(...args6)).toBe(17);
     expect(spy).toHaveBeenCalledTimes(5); // Five unique calls (args1, args3, args4, args5, args6)
   });
 
@@ -192,11 +192,6 @@ describe('weakMapMemoize', () => {
     expect(memoizedFn(object1, 'test', object2, true, object3, 20)).toBe(55); // Cached
 
     // spy should be called for each unique combination
-    // Unique calls: object1-test-object2-true-object3-20, object4-test-object2-true-object3-20,
-    // object1-testing-object2-true-object3-20, object1-test-object5-true-object3-20,
-    // object1-test-object2-false-object3-20, object1-test-object2-true-object6-20,
-    // object1-test-object2-true-object3-30, object1-testing-object2-false-object3-30
-    // Total unique calls: 8
     expect(spy).toHaveBeenCalledTimes(9);
   });
 
@@ -205,8 +200,8 @@ describe('weakMapMemoize', () => {
     const spy = jest.fn((a: number) => a * 2);
     const memoizedFn = weakMapMemoize(spy, {maxEntries: 2});
 
-    expect(memoizedFn(1)).toBe(2); // Cached
-    expect(memoizedFn(2)).toBe(4); // Cached
+    expect(memoizedFn(1)).toBe(2);
+    expect(memoizedFn(2)).toBe(4);
     expect(memoizedFn(3)).toBe(6); // Evicts least recently used (1)
     expect(memoizedFn(2)).toBe(4); // Cached, updates recentness
     expect(memoizedFn(4)).toBe(8); // Evicts least recently used (3)
@@ -225,7 +220,7 @@ describe('weakMapMemoize', () => {
     });
     const memoizedFn = weakMapMemoize(spy, {maxEntries: 3});
 
-    expect(memoizedFn(1, 2)).toBe(3); // Cached
+    expect(memoizedFn(1, 2)).toBe(3);
 
     expect(memoizedFn(1, 2)).toBe(3);
     expect(spy).toHaveBeenCalledTimes(1);

From c9d2ede3781a28d2f2abddda2da1a7eeacb20a9b Mon Sep 17 00:00:00 2001
From: Marco Salazar <marco@elementl.com>
Date: Fri, 20 Dec 2024 15:55:12 -0500
Subject: [PATCH 12/15] rm comment

---
 .../dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts      | 2 --
 1 file changed, 2 deletions(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
index 185341beabb3d..f984c3a81bbf0 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
@@ -1,5 +1,3 @@
-// src/utils/weakMapMemoize.ts
-
 import LRU from 'lru-cache';
 
 type AnyFunction = (...args: any[]) => any;

From f40c8e4d43f187a053e244051e577c57e3a6fb82 Mon Sep 17 00:00:00 2001
From: Marco Salazar <marco@elementl.com>
Date: Fri, 20 Dec 2024 16:09:19 -0500
Subject: [PATCH 13/15] simple

---
 .../dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts  | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
index f984c3a81bbf0..1f8522098ed74 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
@@ -58,12 +58,6 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
           } else {
             parent.map.delete(parentKey);
           }
-
-          // Remove references to help garbage collection
-          delete cacheNode.parent;
-          delete cacheNode.parentKey;
-          delete cacheNode.result;
-          delete cacheNode.lruKey;
         }
       },
       noDisposeOnSet: false, // Ensure dispose is called on eviction

From 80e4ac3c34ee144b8e9418be577c30630ed14420 Mon Sep 17 00:00:00 2001
From: Marco Salazar <marco@elementl.com>
Date: Fri, 20 Dec 2024 16:23:34 -0500
Subject: [PATCH 14/15] fix

---
 .../src/util/__tests__/weakMapMemoize.test.ts |  4 ++--
 .../ui-core/src/util/weakMapMemoize.ts        | 19 +++++++++++++++----
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
index 9e9742f52913c..c99ca26e1f43e 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
@@ -174,8 +174,8 @@ describe('weakMapMemoize', () => {
     const object2 = {b: 10};
     const object3 = {c: 15};
     const object4 = {a: 30};
-    const object5 = {b: 20}; // Corrected to have 'b' property
-    const object6 = {c: 25}; // Corrected to have 'c' property
+    const object5 = {b: 20};
+    const object6 = {c: 25};
 
     expect(memoizedFn(object1, 'test', object2, true, object3, 20)).toBe(5 + 4 + 10 + 1 + 15 + 20); // 55
     expect(memoizedFn(object4, 'test', object2, true, object3, 20)).toBe(30 + 4 + 10 + 1 + 15 + 20); // 80
diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
index 1f8522098ed74..a7d79c71c7307 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
@@ -13,6 +13,7 @@ interface CacheNode {
   parent?: CacheNode;
   parentKey?: any;
   lruKey?: any; // Reference to the key in the LRU cache
+  childCount: number; // Number of child nodes
 }
 
 /**
@@ -39,6 +40,7 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
   const cacheRoot: CacheNode = {
     map: new Map(),
     weakMap: new WeakMap(),
+    childCount: 0,
   };
 
   // Initialize LRU Cache if maxEntries is specified
@@ -48,9 +50,15 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
     lruCache = new LRU<any, CacheNode>({
       max: maxEntries,
       dispose: (key, cacheNode) => {
-        // When an entry is evicted from the LRU cache, remove it from its parent
-        if (cacheNode.parent && cacheNode.parentKey !== undefined) {
+        // Remove the cached entry
+        delete cacheNode.result;
+
+        // If there are no entries in this node then lets remove it.
+        if (cacheNode.childCount === 0 && cacheNode.parent && cacheNode.parentKey !== undefined) {
           const parent = cacheNode.parent;
+
+          // Decrement the parent's child count
+          parent.childCount--;
           const parentKey = cacheNode.parentKey;
 
           if (isObject(parentKey)) {
@@ -82,8 +90,10 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
             weakMap: new WeakMap(),
             parent: currentCache,
             parentKey: arg,
+            childCount: 0,
           };
           currentCache.weakMap.set(arg, newCacheNode);
+          currentCache.childCount++;
         }
         nextCacheNode = currentCache.weakMap.get(arg);
       } else {
@@ -93,8 +103,10 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
             weakMap: new WeakMap(),
             parent: currentCache,
             parentKey: arg,
+            childCount: 0,
           };
           currentCache.map.set(arg, newCacheNode);
+          currentCache.childCount++;
         }
         nextCacheNode = currentCache.map.get(arg);
       }
@@ -118,10 +130,9 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
     currentCache.result = result;
 
     // If LRU cache is enabled, manage the cache entries
-    if (lruCache) {
+    if (lruCache && !currentCache.lruKey) {
       const cacheEntryKey: any[] = path.slice(); // Clone the path to ensure uniqueness
       currentCache.lruKey = cacheEntryKey; // Associate the cache node with the LRU key
-
       lruCache.set(cacheEntryKey, currentCache);
     }
 

From 3715f6855af0bbc37330e57acba4635f57c5111b Mon Sep 17 00:00:00 2001
From: Marco Salazar <marco@elementl.com>
Date: Fri, 20 Dec 2024 17:01:15 -0500
Subject: [PATCH 15/15] fix

---
 .../src/util/__tests__/weakMapMemoize.test.ts | 76 +++++++++++++++++++
 .../ui-core/src/util/weakMapMemoize.ts        | 46 +++++++----
 2 files changed, 107 insertions(+), 15 deletions(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
index c99ca26e1f43e..40b754bd5975b 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
@@ -248,4 +248,80 @@ describe('weakMapMemoize', () => {
     expect(memoizedFn(1, 2)).toBe(3);
     expect(spy).toHaveBeenCalledTimes(6);
   });
+
+  // Test 14: Evicting a cache node without children deletes its result
+  it('should delete the result of an evicted cache node without children', () => {
+    const spy = jest.fn((a: number, b: number) => a + b);
+    const memoizedFn = weakMapMemoize(spy, {maxEntries: 2});
+
+    expect(memoizedFn(1, 1)).toBe(2);
+    expect(memoizedFn(2, 2)).toBe(4);
+    expect(spy).toHaveBeenCalledTimes(2); // Two unique calls
+
+    // This should evict the least recently used entry (1,1)
+    expect(memoizedFn(3, 3)).toBe(6);
+    expect(spy).toHaveBeenCalledTimes(3); // New call for (3,3)
+
+    // Accessing (1,1) again should trigger a new function call since it was evicted, and evict (2, 2)
+    expect(memoizedFn(1, 1)).toBe(2);
+    expect(spy).toHaveBeenCalledTimes(4); // New call for (1,1)
+
+    // evicts 3,3
+    expect(memoizedFn(2, 2)).toBe(4);
+    expect(spy).toHaveBeenCalledTimes(5);
+
+    expect(memoizedFn(3, 3)).toBe(6);
+    expect(spy).toHaveBeenCalledTimes(6);
+  });
+
+  // Test 15: Evicting a cache node with children deletes only its result, children remain cached
+  it('should delete only the result of an evicted cache node with children and keep children cached', () => {
+    const spy = jest.fn((a: number, b: string, c?: number) => a + b.length + (c ?? 0));
+    const memoizedFn = weakMapMemoize(spy, {maxEntries: 3});
+
+    // Creating cache nodes:
+    // (1, 'a')
+    // (1, 'a', 10)
+    // (2, 'b', 30)
+
+    expect(memoizedFn(1, 'a')).toBe(1 + 1);
+    expect(memoizedFn(1, 'a', 10)).toBe(1 + 1 + 10);
+    expect(memoizedFn(2, 'b', 30)).toBe(2 + 1 + 30);
+    expect(spy).toHaveBeenCalledTimes(3); // Three unique calls
+
+    // At this point, maxEntries is 3:
+    // Cache contains:
+    // (1, 'a'), (1, 'a', 10), (2, 'b', 30)
+
+    // Adding a new entry should evict the least recently used, which is (1, 'a')
+    expect(memoizedFn(1, 'c', 40)).toBe(1 + 1 + 40);
+    expect(spy).toHaveBeenCalledTimes(4);
+
+    // Now, cache contains:
+    // (1, 'a', 10), (2, 'b', 30), (1, 'c', 40)
+    expect(memoizedFn(1, 'a', 10)).toBe(1 + 1 + 10);
+    expect(spy).toHaveBeenCalledTimes(4); // no new call
+
+    // Now, cache contains:
+    // (2, 'b', 30), (1, 'c', 40), (1, 'a', 10)
+    // Evicting (2, 'b', 30)
+    expect(memoizedFn(4, 'd', 50)).toBe(4 + 1 + 50);
+    expect(spy).toHaveBeenCalledTimes(5);
+
+    // Now, cache contains:
+    // (1, 'c', 40), (1, 'a', 10), (4, 'd', 50)
+
+    // Accessing (1, 'a', 10) again should trigger a new call since it was evicted and evict (2, b, 30)
+    expect(memoizedFn(2, 'b', 30)).toBe(2 + 1 + 30);
+    expect(spy).toHaveBeenCalledTimes(6);
+
+    // Now, cache contains:
+    // (1, 'a', 10), (4, 'd', 50), (2, 'b', 30)
+
+    expect(memoizedFn(1, 'a', 10)).toBe(12);
+    expect(memoizedFn(4, 'd', 50)).toBe(55);
+    expect(memoizedFn(2, 'b', 30)).toBe(33);
+
+    expect(spy).toHaveBeenCalledTimes(6);
+  });
 });
diff --git a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
index a7d79c71c7307..ec328e445d71d 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/util/weakMapMemoize.ts
@@ -25,6 +25,32 @@ function isObject(value: any): value is object {
   return value !== null && (typeof value === 'object' || typeof value === 'function');
 }
 
+/**
+ * Recursively deletes parent nodes if their childCount reaches zero.
+ * @param cacheNode - The cache node to start deletion from.
+ */
+function recursivelyDeleteParent(cacheNode: CacheNode) {
+  if (cacheNode.parent && cacheNode.parentKey !== undefined) {
+    const parent = cacheNode.parent;
+    const parentKey = cacheNode.parentKey;
+
+    // Remove the current cacheNode from its parent
+    if (isObject(parentKey)) {
+      parent.weakMap.delete(parentKey);
+    } else {
+      parent.map.delete(parentKey);
+    }
+
+    // Decrement the parent's child count
+    parent.childCount--;
+
+    // If the parent's childCount reaches zero and it's not the root, recurse
+    if (parent.childCount === 0 && parent.parent) {
+      recursivelyDeleteParent(parent);
+    }
+  }
+}
+
 /**
  * Memoizes a function using nested Maps and WeakMaps based on the arguments.
  * Optionally limits the number of cached entries using an LRU cache.
@@ -49,23 +75,13 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
   if (maxEntries) {
     lruCache = new LRU<any, CacheNode>({
       max: maxEntries,
-      dispose: (key, cacheNode) => {
-        // Remove the cached entry
+      dispose: (_key, cacheNode) => {
+        // Remove the cached result
         delete cacheNode.result;
 
-        // If there are no entries in this node then lets remove it.
+        // If there are no child nodes, proceed to remove this node from its parent
         if (cacheNode.childCount === 0 && cacheNode.parent && cacheNode.parentKey !== undefined) {
-          const parent = cacheNode.parent;
-
-          // Decrement the parent's child count
-          parent.childCount--;
-          const parentKey = cacheNode.parentKey;
-
-          if (isObject(parentKey)) {
-            parent.weakMap.delete(parentKey);
-          } else {
-            parent.map.delete(parentKey);
-          }
+          recursivelyDeleteParent(cacheNode);
         }
       },
       noDisposeOnSet: false, // Ensure dispose is called on eviction
@@ -131,7 +147,7 @@ export function weakMapMemoize<T extends AnyFunction>(fn: T, options?: WeakMapMe
 
     // If LRU cache is enabled, manage the cache entries
     if (lruCache && !currentCache.lruKey) {
-      const cacheEntryKey: any[] = path.slice(); // Clone the path to ensure uniqueness
+      const cacheEntryKey = Symbol();
       currentCache.lruKey = cacheEntryKey; // Associate the cache node with the LRU key
       lruCache.set(cacheEntryKey, currentCache);
     }