From 3b65e46d8c4532f7ecc2a36e21bf20e576429d80 Mon Sep 17 00:00:00 2001
From: Sebastian Silbermann <sebastian.silbermann@klarna.com>
Date: Sun, 4 Feb 2024 15:45:03 +0100
Subject: [PATCH] DevTools: Add support for use(Context)

---
 .../react-debug-tools/src/ReactDebugHooks.js  | 39 +++++++++++++++----
 .../ReactHooksInspectionIntegration-test.js   | 38 ++++++++++++++++++
 .../app/InspectableElements/CustomHooks.js    |  2 +
 3 files changed, 72 insertions(+), 7 deletions(-)

diff --git a/packages/react-debug-tools/src/ReactDebugHooks.js b/packages/react-debug-tools/src/ReactDebugHooks.js
index b4b0f97e2edaa..e8dc2578be11a 100644
--- a/packages/react-debug-tools/src/ReactDebugHooks.js
+++ b/packages/react-debug-tools/src/ReactDebugHooks.js
@@ -12,6 +12,7 @@ import type {
   ReactContext,
   ReactProviderType,
   StartTransitionOptions,
+  Usable,
 } from 'shared/ReactTypes';
 import type {
   Fiber,
@@ -27,7 +28,10 @@ import {
   ContextProvider,
   ForwardRef,
 } from 'react-reconciler/src/ReactWorkTags';
-import {REACT_MEMO_CACHE_SENTINEL} from 'shared/ReactSymbols';
+import {
+  REACT_MEMO_CACHE_SENTINEL,
+  REACT_CONTEXT_TYPE,
+} from 'shared/ReactSymbols';
 
 type CurrentDispatcherRef = typeof ReactSharedInternals.ReactCurrentDispatcher;
 
@@ -118,11 +122,30 @@ function readContext<T>(context: ReactContext<T>): T {
   return context._currentValue;
 }
 
-function use<T>(): T {
-  // TODO: What should this do if it receives an unresolved promise?
-  throw new Error(
-    'Support for `use` not yet implemented in react-debug-tools.',
-  );
+function use<T>(usable: Usable<T>): T {
+  if (usable !== null && typeof usable === 'object') {
+    // $FlowFixMe[method-unbinding]
+    if (typeof usable.then === 'function') {
+      // TODO: What should this do if it receives an unresolved promise?
+      throw new Error(
+        'Support for `use(Promise)` not yet implemented in react-debug-tools.',
+      );
+    } else if (usable.$$typeof === REACT_CONTEXT_TYPE) {
+      const context: ReactContext<T> = (usable: any);
+      const value = readContext(context);
+
+      hookLog.push({
+        primitive: 'Use',
+        stackError: new Error(),
+        value,
+      });
+
+      return value;
+    }
+  }
+
+  // eslint-disable-next-line react-internal/safe-string-coercion
+  throw new Error('An unsupported type was passed to use(): ' + String(usable));
 }
 
 function useContext<T>(context: ReactContext<T>): T {
@@ -660,7 +683,9 @@ function buildTree(
     // For now, the "id" of stateful hooks is just the stateful hook index.
     // Custom hooks have no ids, nor do non-stateful native hooks (e.g. Context, DebugValue).
     const id =
-      primitive === 'Context' || primitive === 'DebugValue'
+      primitive === 'Context' ||
+      primitive === 'DebugValue' ||
+      primitive === 'Use'
         ? null
         : nativeHookID++;
 
diff --git a/packages/react-debug-tools/src/__tests__/ReactHooksInspectionIntegration-test.js b/packages/react-debug-tools/src/__tests__/ReactHooksInspectionIntegration-test.js
index d185cd4df634c..fc9d338b3f377 100644
--- a/packages/react-debug-tools/src/__tests__/ReactHooksInspectionIntegration-test.js
+++ b/packages/react-debug-tools/src/__tests__/ReactHooksInspectionIntegration-test.js
@@ -1108,6 +1108,44 @@ describe('ReactHooksInspectionIntegration', () => {
     ]);
   });
 
+  it('should support use(Context) hook', () => {
+    const Context = React.createContext('default');
+    function Foo() {
+      const value = React.use(Context);
+      React.useMemo(() => 'memo', []);
+      React.useMemo(() => 'not used', []);
+
+      return value;
+    }
+
+    const renderer = ReactTestRenderer.create(<Foo />);
+    const childFiber = renderer.root.findByType(Foo)._currentFiber();
+    const tree = ReactDebugTools.inspectHooksOfFiber(childFiber);
+    expect(tree).toEqual([
+      {
+        id: null,
+        isStateEditable: false,
+        name: 'Use',
+        value: 'default',
+        subHooks: [],
+      },
+      {
+        id: 0,
+        isStateEditable: false,
+        name: 'Memo',
+        value: 'memo',
+        subHooks: [],
+      },
+      {
+        id: 1,
+        isStateEditable: false,
+        name: 'Memo',
+        value: 'not used',
+        subHooks: [],
+      },
+    ]);
+  });
+
   // @gate enableAsyncActions
   it('should support useOptimistic hook', () => {
     const useOptimistic = React.useOptimistic;
diff --git a/packages/react-devtools-shell/src/app/InspectableElements/CustomHooks.js b/packages/react-devtools-shell/src/app/InspectableElements/CustomHooks.js
index 2709b72ac2bef..aa3a9a6295731 100644
--- a/packages/react-devtools-shell/src/app/InspectableElements/CustomHooks.js
+++ b/packages/react-devtools-shell/src/app/InspectableElements/CustomHooks.js
@@ -19,6 +19,7 @@ import {
   useEffect,
   useOptimistic,
   useState,
+  use,
 } from 'react';
 import {useFormState} from 'react-dom';
 
@@ -76,6 +77,7 @@ function FunctionWithHooks(props: any, ref: React$Ref<any>) {
   // eslint-disable-next-line no-unused-vars
   const contextValueA = useContext(ContextA);
   useOptimistic<number, mixed>(1);
+  use(ContextA);
 
   // eslint-disable-next-line no-unused-vars
   const [_, __] = useState(object);