[Compiler] Align memoization of function expression properties with method shorthand #36151

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

@@ -17,33 +17,51 @@ import {eachInstructionValueOperand} from '../HIR/visitors';
 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.
+ * Align scopes of object method values and function expression property values
+ * to that of their enclosing object expressions. This ensures that objects are
+ * memoized as a single unit regardless of whether properties use method shorthand
+ * syntax or function expression syntax.
+ *
+ * For ObjectMethod values, this is also required for correctness: object methods
+ * and object expressions must be in the same ReactiveBlock as object method
+ * definitions must be inlined.
  */
 
 function findScopesToMerge(fn: HIRFunction): DisjointSet<ReactiveScope> {
   const objectMethodDecls: Set<Identifier> = new Set();
+  const fnExprDecls: Set<Identifier> = new Set();
   const mergeScopesBuilder = new DisjointSet<ReactiveScope>();
 
   for (const [_, block] of fn.body.blocks) {
     for (const {lvalue, value} of block.instructions) {
       if (value.kind === 'ObjectMethod') {
         objectMethodDecls.add(lvalue.identifier);
+      } else if (value.kind === 'FunctionExpression') {
+        fnExprDecls.add(lvalue.identifier);
       } else if (value.kind === 'ObjectExpression') {
         for (const operand of eachInstructionValueOperand(value)) {
-          if (objectMethodDecls.has(operand.identifier)) {
+          if (
+            objectMethodDecls.has(operand.identifier) ||
+            fnExprDecls.has(operand.identifier)
+          ) {
             const operandScope = operand.identifier.scope;
             const lvalueScope = lvalue.identifier.scope;
 
-            CompilerError.invariant(
-              operandScope != null && lvalueScope != null,
-              {
-                reason:
-                  'Internal error: Expected all ObjectExpressions and ObjectMethods to have non-null scope.',
-                loc: GeneratedSource,
-              },
-            );
+            if (operandScope == null || lvalueScope == null) {
+              /*
+               * FunctionExpressions used as object properties may not always
+               * have scopes (e.g. if they were pruned). ObjectMethods must
+               * always have scopes.
+               */
+              if (objectMethodDecls.has(operand.identifier)) {
+                CompilerError.invariant(false, {
+                  reason:
+                    'Internal error: Expected all ObjectExpressions and ObjectMethods to have non-null scope.',
+                  loc: GeneratedSource,
+                });
+              }
+              continue;
+            }
             mergeScopesBuilder.union([operandScope, lvalueScope]);
           }
         }

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

@@ -0,0 +1,57 @@
+
+## Input
+
+```javascript
+function Component({a, b}) {
+  return {
+    test1: () => {
+      console.log(a);
+    },
+    test2: function () {
+      console.log(b);
+    },
+  };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 1, b: 2}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+function Component(t0) {
+  const $ = _c(3);
+  const { a, b } = t0;
+  let t1;
+  if ($[0] !== a || $[1] !== b) {
+    t1 = {
+      test1: () => {
+        console.log(a);
+      },
+      test2: function () {
+        console.log(b);
+      },
+    };
+    $[0] = a;
+    $[1] = b;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: 1, b: 2 }],
+};
+
+```
+
+### 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,15 @@
+function Component({a, b}) {
+  return {
+    test1: () => {
+      console.log(a);
+    },
+    test2: function () {
+      console.log(b);
+    },
+  };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 1, b: 2}],
+};

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

@@ -0,0 +1,57 @@
+
+## Input
+
+```javascript
+function Component({a, b}) {
+  return {
+    test1() {
+      console.log(a);
+    },
+    test2() {
+      console.log(b);
+    },
+  };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 1, b: 2}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+function Component(t0) {
+  const $ = _c(3);
+  const { a, b } = t0;
+  let t1;
+  if ($[0] !== a || $[1] !== b) {
+    t1 = {
+      test1() {
+        console.log(a);
+      },
+      test2() {
+        console.log(b);
+      },
+    };
+    $[0] = a;
+    $[1] = b;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: 1, b: 2 }],
+};
+
+```
+
+### 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-method-shorthand.js

@@ -0,0 +1,15 @@
+function Component({a, b}) {
+  return {
+    test1() {
+      console.log(a);
+    },
+    test2() {
+      console.log(b);
+    },
+  };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 1, b: 2}],
+};

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

@@ -0,0 +1,64 @@
+
+## Input
+
+```javascript
+function Component({a, b, c}) {
+  return {
+    test1() {
+      console.log(a);
+    },
+    test2: () => {
+      console.log(b);
+    },
+    test3: function () {
+      console.log(c);
+    },
+  };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 1, b: 2, c: 3}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+function Component(t0) {
+  const $ = _c(4);
+  const { a, b, c } = t0;
+  let t1;
+  if ($[0] !== a || $[1] !== b || $[2] !== c) {
+    t1 = {
+      test1() {
+        console.log(a);
+      },
+      test2: () => {
+        console.log(b);
+      },
+      test3: function () {
+        console.log(c);
+      },
+    };
+    $[0] = a;
+    $[1] = b;
+    $[2] = c;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: 1, b: 2, c: 3 }],
+};
+
+```
+
+### Eval output
+(kind: ok) {"test1":"[[ function params=0 ]]","test2":"[[ function params=0 ]]","test3":"[[ function params=0 ]]"}

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

@@ -0,0 +1,18 @@
+function Component({a, b, c}) {
+  return {
+    test1() {
+      console.log(a);
+    },
+    test2: () => {
+      console.log(b);
+    },
+    test3: function () {
+      console.log(c);
+    },
+  };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 1, b: 2, c: 3}],
+};

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 = {