@samitouri / QOS-React / commits / 516b76b9ae

[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

Dan Abramov committed Mar 30, 2021 at 16:08 UTC 516b76b9ae177ecc38cc64ba33ddfc7cfe5f5a99
5 files changed +345 -49
packages/react-refresh/src/ReactFreshBabelPlugin.js
+42 -7
@@ -333,6 +333,38 @@ export default function(babel, opts = {}) {
333 return args;
334 }
335
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 +
368 const seenForRegistration = new WeakSet();
369 const seenForSignature = new WeakSet();
370 const seenForOutro = new WeakSet();
@@ -630,13 +662,16 @@ export default function(babel, opts = {}) {
662 // Result: let Foo = () => {}; __signature(Foo, ...);
663 } else {
664 // 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(() => {}, ...)), ...)
675 }
676 },
677 },
packages/react-refresh/src/ReactFreshRuntime.js
+50 -36
@@ -355,12 +355,25 @@ export function setSignature(
355 getCustomHooks?: () => Array<Function>,
356 ): void {
357 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 + }
377 } else {
378 throw new Error(
379 'Unexpected call to React Refresh in a production environment.',
@@ -609,57 +622,58 @@ export function _getMountedRootCount() {
622 // function Hello() {
623 // const [foo, setFoo] = useState(0);
624 // const value = useCustomHook();
612 -// _s(); /* Second call triggers collecting the custom Hook list.
625 +// _s(); /* Call without arguments triggers collecting the custom Hook list.
626 // * This doesn't happen during the module evaluation because we
627 // * don't want to change the module order with inline requires.
628 // * Next calls are noops. */
629 // return <h1>Hi</h1>;
630 // }
631 //
619 -// /* First call specifies the signature: */
632 +// /* Call with arguments attaches the signature to the type: */
633 // _s(
634 // Hello,
635 // 'useState{[foo, setFoo]}(0)',
636 // () => [useCustomHook], /* Lazy to avoid triggering inline requires */
637 // );
625 -type SignatureStatus = 'needsSignature' | 'needsCustomHooks' | 'resolved';
638 export function createSignatureFunctionForTransform() {
639 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';
640 let savedType;
641 let hasCustomHooks;
642 + let didCollectHooks = false;
643 return function<T>(
644 type: T,
645 key: string,
646 forceReset?: boolean,
647 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 + }
676 }
662 - return type;
677 };
678 } else {
679 throw new Error(
packages/react-refresh/src/__tests__/ReactFreshBabelPlugin-test.js
+12
@@ -524,4 +524,16 @@ describe('ReactFreshBabelPlugin', () => {
524 '". If you want to override this check, pass {skipEnvCheck: true} as plugin options.',
525 );
526 });
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 + });
539 });
packages/react-refresh/src/__tests__/ReactFreshIntegration-test.js
+221
@@ -469,6 +469,227 @@ describe('ReactFreshIntegration', () => {
469 }
470 });
471
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 +
693 it('resets effects while preserving state', () => {
694 if (__DEV__) {
695 render(`
packages/react-refresh/src/__tests__/__snapshots__/ReactFreshBabelPlugin-test.js.snap
+20 -6
@@ -37,11 +37,13 @@ const Bar = () => {
37 _s4(Bar, "useContext{}");
38
39 _c2 = Bar;
40 -const Baz = memo(_c3 = _s5(() => {
40 +
41 +const Baz = _s5(memo(_c3 = _s5(() => {
42 _s5();
43
44 return useContext(X);
44 -}, "useContext{}"));
45 +}, "useContext{}")), "useContext{}");
46 +
47 _c4 = Baz;
48
49 const Qux = () => {
@@ -84,6 +86,18 @@ var _c;
86 $RefreshReg$(_c, "App");
87 `;
88
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 +
101 exports[`ReactFreshBabelPlugin generates signatures for function declarations calling hooks 1`] = `
102 var _s = $RefreshSig$();
103
@@ -108,21 +122,21 @@ exports[`ReactFreshBabelPlugin generates signatures for function expressions cal
122 var _s = $RefreshSig$(),
123 _s2 = $RefreshSig$();
124
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) => {
126 _s();
127
128 const [foo, setFoo] = useState(0);
129 React.useEffect(() => {});
130 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{}");
132 _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) {
134 _s2();
135
136 const [foo, setFoo] = useState(0);
137 React.useEffect(() => {});
138 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{}");
140 _c6 = B;
141
142 function hoc() {