Merge pull request #810 from callstack-internal/perf/memoize-useonyx-shallowequal

Memoize shallowEqual verdicts by object identity in useOnyx

Author: youssef-lr

lib/memoizedShallowEqual.ts

@@ -0,0 +1,41 @@
+import {shallowEqual} from 'fast-equals';
+
+/**
+ * Memoizes shallowEqual verdicts by the identity of the compared objects. Onyx values are
+ * treated as immutable (merge/set replace objects, never mutate), so a (prev, next) reference
+ * pair always yields the same verdict. In the hot case — N no-selector hooks on the same big
+ * key — every hook compares the exact same two cache-owned objects, so the first hook pays for
+ * the O(keys) walk and the rest resolve in O(1). WeakMap keys make stale entries impossible to
+ * read (lookup requires holding both exact objects) and let GC reclaim them.
+ */
+const shallowEqualVerdicts = new WeakMap<object, WeakMap<object, boolean>>();
+
+/**
+ * Identity-pair-memoized shallowEqual: same (a, b) references → cached verdict, no walk.
+ */
+function memoizedShallowEqual(a: unknown, b: unknown): boolean {
+    // Only object pairs are memoizable (WeakMap keys) — anything else is O(1) to compare anyway.
+    if (typeof a !== 'object' || a === null || typeof b !== 'object' || b === null) {
+        return shallowEqual(a, b);
+    }
+
+    let verdictsForA = shallowEqualVerdicts.get(a);
+
+    if (!verdictsForA) {
+        verdictsForA = new WeakMap();
+        shallowEqualVerdicts.set(a, verdictsForA);
+    }
+
+    const cachedVerdict = verdictsForA.get(b);
+
+    if (cachedVerdict !== undefined) {
+        return cachedVerdict;
+    }
+
+    const verdict = shallowEqual(a, b);
+    verdictsForA.set(b, verdict);
+
+    return verdict;
+}
+
+export default memoizedShallowEqual;

lib/useOnyx.ts

@@ -7,6 +7,7 @@ import connectionManager from './OnyxConnectionManager';
 import OnyxUtils from './OnyxUtils';
 import type {CollectionKeyBase, OnyxKey, OnyxValue} from './types';
 import onyxSnapshotCache from './OnyxSnapshotCache';
+import memoizedShallowEqual from './memoizedShallowEqual';
 import useLiveRef from './useLiveRef';
 
 type UseOnyxSelector<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>> = (data: OnyxValue<TKey> | undefined) => TReturnValue;
@@ -201,9 +202,11 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(
 
         // shallowEqual checks === first (O(1) for frozen snapshots and stable selector references),
         // then falls back to comparing top-level properties for individual keys that may have
-        // new references with equivalent content.
+        // new references with equivalent content. The comparison is memoized by object identity
+        // (see `memoizedShallowEqual`) so N hooks comparing the same two cache objects pay for
+        // one walk in total instead of one walk each.
         // Normalize null to undefined to ensure consistent comparison (both represent "no value").
-        const areValuesEqual = shallowEqual(previousValueRef.current ?? undefined, newValueRef.current ?? undefined);
+        const areValuesEqual = memoizedShallowEqual(previousValueRef.current ?? undefined, newValueRef.current ?? undefined);
 
         // We update the cached value and the result in the following conditions:
         // We will update the cached value and the result in any of the following situations:

tests/unit/memoizedShallowEqualTest.ts

@@ -0,0 +1,76 @@
+import memoizedShallowEqual from '../../lib/memoizedShallowEqual';
+
+describe('memoizedShallowEqual', () => {
+    describe('shallowEqual semantics', () => {
+        it('should return true for the same reference', () => {
+            const obj = {a: 1};
+            expect(memoizedShallowEqual(obj, obj)).toBe(true);
+        });
+
+        it('should return true for different references with shallowly-equal content', () => {
+            const member = {name: 'John'};
+            expect(memoizedShallowEqual({a: 1, member}, {a: 1, member})).toBe(true);
+        });
+
+        it('should return false when a top-level value differs', () => {
+            expect(memoizedShallowEqual({a: 1}, {a: 2})).toBe(false);
+        });
+
+        it('should return false when key counts differ', () => {
+            expect(memoizedShallowEqual({a: 1}, {a: 1, b: 2})).toBe(false);
+        });
+
+        it('should return false for equal deep content with different nested references', () => {
+            // Shallow, not deep: nested objects are compared by reference.
+            expect(memoizedShallowEqual({member: {name: 'John'}}, {member: {name: 'John'}})).toBe(false);
+        });
+
+        it('should handle non-object inputs', () => {
+            expect(memoizedShallowEqual(undefined, undefined)).toBe(true);
+            expect(memoizedShallowEqual(undefined, {})).toBe(false);
+            expect(memoizedShallowEqual('a', 'a')).toBe(true);
+            expect(memoizedShallowEqual('a', 'b')).toBe(false);
+            expect(memoizedShallowEqual(1, 1)).toBe(true);
+            expect(memoizedShallowEqual(NaN, NaN)).toBe(true);
+        });
+
+        it('should handle arrays', () => {
+            expect(memoizedShallowEqual([1, 2], [1, 2])).toBe(true);
+            expect(memoizedShallowEqual([1, 2], [1, 3])).toBe(false);
+        });
+    });
+
+    describe('memoization', () => {
+        it('should return the cached verdict for the same object pair without re-comparing', () => {
+            const a = {name: 'John'};
+            const b = {name: 'Jane'};
+            expect(memoizedShallowEqual(a, b)).toBe(false);
+
+            // Mutate `b` so the objects are now content-equal. Onyx values are immutable,
+            // so the memo is expected to keep returning the verdict computed for this exact
+            // (a, b) pair — proving the second call resolved from the cache, not a re-compare.
+            b.name = 'John';
+            expect(memoizedShallowEqual(a, b)).toBe(false);
+        });
+
+        it('should cache verdicts per pair, not per object', () => {
+            const a = {x: 1};
+            const equalToA = {x: 1};
+            const differentFromA = {x: 2};
+
+            expect(memoizedShallowEqual(a, equalToA)).toBe(true);
+            expect(memoizedShallowEqual(a, differentFromA)).toBe(false);
+
+            // Both verdicts are retained independently for the same `a`.
+            expect(memoizedShallowEqual(a, equalToA)).toBe(true);
+            expect(memoizedShallowEqual(a, differentFromA)).toBe(false);
+        });
+
+        it('should not memoize non-object inputs', () => {
+            // Primitives cannot be WeakMap keys; these calls must not throw and must compare directly.
+            expect(memoizedShallowEqual(1, {})).toBe(false);
+            expect(memoizedShallowEqual({}, 1)).toBe(false);
+            expect(memoizedShallowEqual(null, null)).toBe(true);
+        });
+    });
+});