Fix suspense replaying forwardRefs with correct secondArg

facebook · Apr 1, 2023 · d111aa9 · d111aa9
1 parent e829146
commit d111aa9
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.js b/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) {

diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.js b/packages/react-reconciler/src/ReactFiberWorkLoop.js
@@ -2387,6 +2387,7 @@ function replaySuspendedUnitOfWork(unitOfWork: Fiber): void {
       // could maybe use this as an opportunity to say `use` doesn't work with
       // `defaultProps` :)
       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;
     }