@samitouri / QOS-React / commits / 4400d6c802

[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.

Sophie Alpert committed Jul 8, 2026 at 09:58 UTC 4400d6c802df45407e7303a6f84a8dcfabde84ff
4 files changed +211 -3
packages/react-reconciler/src/ReactFiber.js
+1
@@ -423,6 +423,7 @@ export function createWorkInProgress(current: Fiber, pendingProps: any): Fiber {
423 switch (workInProgress.tag) {
424 case FunctionComponent:
425 case SimpleMemoComponent:
426 + case MemoComponent:
427 case ClassComponent:
428 case ForwardRef:
429 workInProgress.type = resolveTypeForHotReloading(current.type);
packages/react-reconciler/src/ReactFiberHotReloading.js
-2
@@ -153,8 +153,6 @@ export function isCompatibleFamilyForHotReloading(
153 case MemoComponent:
154 case SimpleMemoComponent: {
155 if ($$typeofNextType === REACT_MEMO_TYPE) {
156 - // TODO: if it was but can no longer be simple,
157 - // we shouldn't set this.
156 needsCompareFamilies = true;
157 } else if ($$typeofNextType === REACT_LAZY_TYPE) {
158 needsCompareFamilies = true;
packages/react-refresh/src/ReactFreshRuntime.js
+11 -1
@@ -158,6 +158,16 @@ function canPreserveStateBetween(prevType: any, nextType: any) {
158 ) {
159 return false;
160 }
161 + // Switching from SimpleMemoComponent to MemoComponent requires a remount;
162 + // for symmetry, remount for the reverse too.
163 + if (getProperty(prevType, '$$typeof') === REACT_MEMO_TYPE) {
164 + if (
165 + (getProperty(prevType, 'compare') === null) !==
166 + (getProperty(nextType, 'compare') === null)
167 + ) {
168 + return false;
169 + }
170 + }
171 }
172 if (haveEqualSignatures(prevType, nextType)) {
173 return true;
@@ -187,7 +197,7 @@ function cloneSet<T>(set: Set<T>): Set<T> {
197 }
198
199 // This is a safety mechanism to protect against rogue getters and Proxies.
190 -function getProperty(object: any, property: string) {
200 +function getProperty(object: any, property: string): any {
201 try {
202 return object[property];
203 } catch (err) {
packages/react-refresh/src/__tests__/ReactFresh-test.js
+199
@@ -970,6 +970,205 @@ describe('ReactFresh', () => {
970 }
971 });
972
973 + it('can remount when adding or removing a memo comparison function', async () => {
974 + if (__DEV__) {
975 + await act(async () => {
976 + await render(() => {
977 + function Test2() {
978 + return <p>hi memo</p>;
979 + }
980 + const Test = React.memo(Test2);
981 + $RefreshReg$(Test2, 'Test$React.memo');
982 + $RefreshReg$(Test, 'Test');
983 + return Test;
984 + });
985 + });
986 +
987 + // Check the initial render
988 + const el = container.firstChild;
989 + expect(el.textContent).toBe('hi memo');
990 +
991 + // Patch to add a custom comparison function.
992 + // The fiber can no longer be a SimpleMemoComponent.
993 + await act(async () => {
994 + await patch(() => {
995 + function Test2() {
996 + return <p>hi memo with compare</p>;
997 + }
998 + const Test = React.memo(Test2, (prevProps, nextProps) => false);
999 + $RefreshReg$(Test2, 'Test$React.memo');
1000 + $RefreshReg$(Test, 'Test');
1001 + return Test;
1002 + });
1003 + });
1004 +
1005 + // Check remount
1006 + expect(container.firstChild).not.toBe(el);
1007 + const nextEl = container.firstChild;
1008 + expect(nextEl.textContent).toBe('hi memo with compare');
1009 +
1010 + // Patch to remove the comparison function again
1011 + await act(async () => {
1012 + await patch(() => {
1013 + function Test2() {
1014 + return <p>hi memo</p>;
1015 + }
1016 + const Test = React.memo(Test2);
1017 + $RefreshReg$(Test2, 'Test$React.memo');
1018 + $RefreshReg$(Test, 'Test');
1019 + return Test;
1020 + });
1021 + });
1022 +
1023 + // Check final remount
1024 + expect(container.firstChild).not.toBe(nextEl);
1025 + const newEl = container.firstChild;
1026 + expect(newEl.textContent).toBe('hi memo');
1027 + }
1028 + });
1029 +
1030 + it('can update a memo comparison function in place', async () => {
1031 + if (__DEV__) {
1032 + await act(async () => {
1033 + await render(() => {
1034 + function Inner({label}) {
1035 + return <p>{label}</p>;
1036 + }
1037 + const InnerMemo = React.memo(Inner, (prevProps, nextProps) => true);
1038 + $RefreshReg$(Inner, 'Inner$React.memo');
1039 + $RefreshReg$(InnerMemo, 'Inner');
1040 +
1041 + function App() {
1042 + const [n, setN] = React.useState(1);
1043 + return (
1044 + <div onClick={() => setN(c => c + 1)}>
1045 + <InnerMemo label={'n:' + n} />
1046 + </div>
1047 + );
1048 + }
1049 + $RefreshReg$(App, 'App');
1050 + return App;
1051 + });
1052 + });
1053 +
1054 + // Check the initial render
1055 + const el = container.firstChild;
1056 + expect(el.textContent).toBe('n:1');
1057 +
1058 + // The comparison function blocks the update.
1059 + await act(async () => {
1060 + el.click();
1061 + });
1062 + expect(el.textContent).toBe('n:1');
1063 +
1064 + // Patch to change only the comparison function implementation.
1065 + await act(async () => {
1066 + await patch(() => {
1067 + function Inner({label}) {
1068 + return <p>{label}</p>;
1069 + }
1070 + const InnerMemo = React.memo(Inner, (prevProps, nextProps) => false);
1071 + $RefreshReg$(Inner, 'Inner$React.memo');
1072 + $RefreshReg$(InnerMemo, 'Inner');
1073 +
1074 + function App() {
1075 + const [n, setN] = React.useState(1);
1076 + return (
1077 + <div onClick={() => setN(c => c + 1)}>
1078 + <InnerMemo label={'n:' + n} />
1079 + </div>
1080 + );
1081 + }
1082 + $RefreshReg$(App, 'App');
1083 + return App;
1084 + });
1085 + });
1086 +
1087 + // No remount, and the previously blocked update shows through
1088 + // because the new comparison function is used.
1089 + expect(container.firstChild).toBe(el);
1090 + expect(el.textContent).toBe('n:2');
1091 +
1092 + // The new comparison function applies to future updates too.
1093 + await act(async () => {
1094 + el.click();
1095 + });
1096 + expect(el.textContent).toBe('n:3');
1097 + }
1098 + });
1099 +
1100 + it('mounts a pre-edit memo element with the latest comparison function', async () => {
1101 + if (__DEV__) {
1102 + let oldElement;
1103 + let newElement;
1104 + let currentChild = null;
1105 +
1106 + await act(async () => {
1107 + await render(() => {
1108 + function Inner({label}) {
1109 + return <p>{label}</p>;
1110 + }
1111 + const InnerMemo = React.memo(Inner);
1112 + $RefreshReg$(Inner, 'Inner$React.memo');
1113 + $RefreshReg$(InnerMemo, 'Inner');
1114 + oldElement = <InnerMemo label="v1" />;
1115 +
1116 + function App() {
1117 + const [, forceUpdate] = React.useState(0);
1118 + return (
1119 + <div onClick={() => forceUpdate(n => n + 1)}>{currentChild}</div>
1120 + );
1121 + }
1122 + $RefreshReg$(App, 'App');
1123 + return App;
1124 + });
1125 + });
1126 +
1127 + // Patch to add a comparison function that blocks all updates,
1128 + // before the memo has ever mounted.
1129 + await act(async () => {
1130 + await patch(() => {
1131 + function Inner({label}) {
1132 + return <p>{label}</p>;
1133 + }
1134 + const InnerMemo = React.memo(Inner, (prevProps, nextProps) => true);
1135 + $RefreshReg$(Inner, 'Inner$React.memo');
1136 + $RefreshReg$(InnerMemo, 'Inner');
1137 + newElement = <InnerMemo label="v2" />;
1138 +
1139 + function App() {
1140 + const [, forceUpdate] = React.useState(0);
1141 + return (
1142 + <div onClick={() => forceUpdate(n => n + 1)}>{currentChild}</div>
1143 + );
1144 + }
1145 + $RefreshReg$(App, 'App');
1146 + return App;
1147 + });
1148 + });
1149 +
1150 + // Mount the element created before the edit. It must resolve to the
1151 + // latest type rather than mounting in the pre-edit shape.
1152 + currentChild = oldElement;
1153 + await act(async () => {
1154 + container.firstChild.click();
1155 + });
1156 + const innerEl = container.firstChild.firstChild;
1157 + expect(innerEl.textContent).toBe('v1');
1158 +
1159 + // Switch to the element created after the edit. It belongs to the
1160 + // same family, so the fiber is reused (no remount)...
1161 + currentChild = newElement;
1162 + await act(async () => {
1163 + container.firstChild.click();
1164 + });
1165 + expect(container.firstChild.firstChild).toBe(innerEl);
1166 + // ...and the comparison function blocks the props update, proving
1167 + // the fiber mounted with the comparison function in effect.
1168 + expect(innerEl.textContent).toBe('v1');
1169 + }
1170 + });
1171 +
1172 it('resets state when switching between different component types', async () => {
1173 if (__DEV__) {
1174 await act(async () => {