Fix:- Improve HOC support and state preservation in React Refresh (#30660)
## Summary This fixes #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` 
BIKI DAS committed
Jan 24, 2025 at 17:18 UTC
de1eaa2655ee29933425695ae0af756bf6713147
2 files changed
+293
-11
packages/react-refresh/src/ReactFreshRuntime.js
+57
-11
@@ -146,6 +146,21 @@ function canPreserveStateBetween(prevType: any, nextType: any) {
146
if (isReactClass(prevType) || isReactClass(nextType)) {
147
return false;
148
}
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
+
164
if (haveEqualSignatures(prevType, nextType)) {
165
return true;
166
}
@@ -183,6 +198,18 @@ function getProperty(object: any, property: string) {
198
}
199
}
200
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
+
213
export function performReactRefresh(): RefreshUpdate | null {
214
if (!__DEV__) {
215
throw new Error(
@@ -200,6 +227,11 @@ export function performReactRefresh(): RefreshUpdate | null {
227
try {
228
const staleFamilies = new Set<Family>();
229
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
+ };
235
236
const updates = pendingUpdates;
237
pendingUpdates = [];
@@ -211,19 +243,33 @@ export function performReactRefresh(): RefreshUpdate | null {
243
updatedFamiliesByType.set(nextType, family);
244
family.current = nextType;
245
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
+ }
268
}
220
- });
269
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
+ });
273
274
helpersByRendererID.forEach(helpers => {
275
// Even if there are no roots, set the handler on first update.
packages/react-refresh/src/__tests__/ReactFresh-test.js
+236
@@ -699,6 +699,242 @@ describe('ReactFresh', () => {
699
}
700
});
701
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
+
938
it('can update simple memo function in isolation', async () => {
939
if (__DEV__) {
940
await render(() => {