Add unstable context bailout for profiling

packages/react-debug-tools/src/ReactDebugHooks.js

@@ -37,6 +37,7 @@ import {
   REACT_CONTEXT_TYPE,
 } from 'shared/ReactSymbols';
 import hasOwnProperty from 'shared/hasOwnProperty';
+import type {ContextDependencyWithSelect} from '../../react-reconciler/src/ReactInternalTypes';
 
 type CurrentDispatcherRef = typeof ReactSharedInternals;
 
@@ -155,7 +156,10 @@ function getPrimitiveStackCache(): Map<string, Array<any>> {
 
 let currentFiber: null | Fiber = null;
 let currentHook: null | Hook = null;
-let currentContextDependency: null | ContextDependency<mixed> = null;
+let currentContextDependency:
+  | null
+  | ContextDependency<mixed>
+  | ContextDependencyWithSelect<mixed> = null;
 
 function nextHook(): null | Hook {
   const hook = currentHook;

packages/react-reconciler/src/ReactFiberHooks.js

@@ -47,6 +47,7 @@ import {
   enableUseDeferredValueInitialArg,
   disableLegacyMode,
   enableNoCloningMemoCache,
+  enableContextProfiling,
 } from 'shared/ReactFeatureFlags';
 import {
   REACT_CONTEXT_TYPE,
@@ -81,7 +82,11 @@ import {
   ContinuousEventPriority,
   higherEventPriority,
 } from './ReactEventPriorities';
-import {readContext, checkIfContextChanged} from './ReactFiberNewContext';
+import {
+  readContext,
+  readContextAndCompare,
+  checkIfContextChanged,
+} from './ReactFiberNewContext';
 import {HostRoot, CacheComponent, HostComponent} from './ReactWorkTags';
 import {
   LayoutStatic as LayoutStaticEffect,
@@ -1053,6 +1058,16 @@ function updateWorkInProgressHook(): Hook {
   return workInProgressHook;
 }
 
+function unstable_useContextWithBailout<T>(
+  context: ReactContext<T>,
+  select: (T => Array<mixed>) | null,
+): T {
+  if (select === null) {
+    return readContext(context);
+  }
+  return readContextAndCompare(context, select);
+}
+
 // NOTE: defining two versions of this function to avoid size impact when this feature is disabled.
 // Previously this function was inlined, the additional `memoCache` property makes it not inlined.
 let createFunctionComponentUpdateQueue: () => FunctionComponentUpdateQueue;
@@ -3689,6 +3704,10 @@ if (enableAsyncActions) {
 if (enableAsyncActions) {
   (ContextOnlyDispatcher: Dispatcher).useOptimistic = throwInvalidHookError;
 }
+if (enableContextProfiling) {
+  (ContextOnlyDispatcher: Dispatcher).unstable_useContextWithBailout =
+    throwInvalidHookError;
+}
 
 const HooksDispatcherOnMount: Dispatcher = {
   readContext,
@@ -3728,6 +3747,10 @@ if (enableAsyncActions) {
 if (enableAsyncActions) {
   (HooksDispatcherOnMount: Dispatcher).useOptimistic = mountOptimistic;
 }
+if (enableContextProfiling) {
+  (HooksDispatcherOnMount: Dispatcher).unstable_useContextWithBailout =
+    unstable_useContextWithBailout;
+}
 
 const HooksDispatcherOnUpdate: Dispatcher = {
   readContext,
@@ -3767,6 +3790,10 @@ if (enableAsyncActions) {
 if (enableAsyncActions) {
   (HooksDispatcherOnUpdate: Dispatcher).useOptimistic = updateOptimistic;
 }
+if (enableContextProfiling) {
+  (HooksDispatcherOnUpdate: Dispatcher).unstable_useContextWithBailout =
+    unstable_useContextWithBailout;
+}
 
 const HooksDispatcherOnRerender: Dispatcher = {
   readContext,
@@ -3806,6 +3833,10 @@ if (enableAsyncActions) {
 if (enableAsyncActions) {
   (HooksDispatcherOnRerender: Dispatcher).useOptimistic = rerenderOptimistic;
 }
+if (enableContextProfiling) {
+  (HooksDispatcherOnRerender: Dispatcher).unstable_useContextWithBailout =
+    unstable_useContextWithBailout;
+}
 
 let HooksDispatcherOnMountInDEV: Dispatcher | null = null;
 let HooksDispatcherOnMountWithHookTypesInDEV: Dispatcher | null = null;
@@ -4019,6 +4050,17 @@ if (__DEV__) {
         return mountOptimistic(passthrough, reducer);
       };
   }
+  if (enableContextProfiling) {
+    (HooksDispatcherOnMountInDEV: Dispatcher).unstable_useContextWithBailout =
+      function <T>(
+        context: ReactContext<T>,
+        select: (T => Array<mixed>) | null,
+      ): T {
+        currentHookNameInDev = 'useContext';
+        mountHookTypesDev();
+        return unstable_useContextWithBailout(context, select);
+      };
+  }
 
   HooksDispatcherOnMountWithHookTypesInDEV = {
     readContext<T>(context: ReactContext<T>): T {
@@ -4200,6 +4242,17 @@ if (__DEV__) {
         return mountOptimistic(passthrough, reducer);
       };
   }
+  if (enableContextProfiling) {
+    (HooksDispatcherOnMountWithHookTypesInDEV: Dispatcher).unstable_useContextWithBailout =
+      function <T>(
+        context: ReactContext<T>,
+        select: (T => Array<mixed>) | null,
+      ): T {
+        currentHookNameInDev = 'useContext';
+        updateHookTypesDev();
+        return unstable_useContextWithBailout(context, select);
+      };
+  }
 
   HooksDispatcherOnUpdateInDEV = {
     readContext<T>(context: ReactContext<T>): T {
@@ -4380,6 +4433,17 @@ if (__DEV__) {
         return updateOptimistic(passthrough, reducer);
       };
   }
+  if (enableContextProfiling) {
+    (HooksDispatcherOnUpdateInDEV: Dispatcher).unstable_useContextWithBailout =
+      function <T>(
+        context: ReactContext<T>,
+        select: (T => Array<mixed>) | null,
+      ): T {
+        currentHookNameInDev = 'useContext';
+        updateHookTypesDev();
+        return unstable_useContextWithBailout(context, select);
+      };
+  }
 
   HooksDispatcherOnRerenderInDEV = {
     readContext<T>(context: ReactContext<T>): T {
@@ -4560,6 +4624,17 @@ if (__DEV__) {
         return rerenderOptimistic(passthrough, reducer);
       };
   }
+  if (enableContextProfiling) {
+    (HooksDispatcherOnUpdateInDEV: Dispatcher).unstable_useContextWithBailout =
+      function <T>(
+        context: ReactContext<T>,
+        select: (T => Array<mixed>) | null,
+      ): T {
+        currentHookNameInDev = 'useContext';
+        updateHookTypesDev();
+        return unstable_useContextWithBailout(context, select);
+      };
+  }
 
   InvalidNestedHooksDispatcherOnMountInDEV = {
     readContext<T>(context: ReactContext<T>): T {
@@ -4766,6 +4841,18 @@ if (__DEV__) {
         return mountOptimistic(passthrough, reducer);
       };
   }
+  if (enableContextProfiling) {
+    (HooksDispatcherOnUpdateInDEV: Dispatcher).unstable_useContextWithBailout =
+      function <T>(
+        context: ReactContext<T>,
+        select: (T => Array<mixed>) | null,
+      ): T {
+        currentHookNameInDev = 'useContext';
+        warnInvalidHookAccess();
+        mountHookTypesDev();
+        return unstable_useContextWithBailout(context, select);
+      };
+  }
 
   InvalidNestedHooksDispatcherOnUpdateInDEV = {
     readContext<T>(context: ReactContext<T>): T {
@@ -4972,6 +5059,18 @@ if (__DEV__) {
         return updateOptimistic(passthrough, reducer);
       };
   }
+  if (enableContextProfiling) {
+    (InvalidNestedHooksDispatcherOnUpdateInDEV: Dispatcher).unstable_useContextWithBailout =
+      function <T>(
+        context: ReactContext<T>,
+        select: (T => Array<mixed>) | null,
+      ): T {
+        currentHookNameInDev = 'useContext';
+        warnInvalidHookAccess();
+        updateHookTypesDev();
+        return unstable_useContextWithBailout(context, select);
+      };
+  }
 
   InvalidNestedHooksDispatcherOnRerenderInDEV = {
     readContext<T>(context: ReactContext<T>): T {
@@ -5178,4 +5277,16 @@ if (__DEV__) {
         return rerenderOptimistic(passthrough, reducer);
       };
   }
+  if (enableContextProfiling) {
+    (InvalidNestedHooksDispatcherOnRerenderInDEV: Dispatcher).unstable_useContextWithBailout =
+      function <T>(
+        context: ReactContext<T>,
+        select: (T => Array<mixed>) | null,
+      ): T {
+        currentHookNameInDev = 'useContext';
+        warnInvalidHookAccess();
+        updateHookTypesDev();
+        return unstable_useContextWithBailout(context, select);
+      };
+  }
 }