Repository: reduxjs/redux-toolkit

Problem statement:
Improve `createSlice.selectors` internals
Per @EskiMojo14 :

> we may still want to pull over the emplace and .unwrapped changes . emplace just simplifies some of the map logic we have, and .unwrapped attaches the original selector to the wrapped version


Hints:

Patch:
diff --git a/packages/toolkit/src/combineSlices.ts b/packages/toolkit/src/combineSlices.ts
--- a/packages/toolkit/src/combineSlices.ts
+++ b/packages/toolkit/src/combineSlices.ts
@@ -8,6 +8,7 @@ import type {
   UnionToIntersection,
   WithOptionalProp,
 } from './tsHelpers'
+import { emplace } from './utils'
 
 type SliceLike<ReducerPath extends string, State> = {
   reducerPath: ReducerPath
@@ -330,37 +331,34 @@ const stateProxyMap = new WeakMap<object, object>()
 const createStateProxy = <State extends object>(
   state: State,
   reducerMap: Partial<Record<string, Reducer>>
-) => {
-  let proxy = stateProxyMap.get(state)
-  if (!proxy) {
-    proxy = new Proxy(state, {
-      get: (target, prop, receiver) => {
-        if (prop === ORIGINAL_STATE) return target
-        const result = Reflect.get(target, prop, receiver)
-        if (typeof result === 'undefined') {
-          const reducer = reducerMap[prop.toString()]
-          if (reducer) {
-            // ensure action type is random, to prevent reducer treating it differently
-            const reducerResult = reducer(undefined, { type: nanoid() })
-            if (typeof reducerResult === 'undefined') {
-              throw new Error(
-                `The slice reducer for key "${prop.toString()}" returned undefined when called for selector(). ` +
-                  `If the state passed to the reducer is undefined, you must ` +
-                  `explicitly return the initial state. The initial state may ` +
-                  `not be undefined. If you don't want to set a value for this reducer, ` +
-                  `you can use null instead of undefined.`
-              )
+) =>
+  emplace(stateProxyMap, state, {
+    insert: () =>
+      new Proxy(state, {
+        get: (target, prop, receiver) => {
+          if (prop === ORIGINAL_STATE) return target
+          const result = Reflect.get(target, prop, receiver)
+          if (typeof result === 'undefined') {
+            const reducer = reducerMap[prop.toString()]
+            if (reducer) {
+              // ensure action type is random, to prevent reducer treating it differently
+              const reducerResult = reducer(undefined, { type: nanoid() })
+              if (typeof reducerResult === 'undefined') {
+                throw new Error(
+                  `The slice reducer for key "${prop.toString()}" returned undefined when called for selector(). ` +
+                    `If the state passed to the reducer is undefined, you must ` +
+                    `explicitly return the initial state. The initial state may ` +
+                    `not be undefined. If you don't want to set a value for this reducer, ` +
+                    `you can use null instead of undefined.`
+                )
+              }
+              return reducerResult
             }
-            return reducerResult
           }
-        }
-        return result
-      },
-    })
-    stateProxyMap.set(state, proxy)
-  }
-  return proxy as State
-}
+          return result
+        },
+      }),
+  }) as State
 
 const original = (state: any) => {
   if (!isStateProxy(state)) {
diff --git a/packages/toolkit/src/createSlice.ts b/packages/toolkit/src/createSlice.ts
--- a/packages/toolkit/src/createSlice.ts
+++ b/packages/toolkit/src/createSlice.ts
@@ -1,4 +1,5 @@
 import type { Action, UnknownAction, Reducer } from 'redux'
+import type { Selector } from 'reselect'
 import type {
   ActionCreatorWithoutPayload,
   PayloadAction,
@@ -25,6 +26,7 @@ import type {
   OverrideThunkApiConfigs,
 } from './createAsyncThunk'
 import { createAsyncThunk as _createAsyncThunk } from './createAsyncThunk'
+import { emplace } from './utils'
 
 const asyncThunkSymbol = Symbol.for('rtk-slice-createasyncthunk')
 // type is annotated because it's too long to infer
@@ -531,6 +533,14 @@ type SliceDefinedCaseReducers<CaseReducers extends SliceCaseReducers<any>> = {
     : never
 }
 
+type RemappedSelector<S extends Selector, NewState> = S extends Selector<
+  any,
+  infer R,
+  infer P
+>
+  ? Selector<NewState, R, P> & { unwrapped: S }
+  : never
+
 /**
  * Extracts the final selector type from the `selectors` object.
  *
@@ -541,10 +551,10 @@ type SliceDefinedSelectors<
   Selectors extends SliceSelectors<State>,
   RootState
 > = {
-  [K in keyof Selectors as string extends K ? never : K]: (
-    rootState: RootState,
-    ...args: Tail<Parameters<Selectors[K]>>
-  ) => ReturnType<Selectors[K]>
+  [K in keyof Selectors as string extends K ? never : K]: RemappedSelector<
+    Selectors[K],
+    RootState
+  >
 }
 
 /**
@@ -756,35 +766,26 @@ export function buildCreateSlice({ creators }: BuildCreateSliceConfig = {}) {
         return _reducer.getInitialState()
       },
       getSelectors(selectState: (rootState: any) => State = selectSelf) {
-        let selectorCache = injectedSelectorCache.get(this)
-        if (!selectorCache) {
-          selectorCache = new WeakMap()
-          injectedSelectorCache.set(this, selectorCache)
-        }
-        let cached = selectorCache.get(selectState)
-        if (!cached) {
-          cached = {}
-          for (const [name, selector] of Object.entries(
-            options.selectors ?? {}
-          )) {
-            cached[name] = (rootState: any, ...args: any[]) => {
-              let sliceState = selectState.call(this, rootState)
-              if (typeof sliceState === 'undefined') {
-                // check if injectInto has been called
-                if (this !== slice) {
-                  sliceState = this.getInitialState()
-                } else if (process.env.NODE_ENV !== 'production') {
-                  throw new Error(
-                    'selectState returned undefined for an uninjected slice reducer'
-                  )
-                }
-              }
-              return selector(sliceState, ...args)
+        const selectorCache = emplace(injectedSelectorCache, this, {
+          insert: () => new WeakMap(),
+        })
+
+        return emplace(selectorCache, selectState, {
+          insert: () => {
+            const map: Record<string, Selector<any, any>> = {}
+            for (const [name, selector] of Object.entries(
+              options.selectors ?? {}
+            )) {
+              map[name] = wrapSelector(
+                this,
+                selector,
+                selectState,
+                this !== slice
+              )
             }
-          }
-          selectorCache.set(selectState, cached)
-        }
-        return cached as any
+            return map
+          },
+        }) as any
       },
       selectSlice(state) {
         let sliceState = state[this.reducerPath]
@@ -816,6 +817,29 @@ export function buildCreateSlice({ creators }: BuildCreateSliceConfig = {}) {
   }
 }
 
+function wrapSelector<State, NewState, S extends Selector<State>>(
+  slice: Slice,
+  selector: S,
+  selectState: Selector<NewState, State>,
+  injected?: boolean
+) {
+  function wrapper(rootState: NewState, ...args: any[]) {
+    let sliceState = selectState.call(slice, rootState)
+    if (typeof sliceState === 'undefined') {
+      if (injected) {
+        sliceState = slice.getInitialState()
+      } else if (process.env.NODE_ENV !== 'production') {
+        throw new Error(
+          'selectState returned undefined for an uninjected slice reducer'
+        )
+      }
+    }
+    return selector(sliceState, ...args)
+  }
+  wrapper.unwrapped = selector
+  return wrapper as RemappedSelector<S, NewState>
+}
+
 /**
  * A function that accepts an initial state, an object full of reducer
  * functions, and a "slice name", and automatically generates
diff --git a/packages/toolkit/src/dynamicMiddleware/index.ts b/packages/toolkit/src/dynamicMiddleware/index.ts
--- a/packages/toolkit/src/dynamicMiddleware/index.ts
+++ b/packages/toolkit/src/dynamicMiddleware/index.ts
@@ -7,7 +7,7 @@ import { compose } from 'redux'
 import { createAction, isAction } from '../createAction'
 import { isAllOf } from '../matchers'
 import { nanoid } from '../nanoid'
-import { find } from '../utils'
+import { emplace, find } from '../utils'
 import type {
   WithMiddleware,
   AddMiddleware,
@@ -69,15 +69,8 @@ export const createDynamicMiddleware = <
   ) as AddMiddleware<State, Dispatch>
 
   const getFinalMiddleware: Middleware<{}, State, Dispatch> = (api) => {
-    const appliedMiddleware = Array.from(middlewareMap.values()).map(
-      (entry) => {
-        let applied = entry.applied.get(api)
-        if (!applied) {
-          applied = entry.middleware(api)
-          entry.applied.set(api, applied)
-        }
-        return applied
-      }
+    const appliedMiddleware = Array.from(middlewareMap.values()).map((entry) =>
+      emplace(entry.applied, api, { insert: () => entry.middleware(api) })
     )
     return compose(...appliedMiddleware)
   }
diff --git a/packages/toolkit/src/utils.ts b/packages/toolkit/src/utils.ts
--- a/packages/toolkit/src/utils.ts
+++ b/packages/toolkit/src/utils.ts
@@ -87,3 +87,82 @@ export class Tuple<Items extends ReadonlyArray<unknown> = []> extends Array<
 export function freezeDraftable<T>(val: T) {
   return isDraftable(val) ? createNextState(val, () => {}) : val
 }
+
+interface WeakMapEmplaceHandler<K extends object, V> {
+  /**
+   * Will be called to get value, if no value is currently in map.
+   */
+  insert?(key: K, map: WeakMap<K, V>): V
+  /**
+   * Will be called to update a value, if one exists already.
+   */
+  update?(previous: V, key: K, map: WeakMap<K, V>): V
+}
+
+interface MapEmplaceHandler<K, V> {
+  /**
+   * Will be called to get value, if no value is currently in map.
+   */
+  insert?(key: K, map: Map<K, V>): V
+  /**
+   * Will be called to update a value, if one exists already.
+   */
+  update?(previous: V, key: K, map: Map<K, V>): V
+}
+
+export function emplace<K, V>(
+  map: Map<K, V>,
+  key: K,
+  handler: MapEmplaceHandler<K, V>
+): V
+export function emplace<K extends object, V>(
+  map: WeakMap<K, V>,
+  key: K,
+  handler: WeakMapEmplaceHandler<K, V>
+): V
+/**
+ * Allow inserting a new value, or updating an existing one
+ * @throws if called for a key with no current value and no `insert` handler is provided
+ * @returns current value in map (after insertion/updating)
+ * ```ts
+ * // return current value if already in map, otherwise initialise to 0 and return that
+ * const num = emplace(map, key, {
+ *   insert: () => 0
+ * })
+ *
+ * // increase current value by one if already in map, otherwise initialise to 0
+ * const num = emplace(map, key, {
+ *   update: (n) => n + 1,
+ *   insert: () => 0,
+ * })
+ *
+ * // only update if value's already in the map - and increase it by one
+ * if (map.has(key)) {
+ *   const num = emplace(map, key, {
+ *     update: (n) => n + 1,
+ *   })
+ * }
+ * ```
+ *
+ * @remarks
+ * Based on https://github.com/tc39/proposal-upsert currently in Stage 2 - maybe in a few years we'll be able to replace this with direct method calls
+ */
+export function emplace<K extends object, V>(
+  map: WeakMap<K, V>,
+  key: K,
+  handler: WeakMapEmplaceHandler<K, V>
+): V {
+  if (map.has(key)) {
+    let value = map.get(key) as V
+    if (handler.update) {
+      value = handler.update(value, key, map)
+      map.set(key, value)
+    }
+    return value
+  }
+  if (!handler.insert)
+    throw new Error('No insert provided for key not already in map')
+  const inserted = handler.insert(key, map)
+  map.set(key, inserted)
+  return inserted
+}


Test patch:
diff --git a/packages/toolkit/src/tests/createSlice.test.ts b/packages/toolkit/src/tests/createSlice.test.ts
--- a/packages/toolkit/src/tests/createSlice.test.ts
+++ b/packages/toolkit/src/tests/createSlice.test.ts
@@ -466,12 +466,15 @@ describe('createSlice', () => {
       reducers: {},
       selectors: {
         selectSlice: (state) => state,
-        selectMultiple: (state, multiplier: number) => state * multiplier,
+        selectMultiple: Object.assign(
+          (state: number, multiplier: number) => state * multiplier,
+          { test: 0 }
+        ),
       },
     })
-    it('expects reducer under slice.name if no selectState callback passed', () => {
+    it('expects reducer under slice.reducerPath if no selectState callback passed', () => {
       const testState = {
-        [slice.name]: slice.getInitialState(),
+        [slice.reducerPath]: slice.getInitialState(),
       }
       const { selectSlice, selectMultiple } = slice.selectors
       expect(selectSlice(testState)).toBe(slice.getInitialState())
@@ -487,6 +490,9 @@ describe('createSlice', () => {
       expect(selectSlice(customState)).toBe(slice.getInitialState())
       expect(selectMultiple(customState, 2)).toBe(slice.getInitialState() * 2)
     })
+    it('allows accessing properties on the selector', () => {
+      expect(slice.selectors.selectMultiple.unwrapped.test).toBe(0)
+    })
   })
   describe('slice injections', () => {
     it('uses injectInto to inject slice into combined reducer', () => {
@@ -580,7 +586,9 @@ describe('createSlice', () => {
           initialState: [] as any[],
           reducers: (create) => ({ thunk: create.asyncThunk(() => {}) }),
         })
-      ).toThrowErrorMatchingInlineSnapshot('"Cannot use `create.asyncThunk` in the built-in `createSlice`. Use `buildCreateSlice({ creators: { asyncThunk: asyncThunkCreator } })` to create a customised version of `createSlice`."')
+      ).toThrowErrorMatchingInlineSnapshot(
+        '"Cannot use `create.asyncThunk` in the built-in `createSlice`. Use `buildCreateSlice({ creators: { asyncThunk: asyncThunkCreator } })` to create a customised version of `createSlice`."'
+      )
     })
     const createThunkSlice = buildCreateSlice({
       creators: { asyncThunk: asyncThunkCreator },
diff --git a/packages/toolkit/src/tests/createSlice.typetest.ts b/packages/toolkit/src/tests/createSlice.typetest.ts
--- a/packages/toolkit/src/tests/createSlice.typetest.ts
+++ b/packages/toolkit/src/tests/createSlice.typetest.ts
@@ -540,7 +540,10 @@ const value = actionCreators.anyKey
     selectors: {
       selectValue: (state) => state.value,
       selectMultiply: (state, multiplier: number) => state.value * multiplier,
-      selectToFixed: (state) => state.value.toFixed(2),
+      selectToFixed: Object.assign(
+        (state: { value: number }) => state.value.toFixed(2),
+        { static: true }
+      ),
     },
   })
 
@@ -555,6 +558,8 @@ const value = actionCreators.anyKey
   expectType<number>(selectMultiply(rootState, 2))
   expectType<string>(selectToFixed(rootState))
 
+  expectType<boolean>(selectToFixed.unwrapped.static)
+
   const nestedState = {
     nested: rootState,
   }