From d0c31e9dd2c43b152b67715beba5181e4d620bf0 Mon Sep 17 00:00:00 2001
From: Mofei Zhang <feifei0@meta.com>
Date: Thu, 24 Oct 2024 12:23:22 -0700
Subject: [PATCH] [compiler][ez] Patch hoistability for ObjectMethods

Extends #31066 to ObjectMethods (somehow missed this before).

'
---
 .../src/HIR/CollectHoistablePropertyLoads.ts  |  8 +-
 .../infer-objectmethod-cond-access.expect.md  | 82 +++++++++++++++++++
 .../infer-objectmethod-cond-access.js         | 26 ++++++
 3 files changed, 114 insertions(+), 2 deletions(-)
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/reduce-reactive-deps/infer-objectmethod-cond-access.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/reduce-reactive-deps/infer-objectmethod-cond-access.js

diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/CollectHoistablePropertyLoads.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/CollectHoistablePropertyLoads.ts
index 80593d6275868..d2e7220159108 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/CollectHoistablePropertyLoads.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/CollectHoistablePropertyLoads.ts
@@ -348,7 +348,8 @@ function collectNonNullsInBlocks(
         assumedNonNullObjects.add(maybeNonNull);
       }
       if (
-        instr.value.kind === 'FunctionExpression' &&
+        (instr.value.kind === 'FunctionExpression' ||
+          instr.value.kind === 'ObjectMethod') &&
         !fn.env.config.enableTreatFunctionDepsAsConditional
       ) {
         const innerFn = instr.value.loweredFunc;
@@ -591,7 +592,10 @@ function collectFunctionExpressionFakeLoads(
 
   for (const [_, block] of fn.body.blocks) {
     for (const {lvalue, value} of block.instructions) {
-      if (value.kind === 'FunctionExpression') {
+      if (
+        value.kind === 'FunctionExpression' ||
+        value.kind === 'ObjectMethod'
+      ) {
         for (const reference of value.loweredFunc.dependencies) {
           let curr: IdentifierId | undefined = reference.identifier.id;
           while (curr != null) {
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/reduce-reactive-deps/infer-objectmethod-cond-access.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/reduce-reactive-deps/infer-objectmethod-cond-access.expect.md
new file mode 100644
index 0000000000000..2c7bf6bbe599e
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/reduce-reactive-deps/infer-objectmethod-cond-access.expect.md
@@ -0,0 +1,82 @@
+
+## Input
+
+```javascript
+// @enablePropagateDepsInHIR
+import {Stringify} from 'shared-runtime';
+
+function Foo({a, shouldReadA}) {
+  return (
+    <Stringify
+      objectMethod={{
+        method() {
+          if (shouldReadA) return a.b.c;
+          return null;
+        },
+      }}
+      shouldInvokeFns={true}
+    />
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{a: null, shouldReadA: true}],
+  sequentialRenders: [
+    {a: null, shouldReadA: true},
+    {a: null, shouldReadA: false},
+    {a: {b: {c: 4}}, shouldReadA: true},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enablePropagateDepsInHIR
+import { Stringify } from "shared-runtime";
+
+function Foo(t0) {
+  const $ = _c(3);
+  const { a, shouldReadA } = t0;
+  let t1;
+  if ($[0] !== shouldReadA || $[1] !== a) {
+    t1 = (
+      <Stringify
+        objectMethod={{
+          method() {
+            if (shouldReadA) {
+              return a.b.c;
+            }
+            return null;
+          },
+        }}
+        shouldInvokeFns={true}
+      />
+    );
+    $[0] = shouldReadA;
+    $[1] = a;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{ a: null, shouldReadA: true }],
+  sequentialRenders: [
+    { a: null, shouldReadA: true },
+    { a: null, shouldReadA: false },
+    { a: { b: { c: 4 } }, shouldReadA: true },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) [[ (exception in render) TypeError: Cannot read properties of null (reading 'b') ]]
+<div>{"objectMethod":{"method":{"kind":"Function","result":null}},"shouldInvokeFns":true}</div>
+<div>{"objectMethod":{"method":{"kind":"Function","result":4}},"shouldInvokeFns":true}</div>
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/reduce-reactive-deps/infer-objectmethod-cond-access.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/reduce-reactive-deps/infer-objectmethod-cond-access.js
new file mode 100644
index 0000000000000..2c8488bb29216
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/reduce-reactive-deps/infer-objectmethod-cond-access.js
@@ -0,0 +1,26 @@
+// @enablePropagateDepsInHIR
+import {Stringify} from 'shared-runtime';
+
+function Foo({a, shouldReadA}) {
+  return (
+    <Stringify
+      objectMethod={{
+        method() {
+          if (shouldReadA) return a.b.c;
+          return null;
+        },
+      }}
+      shouldInvokeFns={true}
+    />
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{a: null, shouldReadA: true}],
+  sequentialRenders: [
+    {a: null, shouldReadA: true},
+    {a: null, shouldReadA: false},
+    {a: {b: {c: 4}}, shouldReadA: true},
+  ],
+};