compiler: Represent pruned scopes instead of inlining

There are a few places where we want to check whether a value actually got memoized, and we currently have to infer this based on values that "should" have a scope and whether a corresponding scope actually exists. This PR adds a new ReactiveStatement variant to model a reactive scope block that was pruned for some reason, and updates all the passes that prune scopes to instead produce this new variant.

ghstack-source-id: aea6dab469acb1f20058b85cb6f9aafab5d167cd
Pull Request resolved: https://github.com/facebook/react/pull/29781

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

@@ -65,12 +65,19 @@ export type ReactiveScopeBlock = {
   instructions: ReactiveBlock;
 };
 
+export type PrunedReactiveScopeBlock = {
+  kind: "pruned-scope";
+  scope: ReactiveScope;
+  instructions: ReactiveBlock;
+};
+
 export type ReactiveBlock = Array<ReactiveStatement>;
 
 export type ReactiveStatement =
   | ReactiveInstructionStatement
   | ReactiveTerminalStatement
-  | ReactiveScopeBlock;
+  | ReactiveScopeBlock
+  | PrunedReactiveScopeBlock;
 
 export type ReactiveInstructionStatement = {
   kind: "instruction";

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

@@ -183,6 +183,7 @@ function visitBlock(context: Context, block: ReactiveBlock): void {
         context.append(stmt, stmt.label);
         break;
       }
+      case "pruned-scope":
       case "scope": {
         CompilerError.invariant(false, {
           reason: "Expected the function to not have scopes already assigned",

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

@@ -400,6 +400,11 @@ function codegenBlockNoReset(
         }
         break;
       }
+      case "pruned-scope": {
+        const scopeBlock = codegenBlockNoReset(cx, item.instructions);
+        statements.push(...scopeBlock.body);
+        break;
+      }
       case "scope": {
         const temp = new Map(cx.temp);
         codegenReactiveScope(cx, statements, item.scope, item.instructions);

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

@@ -34,7 +34,14 @@ class Transform extends ReactiveFunctionTransform<boolean> {
   ): Transformed<ReactiveStatement> {
     this.visitScope(scope, isWithinLoop);
     if (isWithinLoop) {
-      return { kind: "replace-many", value: scope.instructions };
+      return {
+        kind: "replace",
+        value: {
+          kind: "pruned-scope",
+          scope: scope.scope,
+          instructions: scope.instructions,
+        },
+      };
     } else {
       return { kind: "keep" };
     }

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

@@ -66,7 +66,14 @@ class Transform extends ReactiveFunctionTransform<State> {
     this.visitScope(scope, innerState);
     outerState.hasHook ||= innerState.hasHook;
     if (innerState.hasHook) {
-      return { kind: "replace-many", value: scope.instructions };
+      return {
+        kind: "replace",
+        value: {
+          kind: "pruned-scope",
+          scope: scope.scope,
+          instructions: scope.instructions,
+        },
+      };
     } else {
       return { kind: "keep" };
     }

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

@@ -174,6 +174,16 @@ class Transform extends ReactiveFunctionTransform<ReactiveScopeDependencies | nu
           }
           break;
         }
+        case "pruned-scope": {
+          // For now we don't merge across pruned scopes
+          if (current !== null) {
+            log(
+              `Reset scope @${current.block.scope.id} from pruned scope @${instr.scope.id}`
+            );
+            reset();
+          }
+          break;
+        }
         case "instruction": {
           switch (instr.instruction.value.kind) {
             case "ComputedLoad":

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

@@ -7,6 +7,7 @@
 
 import { CompilerError } from "../CompilerError";
 import {
+  PrunedReactiveScopeBlock,
   ReactiveFunction,
   ReactiveScope,
   ReactiveScopeBlock,
@@ -83,6 +84,15 @@ export function writeReactiveBlock(
   writer.writeLine("}");
 }
 
+export function writePrunedScope(
+  writer: Writer,
+  block: PrunedReactiveScopeBlock
+): void {
+  writer.writeLine(`<pruned> ${printReactiveScopeSummary(block.scope)} {`);
+  writeReactiveInstructions(writer, block.instructions);
+  writer.writeLine("}");
+}
+
 export function printDependency(dependency: ReactiveScopeDependency): string {
   const identifier =
     printIdentifier(dependency.identifier) +
@@ -133,6 +143,10 @@ function writeReactiveInstruction(
       writeReactiveBlock(writer, instr);
       break;
     }
+    case "pruned-scope": {
+      writePrunedScope(writer, instr);
+      break;
+    }
     case "terminal": {
       if (instr.label !== null) {
         writer.write(`bb${instr.label.id}: `);

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

@@ -107,7 +107,14 @@ class Transform extends ReactiveFunctionTransform<boolean> {
             this.unmemoizedValues.add(identifier);
           }
         }
-        return { kind: "replace-many", value: scopeBlock.instructions };
+        return {
+          kind: "replace",
+          value: {
+            kind: "pruned-scope",
+            scope: scopeBlock.scope,
+            instructions: scopeBlock.instructions,
+          },
+        };
       }
     }
     return { kind: "keep" };

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

@@ -951,7 +951,14 @@ class PruneScopesTransform extends ReactiveFunctionTransform<
       return { kind: "keep" };
     } else {
       this.prunedScopes.add(scopeBlock.scope.id);
-      return { kind: "replace-many", value: scopeBlock.instructions };
+      return {
+        kind: "replace",
+        value: {
+          kind: "pruned-scope",
+          scope: scopeBlock.scope,
+          instructions: scopeBlock.instructions,
+        },
+      };
     }
   }
 

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

@@ -51,7 +51,14 @@ class Transform extends ReactiveFunctionTransform<State> {
          */
         !hasOwnDeclaration(scopeBlock))
     ) {
-      return { kind: "replace-many", value: scopeBlock.instructions };
+      return {
+        kind: "replace",
+        value: {
+          kind: "pruned-scope",
+          scope: scopeBlock.scope,
+          instructions: scopeBlock.instructions,
+        },
+      };
     } else {
       return { kind: "keep" };
     }

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

@@ -12,6 +12,7 @@ import {
   IdentifierName,
   InstructionId,
   Place,
+  PrunedReactiveScopeBlock,
   ReactiveBlock,
   ReactiveFunction,
   ReactiveScopeBlock,
@@ -84,6 +85,13 @@ class Visitor extends ReactiveFunctionVisitor<Scopes> {
     });
   }
 
+  override visitPrunedScope(
+    scopeBlock: PrunedReactiveScopeBlock,
+    state: Scopes
+  ): void {
+    this.traverseBlock(scopeBlock.instructions, state);
+  }
+
   override visitScope(scope: ReactiveScopeBlock, state: Scopes): void {
     for (const [_, declaration] of scope.scope.declarations) {
       state.visit(declaration.identifier);

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

@@ -9,6 +9,7 @@ import {
   HIRFunction,
   InstructionId,
   Place,
+  PrunedReactiveScopeBlock,
   ReactiveBlock,
   ReactiveFunction,
   ReactiveInstruction,
@@ -196,6 +197,16 @@ export class ReactiveFunctionVisitor<TState = void> {
     this.visitBlock(scope.instructions, state);
   }
 
+  visitPrunedScope(scopeBlock: PrunedReactiveScopeBlock, state: TState): void {
+    this.traversePrunedScope(scopeBlock, state);
+  }
+  traversePrunedScope(
+    scopeBlock: PrunedReactiveScopeBlock,
+    state: TState
+  ): void {
+    this.visitBlock(scopeBlock.instructions, state);
+  }
+
   visitBlock(block: ReactiveBlock, state: TState): void {
     this.traverseBlock(block, state);
   }
@@ -210,6 +221,10 @@ export class ReactiveFunctionVisitor<TState = void> {
           this.visitScope(instr, state);
           break;
         }
+        case "pruned-scope": {
+          this.visitPrunedScope(instr, state);
+          break;
+        }
         case "terminal": {
           this.visitTerminal(instr, state);
           break;
@@ -273,6 +288,10 @@ export class ReactiveFunctionTransform<
           transformed = this.transformScope(instr, state);
           break;
         }
+        case "pruned-scope": {
+          transformed = this.transformPrunedScope(instr, state);
+          break;
+        }
         case "terminal": {
           transformed = this.transformTerminal(instr, state);
           break;
@@ -339,6 +358,14 @@ export class ReactiveFunctionTransform<
     return { kind: "keep" };
   }
 
+  transformPrunedScope(
+    scope: PrunedReactiveScopeBlock,
+    state: TState
+  ): Transformed<ReactiveStatement> {
+    this.visitPrunedScope(scope, state);
+    return { kind: "keep" };
+  }
+
   transformValue(
     id: InstructionId,
     value: ReactiveValue,