Merge pull request #3481 from reduxjs/customise-create-selector

Author: markerikson

docs/api/createEntityAdapter.mdx

@@ -275,9 +275,35 @@ The entity adapter will contain a `getSelectors()` function that returns a set o
 
 Each selector function will be created using the `createSelector` function from Reselect, to enable memoizing calculation of the results.
 
+:::tip
+
+The `createSelector` instance used can be replaced, by passing it as part of the options object (second parameter):
+
+```js
+import {
+  createDraftSafeSelectorCreator,
+  weakMapMemoize,
+} from '@reduxjs/toolkit'
+
+const createWeakMapDraftSafeSelector =
+  createDraftSafeSelectorCreator(weakMapMemoize)
+
+const simpleSelectors = booksAdapter.getSelectors(undefined, {
+  createSelector: createWeakMapDraftSafeSelector,
+})
+
+const globalizedSelectors = booksAdapter.getSelectors((state) => state.books, {
+  createSelector: createWeakMapDraftSafeSelector,
+})
+```
+
+If no instance is passed, it will default to [`createDraftSafeSelector`](./createSelector#createDraftSafeSelector).
+
+:::
+
 Because selector functions are dependent on knowing where in the state tree this specific entity state object is kept, `getSelectors()` can be called in two ways:
 
-- If called without any arguments, it returns an "unglobalized" set of selector functions that assume their `state` argument is the actual entity state object to read from.
+- If called without any arguments (or with undefined as the first parameter), it returns an "unglobalized" set of selector functions that assume their `state` argument is the actual entity state object to read from.
 - It may also be called with a selector function that accepts the entire Redux state tree and returns the correct entity state object.
 
 For example, the entity state for a `Book` type might be kept in the Redux state tree as `state.books`. You can use `getSelectors()` to read from that state in two ways:

docs/api/createSelector.mdx

@@ -37,7 +37,7 @@ All selectors created by `entityAdapter.getSelectors` are "draft safe" selectors
 
 Example:
 
-```js
+```ts no-transpile
 const selectSelf = (state: State) => state
 const unsafeSelector = createSelector(selectSelf, (state) => state.value)
 const draftSafeSelector = createDraftSafeSelector(
@@ -62,3 +62,25 @@ After executing that, `unsafe1` and `unsafe2` will be of the same value, because
 executed on the same object - but `safe2` will actually be different from `safe1` (with the updated value of `2`),
 because the safe selector detected that it was executed on a Immer draft object and recalculated using the current
 value instead of returning a cached value.
+
+:::tip `createDraftSafeSelectorCreator`
+
+RTK also exports a `createDraftSafeSelectorCreator` function, the "draft safe" equivalent of [`createSelectorCreator`](https://github.com/reduxjs/reselect#createselectorcreatormemoize-memoizeoptions).
+
+```ts no-transpile
+import {
+  createDraftSafeSelectorCreator,
+  weakMapMemoize,
+} from '@reduxjs/toolkit'
+
+const createWeakMapDraftSafeSelector =
+  createDraftSafeSelectorCreator(weakMapMemoize)
+
+const selectSelf = (state: State) => state
+const draftSafeSelector = createWeakMapDraftSafeSelector(
+  selectSelf,
+  (state) => state.value
+)
+```
+
+:::

packages/toolkit/src/createDraftSafeSelector.ts

@@ -1,5 +1,17 @@
 import { current, isDraft } from 'immer'
-import { createSelector } from 'reselect'
+import { createSelectorCreator, defaultMemoize } from 'reselect'
+
+export const createDraftSafeSelectorCreator: typeof createSelectorCreator = (
+  ...args: unknown[]
+) => {
+  const createSelector = (createSelectorCreator as any)(...args)
+  return (...args: unknown[]) => {
+    const selector = createSelector(...args)
+    const wrappedSelector = (value: unknown, ...rest: unknown[]) =>
+      selector(isDraft(value) ? current(value) : value, ...rest)
+    return wrappedSelector as any
+  }
+}
 
 /**
  * "Draft-Safe" version of `reselect`'s `createSelector`:
@@ -8,11 +20,5 @@ import { createSelector } from 'reselect'
  * that might be possibly outdated if the draft has been modified since.
  * @public
  */
-export const createDraftSafeSelector: typeof createSelector = (
-  ...args: unknown[]
-) => {
-  const selector = (createSelector as any)(...args)
-  const wrappedSelector = (value: unknown, ...rest: unknown[]) =>
-    selector(isDraft(value) ? current(value) : value, ...rest)
-  return wrappedSelector as any
-}
+export const createDraftSafeSelector =
+  createDraftSafeSelectorCreator(defaultMemoize)

packages/toolkit/src/entities/models.ts

@@ -1,5 +1,6 @@
 import type { UncheckedIndexedAccess } from '../uncheckedindexed'
 import type { PayloadAction } from '../createAction'
+import type { GetSelectorsOptions } from './state_selectors'
 import type { CastAny, Id as Compute } from '../tsHelpers'
 
 /**
@@ -167,8 +168,12 @@ export interface EntityAdapter<T, Id extends EntityId>
   sortComparer: false | Comparer<T>
   getInitialState(): EntityState<T, Id>
   getInitialState<S extends object>(state: S): EntityState<T, Id> & S
-  getSelectors(): EntitySelectors<T, EntityState<T, Id>, Id>
+  getSelectors(
+    selectState?: undefined,
+    options?: GetSelectorsOptions
+  ): EntitySelectors<T, EntityState<T, Id>, Id>
   getSelectors<V>(
-    selectState: (state: V) => EntityState<T, Id>
+    selectState: (state: V) => EntityState<T, Id>,
+    options?: GetSelectorsOptions
   ): EntitySelectors<T, V, Id>
 }

packages/toolkit/src/entities/state_selectors.ts

@@ -1,20 +1,35 @@
-import type { Selector } from 'reselect'
+import type { CreateSelectorFunction, Selector } from 'reselect'
 import { createDraftSafeSelector } from '../createDraftSafeSelector'
 import type { EntityState, EntitySelectors, EntityId } from './models'
 
+export type AnyCreateSelectorFunction = CreateSelectorFunction<
+  (...args: unknown[]) => unknown,
+  <F extends (...args: any[]) => any>(func: F) => F
+>
+
+export interface GetSelectorsOptions {
+  createSelector?: AnyCreateSelectorFunction
+}
+
 export function createSelectorsFactory<T, Id extends EntityId>() {
-  function getSelectors(): EntitySelectors<T, EntityState<T, Id>, Id>
+  function getSelectors(
+    selectState?: undefined,
+    options?: GetSelectorsOptions
+  ): EntitySelectors<T, EntityState<T, Id>, Id>
   function getSelectors<V>(
-    selectState: (state: V) => EntityState<T, Id>
+    selectState: (state: V) => EntityState<T, Id>,
+    options?: GetSelectorsOptions
   ): EntitySelectors<T, V, Id>
   function getSelectors<V>(
-    selectState?: (state: V) => EntityState<T, Id>
+    selectState?: (state: V) => EntityState<T, Id>,
+    options: GetSelectorsOptions = {}
   ): EntitySelectors<T, any, Id> {
+    const { createSelector = createDraftSafeSelector } = options
     const selectIds = (state: EntityState<T, Id>) => state.ids
 
     const selectEntities = (state: EntityState<T, Id>) => state.entities
 
-    const selectAll = createDraftSafeSelector(
+    const selectAll = createSelector(
       selectIds,
       selectEntities,
       (ids, entities): T[] => ids.map((id) => entities[id]!)
@@ -24,33 +39,29 @@ export function createSelectorsFactory<T, Id extends EntityId>() {
 
     const selectById = (entities: Record<Id, T>, id: Id) => entities[id]
 
-    const selectTotal = createDraftSafeSelector(selectIds, (ids) => ids.length)
+    const selectTotal = createSelector(selectIds, (ids) => ids.length)
 
     if (!selectState) {
       return {
         selectIds,
         selectEntities,
         selectAll,
         selectTotal,
-        selectById: createDraftSafeSelector(
-          selectEntities,
-          selectId,
-          selectById
-        ),
+        selectById: createSelector(selectEntities, selectId, selectById),
       }
     }
 
-    const selectGlobalizedEntities = createDraftSafeSelector(
+    const selectGlobalizedEntities = createSelector(
       selectState as Selector<V, EntityState<T, Id>>,
       selectEntities
     )
 
     return {
-      selectIds: createDraftSafeSelector(selectState, selectIds),
+      selectIds: createSelector(selectState, selectIds),
       selectEntities: selectGlobalizedEntities,
-      selectAll: createDraftSafeSelector(selectState, selectAll),
-      selectTotal: createDraftSafeSelector(selectState, selectTotal),
-      selectById: createDraftSafeSelector(
+      selectAll: createSelector(selectState, selectAll),
+      selectTotal: createSelector(selectState, selectTotal),
+      selectById: createSelector(
         selectGlobalizedEntities,
         selectId,
         selectById

packages/toolkit/src/entities/tests/state_selectors.test.ts

@@ -1,10 +1,12 @@
+import { createDraftSafeSelectorCreator } from '../../createDraftSafeSelector'
 import type { EntityAdapter, EntityState } from '../index'
 import { createEntityAdapter } from '../index'
 import type { EntitySelectors } from '../models'
 import type { BookModel } from './fixtures/book'
 import { AClockworkOrange, AnimalFarm, TheGreatGatsby } from './fixtures/book'
 import type { Selector } from 'reselect'
-import { createSelector } from 'reselect'
+import { createSelector, weakMapMemoize } from 'reselect'
+import { vi } from 'vitest'
 
 describe('Entity State Selectors', () => {
   describe('Composed Selectors', () => {
@@ -126,6 +128,22 @@ describe('Entity State Selectors', () => {
       expect(second).toBe(AnimalFarm)
     })
   })
+  describe('custom createSelector instance', () => {
+    it('should use the custom createSelector function if provided', () => {
+      const memoizeSpy = vi.fn(weakMapMemoize)
+      const createCustomSelector = createDraftSafeSelectorCreator(memoizeSpy)
+
+      const adapter = createEntityAdapter({
+        selectId: (book: BookModel) => book.id,
+      })
+
+      adapter.getSelectors(undefined, { createSelector: createCustomSelector })
+
+      expect(memoizeSpy).toHaveBeenCalled()
+
+      memoizeSpy.mockClear()
+    })
+  })
 })
 
 function expectType<T>(t: T) {

packages/toolkit/src/index.ts

@@ -20,7 +20,10 @@ export type {
   OutputSelector,
   ParametricSelector,
 } from 'reselect'
-export { createDraftSafeSelector } from './createDraftSafeSelector'
+export {
+  createDraftSafeSelector,
+  createDraftSafeSelectorCreator,
+} from './createDraftSafeSelector'
 export type { ThunkAction, ThunkDispatch, ThunkMiddleware } from 'redux-thunk'
 
 export {