[Fiber] Push class context providers even if they crash

scripts/fiber/tests-passing.txt

@@ -1157,7 +1157,8 @@ src/renderers/shared/fiber/__tests__/ReactIncremental-test.js
 * provides context when reusing work
 * 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 when providers bail out due to low priority
+* maintains the correct context when unwinding due to an error in render
 
 src/renderers/shared/fiber/__tests__/ReactIncrementalErrorHandling-test.js
 * catches render error in a boundary during full deferred mounting
@@ -1413,6 +1414,8 @@ src/renderers/shared/shared/__tests__/ReactErrorBoundaries-test.js
 * renders an error state if child throws in render
 * renders an error state if child throws in constructor
 * renders an error state if child throws in componentWillMount
+* renders an error state if context provider throws in componentWillMount
+* renders an error state if module-style context provider throws in componentWillMount
 * mounts the error message if mounting fails
 * propagates errors on retry on mounting
 * propagates errors inside boundary during componentWillMount

src/renderers/shared/fiber/ReactFiberBeginWork.js

@@ -31,10 +31,10 @@ var {
 var ReactTypeOfWork = require('ReactTypeOfWork');
 var {
   getMaskedContext,
-  isContextProvider,
   hasContextChanged,
   pushContextProvider,
   pushTopLevelContextObject,
+  invalidateContextProvider,
 } = require('ReactFiberContext');
 var {
   IndeterminateComponent,
@@ -215,6 +215,11 @@ module.exports = function<T, P, I, TI, C, CX>(
   }
 
   function updateClassComponent(current : ?Fiber, workInProgress : Fiber, priorityLevel : PriorityLevel) {
+    // Push context providers early to prevent context stack mismatches.
+    // During mounting we don't know the child context yet as the instance doesn't exist.
+    // We will invalidate the child context in finishClassComponent() right after rendering.
+    const hasContext = pushContextProvider(workInProgress);
+
     let shouldUpdate;
     if (!current) {
       if (!workInProgress.stateNode) {
@@ -229,10 +234,15 @@ module.exports = function<T, P, I, TI, C, CX>(
     } else {
       shouldUpdate = updateClassInstance(current, workInProgress, priorityLevel);
     }
-    return finishClassComponent(current, workInProgress, shouldUpdate);
+    return finishClassComponent(current, workInProgress, shouldUpdate, hasContext);
   }
 
-  function finishClassComponent(current : ?Fiber, workInProgress : Fiber, shouldUpdate : boolean) {
+  function finishClassComponent(
+    current : ?Fiber,
+    workInProgress : Fiber,
+    shouldUpdate : boolean,
+    hasContext : boolean,
+  ) {
     // Schedule side-effects
     if (shouldUpdate) {
       workInProgress.effectTag |= Update;
@@ -246,11 +256,6 @@ module.exports = function<T, P, I, TI, C, CX>(
           workInProgress.effectTag |= Update;
         }
       }
-
-      // Don't forget to push the context before returning.
-      if (isContextProvider(workInProgress)) {
-        pushContextProvider(workInProgress, false);
-      }
       return bailoutOnAlreadyFinishedWork(current, workInProgress);
     }
 
@@ -259,9 +264,10 @@ module.exports = function<T, P, I, TI, C, CX>(
     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);
+
+    // The context might have changed so we need to recalculate it.
+    if (hasContext) {
+      invalidateContextProvider(workInProgress);
     }
     return workInProgress.child;
   }
@@ -417,9 +423,14 @@ module.exports = function<T, P, I, TI, C, CX>(
     if (typeof value === 'object' && value && typeof value.render === 'function') {
       // Proceed under the assumption that this is a class instance
       workInProgress.tag = ClassComponent;
+
+      // Push context providers early to prevent context stack mismatches.
+      // During mounting we don't know the child context yet as the instance doesn't exist.
+      // We will invalidate the child context in finishClassComponent() right after rendering.
+      const hasContext = pushContextProvider(workInProgress);
       adoptClassInstance(workInProgress, value);
       mountClassInstance(workInProgress, priorityLevel);
-      return finishClassComponent(current, workInProgress, true);
+      return finishClassComponent(current, workInProgress, true, hasContext);
     } else {
       // Proceed under the assumption that this is a functional component
       workInProgress.tag = FunctionalComponent;
@@ -535,9 +546,7 @@ module.exports = function<T, P, I, TI, C, CX>(
     // See PR 8590 discussion for context
     switch (workInProgress.tag) {
       case ClassComponent:
-        if (isContextProvider(workInProgress)) {
-          pushContextProvider(workInProgress, false);
-        }
+        pushContextProvider(workInProgress);
         break;
       case HostPortal:
         pushHostContainer(workInProgress, workInProgress.stateNode.containerInfo);

src/renderers/shared/fiber/ReactFiberCompleteWork.js

@@ -21,7 +21,6 @@ import type { ReifiedYield } from 'ReactReifiedYield';
 
 var { reconcileChildFibers } = require('ReactChildFiber');
 var {
-  isContextProvider,
   popContextProvider,
 } = require('ReactFiberContext');
 var ReactTypeOfWork = require('ReactTypeOfWork');
@@ -175,9 +174,7 @@ module.exports = function<T, P, I, TI, C, CX>(
         return null;
       case ClassComponent: {
         // We are leaving this subtree, so pop context if any.
-        if (isContextProvider(workInProgress)) {
-          popContextProvider(workInProgress);
-        }
+        popContextProvider(workInProgress);
         // Don't use the state queue to compute the memoized state. We already
         // merged it and assigned it to the instance. Transfer it from there.
         // Also need to transfer the props, because pendingProps will be null

src/renderers/shared/fiber/ReactFiberContext.js

@@ -35,11 +35,25 @@ if (__DEV__) {
   var checkReactTypeSpec = require('checkReactTypeSpec');
 }
 
-let contextStackCursor : StackCursor<?Object> = createCursor((null: ?Object));
+// A cursor to the current merged context object on the stack.
+let contextStackCursor : StackCursor<Object> = createCursor(emptyObject);
+// A cursor to a boolean indicating whether the context has changed.
 let didPerformWorkStackCursor : StackCursor<boolean> = createCursor(false);
-
-function getUnmaskedContext() {
-  return contextStackCursor.current || emptyObject;
+// Keep track of the previous context object that was on the stack.
+// We use this to get access to the parent context after we have already
+// pushed the next context provider, and now need to merge their contexts.
+let previousContext : Object = emptyObject;
+
+function getUnmaskedContext(workInProgress : Fiber) : Object {
+  const hasOwnContext = isContextProvider(workInProgress);
+  if (hasOwnContext) {
+    // If the fiber is a context provider itself, when we read its context
+    // we have already pushed its own child context on the stack. A context
+    // provider should not "see" its own child context. Therefore we read the
+    // previous (parent) context instead for a context provider.
+    return previousContext;
+  }
+  return contextStackCursor.current;
 }
 
 exports.getMaskedContext = function(workInProgress : Fiber) {
@@ -49,9 +63,8 @@ exports.getMaskedContext = function(workInProgress : Fiber) {
     return emptyObject;
   }
 
-  const unmaskedContext = getUnmaskedContext();
+  const unmaskedContext = getUnmaskedContext(workInProgress);
   const context = {};
-
   for (let key in contextTypes) {
     context[key] = unmaskedContext[key];
   }
@@ -71,14 +84,16 @@ exports.hasContextChanged = function() : boolean {
 function isContextProvider(fiber : Fiber) : boolean {
   return (
     fiber.tag === ClassComponent &&
-    // Instance might be null, if the fiber errored during construction
-    fiber.stateNode &&
-    typeof fiber.stateNode.getChildContext === 'function'
+    fiber.type.childContextTypes != null
   );
 }
 exports.isContextProvider = isContextProvider;
 
 function popContextProvider(fiber : Fiber) : void {
+  if (!isContextProvider(fiber)) {
+    return;
+  }
+
   pop(didPerformWorkStackCursor, fiber);
   pop(contextStackCursor, fiber);
 }
@@ -117,25 +132,48 @@ function processChildContext(fiber : Fiber, parentContext : Object, isReconcilin
 }
 exports.processChildContext = processChildContext;
 
-exports.pushContextProvider = function(workInProgress : Fiber, didPerformWork : boolean) : void {
-  const instance = workInProgress.stateNode;
-  const memoizedMergedChildContext = instance.__reactInternalMemoizedMergedChildContext;
-  const canReuseMergedChildContext = !didPerformWork && memoizedMergedChildContext != null;
-
-  let mergedContext = null;
-  if (canReuseMergedChildContext) {
-    mergedContext = memoizedMergedChildContext;
-  } else {
-    mergedContext = processChildContext(workInProgress, getUnmaskedContext(), true);
-    instance.__reactInternalMemoizedMergedChildContext = mergedContext;
+exports.pushContextProvider = function(workInProgress : Fiber) : boolean {
+  if (!isContextProvider(workInProgress)) {
+    return false;
   }
 
+  const instance = workInProgress.stateNode;
+  // We push the context as early as possible to ensure stack integrity.
+  // If the instance does not exist yet, we will push null at first,
+  // and replace it on the stack later when invalidating the context.
+  const memoizedMergedChildContext = (
+    instance &&
+    instance.__reactInternalMemoizedMergedChildContext
+  ) || emptyObject;
+
+  // Remember the parent context so we can merge with it later.
+  previousContext = contextStackCursor.current;
+  push(contextStackCursor, memoizedMergedChildContext, workInProgress);
+  push(didPerformWorkStackCursor, false, workInProgress);
+
+  return true;
+};
+
+exports.invalidateContextProvider = function(workInProgress : Fiber) : void {
+  const instance = workInProgress.stateNode;
+  invariant(instance, 'Expected to have an instance by this point.');
+
+  // Merge parent and own context.
+  const mergedContext = processChildContext(workInProgress, previousContext, true);
+  instance.__reactInternalMemoizedMergedChildContext = mergedContext;
+
+  // Replace the old (or empty) context with the new one.
+  // It is important to unwind the context in the reverse order.
+  pop(didPerformWorkStackCursor, workInProgress);
+  pop(contextStackCursor, workInProgress);
+  // Now push the new context and mark that it has changed.
   push(contextStackCursor, mergedContext, workInProgress);
-  push(didPerformWorkStackCursor, didPerformWork, workInProgress);
+  push(didPerformWorkStackCursor, true, workInProgress);
 };
 
 exports.resetContext = function() : void {
-  contextStackCursor.current = null;
+  previousContext = emptyObject;
+  contextStackCursor.current = emptyObject;
   didPerformWorkStackCursor.current = false;
 };
 

src/renderers/shared/fiber/ReactFiberScheduler.js

@@ -18,7 +18,6 @@ import type { HostConfig, Deadline } from 'ReactFiberReconciler';
 import type { PriorityLevel } from 'ReactPriorityLevel';
 
 var {
-  isContextProvider,
   popContextProvider,
 } = require('ReactFiberContext');
 const { reset } = require('ReactFiberStack');
@@ -957,9 +956,7 @@ module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C
     while (node && (node !== to) && (node.alternate !== to)) {
       switch (node.tag) {
         case ClassComponent:
-          if (isContextProvider(node)) {
-            popContextProvider(node);
-          }
+          popContextProvider(node);
           break;
         case HostComponent:
           popHostContext(node);

src/renderers/shared/fiber/__tests__/ReactIncremental-test.js

@@ -1932,7 +1932,7 @@ describe('ReactIncremental', () => {
     ]);
   });
 
-  it('maintains the correct context index when context proviers are bailed out due to low priority', () => {
+  it('maintains the correct context when providers bail out due to low priority', () => {
     class Root extends React.Component {
       render() {
         return <Middle {...this.props} />;
@@ -1974,4 +1974,51 @@ describe('ReactIncremental', () => {
     instance.setState({});
     ReactNoop.flush();
   });
+
+  it('maintains the correct context when unwinding due to an error in 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();
+  });
 });

src/renderers/shared/shared/__tests__/ReactErrorBoundaries-test.js

@@ -733,6 +733,58 @@ describe('ReactErrorBoundaries', () => {
     ]);
   });
 
+  it('renders an error state if context provider throws in componentWillMount', () => {
+    class BrokenComponentWillMountWithContext extends React.Component {
+      static childContextTypes = {foo: React.PropTypes.number};
+      getChildContext() {
+        return {foo: 42};
+      }
+      render() {
+        return <div>{this.props.children}</div>;
+      }
+      componentWillMount() {
+        throw new Error('Hello');
+      }
+    }
+
+    var container = document.createElement('div');
+    ReactDOM.render(
+      <ErrorBoundary>
+        <BrokenComponentWillMountWithContext />
+      </ErrorBoundary>,
+      container
+    );
+    expect(container.firstChild.textContent).toBe('Caught an error: Hello.');
+  });
+
+  it('renders an error state if module-style context provider throws in componentWillMount', () => {
+    function BrokenComponentWillMountWithContext() {
+      return {
+        getChildContext() {
+          return {foo: 42};
+        },
+        render() {
+          return <div>{this.props.children}</div>;
+        },
+        componentWillMount() {
+          throw new Error('Hello');
+        },
+      };
+    }
+    BrokenComponentWillMountWithContext.childContextTypes = {
+      foo: React.PropTypes.number,
+    };
+
+    var container = document.createElement('div');
+    ReactDOM.render(
+      <ErrorBoundary>
+        <BrokenComponentWillMountWithContext />
+      </ErrorBoundary>,
+      container
+    );
+    expect(container.firstChild.textContent).toBe('Caught an error: Hello.');
+  });
+
   it('mounts the error message if mounting fails', () => {
     function renderError(error) {
       return (