AlignReactiveScopesToBlockScopes (rewritten against ReactiveFunction)

See the background in #982. This PR reimplements part of InferReactiveScopes, 
aligning reactive scopes to block boundaries, but against ReactiveFunction 
instead of the HIR.

compiler/forget/src/ReactiveScopes/AlignReactiveScopesToBlockScopes.ts

@@ -5,7 +5,18 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import { ReactiveFunction } from "../HIR/HIR";
+import {
+  InstructionId,
+  makeInstructionId,
+  ReactiveBlock,
+  ReactiveFunction,
+  ReactiveScope,
+  ReactiveValueBlock,
+  ScopeId,
+} from "../HIR/HIR";
+import { invariant } from "../Utils/CompilerError";
+import { getInstructionScope } from "./BuildReactiveBlocks";
+import { eachTerminalBlock } from "./visitors";
 
 /**
  * Note: this is the 2nd of 3 passes that determine how to break a function into discrete
@@ -47,4 +58,125 @@ import { ReactiveFunction } from "../HIR/HIR";
  * will be the updated end for that scope).
  */
 
-export function alignReactiveScopesToBlockScopes(fn: ReactiveFunction): void {}
+export function alignReactiveScopesToBlockScopes(fn: ReactiveFunction): void {
+  const context = new Context();
+  context.enter(() => {
+    visitBlock(context, fn.body);
+  });
+}
+
+function visitBlock(context: Context, block: ReactiveBlock): void {
+  for (const stmt of block) {
+    switch (stmt.kind) {
+      case "instruction": {
+        context.visitId(stmt.instruction.id);
+        const scope = getInstructionScope(stmt.instruction);
+        if (scope !== null) {
+          context.visitScope(scope);
+        }
+        break;
+      }
+      case "terminal": {
+        const id = stmt.terminal.id;
+        if (id !== null) {
+          context.visitId(id);
+        }
+        eachTerminalBlock(
+          stmt.terminal,
+          (block) => {
+            context.enter(() => visitBlock(context, block));
+          },
+          (valueBlock) => visitValueBlock(context, valueBlock, id!)
+        );
+        break;
+      }
+      case "scope": {
+        invariant(false, "Expected scopes to be constructed later");
+      }
+    }
+  }
+}
+
+function visitValueBlock(
+  context: Context,
+  block: ReactiveValueBlock,
+  start: InstructionId
+): void {
+  for (const stmt of block.instructions) {
+    switch (stmt.kind) {
+      case "instruction": {
+        context.visitId(stmt.instruction.id);
+        const scope = getInstructionScope(stmt.instruction);
+        if (scope !== null) {
+          scope.range.start = makeInstructionId(
+            Math.min(start, scope.range.start)
+          );
+          context.visitScope(scope);
+        }
+        break;
+      }
+      default: {
+        invariant(false, "Unexpected terminal or scope in value block");
+      }
+    }
+  }
+}
+
+type PendingReactiveScope = { active: boolean; scope: ReactiveScope };
+
+class Context {
+  // For each block scope (outer array) stores a list of ReactiveScopes that start
+  // in that block scope.
+  #blockScopes: Array<{
+    kind: "block" | "value";
+    scopes: Array<PendingReactiveScope>;
+  }> = [];
+
+  // ReactiveScopes whose declaring block scope has ended but may still need to
+  // be "closed" (ie have their range.end be updated). A given scope can be in
+  // blockScopes OR this array but not both.
+  #unclosedScopes: Array<PendingReactiveScope> = [];
+
+  // Set of all scope ids that have been seen so far, regardless of which of
+  // the above data structures they're in, to avoid tracking the same scope twice.
+  #seenScopes: Set<ScopeId> = new Set();
+
+  enter(fn: () => void): void {
+    this.#blockScopes.push({ kind: "block", scopes: [] });
+    fn();
+    const lastScope = this.#blockScopes.pop()!;
+    for (const scope of lastScope.scopes) {
+      if (scope.active) {
+        this.#unclosedScopes.push(scope);
+      }
+    }
+  }
+
+  visitId(id: InstructionId): void {
+    const currentScopes = this.#blockScopes.at(-1)!;
+    if (currentScopes.kind === "value") {
+      return;
+    }
+    const scopes = [...currentScopes.scopes, ...this.#unclosedScopes];
+    for (const pending of scopes) {
+      if (!pending.active) {
+        continue;
+      }
+      if (id >= pending.scope.range.end) {
+        pending.active = false;
+        pending.scope.range.end = id;
+      }
+    }
+  }
+
+  visitScope(scope: ReactiveScope): void {
+    if (!this.#seenScopes.has(scope.id)) {
+      const currentScopes = this.#blockScopes.at(-1)!;
+      this.#seenScopes.add(scope.id);
+      currentScopes.scopes.push({
+        active: true,
+        scope,
+      });
+    }
+  }
+}

compiler/forget/src/ReactiveScopes/BuildReactiveBlocks.ts

@@ -84,13 +84,13 @@ function visitBlock(block: ReactiveBlock): ReactiveBlock {
     }
   }
   while (current.kind === "scope") {
-    invariant(current.scope.range.end === lastId + 1, "Scope ended too soon");
+    // invariant(current.scope.range.end === lastId + 1, "Scope ended too soon");
     current = stack.pop()!;
   }
   return current.instructions;
 }
 
-function getInstructionScope({
+export function getInstructionScope({
   id,
   lvalue,
   value,