[Fresh] Fix a crash with implicit arrow return (#16687)

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(

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();
+  });
 });

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');