Cache initial state in injected scenarios (#4908)

* add cache to injected slices

* add cache to combineSlices

* avoid calling prop.toString so much

* add tests for createSlice cache

* add test for combineSlices cache

Author: EskiMojo14

packages/toolkit/src/combineSlices.ts

@@ -322,7 +322,8 @@ const stateProxyMap = new WeakMap<object, object>()
 
 const createStateProxy = <State extends object>(
   state: State,
-  reducerMap: Partial<Record<string, Reducer>>,
+  reducerMap: Partial<Record<PropertyKey, Reducer>>,
+  initialStateCache: Record<PropertyKey, unknown>,
 ) =>
   getOrInsertComputed(
     stateProxyMap,
@@ -333,7 +334,9 @@ const createStateProxy = <State extends object>(
           if (prop === ORIGINAL_STATE) return target
           const result = Reflect.get(target, prop, receiver)
           if (typeof result === 'undefined') {
-            const reducer = reducerMap[prop.toString()]
+            const cached = initialStateCache[prop]
+            if (typeof cached !== 'undefined') return cached
+            const reducer = reducerMap[prop]
             if (reducer) {
               // ensure action type is random, to prevent reducer treating it differently
               const reducerResult = reducer(undefined, { type: nanoid() })
@@ -346,6 +349,7 @@ const createStateProxy = <State extends object>(
                     `you can use null instead of undefined.`,
                 )
               }
+              initialStateCache[prop] = reducerResult
               return reducerResult
             }
           }
@@ -361,7 +365,8 @@ const original = (state: any) => {
   return state[ORIGINAL_STATE]
 }
 
-const noopReducer: Reducer<Record<string, any>> = (state = {}) => state
+const emptyObject = {}
+const noopReducer: Reducer<Record<string, any>> = (state = emptyObject) => state
 
 export function combineSlices<Slices extends Array<AnySliceLike | ReducerMap>>(
   ...slices: Slices
@@ -382,6 +387,8 @@ export function combineSlices<Slices extends Array<AnySliceLike | ReducerMap>>(
 
   combinedReducer.withLazyLoadedSlices = () => combinedReducer
 
+  const initialStateCache: Record<PropertyKey, unknown> = {}
+
   const inject = (
     slice: AnySliceLike,
     config: InjectConfig = {},
@@ -406,6 +413,10 @@ export function combineSlices<Slices extends Array<AnySliceLike | ReducerMap>>(
       return combinedReducer
     }
 
+    if (config.overrideExisting && currentReducer !== reducerToInject) {
+      delete initialStateCache[reducerPath]
+    }
+
     reducerMap[reducerPath] = reducerToInject
 
     reducer = getReducer()
@@ -423,6 +434,7 @@ export function combineSlices<Slices extends Array<AnySliceLike | ReducerMap>>(
           createStateProxy(
             selectState ? selectState(state as any, ...args) : state,
             reducerMap,
+            initialStateCache,
           ),
           ...args,
         )

packages/toolkit/src/createSlice.ts

@@ -732,6 +732,8 @@ export function buildCreateSlice({ creators }: BuildCreateSliceConfig = {}) {
       >
     >()
 
+    const injectedStateCache = new WeakMap<(rootState: any) => State, State>()
+
     let _reducer: ReducerWithInitialState<State>
 
     function reducer(state: State | undefined, action: UnknownAction) {
@@ -757,7 +759,11 @@ export function buildCreateSlice({ creators }: BuildCreateSliceConfig = {}) {
         let sliceState = state[reducerPath]
         if (typeof sliceState === 'undefined') {
           if (injected) {
-            sliceState = getInitialState()
+            sliceState = getOrInsertComputed(
+              injectedStateCache,
+              selectSlice,
+              getInitialState,
+            )
           } else if (process.env.NODE_ENV !== 'production') {
             throw new Error(
               'selectSlice returned undefined for an uninjected slice reducer',
@@ -766,6 +772,7 @@ export function buildCreateSlice({ creators }: BuildCreateSliceConfig = {}) {
         }
         return sliceState
       }
+
       function getSelectors(
         selectState: (rootState: any) => State = selectSelf,
       ) {
@@ -783,7 +790,12 @@ export function buildCreateSlice({ creators }: BuildCreateSliceConfig = {}) {
             map[name] = wrapSelector(
               selector,
               selectState,
-              getInitialState,
+              () =>
+                getOrInsertComputed(
+                  injectedStateCache,
+                  selectState,
+                  getInitialState,
+                ),
               injected,
             )
           }

packages/toolkit/src/tests/combineSlices.test.ts

@@ -22,6 +22,12 @@ const numberSlice = createSlice({
 
 const booleanReducer = createReducer(false, () => {})
 
+const counterReducer = createSlice({
+  name: 'counter',
+  initialState: () => ({ value: 0 }),
+  reducers: {},
+})
+
 // mimic - we can't use RTKQ here directly
 const api = {
   reducerPath: 'api' as const,
@@ -144,6 +150,7 @@ describe('combineSlices', () => {
   describe('selector', () => {
     const combinedReducer = combineSlices(stringSlice).withLazyLoadedSlices<{
       boolean: boolean
+      counter: { value: number }
     }>()
 
     const uninjectedState = combinedReducer(undefined, dummyAction())
@@ -189,5 +196,20 @@ describe('combineSlices', () => {
         booleanReducer.getInitialState(),
       )
     })
+    it('caches initial state', () => {
+      const beforeInject = combinedReducer(undefined, dummyAction())
+      const injectedReducer = combinedReducer.inject(counterReducer)
+      const selectCounter = injectedReducer.selector((state) => state.counter)
+      const counter = selectCounter(beforeInject)
+      expect(counter).toBe(selectCounter(beforeInject))
+
+      injectedReducer.inject(
+        { reducerPath: 'counter', reducer: () => ({ value: 0 }) },
+        { overrideExisting: true },
+      )
+      const counter2 = selectCounter(beforeInject)
+      expect(counter2).not.toBe(counter)
+      expect(counter2).toBe(selectCounter(beforeInject))
+    })
   })
 })

packages/toolkit/src/tests/createSlice.test.ts

@@ -644,6 +644,38 @@ describe('createSlice', () => {
       // these should be different
       expect(injected.selectors).not.toBe(injected2.selectors)
     })
+    it('caches initial states for selectors', () => {
+      const slice = createSlice({
+        name: 'counter',
+        initialState: () => ({ value: 0 }),
+        reducers: {},
+        selectors: {
+          selectObj: (state) => state,
+        },
+      })
+      // not cached
+      expect(slice.getInitialState()).not.toBe(slice.getInitialState())
+      expect(slice.reducer(undefined, { type: 'dummy' })).not.toBe(
+        slice.reducer(undefined, { type: 'dummy' }),
+      )
+
+      const combinedReducer = combineSlices({
+        static: slice.reducer,
+      }).withLazyLoadedSlices<WithSlice<typeof slice>>()
+
+      const injected = slice.injectInto(combinedReducer)
+
+      // still not cached
+      expect(injected.getInitialState()).not.toBe(injected.getInitialState())
+      expect(injected.reducer(undefined, { type: 'dummy' })).not.toBe(
+        injected.reducer(undefined, { type: 'dummy' }),
+      )
+      // cached
+      expect(injected.selectSlice({})).toBe(injected.selectSlice({}))
+      expect(injected.selectors.selectObj({})).toBe(
+        injected.selectors.selectObj({}),
+      )
+    })
   })
   describe('reducers definition with asyncThunks', () => {
     it('is disabled by default', () => {