Skip to content

Commit de1eaa2

Browse files
Biki-dasBiki-das
authored
Fix:- Improve HOC support and state preservation in React Refresh (facebook#30660)
## Summary This fixes facebook#30659 , the issue was how the state was preserved and needed special cases for the forward and memo, have also added tests related to the same. ## How did you test this change? `yarn test packages/react-refresh/src/__tests__/ReactFresh-test.js` ![Screenshot 2024-08-12 at 4 27 39 PM](https://github.com/user-attachments/assets/2b597a62-c45f-443b-acfc-3232962ba0a3)
1 parent ae9017c commit de1eaa2

File tree

2 files changed

+293
-11
lines changed

2 files changed

+293
-11
lines changed

‎packages/react-refresh/src/ReactFreshRuntime.js‎

Lines changed: 57 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,21 @@ function canPreserveStateBetween(prevType: any, nextType: any) {
146146
if (isReactClass(prevType) || isReactClass(nextType)) {
147147
return false;
148148
}
149+
150+
if (typeof prevType !== typeof nextType) {
151+
return false;
152+
} else if (
153+
typeof prevType === 'object' &&
154+
prevType !== null &&
155+
nextType !== null
156+
) {
157+
if (
158+
getProperty(prevType, '$$typeof') !== getProperty(nextType, '$$typeof')
159+
) {
160+
return false;
161+
}
162+
}
163+
149164
if (haveEqualSignatures(prevType, nextType)) {
150165
return true;
151166
}
@@ -183,6 +198,18 @@ function getProperty(object: any, property: string) {
183198
}
184199
}
185200

201+
function registerRefreshUpdate(
202+
update: RefreshUpdate,
203+
family: Family,
204+
shouldPreserveState: boolean,
205+
) {
206+
if (shouldPreserveState) {
207+
update.updatedFamilies.add(family);
208+
} else {
209+
update.staleFamilies.add(family);
210+
}
211+
}
212+
186213
export function performReactRefresh(): RefreshUpdate | null {
187214
if (!__DEV__) {
188215
throw new Error(
@@ -200,6 +227,11 @@ export function performReactRefresh(): RefreshUpdate | null {
200227
try {
201228
const staleFamilies = new Set<Family>();
202229
const updatedFamilies = new Set<Family>();
230+
// TODO: rename these fields to something more meaningful.
231+
const update: RefreshUpdate = {
232+
updatedFamilies, // Families that will re-render preserving state
233+
staleFamilies, // Families that will be remounted
234+
};
203235

204236
const updates = pendingUpdates;
205237
pendingUpdates = [];
@@ -211,19 +243,33 @@ export function performReactRefresh(): RefreshUpdate | null {
211243
updatedFamiliesByType.set(nextType, family);
212244
family.current = nextType;
213245

214-
// Determine whether this should be a re-render or a re-mount.
215-
if (canPreserveStateBetween(prevType, nextType)) {
216-
updatedFamilies.add(family);
217-
} else {
218-
staleFamilies.add(family);
246+
const shouldPreserveState = canPreserveStateBetween(prevType, nextType);
247+
248+
if (typeof prevType === 'object' && prevType !== null) {
249+
const nextFamily = {
250+
current:
251+
getProperty(nextType, '$$typeof') === REACT_FORWARD_REF_TYPE
252+
? nextType.render
253+
: getProperty(nextType, '$$typeof') === REACT_MEMO_TYPE
254+
? nextType.type
255+
: nextType,
256+
};
257+
switch (getProperty(prevType, '$$typeof')) {
258+
case REACT_FORWARD_REF_TYPE: {
259+
updatedFamiliesByType.set(prevType.render, nextFamily);
260+
registerRefreshUpdate(update, nextFamily, shouldPreserveState);
261+
break;
262+
}
263+
case REACT_MEMO_TYPE:
264+
updatedFamiliesByType.set(prevType.type, nextFamily);
265+
registerRefreshUpdate(update, nextFamily, shouldPreserveState);
266+
break;
267+
}
219268
}
220-
});
221269

222-
// TODO: rename these fields to something more meaningful.
223-
const update: RefreshUpdate = {
224-
updatedFamilies, // Families that will re-render preserving state
225-
staleFamilies, // Families that will be remounted
226-
};
270+
// Determine whether this should be a re-render or a re-mount.
271+
registerRefreshUpdate(update, family, shouldPreserveState);
272+
});
227273

228274
helpersByRendererID.forEach(helpers => {
229275
// Even if there are no roots, set the handler on first update.

‎packages/react-refresh/src/__tests__/ReactFresh-test.js‎

Lines changed: 236 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -699,6 +699,242 @@ describe('ReactFresh', () => {
699699
}
700700
});
701701

702+
it('can remount when change function to memo', async () => {
703+
if (__DEV__) {
704+
await act(async () => {
705+
await render(() => {
706+
function Test() {
707+
return <p>hi test</p>;
708+
}
709+
$RefreshReg$(Test, 'Test');
710+
return Test;
711+
});
712+
});
713+
714+
// Check the initial render
715+
const el = container.firstChild;
716+
expect(el.textContent).toBe('hi test');
717+
718+
// Patch to change function to memo
719+
await act(async () => {
720+
await patch(() => {
721+
function Test2() {
722+
return <p>hi memo</p>;
723+
}
724+
const Test = React.memo(Test2);
725+
$RefreshReg$(Test2, 'Test2');
726+
$RefreshReg$(Test, 'Test');
727+
return Test;
728+
});
729+
});
730+
731+
// Check remount
732+
expect(container.firstChild).not.toBe(el);
733+
const nextEl = container.firstChild;
734+
expect(nextEl.textContent).toBe('hi memo');
735+
736+
// Patch back to original function
737+
await act(async () => {
738+
await patch(() => {
739+
function Test() {
740+
return <p>hi test</p>;
741+
}
742+
$RefreshReg$(Test, 'Test');
743+
return Test;
744+
});
745+
});
746+
747+
// Check final remount
748+
expect(container.firstChild).not.toBe(nextEl);
749+
const newEl = container.firstChild;
750+
expect(newEl.textContent).toBe('hi test');
751+
}
752+
});
753+
754+
it('can remount when change memo to forwardRef', async () => {
755+
if (__DEV__) {
756+
await act(async () => {
757+
await render(() => {
758+
function Test2() {
759+
return <p>hi memo</p>;
760+
}
761+
const Test = React.memo(Test2);
762+
$RefreshReg$(Test2, 'Test2');
763+
$RefreshReg$(Test, 'Test');
764+
return Test;
765+
});
766+
});
767+
// Check the initial render
768+
const el = container.firstChild;
769+
expect(el.textContent).toBe('hi memo');
770+
771+
// Patch to change memo to forwardRef
772+
await act(async () => {
773+
await patch(() => {
774+
function Test2() {
775+
return <p>hi forwardRef</p>;
776+
}
777+
const Test = React.forwardRef(Test2);
778+
$RefreshReg$(Test2, 'Test2');
779+
$RefreshReg$(Test, 'Test');
780+
return Test;
781+
});
782+
});
783+
// Check remount
784+
expect(container.firstChild).not.toBe(el);
785+
const nextEl = container.firstChild;
786+
expect(nextEl.textContent).toBe('hi forwardRef');
787+
788+
// Patch back to memo
789+
await act(async () => {
790+
await patch(() => {
791+
function Test2() {
792+
return <p>hi memo</p>;
793+
}
794+
const Test = React.memo(Test2);
795+
$RefreshReg$(Test2, 'Test2');
796+
$RefreshReg$(Test, 'Test');
797+
return Test;
798+
});
799+
});
800+
// Check final remount
801+
expect(container.firstChild).not.toBe(nextEl);
802+
const newEl = container.firstChild;
803+
expect(newEl.textContent).toBe('hi memo');
804+
}
805+
});
806+
807+
it('can remount when change function to forwardRef', async () => {
808+
if (__DEV__) {
809+
await act(async () => {
810+
await render(() => {
811+
function Test() {
812+
return <p>hi test</p>;
813+
}
814+
$RefreshReg$(Test, 'Test');
815+
return Test;
816+
});
817+
});
818+
819+
// Check the initial render
820+
const el = container.firstChild;
821+
expect(el.textContent).toBe('hi test');
822+
823+
// Patch to change function to forwardRef
824+
await act(async () => {
825+
await patch(() => {
826+
function Test2() {
827+
return <p>hi forwardRef</p>;
828+
}
829+
const Test = React.forwardRef(Test2);
830+
$RefreshReg$(Test2, 'Test2');
831+
$RefreshReg$(Test, 'Test');
832+
return Test;
833+
});
834+
});
835+
836+
// Check remount
837+
expect(container.firstChild).not.toBe(el);
838+
const nextEl = container.firstChild;
839+
expect(nextEl.textContent).toBe('hi forwardRef');
840+
841+
// Patch back to a new function
842+
await act(async () => {
843+
await patch(() => {
844+
function Test() {
845+
return <p>hi test1</p>;
846+
}
847+
$RefreshReg$(Test, 'Test');
848+
return Test;
849+
});
850+
});
851+
852+
// Check final remount
853+
expect(container.firstChild).not.toBe(nextEl);
854+
const newEl = container.firstChild;
855+
expect(newEl.textContent).toBe('hi test1');
856+
}
857+
});
858+
859+
it('resets state when switching between different component types', async () => {
860+
if (__DEV__) {
861+
await act(async () => {
862+
await render(() => {
863+
function Test() {
864+
const [count, setCount] = React.useState(0);
865+
return (
866+
<div onClick={() => setCount(c => c + 1)}>count: {count}</div>
867+
);
868+
}
869+
$RefreshReg$(Test, 'Test');
870+
return Test;
871+
});
872+
});
873+
874+
expect(container.firstChild.textContent).toBe('count: 0');
875+
await act(async () => {
876+
container.firstChild.click();
877+
});
878+
expect(container.firstChild.textContent).toBe('count: 1');
879+
880+
await act(async () => {
881+
await patch(() => {
882+
function Test2() {
883+
const [count, setCount] = React.useState(0);
884+
return (
885+
<div onClick={() => setCount(c => c + 1)}>count: {count}</div>
886+
);
887+
}
888+
const Test = React.memo(Test2);
889+
$RefreshReg$(Test2, 'Test2');
890+
$RefreshReg$(Test, 'Test');
891+
return Test;
892+
});
893+
});
894+
895+
expect(container.firstChild.textContent).toBe('count: 0');
896+
await act(async () => {
897+
container.firstChild.click();
898+
});
899+
expect(container.firstChild.textContent).toBe('count: 1');
900+
901+
await act(async () => {
902+
await patch(() => {
903+
const Test = React.forwardRef((props, ref) => {
904+
const [count, setCount] = React.useState(0);
905+
const handleClick = () => setCount(c => c + 1);
906+
907+
// Ensure ref is extensible
908+
const divRef = React.useRef(null);
909+
React.useEffect(() => {
910+
if (ref) {
911+
if (typeof ref === 'function') {
912+
ref(divRef.current);
913+
} else if (Object.isExtensible(ref)) {
914+
ref.current = divRef.current;
915+
}
916+
}
917+
}, [ref]);
918+
919+
return (
920+
<div ref={divRef} onClick={handleClick}>
921+
count: {count}
922+
</div>
923+
);
924+
});
925+
$RefreshReg$(Test, 'Test');
926+
return Test;
927+
});
928+
});
929+
930+
expect(container.firstChild.textContent).toBe('count: 0');
931+
await act(async () => {
932+
container.firstChild.click();
933+
});
934+
expect(container.firstChild.textContent).toBe('count: 1');
935+
}
936+
});
937+
702938
it('can update simple memo function in isolation', async () => {
703939
if (__DEV__) {
704940
await render(() => {

0 commit comments

Comments
 (0)