From 2e39ef7c24a5155110929c19f918842972c58989 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Mon, 10 Jun 2024 13:22:07 -0700
Subject: [PATCH] [compiler] Add pruned-scope terminal in HIR

Adds the HIR equivalent of a pruned-scope, allowing us to start porting the scope-pruning passes to operate on HIR.

[ghstack-poisoned]
---
 .../babel-plugin-react-compiler/src/HIR/HIR.ts    | 12 +++++++++++-
 .../src/HIR/PrintHIR.ts                           |  9 ++++++++-
 .../src/HIR/visitors.ts                           | 15 ++++++++++-----
 .../src/ReactiveScopes/BuildReactiveFunction.ts   |  3 ++-
 4 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts
index 7ed4109f0720d..d544269869297 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts
@@ -369,7 +369,8 @@ export type Terminal =
   | SequenceTerminal
   | MaybeThrowTerminal
   | TryTerminal
-  | ReactiveScopeTerminal;
+  | ReactiveScopeTerminal
+  | PrunedScopeTerminal;
 
 export type TerminalWithFallthrough = Terminal & { fallthrough: BlockId };
 
@@ -603,6 +604,15 @@ export type ReactiveScopeTerminal = {
   loc: SourceLocation;
 };
 
+export type PrunedScopeTerminal = {
+  kind: "pruned-scope";
+  fallthrough: BlockId;
+  block: BlockId;
+  scope: ReactiveScope;
+  id: InstructionId;
+  loc: SourceLocation;
+};
+
 /*
  * Instructions generally represent expressions but with all nesting flattened away,
  * such that all operands to each instruction are either primitive values OR are
diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts
index 5e1be135b76f1..f1daea1994553 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts
@@ -127,7 +127,8 @@ export function printMixedHIR(
     case "do-while":
     case "for-in":
     case "for-of":
-    case "scope": {
+    case "scope":
+    case "pruned-scope": {
       const terminal = printTerminal(value);
       if (Array.isArray(terminal)) {
         return terminal.join("; ");
@@ -280,6 +281,12 @@ export function printTerminal(terminal: Terminal): Array<string> | string {
       } fallthrough=bb${terminal.fallthrough}`;
       break;
     }
+    case "pruned-scope": {
+      value = `<pruned> Scope ${printReactiveScopeSummary(terminal.scope)} block=bb${
+        terminal.block
+      } fallthrough=bb${terminal.fallthrough}`;
+      break;
+    }
     case "try": {
       value = `Try block=bb${terminal.block} handler=bb${terminal.handler}${
         terminal.handlerBinding !== null
diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts
index 14c158bb109a1..632e665723f44 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts
@@ -851,7 +851,8 @@ export function mapTerminalSuccessors(
         loc: terminal.loc,
       };
     }
-    case "scope": {
+    case "scope":
+    case "pruned-scope": {
       const block = fn(terminal.block);
       const fallthrough = fn(terminal.fallthrough);
       return {
@@ -904,7 +905,8 @@ export function terminalHasFallthrough<
     case "switch":
     case "ternary":
     case "while":
-    case "scope": {
+    case "scope":
+    case "pruned-scope": {
       const _: BlockId = terminal.fallthrough;
       return true;
     }
@@ -1006,7 +1008,8 @@ export function* eachTerminalSuccessor(terminal: Terminal): Iterable<BlockId> {
       yield terminal.block;
       break;
     }
-    case "scope": {
+    case "scope":
+    case "pruned-scope": {
       yield terminal.block;
       break;
     }
@@ -1072,7 +1075,8 @@ export function mapTerminalOperands(
     case "goto":
     case "unreachable":
     case "unsupported":
-    case "scope": {
+    case "scope":
+    case "pruned-scope": {
       // no-op
       break;
     }
@@ -1130,7 +1134,8 @@ export function* eachTerminalOperand(terminal: Terminal): Iterable<Place> {
     case "goto":
     case "unreachable":
     case "unsupported":
-    case "scope": {
+    case "scope":
+    case "pruned-scope": {
       // no-op
       break;
     }
diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/BuildReactiveFunction.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/BuildReactiveFunction.ts
index cc7e917589deb..ebc1d63c6d881 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/BuildReactiveFunction.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/BuildReactiveFunction.ts
@@ -806,6 +806,7 @@ class Driver {
         }
         break;
       }
+      case "pruned-scope":
       case "scope": {
         const fallthroughId = !this.cx.isScheduled(terminal.fallthrough)
           ? terminal.fallthrough
@@ -828,7 +829,7 @@ class Driver {
 
         this.cx.unscheduleAll(scheduleIds);
         blockValue.push({
-          kind: "scope",
+          kind: terminal.kind,
           instructions: block,
           scope: terminal.scope,
         });