Refactor OnyxUpdate type to support generics for improved performance

lib/Onyx.ts

@@ -396,7 +396,7 @@ function clear(keysToPreserve: OnyxKey[] = []): Promise<void> {
  * @param data An array of objects with update expressions
  * @returns resolves when all operations are complete
  */
-function update(data: OnyxUpdate[]): Promise<void> {
+function update<TKey extends OnyxKey>(data: Array<OnyxUpdate<TKey>>): Promise<void> {
     // First, validate the Onyx object is in the format we expect
     data.forEach(({onyxMethod, key, value}) => {
         if (!Object.values(OnyxUtils.METHOD).includes(onyxMethod)) {
@@ -453,7 +453,7 @@ function update(data: OnyxUpdate[]): Promise<void> {
                     collectionKeys.forEach((collectionKey) => enqueueMergeOperation(collectionKey, mergedCollection[collectionKey]));
                 }
             },
-            [OnyxUtils.METHOD.SET_COLLECTION]: (k, v) => promises.push(() => setCollection(k, v as OnyxSetCollectionInput<OnyxKey>)),
+            [OnyxUtils.METHOD.SET_COLLECTION]: (k, v) => promises.push(() => setCollection(k as TKey, v as OnyxSetCollectionInput<TKey>)),
             [OnyxUtils.METHOD.MULTI_SET]: (k, v) => Object.entries(v as Partial<OnyxInputKeyValueMapping>).forEach(([entryKey, entryValue]) => enqueueSetOperation(entryKey, entryValue)),
             [OnyxUtils.METHOD.CLEAR]: () => {
                 clearPromise = clear();

lib/OnyxUtils.ts

@@ -1209,7 +1209,7 @@ function unsubscribeFromKey(subscriptionID: number): void {
     delete callbackToStateMapping[subscriptionID];
 }
 
-function updateSnapshots(data: OnyxUpdate[], mergeFn: typeof Onyx.merge): Array<() => Promise<void>> {
+function updateSnapshots<TKey extends OnyxKey>(data: Array<OnyxUpdate<TKey>>, mergeFn: typeof Onyx.merge): Array<() => Promise<void>> {
     const snapshotCollectionKey = getSnapshotKey();
     if (!snapshotCollectionKey) return [];
 

lib/types.ts

@@ -208,7 +208,7 @@ type NullishObjectDeep<ObjectType extends object> = {
  * Also, the `TMap` type is inferred automatically in `mergeCollection()` method and represents
  * the object of collection keys/values specified in the second parameter of the method.
  */
-type Collection<TKey extends CollectionKeyBase, TValue> = Record<`${TKey}${string}`, TValue> & {[P in TKey]?: never};
+type Collection<TKey extends CollectionKeyBase, TValue> = Record<`${TKey}${string}`, TValue>;
 
 /** Represents the base options used in `Onyx.connect()` method. */
 // NOTE: Any changes to this type like adding or removing options must be accounted in OnyxConnectionManager's `generateConnectionID()` method!
@@ -328,49 +328,23 @@ type OnyxSetCollectionInput<TKey extends OnyxKey> = Collection<TKey, OnyxInput<T
 
 type OnyxMethodMap = typeof OnyxUtils.METHOD;
 
+type ExpandOnyxKeys<TKey extends OnyxKey> = TKey extends CollectionKeyBase ? NoInfer<`${TKey}${string}`> : TKey;
+
 /**
  * OnyxUpdate type includes all onyx methods used in OnyxMethodValueMap.
  * If a new method is added to OnyxUtils.METHOD constant, it must be added to OnyxMethodValueMap type.
  * Otherwise it will show static type errors.
  */
-type OnyxUpdate =
+type OnyxUpdate<TKey extends OnyxKey = OnyxKey> = {
     // ⚠️ DO NOT CHANGE THIS TYPE, UNLESS YOU KNOW WHAT YOU ARE DOING. ⚠️
-    | {
-          [TKey in OnyxKey]:
-              | {
-                    onyxMethod: typeof OnyxUtils.METHOD.SET;
-                    key: TKey;
-                    value: OnyxSetInput<TKey>;
-                }
-              | {
-                    onyxMethod: typeof OnyxUtils.METHOD.MULTI_SET;
-                    key: TKey;
-                    value: OnyxMultiSetInput;
-                }
-              | {
-                    onyxMethod: typeof OnyxUtils.METHOD.MERGE;
-                    key: TKey;
-                    value: OnyxMergeInput<TKey>;
-                }
-              | {
-                    onyxMethod: typeof OnyxUtils.METHOD.CLEAR;
-                    key: TKey;
-                    value?: undefined;
-                };
-      }[OnyxKey]
-    | {
-          [TKey in CollectionKeyBase]:
-              | {
-                    onyxMethod: typeof OnyxUtils.METHOD.MERGE_COLLECTION;
-                    key: TKey;
-                    value: OnyxMergeCollectionInput<TKey>;
-                }
-              | {
-                    onyxMethod: typeof OnyxUtils.METHOD.SET_COLLECTION;
-                    key: TKey;
-                    value: OnyxSetCollectionInput<TKey>;
-                };
-      }[CollectionKeyBase];
+    [K in TKey]:
+        | {onyxMethod: typeof OnyxUtils.METHOD.SET; key: ExpandOnyxKeys<K>; value: OnyxSetInput<K>}
+        | {onyxMethod: typeof OnyxUtils.METHOD.MULTI_SET; key: ExpandOnyxKeys<K>; value: OnyxMultiSetInput}
+        | {onyxMethod: typeof OnyxUtils.METHOD.MERGE; key: ExpandOnyxKeys<K>; value: OnyxMergeInput<K>}
+        | {onyxMethod: typeof OnyxUtils.METHOD.CLEAR; key: ExpandOnyxKeys<K>; value?: never}
+        | {onyxMethod: typeof OnyxUtils.METHOD.MERGE_COLLECTION; key: K; value: OnyxMergeCollectionInput<K>}
+        | {onyxMethod: typeof OnyxUtils.METHOD.SET_COLLECTION; key: K; value: OnyxSetCollectionInput<K>};
+}[TKey];
 
 /**
  * Represents the options used in `Onyx.set()` method.

tests/perf-test/Onyx.perf-test.ts

@@ -1,5 +1,5 @@
 import {measureAsyncFunction} from 'reassure';
-import type {OnyxUpdate} from '../../lib';
+import type {OnyxKey, OnyxUpdate} from '../../lib';
 import Onyx from '../../lib';
 import createRandomReportAction, {getRandomReportActions} from '../utils/collections/reportActions';
 import type GenericCollection from '../utils/GenericCollection';
@@ -115,13 +115,13 @@ describe('Onyx', () => {
 
             const sets = Object.entries(changedReportActions)
                 .filter(([, v]) => Number(v.reportActionID) % 2 === 0)
-                .map(([k, v]): OnyxUpdate => ({key: k, onyxMethod: Onyx.METHOD.SET, value: v}));
+                .map(([k, v]): OnyxUpdate<OnyxKey> => ({key: k, onyxMethod: Onyx.METHOD.SET, value: v}));
 
             const merges = Object.entries(changedReportActions)
                 .filter(([, v]) => Number(v.reportActionID) % 2 !== 0)
-                .map(([k, v]): OnyxUpdate => ({key: k, onyxMethod: Onyx.METHOD.MERGE, value: v}));
+                .map(([k, v]): OnyxUpdate<OnyxKey> => ({key: k, onyxMethod: Onyx.METHOD.MERGE, value: v}));
 
-            const updates = alternateLists(sets, merges) as OnyxUpdate[];
+            const updates = alternateLists(sets, merges) as Array<OnyxUpdate<OnyxKey>>;
 
             await measureAsyncFunction(() => Onyx.update(updates), {
                 beforeEach: async () => {

tests/perf-test/OnyxUtils.perf-test.ts

@@ -9,7 +9,7 @@ import OnyxUtils, {clearOnyxUtilsInternals} from '../../lib/OnyxUtils';
 import type GenericCollection from '../utils/GenericCollection';
 import type {OnyxUpdate} from '../../lib/Onyx';
 import createDeferredTask from '../../lib/createDeferredTask';
-import type {OnyxInputKeyValueMapping, RetriableOnyxOperation} from '../../lib/types';
+import type {OnyxInputKeyValueMapping, OnyxKey, RetriableOnyxOperation} from '../../lib/types';
 
 const ONYXKEYS = {
     TEST_KEY: 'test',
@@ -806,7 +806,7 @@ describe('OnyxUtils', () => {
 
     describe('updateSnapshots', () => {
         test('one call with 100 updates', async () => {
-            const updates: OnyxUpdate[] = [];
+            const updates: Array<OnyxUpdate<OnyxKey>> = [];
             for (let i = 0; i < 100; i++) {
                 updates.push({
                     onyxMethod: OnyxUtils.METHOD.MERGE,

tests/types/OnyxUpdate.ts

@@ -1,20 +1,29 @@
-import type {OnyxUpdate} from '../../dist/types';
+import Onyx from '../../dist/Onyx';
+import type {OnyxMergeCollectionInput, OnyxUpdate} from '../../dist/types';
 import ONYX_KEYS from './setup';
 
-const onyxUpdate: OnyxUpdate = {
+const onyxUpdate: OnyxUpdate<'test'> = {
     onyxMethod: 'set',
     key: ONYX_KEYS.TEST_KEY,
     value: 'string',
 };
 
-const onyxUpdateError: OnyxUpdate = {
+const onyxUpdateCollectionMember: OnyxUpdate<typeof ONYX_KEYS.COLLECTION.TEST_KEY> = {
+    onyxMethod: 'set',
+    key: `${ONYX_KEYS.COLLECTION.TEST_KEY}1`,
+    value: {
+        str: 'test1',
+    },
+};
+
+const onyxUpdateError: OnyxUpdate<'test'> = {
     onyxMethod: 'set',
     key: ONYX_KEYS.TEST_KEY,
     // @ts-expect-error TEST_KEY is a string, not a number
     value: 2,
 };
 
-const onyxUpdateCollection: OnyxUpdate = {
+const onyxUpdateCollection: OnyxUpdate<'test_'> = {
     onyxMethod: 'mergecollection',
     key: ONYX_KEYS.COLLECTION.TEST_KEY,
     value: {
@@ -28,15 +37,15 @@ const onyxUpdateCollection: OnyxUpdate = {
 };
 
 // @ts-expect-error COLLECTION.TEST_KEY is an object, not a number
-const onyxUpdateCollectionError: OnyxUpdate = {
+const onyxUpdateCollectionError: OnyxUpdate<'test_'> = {
     onyxMethod: 'mergecollection',
     key: ONYX_KEYS.COLLECTION.TEST_KEY,
     value: {
         [`${ONYX_KEYS.COLLECTION.TEST_KEY}1`]: 2,
     },
 };
 
-const onyxUpdateCollectionError2: OnyxUpdate = {
+const onyxUpdateCollectionError2: OnyxUpdate<'test_'> = {
     onyxMethod: 'mergecollection',
     key: ONYX_KEYS.COLLECTION.TEST_KEY,
     value: {
@@ -47,13 +56,86 @@ const onyxUpdateCollectionError2: OnyxUpdate = {
     },
 };
 
-// @ts-expect-error COLLECTION.TEST_KEY is invalid key, it is missing the suffix
-const onyxUpdateCollectionError3: OnyxUpdate = {
+const onyxUpdateCollectionError3: OnyxUpdate<'test_'> = {
     onyxMethod: 'mergecollection',
     key: ONYX_KEYS.COLLECTION.TEST_KEY,
     value: {
-        [ONYX_KEYS.COLLECTION.TEST_KEY]: {
+        // @ts-expect-error nonExistingKey is not a valid key
+        ['nonExistingKey']: {
             str: 'test2',
         },
     },
 };
+
+Onyx.update([
+    {
+        onyxMethod: 'mergecollection',
+        key: ONYX_KEYS.COLLECTION.TEST_KEY,
+        value: {
+            [`${ONYX_KEYS.COLLECTION.TEST_KEY}1`]: {
+                str: 'test1',
+            },
+            [`${ONYX_KEYS.COLLECTION.TEST_KEY}2`]: {
+                str: 'test2',
+            },
+        },
+    },
+    {
+        onyxMethod: 'setcollection',
+        key: ONYX_KEYS.COLLECTION.TEST_KEY,
+        value: {
+            [`${ONYX_KEYS.COLLECTION.TEST_KEY}1`]: {
+                str: 'test1',
+            },
+        },
+    },
+    {
+        onyxMethod: 'set',
+        key: ONYX_KEYS.TEST_KEY,
+        value: 'string',
+    },
+]);
+
+Onyx.update([
+    {
+        onyxMethod: 'clear',
+        key: ONYX_KEYS.TEST_KEY,
+    },
+]);
+
+Onyx.update([
+    {
+        onyxMethod: 'merge',
+        key: ONYX_KEYS.COLLECTION.TEST_KEY,
+        value: {
+            str: 'test1',
+        },
+    },
+    {
+        onyxMethod: 'set',
+        key: ONYX_KEYS.TEST_KEY,
+        value: 'string',
+    },
+]);
+
+Onyx.update([
+    {
+        onyxMethod: 'set',
+        key: ONYX_KEYS.TEST_KEY,
+        // @ts-expect-error TEST_KEY is a string, not a boolean
+        value: true,
+    },
+]);
+
+Onyx.update([
+    {
+        onyxMethod: 'mergecollection',
+        key: ONYX_KEYS.COLLECTION.TEST_KEY,
+        value: {
+            // @ts-expect-error nonExistingKey is not a valid key
+            ['nonExistingKey']: {
+                str: 'test1',
+            },
+        },
+    },
+]);

tests/types/mergeCollection.ts

@@ -14,13 +14,16 @@ Onyx.mergeCollection(ONYX_KEYS.COLLECTION.TEST_KEY, {
 });
 
 Onyx.mergeCollection(ONYX_KEYS.COLLECTION.TEST_KEY, {
-    // @ts-expect-error COLLECTION.TEST_KEY is invalid key, it is missing the suffix
-    test_: {
-        str: 'test3',
-    },
     test_2: {
         str: 'test4',
     },
     // @ts-expect-error COLLECTION.TEST_KEY is object, not a number
     test_3: 2,
 });
+
+Onyx.mergeCollection(ONYX_KEYS.COLLECTION.TEST_KEY, {
+    // @ts-expect-error invalidKey is not a valid key
+    invalidKey: {
+        str: 'test3',
+    },
+});

tests/unit/onyxTest.ts

@@ -5,7 +5,7 @@ import waitForPromisesToResolve from '../utils/waitForPromisesToResolve';
 import OnyxUtils from '../../lib/OnyxUtils';
 import type OnyxCache from '../../lib/OnyxCache';
 import StorageMock from '../../lib/storage';
-import type {OnyxCollection, OnyxUpdate} from '../../lib/types';
+import type {OnyxCollection, OnyxKey, OnyxUpdate} from '../../lib/types';
 import type {GenericDeepRecord} from '../types';
 import type GenericCollection from '../utils/GenericCollection';
 import type {Connection} from '../../lib/OnyxConnectionManager';
@@ -997,7 +997,7 @@ describe('Onyx', () => {
                             },
                         },
                     },
-                ] as unknown as OnyxUpdate[]),
+                ] as unknown as Array<OnyxUpdate<OnyxKey>>),
             )
             .then(() => {
                 expect(valuesReceived).toEqual({
@@ -1029,7 +1029,7 @@ describe('Onyx', () => {
                         },
                     ],
                 },
-            ] as unknown as OnyxUpdate[]);
+            ] as unknown as Array<OnyxUpdate<OnyxKey>>);
         } catch (error) {
             if (error instanceof Error) {
                 expect(error.message).toEqual('Invalid value provided in Onyx multiSet. Onyx multiSet value must be of type object.');
@@ -1825,7 +1825,7 @@ describe('Onyx', () => {
                 },
             });
 
-            const queuedUpdates: OnyxUpdate[] = [
+            const queuedUpdates: Array<OnyxUpdate<OnyxKey>> = [
                 {
                     key: `${ONYX_KEYS.COLLECTION.TEST_UPDATE}entry1`,
                     onyxMethod: 'merge',
@@ -1872,7 +1872,7 @@ describe('Onyx', () => {
                 await Onyx.multiSet({[`${ONYX_KEYS.COLLECTION.TEST_UPDATE}entry1`]: entry1});
 
                 const entry1ExpectedResult = lodashCloneDeep(entry1);
-                const queuedUpdates: OnyxUpdate[] = [];
+                const queuedUpdates: Array<OnyxUpdate<OnyxKey>> = [];
 
                 queuedUpdates.push({
                     key: `${ONYX_KEYS.COLLECTION.TEST_UPDATE}entry1`,
@@ -1919,7 +1919,7 @@ describe('Onyx', () => {
                 await Onyx.multiSet({[`${ONYX_KEYS.COLLECTION.TEST_UPDATE}entry1`]: entry1});
 
                 const entry1ExpectedResult = lodashCloneDeep(entry1);
-                const queuedUpdates: OnyxUpdate[] = [];
+                const queuedUpdates: Array<OnyxUpdate<OnyxKey>> = [];
 
                 queuedUpdates.push({
                     key: `${ONYX_KEYS.COLLECTION.TEST_UPDATE}entry1`,
@@ -1991,7 +1991,7 @@ describe('Onyx', () => {
                 await Onyx.multiSet({[`${ONYX_KEYS.COLLECTION.TEST_UPDATE}entry1`]: entry1});
 
                 const entry1ExpectedResult = lodashCloneDeep(entry1);
-                const queuedUpdates: OnyxUpdate[] = [];
+                const queuedUpdates: Array<OnyxUpdate<OnyxKey>> = [];
 
                 queuedUpdates.push({
                     key: `${ONYX_KEYS.COLLECTION.TEST_UPDATE}entry1`,
@@ -2048,7 +2048,7 @@ describe('Onyx', () => {
                 });
 
                 let entry1ExpectedResult = lodashCloneDeep(entry1) as GenericDeepRecord | undefined;
-                const queuedUpdates: OnyxUpdate[] = [];
+                const queuedUpdates: Array<OnyxUpdate<OnyxKey>> = [];
 
                 queuedUpdates.push({
                     key: `${ONYX_KEYS.COLLECTION.TEST_UPDATE}entry1`,
@@ -2121,7 +2121,7 @@ describe('Onyx', () => {
 
                     const entry1ExpectedResult = lodashCloneDeep(entry1);
                     const entry2ExpectedResult = lodashCloneDeep(entry2);
-                    const queuedUpdates: OnyxUpdate[] = [];
+                    const queuedUpdates: Array<OnyxUpdate<OnyxKey>> = [];
 
                     queuedUpdates.push(
                         {
@@ -2663,7 +2663,6 @@ describe('Onyx', () => {
             } as GenericCollection);
 
             await Onyx.setCollection(ONYX_KEYS.COLLECTION.ROUTES, {
-                // @ts-expect-error invalidRoute is not a valid key
                 [invalidRoute]: {name: 'Invalid Route'},
             });