Feat: Call useOnyx selector for entire collection

lib/useOnyx.ts

@@ -1,11 +1,10 @@
 import {deepEqual, shallowEqual} from 'fast-equals';
 import {useCallback, useEffect, useRef, useSyncExternalStore} from 'react';
-import type {IsEqual} from 'type-fest';
 import OnyxCache from './OnyxCache';
 import type {Connection} from './OnyxConnectionManager';
 import connectionManager from './OnyxConnectionManager';
 import OnyxUtils from './OnyxUtils';
-import type {CollectionKeyBase, OnyxCollection, OnyxEntry, OnyxKey, OnyxValue, Selector} from './types';
+import type {CollectionKeyBase, KeyValueMapping, OnyxCollection, OnyxKey, OnyxValue} from './types';
 import useLiveRef from './useLiveRef';
 import usePrevious from './usePrevious';
 
@@ -33,43 +32,78 @@ type UseOnyxInitialValueOption<TInitialValue> = {
     initialValue?: TInitialValue;
 };
 
+type UseOnyxSelector<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>> = (data: OnyxValue<TKey> | undefined) => TReturnValue;
+
 type UseOnyxSelectorOption<TKey extends OnyxKey, TReturnValue> = {
     /**
      * This will be used to subscribe to a subset of an Onyx key's data.
      * Using this setting on `useOnyx` can have very positive performance benefits because the component will only re-render
      * when the subset of data changes. Otherwise, any change of data on any property would normally
      * cause the component to re-render (and that can be expensive from a performance standpoint).
+     * @see `useOnyx` cannot return `null` and so selector will replace `null` with `undefined` to maintain compatibility.
      */
-    selector?: Selector<TKey, unknown, TReturnValue>;
+    selector?: UseOnyxSelector<TKey, TReturnValue>;
 };
 
 type UseOnyxOptions<TKey extends OnyxKey, TReturnValue> = BaseUseOnyxOptions & UseOnyxInitialValueOption<TReturnValue> & UseOnyxSelectorOption<TKey, TReturnValue>;
 
 type FetchStatus = 'loading' | 'loaded';
 
-type SelectedValue<TKey, TValue> = TKey extends CollectionKeyBase ? OnyxCollection<TValue> : OnyxEntry<TValue>;
-
-type CachedValue<TKey extends OnyxKey, TValue> = IsEqual<TValue, OnyxValue<TKey>> extends true ? TValue : SelectedValue<TKey, TValue>;
-
 type ResultMetadata = {
     status: FetchStatus;
 };
 
