From df4bd1b4e99269ac9f3ce3a47ef2e23f76fb1b97 Mon Sep 17 00:00:00 2001
From: Sophie Alpert <git@sophiebits.com>
Date: Wed, 8 Jul 2026 09:11:20 -0700
Subject: [PATCH 1/4] [Fast Refresh] Unify hot reload type resolution (#36962)

Previously resolveForwardRefForHotReloading was treated as a special
case. It's simpler to move this logic to updateForwardRef which
mirrors how updateMemoComponent and mountLazyComponent already are,
and it allows us to unify to a single resolveTypeForHotReloading which
we'll use in the next commit.

Since .type is no longer changing we need to set didReceiveUpdate
manually.

This behavior is already tested in "can update forwardRef render
function in isolation".
---
 packages/react-reconciler/src/ReactFiber.js   | 19 ++------
 .../src/ReactFiberBeginWork.js                | 26 ++++++-----
 .../src/ReactFiberHotReloading.js             | 45 +------------------
 3 files changed, 20 insertions(+), 70 deletions(-)

diff --git a/packages/react-reconciler/src/ReactFiber.js b/packages/react-reconciler/src/ReactFiber.js
index 83b12a04a799..84b345e34dc2 100644
--- a/packages/react-reconciler/src/ReactFiber.js
+++ b/packages/react-reconciler/src/ReactFiber.js
@@ -79,11 +79,7 @@ import {
 } from './ReactWorkTags';
 import {getComponentNameFromOwner} from 'react-reconciler/src/getComponentNameFromFiber';
 import {isDevToolsPresent} from './ReactFiberDevToolsHook';
-import {
-  resolveClassForHotReloading,
-  resolveFunctionForHotReloading,
-  resolveForwardRefForHotReloading,
-} from './ReactFiberHotReloading';
+import {resolveTypeForHotReloading} from './ReactFiberHotReloading';
 import {NoLanes} from './ReactFiberLane';
 import {
   NoMode,
@@ -427,13 +423,9 @@ export function createWorkInProgress(current: Fiber, pendingProps: any): Fiber {
     switch (workInProgress.tag) {
       case FunctionComponent:
       case SimpleMemoComponent:
-        workInProgress.type = resolveFunctionForHotReloading(current.type);
-        break;
       case ClassComponent:
-        workInProgress.type = resolveClassForHotReloading(current.type);
-        break;
       case ForwardRef:
-        workInProgress.type = resolveForwardRefForHotReloading(current.type);
+        workInProgress.type = resolveTypeForHotReloading(current.type);
         break;
       default:
         break;
@@ -573,11 +565,11 @@ export function createFiberFromTypeAndProps(
     if (shouldConstruct(type)) {
       fiberTag = ClassComponent;
       if (__DEV__) {
-        resolvedType = resolveClassForHotReloading(resolvedType);
+        resolvedType = resolveTypeForHotReloading(resolvedType);
       }
     } else {
       if (__DEV__) {
-        resolvedType = resolveFunctionForHotReloading(resolvedType);
+        resolvedType = resolveTypeForHotReloading(resolvedType);
       }
     }
   } else if (typeof type === 'string') {
@@ -664,9 +656,6 @@ export function createFiberFromTypeAndProps(
             // $FlowFixMe[invalid-compare]
             case REACT_FORWARD_REF_TYPE:
               fiberTag = ForwardRef;
-              if (__DEV__) {
-                resolvedType = resolveForwardRefForHotReloading(resolvedType);
-              }
               break getTag;
             // $FlowFixMe[invalid-compare]
             case REACT_MEMO_TYPE:
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.js b/packages/react-reconciler/src/ReactFiberBeginWork.js
index ed040e6d35a3..3674b070d56b 100644
--- a/packages/react-reconciler/src/ReactFiberBeginWork.js
+++ b/packages/react-reconciler/src/ReactFiberBeginWork.js
@@ -133,9 +133,7 @@ import {
 } from 'shared/ReactSymbols';
 import {setCurrentFiber} from './ReactCurrentFiber';
 import {
-  resolveFunctionForHotReloading,
-  resolveForwardRefForHotReloading,
-  resolveClassForHotReloading,
+  resolveTypeForHotReloading,
   resolveRemountTypeForHotReloading,
 } from './ReactFiberHotReloading';
 
@@ -414,7 +412,16 @@ function updateForwardRef(
   // TODO: current can be non-null here even if the component
   // hasn't yet mounted. This happens after the first render suspends.
   // We'll need to figure out if this is fine or can cause issues.
-  const render = Component.render;
+  let render = Component.render;
+  if (__DEV__) {
+    const resolvedRender = resolveTypeForHotReloading(render);
+    if (resolvedRender !== render) {
+      render = resolvedRender;
+      if (current !== null) {
+        didReceiveUpdate = true;
+      }
+    }
+  }
   const ref = workInProgress.ref;
 
   let propsWithoutRef;
@@ -482,7 +489,7 @@ function updateMemoComponent(
     if (isSimpleFunctionComponent(type) && Component.compare === null) {
       let resolvedType = type;
       if (__DEV__) {
-        resolvedType = resolveFunctionForHotReloading(type);
+        resolvedType = resolveTypeForHotReloading(type);
       }
       // If this is a plain function component without default props,
       // and with only the default shallow comparison, we upgrade it
@@ -2103,8 +2110,7 @@ function mountLazyComponent(
       const resolvedProps = resolveClassComponentProps(Component, props);
       workInProgress.tag = ClassComponent;
       if (__DEV__) {
-        workInProgress.type = Component =
-          resolveClassForHotReloading(Component);
+        workInProgress.type = Component = resolveTypeForHotReloading(Component);
       }
       return updateClassComponent(
         null,
@@ -2117,8 +2123,7 @@ function mountLazyComponent(
       workInProgress.tag = FunctionComponent;
       if (__DEV__) {
         validateFunctionComponentInDev(workInProgress, Component);
-        workInProgress.type = Component =
-          resolveFunctionForHotReloading(Component);
+        workInProgress.type = Component = resolveTypeForHotReloading(Component);
       }
       return updateFunctionComponent(
         null,
@@ -2135,8 +2140,7 @@ function mountLazyComponent(
     if ($$typeof === REACT_FORWARD_REF_TYPE) {
       workInProgress.tag = ForwardRef;
       if (__DEV__) {
-        workInProgress.type = Component =
-          resolveForwardRefForHotReloading(Component);
+        workInProgress.type = Component = resolveTypeForHotReloading(Component);
       }
       return updateForwardRef(
         null,
diff --git a/packages/react-reconciler/src/ReactFiberHotReloading.js b/packages/react-reconciler/src/ReactFiberHotReloading.js
index 199e7486f1c6..daa83798ce2a 100644
--- a/packages/react-reconciler/src/ReactFiberHotReloading.js
+++ b/packages/react-reconciler/src/ReactFiberHotReloading.js
@@ -61,7 +61,7 @@ export const setRefreshHandler = (handler: RefreshHandler | null): void => {
   }
 };
 
-export function resolveFunctionForHotReloading(type: any): any {
+export function resolveTypeForHotReloading(type: any): any {
   if (__DEV__) {
     if (resolveFamily === null) {
       // Hot reloading is disabled.
@@ -78,49 +78,6 @@ export function resolveFunctionForHotReloading(type: any): any {
   }
 }
 
-export function resolveClassForHotReloading(type: any): any {
-  // No implementation differences.
-  return resolveFunctionForHotReloading(type);
-}
-
-export function resolveForwardRefForHotReloading(type: any): any {
-  if (__DEV__) {
-    if (resolveFamily === null) {
-      // Hot reloading is disabled.
-      return type;
-    }
-    const family = resolveFamily(type);
-    if (family === undefined) {
-      // Check if we're dealing with a real forwardRef. Don't want to crash early.
-      if (
-        type !== null &&
-        type !== undefined &&
-        typeof type.render === 'function'
-      ) {
-        // ForwardRef is special because its resolved .type is an object,
-        // but it's possible that we only have its inner render function in the map.
-        // If that inner render function is different, we'll build a new forwardRef type.
-        const currentRender = resolveFunctionForHotReloading(type.render);
-        if (type.render !== currentRender) {
-          const syntheticType = {
-            $$typeof: REACT_FORWARD_REF_TYPE,
-            render: currentRender,
-          };
-          if (type.displayName !== undefined) {
-            (syntheticType as any).displayName = type.displayName;
-          }
-          return syntheticType;
-        }
-      }
-      return type;
-    }
-    // Use the latest known implementation.
-    return family.current;
-  } else {
-    return type;
-  }
-}
-
 export function resolveRemountTypeForHotReloading(
   elementType: any,
   type: any,

From 3b9b59f7cb4ab9f00c84d9b8cdf2e45eb4bb4b88 Mon Sep 17 00:00:00 2001
From: Sophie Alpert <git@sophiebits.com>
Date: Wed, 8 Jul 2026 09:56:38 -0700
Subject: [PATCH 2/4] [Fast Refresh] Derive the fiber tag from the resolved
 type when mounting (#36963)

If an edit changed the kind of the type (e.g. a plain function got
wrapped in memo()), the fiber could get the old kind's tag with the
new kind's type, usually throwing an error.
---
 packages/react-reconciler/src/ReactFiber.js   | 22 +++-----
 .../src/ReactFiberBeginWork.js                | 10 +---
 .../src/__tests__/ReactFresh-test.js          | 55 +++++++++++++++++++
 3 files changed, 67 insertions(+), 20 deletions(-)

diff --git a/packages/react-reconciler/src/ReactFiber.js b/packages/react-reconciler/src/ReactFiber.js
index 84b345e34dc2..94a61d04cb83 100644
--- a/packages/react-reconciler/src/ReactFiber.js
+++ b/packages/react-reconciler/src/ReactFiber.js
@@ -561,18 +561,14 @@ export function createFiberFromTypeAndProps(
   let fiberTag: WorkTag = FunctionComponent;
   // The resolved type is set if we know what the final type will be. I.e. it's not lazy.
   let resolvedType = type;
-  if (typeof type === 'function') {
-    if (shouldConstruct(type)) {
+  if (__DEV__) {
+    resolvedType = resolveTypeForHotReloading(type);
+  }
+  if (typeof resolvedType === 'function') {
+    if (shouldConstruct(resolvedType)) {
       fiberTag = ClassComponent;
-      if (__DEV__) {
-        resolvedType = resolveTypeForHotReloading(resolvedType);
-      }
-    } else {
-      if (__DEV__) {
-        resolvedType = resolveTypeForHotReloading(resolvedType);
-      }
     }
-  } else if (typeof type === 'string') {
+  } else if (typeof resolvedType === 'string') {
     // $FlowFixMe[constant-condition]
     if (supportsResources && supportsSingletons) {
       const hostContext = getHostContext();
@@ -594,7 +590,7 @@ export function createFiberFromTypeAndProps(
       fiberTag = HostComponent;
     }
   } else {
-    getTag: switch (type) {
+    getTag: switch (resolvedType) {
       // $FlowFixMe[invalid-compare]
       case REACT_ACTIVITY_TYPE:
         return createFiberFromActivity(pendingProps, mode, lanes, key);
@@ -642,8 +638,8 @@ export function createFiberFromTypeAndProps(
       // Fall through
       default: {
         // $FlowFixMe[invalid-compare]
-        if (typeof type === 'object' && type !== null) {
-          switch (type.$$typeof) {
+        if (typeof resolvedType === 'object' && resolvedType !== null) {
+          switch (resolvedType.$$typeof) {
             // $FlowFixMe[invalid-compare]
             case REACT_CONTEXT_TYPE:
               fiberTag = ContextProvider;
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.js b/packages/react-reconciler/src/ReactFiberBeginWork.js
index 3674b070d56b..0a1bfc427616 100644
--- a/packages/react-reconciler/src/ReactFiberBeginWork.js
+++ b/packages/react-reconciler/src/ReactFiberBeginWork.js
@@ -2102,6 +2102,9 @@ function mountLazyComponent(
   const props = workInProgress.pendingProps;
   const lazyComponent: LazyComponentType<any, any> = elementType;
   let Component = resolveLazy(lazyComponent);
+  if (__DEV__) {
+    Component = resolveTypeForHotReloading(Component);
+  }
   // Store the unwrapped component in the type.
   workInProgress.type = Component;
 
@@ -2109,9 +2112,6 @@ function mountLazyComponent(
     if (isFunctionClassComponent(Component)) {
       const resolvedProps = resolveClassComponentProps(Component, props);
       workInProgress.tag = ClassComponent;
-      if (__DEV__) {
-        workInProgress.type = Component = resolveTypeForHotReloading(Component);
-      }
       return updateClassComponent(
         null,
         workInProgress,
@@ -2123,7 +2123,6 @@ function mountLazyComponent(
       workInProgress.tag = FunctionComponent;
       if (__DEV__) {
         validateFunctionComponentInDev(workInProgress, Component);
-        workInProgress.type = Component = resolveTypeForHotReloading(Component);
       }
       return updateFunctionComponent(
         null,
@@ -2139,9 +2138,6 @@ function mountLazyComponent(
     // $FlowFixMe[invalid-compare]
     if ($$typeof === REACT_FORWARD_REF_TYPE) {
       workInProgress.tag = ForwardRef;
-      if (__DEV__) {
-        workInProgress.type = Component = resolveTypeForHotReloading(Component);
-      }
       return updateForwardRef(
         null,
         workInProgress,
diff --git a/packages/react-refresh/src/__tests__/ReactFresh-test.js b/packages/react-refresh/src/__tests__/ReactFresh-test.js
index 84f5e9a944c3..d1ffabd70504 100644
--- a/packages/react-refresh/src/__tests__/ReactFresh-test.js
+++ b/packages/react-refresh/src/__tests__/ReactFresh-test.js
@@ -915,6 +915,61 @@ describe('ReactFresh', () => {
     }
   });
 
+  it('can mount an element created before its type changed kinds', async () => {
+    if (__DEV__) {
+      let oldElement;
+      let currentChild = null;
+
+      await act(async () => {
+        await render(() => {
+          function Test() {
+            return <p>hi test</p>;
+          }
+          $RefreshReg$(Test, 'Test');
+          oldElement = <Test />;
+
+          function App() {
+            const [, forceUpdate] = React.useState(0);
+            return (
+              <div onClick={() => forceUpdate(n => n + 1)}>{currentChild}</div>
+            );
+          }
+          $RefreshReg$(App, 'App');
+          return App;
+        });
+      });
+
+      // Change the component kind before it has ever mounted.
+      await act(async () => {
+        await patch(() => {
+          function Test2() {
+            return <p>hi memo</p>;
+          }
+          const Test = React.memo(Test2);
+          $RefreshReg$(Test2, 'Test$React.memo');
+          $RefreshReg$(Test, 'Test');
+
+          function App() {
+            const [, forceUpdate] = React.useState(0);
+            return (
+              <div onClick={() => forceUpdate(n => n + 1)}>{currentChild}</div>
+            );
+          }
+          $RefreshReg$(App, 'App');
+          return App;
+        });
+      });
+
+      // Mount the element created before the edit. The fiber must be
+      // created from the latest type, with the tag matching its kind.
+      currentChild = oldElement;
+      await act(async () => {
+        container.firstChild.click();
+      });
+      expect(container.firstChild.textContent).toBe('hi memo');
+    }
+  });
+
   it('resets state when switching between different component types', async () => {
     if (__DEV__) {
       await act(async () => {

From 4400d6c802df45407e7303a6f84a8dcfabde84ff Mon Sep 17 00:00:00 2001
From: Sophie Alpert <git@sophiebits.com>
Date: Wed, 8 Jul 2026 09:58:35 -0700
Subject: [PATCH 3/4] [Fast Refresh] Make edits to a memo comparison function
 take effect (#36964)

Editing the second argument of memo() previously never took effect
until something else remounted the tree, for two separate reasons:

- When adding a comparison function, we need to switch from
  SimpleMemoComponent to MemoComponent so canPreserveStateBetween
  should return false.
- MemoComponent was missing from the hot reload type resolution in
  createWorkInProgress, so existing fibers kept reading .compare from
  the old memo object forever.

New behavior:

- Adding or removing the comparison function remounts; we need to do
  this when adding (SimpleMemoComponent doesn't support a comparison
  function) so let's also do it when removing.
- Editing the comparison function implementation applies in place
  with state preserved

The TODO in isCompatibleFamilyForHotReloading is removed as that
wasn't the right place to do this check.
---
 packages/react-reconciler/src/ReactFiber.js   |   1 +
 .../src/ReactFiberHotReloading.js             |   2 -
 .../react-refresh/src/ReactFreshRuntime.js    |  12 +-
 .../src/__tests__/ReactFresh-test.js          | 199 ++++++++++++++++++
 4 files changed, 211 insertions(+), 3 deletions(-)

diff --git a/packages/react-reconciler/src/ReactFiber.js b/packages/react-reconciler/src/ReactFiber.js
index 94a61d04cb83..b03e806fc2d2 100644
--- a/packages/react-reconciler/src/ReactFiber.js
+++ b/packages/react-reconciler/src/ReactFiber.js
@@ -423,6 +423,7 @@ export function createWorkInProgress(current: Fiber, pendingProps: any): Fiber {
     switch (workInProgress.tag) {
       case FunctionComponent:
       case SimpleMemoComponent:
+      case MemoComponent:
       case ClassComponent:
       case ForwardRef:
         workInProgress.type = resolveTypeForHotReloading(current.type);
diff --git a/packages/react-reconciler/src/ReactFiberHotReloading.js b/packages/react-reconciler/src/ReactFiberHotReloading.js
index daa83798ce2a..905ddc23b516 100644
--- a/packages/react-reconciler/src/ReactFiberHotReloading.js
+++ b/packages/react-reconciler/src/ReactFiberHotReloading.js
@@ -153,8 +153,6 @@ export function isCompatibleFamilyForHotReloading(
       case MemoComponent:
       case SimpleMemoComponent: {
         if ($$typeofNextType === REACT_MEMO_TYPE) {
-          // TODO: if it was but can no longer be simple,
-          // we shouldn't set this.
           needsCompareFamilies = true;
         } else if ($$typeofNextType === REACT_LAZY_TYPE) {
           needsCompareFamilies = true;
diff --git a/packages/react-refresh/src/ReactFreshRuntime.js b/packages/react-refresh/src/ReactFreshRuntime.js
index 1f4944d6992d..f02215b0c858 100644
--- a/packages/react-refresh/src/ReactFreshRuntime.js
+++ b/packages/react-refresh/src/ReactFreshRuntime.js
@@ -158,6 +158,16 @@ function canPreserveStateBetween(prevType: any, nextType: any) {
     ) {
       return false;
     }
+    // Switching from SimpleMemoComponent to MemoComponent requires a remount;
+    // for symmetry, remount for the reverse too.
+    if (getProperty(prevType, '$$typeof') === REACT_MEMO_TYPE) {
+      if (
+        (getProperty(prevType, 'compare') === null) !==
+        (getProperty(nextType, 'compare') === null)
+      ) {
+        return false;
+      }
+    }
   }
   if (haveEqualSignatures(prevType, nextType)) {
     return true;
@@ -187,7 +197,7 @@ function cloneSet<T>(set: Set<T>): Set<T> {
 }
 
 // This is a safety mechanism to protect against rogue getters and Proxies.
-function getProperty(object: any, property: string) {
+function getProperty(object: any, property: string): any {
   try {
     return object[property];
   } catch (err) {
diff --git a/packages/react-refresh/src/__tests__/ReactFresh-test.js b/packages/react-refresh/src/__tests__/ReactFresh-test.js
index d1ffabd70504..d1aca1b2bf31 100644
--- a/packages/react-refresh/src/__tests__/ReactFresh-test.js
+++ b/packages/react-refresh/src/__tests__/ReactFresh-test.js
@@ -970,6 +970,205 @@ describe('ReactFresh', () => {
     }
   });
 
+  it('can remount when adding or removing a memo comparison function', async () => {
+    if (__DEV__) {
+      await act(async () => {
+        await render(() => {
+          function Test2() {
+            return <p>hi memo</p>;
+          }
+          const Test = React.memo(Test2);
+          $RefreshReg$(Test2, 'Test$React.memo');
+          $RefreshReg$(Test, 'Test');
+          return Test;
+        });
+      });
+
+      // Check the initial render
+      const el = container.firstChild;
+      expect(el.textContent).toBe('hi memo');
+
+      // Patch to add a custom comparison function.
+      // The fiber can no longer be a SimpleMemoComponent.
+      await act(async () => {
+        await patch(() => {
+          function Test2() {
+            return <p>hi memo with compare</p>;
+          }
+          const Test = React.memo(Test2, (prevProps, nextProps) => false);
+          $RefreshReg$(Test2, 'Test$React.memo');
+          $RefreshReg$(Test, 'Test');
+          return Test;
+        });
+      });
+
+      // Check remount
+      expect(container.firstChild).not.toBe(el);
+      const nextEl = container.firstChild;
+      expect(nextEl.textContent).toBe('hi memo with compare');
+
+      // Patch to remove the comparison function again
+      await act(async () => {
+        await patch(() => {
+          function Test2() {
+            return <p>hi memo</p>;
+          }
+          const Test = React.memo(Test2);
+          $RefreshReg$(Test2, 'Test$React.memo');
+          $RefreshReg$(Test, 'Test');
+          return Test;
+        });
+      });
+
+      // Check final remount
+      expect(container.firstChild).not.toBe(nextEl);
+      const newEl = container.firstChild;
+      expect(newEl.textContent).toBe('hi memo');
+    }
+  });
+
+  it('can update a memo comparison function in place', async () => {
+    if (__DEV__) {
+      await act(async () => {
+        await render(() => {
+          function Inner({label}) {
+            return <p>{label}</p>;
+          }
+          const InnerMemo = React.memo(Inner, (prevProps, nextProps) => true);
+          $RefreshReg$(Inner, 'Inner$React.memo');
+          $RefreshReg$(InnerMemo, 'Inner');
+
+          function App() {
+            const [n, setN] = React.useState(1);
+            return (
+              <div onClick={() => setN(c => c + 1)}>
+                <InnerMemo label={'n:' + n} />
+              </div>
+            );
+          }
+          $RefreshReg$(App, 'App');
+          return App;
+        });
+      });
+
+      // Check the initial render
+      const el = container.firstChild;
+      expect(el.textContent).toBe('n:1');
+
+      // The comparison function blocks the update.
+      await act(async () => {
+        el.click();
+      });
+      expect(el.textContent).toBe('n:1');
+
+      // Patch to change only the comparison function implementation.
+      await act(async () => {
+        await patch(() => {
+          function Inner({label}) {
+            return <p>{label}</p>;
+          }
+          const InnerMemo = React.memo(Inner, (prevProps, nextProps) => false);
+          $RefreshReg$(Inner, 'Inner$React.memo');
+          $RefreshReg$(InnerMemo, 'Inner');
+
+          function App() {
+            const [n, setN] = React.useState(1);
+            return (
+              <div onClick={() => setN(c => c + 1)}>
+                <InnerMemo label={'n:' + n} />
+              </div>
+            );
+          }
+          $RefreshReg$(App, 'App');
+          return App;
+        });
+      });
+
+      // No remount, and the previously blocked update shows through
+      // because the new comparison function is used.
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('n:2');
+
+      // The new comparison function applies to future updates too.
+      await act(async () => {
+        el.click();
+      });
+      expect(el.textContent).toBe('n:3');
+    }
+  });
+
+  it('mounts a pre-edit memo element with the latest comparison function', async () => {
+    if (__DEV__) {
+      let oldElement;
+      let newElement;
+      let currentChild = null;
+
+      await act(async () => {
+        await render(() => {
+          function Inner({label}) {
+            return <p>{label}</p>;
+          }
+          const InnerMemo = React.memo(Inner);
+          $RefreshReg$(Inner, 'Inner$React.memo');
+          $RefreshReg$(InnerMemo, 'Inner');
+          oldElement = <InnerMemo label="v1" />;
+
+          function App() {
+            const [, forceUpdate] = React.useState(0);
+            return (
+              <div onClick={() => forceUpdate(n => n + 1)}>{currentChild}</div>
+            );
+          }
+          $RefreshReg$(App, 'App');
+          return App;
+        });
+      });
+
+      // Patch to add a comparison function that blocks all updates,
+      // before the memo has ever mounted.
+      await act(async () => {
+        await patch(() => {
+          function Inner({label}) {
+            return <p>{label}</p>;
+          }
+          const InnerMemo = React.memo(Inner, (prevProps, nextProps) => true);
+          $RefreshReg$(Inner, 'Inner$React.memo');
+          $RefreshReg$(InnerMemo, 'Inner');
+          newElement = <InnerMemo label="v2" />;
+
+          function App() {
+            const [, forceUpdate] = React.useState(0);
+            return (
+              <div onClick={() => forceUpdate(n => n + 1)}>{currentChild}</div>
+            );
+          }
+          $RefreshReg$(App, 'App');
+          return App;
+        });
+      });
+
+      // Mount the element created before the edit. It must resolve to the
+      // latest type rather than mounting in the pre-edit shape.
+      currentChild = oldElement;
+      await act(async () => {
+        container.firstChild.click();
+      });
+      const innerEl = container.firstChild.firstChild;
+      expect(innerEl.textContent).toBe('v1');
+
+      // Switch to the element created after the edit. It belongs to the
+      // same family, so the fiber is reused (no remount)...
+      currentChild = newElement;
+      await act(async () => {
+        container.firstChild.click();
+      });
+      expect(container.firstChild.firstChild).toBe(innerEl);
+      // ...and the comparison function blocks the props update, proving
+      // the fiber mounted with the comparison function in effect.
+      expect(innerEl.textContent).toBe('v1');
+    }
+  });
+
   it('resets state when switching between different component types', async () => {
     if (__DEV__) {
       await act(async () => {

From 5123b063a7b3cd5c42ea2be2ccebe4e68d8b12e3 Mon Sep 17 00:00:00 2001
From: Sophie Alpert <git@sophiebits.com>
Date: Wed, 8 Jul 2026 09:59:07 -0700
Subject: [PATCH 4/4] [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.
---
 .../src/ReactFiberBeginWork.js                |  12 +-
 .../src/ReactFiberHotReloading.js             |  34 +--
 .../src/__tests__/ReactFresh-test.js          | 200 ++++++++++++++++++
 3 files changed, 215 insertions(+), 31 deletions(-)

diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.js b/packages/react-reconciler/src/ReactFiberBeginWork.js
index 0a1bfc427616..815dc8b99612 100644
--- a/packages/react-reconciler/src/ReactFiberBeginWork.js
+++ b/packages/react-reconciler/src/ReactFiberBeginWork.js
@@ -132,10 +132,7 @@ import {
   REACT_CONTEXT_TYPE,
 } from 'shared/ReactSymbols';
 import {setCurrentFiber} from './ReactCurrentFiber';
-import {
-  resolveTypeForHotReloading,
-  resolveRemountTypeForHotReloading,
-} from './ReactFiberHotReloading';
+import {resolveTypeForHotReloading} from './ReactFiberHotReloading';
 
 import {
   mountChildFibers,
@@ -4196,10 +4193,9 @@ function beginWork(
     if (workInProgress._debugNeedsRemount && current !== null) {
       // This will restart the begin phase with a new fiber.
       const copiedFiber = createFiberFromTypeAndProps(
-        resolveRemountTypeForHotReloading(
-          workInProgress.elementType,
-          workInProgress.type,
-        ),
+        // Remount from the fiber's outermost identity; mounting resolves
+        // any inner types to their latest implementations.
+        resolveTypeForHotReloading(workInProgress.elementType),
         workInProgress.key,
         workInProgress.pendingProps,
         workInProgress._debugOwner || null,
diff --git a/packages/react-reconciler/src/ReactFiberHotReloading.js b/packages/react-reconciler/src/ReactFiberHotReloading.js
index 905ddc23b516..16c30f3724d5 100644
--- a/packages/react-reconciler/src/ReactFiberHotReloading.js
+++ b/packages/react-reconciler/src/ReactFiberHotReloading.js
@@ -78,29 +78,6 @@ export function resolveTypeForHotReloading(type: any): any {
   }
 }
 
-export function resolveRemountTypeForHotReloading(
-  elementType: any,
-  type: any,
-): any {
-  if (__DEV__) {
-    if (resolveFamily === null) {
-      // Hot reloading is disabled.
-      return type;
-    }
-    // The elementType is the fiber's public identity, so its family tracks
-    // the latest implementation even when an edit changed the kind of the
-    // type (e.g. memo to a plain function) and `type` still points at the
-    // old inner implementation.
-    const family = resolveFamily(elementType);
-    if (family === undefined) {
-      return type;
-    }
-    return family.current;
-  } else {
-    return type;
-  }
-}
-
 export function isCompatibleFamilyForHotReloading(
   fiber: Fiber,
   element: ReactElement,
@@ -294,6 +271,17 @@ function scheduleFibersWithFamiliesRecursively(
         const outerFamily = resolve(outerCandidateType);
         if (outerFamily !== undefined && staleFamilies.has(outerFamily)) {
           needsRemount = true;
+        } else if (
+          typeof outerCandidateType === 'object' &&
+          outerCandidateType.$$typeof === REACT_LAZY_TYPE
+        ) {
+          const payload = outerCandidateType._payload;
+          if (payload._status === 1 /* Resolved; see ReactLazy */) {
+            const middleFamily = resolve(payload._result.default);
+            if (middleFamily !== undefined && staleFamilies.has(middleFamily)) {
+              needsRemount = true;
+            }
+          }
         }
       }
       if (failedBoundaries !== null) {
diff --git a/packages/react-refresh/src/__tests__/ReactFresh-test.js b/packages/react-refresh/src/__tests__/ReactFresh-test.js
index d1aca1b2bf31..53a71538620d 100644
--- a/packages/react-refresh/src/__tests__/ReactFresh-test.js
+++ b/packages/react-refresh/src/__tests__/ReactFresh-test.js
@@ -1169,6 +1169,206 @@ describe('ReactFresh', () => {
     }
   });
 
+  it('can remount lazy(memo()) when adding a comparison function', async () => {
+    if (__DEV__) {
+      let resolve;
+      await render(() => {
+        function Hello() {
+          return <p>hi memo</p>;
+        }
+        const Inner = React.memo(Hello);
+        $RefreshReg$(Hello, 'Hello');
+        $RefreshReg$(Inner, 'Inner');
+
+        const Outer = React.lazy(
+          () =>
+            new Promise(_resolve => {
+              resolve = () => _resolve({default: Inner});
+            }),
+        );
+        $RefreshReg$(Outer, 'Outer');
+
+        function App() {
+          return (
+            <React.Suspense fallback={<p>Loading</p>}>
+              <Outer />
+            </React.Suspense>
+          );
+        }
+        $RefreshReg$(App, 'App');
+        return App;
+      });
+
+      expect(container.textContent).toBe('Loading');
+      await act(() => {
+        resolve();
+      });
+      expect(container.textContent).toBe('hi memo');
+      const el = container.firstChild;
+
+      // Perform a hot update that adds a comparison function. The module
+      // creating the lazy also re-runs, like when an edit propagates.
+      await patch(() => {
+        function Hello() {
+          return <p>hi memo with compare</p>;
+        }
+        const Inner = React.memo(Hello, (prevProps, nextProps) => false);
+        $RefreshReg$(Hello, 'Hello');
+        $RefreshReg$(Inner, 'Inner');
+
+        const Outer = React.lazy(
+          () =>
+            new Promise(_resolve => {
+              resolve = () => _resolve({default: Inner});
+            }),
+        );
+        $RefreshReg$(Outer, 'Outer');
+
+        function App() {
+          return (
+            <React.Suspense fallback={<p>Loading</p>}>
+              <Outer />
+            </React.Suspense>
+          );
+        }
+        $RefreshReg$(App, 'App');
+        return App;
+      });
+
+      // The shape change requires a remount. It goes through the latest
+      // lazy type, which suspends until it resolves. The boundary shows
+      // the fallback while the previous content stays hidden in the DOM.
+      expect(container.textContent).toBe('hi memoLoading');
+      await act(() => {
+        resolve();
+      });
+      expect(container.textContent).toBe('hi memo with compare');
+      expect(container.firstChild).not.toBe(el);
+    }
+  });
+
+  it('can remount lazy(memo()) when adding a comparison function without re-creating the lazy', async () => {
+    if (__DEV__) {
+      let resolve;
+      await render(() => {
+        function Hello() {
+          return <p>hi memo</p>;
+        }
+        const Inner = React.memo(Hello);
+        $RefreshReg$(Hello, 'Hello');
+        $RefreshReg$(Inner, 'Inner');
+
+        const Outer = React.lazy(
+          () =>
+            new Promise(_resolve => {
+              resolve = () => _resolve({default: Inner});
+            }),
+        );
+        $RefreshReg$(Outer, 'Outer');
+
+        function App() {
+          return (
+            <React.Suspense fallback={<p>Loading</p>}>
+              <Outer />
+            </React.Suspense>
+          );
+        }
+        $RefreshReg$(App, 'App');
+        return App;
+      });
+
+      expect(container.textContent).toBe('Loading');
+      await act(() => {
+        resolve();
+      });
+      expect(container.textContent).toBe('hi memo');
+      const el = container.firstChild;
+
+      // Only the lazily loaded module re-runs this time, like when an
+      // edit is contained to it (the lazy type is not re-created, so the
+      // remount cannot go through it).
+      await patch(() => {
+        function Hello() {
+          return <p>hi memo with compare</p>;
+        }
+        const Inner = React.memo(Hello, (prevProps, nextProps) => false);
+        $RefreshReg$(Hello, 'Hello');
+        $RefreshReg$(Inner, 'Inner');
+        return Inner;
+      });
+
+      // The remount goes through the old lazy, whose payload is already
+      // resolved, so it doesn't suspend.
+      expect(container.textContent).toBe('hi memo with compare');
+      expect(container.firstChild).not.toBe(el);
+    }
+  });
+
+  it('can remount an unregistered memo wrapper without losing the wrapper', async () => {
+    if (__DEV__) {
+      let innerRenders = 0;
+      await act(async () => {
+        await render(() => {
+          function Inner({label}) {
+            innerRenders++;
+            return <p>{label}</p>;
+          }
+          $RefreshReg$(Inner, 'Inner');
+          $RefreshSig$(Inner, 'sig1');
+          // The wrapper is deliberately not registered, like a wrapper
+          // created inside a third-party HOC.
+          const InnerMemo = React.memo(Inner);
+
+          function App() {
+            const [, forceUpdate] = React.useState(0);
+            return (
+              <div onClick={() => forceUpdate(n => n + 1)}>
+                <InnerMemo label="hi" />
+              </div>
+            );
+          }
+          $RefreshReg$(App, 'App');
+          return App;
+        });
+      });
+
+      expect(container.textContent).toBe('hi');
+      expect(innerRenders).toBe(1);
+
+      // The memo blocks re-renders with equal props.
+      await act(async () => {
+        container.firstChild.click();
+      });
+      expect(innerRenders).toBe(1);
+
+      // Force a remount by changing the inner function's signature.
+      // Only the inner function's module re-runs; the wrapper and the
+      // element referencing it are not re-created.
+      await act(async () => {
+        await patch(() => {
+          function Inner({label}) {
+            innerRenders++;
+            return <p>{label}</p>;
+          }
+          $RefreshReg$(Inner, 'Inner');
+          $RefreshSig$(Inner, 'sig2');
+          return Inner;
+        });
+      });
+      expect(innerRenders).toBe(2);
+      const innerEl = container.firstChild.firstChild;
+
+      // The remounted fiber must still be a memo: equal props stay
+      // blocked, and the fiber reconciles against the original element
+      // instead of being replaced again.
+      await act(async () => {
+        container.firstChild.click();
+      });
+      expect(innerRenders).toBe(2);
+      expect(container.firstChild.firstChild).toBe(innerEl);
+    }
+  });
+
   it('resets state when switching between different component types', async () => {
     if (__DEV__) {
       await act(async () => {