Add replayForwardRef

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -465,6 +465,52 @@ function updateForwardRef(
   return workInProgress.child;
 }
 
+export function replayForwardRef(
+  current: Fiber | null,
+  workInProgress: Fiber,
+  nextProps: any,
+  Component: any,
+  renderLanes: Lanes,
+): Fiber | null {
+  const secondArg = workInProgress.ref;
+
+  // 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.
+
+  prepareToReadContext(workInProgress, renderLanes);
+  if (enableSchedulingProfiler) {
+    markComponentRenderStarted(workInProgress);
+  }
+  const nextChildren = replaySuspendedComponentWithHooks(
+    current,
+    workInProgress,
+    Component,
+    nextProps,
+    secondArg,
+  );
+
+  // the rest is a fork of updateFunctionComponent
+  const hasId = checkDidRenderIdHook();
+  if (enableSchedulingProfiler) {
+    markComponentRenderStopped();
+  }
+
+  if (current !== null && !didReceiveUpdate) {
+    bailoutHooks(current, workInProgress, renderLanes);
+    return bailoutOnAlreadyFinishedWork(current, workInProgress, renderLanes);
+  }
+
+  if (getIsHydrating() && hasId) {
+    pushMaterializedTreeId(workInProgress);
+  }
+
+  // React DevTools reads this flag.
+  workInProgress.flags |= PerformedWork;
+  reconcileChildren(current, workInProgress, nextChildren, renderLanes);
+  return workInProgress.child;
+}
+
 function updateMemoComponent(
   current: Fiber | null,
   workInProgress: Fiber,
@@ -1169,15 +1215,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.
 
-  if (!disableLegacyContext && secondArg === undefined) {
+  let context: any;
+  if (!disableLegacyContext) {
     const unmaskedContext = getUnmaskedContext(workInProgress, Component, true);
-    secondArg = getMaskedContext(workInProgress, unmaskedContext);
+    context = getMaskedContext(workInProgress, unmaskedContext);
   }
 
   prepareToReadContext(workInProgress, renderLanes);
@@ -1189,7 +1235,7 @@ export function replayFunctionComponent(
     workInProgress,
     Component,
     nextProps,
-    secondArg,
+    context,
   );
   const hasId = checkDidRenderIdHook();
   if (enableSchedulingProfiler) {

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -171,6 +171,7 @@ import {
   SelectiveHydrationException,
   beginWork as originalBeginWork,
   replayFunctionComponent,
+  replayForwardRef,
 } from './ReactFiberBeginWork';
 import {completeWork} from './ReactFiberCompleteWork';
 import {unwindWork, unwindInterruptedWork} from './ReactFiberUnwindWork';
@@ -2380,18 +2381,12 @@ function replaySuspendedUnitOfWork(unitOfWork: Fiber): void {
       // Fallthrough to the next branch.
     }
     // eslint-disable-next-line no-fallthrough
-    case FunctionComponent:
-    case ForwardRef: {
+    case FunctionComponent: {
       // Resolve `defaultProps`. This logic is copied from `beginWork`.
       // 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.tag === FunctionComponent
-          ? unitOfWork.type
-          : unitOfWork.type.render;
-      const secondArg =
-        unitOfWork.tag === FunctionComponent ? undefined : unitOfWork.ref;
+      const Component = unitOfWork.type;
       const unresolvedProps = unitOfWork.pendingProps;
       const resolvedProps =
         unitOfWork.elementType === Component
@@ -2403,7 +2398,26 @@ function replaySuspendedUnitOfWork(unitOfWork: Fiber): void {
         resolvedProps,
         Component,
         workInProgressRootRenderLanes,
-        secondArg,
+      );
+      break;
+    }
+    case ForwardRef: {
+      // Resolve `defaultProps`. This logic is copied from `beginWork`.
+      // 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.render;
+      const unresolvedProps = unitOfWork.pendingProps;
+      const resolvedProps =
+        unitOfWork.elementType === Component
+          ? unresolvedProps
+          : resolveDefaultProps(Component, unresolvedProps);
+      next = replayForwardRef(
+        current,
+        unitOfWork,
+        resolvedProps,
+        Component,
+        workInProgressRootRenderLanes,
       );
       break;
     }