From 30c5ae7b33a1f0015c6b1416239196369d155b57 Mon Sep 17 00:00:00 2001
From: Ben Durrant <ben.j.durrant@gmail.com>
Date: Sun, 29 Oct 2023 12:03:21 +0000
Subject: [PATCH] Allow passing selector creators with different memoize
 options

---
 packages/toolkit/src/entities/state_selectors.ts     | 12 ++++++++++--
 .../src/entities/tests/sorted_state_adapter.test.ts  |  6 +++---
 .../toolkit/src/entities/tests/state_adapter.test.ts |  2 +-
 .../src/entities/tests/state_selectors.test.ts       |  5 ++++-
 4 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/packages/toolkit/src/entities/state_selectors.ts b/packages/toolkit/src/entities/state_selectors.ts
index 3bb3b88872..d2c8d9f833 100644
--- a/packages/toolkit/src/entities/state_selectors.ts
+++ b/packages/toolkit/src/entities/state_selectors.ts
@@ -2,9 +2,15 @@ import type { CreateSelectorFunction, Selector, createSelector } from 'reselect'
 import { createDraftSafeSelector } from '../createDraftSafeSelector'
 import type { EntityState, EntitySelectors, EntityId } from './models'
 
+type AnyFunction = (...args: any) => any
+type AnyCreateSelectorFunction = CreateSelectorFunction<
+  <F extends AnyFunction>(f: F) => F,
+  <F extends AnyFunction>(f: F) => F
+>
+
 export interface GetSelectorsOptions {
   // TODO Review if this causes issues or if we can go back to using `CreateSelectorFunction`
-  createSelector?: typeof createSelector
+  createSelector?: AnyCreateSelectorFunction
 }
 
 export function createSelectorsFactory<T, Id extends EntityId>() {
@@ -20,7 +26,9 @@ export function createSelectorsFactory<T, Id extends EntityId>() {
     selectState?: (state: V) => EntityState<T, Id>,
     options: GetSelectorsOptions = {}
   ): EntitySelectors<T, any, Id> {
-    const { createSelector = createDraftSafeSelector } = options
+    let { createSelector } = options
+    createSelector ??= createDraftSafeSelector
+
     const selectIds = (state: EntityState<T, Id>) => state.ids
 
     const selectEntities = (state: EntityState<T, Id>) => state.entities
diff --git a/packages/toolkit/src/entities/tests/sorted_state_adapter.test.ts b/packages/toolkit/src/entities/tests/sorted_state_adapter.test.ts
index e3287d0e3b..ff691a8326 100644
--- a/packages/toolkit/src/entities/tests/sorted_state_adapter.test.ts
+++ b/packages/toolkit/src/entities/tests/sorted_state_adapter.test.ts
@@ -11,8 +11,8 @@ import {
 import { createNextState } from '../..'
 
 describe('Sorted State Adapter', () => {
-  let adapter: EntityAdapter<string, BookModel>
-  let state: EntityState<string, BookModel>
+  let adapter: EntityAdapter<BookModel, string>
+  let state: EntityState<BookModel, string>
 
   beforeAll(() => {
     //eslint-disable-next-line
@@ -349,7 +349,7 @@ describe('Sorted State Adapter', () => {
       order: number
       ts: number
     }
-    const sortedItemsAdapter = createEntityAdapter<string, OrderedEntity>({
+    const sortedItemsAdapter = createEntityAdapter<OrderedEntity>({
       sortComparer: (a, b) => a.order - b.order,
     })
     const withInitialItems = sortedItemsAdapter.setAll(
diff --git a/packages/toolkit/src/entities/tests/state_adapter.test.ts b/packages/toolkit/src/entities/tests/state_adapter.test.ts
index ec760080ce..a05b715ff9 100644
--- a/packages/toolkit/src/entities/tests/state_adapter.test.ts
+++ b/packages/toolkit/src/entities/tests/state_adapter.test.ts
@@ -6,7 +6,7 @@ import { createSlice } from '../../createSlice'
 import type { BookModel } from './fixtures/book'
 
 describe('createStateOperator', () => {
-  let adapter: EntityAdapter<string, BookModel>
+  let adapter: EntityAdapter<BookModel, string>
 
   beforeEach(() => {
     adapter = createEntityAdapter({
diff --git a/packages/toolkit/src/entities/tests/state_selectors.test.ts b/packages/toolkit/src/entities/tests/state_selectors.test.ts
index 9ef7b32762..f581965be3 100644
--- a/packages/toolkit/src/entities/tests/state_selectors.test.ts
+++ b/packages/toolkit/src/entities/tests/state_selectors.test.ts
@@ -130,7 +130,10 @@ describe('Entity State Selectors', () => {
   })
   describe('custom createSelector instance', () => {
     it('should use the custom createSelector function if provided', () => {
-      const memoizeSpy = vi.fn(weakMapMemoize)
+      const memoizeSpy = vi.fn(
+        // test that we're allowed to pass memoizers with different options, as long as they're optional
+        <F extends (...args: any[]) => any>(fn: F, param?: boolean) => fn
+      )
       const createCustomSelector = createDraftSafeSelectorCreator(memoizeSpy)
 
       const adapter = createEntityAdapter({