[rust] Use new block helper for more passes

compiler/forget/crates/forget_hir/src/function.rs

@@ -62,7 +62,7 @@ impl<'a> Blocks<'a> {
         self.data.keys().cloned().collect()
     }
 
-    pub fn take(&mut self, id: BlockId) -> Box<BasicBlock<'a>> {
+    pub fn remove(&mut self, id: BlockId) -> Box<BasicBlock<'a>> {
         self.data.remove(&id).unwrap().unwrap()
     }
 
@@ -146,7 +146,26 @@ impl<'blocks, 'a> BlockRewriter<'blocks, 'a> {
         }
     }
 
-    pub fn each_block<F>(&mut self, mut f: F) -> Result<(), Diagnostic>
+    pub fn each_block<F>(&mut self, mut f: F) -> ()
+    where
+        F: FnMut(Box<BasicBlock<'a>>, &mut Self) -> BlockRewriterAction<'a>,
+    {
+        let keys = self.blocks.block_ids();
+        for block_id in keys {
+            self.current = block_id;
+            let block = self.blocks.data.get_mut(&block_id).unwrap().take().unwrap();
+            match f(block, self) {
+                BlockRewriterAction::Keep(block) => {
+                    self.blocks.data.insert(block_id, Some(block));
+                }
+                BlockRewriterAction::Remove => {
+                    self.blocks.data.remove(&block_id);
+                }
+            }
+        }
+    }
+
+    pub fn try_each_block<F>(&mut self, mut f: F) -> Result<(), Diagnostic>
     where
         F: FnMut(Box<BasicBlock<'a>>, &mut Self) -> Result<BlockRewriterAction<'a>, Diagnostic>,
     {
@@ -159,13 +178,18 @@ impl<'blocks, 'a> BlockRewriter<'blocks, 'a> {
                     self.blocks.data.insert(block_id, Some(block));
                 }
                 BlockRewriterAction::Remove => {
-                    // nothing to do, already removed from the blocks
+                    self.blocks.data.remove(&block_id);
                 }
             }
         }
         Ok(())
     }
 
+    pub fn contains(&self, block_id: BlockId) -> bool {
+        assert_ne!(block_id, self.current);
+        self.blocks.data.contains_key(&block_id)
+    }
+
     pub fn block(&self, block_id: BlockId) -> &BasicBlock<'a> {
         assert_ne!(block_id, self.current);
         self.blocks.block(block_id)

compiler/forget/crates/forget_hir/src/initialize.rs

@@ -3,7 +3,10 @@ use std::collections::HashSet;
 use forget_diagnostics::{invariant, Diagnostic};
 use thiserror::Error;
 
-use crate::{BlockId, Blocks, GotoKind, GotoTerminal, InstructionIdGenerator, TerminalValue, HIR};
+use crate::{
+    BlockId, BlockRewriter, BlockRewriterAction, Blocks, GotoKind, GotoTerminal,
+    InstructionIdGenerator, TerminalValue, HIR,
+};
 
 /// Runs a variety of passes to put the HIR in canonical form. This should be called
 /// after initial HIR construction and after any transformations that change the
@@ -65,7 +68,7 @@ pub fn reverse_postorder_blocks<'a>(hir: &mut HIR<'a>) {
     // NOTE: could consider sorting the blocks in-place by key
     let mut blocks = Blocks::with_capacity(hir.blocks.len());
     for id in postorder.iter().rev().cloned() {
-        blocks.insert(hir.blocks.take(id));
+        blocks.insert(hir.blocks.remove(id));
     }
 
     hir.blocks = blocks;
@@ -73,52 +76,49 @@ pub fn reverse_postorder_blocks<'a>(hir: &mut HIR<'a>) {
 
 /// Prunes ForTerminal.update values (sets to None) if they are unreachable
 pub fn remove_unreachable_for_updates<'a>(hir: &mut HIR<'a>) {
-    let block_ids = hir.blocks.block_ids();
-
-    for block in hir.blocks.iter_mut() {
+    BlockRewriter::new(&mut hir.blocks, hir.entry).each_block(|mut block, rewriter| {
         if let TerminalValue::For(terminal) = &mut block.terminal.value {
             if let Some(update) = terminal.update {
-                if !block_ids.contains(&update) {
+                if !rewriter.contains(update) {
                     terminal.update = None;
                 }
             }
         }
-    }
+        BlockRewriterAction::Keep(block)
+    });
 }
 
 /// Prunes unreachable fallthrough values, setting them to None if the referenced
 /// block was not otherwise reachable.
 pub fn remove_unreachable_fallthroughs<'a>(hir: &mut HIR<'a>) {
-    let block_ids = hir.blocks.block_ids();
-
-    for block in hir.blocks.iter_mut() {
+    BlockRewriter::new(&mut hir.blocks, hir.entry).each_block(|mut block, rewriter| {
         block
             .terminal
             .value
             .map_optional_fallthroughs(|fallthrough| {
-                if block_ids.contains(&fallthrough) {
+                if rewriter.contains(fallthrough) {
                     Some(fallthrough)
                 } else {
                     None
                 }
-            })
-    }
+            });
+        BlockRewriterAction::Keep(block)
+    });
 }
 
 /// Rewrites DoWhile statements into Gotos if the test block is not reachable
 pub fn remove_unreachable_do_while_statements<'a>(hir: &mut HIR<'a>) {
-    let block_ids = hir.blocks.block_ids();
-
-    for block in hir.blocks.iter_mut() {
+    BlockRewriter::new(&mut hir.blocks, hir.entry).each_block(|mut block, rewriter| {
         if let TerminalValue::DoWhile(terminal) = &mut block.terminal.value {
-            if !block_ids.contains(&terminal.test) {
+            if !rewriter.contains(terminal.test) {
                 block.terminal.value = TerminalValue::Goto(GotoTerminal {
                     block: terminal.body,
                     kind: GotoKind::Break,
                 });
             }
         }
-    }
+        BlockRewriterAction::Keep(block)
+    });
 }
 
 /// Updates the instruction ids for all instructions and blocks

compiler/forget/crates/forget_hir/src/merge_consecutive_blocks.rs

@@ -27,7 +27,7 @@ pub fn merge_consecutive_blocks<'a>(
     let mut rewriter = BlockRewriter::new(blocks, fun.body.entry);
     let mut has_changes = false;
 
-    rewriter.each_block(|mut block, rewriter| {
+    rewriter.try_each_block(|mut block, rewriter| {
         let block_id = block.id;
         // Visit instructions to merge blocks within function expressions
         for instr_ix in &block.instructions {