Rewrite BranchInfo by inst_prop

crates/ir/src/builder/func_builder.rs

@@ -1,15 +1,14 @@
+use super::{
+    ssa::{SsaBuilder, Variable},
+    ModuleBuilder,
+};
 use crate::{
     func_cursor::{CursorLocation, FuncCursor},
-    inst::control_flow::BranchInfo,
+    inst::control_flow::{Branch, BranchInfo},
     module::FuncRef,
     BlockId, Function, GlobalVariable, Immediate, Inst, InstDowncast, InstId, Type, Value, ValueId,
 };
 
-use super::{
-    ssa::{SsaBuilder, Variable},
-    ModuleBuilder,
-};
-
 pub struct FunctionBuilder<C> {
     pub module_builder: ModuleBuilder,
     pub func: Function,
@@ -107,11 +106,12 @@ where
             .declare_struct_type(name, fields, packed)
     }
 
-    /// Inserts an instruction into the current position and returns a `ValueId` for
-    /// the result.
+    /// Inserts an instruction into the current position and returns a `ValueId`
+    /// for the result.
     ///
     /// # Parameters
-    /// - `inst`: The instruction to insert, which must implement the `Inst` trait.
+    /// - `inst`: The instruction to insert, which must implement the `Inst`
+    ///   trait.
     /// - `ret_ty`: The return type of the instruction. A result value will be
     ///   created with this type and associated with the instruction.
     ///
@@ -148,7 +148,8 @@ where
     /// Please refer to [`insert_inst`] if the instruction has a result.
     ///
     /// # Parameters
-    /// - `inst`: The instruction to insert, which must implement the `Inst` trait.
+    /// - `inst`: The instruction to insert, which must implement the `Inst`
+    ///   trait.
     pub fn insert_inst_no_result<I: Inst>(&mut self, inst: I) {
         let inst_id = self.cursor.insert_inst_data(&mut self.func, inst);
         self.append_pred(inst_id);
@@ -209,14 +210,15 @@ where
         };
 
         let current_block = self.cursor.block(&self.func).unwrap();
