brilirs: Add call support

brilirs/src/basic_block.rs

@@ -1,40 +1,64 @@
 use std::collections::HashMap;
 
-// A program composed of basic blocks.
-// (BB index of main program, list of BBs, mapping of label -> BB index)
-pub type BBProgram = (Option<usize>, Vec<BasicBlock>, HashMap<String, usize>);
+pub struct Function {
+  pub args: Vec<bril_rs::Argument>,
+  pub return_type: Option<bril_rs::Type>,
+  pub blocks: Vec<BasicBlock>,
 
-#[derive(Debug)]
-pub struct BasicBlock {
-  pub instrs: Vec<bril_rs::Code>,
-  pub exit: Vec<usize>,
+  // Map from label to the index of the block that is the target of the label.
+  pub label_index: HashMap<String, usize>,
 }
 
-impl BasicBlock {
-  fn new() -> BasicBlock {
-    BasicBlock {
-      instrs: Vec::new(),
-      exit: Vec::new(),
-    }
+impl Function {
+  pub fn new(f: bril_rs::Function) -> Function {
+    let mut func = Function {
+      args: f.args.clone(),
+      return_type: f.return_type.clone(),
+      blocks: vec![],
+      label_index: HashMap::new(),
+    };
+    func.add_blocks(f.instrs);
+    func.build_cfg();
+    func
   }
-}
 
-pub fn find_basic_blocks(prog: bril_rs::Program) -> BBProgram {
-  let mut main_fn = None;
-  let mut blocks = Vec::new();
-  let mut labels = HashMap::new();
+  fn build_cfg(&mut self) {
+    let last_idx = self.blocks.len() - 1;
+    for (i, block) in self.blocks.iter_mut().enumerate() {
+      // If we're before the last block
+      if i < last_idx {
+        // Get the last instruction
+        let last_instr: &bril_rs::Code = block.instrs.last().unwrap();
+        if let bril_rs::Code::Instruction(bril_rs::Instruction::Effect { op, labels, .. }) =
+          last_instr
+        {
+          if let bril_rs::EffectOps::Jump | bril_rs::EffectOps::Branch = op {
+            for l in labels {
+              block.exit.push(
+                *self
+                  .label_index
+                  .get(l)
+                  .expect(&format!("No label {} found.", &l)),
+              );
+            }
+          }
+        } else {
+          block.exit.push(i + 1);
+        }
+      }
+    }
+  }
 
-  let mut bb_helper = |func: bril_rs::Function| -> usize {
+  fn add_blocks(&mut self, instrs: Vec<bril_rs::Code>) {
     let mut curr_block = BasicBlock::new();
-    let root_block = blocks.len();
     let mut curr_label = None;
-    for instr in func.instrs.into_iter() {
+    for instr in instrs {
       match instr {
         bril_rs::Code::Label { ref label } => {
           if !curr_block.instrs.is_empty() {
-            blocks.push(curr_block);
+            self.blocks.push(curr_block);
             if let Some(old_label) = curr_label {
-              labels.insert(old_label, blocks.len() - 1);
+              self.label_index.insert(old_label, self.blocks.len() - 1);
             }
             curr_block = BasicBlock::new();
           }
@@ -46,9 +70,9 @@ pub fn find_basic_blocks(prog: bril_rs::Program) -> BBProgram {
             || op == bril_rs::EffectOps::Return =>
         {
           curr_block.instrs.push(instr);
-          blocks.push(curr_block);
+          self.blocks.push(curr_block);
           if let Some(l) = curr_label {
-            labels.insert(l, blocks.len() - 1);
+            self.label_index.insert(l, self.blocks.len() - 1);
             curr_label = None;
           }
           curr_block = BasicBlock::new();
@@ -58,24 +82,55 @@ pub fn find_basic_blocks(prog: bril_rs::Program) -> BBProgram {
         }
       }
     }
-
     if !curr_block.instrs.is_empty() {
-      blocks.push(curr_block);
+      // If we are here, the function ends without an explicit ret. To make
+      // processing easier, push a Return op onto the last block.
+      curr_block.instrs.push(RET.clone());
+      self.blocks.push(curr_block);
       if let Some(l) = curr_label {
-        labels.insert(l, blocks.len() - 1);
+        self.label_index.insert(l, self.blocks.len() - 1);
       }
     }
+  }
+}
 
-    root_block
-  };
+// A program represented as basic blocks.
+pub struct BBProgram {
+  pub func_index: HashMap<String, Function>,
+}
 
-  for func in prog.functions.into_iter() {
-    let func_name = func.name.clone();
-    let func_block = bb_helper(func);
-    if func_name == "main" {
-      main_fn = Some(func_block);
+impl BBProgram {
+  pub fn new(prog: bril_rs::Program) -> BBProgram {
+    let mut bbprog = BBProgram {
+      func_index: HashMap::new(),
+    };
+    for func in prog.functions {
+      bbprog
+        .func_index
+        .insert(func.name.clone(), Function::new(func));
     }
+    bbprog
   }
+}
+
+#[derive(Debug)]
+pub struct BasicBlock {
+  pub instrs: Vec<bril_rs::Code>,
+  pub exit: Vec<usize>,
+}
 
-  (main_fn, blocks, labels)
+impl BasicBlock {
+  fn new() -> BasicBlock {
+    BasicBlock {
+      instrs: Vec::new(),
+      exit: Vec::new(),
+    }
+  }
 }
+
+const RET: bril_rs::Code = bril_rs::Code::Instruction(bril_rs::Instruction::Effect {
+  op: bril_rs::EffectOps::Return,
+  args: vec![],
+  funcs: vec![],
+  labels: vec![],
+});