Added ReactFiberStack shared by ReactFiberContext and ReactFiberHostContext

ReactFiberStack is the new underlying stack used by ReactFiberContext and ReactFiberHostContext. The goal is to simplify the 2 context managers and to add more errors/dev-warnings when we pop unexpected.
This changeset currently causes a lot of tests to fail (as we are currently popping too many times and/or in the wrong order). I'm pushing this branch now to share with Sebastian as he is working on a related cleanup pass at beginWork.

src/renderers/shared/fiber/ReactFiberBeginWork.js

@@ -270,14 +270,19 @@ module.exports = function<T, P, I, TI, C, CX>(
     const root = (workInProgress.stateNode : FiberRoot);
     if (root.pendingContext) {
       pushTopLevelContextObject(
+        workInProgress,
         root.pendingContext,
         root.pendingContext !== root.context
       );
     } else {
-      pushTopLevelContextObject(root.context, false);
+      pushTopLevelContextObject(
+        workInProgress,
+        root.context,
+        false
+      );
     }
 
-    pushHostContainer(root.containerInfo);
+    pushHostContainer(workInProgress, root.containerInfo);
 
     const updateQueue = workInProgress.updateQueue;
     if (updateQueue) {
@@ -444,7 +449,7 @@ module.exports = function<T, P, I, TI, C, CX>(
   }
 
   function updatePortalComponent(current, workInProgress) {
-    pushHostContainer(workInProgress.stateNode.containerInfo);
+    pushHostContainer(workInProgress, workInProgress.stateNode.containerInfo);
     const priorityLevel = workInProgress.pendingWorkPriority;
     let nextChildren = workInProgress.pendingProps;
     if (hasContextChanged()) {
@@ -535,7 +540,7 @@ module.exports = function<T, P, I, TI, C, CX>(
         }
         break;
       case HostPortal:
-        pushHostContainer(workInProgress.stateNode.containerInfo);
+        pushHostContainer(workInProgress, workInProgress.stateNode.containerInfo);
         break;
     }
     // TODO: What if this is currently in progress?
@@ -628,4 +633,4 @@ module.exports = function<T, P, I, TI, C, CX>(
     beginFailedWork,
   };
 
-};
+};

src/renderers/shared/fiber/ReactFiberCompleteWork.js

@@ -176,7 +176,7 @@ module.exports = function<T, P, I, TI, C, CX>(
       case ClassComponent: {
         // We are leaving this subtree, so pop context if any.
         if (isContextProvider(workInProgress)) {
-          popContextProvider();
+          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.

src/renderers/shared/fiber/ReactFiberContext.js

@@ -13,6 +13,7 @@
 'use strict';
 
 import type { Fiber } from 'ReactFiber';
+import type { StackCursor } from 'ReactFiberStack';
 
 var emptyObject = require('emptyObject');
 var invariant = require('invariant');
@@ -24,20 +25,22 @@ var {
   ClassComponent,
   HostRoot,
 } = require('ReactTypeOfWork');
+const {
+  createCursor,
+  pop,
+  push,
+  reset,
+} = require('ReactFiberStack');
 
 if (__DEV__) {
   var checkReactTypeSpec = require('checkReactTypeSpec');
 }
 
-let index = -1;
-const contextStack : Array<Object> = [];
-const didPerformWorkStack : Array<boolean> = [];
+let contextStackCursor : StackCursor = createCursor();
+let didPerformWorkStackCursor : StackCursor = createCursor();
 
 function getUnmaskedContext() {
-  if (index === -1) {
-    return emptyObject;
-  }
-  return contextStack[index];
+  return contextStackCursor.current;
 }
 
 exports.getMaskedContext = function(workInProgress : Fiber) {
@@ -62,7 +65,7 @@ exports.getMaskedContext = function(workInProgress : Fiber) {
 };
 
 exports.hasContextChanged = function() : boolean {
-  return index > -1 && didPerformWorkStack[index];
+  return Boolean(didPerformWorkStackCursor.current);
 };
 
 function isContextProvider(fiber : Fiber) : boolean {
@@ -75,19 +78,17 @@ function isContextProvider(fiber : Fiber) : boolean {
 }
 exports.isContextProvider = isContextProvider;
 
-function popContextProvider() : void {
-  invariant(index > -1, 'Unexpected context pop');
-  contextStack[index] = emptyObject;
-  didPerformWorkStack[index] = false;
-  index--;
+function popContextProvider(fiber : Fiber) : void {
+  pop(didPerformWorkStackCursor, fiber);
+  pop(contextStackCursor, fiber);
 }
 exports.popContextProvider = popContextProvider;
 
-exports.pushTopLevelContextObject = function(context : Object, didChange : boolean) : void {
-  invariant(index === -1, 'Unexpected context found on stack');
-  index++;
-  contextStack[index] = context;
-  didPerformWorkStack[index] = didChange;
+exports.pushTopLevelContextObject = function(fiber : Fiber, context : Object, didChange : boolean) : void {
+  invariant(contextStackCursor.cursor == null, 'Unexpected context found on stack');
+
+  push(contextStackCursor, context, fiber);
+  push(didPerformWorkStackCursor, didChange, fiber);
 };
 
 function processChildContext(fiber : Fiber, parentContext : Object, isReconciling : boolean): Object {
@@ -129,13 +130,13 @@ exports.pushContextProvider = function(workInProgress : Fiber, didPerformWork :
     instance.__reactInternalMemoizedMergedChildContext = mergedContext;
   }
 
-  index++;
-  contextStack[index] = mergedContext;
-  didPerformWorkStack[index] = didPerformWork;
+  push(contextStackCursor, mergedContext, workInProgress);
+  push(didPerformWorkStackCursor, didPerformWork, workInProgress);
 };
 
 exports.resetContext = function() : void {
-  index = -1;
+  reset(contextStackCursor);
+  reset(didPerformWorkStackCursor);
 };
 
 exports.findCurrentUnmaskedContext = function(fiber: Fiber) : Object {
@@ -162,7 +163,7 @@ exports.unwindContext = function(from : Fiber, to: Fiber) {
   let node = from;
   while (node && (node !== to) && (node.alternate !== to)) {
     if (isContextProvider(node)) {
-      popContextProvider();
+      popContextProvider(node);
     }
     node = node.return;
   }

src/renderers/shared/fiber/ReactFiberHostContext.js

@@ -14,16 +14,22 @@
 
 import type { Fiber } from 'ReactFiber';
 import type { HostConfig } from 'ReactFiberReconciler';
+import type { StackCursor } from 'ReactFiberStack';
+
+const {
+  createCursor,
+  pop,
+  push,
+  reset,
+} = require('ReactFiberStack');
 
 export type HostContext<C, CX> = {
-  getRootHostContainer() : C,
   getHostContext() : CX,
-
-  pushHostContext(fiber : Fiber) : void,
-  popHostContext(fiber : Fiber) : void,
-
-  pushHostContainer(container : C) : void,
+  getRootHostContainer() : C,
   popHostContainer() : void,
+  popHostContext(fiber : Fiber) : void,
+  pushHostContainer(container : C) : void,
+  pushHostContext(fiber : Fiber) : void,
   resetHostContainer() : void,
 };
 
@@ -35,129 +41,77 @@ module.exports = function<T, P, I, TI, C, CX>(
     getRootHostContext,
   } = config;
 
-  // Context stack is reused across the subtrees.
-  // We use a null sentinel on the fiber stack to separate them.
-  let contextFibers : Array<Fiber | null> = [];
-  let contextValues : Array<CX | null> = [];
-  let contextDepth : number = -1;
-  // Current context for fast access.
-  let currentContextValue : CX | null = null;
-  // Current root instance for fast access.
-  let rootInstance : C | null = null;
-  // A stack of outer root instances if we're in a portal.
-  let portalStack : Array<C | null> = [];
-  let portalDepth : number = -1;
+  let contextStackCursor : StackCursor = createCursor();
+  let rootInstanceStackCursor : StackCursor = createCursor();
 
   function getRootHostContainer() : C {
-    if (rootInstance == null) {
+    if (rootInstanceStackCursor.current == null) {
       throw new Error('Expected root container to exist.');
     }
-    return rootInstance;
+    return rootInstanceStackCursor.current;
   }
 
+  // TODO (bvaughn) Pass the host container fiber to push()
   function pushHostContainer(nextRootInstance : C) {
+    // Push current root instance onto the stack;
+    // This allows us to reset root when portals are popped.
+    push(rootInstanceStackCursor, nextRootInstance, null);
+
     const nextRootContext = getRootHostContext(nextRootInstance);
-    if (rootInstance == null) {
-      // We're entering a root.
-      rootInstance = nextRootInstance;
-    } else {
-      // We're entering a portal.
-      // Save the current root to the portal stack.
-      portalDepth++;
-      portalStack[portalDepth] = rootInstance;
-      rootInstance = nextRootInstance;
-    }
-    // Push the next root or portal context.
-    contextDepth++;
-    contextFibers[contextDepth] = null;
-    contextValues[contextDepth] = nextRootContext;
-    currentContextValue = nextRootContext;
+
+    push(contextStackCursor, nextRootContext, null);
   }
 
+  // TODO (bvaughn) Pass the host container fiber to pop()
   function popHostContainer() {
-    if (portalDepth === -1) {
-      // We're popping the root.
-      rootInstance = null;
-      currentContextValue = null;
-      contextDepth = -1;
-    } else {
-      // We're popping a portal.
-      // Restore the root instance.
-      rootInstance = portalStack[portalDepth];
-      portalStack[portalDepth] = null;
-      portalDepth--;
-      // If we pushed any context while in a portal, we need to roll it back.
-      if (contextDepth > -1) {
-        contextDepth--;
-        if (contextDepth > -1) {
-          currentContextValue = contextValues[contextDepth];
-        } else {
-          currentContextValue = null;
-        }
-      }
-    }
+    pop(contextStackCursor, null);
+    pop(rootInstanceStackCursor, null);
   }
 
   function getHostContext() : CX {
-    if (currentContextValue == null) {
+    if (contextStackCursor.current == null) {
       throw new Error('Expected host context to exist.');
     }
-    return currentContextValue;
+
+    return contextStackCursor.current;
   }
 
   function pushHostContext(fiber : Fiber) : void {
-    if (currentContextValue == null) {
+    if (rootInstanceStackCursor.current == null) {
       throw new Error('Expected root host context to exist.');
     }
-    const nextContextValue = getChildHostContext(currentContextValue, fiber.type, rootInstance);
-    if (currentContextValue === nextContextValue) {
+
+    const context = contextStackCursor.current;
+    const rootInstance = rootInstanceStackCursor.current;
+    const nextContext = getChildHostContext(context, fiber.type, rootInstance);
+
+    if (context === nextContext) {
       return;
     }
-    contextDepth++;
-    contextFibers[contextDepth] = fiber;
-    contextValues[contextDepth] = nextContextValue;
-    currentContextValue = nextContextValue;
+
+    push(contextStackCursor, nextContext, fiber);
   }
 
   function popHostContext(fiber : Fiber) : void {
-    if (contextDepth === -1) {
-      return;
-    }
-    if (contextFibers == null || contextValues == null) {
-      throw new Error('Expected host context stacks to exist when index is more than -1.');
-    }
-    if (fiber !== contextFibers[contextDepth]) {
+    if (contextStackCursor.current == null) {
       return;
     }
 
-    contextFibers[contextDepth] = null;
-    contextValues[contextDepth] = null;
-    contextDepth--;
-    if (contextDepth > -1) {
-      currentContextValue = contextValues[contextDepth];
-    } else {
-      currentContextValue = null;
-    }
+    pop(contextStackCursor, fiber);
   }
 
   function resetHostContainer() {
-    // Reset portal stack pointer because we're starting from the very top.
-    portalDepth = -1;
-    // Reset current container state.
-    rootInstance = null;
-    contextDepth = -1;
-    currentContextValue = null;
+    reset(contextStackCursor);
+    reset(rootInstanceStackCursor);
   }
 
   return {
-    getRootHostContainer,
     getHostContext,
-
-    pushHostContext,
-    popHostContext,
-
-    pushHostContainer,
+    getRootHostContainer,
     popHostContainer,
+    popHostContext,
+    pushHostContainer,
+    pushHostContext,
     resetHostContainer,
   };
 };

src/renderers/shared/fiber/ReactFiberStack.js

@@ -0,0 +1,102 @@
+/**
+ * Copyright 2013-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree. An additional grant
+ * of patent rights can be found in the PATENTS file in the same directory.
+ *
+ * @providesModule ReactFiberStack
+ * @flow
+ */
+
+'use strict';
+
+import type { Fiber } from 'ReactFiber';
+
+export type StackCursor = {|
+  current: any
+|};
+
+const emptyObject = require('emptyObject');
+const warning = require('warning');
+
+const valueStack : Array<Object> = [];
+
+if (__DEV__) {
+  var fiberStack : Array<Fiber | null> = [];
+}
+
+let index = -1;
+
+exports.createCursor = function() : StackCursor {
+  return {
+    current: null,
+  };
+};
+
+exports.isEmpty = function() : boolean {
+  return index === -1;
+};
+
+exports.pop = function(
+  cursor : StackCursor,
+  fiber: Fiber | null, // TODO (bvaughn) Tighten up this type to only accept Fiber
+) : void {
+  if (index < 0) {
+    if (__DEV__) {
+      warning(false, 'Unexpected pop.');
+    }
+    return;
+  }
+
+  if (__DEV__) {
+    if (fiber !== fiberStack[index]) {
+      warning(false, 'Unexpected Fiber popped.');
+    }
+  }
+
+  valueStack[index] = emptyObject;
+
+  if (__DEV__) {
+    fiberStack[index] = emptyObject;
+  }
+
+  index--;
+
+  cursor.current = index > -1
+    ? valueStack[index]
+    : null;
+};
+
+exports.push = function(
+  cursor : StackCursor,
+  value : any,
+  fiber: Fiber | null, // TODO (bvaughn) Tighten up this type to only accept Fiber
+) : void {
+  cursor.current = value;
+
+  index++;
+
+  valueStack[index] = value;
+
+  if (__DEV__) {
+    fiberStack[index] = fiber;
+  }
+};
+
+exports.reset = function(
+  cursor : StackCursor,
+) : void {
+  while (index > -1) {
+    valueStack[index] = emptyObject;
+
+    if (__DEV__) {
+      fiberStack[index] = emptyObject;
+    }
+
+    index--;
+  }
+
+  cursor.current = null;
+};