@samitouri / QOS-React-2 / commits / 5123b063a7

[Fast Refresh] Find and remount wrapper edits behind lazy() (#36965)

Two bugs: - Type changes to a lazy's inner component (e.g., adding a `memo` comparison function) were missed by the refresh scan. - If a SimpleMemoComponent needs to be remounted, previously resolveRemountTypeForHotReloading would return the inner function and we would incorrectly recreate the fiber as a FunctionComponent. Instead, we can just use elementType consistently now that the beginWork functions consistently call resolveFamily on the wrapped components.

Sophie Alpert committed Jul 8, 2026 at 09:59 UTC 5123b063a7b3cd5c42ea2be2ccebe4e68d8b12e3
3 files changed +215 -31
packages/react-reconciler/src/ReactFiberBeginWork.js
+4 -8
@@ -132,10 +132,7 @@ import {
132 REACT_CONTEXT_TYPE,
133 } from 'shared/ReactSymbols';
134 import {setCurrentFiber} from './ReactCurrentFiber';
135 -import {
136 - resolveTypeForHotReloading,
137 - resolveRemountTypeForHotReloading,
138 -} from './ReactFiberHotReloading';
135 +import {resolveTypeForHotReloading} from './ReactFiberHotReloading';
136
137 import {
138 mountChildFibers,
@@ -4196,10 +4193,9 @@ function beginWork(
4193 if (workInProgress._debugNeedsRemount && current !== null) {
4194 // This will restart the begin phase with a new fiber.
4195 const copiedFiber = createFiberFromTypeAndProps(
4199 - resolveRemountTypeForHotReloading(
4200 - workInProgress.elementType,
4201 - workInProgress.type,
4202 - ),
4196 + // Remount from the fiber's outermost identity; mounting resolves
4197 + // any inner types to their latest implementations.
4198 + resolveTypeForHotReloading(workInProgress.elementType),
4199 workInProgress.key,
4200 workInProgress.pendingProps,
4201 workInProgress._debugOwner || null,
packages/react-reconciler/src/ReactFiberHotReloading.js
+11 -23
@@ -78,29 +78,6 @@ export function resolveTypeForHotReloading(type: any): any {
78 }
79 }
80
81 -export function resolveRemountTypeForHotReloading(
82 - elementType: any,
83 - type: any,
84 -): any {
85 - if (__DEV__) {
86 - if (resolveFamily === null) {
87 - // Hot reloading is disabled.
88 - return type;
89 - }
90 - // The elementType is the fiber's public identity, so its family tracks
91 - // the latest implementation even when an edit changed the kind of the
92 - // type (e.g. memo to a plain function) and `type` still points at the
93 - // old inner implementation.
94 - const family = resolveFamily(elementType);
95 - if (family === undefined) {
96 - return type;
97 - }
98 - return family.current;
99 - } else {
100 - return type;
101 - }
102 -}
103 -
81 export function isCompatibleFamilyForHotReloading(
82 fiber: Fiber,
83 element: ReactElement,
@@ -294,6 +271,17 @@ function scheduleFibersWithFamiliesRecursively(
271 const outerFamily = resolve(outerCandidateType);
272 if (outerFamily !== undefined && staleFamilies.has(outerFamily)) {
273 needsRemount = true;
274 + } else if (
275 + typeof outerCandidateType === 'object' &&
276 + outerCandidateType.$$typeof === REACT_LAZY_TYPE
277 + ) {
278 + const payload = outerCandidateType._payload;
279 + if (payload._status === 1 /* Resolved; see ReactLazy */) {
280 + const middleFamily = resolve(payload._result.default);
281 + if (middleFamily !== undefined && staleFamilies.has(middleFamily)) {
282 + needsRemount = true;
283 + }
284 + }
285 }
286 }
287 if (failedBoundaries !== null) {
packages/react-refresh/src/__tests__/ReactFresh-test.js
+200
@@ -1169,6 +1169,206 @@ describe('ReactFresh', () => {
1169 }
1170 });
1171
1172 + it('can remount lazy(memo()) when adding a comparison function', async () => {
1173 + if (__DEV__) {
1174 + let resolve;
1175 + await render(() => {
1176 + function Hello() {
1177 + return <p>hi memo</p>;
1178 + }
1179 + const Inner = React.memo(Hello);
1180 + $RefreshReg$(Hello, 'Hello');
1181 + $RefreshReg$(Inner, 'Inner');
1182 +
1183 + const Outer = React.lazy(
1184 + () =>
1185 + new Promise(_resolve => {
1186 + resolve = () => _resolve({default: Inner});
1187 + }),
1188 + );
1189 + $RefreshReg$(Outer, 'Outer');
1190 +
1191 + function App() {
1192 + return (
1193 + <React.Suspense fallback={<p>Loading</p>}>
1194 + <Outer />
1195 + </React.Suspense>
1196 + );
1197 + }
1198 + $RefreshReg$(App, 'App');
1199 + return App;
1200 + });
1201 +
1202 + expect(container.textContent).toBe('Loading');
1203 + await act(() => {
1204 + resolve();
1205 + });
1206 + expect(container.textContent).toBe('hi memo');
1207 + const el = container.firstChild;
1208 +
1209 + // Perform a hot update that adds a comparison function. The module
1210 + // creating the lazy also re-runs, like when an edit propagates.
1211 + await patch(() => {
1212 + function Hello() {
1213 + return <p>hi memo with compare</p>;
1214 + }
1215 + const Inner = React.memo(Hello, (prevProps, nextProps) => false);
1216 + $RefreshReg$(Hello, 'Hello');
1217 + $RefreshReg$(Inner, 'Inner');
1218 +
1219 + const Outer = React.lazy(
1220 + () =>
1221 + new Promise(_resolve => {
1222 + resolve = () => _resolve({default: Inner});
1223 + }),
1224 + );
1225 + $RefreshReg$(Outer, 'Outer');
1226 +
1227 + function App() {
1228 + return (
1229 + <React.Suspense fallback={<p>Loading</p>}>
1230 + <Outer />
1231 + </React.Suspense>
1232 + );
1233 + }
1234 + $RefreshReg$(App, 'App');
1235 + return App;
1236 + });
1237 +
1238 + // The shape change requires a remount. It goes through the latest
1239 + // lazy type, which suspends until it resolves. The boundary shows
1240 + // the fallback while the previous content stays hidden in the DOM.
1241 + expect(container.textContent).toBe('hi memoLoading');
1242 + await act(() => {
1243 + resolve();
1244 + });
1245 + expect(container.textContent).toBe('hi memo with compare');
1246 + expect(container.firstChild).not.toBe(el);
1247 + }
1248 + });
1249 +
1250 + it('can remount lazy(memo()) when adding a comparison function without re-creating the lazy', async () => {
1251 + if (__DEV__) {
1252 + let resolve;
1253 + await render(() => {
1254 + function Hello() {
1255 + return <p>hi memo</p>;
1256 + }
1257 + const Inner = React.memo(Hello);
1258 + $RefreshReg$(Hello, 'Hello');
1259 + $RefreshReg$(Inner, 'Inner');
1260 +
1261 + const Outer = React.lazy(
1262 + () =>
1263 + new Promise(_resolve => {
1264 + resolve = () => _resolve({default: Inner});
1265 + }),
1266 + );
1267 + $RefreshReg$(Outer, 'Outer');
1268 +
1269 + function App() {
1270 + return (
1271 + <React.Suspense fallback={<p>Loading</p>}>
1272 + <Outer />
1273 + </React.Suspense>
1274 + );
1275 + }
1276 + $RefreshReg$(App, 'App');
1277 + return App;
1278 + });
1279 +
1280 + expect(container.textContent).toBe('Loading');
1281 + await act(() => {
1282 + resolve();
1283 + });
1284 + expect(container.textContent).toBe('hi memo');
1285 + const el = container.firstChild;
1286 +
1287 + // Only the lazily loaded module re-runs this time, like when an
1288 + // edit is contained to it (the lazy type is not re-created, so the
1289 + // remount cannot go through it).
1290 + await patch(() => {
1291 + function Hello() {
1292 + return <p>hi memo with compare</p>;
1293 + }
1294 + const Inner = React.memo(Hello, (prevProps, nextProps) => false);
1295 + $RefreshReg$(Hello, 'Hello');
1296 + $RefreshReg$(Inner, 'Inner');
1297 + return Inner;
1298 + });
1299 +
1300 + // The remount goes through the old lazy, whose payload is already
1301 + // resolved, so it doesn't suspend.
1302 + expect(container.textContent).toBe('hi memo with compare');
1303 + expect(container.firstChild).not.toBe(el);
1304 + }
1305 + });
1306 +
1307 + it('can remount an unregistered memo wrapper without losing the wrapper', async () => {
1308 + if (__DEV__) {
1309 + let innerRenders = 0;
1310 + await act(async () => {
1311 + await render(() => {
1312 + function Inner({label}) {
1313 + innerRenders++;
1314 + return <p>{label}</p>;
1315 + }
1316 + $RefreshReg$(Inner, 'Inner');
1317 + $RefreshSig$(Inner, 'sig1');
1318 + // The wrapper is deliberately not registered, like a wrapper
1319 + // created inside a third-party HOC.
1320 + const InnerMemo = React.memo(Inner);
1321 +
1322 + function App() {
1323 + const [, forceUpdate] = React.useState(0);
1324 + return (
1325 + <div onClick={() => forceUpdate(n => n + 1)}>
1326 + <InnerMemo label="hi" />
1327 + </div>
1328 + );
1329 + }
1330 + $RefreshReg$(App, 'App');
1331 + return App;
1332 + });
1333 + });
1334 +
1335 + expect(container.textContent).toBe('hi');
1336 + expect(innerRenders).toBe(1);
1337 +
1338 + // The memo blocks re-renders with equal props.
1339 + await act(async () => {
1340 + container.firstChild.click();
1341 + });
1342 + expect(innerRenders).toBe(1);
1343 +
1344 + // Force a remount by changing the inner function's signature.
1345 + // Only the inner function's module re-runs; the wrapper and the
1346 + // element referencing it are not re-created.
1347 + await act(async () => {
1348 + await patch(() => {
1349 + function Inner({label}) {
1350 + innerRenders++;
1351 + return <p>{label}</p>;
1352 + }
1353 + $RefreshReg$(Inner, 'Inner');
1354 + $RefreshSig$(Inner, 'sig2');
1355 + return Inner;
1356 + });
1357 + });
1358 + expect(innerRenders).toBe(2);
1359 + const innerEl = container.firstChild.firstChild;
1360 +
1361 + // The remounted fiber must still be a memo: equal props stay
1362 + // blocked, and the fiber reconciles against the original element
1363 + // instead of being replaced again.
1364 + await act(async () => {
1365 + container.firstChild.click();
1366 + });
1367 + expect(innerRenders).toBe(2);
1368 + expect(container.firstChild.firstChild).toBe(innerEl);
1369 + }
1370 + });
1371 +
1372 it('resets state when switching between different component types', async () => {
1373 if (__DEV__) {
1374 await act(async () => {