[compiler] extend AlignObjectMethodScopes to also merge FunctionExpression object properties

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/AlignObjectMethodScopes.ts

@@ -20,6 +20,11 @@ import DisjointSet from '../Utils/DisjointSet';
  * Align scopes of object method values to that of their enclosing object expressions.
  * To produce a well-formed JS program in Codegen, object methods and object expressions
  * must be in the same ReactiveBlock as object method definitions must be inlined.
+ *
+ * This also applies to FunctionExpression values (arrow functions and named function
+ * expressions) used as object properties: merging their scopes ensures all captured
+ * dependencies flow into the enclosing object's reactive scope, avoiding stale
+ * memoization when any of those deps change.
  */
 
 function findScopesToMerge(fn: HIRFunction): DisjointSet<ReactiveScope> {
@@ -28,7 +33,10 @@ function findScopesToMerge(fn: HIRFunction): DisjointSet<ReactiveScope> {
 
   for (const [_, block] of fn.body.blocks) {
     for (const {lvalue, value} of block.instructions) {
-      if (value.kind === 'ObjectMethod') {
+      if (
+        value.kind === 'ObjectMethod' ||
+        value.kind === 'FunctionExpression'
+      ) {
         objectMethodDecls.add(lvalue.identifier);
       } else if (value.kind === 'ObjectExpression') {
         for (const operand of eachInstructionValueOperand(value)) {
@@ -40,7 +48,7 @@ function findScopesToMerge(fn: HIRFunction): DisjointSet<ReactiveScope> {
               operandScope != null && lvalueScope != null,
               {
                 reason:
-                  'Internal error: Expected all ObjectExpressions and ObjectMethods to have non-null scope.',
+                  'Internal error: Expected all ObjectExpressions, ObjectMethods, and FunctionExpression object properties to have non-null scope.',
                 loc: GeneratedSource,
               },
             );

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-inconsistent-memoization-object-fn-expression.expect.md

@@ -0,0 +1,68 @@
+
+## Input
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees
+
+import {useState} from 'react';
+
+// Object with function expression properties - should be memoized
+// identically to method shorthand objects
+function useBug() {
+  const [setNodes, setEdges] = useState(null);
+  return {
+    test1: () => {
+      setNodes('test');
+    },
+    test2: function () {
+      setEdges('test');
+    },
+  };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useBug,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @validatePreserveExistingMemoizationGuarantees
+
+import { useState } from "react";
+
+// Object with function expression properties - should be memoized
+// identically to method shorthand objects
+function useBug() {
+  const $ = _c(2);
+  const [setNodes, setEdges] = useState(null);
+  let t0;
+  if ($[0] !== setNodes) {
+    t0 = {
+      test1: () => {
+        setNodes("test");
+      },
+      test2: function () {
+        setEdges("test");
+      },
+    };
+    $[0] = setNodes;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useBug,
+  params: [{}],
+};
+
+```
+
+### Eval output
+(kind: ok) {"test1":"[[ function params=0 ]]","test2":"[[ function params=0 ]]"}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-inconsistent-memoization-object-fn-expression.js

@@ -0,0 +1,22 @@
+// @validatePreserveExistingMemoizationGuarantees
+
+import {useState} from 'react';
+
+// Object with function expression properties - should be memoized
+// identically to method shorthand objects
+function useBug() {
+  const [setNodes, setEdges] = useState(null);
+  return {
+    test1: () => {
+      setNodes('test');
+    },
+    test2: function () {
+      setEdges('test');
+    },
+  };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useBug,
+  params: [{}],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-inconsistent-memoization-object-mixed-properties.expect.md

@@ -0,0 +1,73 @@
+
+## Input
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees
+
+import {useState} from 'react';
+
+// Object mixing method shorthand and function expression properties
+// should memoize entire object as a unit when deps change
+function useMixed() {
+  const [state, setState] = useState(0);
+  return {
+    // method shorthand
+    getValue() {
+      return state;
+    },
+    // arrow function property
+    getValueArrow: () => state,
+    // named function expression property
+    getValueFn: function getValueFn() {
+      return state;
+    },
+  };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useMixed,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @validatePreserveExistingMemoizationGuarantees
+
+import { useState } from "react";
+
+// Object mixing method shorthand and function expression properties
+// should memoize entire object as a unit when deps change
+function useMixed() {
+  const $ = _c(2);
+  const [state] = useState(0);
+  let t0;
+  if ($[0] !== state) {
+    t0 = {
+      getValue() {
+        return state;
+      },
+      getValueArrow: () => state,
+      getValueFn: function getValueFn() {
+        return state;
+      },
+    };
+    $[0] = state;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useMixed,
+  params: [{}],
+};
+
+```
+
+### Eval output
+(kind: ok) {"getValue":"[[ function params=0 ]]","getValueArrow":"[[ function params=0 ]]","getValueFn":"[[ function params=0 ]]"}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-inconsistent-memoization-object-mixed-properties.js

@@ -0,0 +1,26 @@
+// @validatePreserveExistingMemoizationGuarantees
+
+import {useState} from 'react';
+
+// Object mixing method shorthand and function expression properties
+// should memoize entire object as a unit when deps change
+function useMixed() {
+  const [state, setState] = useState(0);
+  return {
+    // method shorthand
+    getValue() {
+      return state;
+    },
+    // arrow function property
+    getValueArrow: () => state,
+    // named function expression property
+    getValueFn: function getValueFn() {
+      return state;
+    },
+  };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useMixed,
+  params: [{}],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-consecutive-scopes-objects.expect.md

@@ -44,7 +44,7 @@ import { Stringify } from "shared-runtime";
 // prevent scome scopes from merging, which concealed a bug with the merging logic.
 // By avoiding JSX we eliminate extraneous instructions and more accurately test the merging.
 function Component(props) {
-  const $ = _c(11);
+  const $ = _c(9);
   const [state, setState] = useState(0);
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
@@ -63,40 +63,36 @@ function Component(props) {
   }
   let t2;
   if ($[3] !== state) {
-    t2 = () => setState(state + 1);
+    let t3;
+    if ($[5] === Symbol.for("react.memo_cache_sentinel")) {
+      t3 = ["increment"];
+      $[5] = t3;
+    } else {
+      t3 = $[5];
+    }
+    t2 = {
+      component: "button",
+      props: {
+        "data-testid": "button",
+        onClick: () => setState(state + 1),
+        children: t3,
+      },
+    };
     $[3] = state;
     $[4] = t2;
   } else {
     t2 = $[4];
   }
   let t3;
-  if ($[5] === Symbol.for("react.memo_cache_sentinel")) {
-    t3 = ["increment"];
-    $[5] = t3;
-  } else {
-    t3 = $[5];
-  }
-  let t4;
-  if ($[6] !== t2) {
-    t4 = {
-      component: "button",
-      props: { "data-testid": "button", onClick: t2, children: t3 },
-    };
-    $[6] = t2;
-    $[7] = t4;
-  } else {
-    t4 = $[7];
-  }
-  let t5;
-  if ($[8] !== t1 || $[9] !== t4) {
-    t5 = [t0, t1, t4];
-    $[8] = t1;
-    $[9] = t4;
-    $[10] = t5;
+  if ($[6] !== t1 || $[7] !== t2) {
+    t3 = [t0, t1, t2];
+    $[6] = t1;
+    $[7] = t2;
+    $[8] = t3;
   } else {
-    t5 = $[10];
+    t3 = $[8];
   }
-  return t5;
+  return t3;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/object-shorthand-method-1.expect.md

@@ -27,33 +27,25 @@ export const FIXTURE_ENTRYPOINT = {
 import { c as _c } from "react/compiler-runtime";
 import { createHookWrapper } from "shared-runtime";
 function useHook(t0) {
-  const $ = _c(5);
+  const $ = _c(3);
   const { a, b } = t0;
   let t1;
-  if ($[0] !== a) {
-    t1 = function () {
-      return [a];
-    };
-    $[0] = a;
-    $[1] = t1;
-  } else {
-    t1 = $[1];
-  }
-  let t2;
-  if ($[2] !== b || $[3] !== t1) {
-    t2 = {
-      x: t1,
+  if ($[0] !== a || $[1] !== b) {
+    t1 = {
+      x: function () {
+        return [a];
+      },
       y() {
         return [b];
       },
     };
-    $[2] = b;
-    $[3] = t1;
-    $[4] = t2;
+    $[0] = a;
+    $[1] = b;
+    $[2] = t1;
   } else {
-    t2 = $[4];
+    t1 = $[2];
   }
-  return t2;
+  return t1;
 }
 
 export const FIXTURE_ENTRYPOINT = {