fix(react-compiler): optimize components declared with arrow function and implicit return and compilationMode: 'infer' (#31792)

fixes https://github.com/facebook/react/issues/31601
https://github.com/facebook/react/issues/31639 cc @josephsavona

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts

@@ -1008,31 +1008,39 @@ function callsHooksOrCreatesJsx(
   return invokesHooks || createsJsx;
 }
 
+function isNonNode(node?: t.Expression | null): boolean {
+  if (!node) {
+    return true;
+  }
+  switch (node.type) {
+    case 'ObjectExpression':
+    case 'ArrowFunctionExpression':
+    case 'FunctionExpression':
+    case 'BigIntLiteral':
+    case 'ClassExpression':
+    case 'NewExpression': // technically `new Array()` is legit, but unlikely
+      return true;
+  }
+  return false;
+}
+
 function returnsNonNode(
   node: NodePath<
     t.FunctionDeclaration | t.ArrowFunctionExpression | t.FunctionExpression
   >,
 ): boolean {
-  let hasReturn = false;
   let returnsNonNode = false;
+  if (
+    // node.traverse#ArrowFunctionExpression isn't called for the root node
+    node.type === 'ArrowFunctionExpression' &&
+    node.node.body.type !== 'BlockStatement'
+  ) {
+    returnsNonNode = isNonNode(node.node.body);
+  }
 
   node.traverse({
     ReturnStatement(ret) {
-      hasReturn = true;
-      const argument = ret.node.argument;
-      if (argument == null) {
-        returnsNonNode = true;
-      } else {
-        switch (argument.type) {
-          case 'ObjectExpression':
-          case 'ArrowFunctionExpression':
-          case 'FunctionExpression':
-          case 'BigIntLiteral':
-          case 'ClassExpression':
-          case 'NewExpression': // technically `new Array()` is legit, but unlikely
-            returnsNonNode = true;
-        }
-      }
+      returnsNonNode = isNonNode(ret.node.argument);
     },
     // Skip traversing all nested functions and their return statements
     ArrowFunctionExpression: skipNestedFunctions(node),
@@ -1041,7 +1049,7 @@ function returnsNonNode(
     ObjectMethod: node => node.skip(),
   });
 
-  return !hasReturn || returnsNonNode;
+  return returnsNonNode;
 }
 
 /*

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/arrow-function-with-implicit-return.expect.md

@@ -0,0 +1,39 @@
+
+## Input
+
+```javascript
+// @compilationMode(infer)
+const Test = () => <div />;
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Test,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @compilationMode(infer)
+const Test = () => {
+  const $ = _c(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <div />;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+};
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Test,
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div></div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/arrow-function-with-implicit-return.js

@@ -0,0 +1,7 @@
+// @compilationMode(infer)
+const Test = () => <div />;
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Test,
+  params: [{}],
+};