compiler: Improve merging of memo scopes that invalidate together

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,11 +19,21 @@ 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";
+import {
+  printReactiveFunction,
+  printReactiveScopeSummary,
+} from "./PrintReactiveFunction";
 import {
   ReactiveFunctionTransform,
   ReactiveFunctionVisitor,
@@ -78,6 +88,8 @@ import {
 export function mergeReactiveScopesThatInvalidateTogether(
   fn: ReactiveFunction
 ): void {
+  log("MergeReactiveScopesThatInvalidateTogether");
+  log(printReactiveFunction(fn));
   const lastUsageVisitor = new FindLastUsageVisitor();
   visitReactiveFunction(fn, lastUsageVisitor, undefined);
   visitReactiveFunction(fn, new Transform(lastUsageVisitor.lastUsage), null);
@@ -430,7 +442,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 +459,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 = {