[wip] inlining

src/ast/src/ir.rs

@@ -188,6 +188,78 @@ impl MuFunctionVersion {
             v: TreeNode_::Instruction(v),
         })
     }
+
+    /// get Map(CallSiteID -> FuncID) that are called by this function
+    pub fn get_static_call_edges(&self) -> HashMap<MuID, MuID> {
+        let mut ret = HashMap::new();
+
+        let f_content = self.content.as_ref().unwrap();
+
+        for (blk_id, block) in f_content.blocks.iter() {
+            let block_content = block.content.as_ref().unwrap();
+
+            for inst in block_content.body.iter() {
+                match inst.v {
+                    TreeNode_::Instruction(ref inst) => {
+                        let ops = inst.ops.read().unwrap();
+
+                        match inst.v {
+                            Instruction_::ExprCall{ref data, ..}
+                            | Instruction_::ExprCCall{ref data, ..}
+                            | Instruction_::Call {ref data, ..}
+                            | Instruction_::CCall {ref data, ..} => {
+                                let ref callee = ops[data.func];
+
+                                match callee.v {
+                                    TreeNode_::Instruction(_) => {},
+                                    TreeNode_::Value(ref pv) => match pv.v {
+                                        Value_::Constant(Constant::FuncRef(id)) => {ret.insert(inst.id(), id);},
+                                        _ => {}
+                                    }
+                                }
+                            },
+                            _ => {
+                                // do nothing
+                            }
+                        }
+                    },
+                    _ => {
+                        unreachable!()
+                    }
+                }
+            }
+        }
+
+        ret
+    }
+
+    pub fn has_throw(&self) -> bool {
+        let f_content = self.content.as_ref().unwrap();
+
+        for (blk_id, block) in f_content.blocks.iter() {
+            let block_content = block.content.as_ref().unwrap();
+
+            for inst in block_content.body.iter() {
+                match inst.v {
+                    TreeNode_::Instruction(ref inst) => {
+                        let ops = inst.ops.read().unwrap();
+
+                        match inst.v {
+                            Instruction_::Throw(_) => {return true;}
+                            _ => {
+                                // do nothing
+                            }
+                        }
+                    },
+                    _ => {
+                        unreachable!()
+                    }
+                }
+            }
+        }
+
+        false
+    }
 }
 
 #[derive(RustcEncodable, RustcDecodable, Clone)]
@@ -310,6 +382,16 @@ impl Block {
     pub fn is_exception_block(&self) -> bool {
         return self.content.as_ref().unwrap().exn_arg.is_some()
     }
+
+    pub fn number_of_irs(&self) -> usize {
+        if self.content.is_none() {
+            0
+        } else {
+            let content = self.content.as_ref().unwrap();
+
+            content.body.len()
+        }
+    }
 }
 
 #[derive(Debug, RustcEncodable, RustcDecodable, Clone)]
@@ -482,6 +564,13 @@ impl TreeNode {
         })
     }
 
