[compiler][rewrite] Patch logic for aligning scopes to non-value blocks

Our previous logic for aligning scopes to block scopes constructs a tree of block and scope nodes. We ensured that blocks always mapped to the same node as their fallthroughs. e.g.
```js
// source
a();
if (...) {
  b();
}
c();

// HIR
bb0:
a()
if test=... consequent=bb1 fallthrough=bb2

bb1:
b()
goto bb2

bb2:
c()

// AlignReactiveScopesToBlockScopesHIR nodes
Root node (maps to both bb0 and bb2)
  |- bb1
  |- ...
```

There are two issues with the existing implementation:
1. Only scopes that overlap with the beginning of a block are aligned correctly. This is because the traversal does not store information about the block-fallthrough pair for scopes that begin *within* the block-fallthrough range.
```
\# This case gets handled correctly
         ┌──────────────┐
         │              │
         block start    block end

scope start     scope end
│               │
└───────────────┘

\# But not this one!
┌──────────────┐
│              │
block start    block end

          scope start     scope end
          │               │
          └───────────────┘
```
2. Only scopes that are directly used by a block is considered. See the `align-scopes-nested-block-structure` fixture for details.

ghstack-source-id: 327dec5
Pull Request resolved: #29891

Author: mofeiZ

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

@@ -22,8 +22,10 @@ import {
   mapTerminalSuccessors,
   terminalFallthrough,
 } from "../HIR/visitors";
+import { retainWhere_Set } from "../Utils/utils";
 import { getPlaceScope } from "./BuildReactiveBlocks";
 
