perf: speed up cache read in useOnyx

lib/OnyxCache.ts

@@ -31,6 +31,9 @@ class OnyxCache {
     /** A map of cached values */
     private storageMap: Record<OnyxKey, OnyxValue<OnyxKey>>;
 
+    /** Index mapping collection keys to their member keys for O(1) lookup */
+    private collectionIndex: Record<OnyxKey, Set<OnyxKey>>;
+
     /**
      * Captured pending tasks for already running storage methods
      * Using a map yields better performance on operations such a delete
@@ -49,11 +52,15 @@ class OnyxCache {
     /** List of keys that have been directly subscribed to or recently modified from least to most recent */
     private recentlyAccessedKeys = new Set<OnyxKey>();
 
+    /** Set of collection keys for fast lookup */
+    private collectionKeys = new Set<OnyxKey>();
+
     constructor() {
         this.storageKeys = new Set();
         this.nullishStorageKeys = new Set();
         this.recentKeys = new Set();
         this.storageMap = {};
+        this.collectionIndex = {};
         this.pendingPromises = new Map();
 
         // bind all public methods to prevent problems with `this`
@@ -83,6 +90,11 @@ class OnyxCache {
             'addLastAccessedKey',
             'addEvictableKeysToRecentlyAccessedList',
             'getKeyForEviction',
+            'setCollectionKeys',
+            'isCollectionKey',
+            'getCollectionKey',
+            'getCollectionData',
+            'getCollectionMemberKeys',
         );
     }
 
@@ -156,19 +168,45 @@ class OnyxCache {
         // since it will either be set to a non nullish value or removed from the cache completely.
         this.nullishStorageKeys.delete(key);
 
+        const collectionKey = this.getCollectionKey(key);
         if (value === null || value === undefined) {
             delete this.storageMap[key];
+
+            // Remove from collection index if it's a collection member
+            if (collectionKey && this.getCollectionMemberKeys(collectionKey)) {
+                this.collectionIndex[collectionKey].delete(key);
+            }
             return undefined;
         }
 
         this.storageMap[key] = value;
 
+        // Update collection index if this is a collection member
+        if (collectionKey) {
+            if (!this.getCollectionMemberKeys(collectionKey)) {
+                this.collectionIndex[collectionKey] = new Set();
+            }
+            this.collectionIndex[collectionKey].add(key);
+        }
+
         return value;
     }
 
     /** Forget the cached value for the given key */
     drop(key: OnyxKey): void {
         delete this.storageMap[key];
+
+        // Update collection index if this is a collection member
+        const collectionKey = this.getCollectionKey(key);
+        if (collectionKey && this.getCollectionMemberKeys(collectionKey)) {
+            this.collectionIndex[collectionKey].delete(key);
+        }
+
+        // If this is a collection key, clear its index
+        if (this.isCollectionKey(key)) {
+            delete this.collectionIndex[key];
+        }
+
         this.storageKeys.delete(key);
         this.recentKeys.delete(key);
     }
@@ -188,10 +226,25 @@ class OnyxCache {
             this.addKey(key);
             this.addToAccessedKeys(key);
 
+            const collectionKey = this.getCollectionKey(key);
+
             if (value === null || value === undefined) {
                 this.addNullishStorageKey(key);
+
+                // Remove from collection index if it's a collection member
+                if (collectionKey && this.getCollectionMemberKeys(collectionKey)) {
+                    this.collectionIndex[collectionKey].delete(key);
+                }
             } else {
                 this.nullishStorageKeys.delete(key);
+
+                // Update collection index if this is a collection member
+                if (collectionKey) {
+                    if (!this.getCollectionMemberKeys(collectionKey)) {
+                        this.collectionIndex[collectionKey] = new Set();
+                    }
+                    this.collectionIndex[collectionKey].add(key);
+                }
             }
         });
     }
@@ -261,6 +314,12 @@ class OnyxCache {
 
         for (const key of keysToRemove) {
             delete this.storageMap[key];
+
+            // Update collection index if this is a collection member
+            const collectionKey = this.getCollectionKey(key);
+            if (collectionKey && this.getCollectionMemberKeys(collectionKey)) {
+                this.collectionIndex[collectionKey].delete(key);
+            }
             this.recentKeys.delete(key);
         }
     }
