perf: preserve references of unchanged collection items

lib/Onyx.ts

@@ -701,9 +701,10 @@ function setCollection<TKey extends CollectionKeyBase, TMap>(collectionKey: TKey
         const keyValuePairs = OnyxUtils.prepareKeyValuePairsForStorage(mutableCollection, true, undefined, true);
         const previousCollection = OnyxUtils.getCachedCollection(collectionKey);
 
-        keyValuePairs.forEach(([key, value]) => cache.set(key, value));
+        // Preserve references for unchanged items in setCollection
+        const preservedCollection = OnyxUtils.preserveCollectionReferences(keyValuePairs);
 
-        const updatePromise = OnyxUtils.scheduleNotifyCollectionSubscribers(collectionKey, mutableCollection, previousCollection);
+        const updatePromise = OnyxUtils.scheduleNotifyCollectionSubscribers(collectionKey, preservedCollection, previousCollection);
 
         return Storage.multiSet(keyValuePairs)
             .catch((error) => OnyxUtils.evictStorageAndRetry(error, setCollection, collectionKey, collection))

lib/OnyxUtils.ts

@@ -544,6 +544,60 @@ function tryGetCachedValue<TKey extends OnyxKey>(key: TKey): OnyxValue<OnyxKey>
     return val;
 }
 
+/**
+ * Utility function to preserve object references for unchanged items in collection operations.
+ * Compares new values with cached values using deep equality and preserves references when data is identical.
+ * @returns The preserved collection with unchanged references maintained
+ */
+function preserveCollectionReferences(keyValuePairs: StorageKeyValuePair[]): OnyxInputKeyValueMapping {
+    const preservedCollection: OnyxInputKeyValueMapping = {};
+
+    keyValuePairs.forEach(([key, value]) => {
+        const cachedValue = cache.get(key, false);
+
+        // If no cached value exists, we need to add the new value (skip expensive deep equality check)
+        // Use deep equality check to preserve references for unchanged items
+        if (cachedValue !== undefined && deepEqual(value, cachedValue)) {
+            // Keep the existing reference
+            preservedCollection[key] = cachedValue;
+        } else {
+            // Update cache only for changed items
+            cache.set(key, value);
+            preservedCollection[key] = value;
+        }
+    });
+
+    return preservedCollection;
+}
+
+/**
+ * Utility function for merge operations that preserves references after cache merge has been performed.
+ * Compares merged values with original cached values and preserves references when data is unchanged.
+ * @returns The preserved collection with unchanged references maintained
+ */
+function preserveCollectionReferencesAfterMerge(
+    collection: Record<string, OnyxValue<OnyxKey>>,
+    originalCachedValues: Record<string, OnyxValue<OnyxKey>>,
+): Record<string, OnyxValue<OnyxKey>> {
+    const preservedCollection: Record<string, OnyxValue<OnyxKey>> = {};
+
+    Object.keys(collection).forEach((key) => {
+        const newMergedValue = cache.get(key, false);
+        const originalValue = originalCachedValues[key];
+
+        // Use deep equality check to preserve references for unchanged items
+        if (originalValue !== undefined && deepEqual(newMergedValue, originalValue)) {
+            // Keep the existing reference and update cache
+            preservedCollection[key] = originalValue;
+            cache.set(key, originalValue);
+        } else {
+            preservedCollection[key] = newMergedValue;
+        }
+    });
+
+    return preservedCollection;
+}
+
 function getCachedCollection<TKey extends CollectionKeyBase>(collectionKey: TKey, collectionMemberKeys?: string[]): NonNullable<OnyxCollection<KeyValueMapping[TKey]>> {
     // Use optimized collection data retrieval when cache is populated
     const collectionData = cache.getCollectionData(collectionKey);
@@ -1343,10 +1397,21 @@ function mergeCollectionWithPatches<TKey extends CollectionKeyBase, TMap>(
             const finalMergedCollection = {...existingKeyCollection, ...newCollection};
 
             // Prefill cache if necessary by calling get() on any existing keys and then merge original data to cache
-            // and update all subscribers
+            // and update all subscribers with reference preservation for unchanged items
             const promiseUpdate = previousCollectionPromise.then((previousCollection) => {
+                // Capture the original cached values before merging
+                const originalCachedValues: Record<string, OnyxValue<OnyxKey>> = {};
+                Object.keys(finalMergedCollection).forEach((key) => {
+                    originalCachedValues[key] = cache.get(key, false);
+                });
+
+                // Then merge all the data into cache as normal
                 cache.merge(finalMergedCollection);
-                return scheduleNotifyCollectionSubscribers(collectionKey, finalMergedCollection, previousCollection);
+
+                // Finally, preserve references for items that didn't actually change
+                const preservedCollection = preserveCollectionReferencesAfterMerge(finalMergedCollection, originalCachedValues);
+
+                return scheduleNotifyCollectionSubscribers(collectionKey, preservedCollection, previousCollection);
             });
 
             return Promise.all(promises)
@@ -1400,9 +1465,10 @@ function partialSetCollection<TKey extends CollectionKeyBase, TMap>(collectionKe
         const previousCollection = getCachedCollection(collectionKey, existingKeys);
         const keyValuePairs = prepareKeyValuePairsForStorage(mutableCollection, true, undefined, true);
 
-        keyValuePairs.forEach(([key, value]) => cache.set(key, value));
+        // Preserve references for unchanged items in partialSetCollection
+        const preservedCollection = preserveCollectionReferences(keyValuePairs);
 
-        const updatePromise = scheduleNotifyCollectionSubscribers(collectionKey, mutableCollection, previousCollection);
+        const updatePromise = scheduleNotifyCollectionSubscribers(collectionKey, preservedCollection, previousCollection);
 
         return Storage.multiSet(keyValuePairs)
             .catch((error) => evictStorageAndRetry(error, partialSetCollection, collectionKey, collection))
@@ -1484,6 +1550,8 @@ const OnyxUtils = {
     updateSnapshots,
     mergeCollectionWithPatches,
     partialSetCollection,
+    preserveCollectionReferences,
+    preserveCollectionReferencesAfterMerge,
     logKeyChanged,
     logKeyRemoved,
 };

tests/unit/onyxTest.ts

@@ -2562,6 +2562,59 @@ describe('Onyx', () => {
                     });
                 });
         });