+type InstructionRange = MutableRange;
 /*
  * Note: this is the 2nd of 4 passes that determine how to break a function into discrete
  * reactive scopes (independently memoizeable units of code):
@@ -66,18 +68,20 @@ import { getPlaceScope } from "./BuildReactiveBlocks";
  * will be the updated end for that scope).
  */
 export function alignReactiveScopesToBlockScopesHIR(fn: HIRFunction): void {
-  const blockNodes = new Map<BlockId, BlockNode>();
-  const rootNode: BlockNode = {
-    kind: "node",
-    valueRange: null,
-    children: [],
-    id: makeInstructionId(0),
-  };
-  blockNodes.set(fn.body.entry, rootNode);
+  const activeBlockFallthroughRanges: Array<{
+    range: InstructionRange;
+    fallthrough: BlockId;
+  }> = [];
+  const activeScopes = new Set<ReactiveScope>();
   const seen = new Set<ReactiveScope>();
+  const valueBlockNodes = new Map<BlockId, ValueBlockNode>();
   const placeScopes = new Map<Place, ReactiveScope>();
 
-  function recordPlace(id: InstructionId, place: Place, node: BlockNode): void {
+  function recordPlace(
+    id: InstructionId,
+    place: Place,
+    node: ValueBlockNode | null
+  ): void {
     if (place.identifier.scope !== null) {
       placeScopes.set(place, place.identifier.scope);
     }
@@ -86,13 +90,14 @@ export function alignReactiveScopesToBlockScopesHIR(fn: HIRFunction): void {
     if (scope == null) {
       return;
     }
-    node.children.push({ kind: "scope", scope, id });
+    activeScopes.add(scope);
+    node?.children.push({ kind: "scope", scope, id });
 
     if (seen.has(scope)) {
       return;
     }
     seen.add(scope);
-    if (node.valueRange !== null) {
+    if (node != null && node.valueRange !== null) {
       scope.range.start = makeInstructionId(
         Math.min(node.valueRange.start, scope.range.start)
       );
@@ -103,16 +108,25 @@ export function alignReactiveScopesToBlockScopesHIR(fn: HIRFunction): void {
   }
 
   for (const [, block] of fn.body.blocks) {
-    const { instructions, terminal } = block;
-    const node = blockNodes.get(block.id);
-    if (node === undefined) {
-      CompilerError.invariant(false, {
-        reason: `Expected a node to be initialized for block`,
-        loc: instructions[0]?.loc ?? terminal.loc,
-        description: `No node for block bb${block.id} (${block.kind})`,
-      });
+    const startingId = block.instructions[0]?.id ?? block.terminal.id;
+    retainWhere_Set(activeScopes, (scope) => scope.range.end > startingId);
+    const top = activeBlockFallthroughRanges.at(-1);
+    if (top?.fallthrough === block.id) {
+      activeBlockFallthroughRanges.pop();
+      /*
+       * All active scopes must have either started before or within the last
+       * block-fallthrough range. In either case, they overlap this block-
+       * fallthrough range and can have their ranges extended.
+       */
+      for (const scope of activeScopes) {
+        scope.range.start = makeInstructionId(
+          Math.min(scope.range.start, top.range.start)
+        );
+      }
     }
 
+    const { instructions, terminal } = block;
+    const node = valueBlockNodes.get(block.id) ?? null;
     for (const instr of instructions) {
       for (const lvalue of eachInstructionLValue(instr)) {
         recordPlace(instr.id, lvalue, node);
@@ -125,36 +139,42 @@ export function alignReactiveScopesToBlockScopesHIR(fn: HIRFunction): void {
       recordPlace(terminal.id, operand, node);
     }
 
-    // Save the current node for the fallback block, where this block scope continues
     const fallthrough = terminalFallthrough(terminal);
-    if (fallthrough !== null && !blockNodes.has(fallthrough)) {
+    if (fallthrough !== null) {
       /*
-       * Any scopes that carried over across a terminal->fallback need their range extended
-       * to at least the first instruction of the fallback
-       *
-       * Note that it's possible for a terminal such as an if or switch to have a null fallback,
-       * indicating that all control-flow paths diverge instead of reaching the fallthrough.
-       * In this case there isn't an instruction id in the program that we can point to for the
-       * updated range. Since the output is correct in this case we leave it, but it would be
-       * more correct to find the maximum instuction id in the whole program and set the range.end
-       * to one greater. Alternatively, we could leave in an unreachable fallthrough (with a new
-       * "unreachable" terminal variant, perhaps) and use that instruction id.
+       * Any currently active scopes that overlaps the block-fallthrough range
+       * need their range extended to at least the first instruction of the
+       * fallthrough
        */
       const fallthroughBlock = fn.body.blocks.get(fallthrough)!;
       const nextId =
         fallthroughBlock.instructions[0]?.id ?? fallthroughBlock.terminal.id;
-      for (const child of node.children) {
-        if (child.kind !== "scope") {
-          continue;
-        }
-        const scope = child.scope;
+      for (const scope of activeScopes) {
         if (scope.range.end > terminal.id) {
           scope.range.end = makeInstructionId(
             Math.max(scope.range.end, nextId)
           );
         }
       }
-      blockNodes.set(fallthrough, node);
+      /**
+       * We also record the block-fallthrough range for future scopes that begin
+       * within the range (and overlap with the range end).
+       */
+      activeBlockFallthroughRanges.push({
+        fallthrough,
+        range: {
+          start: terminal.id,
+          end: nextId,
+        },
+      });
+
+      CompilerError.invariant(!valueBlockNodes.has(fallthrough), {
+        reason: "Expect hir blocks to have unique fallthroughs",
+        loc: terminal.loc,
+      });
+      if (node != null) {
+        valueBlockNodes.set(fallthrough, node);
+      }
     }
 
     /*
@@ -166,48 +186,35 @@ export function alignReactiveScopesToBlockScopesHIR(fn: HIRFunction): void {
      * just those that are direct successors for normal control-flow ordering.
      */
     mapTerminalSuccessors(terminal, (successor) => {
-      if (blockNodes.has(successor)) {
+      if (valueBlockNodes.has(successor)) {
         return successor;
       }
 
       const successorBlock = fn.body.blocks.get(successor)!;
-      /*
-       * we need the block kind check here because the do..while terminal's successor
-       * is a block, and try's successor is a catch block
-       */
       if (successorBlock.kind === "block" || successorBlock.kind === "catch") {
-        const childNode: BlockNode = {
-          kind: "node",
-          id: terminal.id,
-          children: [],
-          valueRange: null,
-        };
-        node.children.push(childNode);
-        blockNodes.set(successor, childNode);
+        /*
+         * we need the block kind check here because the do..while terminal's
+         * successor is a block, and try's successor is a catch block
+         */
       } else if (
-        node.valueRange === null ||
+        node == null ||
         terminal.kind === "ternary" ||
         terminal.kind === "logical" ||
         terminal.kind === "optional"
       ) {
         /**
-         * Create a new scope node whenever we transition from block scope -> value scope.
+         * Create a new node whenever we transition from non-value -> value block.
          *
          * For compatibility with the previous ReactiveFunction-based scope merging logic,
          * we also create new scope nodes for ternary, logical, and optional terminals.
-         * However, inside value blocks we always store a range (valueRange) that is the
+         * Inside value blocks we always store a range (valueRange) that is the
          * start/end instruction ids at the nearest parent block scope level, so that
          * scopes inside the value blocks can be extended to align with block scope
          * instructions.
          */
-        const childNode = {
-          kind: "node",
-          id: terminal.id,
-          children: [],
-          valueRange: null,
-        } as BlockNode;
-        if (node.valueRange === null) {
-          // Transition from block->value scope, derive the outer block scope range
+        let valueRange: MutableRange;
+        if (node == null) {
+          // Transition from block->value block, derive the outer block range
           CompilerError.invariant(fallthrough !== null, {
             reason: `Expected a fallthrough for value block`,
             loc: terminal.loc,
@@ -216,46 +223,50 @@ export function alignReactiveScopesToBlockScopesHIR(fn: HIRFunction): void {
           const nextId =
             fallthroughBlock.instructions[0]?.id ??
             fallthroughBlock.terminal.id;
-          childNode.valueRange = {
+          valueRange = {
             start: terminal.id,
             end: nextId,
           };
         } else {
           // else value->value transition, reuse the range
-          childNode.valueRange = node.valueRange;
+          valueRange = node.valueRange;
         }
-        node.children.push(childNode);
-        blockNodes.set(successor, childNode);
+        const childNode: ValueBlockNode = {
+          kind: "node",
+          id: terminal.id,
+          children: [],
+          valueRange,
+        };
+        node?.children.push(childNode);
+        valueBlockNodes.set(successor, childNode);
       } else {
         // this is a value -> value block transition, reuse the node
-        blockNodes.set(successor, node);
+        valueBlockNodes.set(successor, node);
       }
       return successor;
     });
   }
-
-  // console.log(_debug(rootNode));
 }
 
-type BlockNode = {
+type ValueBlockNode = {
   kind: "node";
   id: InstructionId;
-  valueRange: MutableRange | null;
-  children: Array<BlockNode | ReactiveScopeNode>;
+  valueRange: MutableRange;
+  children: Array<ValueBlockNode | ReactiveScopeNode>;
 };
 type ReactiveScopeNode = {
   kind: "scope";
   id: InstructionId;
   scope: ReactiveScope;
 };
 
-function _debug(node: BlockNode): string {
+function _debug(node: ValueBlockNode): string {
   const buf: Array<string> = [];
   _printNode(node, buf, 0);
   return buf.join("\n");
 }
 function _printNode(
-  node: BlockNode | ReactiveScopeNode,
+  node: ValueBlockNode | ReactiveScopeNode,
   out: Array<string>,
   depth: number = 0
 ): void {
@@ -265,10 +276,7 @@ function _printNode(
       `${prefix}[${node.id}] @${node.scope.id} [${node.scope.range.start}:${node.scope.range.end}]`
     );
   } else {
-    let range =
-      node.valueRange !== null
-        ? ` [${node.valueRange.start}:${node.valueRange.end}]`
-        : "";
+    let range = ` (range=[${node.valueRange.start}:${node.valueRange.end}])`;
     out.push(`${prefix}[${node.id}] node${range} [`);
     for (const child of node.children) {
       _printNode(child, out, depth + 1);

compiler/packages/babel-plugin-react-compiler/src/Utils/utils.ts

@@ -45,6 +45,17 @@ export function retainWhere<T>(
   array.length = writeIndex;
 }
 
+export function retainWhere_Set<T>(
+  items: Set<T>,
+  predicate: (item: T) => boolean
+): void {
+  for (const item of items) {
+    if (!predicate(item)) {
+      items.delete(item);
+    }
+  }
+}
+
 export function getOrInsertWith<U, V>(
   m: Map<U, V>,
   key: U,

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/align-scope-starts-within-cond.expect.md

@@ -0,0 +1,76 @@
+
+## Input
+
+```javascript
+import { mutate } from "shared-runtime";
+
+/**
+ * Similar fixture to `align-scopes-nested-block-structure`, but
+ * a simpler case.
+ */
+function useFoo(cond) {
+  let s = null;
+  if (cond) {
+    s = {};
+  } else {
+    return null;
+  }
+  mutate(s);
+  return s;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [true],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { mutate } from "shared-runtime";
+
+/**
+ * Similar fixture to `align-scopes-nested-block-structure`, but
+ * a simpler case.
+ */
+function useFoo(cond) {
+  const $ = _c(3);
+  let s;
+  let t0;
+  if ($[0] !== cond) {
+    t0 = Symbol.for("react.early_return_sentinel");
+    bb0: {
+      if (cond) {
+        s = {};
+      } else {
+        t0 = null;
+        break bb0;
+      }
+
+      mutate(s);
+    }
+    $[0] = cond;
+    $[1] = t0;
+    $[2] = s;
+  } else {
+    t0 = $[1];
+    s = $[2];
+  }
+  if (t0 !== Symbol.for("react.early_return_sentinel")) {
+    return t0;
+  }
+  return s;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [true],
+};
+
+```
+      
+### Eval output
+(kind: ok) {"wat0":"joe"}