Store nearest provider on context object

Author: acdlite

packages/react-dom/src/server/ReactPartialRenderer.js

@@ -124,33 +124,30 @@ if (__DEV__) {
 }
 
 // Context (new API)
-let providerStack: Array<ReactProvider<mixed>> = []; // Stack of provider objects
+let providerStack: Array<ReactProvider<any>> = []; // Stack of provider objects
 let index = -1;
 
 export function pushProvider<T>(provider: ReactProvider<T>): void {
   index += 1;
   providerStack[index] = provider;
+  const context: ReactContext<any> = provider.type.context;
+  context.currentProvider = provider;
 }
 
 export function popProvider<T>(provider: ReactProvider<T>): void {
   if (__DEV__) {
     warning(index > -1 && provider === providerStack[index], 'Unexpected pop.');
   }
+  // $FlowFixMe - Intentionally unsound
   providerStack[index] = null;
   index -= 1;
-}
-
-// Find the nearest matching provider
-export function getProvider<T>(
-  context: ReactContext<T>,
-): ReactProvider<T> | null {
-  for (let i = index; i > -1; i--) {
-    const provider = providerStack[i];
-    if (provider.type.context === context) {
-      return provider;
-    }
+  const context: ReactContext<any> = provider.type.context;
+  if (index < 0) {
+    context.currentProvider = null;
+  } else {
+    const previousProvider = providerStack[index];
+    context.currentProvider = previousProvider;
   }
-  return null;
 }
 
 let didWarnDefaultInputValue = false;
@@ -731,7 +728,7 @@ class ReactDOMServerRenderer {
       if (typeof elementType === 'object' && elementType !== null) {
         switch (elementType.$$typeof) {
           case REACT_PROVIDER_TYPE: {
-            const provider: ReactProvider<any> = nextChild;
+            const provider: ReactProvider<any> = (nextChild: any);
             const nextProps = provider.props;
             const nextChildren = toArray(nextProps.children);
             const frame: Frame = {
@@ -752,10 +749,10 @@ class ReactDOMServerRenderer {
             return '';
           }
           case REACT_CONSUMER_TYPE: {
-            const consumer: ReactConsumer<any> = nextChild;
+            const consumer: ReactConsumer<any> = (nextChild: any);
             const nextProps = consumer.props;
 
-            const provider = getProvider(consumer.type.context);
+            const provider = consumer.type.context.currentProvider;
             let nextValue;
             if (provider === null) {
               // Detached consumer

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -66,7 +66,7 @@ import {
   pushTopLevelContextObject,
   invalidateContextProvider,
 } from './ReactFiberContext';
-import {pushProvider, getProvider} from './ReactFiberNewContext';
+import {pushProvider} from './ReactFiberNewContext';
 import {NoWork, Never} from './ReactFiberExpirationTime';
 
 let warnedAboutStatelessRefs;
@@ -755,7 +755,7 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
       const oldProps = workInProgress.memoizedProps;
 
       // Get the nearest ancestor provider.
-      const providerFiber: Fiber | null = getProvider(context);
+      const providerFiber: Fiber | null = context.currentProvider;
 
       let newValue;
       let valueDidChange;

packages/react-reconciler/src/ReactFiberNewContext.js

@@ -18,6 +18,8 @@ let index = -1;
 export function pushProvider(providerFiber: Fiber): void {
   index += 1;
   stack[index] = providerFiber;
+  const context: ReactContext<any> = providerFiber.type.context;
+  context.currentProvider = providerFiber;
 }
 
 export function popProvider(providerFiber: Fiber): void {
@@ -26,21 +28,20 @@ export function popProvider(providerFiber: Fiber): void {
   }
   stack[index] = null;
   index -= 1;
-}
-
-// Find the nearest matching provider
-export function getProvider<T>(context: ReactContext<T>): Fiber | null {
-  for (let i = index; i > -1; i--) {
-    const provider = stack[i];
-    if (provider.type.context === context) {
-      return provider;
-    }
+  const context: ReactContext<any> = providerFiber.type.context;
+  if (index < 0) {
+    context.currentProvider = null;
+  } else {
+    const previousProviderFiber = stack[index];
+    context.currentProvider = previousProviderFiber;
   }
-  return null;
 }
 
 export function resetProviderStack(): void {
   for (let i = index; i > -1; i--) {
-    stack[i] = null;
+    const providerFiber = stack[i];
+    const context: ReactContext<any> = providerFiber.type.context;
+    context.currentProvider = null;
+    stack[i] = undefined;
   }
 }

packages/shared/ReactContext.js

@@ -53,6 +53,7 @@ export function createContext<T>(defaultValue: T): ReactContext<T> {
       };
     },
     defaultValue,
+    currentProvider: null,
   };
 
   providerType = {

packages/shared/ReactTypes.js

@@ -1,5 +1,3 @@
-import {Symbol} from './node_modules/typescript/lib/typescript';
-
 /**
  * Copyright (c) 2014-present, Facebook, Inc.
  *
@@ -88,6 +86,7 @@ export type ReactContext<T> = {
   provide(value: T, children: ReactNodeList, key?: string): ReactProvider<T>,
   consume(render: (value: T) => ReactNodeList, key?: string): ReactConsumer<T>,
   defaultValue: T,
+  currentProvider: any, // Fiber | null
 };
 
 export type ReactPortal = {