Move context handling back into Begin and Complete phases

As per discussion, this is better because then this code only runs for the class type of work. We will still need to fix bailouts for deep setState calls.

src/renderers/shared/fiber/ReactFiberBeginWork.js

@@ -27,6 +27,9 @@ var {
 var ReactTypeOfWork = require('ReactTypeOfWork');
 var {
   getMaskedContext,
+  isContextProvider,
+  pushContextProvider,
+  resetContext,
 } = require('ReactFiberContext');
 var {
   IndeterminateComponent,
@@ -195,6 +198,10 @@ module.exports = function<T, P, I, TI, C>(
     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);
+    }
     return workInProgress.child;
   }
 
@@ -352,6 +359,10 @@ module.exports = function<T, P, I, TI, C>(
 
     cloneChildFibers(current, workInProgress);
     markChildAsProgressed(current, workInProgress, priorityLevel);
+    // Put context on the stack because we will work on children
+    if (isContextProvider(workInProgress)) {
+      pushContextProvider(workInProgress);
+    }
     return workInProgress.child;
   }
 
@@ -362,6 +373,11 @@ module.exports = function<T, P, I, TI, C>(
   }
 
   function beginWork(current : ?Fiber, workInProgress : Fiber, priorityLevel : PriorityLevel) : ?Fiber {
+    if (!workInProgress.return) {
+      // Don't start new work with context on the stack.
+      resetContext();
+    }
+
     if (workInProgress.pendingWorkPriority === NoWork ||
         workInProgress.pendingWorkPriority > priorityLevel) {
       return bailoutOnLowPriority(current, workInProgress);

src/renderers/shared/fiber/ReactFiberCompleteWork.js

@@ -18,6 +18,10 @@ import type { HostConfig } from 'ReactFiberReconciler';
 import type { ReifiedYield } from 'ReactReifiedYield';
 
 var { reconcileChildFibers } = require('ReactChildFiber');
+var {
+  isContextProvider,
+  popContextProvider,
+} = require('ReactFiberContext');
 var ReactTypeOfWork = require('ReactTypeOfWork');
 var ReactTypeOfSideEffect = require('ReactTypeOfSideEffect');
 var {
@@ -125,6 +129,10 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
         return null;
       case ClassComponent:
         transferOutput(workInProgress.child, workInProgress);
+        // We are leaving this subtree, so pop context if any.
+        if (isContextProvider(workInProgress)) {
+          popContextProvider();
+        }
         // 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/ReactFiberScheduler.js

@@ -24,13 +24,6 @@ var ReactCurrentOwner = require('ReactCurrentOwner');
 
 var { cloneFiber } = require('ReactFiber');
 
-var {
-  isContextProvider,
-  pushContextProvider,
-  popContextProvider,
-  resetContext,
-} = require('ReactFiberContext');
-
 var {
   NoWork,
   SynchronousPriority,
@@ -274,11 +267,6 @@ 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
@@ -357,11 +345,6 @@ 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
@@ -377,12 +360,7 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
       ReactFiberInstrumentation.debugTool.onDidBeginWork(workInProgress);
     }
 
-    if (next) {
-      // There is work deeper in the tree, so push the context if it exists.
-      if (isContextProvider(workInProgress)) {
-        pushContextProvider(workInProgress);
-      }
-    } else {
+    if (!next) {
       if (__DEV__ && ReactFiberInstrumentation.debugTool) {
         ReactFiberInstrumentation.debugTool.onWillCompleteWork(workInProgress);
       }