[compiler] Improve merging of scopes that invalidate together

We try to merge consecutive reactive scopes that will always invalidate together, but there's one common case that isn't handled.

```js
const y = [[x]];
```

Here we'll create two consecutive scopes for the inner and outer array expressions. Because the input to the second scope is a temporary, they'll merge into one scope.

But if we name the inner array, the merging stops:

```js
const array = [x];
const y = [array];
```

This is because the merging logic checks if all the dependencies of the second scope are outputs of the first scope, but doesn't account for renaming due to LoadLocal/StoreLocal. The fix is to track these temporaries.

Author: josephsavona

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

@@ -119,6 +119,7 @@ class FindLastUsageVisitor extends ReactiveFunctionVisitor<void> {
 
 class Transform extends ReactiveFunctionTransform<ReactiveScopeDependencies | null> {
   lastUsage: Map<DeclarationId, InstructionId>;
+  temporaries: Map<DeclarationId, DeclarationId> = new Map();
 
   constructor(lastUsage: Map<DeclarationId, InstructionId>) {
     super();
@@ -215,6 +216,12 @@ class Transform extends ReactiveFunctionTransform<ReactiveScopeDependencies | nu
                 current.lvalues.add(
                   instr.instruction.lvalue.identifier.declarationId,
                 );
+                if (instr.instruction.value.kind === 'LoadLocal') {
+                  this.temporaries.set(
+                    instr.instruction.lvalue.identifier.declarationId,
+                    instr.instruction.value.place.identifier.declarationId,
+                  );
+                }
               }
               break;
             }
@@ -236,6 +243,13 @@ class Transform extends ReactiveFunctionTransform<ReactiveScopeDependencies | nu
                   )) {
                     current.lvalues.add(lvalue.identifier.declarationId);
                   }
