Fix suspense replaying forwardRefs

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -1169,15 +1169,15 @@ export function replayFunctionComponent(
   nextProps: any,
   Component: any,
   renderLanes: Lanes,
+  secondArg: any
 ): Fiber | null {
   // This function is used to replay a component that previously suspended,
   // after its data resolves. It's a simplified version of
   // updateFunctionComponent that reuses the hooks from the previous attempt.
 
-  let context;
-  if (!disableLegacyContext) {
+  if (!disableLegacyContext && secondArg === undefined) {
     const unmaskedContext = getUnmaskedContext(workInProgress, Component, true);
-    context = getMaskedContext(workInProgress, unmaskedContext);
+    secondArg = getMaskedContext(workInProgress, unmaskedContext);
   }
 
   prepareToReadContext(workInProgress, renderLanes);
@@ -1189,7 +1189,7 @@ export function replayFunctionComponent(
     workInProgress,
     Component,
     nextProps,
-    context,
+    secondArg,
   );
   const hasId = checkDidRenderIdHook();
   if (enableSchedulingProfiler) {

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -2386,7 +2386,8 @@ function replaySuspendedUnitOfWork(unitOfWork: Fiber): void {
       // TODO: Consider moving this switch statement into that module. Also,
       // could maybe use this as an opportunity to say `use` doesn't work with
       // `defaultProps` :)
-      const Component = unitOfWork.type;
+      const Component = unitOfWork.tag === FunctionComponent ? unitOfWork.type : unitOfWork.type.render;
+      const secondArg = unitOfWork.tag === FunctionComponent ? undefined : unitOfWork.ref;
       const unresolvedProps = unitOfWork.pendingProps;
       const resolvedProps =
         unitOfWork.elementType === Component
@@ -2398,18 +2399,21 @@ function replaySuspendedUnitOfWork(unitOfWork: Fiber): void {
         resolvedProps,
         Component,
         workInProgressRootRenderLanes,
+        secondArg
       );
       break;
     }
     case SimpleMemoComponent: {
       const Component = unitOfWork.type;
       const nextProps = unitOfWork.pendingProps;
+      const secondArg = undefined;
       next = replayFunctionComponent(
         current,
         unitOfWork,
         nextProps,
         Component,
         workInProgressRootRenderLanes,
+        secondArg
       );
       break;
     }

packages/react-reconciler/src/__tests__/ReactUse-test.js

@@ -1472,4 +1472,41 @@ describe('ReactUse', () => {
     assertLog(['Hi']);
     expect(root).toMatchRenderedOutput('Hi');
   });
+
+  it('basic use(promise) with forwardRefs', async () => {
+    const promiseA = Promise.resolve('A');
+    const promiseB = Promise.resolve('B');
+    const promiseC = Promise.resolve('C');
+
+    const refSymbol = {};
+
+    const Async = React.forwardRef((props, ref) => {
+      React.useImperativeHandle(ref, () => refSymbol);
+      const text = use(promiseA) + use(promiseB) + use(promiseC);
+      return <Text text={text} />;
+    });
+
+    let _ref;
+    function App() {
+      const ref = (arg) => {
+        _ref = arg
+      };
+
+      return (
+        <Suspense fallback={<Text text="Loading..." />}>
+          <Async ref={ref} />
+        </Suspense>
+      );
+    }
+
+    const root = ReactNoop.createRoot();
+    await act(() => {
+      startTransition(() => {
+        root.render(<App />);
+      });
+    });
+    assertLog(['ABC']);
+    expect(root).toMatchRenderedOutput('ABC');
+    expect(_ref).toBe(refSymbol);
+  });
 });