-type UseOnyxResult<TKey extends OnyxKey, TValue> = [CachedValue<TKey, TValue>, ResultMetadata];
+type UseOnyxResult<TValue> = [NonNullable<TValue> | undefined, ResultMetadata];
+
+/**
+ * Gets the cached value from the Onyx cache. If the key is a collection key, it will return all the values in the collection.
+ * It is a fork of `tryGetCachedValue` from `OnyxUtils` caused by different selector logic in `useOnyx`. It should be unified in the future, when `withOnyx` is removed.
+ */
+function tryGetCachedValue<TKey extends OnyxKey>(key: TKey): OnyxValue<OnyxKey> {
+    if (!OnyxUtils.isCollectionKey(key)) {
+        return OnyxCache.get(key);
+    }
+
+    const allCacheKeys = OnyxCache.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)) {
+            return;
+        }
+
+        values[cacheKey] = OnyxCache.get(cacheKey);
+    });
+
+    return values;
+}
+
+/**
+ * Gets the value from cache and maps it with selector. It changes `null` to `undefined` for `useOnyx` compatibility.
+ */
+function getCachedValue<TKey extends OnyxKey, TValue>(key: TKey, selector?: UseOnyxSelector<TKey, TValue>) {
+    const value = tryGetCachedValue(key) as OnyxValue<TKey>;
+
+    const selectedValue = selector ? selector(value) : (value as TValue);
 
-function getCachedValue<TKey extends OnyxKey, TValue>(key: TKey, selector?: Selector<TKey, unknown, unknown>): CachedValue<TKey, TValue> | undefined {
-    return (OnyxUtils.tryGetCachedValue(key, {selector}) ?? undefined) as CachedValue<TKey, TValue> | undefined;
+    return selectedValue ?? undefined;
 }
 
 function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(
     key: TKey,
     options?: BaseUseOnyxOptions & UseOnyxInitialValueOption<TReturnValue> & Required<UseOnyxSelectorOption<TKey, TReturnValue>>,
-): UseOnyxResult<TKey, TReturnValue>;
+): UseOnyxResult<TReturnValue>;
 function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(
     key: TKey,
     options?: BaseUseOnyxOptions & UseOnyxInitialValueOption<NoInfer<TReturnValue>>,
-): UseOnyxResult<TKey, TReturnValue>;
-function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey, options?: UseOnyxOptions<TKey, TReturnValue>): UseOnyxResult<TKey, TReturnValue> {
+): UseOnyxResult<TReturnValue>;
+function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey, options?: UseOnyxOptions<TKey, TReturnValue>): UseOnyxResult<TReturnValue> {
     const connectionRef = useRef<Connection | null>(null);
     const previousKey = usePrevious(key);
 
@@ -78,16 +112,16 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey
 
     // Stores the previous cached value as it's necessary to compare with the new value in `getSnapshot()`.
     // We initialize it to `null` to simulate that we don't have any value from cache yet.
-    const previousValueRef = useRef<CachedValue<TKey, TReturnValue> | undefined | null>(null);
+    const previousValueRef = useRef<TReturnValue | undefined | null>(null);
 
     // Stores the newest cached value in order to compare with the previous one and optimize `getSnapshot()` execution.
-    const newValueRef = useRef<CachedValue<TKey, TReturnValue> | undefined | null>(null);
+    const newValueRef = useRef<TReturnValue | undefined | null>(null);
 
     // Stores the previously result returned by the hook, containing the data from cache and the fetch status.
     // We initialize it to `undefined` and `loading` fetch status to simulate the initial result when the hook is loading from the cache.
     // However, if `initWithStoredValues` is `false` we set the fetch status to `loaded` since we want to signal that data is ready.
-    const resultRef = useRef<UseOnyxResult<TKey, TReturnValue>>([
-        undefined as CachedValue<TKey, TReturnValue>,
+    const resultRef = useRef<UseOnyxResult<TReturnValue>>([
+        undefined,
         {
             status: options?.initWithStoredValues === false ? 'loaded' : 'loading',
         },
@@ -159,7 +193,7 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey
             // If `newValueRef.current` is `null` or any other value it means that the cache does have a value for that key.
             // This difference between `undefined` and other values is crucial and it's used to address the following
             // conditions and use cases.
-            newValueRef.current = getCachedValue<TKey, TReturnValue>(key, selectorRef.current);
+            newValueRef.current = getCachedValue(key, selectorRef.current);
 
             // We set this flag to `false` again since we don't want to get the newest cached value every time `getSnapshot()` is executed,
             // and only when `Onyx.connect()` callback is fired.
@@ -182,7 +216,7 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey
         // If data is not present in cache and `initialValue` is set during the first connection,
         // we set the new value to `initialValue` and fetch status to `loaded` since we already have some data to return to the consumer.
         if (isFirstConnectionRef.current && !hasCacheForKey && options?.initialValue !== undefined) {
-            newValueRef.current = (options?.initialValue ?? undefined) as CachedValue<TKey, TReturnValue>;
+            newValueRef.current = (options?.initialValue ?? undefined) as TReturnValue;
             newFetchStatus = 'loaded';
         }
 
@@ -204,11 +238,11 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey
             previousValueRef.current = newValueRef.current;
 
             // If the new value is `null` we default it to `undefined` to ensure the consumer gets a consistent result from the hook.
-            resultRef.current = [previousValueRef.current as CachedValue<TKey, TReturnValue>, {status: newFetchStatus ?? 'loaded'}];
+            resultRef.current = [previousValueRef.current ?? undefined, {status: newFetchStatus ?? 'loaded'}];
         }
 
         return resultRef.current;
-    }, [key, selectorRef, options?.initWithStoredValues, options?.allowStaleData, options?.initialValue]);
+    }, [options?.initWithStoredValues, options?.allowStaleData, options?.initialValue, key, selectorRef]);
 
     const subscribe = useCallback(
         (onStoreChange: () => void) => {
@@ -247,7 +281,7 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey
         checkEvictableKey();
     }, [checkEvictableKey]);
 
-    const result = useSyncExternalStore<UseOnyxResult<TKey, TReturnValue>>(subscribe, getSnapshot);
+    const result = useSyncExternalStore<UseOnyxResult<TReturnValue>>(subscribe, getSnapshot);
 
     return result;
 }

tests/unit/useOnyxTest.ts

@@ -1,5 +1,5 @@
 import {act, renderHook} from '@testing-library/react-native';
-import type {OnyxEntry} from '../../lib';
+import type {OnyxCollection, OnyxEntry} from '../../lib';
 import Onyx, {useOnyx} from '../../lib';
 import OnyxUtils from '../../lib/OnyxUtils';
 import StorageMock from '../../lib/storage';
@@ -176,7 +176,11 @@ describe('useOnyx', () => {
             const {result} = renderHook(() =>
                 useOnyx(ONYXKEYS.COLLECTION.TEST_KEY, {
                     // @ts-expect-error bypass
-                    selector: (entry: OnyxEntry<{id: string; name: string}>) => entry?.id,
+                    selector: (entries: OnyxCollection<{id: string; name: string}>) =>
+                        Object.entries(entries ?? {}).reduce<NonNullable<OnyxCollection<string>>>((acc, [key, value]) => {
+                            acc[key] = value?.id;
+                            return acc;
+                        }, {}),
                 }),
             );