-        for dest in branch_info.iter_dests() {
+        for dest in branch_info.dests() {
             self.ssa_builder.append_pred(dest, current_block)
         }
     }
 }
 
 #[cfg(test)]
 mod tests {
+    use super::{super::test_util::*, *};
     use crate::{
         inst::{
             arith::{Add, Mul, Sub},
@@ -226,8 +228,6 @@ mod tests {
         isa::Isa,
     };
 
-    use super::{super::test_util::*, *};
-
     #[test]
     fn entry_block() {
         let (evm, mut builder) = test_func_builder(&[], Type::Void);

crates/ir/src/cfg.rs

@@ -2,7 +2,7 @@ use std::collections::BTreeSet;
 
 use cranelift_entity::{packed_option::PackedOption, SecondaryMap};
 
-use crate::{BlockId, Function};
+use crate::{inst::control_flow::Branch, BlockId, Function};
 
 #[derive(Default, Debug, Clone, PartialEq, Eq)]
 pub struct ControlFlowGraph {
@@ -86,7 +86,7 @@ impl ControlFlowGraph {
             return;
         };
 
-        for dest in branch_info.iter_dests() {
+        for dest in branch_info.dests() {
             let block = func.layout.inst_block(inst);
             self.add_edge(block, dest);
         }

crates/ir/src/dfg.rs

@@ -4,18 +4,17 @@ use std::{collections::BTreeSet, fmt};
 use cranelift_entity::{entity_impl, packed_option::PackedOption, PrimaryMap, SecondaryMap};
 use rustc_hash::FxHashMap;
 
+use super::{Immediate, Type, Value, ValueId};
 use crate::{
     inst::{
-        control_flow::{self, BranchInfo, BranchInfoMut, Jump, Phi},
+        control_flow::{self, Branch, BranchInfo, BranchInfoMut, Jump, Phi},
         InstId,
     },
     ir_writer::DisplayWithFunc,
     module::ModuleCtx,
     Function, GlobalVariable, Inst, InstDowncast, InstDowncastMut, InstSetBase,
 };
 
-use super::{Immediate, Type, Value, ValueId};
-
 pub struct DataFlowGraph {
     pub ctx: ModuleCtx,
     #[doc(hidden)]
@@ -299,9 +298,9 @@ impl DataFlowGraph {
                 }
 
                 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);
+                let dests = bi.dests();
+                if dests.len() == 1 {
+                    let jump = self.make_jump(dests[0]);
                     self.insts[inst] = Box::new(jump);
                 }
             }
@@ -321,8 +320,8 @@ impl DisplayWithFunc for BlockId {
 }
 
 /// A block data definition.
-/// A Block data doesn't hold any information for layout of a program. It is managed by
-/// [`super::layout::Layout`].
+/// A Block data doesn't hold any information for layout of a program. It is
+/// managed by [`super::layout::Layout`].
 #[derive(Debug, Clone, Default)]
 pub struct Block {}
 

crates/ir/src/inst/control_flow.rs

@@ -1,16 +1,15 @@
 use std::fmt;
 
-use macros::Inst;
+use macros::{inst_prop, Inst};
 use smallvec::SmallVec;
 
+use super::InstDowncastMut;
 use crate::{
     ir_writer::{DisplayWithFunc, DisplayableWithFunc},
     module::FuncRef,
     BlockId, Function, Inst, Type, ValueId,
 };
 
-use super::{InstDowncast, InstDowncastMut};
-
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Inst)]
 #[inst(terminator)]
 pub struct Jump {
@@ -92,6 +91,13 @@ impl Phi {
         let pos = self.args.iter().position(|(_, b)| *b == block)?;
         Some(self.args.remove(pos).0)
     }
+
+    pub fn retain<F>(&mut self, mut f: F)
+    where
+        F: FnMut(BlockId) -> bool,
+    {
+        self.args.retain(|(_, block)| f(*block))
+    }
 }
 impl DisplayWithFunc for Phi {
     fn fmt(&self, func: &Function, formatter: &mut fmt::Formatter) -> fmt::Result {
@@ -148,77 +154,28 @@ impl DisplayWithFunc for Return {
     }
 }
 
-#[derive(Clone, Copy)]
-pub enum BranchInfo<'i> {
-    Jump(&'i Jump),
-    Br(&'i Br),
-    BrTable(&'i BrTable),
-}
-
-impl<'i> BranchInfo<'i> {
-    pub fn iter_dests(self) -> impl Iterator<Item = BlockId> + 'i {
-        BranchDestIter {
-            branch_info: self,
-            idx: 0,
-        }
-    }
+#[inst_prop(Subset = "BranchInfo")]
+pub trait Branch {
+    fn dests(&self) -> Vec<BlockId>;
 
-    pub fn num_dests(self) -> usize {
-        match self {
-            Self::Jump(_) => 1,
-            Self::Br(_) => 2,
-            Self::BrTable(brt) => {
-                let ent_len = brt.table.len();
-                if brt.default.is_some() {
-                    ent_len + 1
-                } else {
-                    ent_len
-                }
-            }
-        }
-    }
+    type Members = (Jump, Br, BrTable);
 }
 
-impl<'i> InstDowncast for BranchInfo<'i> {
-    fn downcast(isb: &dyn crate::InstSetBase, inst: &dyn super::Inst) -> Option<Self> {
-        InstDowncast::map(isb, inst, BranchInfo::Jump)
-            .or_else(|| InstDowncast::map(isb, inst, BranchInfo::Br))
-            .or_else(|| InstDowncast::map(isb, inst, BranchInfo::BrTable))
+impl Branch for Jump {
+    fn dests(&self) -> Vec<BlockId> {
+        todo!()
     }
 }
 
-#[derive(Clone, Copy)]
-struct BranchDestIter<'i> {
-    branch_info: BranchInfo<'i>,
-    idx: usize,
+impl Branch for Br {
+    fn dests(&self) -> Vec<BlockId> {
+        todo!()
+    }
 }
 
-impl<'i> Iterator for BranchDestIter<'i> {
-    type Item = BlockId;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        let idx = self.idx;
-        if idx >= self.branch_info.num_dests() {
-            return None;
-        }
-        self.idx += 1;
-
-        match self.branch_info {
-            BranchInfo::Jump(j) => Some(j.dest),
-
-            BranchInfo::Br(br) => {
-                let dest = if idx == 0 { br.z_dest } else { br.nz_dest };
-                Some(dest)
-            }
-
-            BranchInfo::BrTable(brt) => {
-                if idx < brt.table.len() {
-                    Some(brt.table[idx].1)
-                } else {
-                    brt.default
-                }
-            }
-        }
+impl Branch for BrTable {
+    fn dests(&self) -> Vec<BlockId> {
+        todo!()
     }
 }
 

crates/ir/src/isa/evm.rs

@@ -1,10 +1,9 @@
 use std::sync::LazyLock;
 
-use crate::inst::evm::inst_set::EvmInstSet;
+use sonatina_triple::{Architecture, Chain, EvmVersion, TargetTriple, Version};
 
 use super::{Isa, TypeLayout};
-
-use sonatina_triple::{Architecture, Chain, EvmVersion, TargetTriple, Version};
+use crate::inst::evm::inst_set::EvmInstSet;
 
 #[derive(Debug, Clone, Copy)]
 pub struct Evm {

crates/ir/src/lib.rs

@@ -27,6 +27,7 @@ pub use dfg::{Block, BlockId, DataFlowGraph};
 pub use function::{Function, Signature};
 pub use global_variable::{GlobalVariable, GlobalVariableData};
 pub use graphviz::render_to;
+pub(crate) use inst::ValueVisitable;
 pub use inst::{
     inst_set::{InstSetBase, InstSetExt},
     HasInst, Inst, InstDowncast, InstDowncastMut, InstId,
@@ -37,8 +38,6 @@ pub use module::Module;
 pub use types::Type;
 pub use value::{Immediate, Value, ValueId};
 
-pub(crate) use inst::ValueVisitable;
-
 pub mod prelude {
     pub use crate::{
         inst::{