From e64ef9e1e04f333f0a820cd698ff039a11efd959 Mon Sep 17 00:00:00 2001
From: Lauren Tan <poteto@users.noreply.github.com>
Date: Wed, 9 Nov 2022 16:37:14 -0500
Subject: [PATCH] Extract eachBlockOperand to its own visitor

Follow up for #757:

- Adds a new visitor which iterates over every Place within a BasicBlock -
Remove unused entryBlock binding - Comments
---
 compiler/forget/src/HIR/LeaveSSA.ts | 35 ++++++++++-------------------
 compiler/forget/src/HIR/visitors.ts | 15 ++++++++++++-
 2 files changed, 26 insertions(+), 24 deletions(-)

diff --git a/compiler/forget/src/HIR/LeaveSSA.ts b/compiler/forget/src/HIR/LeaveSSA.ts
index ae1684483a..2fe73c53b8 100644
--- a/compiler/forget/src/HIR/LeaveSSA.ts
+++ b/compiler/forget/src/HIR/LeaveSSA.ts
@@ -1,11 +1,14 @@
-import { invariant } from "../CompilerError";
-import { HIRFunction, Identifier, Place } from "./HIR";
-import { eachInstructionOperand, eachTerminalOperand } from "./visitors";
+import { BasicBlock, HIRFunction, Identifier, Phi, Place } from "./HIR";
+import { eachBlockOperand } from "./visitors";
 
+/**
+ * Leaves SSA form by building up a mapping of SSA'd {@link Identifier}s to their original
+ * {@link Identifier}, then rewriting all {@link Place}s within a {@link BasicBlock} to reference
+ * the original id. This allows us to skip adding instruction copies when removing {@link Phi}s,
+ * while still allowing shadowing to work.
+ */
 export default function leaveSSA(fn: HIRFunction) {
   const ir = fn.body;
-  const entryBlock = ir.blocks.get(ir.entry);
-  invariant(entryBlock, "expected to find the entry basic block");
   const originalIdMap = new Map<
     /* SSA'd id */ Identifier,
     /* original id*/ Identifier
@@ -25,26 +28,12 @@ export default function leaveSSA(fn: HIRFunction) {
     return;
   }
 
-  function tryRewrite(place: Place) {
-    const originalId = originalIdMap.get(place.identifier);
-    if (originalId != null) {
-      place.identifier = originalId;
-    }
-  }
-
   for (const [, block] of ir.blocks) {
-    for (const instr of block.instructions) {
-      // LValues also need to be rewritten as they might be declaring or reassigning an identifier
-      // that was previously SSA'd.
-      if (instr.lvalue != null) {
-        tryRewrite(instr.lvalue.place);
+    for (const place of eachBlockOperand(block)) {
+      const originalId = originalIdMap.get(place.identifier);
+      if (originalId != null) {
+        place.identifier = originalId;
       }
-      for (const place of eachInstructionOperand(instr)) {
-        tryRewrite(place);
-      }
-    }
-    for (const place of eachTerminalOperand(block.terminal)) {
-      tryRewrite(place);
     }
   }
 }
diff --git a/compiler/forget/src/HIR/visitors.ts b/compiler/forget/src/HIR/visitors.ts
index d19bfe512e..72c790efd7 100644
--- a/compiler/forget/src/HIR/visitors.ts
+++ b/compiler/forget/src/HIR/visitors.ts
@@ -6,7 +6,7 @@
  */
 
 import { assertExhaustive } from "../Common/utils";
-import { BlockId, Instruction, Place, Terminal } from "./HIR";
+import { BasicBlock, BlockId, Instruction, Place, Terminal } from "./HIR";
 
 export function* eachInstructionOperand(instr: Instruction): Iterable<Place> {
   const instrValue = instr.value;
@@ -304,3 +304,16 @@ export function* eachTerminalOperand(terminal: Terminal): Iterable<Place> {
     }
   }
 }
+
+/**
+ * Iterates over all {@link Place}s within a {@link BasicBlock}.
+ */
+export function* eachBlockOperand(block: BasicBlock): Iterable<Place> {
+  for (const instr of block.instructions) {
+    yield* eachInstructionOperand(instr);
+    if (instr.lvalue != null) {
+      yield instr.lvalue.place;
+    }
+  }
+  yield* eachTerminalOperand(block.terminal);
+}