Move context management into scheduler

It is error-prone to push and pop context in begin or complete phases because of all the bailouts. Scheduler has enough knowledge to see if pushing is necessary because it knows when we go inside or outside the tree.

Author: gaearon

scripts/fiber/tests-passing.txt

@@ -820,6 +820,8 @@ src/renderers/shared/fiber/__tests__/ReactIncremental-test.js
 * can nest batchedUpdates
 * can handle if setState callback throws
 * merges and masks context
+* does not leak own context into context provider
+* provides context when reusing work
 
 src/renderers/shared/fiber/__tests__/ReactIncrementalErrorHandling-test.js
 * catches render error in a boundary during mounting

src/renderers/shared/fiber/ReactFiberBeginWork.js

@@ -27,8 +27,6 @@ var {
 var ReactTypeOfWork = require('ReactTypeOfWork');
 var {
   getMaskedContext,
-  pushContextProvider,
-  resetContext,
 } = require('ReactFiberContext');
 var {
   IndeterminateComponent,
@@ -189,14 +187,11 @@ module.exports = function<T, P, I, TI, C>(
     } else {
       shouldUpdate = updateClassInstance(current, workInProgress);
     }
-    const instance = workInProgress.stateNode;
-    if (typeof instance.getChildContext === 'function') {
-      pushContextProvider(workInProgress);
-    }
     if (!shouldUpdate) {
       return bailoutOnAlreadyFinishedWork(current, workInProgress);
     }
     // Rerender
+    const instance = workInProgress.stateNode;
     ReactCurrentOwner.current = workInProgress;
     const nextChildren = instance.render();
     reconcileChildren(current, workInProgress, nextChildren);
@@ -367,9 +362,6 @@ module.exports = function<T, P, I, TI, C>(
   }
 
   function beginWork(current : ?Fiber, workInProgress : Fiber, priorityLevel : PriorityLevel) : ?Fiber {
-    if (!workInProgress.return) {
-      resetContext();
-    }
     if (workInProgress.pendingWorkPriority === NoWork ||
         workInProgress.pendingWorkPriority > priorityLevel) {
       return bailoutOnLowPriority(current, workInProgress);

src/renderers/shared/fiber/ReactFiberCompleteWork.js

@@ -18,7 +18,6 @@ import type { HostConfig } from 'ReactFiberReconciler';
 import type { ReifiedYield } from 'ReactReifiedYield';
 
 var { reconcileChildFibers } = require('ReactChildFiber');
-var { popContextProvider } = require('ReactFiberContext');
 var ReactTypeOfWork = require('ReactTypeOfWork');
 var ReactTypeOfSideEffect = require('ReactTypeOfSideEffect');
 var {
@@ -121,11 +120,10 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
 
   function completeWork(current : ?Fiber, workInProgress : Fiber) : ?Fiber {
     switch (workInProgress.tag) {
-      case FunctionalComponent: {
+      case FunctionalComponent:
         transferOutput(workInProgress.child, workInProgress);
         return null;
-      }
-      case ClassComponent: {
+      case ClassComponent:
         transferOutput(workInProgress.child, 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.
@@ -150,13 +148,8 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
           workInProgress.callbackList = updateQueue;
           markCallback(workInProgress);
         }
-        const instance = workInProgress.stateNode;
-        if (typeof instance.getChildContext === 'function') {
-          popContextProvider();
-        }
         return null;
-      }
-      case HostContainer: {
+      case HostContainer:
         transferOutput(workInProgress.child, workInProgress);
         // We don't know if a container has updated any children so we always
         // need to update it right now. We schedule this side-effect before
@@ -165,8 +158,7 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
         // are invoked.
         markUpdate(workInProgress);
         return null;
-      }
-      case HostComponent: {
+      case HostComponent:
         let newProps = workInProgress.pendingProps;
         if (current && workInProgress.stateNode != null) {
           // If we have an alternate, that means this is an update and we need to
@@ -208,8 +200,7 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
         }
         workInProgress.memoizedProps = newProps;
         return null;
-      }
-      case HostText: {
+      case HostText:
         let newText = workInProgress.pendingProps;
         if (current && workInProgress.stateNode != null) {
           // If we have an alternate, that means this is an update and we need to
@@ -231,32 +222,25 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
         }
         workInProgress.memoizedProps = newText;
         return null;
-      }
-      case CoroutineComponent: {
+      case CoroutineComponent:
         return moveCoroutineToHandlerPhase(current, workInProgress);
-      }
-      case CoroutineHandlerPhase: {
+      case CoroutineHandlerPhase:
         transferOutput(workInProgress.stateNode, workInProgress);
         // Reset the tag to now be a first phase coroutine.
         workInProgress.tag = CoroutineComponent;
         return null;
-      }
-      case YieldComponent: {
+      case YieldComponent:
         // Does nothing.
         return null;
-      }
-      case Fragment: {
+      case Fragment:
         transferOutput(workInProgress.child, workInProgress);
         return null;
-      }
 
       // Error cases
-      case IndeterminateComponent: {
+      case IndeterminateComponent:
         throw new Error('An indeterminate component should have become determinate before completing.');
-      }
-      default: {
+      default:
         throw new Error('Unknown unit of work tag');
-      }
     }
   }
 

src/renderers/shared/fiber/ReactFiberContext.js

@@ -19,6 +19,9 @@ var invariant = require('invariant');
 var {
   getComponentName,
 } = require('ReactFiberTreeReflection');
+var {
+  ClassComponent,
+} = require('ReactTypeOfWork');
 
 if (__DEV__) {
   var checkReactTypeSpec = require('checkReactTypeSpec');
@@ -56,12 +59,19 @@ exports.getMaskedContext = function(fiber : Fiber) {
   return context;
 };
 
-exports.popContextProvider = function() {
+exports.isContextProvider = function(fiber : Fiber) : boolean {
+  return (
+    fiber.tag === ClassComponent &&
+    typeof fiber.stateNode.getChildContext === 'function'
+  );
+};
+
+exports.popContextProvider = function() : void {
   stack[index] = emptyObject;
   index--;
 };
 
-exports.pushContextProvider = function(fiber : Fiber) {
+exports.pushContextProvider = function(fiber : Fiber) : void {
   const instance = fiber.stateNode;
   const childContextTypes = fiber.type.childContextTypes;
   const childContext = instance.getChildContext();
@@ -85,7 +95,7 @@ exports.pushContextProvider = function(fiber : Fiber) {
   stack[index] = mergedContext;
 };
 
-exports.resetContext = function() {
+exports.resetContext = function() : void {
   index = -1;
 };
 

src/renderers/shared/fiber/ReactFiberScheduler.js

@@ -24,6 +24,13 @@ var ReactCurrentOwner = require('ReactCurrentOwner');
 
 var { cloneFiber } = require('ReactFiber');
 
+var {
+  isContextProvider,
+  pushContextProvider,
+  popContextProvider,
+  resetContext,
+} = require('ReactFiberContext');
+
 var {
   NoWork,
   SynchronousPriority,
@@ -267,6 +274,11 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
       const current = workInProgress.alternate;
       const next = completeWork(current, workInProgress);
 
+      // We are leaving this subtree, so pop context if any.
+      if (isContextProvider(workInProgress)) {
+        popContextProvider();
+      }
+
       resetWorkPriority(workInProgress);
 
       // The work is now done. We don't need this anymore. This flags
@@ -345,6 +357,11 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
   }
 
   function performUnitOfWork(workInProgress : Fiber) : ?Fiber {
+    if (!workInProgress.return) {
+      // Don't start new work with context on the stack.
+      resetContext();
+    }
+
     // The current, flushed, state of this fiber is the alternate.
     // Ideally nothing should rely on this, but relying on it here
     // means that we don't need an additional field on the work in
@@ -360,7 +377,12 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
       ReactFiberInstrumentation.debugTool.onDidBeginWork(workInProgress);
     }
 
-    if (!next) {
+    if (next) {
+      // There is work deeper in the tree, so push the context if it exists.
+      if (isContextProvider(workInProgress)) {
+        pushContextProvider(workInProgress);
+      }
+    } else {
       if (__DEV__ && ReactFiberInstrumentation.debugTool) {
         ReactFiberInstrumentation.debugTool.onWillCompleteWork(workInProgress);
       }