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: 720b9288bb501006329ead733cacbc67df487e16 Pull Request resolved: #29156
facebook · May 18, 2024 · 296e703 · 296e703
1 parent 2d8d37c
commit 296e703
Show file tree

Hide file tree

Showing 26 changed files with 187 additions and 396 deletions.
diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/ObjectShape.ts b/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";

diff --git a/...bel-plugin-react-compiler/src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts b/...bel-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>

diff --git a/compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts b/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":

diff --git a/...babel-plugin-react-compiler/src/__tests__/fixtures/compiler/component.expect.md b/...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 = {

diff --git a/...compiler/error.invalid-aliased-ref-in-callback-invoked-during-render-.expect.md b/...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 |
 ```

diff --git a/...fixtures/compiler/error.invalid-ref-in-callback-invoked-during-render.expect.md b/...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 |
 ```

diff --git a/...mpiler/src/__tests__/fixtures/compiler/existing-variables-with-c-name.expect.md b/...mpiler/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 = {

diff --git a/.../__tests__/fixtures/compiler/for-of-nonmutating-loop-local-collection.expect.md b/.../__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 = {

diff --git a/...n-react-compiler/src/__tests__/fixtures/compiler/hot-module-reloading.expect.md b/...n-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 = {

diff --git a/...ixtures/compiler/inner-memo-value-not-promoted-to-outer-scope-dynamic.expect.md b/...ixtures/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;
 }
 
 ```

diff --git a/...abel-plugin-react-compiler/src/__tests__/fixtures/compiler/jsx-spread.expect.md b/...abel-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;
 }
 
 ```