Exhaustive switches for some terminal handling

Updates two points in the compiler that were easy to miss when adding new 
terminals: 

* HIRBuilder's `removeUnreachableFallthroughs()` nulls out unreachable 
fallthroughs, but this had a non-exhaustive `if` statement. It now uses a helper 
function which internally has an exhaustive switch. 

* LeaveSSA needs to schedule block fallthroughs, but had a non-exhaustive `if` 
statement. It also uses a helper function which internally has an exhaustive 
switch. 

cc @poteto since you ran into this (ie the compiler not alerting you to update 
these places) w your diffs.

compiler/forget/src/HIR/HIRBuilder.ts

@@ -28,7 +28,11 @@ import {
   Terminal,
 } from "./HIR";
 import { printInstruction } from "./PrintHIR";
-import { eachTerminalSuccessor, mapTerminalSuccessors } from "./visitors";
+import {
+  eachTerminalSuccessor,
+  mapOptionalFallthroughs,
+  mapTerminalSuccessors,
+} from "./visitors";
 
 // *******************************************************************************************
 // *******************************************************************************************
@@ -535,19 +539,13 @@ export function removeUnreachableFallthroughs(func: HIR): void {
 
   // Cleanup any fallthrough blocks that weren't visited
   for (const [_, block] of func.blocks) {
-    if (
-      block.terminal.kind === "if" ||
-      block.terminal.kind === "switch" ||
-      block.terminal.kind === "while" ||
-      block.terminal.kind === "label"
-    ) {
-      if (
-        block.terminal.fallthrough !== null &&
-        !visited.has(block.terminal.fallthrough)
-      ) {
-        block.terminal.fallthrough = null;
+    mapOptionalFallthroughs(block.terminal, (fallthrough) => {
+      if (visited.has(fallthrough)) {
+        return fallthrough;
+      } else {
+        return null;
       }
-    }
+    });
   }
 }
 

compiler/forget/src/HIR/visitors.ts

@@ -701,6 +701,112 @@ export function mapTerminalSuccessors(
   }
 }
 
+/**
+ * Helper to get a terminal's fallthrough. The main reason to extract this as a helper
+ * function is to ensure that we use an exhaustive switch to ensure that we add new terminal
+ * variants as appropriate.
+ */
+export function terminalFallthrough(terminal: Terminal): BlockId | null {
+  switch (terminal.kind) {
+    case "branch":
+    case "goto":
+    case "return":
+    case "throw":
+    case "unsupported": {
+      return null;
+    }
+    case "do-while":
+    case "for-of":
+    case "for":
+    case "if":
+    case "label":
+    case "logical":
+    case "optional-call":
+    case "switch":
+    case "ternary":
+    case "while": {
+      return terminal.fallthrough;
+    }
+    default: {
+      assertExhaustive(
+        terminal,
+        `Unexpected terminal kind '${(terminal as any).kind}'`
+      );
+    }
+  }
+}
+
+export function mapOptionalFallthroughs(
+  terminal: Terminal,
+  fn: (block: BlockId) => BlockId | null
+): void {
+  switch (terminal.kind) {
+    case "branch":
+    case "goto":
+    case "return":
+    case "throw":
+    case "unsupported": {
+      return;
+    }
+    // NOTE: TypeScript has a bug where it does not correctly model properties whose values are
+    // non-null in some cases and nullable in other cases, if those cases are joined together.
+    // Thus we use one block per case here to ensure that any changes to the types will cause
+    // a compiler error.
+    case "do-while": {
+      const _: BlockId = terminal.fallthrough;
+      break;
+    }
+    case "for-of": {
+      const _: BlockId = terminal.fallthrough;
+      break;
+    }
+    case "for": {
+      const _: BlockId = terminal.fallthrough;
+      break;
+    }
+    case "logical": {
+      const _: BlockId = terminal.fallthrough;
+      break;
+    }
+    case "optional-call": {
+      const _: BlockId = terminal.fallthrough;
+      break;
+    }
+    case "ternary": {
+      const _: BlockId = terminal.fallthrough;
+      break;
+    }
+    case "while": {
+      const _: BlockId = terminal.fallthrough;
+      break;
+    }
+    case "switch": {
+      if (terminal.fallthrough !== null) {
+        terminal.fallthrough = fn(terminal.fallthrough);
+      }
+      break;
+    }
+    case "if": {
+      if (terminal.fallthrough !== null) {
+        terminal.fallthrough = fn(terminal.fallthrough);
+      }
+      break;
+    }
+    case "label": {
+      if (terminal.fallthrough !== null) {
+        terminal.fallthrough = fn(terminal.fallthrough);
+      }
+      break;
+    }
+    default: {
+      assertExhaustive(
+        terminal,
+        `Unexpected terminal kind '${(terminal as any).kind}'`
+      );
+    }
+  }
+}
+
 /**
  * Iterates over the successor block ids of the provided terminal. The function is called
  * specifically for the successors that define the standard control flow, and not

compiler/forget/src/SSA/LeaveSSA.ts

@@ -28,6 +28,7 @@ import {
   eachInstructionValueOperand,
   eachPatternOperand,
   eachTerminalOperand,
+  terminalFallthrough,
 } from "../HIR/visitors";
 
 /**
@@ -294,17 +295,9 @@ export function leaveSSA(fn: HIRFunction): void {
         }
       }
     }
-    if (
-      (terminal.kind === "if" ||
-        terminal.kind === "switch" ||
-        terminal.kind === "while" ||
-        terminal.kind === "do-while" ||
-        terminal.kind === "for" ||
-        terminal.kind === "for-of" ||
-        terminal.kind === "label") &&
-      terminal.fallthrough !== null
-    ) {
-      const fallthrough = fn.body.blocks.get(terminal.fallthrough)!;
+    const fallthroughId = terminalFallthrough(terminal);
+    if (fallthroughId !== null) {
+      const fallthrough = fn.body.blocks.get(fallthroughId)!;
       pushPhis(fallthrough);
     }
     if (terminal.kind === "while" || terminal.kind === "for") {
@@ -345,14 +338,6 @@ export function leaveSSA(fn: HIRFunction): void {
         pushPhis(update);
       }
     }
-    if (
-      terminal.kind === "logical" ||
-      terminal.kind === "ternary" ||
-      terminal.kind === "optional-call"
-    ) {
-      const fallthrough = fn.body.blocks.get(terminal.fallthrough)!;
-      pushPhis(fallthrough);
-    }
 
     for (const { phi, block: phiBlock } of reassignmentPhis) {
       // In some cases one of the phi operands can be defined *before* the let binding