+                  this.temporaries.set(
+                    instr.instruction.value.lvalue.place.identifier
+                      .declarationId,
+                    this.temporaries.get(
+                      instr.instruction.value.value.identifier.declarationId,
+                    ) ?? instr.instruction.value.value.identifier.declarationId,
+                  );
                 } else {
                   log(
                     `Reset scope @${current.block.scope.id} from StoreLocal in [${instr.instruction.id}]`,
@@ -260,7 +274,7 @@ class Transform extends ReactiveFunctionTransform<ReactiveScopeDependencies | nu
         case 'scope': {
           if (
             current !== null &&
-            canMergeScopes(current.block, instr) &&
+            canMergeScopes(current.block, instr, this.temporaries) &&
             areLValuesLastUsedByScope(
               instr.scope,
               current.lvalues,
@@ -426,6 +440,7 @@ function areLValuesLastUsedByScope(
 function canMergeScopes(
   current: ReactiveScopeBlock,
   next: ReactiveScopeBlock,
+  temporaries: Map<DeclarationId, DeclarationId>,
 ): boolean {
   // Don't merge scopes with reassignments
   if (
@@ -469,16 +484,22 @@ function canMergeScopes(
           Iterable_some(
             current.scope.declarations.values(),
             decl =>
-              decl.identifier.declarationId === dep.identifier.declarationId,
+              decl.identifier.declarationId === dep.identifier.declarationId ||
+              decl.identifier.declarationId ===
+                temporaries.get(dep.identifier.declarationId),
           ),
       ))
   ) {
     log(`  outputs of prev are input to current`);
     return true;
   }
   log(`  cannot merge scopes:`);
-  log(`  ${printReactiveScopeSummary(current.scope)}`);
-  log(`  ${printReactiveScopeSummary(next.scope)}`);
+  log(
+    `  ${printReactiveScopeSummary(current.scope)} ${[...current.scope.declarations.values()].map(decl => decl.identifier.declarationId)}`,
+  );
+  log(
+    `  ${printReactiveScopeSummary(next.scope)} ${[...next.scope.dependencies].map(dep => `${dep.identifier.declarationId} ${temporaries.get(dep.identifier.declarationId) ?? dep.identifier.declarationId}`)}`,
+  );
   return false;
 }
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/array-at-closure.expect.md

@@ -19,7 +19,7 @@ function Component(props) {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function Component(props) {
-  const $ = _c(7);
+  const $ = _c(5);
   let t0;
   if ($[0] !== props.x) {
     t0 = foo(props.x);
@@ -31,26 +31,19 @@ function Component(props) {
   const x = t0;
   let t1;
   if ($[2] !== props || $[3] !== x) {
-    t1 = function () {
+    const fn = function () {
       const arr = [...bar(props)];
       return arr.at(x);
     };
+
+    t1 = fn();
     $[2] = props;
     $[3] = x;
     $[4] = t1;
   } else {
     t1 = $[4];
   }
-  const fn = t1;
-  let t2;
-  if ($[5] !== fn) {
-    t2 = fn();
-    $[5] = fn;
-    $[6] = t2;
-  } else {
-    t2 = $[6];
-  }
-  const fnResult = t2;
+  const fnResult = t1;
   return fnResult;
 }
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/array-map-captures-receiver-noAlias.expect.md

@@ -23,34 +23,18 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function Component(props) {
-  const $ = _c(6);
+  const $ = _c(2);
   let t0;
   if ($[0] !== props.a) {
-    t0 = { a: props.a };
+    const item = { a: props.a };
+    const items = [item];
+    t0 = items.map(_temp);
     $[0] = props.a;
     $[1] = t0;
   } else {
     t0 = $[1];
   }
-  const item = t0;
-  let t1;
-  if ($[2] !== item) {
-    t1 = [item];
-    $[2] = item;
-    $[3] = t1;
-  } else {
-    t1 = $[3];
-  }
-  const items = t1;
-  let t2;
-  if ($[4] !== items) {
-    t2 = items.map(_temp);
-    $[4] = items;
-    $[5] = t2;
-  } else {
-    t2 = $[5];
-  }
-  const mapped = t2;
+  const mapped = t0;
   return mapped;
 }
 function _temp(item_0) {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/array-map-noAlias-escaping-function.expect.md

@@ -21,26 +21,18 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function Component(props) {
-  const $ = _c(4);
+  const $ = _c(2);
   const f = _temp;
   let t0;
   if ($[0] !== props.items) {
-    t0 = [...props.items].map(f);
+    const x = [...props.items].map(f);
+    t0 = [x, f];
     $[0] = props.items;
     $[1] = t0;
   } else {
     t0 = $[1];
   }
-  const x = t0;
-  let t1;
-  if ($[2] !== x) {
-    t1 = [x, f];
-    $[2] = x;
-    $[3] = t1;
-  } else {
-    t1 = $[3];
-  }
-  return t1;
+  return t0;
 }
 function _temp(item) {
   return item;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/capturing-arrow-function-1.expect.md

@@ -23,27 +23,19 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function component(a) {
-  const $ = _c(4);
+  const $ = _c(2);
   let t0;
   if ($[0] !== a) {
-    t0 = { a };
+    const z = { a };
+    t0 = () => {
+      console.log(z);
+    };
     $[0] = a;
     $[1] = t0;
   } else {
     t0 = $[1];
   }
-  const z = t0;
-  let t1;
-  if ($[2] !== z) {
-    t1 = () => {
-      console.log(z);
-    };
-    $[2] = z;
-    $[3] = t1;
-  } else {
-    t1 = $[3];
-  }
-  const x = t1;
+  const x = t0;
   return x;
 }
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/capturing-function-1.expect.md

@@ -23,27 +23,19 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function component(a) {
-  const $ = _c(4);
+  const $ = _c(2);
   let t0;
   if ($[0] !== a) {
-    t0 = { a };
+    const z = { a };
+    t0 = function () {
+      console.log(z);
+    };
     $[0] = a;
     $[1] = t0;
   } else {
     t0 = $[1];
   }
-  const z = t0;
-  let t1;
-  if ($[2] !== z) {
-    t1 = function () {
-      console.log(z);
-    };
-    $[2] = z;
-    $[3] = t1;
-  } else {
-    t1 = $[3];
-  }
-  const x = t1;
+  const x = t0;
   return x;
 }
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/capturing-function-alias-computed-load-2-iife.expect.md

@@ -25,25 +25,17 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function bar(a) {
-  const $ = _c(4);
-  let t0;
-  if ($[0] !== a) {
-    t0 = [a];
-    $[0] = a;
-    $[1] = t0;
-  } else {
-    t0 = $[1];
-  }
-  const x = t0;
+  const $ = _c(2);
   let y;
-  if ($[2] !== x[0][1]) {
+  if ($[0] !== a) {
+    const x = [a];
     y = {};
 
     y = x[0][1];
-    $[2] = x[0][1];
-    $[3] = y;
+    $[0] = a;
+    $[1] = y;
   } else {
-    y = $[3];
+    y = $[1];
   }
   return y;
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/capturing-function-alias-computed-load-3-iife.expect.md

@@ -29,29 +29,20 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function bar(a, b) {
-  const $ = _c(6);
-  let t0;
-  if ($[0] !== a || $[1] !== b) {
-    t0 = [a, b];
-    $[0] = a;
-    $[1] = b;
-    $[2] = t0;
-  } else {
-    t0 = $[2];
-  }
-  const x = t0;
+  const $ = _c(3);
   let y;
-  if ($[3] !== x[0][1] || $[4] !== x[1][0]) {
+  if ($[0] !== a || $[1] !== b) {
+    const x = [a, b];
     y = {};
     let t = {};
 
     y = x[0][1];
     t = x[1][0];
-    $[3] = x[0][1];
-    $[4] = x[1][0];
-    $[5] = y;
+    $[0] = a;
+    $[1] = b;
+    $[2] = y;
   } else {
-    y = $[5];
+    y = $[2];
   }
   return y;
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/capturing-function-alias-computed-load-4-iife.expect.md

@@ -25,25 +25,17 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function bar(a) {
-  const $ = _c(4);
-  let t0;
-  if ($[0] !== a) {
-    t0 = [a];
-    $[0] = a;
-    $[1] = t0;
-  } else {
-    t0 = $[1];
-  }
-  const x = t0;
+  const $ = _c(2);
   let y;
-  if ($[2] !== x[0].a[1]) {
+  if ($[0] !== a) {
+    const x = [a];
     y = {};
 
     y = x[0].a[1];
-    $[2] = x[0].a[1];
-    $[3] = y;
+    $[0] = a;
+    $[1] = y;
   } else {
-    y = $[3];
+    y = $[1];
   }
   return y;
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/capturing-function-alias-computed-load-iife.expect.md

@@ -24,25 +24,17 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function bar(a) {
-  const $ = _c(4);
-  let t0;
-  if ($[0] !== a) {
-    t0 = [a];
-    $[0] = a;
-    $[1] = t0;
-  } else {
-    t0 = $[1];
-  }
-  const x = t0;
+  const $ = _c(2);
   let y;
-  if ($[2] !== x[0]) {
+  if ($[0] !== a) {
+    const x = [a];
     y = {};
 
     y = x[0];
-    $[2] = x[0];
-    $[3] = y;
+    $[0] = a;
+    $[1] = y;
   } else {
-    y = $[3];
+    y = $[1];
   }
   return y;
 }