brilirs: Add call support

This is a simple recursive implementation of the `call` operation (both
in the value and effect contexts). I originally tried an iterative version,
but the borrow checker gave me a tough fight. This is much more readable
than tracking a stack of function call environments directly anyway.

Author: yati-sagade

brilirs/src/basic_block.rs

@@ -1,40 +1,61 @@
 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>);
-
-#[derive(Debug)]
-pub struct BasicBlock {
-  pub instrs: Vec<bril_rs::Code>,
-  pub exit: Vec<usize>,
+pub struct FuncInfo {
+  pub args: Vec<bril_rs::Argument>,
+  pub return_type: Option<bril_rs::Type>,
+  pub start_block: usize,
 }
 
-impl BasicBlock {
-  fn new() -> BasicBlock {
-    BasicBlock {
-      instrs: Vec::new(),
-      exit: Vec::new(),
+impl FuncInfo {
+  fn new(f: &bril_rs::Function, start_block: usize) -> FuncInfo {
+    FuncInfo {
+      args: f.args.clone(),
+      return_type: f.return_type.clone(),
+      start_block: start_block,
     }
   }
 }
 
-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();
+// A program represented as basic blocks.
+pub struct BBProgram {
+  pub blocks: Vec<BasicBlock>,
+
+  // Map from label name to index in `blocks` of the block named by that label.
+  pub label_index: HashMap<String, usize>,
+
+  // Map from function name to information needed during runtime about the
+  // function.
+  pub func_index: HashMap<String, FuncInfo>,
+}
+
+impl BBProgram {
+  pub fn new(prog: bril_rs::Program) -> BBProgram {
+    let mut bbprog = BBProgram {
+      blocks: vec![],
+      label_index: HashMap::new(),
+      func_index: HashMap::new(),
+    };
+    for func in prog.functions {
+      bbprog.add_func_bbs(func);
+    }
+    bbprog
+  }
 
-  let mut bb_helper = |func: bril_rs::Function| -> usize {
+  fn add_func_bbs(&mut self, func: bril_rs::Function) -> usize {
+    self
+      .func_index
+      .insert(func.name.clone(), FuncInfo::new(&func, self.blocks.len()));
     let mut curr_block = BasicBlock::new();
-    let root_block = blocks.len();
+    let root_block = self.blocks.len();
     let mut curr_label = None;
+
     for instr in func.instrs.into_iter() {
       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 +67,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();
@@ -60,22 +81,36 @@ 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
-  };
+  }
+}
+
+#[derive(Debug)]
+pub struct BasicBlock {
+  pub instrs: Vec<bril_rs::Code>,
+  pub exit: Vec<usize>,
+}
 
-  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 BasicBlock {
+  fn new() -> BasicBlock {
+    BasicBlock {
+      instrs: Vec::new(),
+      exit: Vec::new(),
     }
   }
-
-  (main_fn, blocks, labels)
 }
+
+const RET: bril_rs::Code = bril_rs::Code::Instruction(bril_rs::Instruction::Effect {
+  op: bril_rs::EffectOps::Return,
+  args: vec![],
+  funcs: vec![],
+  labels: vec![],
+});

brilirs/src/cfg.rs

@@ -1,34 +1,23 @@
-use crate::basic_block::BasicBlock;
+use crate::basic_block::BBProgram;
 
-use std::collections::HashMap;
-
-type CFG = Vec<BasicBlock>;
-
-pub fn build_cfg(mut blocks: Vec<BasicBlock>, label_to_block_idx: &HashMap<String, usize>) -> CFG {
-  let last_idx = blocks.len() - 1;
-  for (i, block) in blocks.iter_mut().enumerate() {
+pub fn build_cfg(prog: &mut BBProgram) {
+  let last_idx = prog.blocks.len() - 1;
+  for (i, block) in prog.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
       {
-        match op {
-          bril_rs::EffectOps::Jump | bril_rs::EffectOps::Branch => {
-            for l in labels {
-              block.exit.push(label_to_block_idx[l]);
-            }
+        if let bril_rs::EffectOps::Jump | bril_rs::EffectOps::Branch = op {
+          for l in labels {
+            block.exit.push(prog.label_index[l]);
           }
-          bril_rs::EffectOps::Return => {}
-          // TODO(yati): Do all effect ops end a BB?
-          _ => {}
         }
       } else {
         block.exit.push(i + 1);
       }
     }
   }
-
-  blocks
 }