@@ -272,7 +331,8 @@ class OnyxCache {
 
     /** Check if the value has changed */
     hasValueChanged(key: OnyxKey, value: OnyxValue<OnyxKey>): boolean {
-        return !deepEqual(this.storageMap[key], value);
+        const currentValue = this.get(key, false);
+        return !deepEqual(currentValue, value);
     }
 
     /**
@@ -363,6 +423,67 @@ class OnyxCache {
         }
         return undefined;
     }
+
+    /**
+     * Set the collection keys for optimized storage
+     */
+    setCollectionKeys(collectionKeys: Set<OnyxKey>): void {
+        this.collectionKeys = collectionKeys;
+
+        // Initialize collection indexes for existing collection keys
+        collectionKeys.forEach((collectionKey) => {
+            if (this.getCollectionMemberKeys(collectionKey)) {
+                return;
+            }
+            this.collectionIndex[collectionKey] = new Set();
+        });
+    }
+
+    /**
+     * Check if a key is a collection key
+     */
+    isCollectionKey(key: OnyxKey): boolean {
+        return this.collectionKeys.has(key);
+    }
+
+    /**
+     * Get the collection key for a given member key
+     */
+    getCollectionKey(key: OnyxKey): OnyxKey | null {
+        for (const collectionKey of this.collectionKeys) {
+            if (key.startsWith(collectionKey) && key.length > collectionKey.length) {
+                return collectionKey;
+            }
+        }
+        return null;
+    }
+
+    /**
+     * Get all data for a collection key
+     */
+    getCollectionData(collectionKey: OnyxKey): Record<OnyxKey, OnyxValue<OnyxKey>> | undefined {
+        const memberKeys = this.getCollectionMemberKeys(collectionKey);
+        if (!memberKeys || memberKeys.size === 0) {
+            return undefined;
+        }
+
+        const collectionData: Record<OnyxKey, OnyxValue<OnyxKey>> = {};
+        memberKeys.forEach((memberKey) => {
+            const value = this.storageMap[memberKey];
+            if (value !== undefined) {
+                collectionData[memberKey] = value;
+            }
+        });
+
+        return collectionData;
+    }
+
+    /**
+     * Get all member keys for a collection key
+     */
+    getCollectionMemberKeys(collectionKey: OnyxKey): Set<OnyxKey> | undefined {
+        return this.collectionIndex[collectionKey];
+    }
 }
 
 const instance = new OnyxCache();

lib/OnyxUtils.ts

@@ -154,6 +154,9 @@ function initStoreValues(keys: DeepRecord<string, OnyxKey>, initialKeyStates: Pa
     // Let Onyx know about which keys are safe to evict
     cache.setEvictionAllowList(evictableKeys);
 
+    // Set collection keys in cache for optimized storage
+    cache.setCollectionKeys(onyxCollectionKeySet);
+
     if (typeof keys.COLLECTION === 'object' && typeof keys.COLLECTION.SNAPSHOT === 'string') {
         snapshotKey = keys.COLLECTION.SNAPSHOT;
         fullyMergedSnapshotKeys = new Set(fullyMergedSnapshotKeysParam ?? []);
@@ -527,23 +530,28 @@ function tryGetCachedValue<TKey extends OnyxKey>(key: TKey, mapping?: Partial<Ma
     let val = cache.get(key);
 
     if (isCollectionKey(key)) {
+        const collectionData = cache.getCollectionData(key);
         const allCacheKeys = cache.getAllKeys();
-
-        // It is possible we haven't loaded all keys yet so we do not know if the
-        // collection actually exists.
-        if (allCacheKeys.size === 0) {
-            return;
-        }
-
-        const values: OnyxCollection<KeyValueMapping[TKey]> = {};
-        allCacheKeys.forEach((cacheKey) => {
-            if (!cacheKey.startsWith(key)) {
+        if (collectionData !== undefined && allCacheKeys.size > 0) {
+            val = collectionData;
+        } else {
+            // Fallback to original logic
+            // It is possible we haven't loaded all keys yet so we do not know if the
+            // collection actually exists.
+            if (allCacheKeys.size === 0) {
                 return;
             }
 
-            values[cacheKey] = cache.get(cacheKey);
-        });
-        val = values;
+            const values: OnyxCollection<KeyValueMapping[TKey]> = {};
+            allCacheKeys.forEach((cacheKey) => {
+                if (!cacheKey.startsWith(key)) {
+                    return;
+                }
+
+                values[cacheKey] = cache.get(cacheKey);
+            });
+            val = values;
+        }
     }
 
     if (mapping?.selector) {
@@ -558,7 +566,28 @@ function tryGetCachedValue<TKey extends OnyxKey>(key: TKey, mapping?: Partial<Ma
 }
 
 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);
     const allKeys = collectionMemberKeys || cache.getAllKeys();
+    if (collectionData !== undefined && (Array.isArray(allKeys) ? allKeys.length > 0 : allKeys.size > 0)) {
+        // If we have specific member keys, filter the collection
+        if (collectionMemberKeys) {
+            const filteredCollection: OnyxCollection<KeyValueMapping[TKey]> = {};
+            collectionMemberKeys.forEach((key) => {
+                if (collectionData[key] !== undefined) {
+                    filteredCollection[key] = collectionData[key];
+                } else if (cache.hasNullishStorageKey(key)) {
+                    filteredCollection[key] = cache.get(key);
+                }
+            });
+            return filteredCollection;
+        }
+
+        // Return a copy to avoid mutations affecting the cache
+        return {...collectionData};
+    }
+
+    // Fallback to original implementation if collection data not available
     const collection: OnyxCollection<KeyValueMapping[TKey]> = {};
 
     // forEach exists on both Set and Array