diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.js b/packages/react-reconciler/src/ReactFiberBeginWork.js
index ef947a15a33e2..f14f63a86704f 100644
--- 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
index 71e6f4feabe41..16c29df5a19d3 100644
--- a/packages/react-reconciler/src/ReactFiberWorkLoop.js
+++ b/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;
}
diff --git a/packages/react-reconciler/src/__tests__/ReactUse-test.js b/packages/react-reconciler/src/__tests__/ReactUse-test.js
index 86bb586ec5829..e0b9e83d6a346 100644
--- a/packages/react-reconciler/src/__tests__/ReactUse-test.js
+++ b/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 ;
+ });
+
+ let _ref;
+ function App() {
+ const ref = (arg) => {
+ _ref = arg
+ };
+
+ return (
+ }>
+
+
+ );
+ }
+
+ const root = ReactNoop.createRoot();
+ await act(() => {
+ startTransition(() => {
+ root.render();
+ });
+ });
+ assertLog(['ABC']);
+ expect(root).toMatchRenderedOutput('ABC');
+ expect(_ref).toBe(refSymbol);
+ });
});