From 070262b092c0925d255f907c5d187e35b4ba67b5 Mon Sep 17 00:00:00 2001
From: Dan Abramov <dan.abramov@me.com>
Date: Tue, 30 Mar 2021 14:55:12 +0100
Subject: [PATCH] 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.
---
 .../src/ReactFreshBabelPlugin.js              | 24 ++++++++++++++-----
 .../ReactFreshBabelPlugin-test.js.snap        |  4 ++--
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/packages/react-refresh/src/ReactFreshBabelPlugin.js b/packages/react-refresh/src/ReactFreshBabelPlugin.js
index 7df35f967b1d3..306bb3ff99191 100644
--- 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
index 4e2f8eec0518c..7c1276f817880 100644
--- 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);
 }
 `;