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: 91ba77d
Pull Request resolved: facebook/react#29781

Author: josephsavona

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

@@ -950,7 +950,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" };
     }