Update createReducer to accept a lazy state init function

.gitignore

@@ -16,6 +16,7 @@ temp/
 .tmp-projections
 build/
 .rts2*
+coverage/
 
 typesversions
 .cache

docs/api/createReducer.mdx

@@ -121,6 +121,21 @@ so we recommend the "builder callback" notation in most cases.
 
 [params](docblock://createReducer.ts?token=createReducer&overload=1)
 
+### Returns
+
+The generated reducer function.
+
+The reducer will have a `getInitialState` function attached that will return the initial state when called. This may be useful for tests or usage with React's `useReducer` hook:
+
+```js
+const counterReducer = createReducer(0, {
+  increment: (state, action) => state + action.payload,
+  decrement: (state, action) => state - action.payload,
+})
+
+console.log(counterReducer.getInitialState()) // 0
+```
+
 ### Example Usage
 
 [examples](docblock://createReducer.ts?token=createReducer&overload=1)

docs/api/createSlice.mdx

@@ -71,6 +71,8 @@ function createSlice({
 
 The initial state value for this slice of state.
 
+This may also be a "lazy initializer" function, which should return an initial state value when called. This will be used whenever the reducer is called with `undefined` as its state value, and is primarily useful for cases like reading initial state from `localStorage`.
+
 ### `name`
 
 A string name for this slice of state. Generated action type constants will use this as a prefix.
@@ -196,7 +198,8 @@ We recommend using the `builder callback` API as the default, especially if you
     name : string,
     reducer : ReducerFunction,
     actions : Record<string, ActionCreator>,
-    caseReducers: Record<string, CaseReducer>
+    caseReducers: Record<string, CaseReducer>.
+    getInitialState: () => State
 }
 ```
 

packages/toolkit/src/createReducer.ts

@@ -66,6 +66,16 @@ export type CaseReducers<S, AS extends Actions> = {
   [T in keyof AS]: AS[T] extends Action ? CaseReducer<S, AS[T]> : void
 }
 
+export type NotFunction<T> = T extends Function ? never : T
+
+function isStateFunction<S>(x: unknown): x is () => S {
+  return typeof x === 'function'
+}
+
+export type ReducerWithInitialState<S extends NotFunction<any>> = Reducer<S> & {
+  getInitialState: () => S
+}
+
 /**
  * A utility function that allows defining a reducer as a mapping from action
  * type to *case reducer* functions that handle these action types. The
@@ -84,8 +94,8 @@ export type CaseReducers<S, AS extends Actions> = {
  * That builder provides `addCase`, `addMatcher` and `addDefaultCase` functions that may be
  * called to define what actions this reducer will handle.
  *
- * @param initialState - The initial state that should be used when the reducer is called the first time.
- * @param builderCallback - A callback that receives a *builder* object to define
+ * @param initialState - `State | (() => State)`: The initial state that should be used when the reducer is called the first time. This may also be a "lazy initializer" function, which should return an initial state value when called. This will be used whenever the reducer is called with `undefined` as its state value, and is primarily useful for cases like reading initial state from `localStorage`.
+ * @param builderCallback - `(builder: Builder) => void` A callback that receives a *builder* object to define
  *   case reducers via calls to `builder.addCase(actionCreatorOrType, reducer)`.
  * @example
 ```ts
@@ -105,7 +115,7 @@ function isActionWithNumberPayload(
   return typeof action.payload === "number";
 }
 
-createReducer(
+const reducer = createReducer(
   {
     counter: 0,
     sumOfNumberPayloads: 0,
@@ -130,10 +140,10 @@ createReducer(
 ```
  * @public
  */
-export function createReducer<S>(
-  initialState: S,
+export function createReducer<S extends NotFunction<any>>(
+  initialState: S | (() => S),
   builderCallback: (builder: ActionReducerMapBuilder<S>) => void
-): Reducer<S>
+): ReducerWithInitialState<S>
 
 /**
  * A utility function that allows defining a reducer as a mapping from action
@@ -151,7 +161,7 @@ export function createReducer<S>(
  * This overload accepts an object where the keys are string action types, and the values
  * are case reducer functions to handle those action types.
  *
- * @param initialState - The initial state that should be used when the reducer is called the first time.
+ * @param initialState - `State | (() => State)`: The initial state that should be used when the reducer is called the first time. This may also be a "lazy initializer" function, which should return an initial state value when called. This will be used whenever the reducer is called with `undefined` as its state value, and is primarily useful for cases like reading initial state from `localStorage`.
  * @param actionsMap - An object mapping from action types to _case reducers_, each of which handles one specific action type.
  * @param actionMatchers - An array of matcher definitions in the form `{matcher, reducer}`.
  *   All matching reducers will be executed in order, independently if a case reducer matched or not.
@@ -164,6 +174,14 @@ const counterReducer = createReducer(0, {
   increment: (state, action) => state + action.payload,
   decrement: (state, action) => state - action.payload
 })
+
+// Alternately, use a "lazy initializer" to provide the initial state
+// (works with either form of createReducer)
+const initialState = () => 0
+const counterReducer = createReducer(initialState, {
+  increment: (state, action) => state + action.payload,
+  decrement: (state, action) => state - action.payload
+})
 ```
 
  * Action creators that were generated using [`createAction`](./createAction) may be used directly as the keys here, using computed property syntax:
@@ -180,31 +198,38 @@ const counterReducer = createReducer(0, {
  * @public
  */
 export function createReducer<
-  S,
+  S extends NotFunction<any>,
   CR extends CaseReducers<S, any> = CaseReducers<S, any>
 >(
-  initialState: S,
+  initialState: S | (() => S),
   actionsMap: CR,
   actionMatchers?: ActionMatcherDescriptionCollection<S>,
   defaultCaseReducer?: CaseReducer<S>
-): Reducer<S>
+): ReducerWithInitialState<S>
 
-export function createReducer<S>(
-  initialState: S,
+export function createReducer<S extends NotFunction<any>>(
+  initialState: S | (() => S),
   mapOrBuilderCallback:
     | CaseReducers<S, any>
     | ((builder: ActionReducerMapBuilder<S>) => void),
   actionMatchers: ReadonlyActionMatcherDescriptionCollection<S> = [],
   defaultCaseReducer?: CaseReducer<S>
-): Reducer<S> {
+): ReducerWithInitialState<S> {
   let [actionsMap, finalActionMatchers, finalDefaultCaseReducer] =
     typeof mapOrBuilderCallback === 'function'
       ? executeReducerBuilderCallback(mapOrBuilderCallback)
       : [mapOrBuilderCallback, actionMatchers, defaultCaseReducer]
 
-  const frozenInitialState = createNextState(initialState, () => {})
+  // Ensure the initial state gets frozen either way
+  let getInitialState: () => S
+  if (isStateFunction(initialState)) {
+    getInitialState = () => createNextState(initialState(), () => {})
+  } else {
+    const frozenInitialState = createNextState(initialState, () => {})
+    getInitialState = () => frozenInitialState
+  }
 
-  return function (state = frozenInitialState, action): S {
+  function reducer(state = getInitialState(), action: any): S {
     let caseReducers = [
       actionsMap[action.type],
       ...finalActionMatchers
@@ -257,4 +282,8 @@ export function createReducer<S>(
       return previousState
     }, state)
   }
+
+  reducer.getInitialState = getInitialState
+
+  return reducer as ReducerWithInitialState<S>
 }

packages/toolkit/src/createSlice.ts

@@ -8,7 +8,7 @@ import type {
 } from './createAction'
 import { createAction } from './createAction'
 import type { CaseReducer, CaseReducers } from './createReducer'
-import { createReducer } from './createReducer'
+import { createReducer, NotFunction } from './createReducer'
 import type { ActionReducerMapBuilder } from './mapBuilders'
 import { executeReducerBuilderCallback } from './mapBuilders'
 import type { NoInfer } from './tsHelpers'
@@ -53,6 +53,12 @@ export interface Slice<
    * This enables reuse and testing if they were defined inline when calling `createSlice`.
    */
   caseReducers: SliceDefinedCaseReducers<CaseReducers>
+
+  /**
+   * Provides access to the initial state value given to the slice.
+   * If a lazy state initializer was provided, it will be called and a fresh value returned.
+   */
+  getInitialState: () => State
 }
 
 /**
@@ -71,9 +77,9 @@ export interface CreateSliceOptions<
   name: Name
 
   /**
-   * The initial state to be returned by the slice reducer.
+   * The initial state that should be used when the reducer is called the first time. This may also be a "lazy initializer" function, which should return an initial state value when called. This will be used whenever the reducer is called with `undefined` as its state value, and is primarily useful for cases like reading initial state from `localStorage`.
    */
-  initialState: State
+  initialState: State | (() => State)
 
   /**
    * A mapping from action types to action-type-specific *case reducer*
@@ -301,5 +307,6 @@ export function createSlice<
     reducer,
     actions: actionCreators as any,
     caseReducers: sliceCaseReducersByName as any,
+    getInitialState: reducer.getInitialState,
   }
 }

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

@@ -98,8 +98,10 @@ describe('createReducer', () => {
     test('Freezes initial state', () => {
       const initialState = [{ text: 'Buy milk' }]
       const todosReducer = createReducer(initialState, {})
+      const frozenInitialState = todosReducer(undefined, { type: 'dummy' })
 
-      const mutateStateOutsideReducer = () => (initialState[0].text = 'edited')
+      const mutateStateOutsideReducer = () =>
+        (frozenInitialState[0].text = 'edited')
       expect(mutateStateOutsideReducer).toThrowError(
         /Cannot assign to read only property/
       )
@@ -132,6 +134,41 @@ describe('createReducer', () => {
     behavesLikeReducer(todosReducer)
   })
 
+  describe('Accepts a lazy state init function to generate initial state', () => {
+    const addTodo: AddTodoReducer = (state, action) => {
+      const { newTodo } = action.payload
+      state.push({ ...newTodo, completed: false })
+    }
+
+    const toggleTodo: ToggleTodoReducer = (state, action) => {
+      const { index } = action.payload
+      const todo = state[index]
+      todo.completed = !todo.completed
+    }
+
+    const lazyStateInit = () => [] as TodoState
+
+    const todosReducer = createReducer(lazyStateInit, {
+      ADD_TODO: addTodo,
+      TOGGLE_TODO: toggleTodo,
+    })
+
+    behavesLikeReducer(todosReducer)
+
+    it('Should only call the init function when `undefined` state is passed in', () => {
+      const spy = jest.fn().mockReturnValue(42)
+
+      const dummyReducer = createReducer(spy, {})
+      expect(spy).not.toHaveBeenCalled()
+
+      dummyReducer(123, { type: 'dummy' })
+      expect(spy).not.toHaveBeenCalled()
+
+      const initialState = dummyReducer(undefined, { type: 'dummy' })
+      expect(spy).toHaveBeenCalledTimes(1)
+    })
+  })
+
   describe('given draft state from immer', () => {
     const addTodo: AddTodoReducer = (state, action) => {
       const { newTodo } = action.payload

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

@@ -63,6 +63,48 @@ describe('createSlice', () => {
       expect(caseReducers.increment).toBeTruthy()
       expect(typeof caseReducers.increment).toBe('function')
     })
+
+    it('getInitialState should return the state', () => {
+      const initialState = 42
+      const slice = createSlice({
+        name: 'counter',
+        initialState,
+        reducers: {},
+      })
+
+      expect(slice.getInitialState()).toBe(initialState)
+    })
+  })
+
+  describe('when initialState is a function', () => {
+    const initialState = () => ({ user: '' })
+
+    const { actions, reducer } = createSlice({
+      reducers: {
+        setUserName: (state, action) => {
+          state.user = action.payload
+        },
+      },
+      initialState,
+      name: 'user',
+    })
+
+    it('should set the username', () => {
+      expect(reducer(undefined, actions.setUserName('eric'))).toEqual({
+        user: 'eric',
+      })
+    })
+
+    it('getInitialState should return the state', () => {
+      const initialState = () => 42
+      const slice = createSlice({
+        name: 'counter',
+        initialState,
+        reducers: {},
+      })
+
+      expect(slice.getInitialState()).toBe(42)
+    })
   })
 
   describe('when mutating state object', () => {