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,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,6 +437,11 @@ function canMergeScopes(
         }))
       ),
       next.scope.dependencies
+    ) ||
+    [...next.scope.dependencies].some(
+      (dep) =>
+        current.scope.declarations.has(dep.identifier.id) &&
+        isAlwaysInvalidatingType(dep.identifier.type)
     )
   ) {
     log(`  outputs of prev are input to current`);
@@ -441,6 +453,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/component.expect.md

@@ -42,7 +42,7 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function Component(props) {
-  const $ = _c(8);
+  const $ = _c(5);
   const items = props.items;
   const maxItems = props.maxItems;
   let renderedItems;
@@ -72,27 +72,18 @@ function Component(props) {
   const count = renderedItems.length;
   let t0;
   if ($[3] !== count) {
-    t0 = <h1>{count} Items</h1>;
-    $[3] = count;
-    $[4] = t0;
-  } else {
-    t0 = $[4];
-  }
-  let t1;
-  if ($[5] !== t0 || $[6] !== renderedItems) {
-    t1 = (
+    t0 = (
       <div>
-        {t0}
+        <h1>{count} Items</h1>
         {renderedItems}
       </div>
     );
-    $[5] = t0;
-    $[6] = renderedItems;
-    $[7] = t1;
+    $[3] = count;
+    $[4] = t0;
   } else {
-    t1 = $[7];
+    t0 = $[4];
   }
-  return t1;
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {

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/existing-variables-with-c-name.expect.md

@@ -33,7 +33,7 @@ import { useMemo, useState } from "react";
 import { ValidateMemoization } from "shared-runtime";
 
 function Component(props) {
-  const $ = _c2(7);
+  const $ = _c2(4);
   const [state] = useState(0);
 
   const c = state;
@@ -54,22 +54,13 @@ function Component(props) {
   const array = t0;
   let t2;
   if ($[2] !== state) {
-    t2 = [state];
+    t2 = <ValidateMemoization inputs={[state]} output={array} />;
     $[2] = state;
     $[3] = t2;
   } else {
     t2 = $[3];
   }
-  let t3;
-  if ($[4] !== t2 || $[5] !== array) {
-    t3 = <ValidateMemoization inputs={t2} output={array} />;
-    $[4] = t2;
-    $[5] = array;
-    $[6] = t3;
-  } else {
-    t3 = $[6];
-  }
-  return t3;
+  return t2;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/for-of-nonmutating-loop-local-collection.expect.md

@@ -44,7 +44,7 @@ import { useMemo } from "react";
 import { ValidateMemoization } from "shared-runtime";
 
 function Component(t0) {
-  const $ = _c(19);
+  const $ = _c(10);
   const { a, b } = t0;
   let t1;
   let t2;
@@ -75,54 +75,27 @@ function Component(t0) {
   const y = t3;
   let t4;
   if ($[5] !== a) {
-    t4 = [a];
+    t4 = <ValidateMemoization inputs={[a]} output={x} />;
     $[5] = a;
     $[6] = t4;
   } else {
     t4 = $[6];
   }
   let t5;
-  if ($[7] !== t4 || $[8] !== x) {
-    t5 = <ValidateMemoization inputs={t4} output={x} />;
-    $[7] = t4;
-    $[8] = x;
-    $[9] = t5;
-  } else {
-    t5 = $[9];
-  }
-  let t6;
-  if ($[10] !== x || $[11] !== b) {
-    t6 = [x, b];
-    $[10] = x;
-    $[11] = b;
-    $[12] = t6;
-  } else {
-    t6 = $[12];
-  }
-  let t7;
-  if ($[13] !== t6 || $[14] !== y) {
-    t7 = <ValidateMemoization inputs={t6} output={y} />;
-    $[13] = t6;
-    $[14] = y;
-    $[15] = t7;
-  } else {
-    t7 = $[15];
-  }
-  let t8;
-  if ($[16] !== t5 || $[17] !== t7) {
-    t8 = (
+  if ($[7] !== x || $[8] !== b) {
+    t5 = (
       <>
-        {t5}
-        {t7}
+        {t4}
+        <ValidateMemoization inputs={[x, b]} output={y} />
       </>
     );
-    $[16] = t5;
-    $[17] = t7;
-    $[18] = t8;
+    $[7] = x;
+    $[8] = b;
+    $[9] = t5;
   } else {
-    t8 = $[18];
+    t5 = $[9];
   }
-  return t8;
+  return t5;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/hot-module-reloading.expect.md

@@ -28,11 +28,11 @@ import { useMemo, useState } from "react";
 import { ValidateMemoization } from "shared-runtime";
 
 function Component(props) {
-  const $ = _c(8);
+  const $ = _c(5);
   if (
     $[0] !== "bb6936608c0afe8e313aa547ca09fbc8451f24664284368812127c7e9bc2bca9"
   ) {
-    for (let $i = 0; $i < 8; $i += 1) {
+    for (let $i = 0; $i < 5; $i += 1) {
       $[$i] = Symbol.for("react.memo_cache_sentinel");
     }
     $[0] = "bb6936608c0afe8e313aa547ca09fbc8451f24664284368812127c7e9bc2bca9";
@@ -52,22 +52,13 @@ function Component(props) {
   const doubled = t0;
   let t3;
   if ($[3] !== state) {
-    t3 = [state];
+    t3 = <ValidateMemoization inputs={[state]} output={doubled} />;
     $[3] = state;
     $[4] = t3;
   } else {
     t3 = $[4];
   }
-  let t4;
-  if ($[5] !== t3 || $[6] !== doubled) {
-    t4 = <ValidateMemoization inputs={t3} output={doubled} />;
-    $[5] = t3;
-    $[6] = doubled;
-    $[7] = t4;
-  } else {
-    t4 = $[7];
-  }
-  return t4;
+  return t3;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inner-memo-value-not-promoted-to-outer-scope-dynamic.expect.md

@@ -24,7 +24,7 @@ function Component(props) {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function Component(props) {
-  const $ = _c(15);
+  const $ = _c(8);
   const item = useFragment(FRAGMENT, props.item);
   useFreeze(item);
   let t0;
@@ -56,37 +56,20 @@ function Component(props) {
   }
   let t2;
   if ($[6] !== t0) {
-    t2 = <span>{t0}</span>;
+    t2 = (
+      <T1>
+        <T0>
+          {t1}
+          <span>{t0}</span>
+        </T0>
+      </T1>
+    );
     $[6] = t0;
     $[7] = t2;
   } else {
     t2 = $[7];
   }
-  let t3;
-  if ($[8] !== T0 || $[9] !== t1 || $[10] !== t2) {
-    t3 = (
-      <T0>
-        {t1}
-        {t2}
-      </T0>
-    );
-    $[8] = T0;
-    $[9] = t1;
-    $[10] = t2;
-    $[11] = t3;
-  } else {
-    t3 = $[11];
-  }
-  let t4;
-  if ($[12] !== T1 || $[13] !== t3) {
-    t4 = <T1>{t3}</T1>;
-    $[12] = T1;
-    $[13] = t3;
-    $[14] = t4;
-  } else {
-    t4 = $[14];
-  }
-  return t4;
+  return t2;
 }
 
 ```

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/jsx-spread.expect.md

@@ -15,27 +15,18 @@ function Component(props) {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function Component(props) {
-  const $ = _c(5);
+  const $ = _c(2);
 
   const t0 = props.cond ? props.foo : props.bar;
   let t1;
   if ($[0] !== t0) {
-    t1 = { bar: t0 };
+    t1 = <Component {...props} {...{ bar: t0 }} />;
     $[0] = t0;
     $[1] = t1;
   } else {
     t1 = $[1];
   }
-  let t2;
-  if ($[2] !== props || $[3] !== t1) {
-    t2 = <Component {...props} {...t1} />;
-    $[2] = props;
-    $[3] = t1;
-    $[4] = t2;
-  } else {
-    t2 = $[4];
-  }
-  return t2;
+  return t1;
 }
 
 ```