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.
facebook · Mar 30, 2021 · 070262b · 070262b
1 parent d2a2045
commit 070262b
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 8 deletions.
diff --git a/packages/react-refresh/src/ReactFreshBabelPlugin.js b/packages/react-refresh/src/ReactFreshBabelPlugin.js
@@ -339,14 +339,26 @@ export default function(babel, opts = {}) {
       if (!path) {
         return calls;
       }
-      if (path.node.type === 'AssignmentExpression') {
+      let parentPath = path.parentPath;
+      if (!parentPath) {
+        return calls;
+      }
+      if (
+        // hoc(_c = function() { })
+        parentPath.node.type === 'AssignmentExpression' &&
+        path.node === parentPath.node.right
+      ) {
         // Ignore registrations.
-        path = path.parentPath;
+        path = parentPath;
         continue;
       }
-      if (path.node.type === 'CallExpression') {
-        calls.push(path);
-        path = path.parentPath;
+      if (
+        // hoc1(hoc2(...))
+        parentPath.node.type === 'CallExpression' &&
+        path.node !== parentPath.node.callee
+      ) {
+        calls.push(parentPath);
+        path = parentPath;
         continue;
       }
       return calls; // Stop at other types.
@@ -650,7 +662,7 @@ export default function(babel, opts = {}) {
             // Result: let Foo = () => {}; __signature(Foo, ...);
           } else {
             // let Foo = hoc(() => {})
-            const paths = [path, ...findHOCCallPathsAbove(path.parentPath)];
+            const paths = [path, ...findHOCCallPathsAbove(path)];
             paths.forEach(p => {
               p.replaceWith(
                 t.callExpression(

diff --git a/packages/react-refresh/src/__tests__/__snapshots__/ReactFreshBabelPlugin-test.js.snap b/packages/react-refresh/src/__tests__/__snapshots__/ReactFreshBabelPlugin-test.js.snap
@@ -90,11 +90,11 @@ exports[`ReactFreshBabelPlugin does not get tripped by IIFEs 1`] = `
 while (item) {
   var _s = $RefreshSig$();
 
-  _s(_s(item => {
+  _s(item => {
     _s();
 
     useFoo();
-  }, "useFoo{}", true)(item), "useFoo{}", true);
+  }, "useFoo{}", true)(item);
 }
 `;