Revert "Fix suspense replaying forward refs (#26535)"

This reverts commit 7329ea8.

Author: kassens

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -1168,13 +1168,18 @@ export function replayFunctionComponent(
   workInProgress: Fiber,
   nextProps: any,
   Component: any,
-  secondArg: any,
   renderLanes: Lanes,
 ): 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) {
+    const unmaskedContext = getUnmaskedContext(workInProgress, Component, true);
+    context = getMaskedContext(workInProgress, unmaskedContext);
+  }
+
   prepareToReadContext(workInProgress, renderLanes);
   if (enableSchedulingProfiler) {
     markComponentRenderStarted(workInProgress);
@@ -1184,7 +1189,7 @@ export function replayFunctionComponent(
     workInProgress,
     Component,
     nextProps,
-    secondArg,
+    context,
   );
   const hasId = checkDidRenderIdHook();
   if (enableSchedulingProfiler) {

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -39,7 +39,6 @@ import {
   enableTransitionTracing,
   useModernStrictMode,
   revertRemovalOfSiblingPrerendering,
-  disableLegacyContext,
 } from 'shared/ReactFeatureFlags';
 import ReactSharedInternals from 'shared/ReactSharedInternals';
 import is from 'shared/objectIs';
@@ -282,7 +281,6 @@ import {
   flushSyncWorkOnLegacyRootsOnly,
   getContinuationForRoot,
 } from './ReactFiberRootScheduler';
-import {getMaskedContext, getUnmaskedContext} from './ReactFiberContext';
 
 const ceil = Math.ceil;
 
@@ -2382,8 +2380,8 @@ function replaySuspendedUnitOfWork(unitOfWork: Fiber): void {
       // Fallthrough to the next branch.
     }
     // eslint-disable-next-line no-fallthrough
-    case SimpleMemoComponent:
-    case FunctionComponent: {
+    case FunctionComponent:
+    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
@@ -2394,39 +2392,23 @@ function replaySuspendedUnitOfWork(unitOfWork: Fiber): void {
         unitOfWork.elementType === Component
           ? unresolvedProps
           : resolveDefaultProps(Component, unresolvedProps);
-      let context: any;
-      if (!disableLegacyContext) {
-        const unmaskedContext = getUnmaskedContext(unitOfWork, Component, true);
-        context = getMaskedContext(unitOfWork, unmaskedContext);
-      }
       next = replayFunctionComponent(
         current,
         unitOfWork,
         resolvedProps,
         Component,
-        context,
         workInProgressRootRenderLanes,
       );
       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);
-
+    case SimpleMemoComponent: {
+      const Component = unitOfWork.type;
+      const nextProps = unitOfWork.pendingProps;
       next = replayFunctionComponent(
         current,
         unitOfWork,
-        resolvedProps,
+        nextProps,
         Component,
-        unitOfWork.ref,
         workInProgressRootRenderLanes,
       );
       break;

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

@@ -1472,132 +1472,4 @@ describe('ReactUse', () => {
     assertLog(['Hi']);
     expect(root).toMatchRenderedOutput('Hi');
   });
-
-  test('unwrap uncached promises inside forwardRef', async () => {
-    const asyncInstance = {};
-    const Async = React.forwardRef((props, ref) => {
-      React.useImperativeHandle(ref, () => asyncInstance);
-      const text = use(Promise.resolve('Async'));
-      return <Text text={text} />;
-    });
-
-    const ref = React.createRef();
-    function App() {
-      return (
-        <Suspense fallback={<Text text="Loading..." />}>
-          <Async ref={ref} />
-        </Suspense>
-      );
-    }
-
-    const root = ReactNoop.createRoot();
-    await act(() => {
-      startTransition(() => {
-        root.render(<App />);
-      });
-    });
-    assertLog(['Async']);
-    expect(root).toMatchRenderedOutput('Async');
-    expect(ref.current).toBe(asyncInstance);
-  });
-
-  test('unwrap uncached promises inside memo', async () => {
-    const Async = React.memo(
-      props => {
-        const text = use(Promise.resolve(props.text));
-        return <Text text={text} />;
-      },
-      (a, b) => a.text === b.text,
-    );
-
-    function App({text}) {
-      return (
-        <Suspense fallback={<Text text="Loading..." />}>
-          <Async text={text} />
-        </Suspense>
-      );
-    }
-
-    const root = ReactNoop.createRoot();
-    await act(() => {
-      startTransition(() => {
-        root.render(<App text="Async" />);
-      });
-    });
-    assertLog(['Async']);
-    expect(root).toMatchRenderedOutput('Async');
-
-    // Update to the same value
-    await act(() => {
-      startTransition(() => {
-        root.render(<App text="Async" />);
-      });
-    });
-    // Should not have re-rendered, because it's memoized
-    assertLog([]);
-    expect(root).toMatchRenderedOutput('Async');
-
-    // Update to a different value
-    await act(() => {
-      startTransition(() => {
-        root.render(<App text="Async!" />);
-      });
-    });
-    assertLog(['Async!']);
-    expect(root).toMatchRenderedOutput('Async!');
-  });
-
-  // @gate !disableLegacyContext
-  test('unwrap uncached promises in component that accesses legacy context', async () => {
-    class ContextProvider extends React.Component {
-      static childContextTypes = {
-        legacyContext() {},
-      };
-      getChildContext() {
-        return {legacyContext: 'Async'};
-      }
-      render() {
-        return this.props.children;
-      }
-    }
-
-    function Async({label}, context) {
-      const text = use(Promise.resolve(context.legacyContext + ` (${label})`));
-      return <Text text={text} />;
-    }
-    Async.contextTypes = {
-      legacyContext: () => {},
-    };
-
-    const AsyncMemo = React.memo(Async, (a, b) => a.label === b.label);
-
-    function App() {
-      return (
-        <ContextProvider>
-          <Suspense fallback={<Text text="Loading..." />}>
-            <div>
-              <Async label="function component" />
-            </div>
-            <div>
-              <AsyncMemo label="memo component" />
-            </div>
-          </Suspense>
-        </ContextProvider>
-      );
-    }
-
-    const root = ReactNoop.createRoot();
-    await act(() => {
-      startTransition(() => {
-        root.render(<App />);
-      });
-    });
-    assertLog(['Async (function component)', 'Async (memo component)']);
-    expect(root).toMatchRenderedOutput(
-      <>
-        <div>Async (function component)</div>
-        <div>Async (memo component)</div>
-      </>,
-    );
-  });
 });