[Fiber] Push class context providers even if they crash (#8627)

* Push class context providers early

Previously we used to push them only after the instance was available. This caused issues in cases an error is thrown during componentWillMount().

In that case we never got to pushing the provider in the begin phase, but in complete phase the provider check returned true since the instance existed by that point. As a result we got mismatching context pops.

We solve the issue by making the context check independent of whether the instance actually exists. Instead we're checking the type itself.

This lets us push class context early. However there's another problem: we might not know the context value. If the instance is not yet created, we can't call getChildContext on it.

To fix this, we are introducing a way to replace current value on the stack, and a way to read the previous value. This also helps remove some branching and split the memoized from invalidated code paths.

* Add a now-passing test from #8604

Also rename another test to have a shorter name.

* Move isContextProvider() checks into push() and pop()

All uses of push() and pop() are guarded by it anyway.

This makes it more similar to how we use host context.

There is only one other place where isContextProvider() is used and that's legacy code needed for renderSubtree().

* Clarify why we read the previous context

* Use invariant instead of throwing an error

* Fix off-by-one in ReactFiberStack

* Manually keep track of the last parent context

The previous algorithm was flawed and worked by accident, as shown by the failing tests after an off-by-one was fixed.

The implementation of getPrevious() was incorrect because the context stack currently has no notion of a previous value per cursor.

Instead, we are caching the previous value directly in the ReactFiberContext in a local variable.

Additionally, we are using push() and pop() instead of adding a new replace() method.

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 (