Optimize method calls w props receiver (#31775)

Redo of #31771 without ghstack

Author: josephsavona

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

@@ -104,6 +104,7 @@ import {validateNoSetStateInPassiveEffects} from '../Validation/ValidateNoSetSta
 import {validateNoJSXInTryStatement} from '../Validation/ValidateNoJSXInTryStatement';
 import {propagateScopeDependenciesHIR} from '../HIR/PropagateScopeDependenciesHIR';
 import {outlineJSX} from '../Optimization/OutlineJsx';
+import {optimizePropsMethodCalls} from '../Optimization/OptimizePropsMethodCalls';
 
 export type CompilerPipelineValue =
   | {kind: 'ast'; name: string; value: CodegenFunction}
@@ -209,6 +210,9 @@ function* runWithEnvironment(
     lowerContextAccess(hir, env.config.lowerContextAccess);
   }
 
+  optimizePropsMethodCalls(hir);
+  yield log({kind: 'hir', name: 'OptimizePropsMethodCalls', value: hir});
+
   analyseFunctions(hir);
   yield log({kind: 'hir', name: 'AnalyseFunctions', value: hir});
 

compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts

@@ -1644,6 +1644,10 @@ export function isArrayType(id: Identifier): boolean {
   return id.type.kind === 'Object' && id.type.shapeId === 'BuiltInArray';
 }
 
+export function isPropsType(id: Identifier): boolean {
+  return id.type.kind === 'Object' && id.type.shapeId === 'BuiltInProps';
+}
+
 export function isRefValueType(id: Identifier): boolean {
   return id.type.kind === 'Object' && id.type.shapeId === 'BuiltInRefValue';
 }

compiler/packages/babel-plugin-react-compiler/src/Optimization/OptimizePropsMethodCalls.ts

@@ -0,0 +1,52 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import {HIRFunction, isPropsType} from '../HIR';
+
+/**
+ * Converts method calls into regular calls where the receiver is the props object:
+ *
+ * Example:
+ *
+ * ```
+ * // INPUT
+ * props.foo();
+ *
+ * // OUTPUT
+ * const t0 = props.foo;
+ * t0();
+ * ```
+ *
+ * Counter example:
+ *
+ * Here the receiver is `props.foo`, not the props object, so we don't rewrite it:
+ *
+ * // INPUT
+ * props.foo.bar();
+ *
+ * // OUTPUT
+ * props.foo.bar();
+ * ```
+ */
+export function optimizePropsMethodCalls(fn: HIRFunction): void {
+  for (const [, block] of fn.body.blocks) {
+    for (let i = 0; i < block.instructions.length; i++) {
+      const instr = block.instructions[i]!;
+      if (
+        instr.value.kind === 'MethodCall' &&
+        isPropsType(instr.value.receiver.identifier)
+      ) {
+        instr.value = {
+          kind: 'CallExpression',
+          callee: instr.value.property,
+          args: instr.value.args,
+          loc: instr.value.loc,
+        };
+      }
+    }
+  }
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/jsx-outlining-child-stored-in-id.expect.md

@@ -3,7 +3,7 @@
 
 ```javascript
 // @enableJsxOutlining
-function Component(arr) {
+function Component({arr}) {
   const x = useX();
   return arr.map(i => {
     <>
@@ -49,12 +49,13 @@ export const FIXTURE_ENTRYPOINT = {
 
 ```javascript
 import { c as _c } from "react/compiler-runtime"; // @enableJsxOutlining
-function Component(arr) {
+function Component(t0) {
   const $ = _c(3);
+  const { arr } = t0;
   const x = useX();
-  let t0;
+  let t1;
   if ($[0] !== arr || $[1] !== x) {
-    t0 = arr.map((i) => {
+    t1 = arr.map((i) => {
       arr.map((i_0, id) => {
         const T0 = _temp;
         const child = <T0 i={i_0} x={x} />;
@@ -65,11 +66,11 @@ function Component(arr) {
     });
     $[0] = arr;
     $[1] = x;
-    $[2] = t0;
+    $[2] = t1;
   } else {
-    t0 = $[2];
+    t1 = $[2];
   }
-  return t0;
+  return t1;
 }
 function _temp(t0) {
   const $ = _c(5);
@@ -140,4 +141,4 @@ export const FIXTURE_ENTRYPOINT = {
 ```
       
 ### Eval output
-(kind: exception) arr.map is not a function
+(kind: ok) [null,null]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/jsx-outlining-child-stored-in-id.js

@@ -1,5 +1,5 @@
 // @enableJsxOutlining
-function Component(arr) {
+function Component({arr}) {
   const x = useX();
   return arr.map(i => {
     <>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/props-method-dependency.expect.md

@@ -0,0 +1,78 @@
+
+## Input
+
+```javascript
+// @compilationMode(infer)
+import {useMemo} from 'react';
+import {ValidateMemoization} from 'shared-runtime';
+
+function Component(props) {
+  const x = useMemo(() => props.x(), [props.x]);
+  return <ValidateMemoization inputs={[props.x]} output={x} />;
+}
+
+const f = () => ['React'];
+const g = () => ['Compiler'];
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{x: () => ['React']}],
+  sequentialRenders: [{x: f}, {x: g}, {x: g}, {x: f}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @compilationMode(infer)
+import { useMemo } from "react";
+import { ValidateMemoization } from "shared-runtime";
+
+function Component(props) {
+  const $ = _c(7);
+  let t0;
+  let t1;
+  if ($[0] !== props.x) {
+    t1 = props.x();
+    $[0] = props.x;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  t0 = t1;
+  const x = t0;
+  let t2;
+  if ($[2] !== props.x) {
+    t2 = [props.x];
+    $[2] = props.x;
+    $[3] = t2;
+  } else {
+    t2 = $[3];
+  }
+  let t3;
+  if ($[4] !== t2 || $[5] !== x) {
+    t3 = <ValidateMemoization inputs={t2} output={x} />;
+    $[4] = t2;
+    $[5] = x;
+    $[6] = t3;
+  } else {
+    t3 = $[6];
+  }
+  return t3;
+}
+
+const f = () => ["React"];
+const g = () => ["Compiler"];
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: () => ["React"] }],
+  sequentialRenders: [{ x: f }, { x: g }, { x: g }, { x: f }],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"inputs":["[[ function params=0 ]]"],"output":["React"]}</div>
+<div>{"inputs":["[[ function params=0 ]]"],"output":["Compiler"]}</div>
+<div>{"inputs":["[[ function params=0 ]]"],"output":["Compiler"]}</div>
+<div>{"inputs":["[[ function params=0 ]]"],"output":["React"]}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/props-method-dependency.js

@@ -0,0 +1,16 @@
+// @compilationMode(infer)
+import {useMemo} from 'react';
+import {ValidateMemoization} from 'shared-runtime';
+
+function Component(props) {
+  const x = useMemo(() => props.x(), [props.x]);
+  return <ValidateMemoization inputs={[props.x]} output={x} />;
+}
+
+const f = () => ['React'];
+const g = () => ['Compiler'];
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{x: () => ['React']}],
+  sequentialRenders: [{x: f}, {x: g}, {x: g}, {x: f}],
+};