[compiler] Improve merging of scopes that invalidate together

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

@@ -114,6 +114,99 @@ const TYPED_GLOBALS: Array<[string, BuiltInType]> = [
           returnValueKind: ValueKind.Mutable,
         }),
       ],
+      [
+        'entries',
+        addFunction(DEFAULT_SHAPES, [], {
+          positionalParams: [Effect.Capture],
+          restParam: null,
+          returnType: {kind: 'Object', shapeId: BuiltInArrayId},
+          calleeEffect: Effect.Read,
+          returnValueKind: ValueKind.Mutable,
+          aliasing: {
+            receiver: '@receiver',
+            params: ['@object'],
+            rest: null,
+            returns: '@returns',
+            temporaries: [],
+            effects: [
+              {
+                kind: 'Create',
+                into: '@returns',
+                reason: ValueReason.KnownReturnSignature,
+                value: ValueKind.Mutable,
+              },
+              // Object values are captured into the return
+              {
+                kind: 'Capture',
+                from: '@object',
+                into: '@returns',
+              },
+            ],
+          },
+        }),
+      ],
+      [
+        'keys',
+        addFunction(DEFAULT_SHAPES, [], {
+          positionalParams: [Effect.Read],
+          restParam: null,
+          returnType: {kind: 'Object', shapeId: BuiltInArrayId},
+          calleeEffect: Effect.Read,
+          returnValueKind: ValueKind.Mutable,
+          aliasing: {
+            receiver: '@receiver',
+            params: ['@object'],
+            rest: null,
+            returns: '@returns',
+            temporaries: [],
+            effects: [
+              {
+                kind: 'Create',
+                into: '@returns',
+                reason: ValueReason.KnownReturnSignature,
+                value: ValueKind.Mutable,
+              },
+              // Only keys are captured, and keys are immutable
+              {
+                kind: 'ImmutableCapture',
+                from: '@object',
+                into: '@returns',
+              },
+            ],
+          },
+        }),
+      ],
+      [
+        'values',
+        addFunction(DEFAULT_SHAPES, [], {
+          positionalParams: [Effect.Capture],
+          restParam: null,
+          returnType: {kind: 'Object', shapeId: BuiltInArrayId},
+          calleeEffect: Effect.Read,
+          returnValueKind: ValueKind.Mutable,
+          aliasing: {
+            receiver: '@receiver',
+            params: ['@object'],
+            rest: null,
+            returns: '@returns',
+            temporaries: [],
+            effects: [
+              {
+                kind: 'Create',
+                into: '@returns',
+                reason: ValueReason.KnownReturnSignature,
+                value: ValueKind.Mutable,
+              },
+              // Object values are captured into the return
+              {
+                kind: 'Capture',
+                from: '@object',
+                into: '@returns',
+              },
+            ],
+          },
+        }),
+      ],
     ]),
   ],
   [

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

@@ -142,6 +142,7 @@ function parseAliasingSignatureConfig(
   const effects = typeConfig.effects.map(
     (effect: AliasingEffectConfig): AliasingEffect => {
       switch (effect.kind) {
+        case 'ImmutableCapture':
         case 'CreateFrom':
         case 'Capture':
         case 'Alias':

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

@@ -111,6 +111,19 @@ export const AliasEffectSchema: z.ZodType<AliasEffectConfig> = z.object({
   into: LifetimeIdSchema,
 });
 
+export type ImmutableCaptureEffectConfig = {
+  kind: 'ImmutableCapture';
+  from: string;
+  into: string;
+};
+
+export const ImmutableCaptureEffectSchema: z.ZodType<ImmutableCaptureEffectConfig> =
+  z.object({
+    kind: z.literal('ImmutableCapture'),
+    from: LifetimeIdSchema,
+    into: LifetimeIdSchema,
+  });
+
 export type CaptureEffectConfig = {
   kind: 'Capture';
   from: string;
@@ -187,6 +200,7 @@ export type AliasingEffectConfig =
   | AssignEffectConfig
   | AliasEffectConfig
   | CaptureEffectConfig
+  | ImmutableCaptureEffectConfig
   | ImpureEffectConfig
   | MutateEffectConfig
   | MutateTransitiveConditionallyConfig
@@ -199,6 +213,7 @@ export const AliasingEffectSchema: z.ZodType<AliasingEffectConfig> = z.union([
   AssignEffectSchema,
   AliasEffectSchema,
   CaptureEffectSchema,
+  ImmutableCaptureEffectSchema,
   ImpureEffectSchema,
   MutateEffectSchema,
   MutateTransitiveConditionallySchema,

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 (
@@ -465,20 +480,27 @@ function canMergeScopes(
     (next.scope.dependencies.size !== 0 &&
       [...next.scope.dependencies].every(
         dep =>
+          dep.path.length === 0 &&
           isAlwaysInvalidatingType(dep.identifier.type) &&
           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;
 }