Skip to content

Commit

Permalink
Fix:- Improve HOC support and state preservation in React Refresh (fa…
Browse files Browse the repository at this point in the history 
…cebook#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)
  • Loading branch information
@Biki-das
Biki-das authored Jan 24, 2025
1 parent ae9017c commit de1eaa2
Show file tree
Hide file tree
Showing 2 changed files with 293 additions and 11 deletions.
68 changes: 57 additions & 11 deletions packages/react-refresh/src/ReactFreshRuntime.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,21 @@ function canPreserveStateBetween(prevType: any, nextType: any) {
if (isReactClass(prevType) || isReactClass(nextType)) {
return false;
}

if (typeof prevType !== typeof nextType) {
return false;
} else if (
typeof prevType === 'object' &&
prevType !== null &&
nextType !== null
) {
if (
getProperty(prevType, '$$typeof') !== getProperty(nextType, '$$typeof')
) {
return false;
}
}

if (haveEqualSignatures(prevType, nextType)) {
return true;
}
Expand Down Expand Up @@ -183,6 +198,18 @@ function getProperty(object: any, property: string) {
}
}

function registerRefreshUpdate(
update: RefreshUpdate,
family: Family,
shouldPreserveState: boolean,
) {
if (shouldPreserveState) {
update.updatedFamilies.add(family);
} else {
update.staleFamilies.add(family);
}
}

export function performReactRefresh(): RefreshUpdate | null {
if (!__DEV__) {
throw new Error(
Expand All @@ -200,6 +227,11 @@ export function performReactRefresh(): RefreshUpdate | null {
try {
const staleFamilies = new Set<Family>();
const updatedFamilies = new Set<Family>();
// TODO: rename these fields to something more meaningful.
const update: RefreshUpdate = {
updatedFamilies, // Families that will re-render preserving state
staleFamilies, // Families that will be remounted
};

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

// Determine whether this should be a re-render or a re-mount.
if (canPreserveStateBetween(prevType, nextType)) {
updatedFamilies.add(family);
} else {
staleFamilies.add(family);
const shouldPreserveState = canPreserveStateBetween(prevType, nextType);

if (typeof prevType === 'object' && prevType !== null) {
const nextFamily = {
current:
getProperty(nextType, '$$typeof') === REACT_FORWARD_REF_TYPE
? nextType.render
: getProperty(nextType, '$$typeof') === REACT_MEMO_TYPE
? nextType.type
: nextType,
};
switch (getProperty(prevType, '$$typeof')) {
case REACT_FORWARD_REF_TYPE: {
updatedFamiliesByType.set(prevType.render, nextFamily);
registerRefreshUpdate(update, nextFamily, shouldPreserveState);
break;
}
case REACT_MEMO_TYPE:
updatedFamiliesByType.set(prevType.type, nextFamily);
registerRefreshUpdate(update, nextFamily, shouldPreserveState);
break;
}
}
});

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

helpersByRendererID.forEach(helpers => {
// Even if there are no roots, set the handler on first update.
Expand Down
236 changes: 236 additions & 0 deletions packages/react-refresh/src/__tests__/ReactFresh-test.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -699,6 +699,242 @@ describe('ReactFresh', () => {
}
});

it('can remount when change function to memo', async () => {
if (__DEV__) {
await act(async () => {
await render(() => {
function Test() {
return <p>hi test</p>;
}
$RefreshReg$(Test, 'Test');
return Test;
});
});

// Check the initial render
const el = container.firstChild;
expect(el.textContent).toBe('hi test');

// Patch to change function to memo
await act(async () => {
await patch(() => {
function Test2() {
return <p>hi memo</p>;
}
const Test = React.memo(Test2);
$RefreshReg$(Test2, 'Test2');
$RefreshReg$(Test, 'Test');
return Test;
});
});

// Check remount
expect(container.firstChild).not.toBe(el);
const nextEl = container.firstChild;
expect(nextEl.textContent).toBe('hi memo');

// Patch back to original function
await act(async () => {
await patch(() => {
function Test() {
return <p>hi test</p>;
}
$RefreshReg$(Test, 'Test');
return Test;
});
});

// Check final remount
expect(container.firstChild).not.toBe(nextEl);
const newEl = container.firstChild;
expect(newEl.textContent).toBe('hi test');
}
});

it('can remount when change memo to forwardRef', async () => {
if (__DEV__) {
await act(async () => {
await render(() => {
function Test2() {
return <p>hi memo</p>;
}
const Test = React.memo(Test2);
$RefreshReg$(Test2, 'Test2');
$RefreshReg$(Test, 'Test');
return Test;
});
});
// Check the initial render
const el = container.firstChild;
expect(el.textContent).toBe('hi memo');

// Patch to change memo to forwardRef
await act(async () => {
await patch(() => {
function Test2() {
return <p>hi forwardRef</p>;
}
const Test = React.forwardRef(Test2);
$RefreshReg$(Test2, 'Test2');
$RefreshReg$(Test, 'Test');
return Test;
});
});
// Check remount
expect(container.firstChild).not.toBe(el);
const nextEl = container.firstChild;
expect(nextEl.textContent).toBe('hi forwardRef');

// Patch back to memo
await act(async () => {
await patch(() => {
function Test2() {
return <p>hi memo</p>;
}
const Test = React.memo(Test2);
$RefreshReg$(Test2, 'Test2');
$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 remount when change function to forwardRef', async () => {
if (__DEV__) {
await act(async () => {
await render(() => {
function Test() {
return <p>hi test</p>;
}
$RefreshReg$(Test, 'Test');
return Test;
});
});

// Check the initial render
const el = container.firstChild;
expect(el.textContent).toBe('hi test');

// Patch to change function to forwardRef
await act(async () => {
await patch(() => {
function Test2() {
return <p>hi forwardRef</p>;
}
const Test = React.forwardRef(Test2);
$RefreshReg$(Test2, 'Test2');
$RefreshReg$(Test, 'Test');
return Test;
});
});

// Check remount
expect(container.firstChild).not.toBe(el);
const nextEl = container.firstChild;
expect(nextEl.textContent).toBe('hi forwardRef');

// Patch back to a new function
await act(async () => {
await patch(() => {
function Test() {
return <p>hi test1</p>;
}
$RefreshReg$(Test, 'Test');
return Test;
});
});

// Check final remount
expect(container.firstChild).not.toBe(nextEl);
const newEl = container.firstChild;
expect(newEl.textContent).toBe('hi test1');
}
});

it('resets state when switching between different component types', async () => {
if (__DEV__) {
await act(async () => {
await render(() => {
function Test() {
const [count, setCount] = React.useState(0);
return (
<div onClick={() => setCount(c => c + 1)}>count: {count}</div>
);
}
$RefreshReg$(Test, 'Test');
return Test;
});
});

expect(container.firstChild.textContent).toBe('count: 0');
await act(async () => {
container.firstChild.click();
});
expect(container.firstChild.textContent).toBe('count: 1');

await act(async () => {
await patch(() => {
function Test2() {
const [count, setCount] = React.useState(0);
return (
<div onClick={() => setCount(c => c + 1)}>count: {count}</div>
);
}
const Test = React.memo(Test2);
$RefreshReg$(Test2, 'Test2');
$RefreshReg$(Test, 'Test');
return Test;
});
});

expect(container.firstChild.textContent).toBe('count: 0');
await act(async () => {
container.firstChild.click();
});
expect(container.firstChild.textContent).toBe('count: 1');

await act(async () => {
await patch(() => {
const Test = React.forwardRef((props, ref) => {
const [count, setCount] = React.useState(0);
const handleClick = () => setCount(c => c + 1);

// Ensure ref is extensible
const divRef = React.useRef(null);
React.useEffect(() => {
if (ref) {
if (typeof ref === 'function') {
ref(divRef.current);
} else if (Object.isExtensible(ref)) {
ref.current = divRef.current;
}
}
}, [ref]);

return (
<div ref={divRef} onClick={handleClick}>
count: {count}
</div>
);
});
$RefreshReg$(Test, 'Test');
return Test;
});
});

expect(container.firstChild.textContent).toBe('count: 0');
await act(async () => {
container.firstChild.click();
});
expect(container.firstChild.textContent).toBe('count: 1');
}
});

it('can update simple memo function in isolation', async () => {
if (__DEV__) {
await render(() => {
Expand Down

0 comments on commit de1eaa2

Please sign in to comment.