-
Notifications
You must be signed in to change notification settings - Fork 1.5k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add weakMapMemoize utility. (#26605)
## Summary & Motivation Introduces a memoize function thats more powerful than what lodash offers. Lodash memoize only works with string keys, and so it requires you to take your arguments and somehow serialize it to a string. Depending on the object this can be very slow if the object is big. To avoid the headache of serialization I'm adding a weakMapMemoize which is much more flexible and can cache any function mixing both primitives and objects. Going to use this to memoize the `filterByQuery` helpers for the GanttChart since it looks like we call the utility in a lot of different places separately and having them coordinate so that there's only 1 call is a bit tricky. ## How I Tested These Changes jest tests
- Loading branch information
Showing
2 changed files
with
484 additions
and
0 deletions.
There are no files selected for viewing
327 changes: 327 additions & 0 deletions
327
js_modules/dagster-ui/packages/ui-core/src/util/__tests__/weakMapMemoize.test.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,327 @@ | ||
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, {maxEntries: 3}); | ||
|
||
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 | ||
}); | ||
|
||
// 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, {maxEntries: 2}); | ||
|
||
const obj1 = {x: 10}; | ||
const obj2 = {x: 20}; | ||
|
||
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 | ||
}); | ||
|
||
// 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, {maxEntries: 3}); | ||
|
||
const obj1 = {y: 100}; | ||
const obj2 = {y: 200}; | ||
|
||
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 | ||
}); | ||
|
||
// 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, {maxEntries: 1}); | ||
|
||
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, {maxEntries: 2}); | ||
|
||
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 | ||
}); | ||
|
||
// 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, {maxEntries: 3}); | ||
|
||
const func1 = (x: number) => x * 2; | ||
const func2 = (x: number) => x * 3; | ||
|
||
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 | ||
}); | ||
|
||
// 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, {maxEntries: 4}); | ||
|
||
const obj1 = {key: 'value1'}; | ||
const obj2 = {key: 'value2'}; | ||
|
||
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 | ||
}); | ||
|
||
// 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, {maxEntries: 3}); | ||
|
||
const array1 = [1, 2, 3]; | ||
const array2 = [4, 5, 6]; | ||
|
||
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 | ||
}); | ||
|
||
// 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, {maxEntries: 4}); | ||
|
||
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 | ||
}); | ||
|
||
// 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, {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]; | ||
|
||
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) | ||
}); | ||
|
||
// 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, {maxEntries: 7}); | ||
|
||
const object1 = {a: 5}; | ||
const object2 = {b: 10}; | ||
const object3 = {c: 15}; | ||
const object4 = {a: 30}; | ||
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 | ||
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 | ||
expect(spy).toHaveBeenCalledTimes(9); | ||
}); | ||
|
||
// 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}); | ||
|
||
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) | ||
expect(spy).toHaveBeenCalledTimes(4); // Calls for 1,2,3,4 | ||
|
||
// Accessing 1 again should trigger a new call since it was evicted | ||
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); | ||
|
||
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); | ||
}); | ||
|
||
// 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); | ||
}); | ||
}); |
Oops, something went wrong.
d5d70ca
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Deploy preview for dagit-core-storybook ready!
✅ Preview
https://dagit-core-storybook-iq34k15py-elementl.vercel.app
Built with commit d5d70ca.
This pull request is being automatically deployed with vercel-action