compiler: Improve merging of memo scopes that invalidate together

Improves merging of consecutive scopes so that we now merge two scopes if the dependencies of the second scope are a subset of the previous scope's output *and* that dependency has a type that will always produce a new value (array, object, jsx, function) if it is re-evaluated.

To make this easier, we extend the set of builtin types to include ones for function expressions and JSX and to infer these types in InferTypes. This allows using the already inferred types in MergeReactiveScopesThatInvalidateTogether.

ghstack-source-id: e9119fc
Pull Request resolved: #29156

Author: josephsavona

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

@@ -188,6 +188,8 @@ export type ObjectShape = {
  */
 export type ShapeRegistry = Map<string, ObjectShape>;
 export const BuiltInArrayId = "BuiltInArray";
+export const BuiltInFunctionId = "BuiltInFunction";
+export const BuiltInJsxId = "BuiltInJsx";
 export const BuiltInObjectId = "BuiltInObject";
 export const BuiltInUseStateId = "BuiltInUseState";
 export const BuiltInSetStateId = "BuiltInSetState";

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

@@ -19,8 +19,15 @@ import {
   ReactiveScopeDependencies,
   ReactiveScopeDependency,
   ReactiveStatement,
+  Type,
   makeInstructionId,
 } from "../HIR";
+import {
+  BuiltInArrayId,
+  BuiltInFunctionId,
+  BuiltInJsxId,
+  BuiltInObjectId,
+} from "../HIR/ObjectShape";
 import { eachInstructionLValue } from "../HIR/visitors";
 import { assertExhaustive } from "../Utils/utils";
 import { printReactiveScopeSummary } from "./PrintReactiveFunction";
@@ -430,7 +437,13 @@ function canMergeScopes(
         }))
       ),
       next.scope.dependencies
-    )
+    ) ||
+    (next.scope.dependencies.size !== 0 &&
+      [...next.scope.dependencies].every(
+        (dep) =>
+          current.scope.declarations.has(dep.identifier.id) &&
+          isAlwaysInvalidatingType(dep.identifier.type)
+      ))
   ) {
     log(`  outputs of prev are input to current`);
     return true;
@@ -441,6 +454,20 @@ function canMergeScopes(
   return false;
 }
 
+function isAlwaysInvalidatingType(type: Type): boolean {
+  if (type.kind === "Object") {
+    switch (type.shapeId) {
+      case BuiltInArrayId:
+      case BuiltInObjectId:
+      case BuiltInFunctionId:
+      case BuiltInJsxId: {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
 function areEqualDependencies(
   a: Set<ReactiveScopeDependency>,
   b: Set<ReactiveScopeDependency>

compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts

@@ -20,6 +20,8 @@ import {
 } from "../HIR/HIR";
 import {
   BuiltInArrayId,
+  BuiltInFunctionId,
+  BuiltInJsxId,
   BuiltInObjectId,
   BuiltInUseRefId,
 } from "../HIR/ObjectShape";
@@ -313,6 +315,7 @@ function* generateInstructionTypes(
 
     case "FunctionExpression": {
       yield* generate(value.loweredFunc.func);
+      yield equation(left, { kind: "Object", shapeId: BuiltInFunctionId });
       break;
     }
 
@@ -327,10 +330,13 @@ function* generateInstructionTypes(
       break;
     }
 
+    case "JsxExpression":
+    case "JsxFragment": {
+      yield equation(left, { kind: "Object", shapeId: BuiltInJsxId });
+      break;
+    }
     case "DeclareLocal":
     case "NewExpression":
-    case "JsxExpression":
-    case "JsxFragment":
     case "RegExpLiteral":
     case "PropertyStore":
     case "ComputedStore":

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-aliased-ref-in-callback-invoked-during-render-.expect.md

@@ -22,7 +22,7 @@ function Component(props) {
    7 |     return <Foo item={item} current={current} />;
    8 |   };
 >  9 |   return <Items>{props.items.map((item) => renderItem(item))}</Items>;
-     |                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^ InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef). Cannot access ref value at mutate? $64[13:15] (9:9)
+     |                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^ InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef). Cannot access ref value at mutate? $64[13:15]:TObject<BuiltInFunction> (9:9)
   10 | }
   11 |
 ```

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-ref-in-callback-invoked-during-render.expect.md

@@ -21,7 +21,7 @@ function Component(props) {
    6 |     return <Foo item={item} current={current} />;
    7 |   };
 >  8 |   return <Items>{props.items.map((item) => renderItem(item))}</Items>;
-     |                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^ InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef). Cannot access ref value at mutate? $60[14:16] (8:8)
+     |                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^ InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef). Cannot access ref value at mutate? $60[14:16]:TObject<BuiltInFunction> (8:8)
    9 | }
   10 |
 ```

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-consecutive-scopes-deps-subset-of-decls.expect.md

@@ -33,44 +33,23 @@ import { c as _c } from "react/compiler-runtime";
 import { useState } from "react";
 
 function Component() {
-  const $ = _c(8);
+  const $ = _c(2);
   const [count, setCount] = useState(0);
   let t0;
-  let t1;
   if ($[0] !== count) {
-    t0 = <button onClick={() => setCount(count - 1)}>Decrement</button>;
+    t0 = (
+      <div>
+        <button onClick={() => setCount(count - 1)}>Decrement</button>
 
-    t1 = () => setCount(count + 1);
+        <button onClick={() => setCount(count + 1)}>Increment</button>
+      </div>
+    );
     $[0] = count;
     $[1] = t0;
-    $[2] = t1;
   } else {
     t0 = $[1];
-    t1 = $[2];
-  }
-  let t2;
-  if ($[3] !== t1) {
-    t2 = <button onClick={t1}>Increment</button>;
-    $[3] = t1;
-    $[4] = t2;
-  } else {
-    t2 = $[4];
-  }
-  let t3;
-  if ($[5] !== t0 || $[6] !== t2) {
-    t3 = (
-      <div>
-        {t0}
-        {t2}
-      </div>
-    );
-    $[5] = t0;
-    $[6] = t2;
-    $[7] = t3;
-  } else {
-    t3 = $[7];
   }
-  return t3;
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {