From 9044bb0fa3c60ffe723034486b02168bdf63f786 Mon Sep 17 00:00:00 2001
From: Dan Abramov <dan.abramov@gmail.com>
Date: Fri, 6 Sep 2019 19:58:07 +0100
Subject: [PATCH] [Fresh] Fix a crash with implicit arrow return (#16687)

---
 .../src/ReactFreshBabelPlugin.js              |  5 ++
 .../__tests__/ReactFreshBabelPlugin-test.js   | 13 ++++
 .../ReactFreshBabelPlugin-test.js.snap        | 63 +++++++++++++++++++
 3 files changed, 81 insertions(+)

diff --git a/packages/react-refresh/src/ReactFreshBabelPlugin.js b/packages/react-refresh/src/ReactFreshBabelPlugin.js
index 3b4c1e056b45b..b8474af1a0be6 100644
--- a/packages/react-refresh/src/ReactFreshBabelPlugin.js
+++ b/packages/react-refresh/src/ReactFreshBabelPlugin.js
@@ -571,6 +571,11 @@ export default function(babel) {
 
           // The signature call is split in two parts. One part is called inside the function.
           // This is used to signal when first render happens.
+          if (path.node.body.type !== 'BlockStatement') {
+            path.node.body = t.blockStatement([
+              t.returnStatement(path.node.body),
+            ]);
+          }
           path
             .get('body')
             .unshiftContainer(
diff --git a/packages/react-refresh/src/__tests__/ReactFreshBabelPlugin-test.js b/packages/react-refresh/src/__tests__/ReactFreshBabelPlugin-test.js
index fa5e152fbdc7a..2f47f1491b021 100644
--- a/packages/react-refresh/src/__tests__/ReactFreshBabelPlugin-test.js
+++ b/packages/react-refresh/src/__tests__/ReactFreshBabelPlugin-test.js
@@ -460,4 +460,17 @@ describe('ReactFreshBabelPlugin', () => {
     `),
     ).toMatchSnapshot();
   });
+
+  it('can handle implicit arrow returns', () => {
+    expect(
+      transform(`
+        export default () => useContext(X);
+        export const Foo = () => useContext(X);
+        module.exports = () => useContext(X);
+        const Bar = () => useContext(X);
+        const Baz = memo(() => useContext(X));
+        const Qux = () => (0, useContext(X));
+      `),
+    ).toMatchSnapshot();
+  });
 });
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 8035509814521..5fae0ba2b9ad9 100644
--- a/packages/react-refresh/src/__tests__/__snapshots__/ReactFreshBabelPlugin-test.js.snap
+++ b/packages/react-refresh/src/__tests__/__snapshots__/ReactFreshBabelPlugin-test.js.snap
@@ -1,5 +1,68 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
+exports[`ReactFreshBabelPlugin can handle implicit arrow returns 1`] = `
+var _s = $RefreshSig$(),
+    _s2 = $RefreshSig$(),
+    _s3 = $RefreshSig$(),
+    _s4 = $RefreshSig$(),
+    _s5 = $RefreshSig$(),
+    _s6 = $RefreshSig$();
+
+export default _s(() => {
+  _s();
+
+  return useContext(X);
+}, "useContext{}");
+export const Foo = () => {
+  _s2();
+
+  return useContext(X);
+};
+
+_s2(Foo, "useContext{}");
+
+_c = Foo;
+module.exports = _s3(() => {
+  _s3();
+
+  return useContext(X);
+}, "useContext{}");
+
+const Bar = () => {
+  _s4();
+
+  return useContext(X);
+};
+
+_s4(Bar, "useContext{}");
+
+_c2 = Bar;
+
+const Baz = _c4 = memo(_c3 = _s5(() => {
+  _s5();
+
+  return useContext(X);
+}, "useContext{}"));
+
+const Qux = () => {
+  _s6();
+
+  return 0, useContext(X);
+};
+
+_s6(Qux, "useContext{}");
+
+_c5 = Qux;
+
+var _c, _c2, _c3, _c4, _c5;
+
+$RefreshReg$(_c, "Foo");
+$RefreshReg$(_c2, "Bar");
+$RefreshReg$(_c3, "Baz$memo");
+$RefreshReg$(_c4, "Baz");
+$RefreshReg$(_c5, "Qux");
+`;
+
 exports[`ReactFreshBabelPlugin does not consider require-like methods to be HOCs 1`] = `
 const A = require('A');