compiler: Use types to decide which scopes are eligible for merging

In MergeReactiveScopesThatInvalidateTogether when deciding which scopes were eligible for mergin at all, we looked specifically at the instructions whose lvalue produces the declaration. So if a scope declaration was `t0`, we'd love for the instruction where `t0` was the lvalue and look at the instruction type to decide if it is eligible for merging.

Here, we use the inferred type instead (now that the inferred types support the same set of types of instructions we looked at before). This allows us to find more cases where scopes can be merged.

ghstack-source-id: f36a92a
Pull Request resolved: #29157

Author: josephsavona

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

@@ -13,7 +13,6 @@ import {
   Place,
   ReactiveBlock,
   ReactiveFunction,
-  ReactiveInstruction,
   ReactiveScope,
   ReactiveScopeBlock,
   ReactiveScopeDependencies,
@@ -520,64 +519,7 @@ function scopeIsEligibleForMerging(scopeBlock: ReactiveScopeBlock): boolean {
      */
     return true;
   }
-  const visitor = new DeclarationTypeVisitor(scopeBlock.scope);
-  visitor.visitScope(scopeBlock, undefined);
-  return visitor.alwaysInvalidatesOnInputChange;
-}
-
-class DeclarationTypeVisitor extends ReactiveFunctionVisitor<void> {
-  scope: ReactiveScope;
-  alwaysInvalidatesOnInputChange: boolean = false;
-
-  constructor(scope: ReactiveScope) {
-    super();
-    this.scope = scope;
-  }
-
-  override visitScope(scopeBlock: ReactiveScopeBlock, state: void): void {
-    if (scopeBlock.scope.id !== this.scope.id) {
-      return;
-    }
-    this.traverseScope(scopeBlock, state);
-  }
-
-  override visitInstruction(
-    instruction: ReactiveInstruction,
-    state: void
-  ): void {
-    this.traverseInstruction(instruction, state);
-    if (
-      instruction.lvalue === null ||
-      !this.scope.declarations.has(instruction.lvalue.identifier.id)
-    ) {
-      /*
-       * no lvalue or this instruction isn't directly constructing a
-       * scope output value, skip
-       */
-      log(
-        `    skip instruction lvalue=${
-          instruction.lvalue?.identifier.id
-        } declaration?=${
-          instruction.lvalue != null &&
-          this.scope.declarations.has(instruction.lvalue.identifier.id)
-        } scope=${printReactiveScopeSummary(this.scope)}`
-      );
-      return;
-    }
-    switch (instruction.value.kind) {
-      case "FunctionExpression":
-      case "ArrayExpression":
-      case "JsxExpression":
-      case "JsxFragment":
-      case "ObjectExpression": {
-        /*
-         * These instruction types *always* allocate. If they execute
-         * they will produce a new value, triggering downstream reactive
-         * updates
-         */
-        this.alwaysInvalidatesOnInputChange = true;
-        break;
-      }
-    }
-  }
+  return [...scopeBlock.scope.declarations].some(([, decl]) =>
+    isAlwaysInvalidatingType(decl.identifier.type)
+  );
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allocating-primitive-as-dep-nested-scope.expect.md

@@ -5,15 +5,33 @@
 // bar(props.b) is an allocating expression that produces a primitive, which means
 // that Forget should memoize it.
 // Correctness:
+
+import { identity, mutate, setProperty } from "shared-runtime";
+
 //   - y depends on either bar(props.b) or bar(props.b) + 1
 function AllocatingPrimitiveAsDepNested(props) {
   let x = {};
   mutate(x);
-  let y = foo(bar(props.b) + 1);
-  mutate(x, props.a);
+  let y = identity(identity(props.b) + 1);
+  setProperty(x, props.a);
   return [x, y];
 }
 
+export const FIXTURE_ENTRYPOINT = {
+  fn: AllocatingPrimitiveAsDepNested,
+  params: [{ a: 1, b: 2 }],
+  sequentialRenders: [
+    // change b
+    { a: 1, b: 3 },
+    // change b
+    { a: 1, b: 4 },
+    // change a
+    { a: 2, b: 4 },
+    // change a
+    { a: 3, b: 4 },
+  ],
+};
+
 ```
 
 ## Code
@@ -22,44 +40,56 @@ function AllocatingPrimitiveAsDepNested(props) {
 import { c as _c } from "react/compiler-runtime"; // bar(props.b) is an allocating expression that produces a primitive, which means
 // that Forget should memoize it.
 // Correctness:
+
+import { identity, mutate, setProperty } from "shared-runtime";
+
 //   - y depends on either bar(props.b) or bar(props.b) + 1
 function AllocatingPrimitiveAsDepNested(props) {
-  const $ = _c(9);
-  let x;
-  let y;
+  const $ = _c(5);
+  let t0;
   if ($[0] !== props.b || $[1] !== props.a) {
-    x = {};
+    const x = {};
     mutate(x);
-    const t0 = bar(props.b) + 1;
-    let t1;
-    if ($[4] !== t0) {
-      t1 = foo(t0);
-      $[4] = t0;
-      $[5] = t1;
+    const t1 = identity(props.b) + 1;
+    let t2;
+    if ($[3] !== t1) {
+      t2 = identity(t1);
+      $[3] = t1;
+      $[4] = t2;
     } else {
-      t1 = $[5];
+      t2 = $[4];
     }
-    y = t1;
-    mutate(x, props.a);
+    const y = t2;
+    setProperty(x, props.a);
+    t0 = [x, y];
     $[0] = props.b;
     $[1] = props.a;
-    $[2] = x;
-    $[3] = y;
-  } else {
-    x = $[2];
-    y = $[3];
-  }
-  let t0;
-  if ($[6] !== x || $[7] !== y) {
-    t0 = [x, y];
-    $[6] = x;
-    $[7] = y;
-    $[8] = t0;
+    $[2] = t0;
   } else {
-    t0 = $[8];
+    t0 = $[2];
   }
   return t0;
 }
 
+export const FIXTURE_ENTRYPOINT = {
+  fn: AllocatingPrimitiveAsDepNested,
+  params: [{ a: 1, b: 2 }],
+  sequentialRenders: [
+    // change b
+    { a: 1, b: 3 },
+    // change b
+    { a: 1, b: 4 },
+    // change a
+    { a: 2, b: 4 },
+    // change a
+    { a: 3, b: 4 },
+  ],
+};
+
 ```
-      
+      
+### Eval output
+(kind: ok) [{"wat0":"joe","wat1":1},4]
+[{"wat0":"joe","wat1":1},5]
+[{"wat0":"joe","wat1":2},5]
+[{"wat0":"joe","wat1":3},5]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allocating-primitive-as-dep-nested-scope.js

@@ -1,11 +1,29 @@
 // bar(props.b) is an allocating expression that produces a primitive, which means
 // that Forget should memoize it.
 // Correctness:
+
+import { identity, mutate, setProperty } from "shared-runtime";
+
 //   - y depends on either bar(props.b) or bar(props.b) + 1
 function AllocatingPrimitiveAsDepNested(props) {
   let x = {};
   mutate(x);
-  let y = foo(bar(props.b) + 1);
-  mutate(x, props.a);
+  let y = identity(identity(props.b) + 1);
+  setProperty(x, props.a);
   return [x, y];
 }
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: AllocatingPrimitiveAsDepNested,
+  params: [{ a: 1, b: 2 }],
+  sequentialRenders: [
+    // change b
+    { a: 1, b: 3 },
+    // change b
+    { a: 1, b: 4 },
+    // change a
+    { a: 2, b: 4 },
+    // change a
+    { a: 3, b: 4 },
+  ],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/array-access-assignment.expect.md

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-function foo(a, b, c) {
+function Component({ a, b, c }) {
   const x = [a];
   const y = [null, b];
   const z = [[], [], [c]];
@@ -12,9 +12,15 @@ function foo(a, b, c) {
 }
 
 export const FIXTURE_ENTRYPOINT = {
-  fn: foo,
-  params: [1, 2, 3],
-  isComponent: false,
+  fn: Component,
+  params: [{ a: 1, b: 20, c: 300 }],
+  sequentialRenders: [
+    { a: 2, b: 20, c: 300 },
+    { a: 3, b: 20, c: 300 },
+    { a: 3, b: 21, c: 300 },
+    { a: 3, b: 22, c: 300 },
+    { a: 3, b: 22, c: 301 },
+  ],
 };
 
 ```
@@ -23,52 +29,52 @@ export const FIXTURE_ENTRYPOINT = {
 
 ```javascript
 import { c as _c } from "react/compiler-runtime";
-function foo(a, b, c) {
-  const $ = _c(10);
-  let x;
-  let z;
+function Component(t0) {
+  const $ = _c(6);
+  const { a, b, c } = t0;
+  let t1;
   if ($[0] !== a || $[1] !== b || $[2] !== c) {
-    x = [a];
-    let t0;
-    if ($[5] !== b) {
-      t0 = [null, b];
-      $[5] = b;
-      $[6] = t0;
+    const x = [a];
+    let t2;
+    if ($[4] !== b) {
+      t2 = [null, b];
+      $[4] = b;
+      $[5] = t2;
     } else {
-      t0 = $[6];
+      t2 = $[5];
     }
-    const y = t0;
-    z = [[], [], [c]];
+    const y = t2;
+    const z = [[], [], [c]];
     x[0] = y[1];
     z[0][0] = x[0];
+    t1 = [x, z];
     $[0] = a;
     $[1] = b;
     $[2] = c;
-    $[3] = x;
-    $[4] = z;
+    $[3] = t1;
   } else {
-    x = $[3];
-    z = $[4];
+    t1 = $[3];
   }
-  let t0;
-  if ($[7] !== x || $[8] !== z) {
-    t0 = [x, z];
-    $[7] = x;
-    $[8] = z;
-    $[9] = t0;
-  } else {
-    t0 = $[9];
-  }
-  return t0;
+  return t1;
 }
 
 export const FIXTURE_ENTRYPOINT = {
-  fn: foo,
-  params: [1, 2, 3],
-  isComponent: false,
+  fn: Component,
+  params: [{ a: 1, b: 20, c: 300 }],
+  sequentialRenders: [
+    { a: 2, b: 20, c: 300 },
+    { a: 3, b: 20, c: 300 },
+    { a: 3, b: 21, c: 300 },
+    { a: 3, b: 22, c: 300 },
+    { a: 3, b: 22, c: 301 },
+  ],
 };
 
 ```
       
 ### Eval output
-(kind: ok) [[2],[[2],[],[3]]]
+(kind: ok) [[20],[[20],[],[300]]]
+[[20],[[20],[],[300]]]
+[[21],[[21],[],[300]]]
+[[22],[[22],[],[300]]]
+[[22],[[22],[],[301]]]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/array-access-assignment.js

@@ -1,4 +1,4 @@
-function foo(a, b, c) {
+function Component({ a, b, c }) {
   const x = [a];
   const y = [null, b];
   const z = [[], [], [c]];
@@ -8,7 +8,13 @@ function foo(a, b, c) {
 }
 
 export const FIXTURE_ENTRYPOINT = {
-  fn: foo,
-  params: [1, 2, 3],
-  isComponent: false,
+  fn: Component,
+  params: [{ a: 1, b: 20, c: 300 }],
+  sequentialRenders: [
+    { a: 2, b: 20, c: 300 },
+    { a: 3, b: 20, c: 300 },
+    { a: 3, b: 21, c: 300 },
+    { a: 3, b: 22, c: 300 },
+    { a: 3, b: 22, c: 301 },
+  ],
 };