From 36853d119e5d999caaa02c04398f2237563bc648 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Fri, 19 Jul 2024 10:23:47 +0900
Subject: [PATCH] [compiler] Maintain RPO and unique instruction ids when
 constructing scope terminals

Later passes may rely on HIR invariants such as blocks being in RPO or instructions having unique, ascending InstructionIds. However, BuildReactiveScopeTerminalsHIR doesn't currently gurantee this.

This PR updates that pass to first restore RPO, fixup predecessors (the previous logic tried to do this but failed on unreachable blocks, where `markPredecessors()` handles that case), and renumber instructions. Then it walks instructions and scopes to update identifier and scope ranges given the new instruction ids.

ghstack-source-id: 2127da749ff7d2f9f27144f6b6e484250389b0de
Pull Request resolved: https://github.com/facebook/react/pull/30399
---
 .../src/HIR/AssertTerminalBlocksExist.ts      |  1 +
 .../src/HIR/BuildReactiveScopeTerminalsHIR.ts | 64 ++++++++++++++++---
 2 files changed, 57 insertions(+), 8 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/AssertTerminalBlocksExist.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/AssertTerminalBlocksExist.ts
index 1708f7beefea6..e696fddbdeb26 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/AssertTerminalBlocksExist.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/AssertTerminalBlocksExist.ts
@@ -39,6 +39,7 @@ export function assertTerminalPredsExist(fn: HIRFunction): void {
         [...eachTerminalSuccessor(predBlock.terminal)].includes(block.id),
         {
           reason: 'Terminal successor does not reference correct predecessor',
+          description: `Block bb${block.id} has bb${predBlock.id} as a predecessor, but bb${predBlock.id}'s successors do not include bb${block.id}`,
           loc: GeneratedSource,
         },
       );
diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildReactiveScopeTerminalsHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildReactiveScopeTerminalsHIR.ts
index bf12bfd86fa87..6819c2c6c5bfb 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildReactiveScopeTerminalsHIR.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildReactiveScopeTerminalsHIR.ts
@@ -9,10 +9,17 @@ import {
   GotoVariant,
   HIRFunction,
   InstructionId,
+  makeInstructionId,
   ReactiveScope,
   ReactiveScopeTerminal,
   ScopeId,
 } from './HIR';
+import {
+  markInstructionIds,
+  markPredecessors,
+  reversePostorderBlocks,
+} from './HIRBuilder';
+import {eachInstructionLValue} from './visitors';
 
 /**
  * This pass assumes that all program blocks are properly nested with respect to fallthroughs
@@ -142,16 +149,9 @@ export function buildReactiveScopeTerminalsHIR(fn: HIRFunction): void {
 
   /**
    * Step 3:
-   * Repoint preds and phis when they refer to a rewritten block.
+   * Repoint phis when they refer to a rewritten block.
    */
   for (const [, block] of originalBlocks) {
-    for (const pred of block.preds) {
-      const newId = rewrittenFinalBlocks.get(pred);
-      if (newId != null) {
-        block.preds.delete(pred);
-        block.preds.add(newId);
-      }
-    }
     for (const phi of block.phis) {
       for (const [originalId, value] of phi.operands) {
         const newId = rewrittenFinalBlocks.get(originalId);
@@ -162,6 +162,54 @@ export function buildReactiveScopeTerminalsHIR(fn: HIRFunction): void {
       }
     }
   }
+
+  /**
+   * Step 4:
+   * Fixup the HIR to restore RPO, ensure correct predecessors, and
+   * renumber instructions. Note that the renumbering instructions
+   * invalidates scope and identifier ranges, so we fix them in the
+   * next step.
+   */
+  reversePostorderBlocks(fn.body);
+  markPredecessors(fn.body);
+  markInstructionIds(fn.body);
+
+  /**
+   * Step 5:
+   * Fix scope and identifier ranges to account for renumbered instructions
+   */
+  for (const [, block] of fn.body.blocks) {
+    for (const instruction of block.instructions) {
+      for (const lvalue of eachInstructionLValue(instruction)) {
+        /*
+         * Any lvalues whose mutable range was a single instruction must have
+         * started at the current instruction, so update the range to match
+         * the instruction's new id
+         */
+        if (
+          lvalue.identifier.mutableRange.end ===
+          lvalue.identifier.mutableRange.start + 1
+        ) {
+          lvalue.identifier.mutableRange.start = instruction.id;
+          lvalue.identifier.mutableRange.end = makeInstructionId(
+            instruction.id + 1,
+          );
+        }
+      }
+    }
+    const terminal = block.terminal;
+    if (terminal.kind === 'scope' || terminal.kind === 'pruned-scope') {
+      /*
+       * Scope ranges should always align to start at the 'scope' terminal
+       * and end at the first instruction of the fallthrough block
+       */
+      const fallthroughBlock = fn.body.blocks.get(terminal.fallthrough)!;
+      const firstId =
+        fallthroughBlock.instructions[0]?.id ?? fallthroughBlock.terminal.id;
+      terminal.scope.range.start = terminal.id;
+      terminal.scope.range.end = firstId;
+    }
+  }
 }
 
 type TerminalRewriteInfo =