+
+        it('should preserve object references for unchanged items when merging collections', async () => {
+            const item1 = {id: 1, name: 'Item 1'};
+            const item2 = {id: 2, name: 'Item 2'};
+
+            // Set initial data using separate multiSet calls like existing tests
+            await Onyx.multiSet({[`${ONYX_KEYS.COLLECTION.TEST_KEY}1`]: item1});
+            await Onyx.multiSet({[`${ONYX_KEYS.COLLECTION.TEST_KEY}2`]: item2});
+
+            // Get references to the cached objects
+            const cachedItem1 = cache.get(`${ONYX_KEYS.COLLECTION.TEST_KEY}1`);
+            const cachedItem2 = cache.get(`${ONYX_KEYS.COLLECTION.TEST_KEY}2`);
+
+            // Merge collection with same data for item1 and changed data for item2
+            await Onyx.mergeCollection(ONYX_KEYS.COLLECTION.TEST_KEY, {
+                [`${ONYX_KEYS.COLLECTION.TEST_KEY}1`]: {id: 1, name: 'Item 1'}, // Same data
+                [`${ONYX_KEYS.COLLECTION.TEST_KEY}2`]: {id: 2, name: 'Item 2 Updated'}, // Changed data
+            } as GenericCollection);
+
+            // Check that unchanged item keeps its reference
+            const newCachedItem1 = cache.get(`${ONYX_KEYS.COLLECTION.TEST_KEY}1`);
+            const newCachedItem2 = cache.get(`${ONYX_KEYS.COLLECTION.TEST_KEY}2`);
+
+            expect(newCachedItem1).toBe(cachedItem1); // Same reference
+            expect(newCachedItem2).not.toBe(cachedItem2); // Different reference
+            expect(newCachedItem2).toEqual({id: 2, name: 'Item 2 Updated'}); // Correct data
+        });
+
+        it('should preserve all references when no data changes', async () => {
+            const item1 = {id: 1, name: 'Item 1', nested: {value: 'test'}};
+            const item2 = {id: 2, name: 'Item 2', nested: {value: 'test2'}};
+
+            // Set initial data using separate calls
+            await Onyx.multiSet({[`${ONYX_KEYS.COLLECTION.TEST_KEY}1`]: item1});
+            await Onyx.multiSet({[`${ONYX_KEYS.COLLECTION.TEST_KEY}2`]: item2});
+
+            // Get references to the cached objects
+            const cachedItem1 = cache.get(`${ONYX_KEYS.COLLECTION.TEST_KEY}1`);
+            const cachedItem2 = cache.get(`${ONYX_KEYS.COLLECTION.TEST_KEY}2`);
+
+            // Merge collection with identical data
+            await Onyx.mergeCollection(ONYX_KEYS.COLLECTION.TEST_KEY, {
+                [`${ONYX_KEYS.COLLECTION.TEST_KEY}1`]: {id: 1, name: 'Item 1', nested: {value: 'test'}},
+                [`${ONYX_KEYS.COLLECTION.TEST_KEY}2`]: {id: 2, name: 'Item 2', nested: {value: 'test2'}},
+            } as GenericCollection);
+
+            // Both items should keep their references since data is identical
+            const newCachedItem1 = cache.get(`${ONYX_KEYS.COLLECTION.TEST_KEY}1`);
+            const newCachedItem2 = cache.get(`${ONYX_KEYS.COLLECTION.TEST_KEY}2`);
+
+            expect(newCachedItem1).toBe(cachedItem1); // Same reference
+            expect(newCachedItem2).toBe(cachedItem2); // Same reference
+        });
     });
 
     describe('set', () => {
@@ -2703,6 +2756,35 @@ describe('Onyx', () => {
                 [routeB]: {name: 'Route B'},
             });
         });
