Fix suspense replaying forwardRefs with correct secondArg

Author: hansottowirtz

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

@@ -2390,6 +2390,8 @@ function replaySuspendedUnitOfWork(unitOfWork: Fiber): void {
         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
@@ -2401,18 +2403,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;
     }