From af8532f25121a6b4060e3edd7186f0a58812d771 Mon Sep 17 00:00:00 2001
From: mofeiZ <34200447+mofeiZ@users.noreply.github.com>
Date: Mon, 13 Jan 2025 12:18:59 -0500
Subject: [PATCH] [compiler][ez] Patch compilationMode:infer object method edge
 case (#32055)

Fix for  https://github.com/facebook/react/issues/31180
---
 .../src/Entrypoint/Program.ts                 |  2 +
 .../infer-nested-object-method.expect.md      | 63 +++++++++++++++++++
 .../compiler/infer-nested-object-method.jsx   | 19 ++++++
 3 files changed, 84 insertions(+)
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-nested-object-method.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-nested-object-method.jsx

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts
index bb0d662c4f67e..6ab9ee79c7412 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts
@@ -992,9 +992,11 @@ function returnsNonNode(
         }
       }
     },
+    // Skip traversing all nested functions and their return statements
     ArrowFunctionExpression: skipNestedFunctions(node),
     FunctionExpression: skipNestedFunctions(node),
     FunctionDeclaration: skipNestedFunctions(node),
+    ObjectMethod: node => node.skip(),
   });
 
   return !hasReturn || returnsNonNode;
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-nested-object-method.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-nested-object-method.expect.md
new file mode 100644
index 0000000000000..4c89f59919f0f
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-nested-object-method.expect.md
@@ -0,0 +1,63 @@
+
+## Input
+
+```javascript
+// @compilationMode(infer)
+
+import {Stringify} from 'shared-runtime';
+
+function Test() {
+  const context = {
+    testFn() {
+      // if it is an arrow function its work
+      return () => 'test'; // it will break compile if returns an arrow fn
+    },
+  };
+
+  return <Stringify value={context} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Test,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @compilationMode(infer)
+
+import { Stringify } from "shared-runtime";
+
+function Test() {
+  const $ = _c(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    const context = {
+      testFn() {
+        return _temp;
+      },
+    };
+
+    t0 = <Stringify value={context} shouldInvokeFns={true} />;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+function _temp() {
+  return "test";
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Test,
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"value":{"testFn":{"kind":"Function","result":{"kind":"Function","result":"test"}}},"shouldInvokeFns":true}</div>
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-nested-object-method.jsx b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-nested-object-method.jsx
new file mode 100644
index 0000000000000..c8e396b057f8f
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-nested-object-method.jsx
@@ -0,0 +1,19 @@
+// @compilationMode(infer)
+
+import {Stringify} from 'shared-runtime';
+
+function Test() {
+  const context = {
+    testFn() {
+      // if it is an arrow function its work
+      return () => 'test'; // it will break compile if returns an arrow fn
+    },
+  };
+
+  return <Stringify value={context} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Test,
+  params: [{}],
+};