Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 516b76b

Browse files
gaearongaearon
authoredMar 30, 2021
[Fast Refresh] Support callthrough HOCs (#21104)
* [Fast Refresh] Support callthrough HOCs * Add a newly failing testing to demonstrate the flaw This shows why my initial approach doesn't make sense. * Attach signatures at every nesting level * Sign nested memo/forwardRef too * Add an IIFE test This is not a case that is important for Fast Refresh, but we shouldn't change the code semantics. This case shows the transform isn't quite correct. It's wrapping the call at the wrong place. * Find HOCs above more precisely This fixes a false positive that was causing an IIFE to be wrapped in the wrong place, which made the wrapping unsafe. * Be defensive against non-components being passed to setSignature * Fix lint
1 parent 0853aab commit 516b76b

File tree

5 files changed

+345
-49
lines changed

5 files changed

+345
-49
lines changed
 

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

Lines changed: 42 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -333,6 +333,38 @@ export default function(babel, opts = {}) {
333333
return args;
334334
}
335335

336+
function findHOCCallPathsAbove(path) {
337+
const calls = [];
338+
while (true) {
339+
if (!path) {
340+
return calls;
341+
}
342+
const parentPath = path.parentPath;
343+
if (!parentPath) {
344+
return calls;
345+
}
346+
if (
347+
// hoc(_c = function() { })
348+
parentPath.node.type === 'AssignmentExpression' &&
349+
path.node === parentPath.node.right
350+
) {
351+
// Ignore registrations.
352+
path = parentPath;
353+
continue;
354+
}
355+
if (
356+
// hoc1(hoc2(...))
357+
parentPath.node.type === 'CallExpression' &&
358+
path.node !== parentPath.node.callee
359+
) {
360+
calls.push(parentPath);
361+
path = parentPath;
362+
continue;
363+
}
364+
return calls; // Stop at other types.
365+
}
366+
}
367+
336368
const seenForRegistration = new WeakSet();
337369
const seenForSignature = new WeakSet();
338370
const seenForOutro = new WeakSet();
@@ -630,13 +662,16 @@ export default function(babel, opts = {}) {
630662
// Result: let Foo = () => {}; __signature(Foo, ...);
631663
} else {
632664
// let Foo = hoc(() => {})
633-
path.replaceWith(
634-
t.callExpression(
635-
sigCallID,
636-
createArgumentsForSignature(node, signature, path.scope),
637-
),
638-
);
639-
// Result: let Foo = hoc(__signature(() => {}, ...))
665+
const paths = [path, ...findHOCCallPathsAbove(path)];
666+
paths.forEach(p => {
667+
p.replaceWith(
668+
t.callExpression(
669+
sigCallID,
670+
createArgumentsForSignature(p.node, signature, p.scope),
671+
),
672+
);
673+
});
674+
// Result: let Foo = __signature(hoc(__signature(() => {}, ...)), ...)
640675
}
641676
},
642677
},

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

