Inst rewrite in ADCE

crates/codegen/src/optim/adce.rs

@@ -6,23 +6,22 @@ use std::collections::BTreeSet;
 use crate::post_domtree::{PDFSet, PDTIdom, PostDomTree};
 
 use sonatina_ir::{
-    func_cursor::{CursorLocation, FuncCursor, InsnInserter},
-    inst::InsnData,
-    BlockId, Function, Insn,
+    func_cursor::{CursorLocation, FuncCursor, InstInserter},
+    BlockId, Function, InstId,
 };
 
 pub struct AdceSolver {
-    live_insns: SecondaryMap<Insn, bool>,
+    live_insts: SecondaryMap<InstId, bool>,
     live_blocks: SecondaryMap<BlockId, bool>,
     empty_blocks: BTreeSet<BlockId>,
     post_domtree: PostDomTree,
-    worklist: Vec<Insn>,
+    worklist: Vec<InstId>,
 }
 
 impl AdceSolver {
     pub fn new() -> Self {
         Self {
-            live_insns: SecondaryMap::default(),
+            live_insts: SecondaryMap::default(),
             live_blocks: SecondaryMap::default(),
             empty_blocks: BTreeSet::default(),
             post_domtree: PostDomTree::default(),
@@ -31,7 +30,7 @@ impl AdceSolver {
     }
 
     pub fn clear(&mut self) {
-        self.live_insns.clear();
+        self.live_insts.clear();
         self.live_blocks.clear();
         self.empty_blocks.clear();
         self.post_domtree.clear();
@@ -42,7 +41,7 @@ impl AdceSolver {
         while self.run_dce(func) {}
     }
 
-    /// Returns `true` if branch insn is modified while dead code elimination.
+    /// Returns `true` if branch inst is modified while dead code elimination.
     fn run_dce(&mut self, func: &mut Function) -> bool {
         self.clear();
 
@@ -56,15 +55,15 @@ impl AdceSolver {
         }
 
         for block in func.layout.iter_block() {
-            for insn in func.layout.iter_inst(block) {
-                if func.dfg.has_side_effect(insn) {
-                    self.mark_insn(func, insn);
+            for inst in func.layout.iter_inst(block) {
+                if func.dfg.has_side_effect(inst) {
+                    self.mark_inst(func, inst);
                 }
             }
         }
 
-        while let Some(insn) = self.worklist.pop() {
-            self.mark_by_insn(func, insn, &pdf_set);
+        while let Some(inst) = self.worklist.pop() {
+            self.mark_by_inst(func, inst, &pdf_set);
         }
 
         self.eliminate_dead_code(func)
@@ -80,48 +79,49 @@ impl AdceSolver {
         false
     }
 
-    fn mark_insn(&mut self, func: &Function, insn: Insn) {
-        let mut mark_insn = |insn, block| {
-            if !self.does_insn_live(insn) {
-                self.live_insns[insn] = true;
-                self.worklist.push(insn);
+    fn mark_inst(&mut self, func: &Function, inst: InstId) {
+        let mut mark_inst = |inst, block| {
+            if !self.does_inst_live(inst) {
+                self.live_insts[inst] = true;
+                self.worklist.push(inst);
                 self.mark_block(block);
                 true
             } else {
                 false
             }
         };
 
-        let insn_block = func.layout.inst_block(insn);
-        if mark_insn(insn, insn_block) {
-            let last_insn = func.layout.last_inst_of(insn_block).unwrap();
-            mark_insn(last_insn, insn_block);
+        let inst_block = func.layout.inst_block(inst);
+        if mark_inst(inst, inst_block) {
+            let last_inst = func.layout.last_inst_of(inst_block).unwrap();
+            mark_inst(last_inst, inst_block);
         }
     }
 
     fn mark_block(&mut self, block: BlockId) {
         self.live_blocks[block] = true;
     }
 
-    fn does_insn_live(&self, insn: Insn) -> bool {
-        self.live_insns[insn]
+    fn does_inst_live(&self, inst: InstId) -> bool {
+        self.live_insts[inst]
     }
 
     fn does_block_live(&self, block: BlockId) -> bool {
         self.live_blocks[block]
     }
 
-    fn mark_by_insn(&mut self, func: &Function, insn: Insn, pdf_set: &PDFSet) {
-        for &value in func.dfg.insn_args(insn) {
-            if let Some(value_insn) = func.dfg.value_insn(value) {
-                self.mark_insn(func, value_insn);
+    fn mark_by_inst(&mut self, func: &Function, inst_id: InstId, pdf_set: &PDFSet) {
+        let inst = func.dfg.inst(inst_id);
+        inst.visit_values(&mut |value| {
+            if let Some(value_inst) = func.dfg.value_inst(value) {
+                self.mark_inst(func, value_inst);
             }
-        }
+        });
 
-        let insn_block = func.layout.inst_block(insn);
-        for &block in pdf_set.frontiers(insn_block) {
-            if let Some(last_insn) = func.layout.last_inst_of(block) {
-                self.mark_insn(func, last_insn)
+        let inst_block = func.layout.inst_block(inst_id);
+        for &block in pdf_set.frontiers(inst_block) {
+            if let Some(last_inst) = func.layout.last_inst_of(block) {
+                self.mark_inst(func, last_inst)
             }
         }
     }
@@ -133,11 +133,11 @@ impl AdceSolver {
             return false;
         };
 
-        let mut inserter = InsnInserter::at_location(CursorLocation::BlockTop(entry));
+        let mut inserter = InstInserter::at_location(CursorLocation::BlockTop(entry));
         loop {
             match inserter.loc() {
-                CursorLocation::At(insn) => {
-                    if self.does_insn_live(insn) {
+                CursorLocation::At(inst) => {
+                    if self.does_inst_live(inst) {
                         inserter.proceed(func);
                     } else {
                         inserter.remove_inst(func)
@@ -160,15 +160,15 @@ impl AdceSolver {
             }
         }
 
-        // Modify branch insns to remove unreachable edges.
+        // Modify branch insts to remove unreachable edges.
         inserter.set_to_entry(func);
-        let mut br_insn_modified = false;
+        let mut br_inst_modified = false;
         while let Some(block) = inserter.block(func) {
-            br_insn_modified |= self.modify_branch(func, &mut inserter, block);
+            br_inst_modified |= self.modify_branch(func, &mut inserter, block);
             inserter.proceed_block(func);
         }
 
-        br_insn_modified
+        br_inst_modified
     }
 
     fn living_post_dom(&self, mut block: BlockId) -> Option<BlockId> {
@@ -182,20 +182,24 @@ impl AdceSolver {
         }
     }
 
-    /// Returns `true` if branch insn is modified.
+    /// Returns `true` if branch inst is modified.
     fn modify_branch(
         &self,
         func: &mut Function,
-        inserter: &mut InsnInserter,
+        inserter: &mut InstInserter,
         block: BlockId,
     ) -> bool {
-        let last_insn = match func.layout.last_inst_of(block) {
-            Some(insn) => insn,
+        let last_inst = match func.layout.last_inst_of(block) {
+            Some(inst) => inst,
             None => return false,
         };
-        inserter.set_location(CursorLocation::At(last_insn));
+        inserter.set_location(CursorLocation::At(last_inst));
 
-        let dests: Vec<_> = func.dfg.analyze_branch(last_insn).iter_dests().collect();
+        let dests: Vec<_> = func
+            .dfg
+            .branch_info(last_inst)
+            .map(|bi| bi.iter_dests().collect())
+            .unwrap_or_default();
 
         let mut changed = false;
         for dest in dests {
@@ -206,25 +210,29 @@ impl AdceSolver {
             match self.living_post_dom(dest) {
                 // If the destination is dead but its post dominator is living, then change the
                 // destination to the post dominator.
-                Some(postdom) => func.dfg.rewrite_branch_dest(last_insn, dest, postdom),
+                Some(postdom) => func.dfg.rewrite_branch_dest(last_inst, dest, postdom),
 
                 // If the block doesn't have post dominator, then remove the dest.
                 None => {
-                    func.dfg.remove_branch_dest(last_insn, dest);
+                    func.dfg.remove_branch_dest(last_inst, dest);
                 }
             }
 
             changed = true;
         }
 
-        // Turn branch insn to `jump` if all dests is the same.
-        let branch_info = func.dfg.analyze_branch(last_insn);
-        if branch_info.dests_num() > 1 {
+        // Turn branch inst to `jump` if all dests is the same.
+        let Some(branch_info) = func.dfg.branch_info(last_inst) else {
+            return changed;
+        };
+        if branch_info.num_dests() > 1 {
             let mut branch_dests = branch_info.iter_dests();
             let first_dest = branch_dests.next().unwrap();
             if branch_dests.all(|dest| dest == first_dest) {
                 changed = true;
-                inserter.replace(func, InsnData::jump(first_dest));
+                let jump = func.dfg.make_jump(first_dest);
+                drop(branch_dests);
+                inserter.replace(func, jump);
             }
         }
 

crates/ir/src/dfg.rs

@@ -255,6 +255,50 @@ impl DataFlowGraph {
         self.branch_info_mut(inst)
             .map(|mut bi| bi.rewrite_branch_dest(from, to));
     }
+
+    pub fn remove_branch_dest(&mut self, inst: InstId, dest: BlockId) {
+        let Some(bi) = self.branch_info_mut(inst) else {
+            panic!("not a branch");
+        };
+
+        match bi {
+            BranchInfoMut::Jump(_) => panic!("can't remove destination from `Jump` insn"),
+
+            BranchInfoMut::Br(br) => {
+                let remain = if *br.z_dest() == dest {
+                    *br.nz_dest()
+                } else if *br.nz_dest() == dest {
+                    *br.z_dest()
+                } else {
+                    panic!("no dests found in the branch destination")
+                };
+
+                let cond = *br.cond();
+                self.users[cond].remove(&inst);
+                let jump = self.make_jump(remain);
+                self.insts[inst] = Box::new(jump);
+            }
+
+            BranchInfoMut::BrTable(brt) => {
+                if Some(dest) == *brt.default() {
+                    *brt.default_mut() = None;
+                }
+
+                let (keep, drop) = brt.table().iter().copied().partition(|(_, b)| *b != dest);
+                *brt.table_mut() = keep;
+                for (val, _) in drop {
+                    self.users[val].remove(&inst);
+                }
+
+                let bi = self.branch_info(inst).unwrap();
+                if bi.num_dests() == 1 {
+                    let remain = bi.iter_dests().next().unwrap();
+                    let jump = self.make_jump(remain);
+                    self.insts[inst] = Box::new(jump);
+                }
+            }
+        }
+    }
 }
 
 /// An opaque reference to [`Block`]