+
+        it('should preserve object references for unchanged items when setting collections', async () => {
+            const item1 = {id: 1, name: 'Item 1'};
+            const item2 = {id: 2, name: 'Item 2'};
+
+            // Set initial data
+            await Onyx.setCollection(ONYX_KEYS.COLLECTION.TEST_KEY, {
+                [`${ONYX_KEYS.COLLECTION.TEST_KEY}1`]: item1,
+                [`${ONYX_KEYS.COLLECTION.TEST_KEY}2`]: item2,
+            } as GenericCollection);
+
+            // Get references to the cached objects
+            const cachedItem1 = cache.get(`${ONYX_KEYS.COLLECTION.TEST_KEY}1`);
+            const cachedItem2 = cache.get(`${ONYX_KEYS.COLLECTION.TEST_KEY}2`);
+
+            // Set collection with same data for item1 and changed data for item2
+            await Onyx.setCollection(ONYX_KEYS.COLLECTION.TEST_KEY, {
+                [`${ONYX_KEYS.COLLECTION.TEST_KEY}1`]: {id: 1, name: 'Item 1'}, // Same data
+                [`${ONYX_KEYS.COLLECTION.TEST_KEY}2`]: {id: 2, name: 'Item 2 Updated'}, // Changed data
+            } as GenericCollection);
+
+            // Check that unchanged item keeps its reference
+            const newCachedItem1 = cache.get(`${ONYX_KEYS.COLLECTION.TEST_KEY}1`);
+            const newCachedItem2 = cache.get(`${ONYX_KEYS.COLLECTION.TEST_KEY}2`);
+
+            expect(newCachedItem1).toBe(cachedItem1); // Same reference
+            expect(newCachedItem2).not.toBe(cachedItem2); // Different reference
+            expect(newCachedItem2).toEqual({id: 2, name: 'Item 2 Updated'}); // Correct data
+        });
     });
 
     describe('skippable collection member ids', () => {