From 343d7d9d1234bd2f29ee3be1c55d91947fab9b0d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?F=C3=A1bio=20Henriques?= <contact@fabiohenriques.dev>
Date: Thu, 27 Jul 2023 14:40:15 +0100
Subject: [PATCH] Second attempt to solve collection key issue

---
 lib/types.d.ts    |  2 +-
 lib/withOnyx.d.ts | 73 ++++++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 67 insertions(+), 8 deletions(-)

diff --git a/lib/types.d.ts b/lib/types.d.ts
index d74cbe82..a61c5706 100644
--- a/lib/types.d.ts
+++ b/lib/types.d.ts
@@ -40,4 +40,4 @@ type Selector<TKey extends OnyxKey, TReturnType> = IsEqual<KeyValueMapping[TKey]
 
 // type Selector<TKey extends OnyxKey, TReturnType> = IsEqual<GetOnyxValue<TKey>, unknown> extends true ? (value: unknown) => unknown : (value: GetOnyxValue<TKey> | null) => TReturnType;
 
-export {CollectionKey, CustomTypeOptions, DeepRecord, Key, MergeBy, OnyxKey, KeyValueMapping, Selector, GetOnyxValue, TypeOptions};
+export {CollectionKey, CollectionKeyBase, CustomTypeOptions, DeepRecord, Key, MergeBy, OnyxKey, KeyValueMapping, Selector, GetOnyxValue, TypeOptions};
diff --git a/lib/withOnyx.d.ts b/lib/withOnyx.d.ts
index 95b06b0c..faedbc56 100644
--- a/lib/withOnyx.d.ts
+++ b/lib/withOnyx.d.ts
@@ -1,14 +1,17 @@
 import {IsEqual} from 'type-fest';
-import {GetOnyxValue, OnyxKey, Selector} from './types';
+import {CollectionKey, CollectionKeyBase, GetOnyxValue, Key, KeyValueMapping, OnyxKey, Selector} from './types';
 
-type BaseMapping<TComponentProps, TOnyxProps> = {
+type BaseMapping<TComponentProps, TOnyxProps, TOnyxProp extends keyof TOnyxProps, TOnyxKey extends OnyxKey> = {
     canEvict?: boolean | ((props: Omit<TComponentProps, keyof TOnyxProps>) => boolean);
     initWithStoredValues?: boolean;
+    // TODO: Remove this when types are finished.
+    onyxValue?: GetOnyxValue<TOnyxKey> | null;
 };
 
-type Mapping<TComponentProps, TOnyxProps, TOnyxProp extends keyof TOnyxProps, TOnyxKey extends OnyxKey> = BaseMapping<TComponentProps, TOnyxProps> &
+// TODO: Still being worked on.
+type Mapping<TComponentProps, TOnyxProps, TOnyxProp extends keyof TOnyxProps, TOnyxKey extends Key> = BaseMapping<TComponentProps, TOnyxProps, TOnyxProp, TOnyxKey> &
     (
-        | (IsEqual<GetOnyxValue<TOnyxKey> | null, TOnyxProps[TOnyxProp]> extends true
+        | (IsEqual<KeyValueMapping[TOnyxKey] | null, TOnyxProps[TOnyxProp]> extends true
               ? {
                     key: TOnyxKey | ((props: Omit<TComponentProps, keyof TOnyxProps>) => TOnyxKey);
                 }
@@ -23,14 +26,70 @@ type Mapping<TComponentProps, TOnyxProps, TOnyxProp extends keyof TOnyxProps, TO
           }
     );
 
+// TODO: Still being worked on.
+type CollectionMapping<TComponentProps, TOnyxProps, TOnyxProp extends keyof TOnyxProps, TOnyxKey extends CollectionKeyBase> = BaseMapping<TComponentProps, TOnyxProps, TOnyxProp, TOnyxKey> &
+    (
+        | (IsEqual<Record<string, KeyValueMapping[TOnyxKey] | null> | null, TOnyxProps[TOnyxProp]> extends true
+              ? {
+                    key: TOnyxKey | ((props: Omit<TComponentProps, keyof TOnyxProps>) => TOnyxKey);
+                }
+              : never)
+        | {
+              key: TOnyxKey;
+              selector: Selector<TOnyxKey, TOnyxProps[TOnyxProp]>;
+          }
+        | {
+              key: (props: Omit<TComponentProps, keyof TOnyxProps>) => TOnyxKey;
+              selector: Selector<TOnyxKey, TOnyxProps[TOnyxProp]>;
+          }
+    );
+
+// TODO: Still being worked on.
+type CollectionRecordMapping<TComponentProps, TOnyxProps, TOnyxProp extends keyof TOnyxProps, TOnyxKey extends CollectionKey> = BaseMapping<
+    TComponentProps,
+    TOnyxProps,
+    TOnyxProp,
+    TOnyxKey
+> &
+    (
+        | (IsEqual<KeyValueMapping[TOnyxKey] | null, TOnyxProps[TOnyxProp]> extends true
+              ? {
+                    key: TOnyxKey | ((props: Omit<TComponentProps, keyof TOnyxProps>) => TOnyxKey);
+                }
+              : never)
+        | {
+              key: TOnyxKey;
+              selector: Selector<TOnyxKey, TOnyxProps[TOnyxProp]>;
+          }
+        | {
+              key: (props: Omit<TComponentProps, keyof TOnyxProps>) => TOnyxKey;
+              selector: Selector<TOnyxKey, TOnyxProps[TOnyxProp]>;
+          }
+    );
+
+// TODO: Still being worked on.
 type OnyxPropMapping<TComponentProps, TOnyxProps, TOnyxProp extends keyof TOnyxProps> = {
-    [TOnyxKey in OnyxKey]: Mapping<TComponentProps, TOnyxProps, TOnyxProp, TOnyxKey>;
-}[OnyxKey];
+    [TOnyxKey in Key]: Mapping<TComponentProps, TOnyxProps, TOnyxProp, TOnyxKey>;
+}[Key];
+
+// TODO: Still being worked on.
+type OnyxPropCollectionMapping<TComponentProps, TOnyxProps, TOnyxProp extends keyof TOnyxProps> = {
+    [TOnyxKey in CollectionKeyBase]: CollectionMapping<TComponentProps, TOnyxProps, TOnyxProp, TOnyxKey>;
+}[CollectionKeyBase];
+
+// TODO: Still being worked on.
+type OnyxPropCollectionRecordMapping<TComponentProps, TOnyxProps, TOnyxProp extends keyof TOnyxProps> = {
+    [TOnyxKey in CollectionKey]: CollectionRecordMapping<TComponentProps, TOnyxProps, TOnyxProp, TOnyxKey>;
+}[CollectionKey];
 
 declare function withOnyx<TComponentProps, TOnyxProps>(
     mapping: {
-        [TOnyxProp in keyof TOnyxProps]: OnyxPropMapping<TComponentProps, TOnyxProps, TOnyxProp>;
+        [TOnyxProp in keyof TOnyxProps]:
+            | OnyxPropMapping<TComponentProps, TOnyxProps, TOnyxProp>
+            | OnyxPropCollectionMapping<TComponentProps, TOnyxProps, TOnyxProp>
+            | OnyxPropCollectionRecordMapping<TComponentProps, TOnyxProps, TOnyxProp>;
     },
 ): (component: React.ComponentType<TComponentProps>) => React.ComponentType<Omit<TComponentProps, keyof TOnyxProps>>;
 
 export default withOnyx;
+export {Mapping, CollectionMapping, CollectionRecordMapping, OnyxPropMapping, GetOnyxValue, Selector};