+    pub fn new_boxed_inst(v: Instruction) -> Box<TreeNode> {
+        Box::new(TreeNode{
+            op: pick_op_code_for_inst(&v),
+            v: TreeNode_::Instruction(v),
+        })
+    }
+
     pub fn extract_ssa_id(&self) -> Option<MuID> {
         match self.v {
             TreeNode_::Value(ref pv) => {
@@ -514,6 +603,13 @@ impl TreeNode {
             _ => None
         }
     }
+
+    pub fn into_inst(self) -> Option<Instruction> {
+        match self.v {
+            TreeNode_::Instruction(inst) => Some(inst),
+            _ => None
+        }
+    }
 }
 
 /// use +() to display a node

src/compiler/mod.rs

@@ -64,6 +64,7 @@ impl CompilerPolicy {
 impl Default for CompilerPolicy {
     fn default() -> Self {
         let mut passes : Vec<Box<CompilerPass>> = vec![];
+        passes.push(Box::new(passes::Inlining::new()));
         // ir level passes
         passes.push(Box::new(passes::DefUse::new()));
         passes.push(Box::new(passes::TreeGen::new()));

src/compiler/passes/inlining.rs

@@ -24,11 +24,243 @@ impl Inlining {
     }
 
     fn check(&mut self, vm: &VM, func: &mut MuFunctionVersion) -> bool {
-        unimplemented!()
+        debug!("check inline");
+        let mut inline_something = false;
+
+        for func_id in func.get_static_call_edges().values() {
+            let should_inline_this = self.check_should_inline_func(*func_id, vm);
+
+            inline_something = inline_something || should_inline_this;
+        }
+
+        inline_something
+    }
+
+    #[allow(unused_variables)]
+    fn check_should_inline_func(&mut self, func_id: MuID, vm: &VM) -> bool {
+        let funcs_guard = vm.funcs().read().unwrap();
+        let func = match funcs_guard.get(&func_id) {
+            Some(func) => func.read().unwrap(),
+            None => panic!("callee {} is undeclared", func_id)
+        };
+
+        let fv_id = match func.cur_ver {
+            Some(fv_id) => fv_id,
+            None => {
+                // the funtion is not defined
+                info!("the function is undefined, we cannot inline it. ");
+                return false;
+            }
+        };
+
+        match self.should_inline.get(&fv_id) {
+            Some(flag) => {
+                trace!("func {} should be inlined (checked before)", func_id);
+                return *flag;
+            }
+            None => {}
+        }
+
+        let fv_guard = vm.func_vers().read().unwrap();
+        let fv = fv_guard.get(&fv_id).unwrap().read().unwrap();
+
+        // if the function is forced inline, then we inline it
+        if fv.force_inline {
+            trace!("func {} is forced as inline function", func_id);
+            return true;
+        }
+
+        // some heuristics here to decide if we should inline the function
+        // to be more precise. we should be target specific
+        let n_params = fv.sig.arg_tys.len();
+        let n_insts  = fv.content.as_ref().unwrap().blocks.values().fold(0usize, |mut sum, ref block| {sum += block.number_of_irs(); sum});
+        let out_calls = fv.get_static_call_edges();
+        let has_throw = fv.has_throw();
+
+        // now we use a simple heuristic here:
+        // insts fewer than 10, no static out calls, no throw
+        let should_inline = n_insts <= 10 && out_calls.len() == 0 && !has_throw;
+
+        trace!("func has {} insts", n_insts);
+        trace!("     has {} out calls", out_calls.len());
+        trace!("     has throws? {}", has_throw);
+        trace!("SO func should be inlined? {}", should_inline);
+
+        self.should_inline.insert(fv_id, should_inline);
+
+        should_inline
     }
 
     fn inline(&mut self, vm: &VM, func: &mut MuFunctionVersion) {
-        unimplemented!()
+        debug!("inlining for Function {}", func);
+
+        let call_edges = func.get_static_call_edges();
+
+        let mut f_content = func.content.as_mut().unwrap();
+        let ref mut f_context = func.context;
+
+        let mut new_blocks : Vec<Block> = vec![];
+
+        for (blk_id, mut block) in f_content.blocks.drain() {
+            // clone curent block, and clear its instructions
+            let mut cur_block = block.clone();
+            cur_block.content.as_mut().unwrap().body.clear();
+
+            // iterate through instructions
+            for inst in block.content.unwrap().body {
+                let inst_id = inst.id();
+                if call_edges.contains_key(&inst_id) {
+                    trace!("inserting inlined function at {}", inst);
+
+                    // from TreeNode into Inst (we do not need old TreeNode)
+                    let inst = inst.into_inst().unwrap();
+
+                    // (inline expansion)
+
+                    let inlined_func = *call_edges.get(&inst.id()).unwrap();
+                    trace!("function being inlined is {}", inlined_func);
+
+                    let inlined_fvid = match vm.get_cur_version_of(inlined_func) {
+                        Some(fvid) => fvid,
+                        None => panic!("cannot resolve current version of Func {}, which is supposed to be inlined", inlined_func)
+                    };
+
+                    let inlined_fvs_guard = vm.func_vers().read().unwrap();
+                    let inlined_fv_lock   = inlined_fvs_guard.get(&inlined_fvid).unwrap();
+                    let inlined_fv_guard  = inlined_fv_lock.read().unwrap();
+
+                    let inlined_entry = inlined_fv_guard.content.as_ref().unwrap().entry;
+
+                    // change current call insts to a branch
+                    trace!("turning CALL instruction into a branch");
+                    let ops = inst.ops.read().unwrap();
+
+                    match inst.v {
+                        Instruction_::ExprCall {ref data, ..}
+                        | Instruction_::ExprCCall {ref data, ..} => {
+                            let arg_nodes  : Vec<P<TreeNode>> = data.args.iter().map(|x| ops[*x].clone()).collect();
+                            let arg_indices: Vec<OpIndex> = (0..arg_nodes.len()).collect();
+
+                            let branch = Instruction{
+                                hdr: inst.hdr.clone(),
+                                value: inst.value.clone(),
+                                ops: RwLock::new(arg_nodes),
+                                v: Instruction_::Branch1(Destination{
+                                    target: inlined_entry,
+                                    args: arg_indices.iter().map(|x| DestArg::Normal(*x)).collect()
+                                })
+                            };
+
+                            // add branch to current block
+                            cur_block.content.as_mut().unwrap().body.push(TreeNode::new_boxed_inst(branch));
+
+                            // finish current block
+                            new_blocks.push(cur_block.clone());
+                            let old_name = cur_block.name().unwrap();
+
+                            // start a new block
+                            cur_block = Block::new(vm.next_id());
+                            let new_name = format!("{}_cont_after_inline_{}", old_name, inst_id);
+                            vm.set_name(cur_block.as_entity(), new_name);
+
+                            // deal with the inlined function
+                            copy_inline_blocks(&mut new_blocks, cur_block.id(), inlined_fv_guard.content.as_ref().unwrap());
+                            copy_inline_context(f_context, &inlined_fv_guard.context);
+                        },
+
+                        Instruction_::Call {ref data, ref resume}
+                        | Instruction_::CCall {ref data, ref resume} => {
+                            let arg_nodes  : Vec<P<TreeNode>> = data.args.iter().map(|x| ops[*x].clone()).collect();
+                            let arg_indices: Vec<OpIndex> = (0..arg_nodes.len()).collect();
+
+                            let branch = Instruction{
+                                hdr: inst.hdr.clone(),
+                                value: inst.value.clone(),
+                                ops: RwLock::new(arg_nodes),
+                                v: Instruction_::Branch1(Destination{
+                                    target: inlined_entry,
+                                    args: arg_indices.iter().map(|x| DestArg::Normal(*x)).collect()
+                                })
+                            };
+
+                            // add branch to current block
+                            cur_block.content.as_mut().unwrap().body.push(TreeNode::new_boxed_inst(branch));
+
+                            // deal with inlined function
+                            let next_block = resume.normal_dest.target;
+
+                            copy_inline_blocks(&mut new_blocks, next_block, inlined_fv_guard.content.as_ref().unwrap());
+                            copy_inline_context(f_context, &inlined_fv_guard.context);
+                        },
+
+                        _ => panic!("unexpected callsite: {}", inst)
+                    }
+
+                    // change return in inlined function to a branch to the new block
+                } else {
+                    cur_block.content.as_mut().unwrap().body.push(inst.clone());
+                }
+            }
+
+            new_blocks.push(cur_block);
+        }
+
+        f_content.blocks.clear();
+        for blk in new_blocks {
+            f_content.blocks.insert(blk.id(), blk);
+        }
+    }
+}
+
+fn copy_inline_blocks(caller: &mut Vec<Block>, ret_block: MuID, callee: &FunctionContent) {
+    for block in callee.blocks.values() {
+        let mut block = block.clone();
+
+        // check its last instruction
+        {
+            let block_content = block.content.as_mut().unwrap();
+            let last_inst = block_content.body.pop().unwrap();
+            let last_inst_clone = last_inst.clone();
+
+            match last_inst.v {
+                TreeNode_::Instruction(inst) => {
+                    let hdr = inst.hdr;
+                    let value = inst.value;
+                    let ops = inst.ops;
+                    let v = inst.v;
+
+                    match v {
+                        Instruction_::Return(vec) => {
+                            // change RET to a branch
+                            let branch = Instruction {
+                                hdr: hdr,
+                                value: value,
+                                ops: ops,
+                                v: Instruction_::Branch1(Destination {
+                                    target: ret_block,
+                                    args: vec.iter().map(|x| DestArg::Normal(*x)).collect()
+                                })
+                            };
+
+                            block_content.body.push(TreeNode::new_boxed_inst(branch));
+                        },
+                        _ => {block_content.body.push(last_inst_clone);}
+                    }
+                },
+                _ => {
+                    // do nothing, and directly push the instruction back
+                    block_content.body.push(last_inst_clone)
+                }
+            }
+        }
+
+        caller.push(block);
+    }
+}
+
+fn copy_inline_context(caller: &mut FunctionContext, callee: &FunctionContext) {
+    for (id, entry) in callee.values.iter() {
+        caller.values.insert(*id, SSAVarEntry::new(entry.value().clone()));
     }
 }
 

src/testutil/mod.rs

@@ -66,3 +66,25 @@ pub fn compile_fnc<'a>(fnc_name: &'static str, build_fnc: &'a Fn() -> VM) -> ll:
     let dylib = aot::link_dylib(vec![Mu(fnc_name)], libname, &vm);
     ll::Library::new(dylib.as_os_str()).unwrap()
 }
+
+pub fn compile_fncs<'a>(entry: &'static str, fnc_names: Vec<&'static str>, build_fnc: &'a Fn() -> VM) -> ll::Library {
+    VM::start_logging_trace;
+
+    let vm = Arc::new(build_fnc());
+    let compiler = Compiler::new(CompilerPolicy::default(), vm.clone());
+
+    for func in fnc_names.iter() {
+        let func_id = vm.id_of(func);
+        let funcs = vm.funcs().read().unwrap();
+        let func = funcs.get(&func_id).unwrap().read().unwrap();
+        let func_vers = vm.func_vers().read().unwrap();
+        let mut func_ver = func_vers.get(&func.cur_ver.unwrap()).unwrap().write().unwrap();
+        compiler.compile(&mut func_ver);
+    }
+
+    backend::emit_context(&vm);
+
+    let libname = &format!("lib{}.dylib", entry);
+    let dylib = aot::link_dylib(fnc_names.iter().map(|x| Mu(x)).collect(), libname, &vm);
+    ll::Library::new(dylib.as_os_str()).unwrap()
+}
\ No newline at end of file

src/vm/vm.rs

@@ -736,6 +736,17 @@ impl <'a> VM {
     pub fn func_vers(&self) -> &RwLock<HashMap<MuID, RwLock<MuFunctionVersion>>> {
         &self.func_vers
     }
+
+    pub fn get_cur_version_of(&self, fid: MuID) -> Option<MuID> {
+        let funcs_guard = self.funcs.read().unwrap();
+        match funcs_guard.get(&fid) {
+            Some(rwlock_func) => {
+                let func_guard = rwlock_func.read().unwrap();
+                func_guard.cur_ver
+            },
+            None => None
+        }
+    }
     
     pub fn compiled_funcs(&self) -> &RwLock<HashMap<MuID, RwLock<CompiledFunction>>> {
         &self.compiled_funcs

tests/ir_macros.rs

@@ -38,6 +38,11 @@ macro_rules! typedef {
         let $name = $vm.declare_type($vm.next_id(), MuType_::hybrid(Mu(stringify!($name)), vec![], $var_ty.clone()));
         $vm.set_name($name.as_entity(), Mu(stringify!($name)));
     };
+
+    (($vm: expr) $name: ident = mu_funcref($sig: ident)) => {
+        let $name = $vm.declare_type($vm.next_id(), MuType_::funcref($sig.clone()));
+        $vm.set_name($name.as_entity(), Mu(stringify!($name)));
+    }
 }
 
 macro_rules! constdef {
@@ -309,6 +314,25 @@ macro_rules! inst {
         });
     };
 
