diff --git a/src/renderers/shared/fiber/ReactFiberContext.js b/src/renderers/shared/fiber/ReactFiberContext.js
index 6caea831b3..33766583a5 100644
--- a/src/renderers/shared/fiber/ReactFiberContext.js
+++ b/src/renderers/shared/fiber/ReactFiberContext.js
@@ -159,13 +159,3 @@ exports.findCurrentUnmaskedContext = function(fiber: Fiber) : Object {
   }
   return node.stateNode.context;
 };
-
-exports.unwindContext = function(from : Fiber, to: Fiber) {
-  let node = from;
-  while (node && (node !== to) && (node.alternate !== to)) {
-    if (isContextProvider(node)) {
-      popContextProvider(node);
-    }
-    node = node.return;
-  }
-};
diff --git a/src/renderers/shared/fiber/ReactFiberScheduler.js b/src/renderers/shared/fiber/ReactFiberScheduler.js
index defece414b..10ccf0098d 100644
--- a/src/renderers/shared/fiber/ReactFiberScheduler.js
+++ b/src/renderers/shared/fiber/ReactFiberScheduler.js
@@ -17,6 +17,11 @@ import type { FiberRoot } from 'ReactFiberRoot';
 import type { HostConfig, Deadline } from 'ReactFiberReconciler';
 import type { PriorityLevel } from 'ReactPriorityLevel';
 
+var {
+  isContextProvider,
+  popContextProvider,
+} = require('ReactFiberContext');
+
 var ReactFiberBeginWork = require('ReactFiberBeginWork');
 var ReactFiberCompleteWork = require('ReactFiberCompleteWork');
 var ReactFiberCommitWork = require('ReactFiberCommitWork');
@@ -59,7 +64,6 @@ var {
 
 var {
   resetContext,
-  unwindContext,
 } = require('ReactFiberContext');
 
 if (__DEV__) {
@@ -724,8 +728,7 @@ module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C
           // from the root. Can't do that until then because without memoized
           // props, the nodes higher up in the tree will rerender unnecessarily.
           if (failedWork) {
-            unwindContext(failedWork, boundary);
-            unwindHostContext(failedWork, boundary);
+            unwindContexts(failedWork, boundary);
           }
           nextUnitOfWork = completeUnitOfWork(boundary);
         }
@@ -945,18 +948,23 @@ module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C
     }
   }
 
-  function unwindHostContext(from : Fiber, to: Fiber) {
+  function unwindContexts(from : Fiber, to: Fiber) {
     let node = from;
     while (node && (node !== to) && (node.alternate !== to)) {
       switch (node.tag) {
+        case ClassComponent:
+          if (isContextProvider(node)) {
+            popContextProvider(node);
+          }
+          break;
         case HostComponent:
           popHostContext(node);
           break;
         case HostRoot:
-          popHostContainer();
+          popHostContainer(node);
           break;
         case HostPortal:
-          popHostContainer();
+          popHostContainer(node);
           break;
       }
       node = node.return;
@@ -1118,4 +1126,4 @@ module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C
     syncUpdates: syncUpdates,
     deferredUpdates: deferredUpdates,
   };
-};
+};
\ No newline at end of file