From f6bc48dad5386a2bd88aab42341ac303381a40b1 Mon Sep 17 00:00:00 2001
From: markostanimirovic <markostanimirovic95@gmail.com>
Date: Sat, 4 Feb 2023 20:32:59 +0100
Subject: [PATCH] fix(store): support using factory selectors as extra
 selectors

---
 modules/store/spec/feature_creator.spec.ts    | 10 ++++++++++
 .../store/spec/types/feature_creator.spec.ts  | 20 ++++++++++++++++---
 modules/store/src/feature_creator.ts          |  3 ++-
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/modules/store/spec/feature_creator.spec.ts b/modules/store/spec/feature_creator.spec.ts
index d2b208bb48..31695fc676 100644
--- a/modules/store/spec/feature_creator.spec.ts
+++ b/modules/store/spec/feature_creator.spec.ts
@@ -124,6 +124,12 @@ describe('createFeature()', () => {
             selectCount2,
             (count1, count2) => count1 + count2
           ),
+          selectCount3: (count: number) =>
+            createSelector(
+              selectCount1,
+              selectCount2,
+              (count1, count2) => count1 + count2 + count
+            ),
         }),
       });
 
@@ -142,6 +148,9 @@ describe('createFeature()', () => {
       expect(counterFeature.selectTotalCount({ counter: initialState })).toBe(
         initialState.count1 + initialState.count2
       );
+      expect(counterFeature.selectCount3(1)({ counter: initialState })).toBe(
+        initialState.count1 + initialState.count2 + 1
+      );
       expect(Object.keys(counterFeature)).toEqual([
         'name',
         'reducer',
@@ -150,6 +159,7 @@ describe('createFeature()', () => {
         'selectCount2',
         'selectSquaredCount2',
         'selectTotalCount',
+        'selectCount3',
       ]);
     });
 
diff --git a/modules/store/spec/types/feature_creator.spec.ts b/modules/store/spec/types/feature_creator.spec.ts
index c4e105bc92..339bb96e78 100644
--- a/modules/store/spec/types/feature_creator.spec.ts
+++ b/modules/store/spec/types/feature_creator.spec.ts
@@ -425,6 +425,8 @@ describe('createFeature()', () => {
               selectCount,
               (count) => count + 1
             ),
+            selectCountPlusNum: (num: number) =>
+              createSelector(selectCount, (count) => count + num),
           }),
         });
 
@@ -435,6 +437,7 @@ describe('createFeature()', () => {
           selectCount,
           selectCounterState2,
           selectCountPlus1,
+          selectCountPlusNum,
         } = counterFeature;
         let counterFeatureKeys: keyof typeof counterFeature;
       `);
@@ -457,9 +460,13 @@ describe('createFeature()', () => {
         'selectCountPlus1',
         'MemoizedSelector<Record<string, any>, number, (s1: number) => number>'
       );
+      snippet.toInfer(
+        'selectCountPlusNum',
+        '(num: number) => MemoizedSelector<Record<string, any>, number, (s1: number) => number>'
+      );
       snippet.toInfer(
         'counterFeatureKeys',
-        '"name" | "reducer" | "selectCounterState" | "selectCount" | "selectCounterState2" | "selectCountPlus1"'
+        '"name" | "reducer" | "selectCounterState" | "selectCount" | "selectCounterState2" | "selectCountPlus1" | "selectCountPlusNum"'
       );
     });
 
@@ -633,6 +640,7 @@ describe('createFeature()', () => {
       const snippet = expectSnippet(`
         type ExtraSelectors = {
           selectCountStr: Selector<Record<string, any>, string>;
+          selectCountPlusNum: (num: number) => Selector<Record<string, any>, number>;
         }
 
         function getExtraSelectors(
@@ -643,6 +651,8 @@ describe('createFeature()', () => {
               selectCount,
               (count) => count + ''
             ),
+            selectCountPlusNum: (num: number) =>
+              createSelector(selectCount, (count) => count + num)
           };
         }
 
@@ -653,7 +663,7 @@ describe('createFeature()', () => {
             getExtraSelectors(selectCounterState),
         });
 
-        const { selectCountStr } = counterFeature;
+        const { selectCountStr, selectCountPlusNum } = counterFeature;
         let counterFeatureKeys: keyof typeof counterFeature;
       `);
 
@@ -661,9 +671,13 @@ describe('createFeature()', () => {
         'selectCountStr',
         'Selector<Record<string, any>, string>'
       );
+      snippet.toInfer(
+        'selectCountPlusNum',
+        '(num: number) => Selector<Record<string, any>, number>'
+      );
       snippet.toInfer(
         'counterFeatureKeys',
-        '"name" | "reducer" | "selectCounterState" | "selectCountStr"'
+        '"name" | "reducer" | "selectCounterState" | keyof ExtraSelectors'
       );
     });
 
diff --git a/modules/store/src/feature_creator.ts b/modules/store/src/feature_creator.ts
index 46597982db..87f5924cfa 100644
--- a/modules/store/src/feature_creator.ts
+++ b/modules/store/src/feature_creator.ts
@@ -41,7 +41,8 @@ type BaseSelectors<
 
 type SelectorsDictionary = Record<
   string,
-  Selector<Record<string, any>, unknown>
+  | Selector<Record<string, any>, unknown>
+  | ((...args: any[]) => Selector<Record<string, any>, unknown>)
 >;
 
 type ExtraSelectorsFactory<