+    // CALL
+    (($vm: expr, $fv: ident) $name: ident: $res: ident = EXPRCALL ($cc: expr, is_abort: $is_abort: expr) $func: ident ($($val: ident), +)) => {
+        let ops = vec![$func, $($val.clone()), *];
+        let ops_len = ops.len();
+        let $name = $fv.new_inst(Instruction{
+            hdr:    MuEntityHeader::unnamed($vm.next_id()),
+            value:  Some(vec![$res.clone_value()]),
+            ops:    RwLock::new(ops),
+            v:      Instruction_::ExprCall {
+                        data: CallData {
+                            func: 0,
+                            args: (1..ops_len).collect(),
+                            convention: $cc
+                        },
+                        is_abort: $is_abort
+                    }
+        });
+    };
+
     // RET
     (($vm: expr, $fv: ident) $name: ident: RET ($($val: ident), +)) => {
         let $name = $fv.new_inst(Instruction{

tests/test_compiler/mod.rs

@@ -11,4 +11,5 @@ mod test_int;
 mod test_binop;
 mod test_controlflow;
 mod test_call;
-mod test_mem_inst;
\ No newline at end of file
+mod test_mem_inst;
+mod test_inline;
\ No newline at end of file

tests/test_compiler/test_inline.rs

+extern crate libloading;
+
+use mu::ast::types::*;
+use mu::ast::ir::*;
+use mu::ast::inst::*;
+use mu::ast::op::*;
+use mu::vm::*;
+use mu::compiler::*;
+use mu::testutil;
+
+use std::sync::Arc;
+use std::sync::RwLock;
+use mu::testutil::aot;
+
+#[test]
+fn test_inline_add() {
+    let lib = testutil::compile_fncs("add_trampoline", vec!["add_trampoline", "add"], &inline_add);
+
+    unsafe {
+        let inline_add : libloading::Symbol<unsafe extern fn(u64, u64) -> u64> = lib.get(b"add_trampoline").unwrap();
+
+        let inline_add_1_1 = inline_add(1, 1);
+        println!("add(1, 1) = {}", inline_add_1_1);
+        assert!(inline_add_1_1 == 2);
+    }
+}
+
+fn inline_add() -> VM {
+    let vm = VM::new();
+
+    typedef!    ((vm) int64 = mu_int(64));
+    funcsig!    ((vm) sig = (int64, int64) -> (int64));
+
+    funcdecl!   ((vm) <sig> add);
+    {
+        // add
+        funcdef!    ((vm) <sig> add VERSION add_v1);
+
+        block!      ((vm, add_v1) blk_entry);
+        ssa!        ((vm, add_v1) <int64> x);
+        ssa!        ((vm, add_v1) <int64> y);
+
+        ssa!        ((vm, add_v1) <int64> res);
+        inst!       ((vm, add_v1) blk_entry_add:
+            res = BINOP (BinOp::Add) x y
+        );
+
+        inst!       ((vm, add_v1) blk_entry_ret:
+            RET (res)
+        );
+
+        define_block!   ((vm, add_v1) blk_entry(x, y) {blk_entry_add, blk_entry_ret});
+
+        define_func_ver!((vm) add_v1 (entry: blk_entry) {blk_entry});
+    }
+
+    {
+        // add_trampoline
+        typedef!    ((vm) funcref_to_sig = mu_funcref(sig));
+        constdef!   ((vm) <funcref_to_sig> funcref_add = Constant::FuncRef(add));
+
+        funcdecl!   ((vm) <sig> add_trampoline);
+        funcdef!    ((vm) <sig> add_trampoline VERSION add_trampoline_v1);
+
+        block!      ((vm, add_trampoline_v1) tramp_blk_entry);
+        ssa!        ((vm, add_trampoline_v1) <int64> tramp_x);
+        ssa!        ((vm, add_trampoline_v1) <int64> tramp_y);
+
+        consta!     ((vm, add_trampoline_v1) funcref_add_local = funcref_add);
+        ssa!        ((vm, add_trampoline_v1) <int64> tramp_res);
+        inst!       ((vm, add_trampoline_v1) tramp_blk_call:
+            tramp_res = EXPRCALL (CallConvention::Mu, is_abort: false) funcref_add_local (tramp_x, tramp_y)
+        );
+
+        inst!       ((vm, add_trampoline_v1) tramp_blk_ret:
+            RET (tramp_res)
+        );
+
+        define_block!   ((vm, add_trampoline_v1) tramp_blk_entry(tramp_x, tramp_y) {tramp_blk_call, tramp_blk_ret});
+
+        define_func_ver!((vm) add_trampoline_v1 (entry: tramp_blk_entry) {tramp_blk_entry});
+    }
+
+    vm
+}
\ No newline at end of file