From b86cdde13f94062b6290cb0a463fc9954a0b7b0e Mon Sep 17 00:00:00 2001
From: Brian Vaughn <bvaughn@fb.com>
Date: Mon, 19 Dec 2016 16:14:19 -0800
Subject: [PATCH] Fix context handling for errors thrown in ClassComponent
 render Context providers are currently popped one too many times if an error
 is thrown by a context-provider's render method. In this case, the begin
 phase is exited before context is pushed but we still pop it in. This PR adds
 a test for this behavior (failing before) and fixes it by ensuring that we
 always push context before continuing (or unwinding).

---
 scripts/fiber/tests-passing.txt               |  1 +
 .../shared/fiber/ReactFiberBeginWork.js       | 18 ++++---
 .../fiber/__tests__/ReactIncremental-test.js  | 47 +++++++++++++++++++
 3 files changed, 59 insertions(+), 7 deletions(-)

diff --git a/scripts/fiber/tests-passing.txt b/scripts/fiber/tests-passing.txt
index 586dd4f7e35e..de50e0a963d3 100644
--- a/scripts/fiber/tests-passing.txt
+++ b/scripts/fiber/tests-passing.txt
@@ -1157,6 +1157,7 @@ src/renderers/shared/fiber/__tests__/ReactIncremental-test.js
 * reads context when setState is below the provider
 * reads context when setState is above the provider
 * maintains the correct context index when context proviers are bailed out due to low priority
+* maintains the correct context index when unwinding due to an error ocurring during render
 
 src/renderers/shared/fiber/__tests__/ReactIncrementalErrorHandling-test.js
 * catches render error in a boundary during full deferred mounting
diff --git a/src/renderers/shared/fiber/ReactFiberBeginWork.js b/src/renderers/shared/fiber/ReactFiberBeginWork.js
index 35b3570db058..f1b357ab752c 100644
--- a/src/renderers/shared/fiber/ReactFiberBeginWork.js
+++ b/src/renderers/shared/fiber/ReactFiberBeginWork.js
@@ -240,13 +240,17 @@ module.exports = function<T, P, I, TI, C, CX>(
     }
 
     // Rerender
-    const instance = workInProgress.stateNode;
-    ReactCurrentOwner.current = workInProgress;
-    const nextChildren = instance.render();
-    reconcileChildren(current, workInProgress, nextChildren);
-    // Put context on the stack because we will work on children
-    if (isContextProvider(workInProgress)) {
-      pushContextProvider(workInProgress, true);
+    try {
+      const instance = workInProgress.stateNode;
+      ReactCurrentOwner.current = workInProgress;
+      const nextChildren = instance.render();
+      reconcileChildren(current, workInProgress, nextChildren);
+    } finally {
+      // Put context on the stack because we will work on children.
+      // Push even if Error in render() because unwindContext() always pops.
+      if (isContextProvider(workInProgress)) {
+        pushContextProvider(workInProgress, true);
+      }
     }
     return workInProgress.child;
   }
diff --git a/src/renderers/shared/fiber/__tests__/ReactIncremental-test.js b/src/renderers/shared/fiber/__tests__/ReactIncremental-test.js
index 3a91f6c3cc4a..513c21546d4d 100644
--- a/src/renderers/shared/fiber/__tests__/ReactIncremental-test.js
+++ b/src/renderers/shared/fiber/__tests__/ReactIncremental-test.js
@@ -1975,4 +1975,51 @@ describe('ReactIncremental', () => {
     instance.setState({});
     ReactNoop.flush();
   });
+
+  it('maintains the correct context index when unwinding due to an error ocurring during render', () => {
+    class Root extends React.Component {
+      unstable_handleError(error) {
+        // If context is pushed/popped correctly,
+        // This method will be used to handle the intentionally-thrown Error.
+      }
+      render() {
+        return <ContextProvider depth={1} />;
+      }
+    }
+
+    let instance;
+
+    class ContextProvider extends React.Component {
+      constructor(props, context) {
+        super(props, context);
+        this.state = {};
+        if (props.depth === 1) {
+          instance = this;
+        }
+      }
+      static childContextTypes = {};
+      getChildContext() {
+        return {};
+      }
+      render() {
+        if (this.state.throwError) {
+          throw Error();
+        }
+        return this.props.depth < 4
+          ? <ContextProvider depth={this.props.depth + 1} />
+          : <div/>;
+      }
+    }
+
+    // Init
+    ReactNoop.render(<Root />);
+    ReactNoop.flush();
+
+    // Trigger an update in the middle of the tree
+    // This is necessary to reproduce the error as it curently exists.
+    instance.setState({
+      throwError: true,
+    });
+    ReactNoop.flush();
+  });
 });