From 09bdf3553f9de20ee17fe2e1fc3f6e8ba11efaef Mon Sep 17 00:00:00 2001
From: Sathya Gunasekaran <gsathya.ceg@gmail.com>
Date: Wed, 4 Jan 2023 18:09:24 +0000
Subject: [PATCH] [hir] Remove unreachable fallthroughs separately from shrink

shrink visits all the fallthroughs even if they are unreachable so this isn't
the right place to prune unreachable blocks.

This PR moves pruning into a separate pass.
---
 compiler/forget/src/HIR/HIRBuilder.ts         | 21 ++++++++++-----
 .../fixtures/hir/return-conditional.expect.md | 27 +++++++++++++++++++
 .../fixtures/hir/return-conditional.js        |  7 +++++
 3 files changed, 49 insertions(+), 6 deletions(-)
 create mode 100644 compiler/forget/src/__tests__/fixtures/hir/return-conditional.expect.md
 create mode 100644 compiler/forget/src/__tests__/fixtures/hir/return-conditional.js

diff --git a/compiler/forget/src/HIR/HIRBuilder.ts b/compiler/forget/src/HIR/HIRBuilder.ts
index 5a80f54650..cb9381d347 100644
--- a/compiler/forget/src/HIR/HIRBuilder.ts
+++ b/compiler/forget/src/HIR/HIRBuilder.ts
@@ -162,13 +162,15 @@ export default class HIRBuilder {
       entry: this.#entry,
     };
     logHIR("Build (pre-shrink)", ir);
-    // First reduce indirections and prune unreachable blocks
+    // First reduce indirections
     let shrunk = shrink(ir);
     logHIR("Build (shrunk)", shrunk);
     // then convert to reverse postorder
     const rpo = reversePostorderBlocks(shrunk);
+    removeUnreachableFallthroughs(rpo);
     markInstructionIds(rpo);
     markPredecessors(rpo);
+
     return rpo;
   }
 
@@ -356,7 +358,7 @@ export default class HIRBuilder {
 }
 
 /**
- * Helper to shrink a CFG to eliminate unreachable node and eliminate jump-only blocks.
+ * Helper to shrink a CFG eliminate jump-only blocks.
  */
 function shrink(func: HIR): HIR {
   const gotos = new Map();
@@ -406,8 +408,17 @@ function shrink(func: HIR): HIR {
     });
   }
 
+  return { blocks, entry: func.entry };
+}
+
+function removeUnreachableFallthroughs(func: HIR) {
+  const visited: Set<BlockId> = new Set();
+  for (const [_, block] of func.blocks) {
+    visited.add(block.id);
+  }
+
   // Cleanup any fallthrough blocks that weren't visited
-  for (const block of blocks.values()) {
+  for (const [_, block] of func.blocks) {
     if (
       block.terminal.kind === "if" ||
       block.terminal.kind === "switch" ||
@@ -415,15 +426,13 @@ function shrink(func: HIR): HIR {
     ) {
       if (
         block.terminal.fallthrough !== null &&
-        !blocks.has(block.terminal.fallthrough)
+        !visited.has(block.terminal.fallthrough)
       ) {
         block.terminal.fallthrough = null;
       }
     }
   }
-  return { blocks, entry: func.entry };
 }
-
 /**
  * Converts the graph to reverse-postorder, with predecessor blocks appearing
  * before successors except in the case of back links (ie loops).
diff --git a/compiler/forget/src/__tests__/fixtures/hir/return-conditional.expect.md b/compiler/forget/src/__tests__/fixtures/hir/return-conditional.expect.md
new file mode 100644
index 0000000000..6244ddef97
--- /dev/null
+++ b/compiler/forget/src/__tests__/fixtures/hir/return-conditional.expect.md
@@ -0,0 +1,27 @@
+
+## Input
+
+```javascript
+function foo(a, b) {
+  if (a == null) {
+    return null;
+  } else {
+    return b;
+  }
+}
+
+```
+
+## Code
+
+```javascript
+function foo(a, b) {
+  if (a == null) {
+    return null;
+  } else {
+    return b;
+  }
+}
+
+```
+      
\ No newline at end of file
diff --git a/compiler/forget/src/__tests__/fixtures/hir/return-conditional.js b/compiler/forget/src/__tests__/fixtures/hir/return-conditional.js
new file mode 100644
index 0000000000..81d81e12c1
--- /dev/null
+++ b/compiler/forget/src/__tests__/fixtures/hir/return-conditional.js
@@ -0,0 +1,7 @@
+function foo(a, b) {
+  if (a == null) {
+    return null;
+  } else {
+    return b;
+  }
+}