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); + }); });