Fix context memory leak (#10680)

* Fix context memory leak

* Fix Flow

Author: gaearon

src/renderers/shared/fiber/ReactFiberBeginWork.js

@@ -305,7 +305,7 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
     return workInProgress.child;
   }
 
-  function updateHostRoot(current, workInProgress, priorityLevel) {
+  function pushHostRootContext(workInProgress) {
     const root = (workInProgress.stateNode: FiberRoot);
     if (root.pendingContext) {
       pushTopLevelContextObject(
@@ -317,9 +317,11 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
       // Should always be set
       pushTopLevelContextObject(workInProgress, root.context, false);
     }
-
     pushHostContainer(workInProgress, root.containerInfo);
+  }
 
+  function updateHostRoot(current, workInProgress, priorityLevel) {
+    pushHostRootContext(workInProgress);
     const updateQueue = workInProgress.updateQueue;
     if (updateQueue !== null) {
       const prevState = workInProgress.memoizedState;
@@ -777,8 +779,7 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
         pushContextProvider(workInProgress);
         break;
       case HostRoot:
-        const root: FiberRoot = workInProgress.stateNode;
-        pushHostContainer(workInProgress, root.containerInfo);
+        pushHostRootContext(workInProgress);
         break;
       default:
         invariant(

src/renderers/shared/fiber/ReactFiberCompleteWork.js

@@ -21,7 +21,10 @@ import type {FiberRoot} from 'ReactFiberRoot';
 import type {HostConfig} from 'ReactFiberReconciler';
 
 var {reconcileChildFibers} = require('ReactChildFiber');
-var {popContextProvider} = require('ReactFiberContext');
+var {
+  popContextProvider,
+  popTopLevelContextObject,
+} = require('ReactFiberContext');
 var ReactTypeOfWork = require('ReactTypeOfWork');
 var ReactTypeOfSideEffect = require('ReactTypeOfSideEffect');
 var ReactPriorityLevel = require('ReactPriorityLevel');
@@ -211,7 +214,8 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
         return null;
       }
       case HostRoot: {
-        // TODO: Pop the host container after #8607 lands.
+        popHostContainer(workInProgress);
+        popTopLevelContextObject(workInProgress);
         const fiberRoot = (workInProgress.stateNode: FiberRoot);
         if (fiberRoot.pendingContext) {
           fiberRoot.context = fiberRoot.pendingContext;

src/renderers/shared/fiber/ReactFiberContext.js

@@ -135,6 +135,11 @@ function popContextProvider(fiber: Fiber): void {
 }
 exports.popContextProvider = popContextProvider;
 
+exports.popTopLevelContextObject = function(fiber: Fiber) {
+  pop(didPerformWorkStackCursor, fiber);
+  pop(contextStackCursor, fiber);
+};
+
 exports.pushTopLevelContextObject = function(
   fiber: Fiber,
   context: Object,