Don't discard render phase state updates with the eager reducer optimization (facebook#14852)

gaearon · web-flow · commit c506ded3b275 · 2019-02-19T18:40:10.000Z
* Add test cases for setState(fn) + render phase updates

* Update eager state and reducer for render phase updates

* Fix a newly firing warning
diff --git a/packages/react-reconciler/src/ReactFiberHooks.js b/packages/react-reconciler/src/ReactFiberHooks.js
@@ -607,7 +607,6 @@ function updateReducer<S, I, A>(
         }
 
         hook.memoizedState = newState;
-
         // Don't persist the state accumlated from the render phase updates to
         // the base state unless the queue is empty.
         // TODO: Not sure if this is the desired semantics, but it's what we
@@ -616,6 +615,9 @@ function updateReducer<S, I, A>(
           hook.baseState = newState;
         }
 
+        queue.eagerReducer = reducer;
+        queue.eagerState = newState;
+
         return [newState, dispatch];
       }
     }
diff --git a/packages/react-reconciler/src/__tests__/ReactHooks-test.internal.js b/packages/react-reconciler/src/__tests__/ReactHooks-test.internal.js
@@ -669,6 +669,76 @@ describe('ReactHooks', () => {
     }).toThrow('is not a function');
   });
 
+  it('does not forget render phase useState updates inside an effect', () => {
+    const {useState, useEffect} = React;
+
+    function Counter() {
+      const [counter, setCounter] = useState(0);
+      if (counter === 0) {
+        setCounter(x => x + 1);
+        setCounter(x => x + 1);
+      }
+      useEffect(() => {
+        setCounter(x => x + 1);
+        setCounter(x => x + 1);
+      }, []);
+      return counter;
+    }
+
+    const root = ReactTestRenderer.create(null);
+    ReactTestRenderer.act(() => {
+      root.update(<Counter />);
+    });
+    expect(root).toMatchRenderedOutput('4');
+  });
+
+  it('does not forget render phase useReducer updates inside an effect with hoisted reducer', () => {
+    const {useReducer, useEffect} = React;
+
+    const reducer = x => x + 1;
+    function Counter() {
+      const [counter, increment] = useReducer(reducer, 0);
+      if (counter === 0) {
+        increment();
+        increment();
+      }
+      useEffect(() => {
+        increment();
+        increment();
+      }, []);
+      return counter;
+    }
+
+    const root = ReactTestRenderer.create(null);
+    ReactTestRenderer.act(() => {
+      root.update(<Counter />);
+    });
+    expect(root).toMatchRenderedOutput('4');
+  });
+
+  it('does not forget render phase useReducer updates inside an effect with inline reducer', () => {
+    const {useReducer, useEffect} = React;
+
+    function Counter() {
+      const [counter, increment] = useReducer(x => x + 1, 0);
+      if (counter === 0) {
+        increment();
+        increment();
+      }
+      useEffect(() => {
+        increment();
+        increment();
+      }, []);
+      return counter;
+    }
+
+    const root = ReactTestRenderer.create(null);
+    ReactTestRenderer.act(() => {
+      root.update(<Counter />);
+    });
+    expect(root).toMatchRenderedOutput('4');
+  });
+
   it('warns for bad useImperativeHandle first arg', () => {
     const {useImperativeHandle} = React;
     function App() {
diff --git a/packages/react-reconciler/src/__tests__/ReactHooksWithNoopRenderer-test.internal.js b/packages/react-reconciler/src/__tests__/ReactHooksWithNoopRenderer-test.internal.js
@@ -454,7 +454,9 @@ describe('ReactHooksWithNoopRenderer', () => {
 
       // Test that it works on update, too. This time the log is a bit different
       // because we started with reducerB instead of reducerA.
-      counter.current.dispatch('reset');
+      ReactNoop.act(() => {
+        counter.current.dispatch('reset');
+      });
       ReactNoop.render(<Counter ref={counter} />);
       expect(ReactNoop.flush()).toEqual([
         'Render: 0',

Original file line number	Diff line number	Diff line change
`@@ -607,7 +607,6 @@ function updateReducer<S, I, A>(`
`607`	`607`	`}`
`608`	`608`
`609`	`609`	`hook.memoizedState = newState;`
`610`		`-`
`611`	`610`	`// Don't persist the state accumlated from the render phase updates to`
`612`	`611`	`// the base state unless the queue is empty.`
`613`	`612`	`// TODO: Not sure if this is the desired semantics, but it's what we`
`@@ -616,6 +615,9 @@ function updateReducer<S, I, A>(`
`616`	`615`	`hook.baseState = newState;`
`617`	`616`	`}`
`618`	`617`
	`618`	`+ queue.eagerReducer = reducer;`
	`619`	`+ queue.eagerState = newState;`
	`620`	`+`
`619`	`621`	`return [newState, dispatch];`
`620`	`622`	`}`
`621`	`623`	`}`