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.

brilirs/src/basic_block.rs

@@ -1,40 +1,43 @@
 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>);
+// A program represented as basic blocks.
+pub struct BBProgram {
+  pub blocks: Vec<BasicBlock>,
 
-#[derive(Debug)]
-pub struct BasicBlock {
-  pub instrs: Vec<bril_rs::Code>,
-  pub exit: Vec<usize>,
+  // 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 the index into `blocks` of the starting block of
+  // the function.
+  pub func_index: HashMap<String, usize>,
 }
 
-impl BasicBlock {
-  fn new() -> BasicBlock {
-    BasicBlock {
-      instrs: Vec::new(),
-      exit: Vec::new(),
+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
   }
-}
-
-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();
 
-  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(), 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 +49,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 +63,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
 }

brilirs/src/interp.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 use std::convert::TryFrom;
 use std::fmt;
+use std::iter::FromIterator;
 
 use crate::basic_block::{BBProgram, BasicBlock};
 
@@ -9,6 +10,8 @@ pub enum InterpError {
   BadJsonInt,
   BadJsonBool,
   NoMainFunction,
+  FuncNotFound(String),
+  NoRetValForfunc(String),
   BadNumArgs(usize, usize),   // (expected, actual)
   BadNumLabels(usize, usize), // (expected, actual)
   VarNotFound(String),
@@ -151,12 +154,15 @@ impl TryFrom<&Value> for f64 {
 }
 
 #[allow(clippy::float_cmp)]
-fn execute_value_op(
+fn execute_value_op<W: std::io::Write>(
+  prog: &BBProgram,
   op: &bril_rs::ValueOps,
   dest: &str,
   op_type: &bril_rs::Type,
   args: &[String],
+  funcs: &[String],
   value_store: &mut HashMap<String, Value>,
+  out: &mut W,
 ) -> Result<(), InterpError> {
   use bril_rs::ValueOps::*;
   match *op {
@@ -271,7 +277,19 @@ fn execute_value_op(
       let args = get_args::<f64>(value_store, 2, args)?;
       value_store.insert(String::from(dest), Value::Bool(args[0] >= args[1]));
     }
-    Call => unreachable!(), // TODO(yati): Why is Call a ValueOp as well?
+    Call => {
+      assert!(funcs.len() == 1);
+      let vals = get_values(value_store, args.len(), args)?;
+      let vars: HashMap<String, Value> =
+        HashMap::from_iter(args.iter().cloned().zip(vals.into_iter().cloned()));
+      if let Some(val) = execute_func(&prog, &funcs[0], vars, out)? {
+        check_asmt_type(&val.get_type(), op_type)?;
+        value_store.insert(String::from(dest), val);
+      } else {
+        // This is a value-op call, so the target func must return a result.
+        return Err(InterpError::NoRetValForfunc(funcs[0].clone()));
+      }
+    }
     Phi | Alloc | Load | PtrAdd => unimplemented!(),
   }
   Ok(())
@@ -293,17 +311,29 @@ fn check_num_labels(expected: usize, labels: &[String]) -> Result<(), InterpErro
   }
 }
 
-// Returns whether the program should continue running (i.e., if a Return was
-// *not* executed).
+// Result of executing an effect operation.
+enum EffectResult {
+  // Return from the current function without any value.
+  Return,
+
+  // Return a given value from the current function.
+  ReturnWithVal(Value),
+
+  // Continue execution of the current function.
+  Continue,
+}
+
 fn execute_effect_op<T: std::io::Write>(
+  prog: &BBProgram,
   op: &bril_rs::EffectOps,
   args: &[String],
   labels: &[String],
+  funcs: &[String],
   curr_block: &BasicBlock,
   value_store: &HashMap<String, Value>,
-  mut out: T,
+  out: &mut T,
   next_block_idx: &mut Option<usize>,
-) -> Result<bool, InterpError> {
+) -> Result<EffectResult, InterpError> {
   use bril_rs::EffectOps::*;
   match op {
     Jump => {
@@ -320,36 +350,53 @@ fn execute_effect_op<T: std::io::Write>(
     Return => {
       out.flush().map_err(|e| InterpError::IoError(Box::new(e)))?;
       // NOTE: This only works so long as `main` is the only function
-      return Ok(false);
+      if args.is_empty() {
+        return Ok(EffectResult::Return);
+      }
+      let retval = value_store
+        .get(&args[0])
+        .ok_or(InterpError::VarNotFound(args[0].clone()))?;
+      return Ok(EffectResult::ReturnWithVal(retval.clone()));
     }
     Print => {
+      let vals = get_values(value_store, args.len(), args)?;
       writeln!(
         out,
         "{}",
-        args
+        vals
           .iter()
-          .map(|a| format!("{}", value_store[a]))
+          .map(|v| format!("{}", v))
           .collect::<Vec<_>>()
           .join(", ")
       )
       .map_err(|e| InterpError::IoError(Box::new(e)))?;
     }
     Nop => {}
-    Call => unreachable!(),
+    Call => {
+      assert!(funcs.len() == 1);
+      let vals = get_values(value_store, args.len(), args)?;
+      let vars: HashMap<String, Value> =
+        HashMap::from_iter(args.iter().cloned().zip(vals.into_iter().cloned()));
+      execute_func(&prog, &funcs[0], vars, out)?;
+    }
     Store | Free | Speculate | Commit | Guard => unimplemented!(),
   }
-  Ok(true)
+  Ok(EffectResult::Continue)
 }
 
-pub fn execute<T: std::io::Write>(prog: BBProgram, mut out: T) -> Result<(), InterpError> {
-  let (main_fn, blocks, _labels) = prog;
-  let mut curr_block_idx: usize = main_fn.ok_or(InterpError::NoMainFunction)?;
-
-  // Map from variable name to value.
-  let mut value_store: HashMap<String, Value> = HashMap::new();
+fn execute_func<T: std::io::Write>(
+  prog: &BBProgram,
+  func: &str,
+  mut vars: HashMap<String, Value>,
+  out: &mut T,
+) -> Result<Option<Value>, InterpError> {
+  let mut curr_block_idx = *prog
+    .func_index
+    .get(func)
+    .ok_or(InterpError::FuncNotFound(String::from(func)))?;
 
   loop {
-    let curr_block = &blocks[curr_block_idx];
+    let curr_block = &prog.blocks[curr_block_idx];
     let curr_instrs = &curr_block.instrs;
     let mut next_block_idx = if curr_block.exit.len() == 1 {
       Some(curr_block.exit[0])
@@ -367,44 +414,58 @@ pub fn execute<T: std::io::Write>(prog: BBProgram, mut out: T) -> Result<(), Int
             value,
           } => {
             check_asmt_type(const_type, &value.get_type())?;
-            value_store.insert(dest.clone(), Value::from(value));
+            vars.insert(dest.clone(), Value::from(value));
           }
           bril_rs::Instruction::Value {
             op,
             dest,
             op_type,
             args,
+            funcs,
             ..
           } => {
-            execute_value_op(op, dest, op_type, args, &mut value_store)?;
+            execute_value_op(&prog, op, dest, op_type, args, funcs, &mut vars, out)?;
           }
           bril_rs::Instruction::Effect {
-            op, args, labels, ..
+            op,
+            args,
+            labels,
+            funcs,
+            ..
           } => {
-            let should_continue = execute_effect_op(
+            match execute_effect_op(
+              prog,
               op,
               args,
               labels,
+              funcs,
               &curr_block,
-              &value_store,
-              &mut out,
+              &vars,
+              out,
               &mut next_block_idx,
-            )?;
-
-            // TODO(yati): Correct only when main is the only function.
-            if !should_continue {
-              return Ok(());
-            }
+            )? {
+              EffectResult::Continue => {}
+              EffectResult::Return => {
+                return Ok(None);
+              }
+              EffectResult::ReturnWithVal(val) => {
+                return Ok(Some(val));
+              }
+            };
           }
         }
       }
     }
-
     if let Some(idx) = next_block_idx {
       curr_block_idx = idx;
     } else {
       out.flush().map_err(|e| InterpError::IoError(Box::new(e)))?;
-      return Ok(());
+      return Ok(None);
     }
   }
 }
+
+pub fn execute<T: std::io::Write>(prog: BBProgram, out: &mut T) -> Result<(), InterpError> {
+  // Ignore return value of @main.
+  execute_func(&prog, "main", HashMap::new(), out).map(|_| ())
+}