added another test for CFA

src/ast/ir.rs

@@ -69,19 +69,27 @@ pub struct FunctionContent {
 
 impl FunctionContent {
     pub fn get_entry_block(&self) -> &Block {
-        self.get_block(self.entry).unwrap()
+        self.get_block(self.entry)
     }
     
     pub fn get_entry_block_mut(&mut self) -> &mut Block {
-        self.get_block_mut(self.entry).unwrap()
+        self.get_block_mut(self.entry)
     }
     
-    pub fn get_block(&self, tag: MuTag) -> Option<&Block> {
-        self.blocks.get(tag)
+    pub fn get_block(&self, tag: MuTag) -> &Block {
+        let ret = self.blocks.get(tag);
+        match ret {
+            Some(b) => b,
+            None => panic!("cannot find block {}", tag)
+        }
     }
     
-    pub fn get_block_mut(&mut self, tag: MuTag) -> Option<&mut Block> {
-        self.blocks.get_mut(tag)
+    pub fn get_block_mut(&mut self, tag: MuTag) -> &mut Block {
+        let ret = self.blocks.get_mut(tag);
+        match ret {
+            Some(b) => b,
+            None => panic!("cannot find block {}", tag)
+        }
     } 
 }
 

src/compiler/passes/control_flow.rs

 use ast::ir::*;
 use ast::inst::Instruction_::*;
+use common::vector_as_str;
 use vm::context::VMContext;
 
 use compiler::CompilerPass;
@@ -22,18 +23,23 @@ fn check_edge_kind(target: MuTag, stack: &Vec<MuTag>) -> EdgeKind {
     }
 }
 
-fn new_edge(cur: MuTag, edge: BlockEdge, func: &mut MuFunction) {
+fn new_edge(cur: MuTag, edge: BlockEdge, stack: &mut Vec<MuTag>, visited: &mut Vec<MuTag>, func: &mut MuFunction) {
     // add current block to target's predecessors
     {
-        let target = func.content.as_mut().unwrap().get_block_mut(edge.target).unwrap();
+        let target = func.content.as_mut().unwrap().get_block_mut(edge.target);
         target.control_flow.preds.push(cur);
     }
     
-    // add target as current block's successors
+    // add target as current block's successors and start dfs
+    let succ = edge.target;
     {
-        let cur = func.content.as_mut().unwrap().get_block_mut(cur).unwrap();
+        let cur = func.content.as_mut().unwrap().get_block_mut(cur);
         cur.control_flow.succs.push(edge);
     }
+    if !visited.contains(&succ) {
+        dfs(succ, stack, visited, func);
+    }
+    
 }
 
 const WATCHPOINT_DISABLED_CHANCE : f32 = 0.9f32;
@@ -42,12 +48,16 @@ const NORMAL_RESUME_CHANCE       : f32 = 0.6f32;
 const EXN_RESUME_CHANCE          : f32 = 1f32 - NORMAL_RESUME_CHANCE;
 
 fn dfs(cur: MuTag, stack: &mut Vec<MuTag>, visited: &mut Vec<MuTag>, func: &mut MuFunction) {
+    trace!("dfs visiting block {}", cur);
+    trace!("current stack: {:?}", stack);
+    trace!("current visited: {:?}", visited);
+    
     stack.push(cur);
     visited.push(cur);
     
     // find all the successors for current block, and push them to the stack    
     let out_edges : Vec<BlockEdge> = {
-        let cur = func.content.as_mut().unwrap().get_block_mut(cur).unwrap(); 
+        let cur = func.content.as_mut().unwrap().get_block_mut(cur); 
         let ref body = cur.content.as_ref().unwrap().body;
         let last_inst = body.last().unwrap();
         
@@ -171,8 +181,10 @@ fn dfs(cur: MuTag, stack: &mut Vec<MuTag>, visited: &mut Vec<MuTag>, func: &mut
         }
     };
     
+    trace!("out edges for {}: {}", cur, vector_as_str(&out_edges));
+    
     for edge in out_edges {
-        new_edge(cur, edge, func);
+        new_edge(cur, edge, stack, visited, func);
     }
     
     stack.pop();

tests/lib.rs

 mod test_ir;
 mod test_compiler;
 
-#[macro_export]
-macro_rules! init_logger {
-    ($level : expr) => {
-        match simple_logger::init_with_level($level) {
-            Ok(_) => {},
-            Err(_) => {}
-        }
+#[macro_use]
+mod common {
+    use std::fmt;
+    
+    pub fn assert_str_vector (left: &Vec<&str>, right: &Vec<&str>) {
+        left.clone().sort();
+        right.clone().sort();
+        
+        assert_debug_str(left, right);
+    }
+    
+    pub fn assert_debug_str<T: fmt::Debug, U: fmt::Debug> (left: T, right: U) {
+        assert_eq!(format!("{:?}", left), format!("{:?}", right))
     }
 }
\ No newline at end of file

tests/test_compiler/test_pre_instsel.rs

@@ -2,7 +2,10 @@ extern crate mu;
 extern crate log;
 extern crate simple_logger;
 
+use common::*;
 use test_ir::test_ir::factorial;
+use test_ir::test_ir::sum;
+use self::mu::ast::ir::*;
 use self::mu::compiler::*;
 use self::mu::vm::context::VMContext;
 
@@ -49,7 +52,7 @@ fn test_build_tree() {
 }
 
 #[test]
-fn test_cfa() {
+fn test_cfa_factorial() {
     simple_logger::init_with_level(log::LogLevel::Trace).ok();
     
     let vm_context : VMContext = factorial();
@@ -63,5 +66,67 @@ fn test_cfa() {
         vm_context.get_func("fac").unwrap().borrow_mut()
     };
     
-    compiler.compile(&vm_context, &mut factorial_func);    
+    compiler.compile(&vm_context, &mut factorial_func);
+    
+    // assert cfa
+    let content = factorial_func.content.as_ref().unwrap();
+    
+    // blk_0: preds=[], succs=[blk_2, blk_1]
+    let blk_0 = content.get_block("blk_0");
+    assert_str_vector(&blk_0.control_flow.preds, &vec![]);
+    assert_str_vector(&block_edges_into_vec(&blk_0.control_flow.succs), &vec!["blk_2", "blk_1"]);
+    
+    // blk_2: preds=[blk_0, blk_1], succs=[]
+    let blk_2 = content.get_block("blk_2");
+    assert_str_vector(&blk_2.control_flow.preds, &vec!["blk_0", "blk_1"]);
+    assert_str_vector(&block_edges_into_vec(&blk_2.control_flow.succs), &vec![]);
+    
+    // blk_1: preds=[blk_0], succs=[blk_2]
+    let blk_1 = content.get_block("blk_1");
+    assert_str_vector(&blk_1.control_flow.preds, &vec!["blk_0"]);
+    assert_str_vector(&block_edges_into_vec(&blk_1.control_flow.succs), &vec!["blk_2"]);
+}
+
+#[test]
+fn test_cfa_sum() {
+    simple_logger::init_with_level(log::LogLevel::Trace).ok();
+    
+    let vm_context : VMContext = sum();
+    let compiler = Compiler::new(CompilerPolicy::new(vec![
+            Box::new(passes::DefUse::new()),
+            Box::new(passes::TreeGen::new()),
+            Box::new(passes::ControlFlowAnalysis::new())
+    ]));
+    
+    let mut sum_func = {
+        vm_context.get_func("sum").unwrap().borrow_mut()
+    };
+    
+    compiler.compile(&vm_context, &mut sum_func);
+    
+    // assert cfa
+    let content = sum_func.content.as_ref().unwrap();
+    
+    // entry: preds=[], succs=[head]
+    let entry = content.get_block("entry");
+    assert_str_vector(&entry.control_flow.preds, &vec![]);
+    assert_str_vector(&block_edges_into_vec(&entry.control_flow.succs), &vec!["head"]);
+    
+    // head: preds=[entry, head], succs=[head, ret]
+    let head = content.get_block("head");
+    assert_str_vector(&head.control_flow.preds, &vec!["entry", "head"]);
+    assert_str_vector(&block_edges_into_vec(&head.control_flow.succs), &vec!["ret", "head"]);
+    
+    // ret: preds=[head], succs=[]
+    let ret = content.get_block("ret");
+    assert_str_vector(&ret.control_flow.preds, &vec!["head"]);
+    assert_str_vector(&block_edges_into_vec(&ret.control_flow.succs), &vec![]);
+}
+
+fn block_edges_into_vec(edges: &Vec<BlockEdge>) -> Vec<&str> {
+    let mut ret = vec![];
+    for edge in edges {
+        ret.push(edge.target);
+    }
+    ret
 }
\ No newline at end of file

tests/test_ir/test_ir.rs

@@ -16,6 +16,143 @@ fn test_factorial() {
     let vm = factorial();
 }
 
+#[test]
+#[allow(unused_variables)]
+fn test_sum() {
+    let vm = sum();
+}
+
+pub fn sum() -> VMContext {
+    let mut vm = VMContext::new();
+    
+    // .typedef @int_64 = int<64>
+    let type_def_int64 = vm.declare_type("int_64", P(MuType::int(64)));
+    let type_def_int1  = vm.declare_type("int_1", P(MuType::int(1)));
+    
+    // .const @int_64_0 <@int_64> = 0
+    // .const @int_64_1 <@int_64> = 1
+    let const_def_int64_0 = vm.declare_const("int64_0", type_def_int64.clone(), Constant::Int(0));
+    let const_def_int64_1 = vm.declare_const("int64_1", type_def_int64.clone(), Constant::Int(1));
+    
+    // .funcsig @sum_sig = (@int_64) -> (@int_64)
+    let sum_sig = vm.declare_func_sig("sum_sig", vec![type_def_int64.clone()], vec![type_def_int64.clone()]);
+    
+    // .funcdef @sum VERSION @sum_v1 <@sum_sig>
+    let mut func = MuFunction::new("sum", sum_sig.clone());
+    
+    // %entry(<@int_64> %n):
+    let mut blk_entry = Block::new("entry");
+    let blk_entry_n = func.new_ssa(0, "blk_entry_n", type_def_int64.clone());
+    let const_def_int64_0_local = func.new_value(const_def_int64_0.clone()); // FIXME: why we need a local version?
+    let const_def_int64_1_local = func.new_value(const_def_int64_1.clone());
+    
+    // BRANCH %head
+    let blk_entry_term = TreeNode::new_inst(Instruction {
+        value: None,
+        ops: RefCell::new(vec![blk_entry_n.clone(), const_def_int64_0_local.clone(), const_def_int64_0_local.clone()]),
+        v: Instruction_::Branch1(Destination{
+            target: "head",
+            args: vec![DestArg::Normal(0), DestArg::Normal(1), DestArg::Normal(2)]
+        })
+    });
+    
+    let blk_entry_content = BlockContent {
+        args: vec![blk_entry_n.clone()],
+        body: vec![blk_entry_term],
+        keepalives: None
+    };
+    blk_entry.content = Some(blk_entry_content);
+    
+    // %head(<@int_64> %n, <@int_64> %s, <@int_64> %i):
+    let mut blk_head = Block::new("head");
+    let blk_head_n = func.new_ssa(1, "blk_head_n", type_def_int64.clone());
+    let blk_head_s = func.new_ssa(2, "blk_head_s", type_def_int64.clone());
+    let blk_head_i = func.new_ssa(3, "blk_head_i", type_def_int64.clone());
+    
+    // %s2 = ADD %s %i
+    let blk_head_s2 = func.new_ssa(4, "blk_head_s2", type_def_int64.clone());
+    let blk_head_inst0 = TreeNode::new_inst(Instruction {
+        value: Some(vec![blk_head_s2.clone()]),
+        ops: RefCell::new(vec![blk_head_s.clone(), blk_head_i.clone()]),
+        v: Instruction_::BinOp(BinOp::Add, 0, 1)
+    });
+    
+    // %i2 = ADD %i 1
+    let blk_head_i2 = func.new_ssa(5, "blk_head_i2", type_def_int64.clone());
+    let blk_head_inst1 = TreeNode::new_inst(Instruction {
+        value: Some(vec![blk_head_i2.clone()]),
+        ops: RefCell::new(vec![blk_head_i.clone(), const_def_int64_1_local.clone()]),
+        v: Instruction_::BinOp(BinOp::Add, 0, 1)
+    });
+    
+    // %cond = UGT %i %n
+    let blk_head_cond = func.new_ssa(6, "blk_head_cond", type_def_int1.clone());
+    let blk_head_inst2 = TreeNode::new_inst(Instruction {
+        value: Some(vec![blk_head_cond.clone()]),
+        ops: RefCell::new(vec![blk_head_i.clone(), blk_head_n.clone()]),
+        v: Instruction_::CmpOp(CmpOp::UGT, 0, 1)
+    });
+
+    // BRANCH2 %cond %ret(%s2) %head(%n %s2 %i2)
+    let blk_head_term = TreeNode::new_inst(Instruction{
+        value: None,
+        ops: RefCell::new(vec![blk_head_cond.clone(), blk_head_n.clone(), blk_head_s2.clone(), blk_head_i2.clone()]),
+        v: Instruction_::Branch2 {
+            cond: 0,
+            true_dest: Destination {
+                target: "ret",
+                args: vec![DestArg::Normal(2)]
+            },
+            false_dest: Destination {
+                target: "head",
+                args: vec![DestArg::Normal(1), DestArg::Normal(2), DestArg::Normal(3)]
+            },
+            true_prob: 0.6f32
+        }
+    });
+
+    let blk_head_content = BlockContent {
+        args: vec![blk_head_n.clone(), blk_head_s.clone(), blk_head_i.clone()],
+        body: vec![blk_head_inst0, blk_head_inst1, blk_head_inst2, blk_head_term],
+        keepalives: None
+    };
+    blk_head.content = Some(blk_head_content);
+    
+    // %ret(<@int_64> %s):
+    let mut blk_ret = Block::new("ret");
+    let blk_ret_s = func.new_ssa(7, "blk_ret_s", type_def_int64.clone());
+    
+    // RET %s
+    let blk_ret_term = TreeNode::new_inst(Instruction{
+        value: None,
+        ops: RefCell::new(vec![blk_ret_s.clone()]),
+        v: Instruction_::Return(vec![0])
+    });
+    
+    let blk_ret_content = BlockContent {
+        args: vec![blk_ret_s.clone()],
+        body: vec![blk_ret_term],
+        keepalives: None
+    };
+    blk_ret.content = Some(blk_ret_content);
+
+    // wrap into a function
+    func.define(FunctionContent{
+            entry: "entry", 
+            blocks: {
+                let mut blocks = HashMap::new();
+                blocks.insert("entry", blk_entry);
+                blocks.insert("head", blk_head);
+                blocks.insert("ret", blk_ret);
+                blocks
+            }
+    });
+    
+    vm.declare_func(func);
+    
+    vm
+}
+
 #[allow(unused_variables)]
 pub fn factorial() -> VMContext {
     let mut vm = VMContext::new();