Lines changed: 50 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -355,12 +355,25 @@ export function setSignature(
355355
getCustomHooks?: () => Array<Function>,
356356
): void {
357357
if (__DEV__) {
358-
allSignaturesByType.set(type, {
359-
forceReset,
360-
ownKey: key,
361-
fullKey: null,
362-
getCustomHooks: getCustomHooks || (() => []),
363-
});
358+
if (!allSignaturesByType.has(type)) {
359+
allSignaturesByType.set(type, {
360+
forceReset,
361+
ownKey: key,
362+
fullKey: null,
363+
getCustomHooks: getCustomHooks || (() => []),
364+
});
365+
}
366+
// Visit inner types because we might not have signed them.
367+
if (typeof type === 'object' && type !== null) {
368+
switch (getProperty(type, '$$typeof')) {
369+
case REACT_FORWARD_REF_TYPE:
370+
setSignature(type.render, key, forceReset, getCustomHooks);
371+
break;
372+
case REACT_MEMO_TYPE:
373+
setSignature(type.type, key, forceReset, getCustomHooks);
374+
break;
375+
}
376+
}
364377
} else {
365378
throw new Error(
366379
'Unexpected call to React Refresh in a production environment.',
@@ -609,57 +622,58 @@ export function _getMountedRootCount() {
609622
// function Hello() {
610623
// const [foo, setFoo] = useState(0);
611624
// const value = useCustomHook();
612-
// _s(); /* Second call triggers collecting the custom Hook list.
625+
// _s(); /* Call without arguments triggers collecting the custom Hook list.
613626
// * This doesn't happen during the module evaluation because we
614627
// * don't want to change the module order with inline requires.
615628
// * Next calls are noops. */
616629
// return <h1>Hi</h1>;
617630
// }
618631
//
619-
// /* First call specifies the signature: */
632+
// /* Call with arguments attaches the signature to the type: */
620633
// _s(
621634
// Hello,
622635
// 'useState{[foo, setFoo]}(0)',
623636
// () => [useCustomHook], /* Lazy to avoid triggering inline requires */
624637
// );
625-
type SignatureStatus = 'needsSignature' | 'needsCustomHooks' | 'resolved';
626638
export function createSignatureFunctionForTransform() {
627639
if (__DEV__) {
628-
// We'll fill in the signature in two steps.
629-
// First, we'll know the signature itself. This happens outside the component.
630-
// Then, we'll know the references to custom Hooks. This happens inside the component.
631-
// After that, the returned function will be a fast path no-op.
632-
let status: SignatureStatus = 'needsSignature';
633640
let savedType;
634641
let hasCustomHooks;
642+
let didCollectHooks = false;
635643
return function<T>(
636644
type: T,
637645
key: string,
638646
forceReset?: boolean,
639647
getCustomHooks?: () => Array<Function>,
640-
): T {
641-
switch (status) {
642-
case 'needsSignature':
643-
if (type !== undefined) {
644-
// If we received an argument, this is the initial registration call.
645-
savedType = type;
646-
hasCustomHooks = typeof getCustomHooks === 'function';
647-
setSignature(type, key, forceReset, getCustomHooks);
648-
// The next call we expect is from inside a function, to fill in the custom Hooks.
649-
status = 'needsCustomHooks';
650-
}
651-
break;
652-
case 'needsCustomHooks':
653-
if (hasCustomHooks) {
654-
collectCustomHooksForSignature(savedType);
655-
}
656-
status = 'resolved';
657-
break;
658-
case 'resolved':
659-
// Do nothing. Fast path for all future renders.
660-
break;
648+
): T | void {
649+
if (typeof key === 'string') {
650+
// We're in the initial phase that associates signatures
651+
// with the functions. Note this may be called multiple times
652+
// in HOC chains like _s(hoc1(_s(hoc2(_s(actualFunction))))).
653+
if (!savedType) {
654+
// We're in the innermost call, so this is the actual type.
655+
savedType = type;
656+
hasCustomHooks = typeof getCustomHooks === 'function';
657+
}
658+
// Set the signature for all types (even wrappers!) in case
659+
// they have no signatures of their own. This is to prevent
660+
// problems like https://github.com/facebook/react/issues/20417.
661+
if (
662+
type != null &&
663+
(typeof type === 'function' || typeof type === 'object')
664+
) {
665+
setSignature(type, key, forceReset, getCustomHooks);
666+
}
667+
return type;
668+
} else {
669+
// We're in the _s() call without arguments, which means
670+
// this is the time to collect custom Hook signatures.
671+
// Only do this once. This path is hot and runs *inside* every render!
672+
if (!didCollectHooks && hasCustomHooks) {
673+
didCollectHooks = true;
674+
collectCustomHooksForSignature(savedType);
675+
}
661676
}
662-
return type;
663677
};
664678
} else {
665679
throw new Error(

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -524,4 +524,16 @@ describe('ReactFreshBabelPlugin', () => {
524524
'". If you want to override this check, pass {skipEnvCheck: true} as plugin options.',
525525
);
526526
});
527+
528+
it('does not get tripped by IIFEs', () => {
529+
expect(
530+
transform(`
531+
while (item) {
532+
(item => {
533+
useFoo();
534+
})(item);
535+
}
536+
`),
537+
).toMatchSnapshot();
538+
});
527539
});

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

Lines changed: 221 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -469,6 +469,227 @@ describe('ReactFreshIntegration', () => {
469469
}
470470
});
471471

472+
it('resets state when renaming a state variable inside a HOC with direct call', () => {
473+
if (__DEV__) {
474+
render(`
475+
const {useState} = React;
476+
const S = 1;
477+
478+
function hocWithDirectCall(Wrapped) {
479+
return function Generated() {
480+
return Wrapped();
481+
};
482+
}
483+
484+
export default hocWithDirectCall(() => {
485+
const [foo, setFoo] = useState(S);
486+
return <h1>A{foo}</h1>;
487+
});
488+
`);
489+
const el = container.firstChild;
490+
expect(el.textContent).toBe('A1');
491+
492+
patch(`
493+
const {useState} = React;
494+
const S = 2;
495+
496+
function hocWithDirectCall(Wrapped) {
497+
return function Generated() {
498+
return Wrapped();
499+
};
500+
}
501+
502+
export default hocWithDirectCall(() => {
503+
const [foo, setFoo] = useState(S);
504+
return <h1>B{foo}</h1>;
505+
});
506+
`);
507+
// Same state variable name, so state is preserved.
508+
expect(container.firstChild).toBe(el);
509+
expect(el.textContent).toBe('B1');
510+
511+
patch(`
512+
const {useState} = React;
513+
const S = 3;
514+
515+
function hocWithDirectCall(Wrapped) {
516+
return function Generated() {
517+
return Wrapped();
518+
};
519+
}
520+
521+
export default hocWithDirectCall(() => {
522+
const [bar, setBar] = useState(S);
523+
return <h1>C{bar}</h1>;
524+
});
525+
`);
526+
// Different state variable name, so state is reset.
527+
expect(container.firstChild).not.toBe(el);
528+
const newEl = container.firstChild;
529+
expect(newEl.textContent).toBe('C3');
530+
}
531+
});
532+
533+
it('does not crash when changing Hook order inside a HOC with direct call', () => {
534+
if (__DEV__) {
535+
render(`
536+
const {useEffect} = React;
537+
538+
function hocWithDirectCall(Wrapped) {
539+
return function Generated() {
540+
return Wrapped();
541+
};
542+
}
543+
544+
export default hocWithDirectCall(() => {
545+
useEffect(() => {}, []);
546+
return <h1>A</h1>;
547+
});
548+
`);
549+
const el = container.firstChild;
550+
expect(el.textContent).toBe('A');
551+
552+
patch(`
553+
const {useEffect} = React;
554+
555+
function hocWithDirectCall(Wrapped) {
556+
return function Generated() {
557+
return Wrapped();
558+
};
559+
}
560+
561+
export default hocWithDirectCall(() => {
562+
useEffect(() => {}, []);
563+
useEffect(() => {}, []);
564+
return <h1>B</h1>;
565+
});
566+
`);
567+
// Hook order changed, so we remount.
568+
expect(container.firstChild).not.toBe(el);
569+
const newEl = container.firstChild;
570+
expect(newEl.textContent).toBe('B');
571+
}
572+
});
573+
574+
it('does not crash when changing Hook order inside a memo-ed HOC with direct call', () => {
575+
if (__DEV__) {
576+
render(`
577+
const {useEffect, memo} = React;
578+
579+
function hocWithDirectCall(Wrapped) {
580+
return memo(function Generated() {
581+
return Wrapped();
582+
});
583+
}
584+
585+
export default hocWithDirectCall(() => {
586+
useEffect(() => {}, []);
587+
return <h1>A</h1>;
588+
});
589+
`);
590+
const el = container.firstChild;
591+
expect(el.textContent).toBe('A');
592+
593+
patch(`
594+
const {useEffect, memo} = React;
595+
596+
function hocWithDirectCall(Wrapped) {
597+
return memo(function Generated() {
598+
return Wrapped();
599+
});
600+
}
601+
602+
export default hocWithDirectCall(() => {
603+
useEffect(() => {}, []);
604+
useEffect(() => {}, []);
605+
return <h1>B</h1>;
606+
});
607+
`);
608+
// Hook order changed, so we remount.
609+
expect(container.firstChild).not.toBe(el);
610+
const newEl = container.firstChild;
611+
expect(newEl.textContent).toBe('B');
612+
}
613+
});
614+
615+
it('does not crash when changing Hook order inside a memo+forwardRef-ed HOC with direct call', () => {
616+
if (__DEV__) {
617+
render(`
618+
const {useEffect, memo, forwardRef} = React;
619+
620+
function hocWithDirectCall(Wrapped) {
621+
return memo(forwardRef(function Generated() {
622+
return Wrapped();
623+
}));
624+
}
625+
626+
export default hocWithDirectCall(() => {
627+
useEffect(() => {}, []);
628+
return <h1>A</h1>;
629+
});
630+
`);
631+
const el = container.firstChild;
632+
expect(el.textContent).toBe('A');
633+
634+
patch(`
635+
const {useEffect, memo, forwardRef} = React;
636+
637+
function hocWithDirectCall(Wrapped) {
638+
return memo(forwardRef(function Generated() {
639+
return Wrapped();
640+
}));
641+
}
642+
643+
export default hocWithDirectCall(() => {
644+
useEffect(() => {}, []);
645+
useEffect(() => {}, []);
646+
return <h1>B</h1>;
647+
});
648+
`);
649+
// Hook order changed, so we remount.
650+
expect(container.firstChild).not.toBe(el);
651+
const newEl = container.firstChild;
652+
expect(newEl.textContent).toBe('B');
653+
}
654+
});
655+
656+
it('does not crash when changing Hook order inside a HOC returning an object', () => {
657+
if (__DEV__) {
658+
render(`
659+
const {useEffect} = React;
660+
661+
function hocWithDirectCall(Wrapped) {
662+
return {Wrapped: Wrapped};
663+
}
664+
665+
export default hocWithDirectCall(() => {
666+
useEffect(() => {}, []);
667+
return <h1>A</h1>;
668+
}).Wrapped;
669+
`);
670+
const el = container.firstChild;
671+
expect(el.textContent).toBe('A');
672+
673+
patch(`
674+
const {useEffect} = React;
675+
676+
function hocWithDirectCall(Wrapped) {
677+
return {Wrapped: Wrapped};
678+
}
679+
680+
export default hocWithDirectCall(() => {
681+
useEffect(() => {}, []);
682+
useEffect(() => {}, []);
683+
return <h1>B</h1>;
684+
}).Wrapped;
685+
`);
686+
// Hook order changed, so we remount.
687+
expect(container.firstChild).not.toBe(el);
688+
const newEl = container.firstChild;
689+
expect(newEl.textContent).toBe('B');
690+
}
691+
});
692+
472693
it('resets effects while preserving state', () => {
473694
if (__DEV__) {
474695
render(`

‎packages/react-refresh/src/__tests__/__snapshots__/ReactFreshBabelPlugin-test.js.snap

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -37,11 +37,13 @@ const Bar = () => {
3737
_s4(Bar, "useContext{}");
3838
3939
_c2 = Bar;
40-
const Baz = memo(_c3 = _s5(() => {
40+
41+
const Baz = _s5(memo(_c3 = _s5(() => {
4142
_s5();
4243
4344
return useContext(X);
44-
}, "useContext{}"));
45+
}, "useContext{}")), "useContext{}");
46+
4547
_c4 = Baz;
4648
4749
const Qux = () => {
@@ -84,6 +86,18 @@ var _c;
8486
$RefreshReg$(_c, "App");
8587
`;
8688

89+
exports[`ReactFreshBabelPlugin does not get tripped by IIFEs 1`] = `
90+
while (item) {
91+
var _s = $RefreshSig$();
92+
93+
_s(item => {
94+
_s();
95+
96+
useFoo();
97+
}, "useFoo{}", true)(item);
98+
}
99+
`;
100+
87101
exports[`ReactFreshBabelPlugin generates signatures for function declarations calling hooks 1`] = `
88102
var _s = $RefreshSig$();
89103
@@ -108,21 +122,21 @@ exports[`ReactFreshBabelPlugin generates signatures for function expressions cal
108122
var _s = $RefreshSig$(),
109123
_s2 = $RefreshSig$();
110124
111-
export const A = React.memo(_c2 = React.forwardRef(_c = _s((props, ref) => {
125+
export const A = _s(React.memo(_c2 = _s(React.forwardRef(_c = _s((props, ref) => {
112126
_s();
113127
114128
const [foo, setFoo] = useState(0);
115129
React.useEffect(() => {});
116130
return <h1 ref={ref}>{foo}</h1>;
117-
}, "useState{[foo, setFoo](0)}\\nuseEffect{}")));
131+
}, "useState{[foo, setFoo](0)}\\nuseEffect{}")), "useState{[foo, setFoo](0)}\\nuseEffect{}")), "useState{[foo, setFoo](0)}\\nuseEffect{}");
118132
_c3 = A;
119-
export const B = React.memo(_c5 = React.forwardRef(_c4 = _s2(function (props, ref) {
133+
export const B = _s2(React.memo(_c5 = _s2(React.forwardRef(_c4 = _s2(function (props, ref) {
120134
_s2();
121135
122136
const [foo, setFoo] = useState(0);
123137
React.useEffect(() => {});
124138
return <h1 ref={ref}>{foo}</h1>;
125-
}, "useState{[foo, setFoo](0)}\\nuseEffect{}")));
139+
}, "useState{[foo, setFoo](0)}\\nuseEffect{}")), "useState{[foo, setFoo](0)}\\nuseEffect{}")), "useState{[foo, setFoo](0)}\\nuseEffect{}");
126140
_c6 = B;
127141
128142
function hoc() {

0 commit comments

Comments
 (0)
Please sign in to comment.