use(context) and basic unit test

mofeiZ · mofeiZ · commit 6a0bb46effad · 2022-09-07T11:58:23.000-04:00
diff --git a/packages/react-reconciler/src/ReactFiberHooks.new.js b/packages/react-reconciler/src/ReactFiberHooks.new.js
@@ -36,6 +36,7 @@ import {
   enableUseHook,
   enableUseMemoCacheHook,
 } from 'shared/ReactFeatureFlags';
+import {REACT_CONTEXT_TYPE} from 'shared/ReactSymbols';
 
 import {NoMode, ConcurrentMode, DebugTracingMode} from './ReactTypeOfMode';
 import {
@@ -714,68 +715,70 @@ function createFunctionComponentUpdateQueue(): FunctionComponentUpdateQueue {
 }
 
 function use<T>(usable: Usable<T>): T {
-  if (
-    usable !== null &&
-    typeof usable === 'object' &&
-    typeof usable.then === 'function'
-  ) {
-    // This is a thenable.
-    const thenable: Thenable<T> = (usable: any);
-
-    // Track the position of the thenable within this fiber.
-    const index = thenableIndexCounter;
-    thenableIndexCounter += 1;
-
-    switch (thenable.status) {
-      case 'fulfilled': {
-        const fulfilledValue: T = thenable.value;
-        return fulfilledValue;
-      }
-      case 'rejected': {
-        const rejectedError = thenable.reason;
-        throw rejectedError;
-      }
-      default: {
-        const prevThenableAtIndex: Thenable<T> | null = getPreviouslyUsedThenableAtIndex(
-          index,
-        );
-        if (prevThenableAtIndex !== null) {
-          switch (prevThenableAtIndex.status) {
-            case 'fulfilled': {
-              const fulfilledValue: T = prevThenableAtIndex.value;
-              return fulfilledValue;
-            }
-            case 'rejected': {
-              const rejectedError: mixed = prevThenableAtIndex.reason;
-              throw rejectedError;
-            }
-            default: {
-              // The thenable still hasn't resolved. Suspend with the same
-              // thenable as last time to avoid redundant listeners.
-              throw prevThenableAtIndex;
+  if (usable !== null && typeof usable === 'object') {
+    if (typeof usable.then === 'function') {
+      // This is a thenable.
+      const thenable: Thenable<T> = (usable: any);
+
+      // Track the position of the thenable within this fiber.
+      const index = thenableIndexCounter;
+      thenableIndexCounter += 1;
+
+      switch (thenable.status) {
+        case 'fulfilled': {
+          const fulfilledValue: T = thenable.value;
+          return fulfilledValue;
+        }
+        case 'rejected': {
+          const rejectedError = thenable.reason;
+          throw rejectedError;
+        }
+        default: {
+          const prevThenableAtIndex: Thenable<T> | null = getPreviouslyUsedThenableAtIndex(
+            index,
+          );
+          if (prevThenableAtIndex !== null) {
+            switch (prevThenableAtIndex.status) {
+              case 'fulfilled': {
+                const fulfilledValue: T = prevThenableAtIndex.value;
+                return fulfilledValue;
+              }
+              case 'rejected': {
+                const rejectedError: mixed = prevThenableAtIndex.reason;
+                throw rejectedError;
+              }
+              default: {
+                // The thenable still hasn't resolved. Suspend with the same
+                // thenable as last time to avoid redundant listeners.
+                throw prevThenableAtIndex;
+              }
             }
+          } else {
+            // This is the first time something has been used at this index.
+            // Stash the thenable at the current index so we can reuse it during
+            // the next attempt.
+            trackUsedThenable(thenable, index);
+
+            // Suspend.
+            // TODO: Throwing here is an implementation detail that allows us to
+            // unwind the call stack. But we shouldn't allow it to leak into
+            // userspace. Throw an opaque placeholder value instead of the
+            // actual thenable. If it doesn't get captured by the work loop, log
+            // a warning, because that means something in userspace must have
+            // caught it.
+            throw thenable;
           }
-        } else {
-          // This is the first time something has been used at this index.
-          // Stash the thenable at the current index so we can reuse it during
-          // the next attempt.
-          trackUsedThenable(thenable, index);
-
-          // Suspend.
-          // TODO: Throwing here is an implementation detail that allows us to
-          // unwind the call stack. But we shouldn't allow it to leak into
-          // userspace. Throw an opaque placeholder value instead of the
-          // actual thenable. If it doesn't get captured by the work loop, log
-          // a warning, because that means something in userspace must have
-          // caught it.
-          throw thenable;
         }
       }
+    } else if (
+      usable.$$typeof != null &&
+      usable.$$typeof === REACT_CONTEXT_TYPE
+    ) {
+      const context: ReactContext<T> = (usable: any);
+      return readContext(context);
     }
   }
 
-  // TODO: Add support for Context
-
   // eslint-disable-next-line react-internal/safe-string-coercion
   throw new Error('An unsupported type was passed to use(): ' + String(usable));
 }
diff --git a/packages/react-reconciler/src/ReactFiberHooks.old.js b/packages/react-reconciler/src/ReactFiberHooks.old.js
@@ -36,6 +36,7 @@ import {
   enableUseHook,
   enableUseMemoCacheHook,
 } from 'shared/ReactFeatureFlags';
+import {REACT_CONTEXT_TYPE} from 'shared/ReactSymbols';
 
 import {NoMode, ConcurrentMode, DebugTracingMode} from './ReactTypeOfMode';
 import {
@@ -714,68 +715,70 @@ function createFunctionComponentUpdateQueue(): FunctionComponentUpdateQueue {
 }
 
 function use<T>(usable: Usable<T>): T {
-  if (
-    usable !== null &&
-    typeof usable === 'object' &&
-    typeof usable.then === 'function'
-  ) {
-    // This is a thenable.
-    const thenable: Thenable<T> = (usable: any);
-
-    // Track the position of the thenable within this fiber.
-    const index = thenableIndexCounter;
-    thenableIndexCounter += 1;
-
-    switch (thenable.status) {
-      case 'fulfilled': {
-        const fulfilledValue: T = thenable.value;
-        return fulfilledValue;
-      }
-      case 'rejected': {
-        const rejectedError = thenable.reason;
-        throw rejectedError;
-      }
-      default: {
-        const prevThenableAtIndex: Thenable<T> | null = getPreviouslyUsedThenableAtIndex(
-          index,
-        );
-        if (prevThenableAtIndex !== null) {
-          switch (prevThenableAtIndex.status) {
-            case 'fulfilled': {
-              const fulfilledValue: T = prevThenableAtIndex.value;
-              return fulfilledValue;
-            }
-            case 'rejected': {
-              const rejectedError: mixed = prevThenableAtIndex.reason;
-              throw rejectedError;
-            }
-            default: {
-              // The thenable still hasn't resolved. Suspend with the same
-              // thenable as last time to avoid redundant listeners.
-              throw prevThenableAtIndex;
+  if (usable !== null && typeof usable === 'object') {
+    if (typeof usable.then === 'function') {
+      // This is a thenable.
+      const thenable: Thenable<T> = (usable: any);
+
+      // Track the position of the thenable within this fiber.
+      const index = thenableIndexCounter;
+      thenableIndexCounter += 1;
+
+      switch (thenable.status) {
+        case 'fulfilled': {
+          const fulfilledValue: T = thenable.value;
+          return fulfilledValue;
+        }
+        case 'rejected': {
+          const rejectedError = thenable.reason;
+          throw rejectedError;
+        }
+        default: {
+          const prevThenableAtIndex: Thenable<T> | null = getPreviouslyUsedThenableAtIndex(
+            index,
+          );
+          if (prevThenableAtIndex !== null) {
+            switch (prevThenableAtIndex.status) {
+              case 'fulfilled': {
+                const fulfilledValue: T = prevThenableAtIndex.value;
+                return fulfilledValue;
+              }
+              case 'rejected': {
+                const rejectedError: mixed = prevThenableAtIndex.reason;
+                throw rejectedError;
+              }
+              default: {
+                // The thenable still hasn't resolved. Suspend with the same
+                // thenable as last time to avoid redundant listeners.
+                throw prevThenableAtIndex;
+              }
             }
+          } else {
+            // This is the first time something has been used at this index.
+            // Stash the thenable at the current index so we can reuse it during
+            // the next attempt.
+            trackUsedThenable(thenable, index);
+
+            // Suspend.
+            // TODO: Throwing here is an implementation detail that allows us to
+            // unwind the call stack. But we shouldn't allow it to leak into
+            // userspace. Throw an opaque placeholder value instead of the
+            // actual thenable. If it doesn't get captured by the work loop, log
+            // a warning, because that means something in userspace must have
+            // caught it.
+            throw thenable;
           }
-        } else {
-          // This is the first time something has been used at this index.
-          // Stash the thenable at the current index so we can reuse it during
-          // the next attempt.
-          trackUsedThenable(thenable, index);
-
-          // Suspend.
-          // TODO: Throwing here is an implementation detail that allows us to
-          // unwind the call stack. But we shouldn't allow it to leak into
-          // userspace. Throw an opaque placeholder value instead of the
-          // actual thenable. If it doesn't get captured by the work loop, log
-          // a warning, because that means something in userspace must have
-          // caught it.
-          throw thenable;
         }
       }
+    } else if (
+      usable.$$typeof != null &&
+      usable.$$typeof === REACT_CONTEXT_TYPE
+    ) {
+      const context: ReactContext<T> = (usable: any);
+      return readContext(context);
     }
   }
 
-  // TODO: Add support for Context
-
   // eslint-disable-next-line react-internal/safe-string-coercion
   throw new Error('An unsupported type was passed to use(): ' + String(usable));
 }
diff --git a/packages/react-reconciler/src/__tests__/ReactWakeable-test.js b/packages/react-reconciler/src/__tests__/ReactWakeable-test.js
@@ -315,4 +315,28 @@ describe('ReactWakeable', () => {
     ]);
     expect(root).toMatchRenderedOutput('Caught an error: Oops!');
   });
+
+  // @gate enableUseHook
+  test('basic use(context)', () => {
+    const ContextA = React.createContext('');
+    const ContextB = React.createContext('B');
+
+    function Sync() {
+      const text = use(ContextA) + use(ContextB);
+      return text;
+    }
+
+    function App() {
+      return (
+        <ContextA.Provider value="A">
+          <Sync />
+        </ContextA.Provider>
+      );
+    }
+
+    const root = ReactNoop.createRoot();
+    root.render(<App />);
+    expect(Scheduler).toFlushWithoutYielding();
+    expect(root).toMatchRenderedOutput('AB');
+  });
 });
diff --git a/packages/shared/ReactTypes.js b/packages/shared/ReactTypes.js
@@ -214,4 +214,4 @@ export type StartTransitionOptions = {
 };
 
 // TODO: Add Context support
-export type Usable<T> = Thenable<T>;
+export type Usable<T> = Thenable<T> | ReactContext<T>;
