Fix unwinding starting with a wrong Fiber on error in the complete ph…

…ase (#13237)

* Add a repro case for profiler unwinding

This currently fails the tests due to an unexpected warning.

* Add a regression test for context stack

* Simplify the first test case

* Update nextUnitOfWork inside completeUnitOfWork()

The bug was caused by a structure like this:

    </Provider>
  </div>
</errorInCompletePhase>

We forgot to update nextUnitOfWork so it was still pointing at Provider when errorInCompletePhase threw. As a result, we would try to unwind from Provider (rather than from errorInCompletePhase), and thus pop the Provider twice.

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -791,11 +791,12 @@ function completeUnitOfWork(workInProgress: Fiber): Fiber | null {
 
     if ((workInProgress.effectTag & Incomplete) === NoEffect) {
       // This fiber completed.
-      let next = completeWork(
+      nextUnitOfWork = completeWork(
         current,
         workInProgress,
         nextRenderExpirationTime,
       );
+      let next = nextUnitOfWork;
       stopWorkTimer(workInProgress);
       resetExpirationTime(workInProgress, nextRenderExpirationTime);
       if (__DEV__) {

packages/react-reconciler/src/__tests__/ReactNewContext-test.internal.js

@@ -1128,6 +1128,27 @@ describe('ReactNewContext', () => {
     ]);
   });
 
+  it('unwinds after errors in complete phase', () => {
+    const Context = React.createContext(0);
+
+    // This is a regression test for stack misalignment
+    // caused by unwinding the context from wrong point.
+    ReactNoop.render(
+      <errorInCompletePhase>
+        <Context.Provider />
+      </errorInCompletePhase>,
+    );
+    expect(ReactNoop.flush).toThrow('Error in host config.');
+
+    ReactNoop.render(
+      <Context.Provider value={10}>
+        <Context.Consumer>{value => <span prop={value} />}</Context.Consumer>
+      </Context.Provider>,
+    );
+    ReactNoop.flush();
+    expect(ReactNoop.getChildren()).toEqual([span(10)]);
+  });
+
   describe('fuzz test', () => {
     const Fragment = React.Fragment;
     const contextKeys = ['A', 'B', 'C', 'D', 'E', 'F', 'G'];

packages/react/src/__tests__/ReactProfiler-test.internal.js

@@ -999,6 +999,17 @@ describe('Profiler', () => {
             );
             expect(ReactNoop.flush).toThrow('Error in host config.');
 
+            // A similar case we've seen caused by an invariant in ReactDOM.
+            // It didn't reproduce without a host component inside.
+            ReactNoop.render(
+              <React.unstable_Profiler id="profiler" onRender={jest.fn()}>
+                <errorInCompletePhase>
+                  <span>hi</span>
+                </errorInCompletePhase>
+              </React.unstable_Profiler>,
+            );
+            expect(ReactNoop.flush).toThrow('Error in host config.');
+
             // So long as the profiler timer's fiber stack is reset correctly,
             // Subsequent renders should not error.
             ReactNoop.render(