Allow nested Immer produce statements with createReducer (#509)

* Allow nested produce statements without swallowing in the innermost produce's changes

* Resolve linting errors

* Update src/createReducer.ts

Co-Authored-By: Jonas Holst Damtoft <[email protected]>

* Revised code comments, reverted nullish coalesce

* Add unit tests for nested immer produce statements

Co-authored-by: Andrew Craswell <[email protected]>
Co-authored-by: Jonas Holst Damtoft <[email protected]>

Author: 3 people

src/createReducer.test.ts

@@ -1,5 +1,6 @@
 import { createReducer, CaseReducer } from './createReducer'
 import { PayloadAction, createAction } from './createAction'
+import { createNextState, Draft } from './'
 import { Reducer } from 'redux'
 
 interface Todo {
@@ -75,6 +76,36 @@ describe('createReducer', () => {
     behavesLikeReducer(todosReducer)
   })
 
+  describe('given draft state from immer', () => {
+    const addTodo: AddTodoReducer = (state, action) => {
+      const { newTodo } = action.payload
+
+      // Can safely call state.push() here
+      state.push({ ...newTodo, completed: false })
+    }
+
+    const toggleTodo: ToggleTodoReducer = (state, action) => {
+      const { index } = action.payload
+
+      const todo = state[index]
+      // Can directly modify the todo object
+      todo.completed = !todo.completed
+    }
+
+    const todosReducer = createReducer([] as TodoState, {
+      ADD_TODO: addTodo,
+      TOGGLE_TODO: toggleTodo
+    })
+
+    const wrappedReducer: TodosReducer = (state = [], action) => {
+      return createNextState(state, (draft: Draft<TodoState>) => {
+        todosReducer(draft, action)
+      })
+    }
+
+    behavesLikeReducer(wrappedReducer)
+  })
+
   describe('alternative builder callback for actionMap', () => {
     const increment = createAction<number, 'increment'>('increment')
     const decrement = createAction<number, 'decrement'>('decrement')

src/createReducer.ts

@@ -1,4 +1,4 @@
-import createNextState, { Draft } from 'immer'
+import createNextState, { Draft, isDraft } from 'immer'
 import { AnyAction, Action, Reducer } from 'redux'
 import {
   executeReducerBuilderCallback,
@@ -105,12 +105,24 @@ export function createReducer<S>(
       : mapOrBuilderCallback
 
   return function(state = initialState, action): S {
-    // @ts-ignore createNextState() produces an Immutable<Draft<S>> rather
-    // than an Immutable<S>, and TypeScript cannot find out how to reconcile
-    // these two types.
-    return createNextState(state, (draft: Draft<S>) => {
-      const caseReducer = actionsMap[action.type]
-      return caseReducer ? caseReducer(draft, action) : undefined
-    })
+    const caseReducer = actionsMap[action.type]
+    if (caseReducer) {
+      if (isDraft(state)) {
+        // we must already be inside a `createNextState` call, likely because
+        // this is being wrapped in `createReducer`, `createSlice`, or nested
+        // inside an existing draft. It's safe to just pass the draft to the mutator.
+        const draft = state as Draft<S> // We can aassume this is already a draft
+        return caseReducer(draft, action) || state
+      } else {
+        // @ts-ignore createNextState() produces an Immutable<Draft<S>> rather
+        // than an Immutable<S>, and TypeScript cannot find out how to reconcile
+        // these two types.
+        return createNextState(state, (draft: Draft<S>) => {
+          return caseReducer(draft, action)
+        })
+      }
+    }
+
+    return state
   }
 }