diff --git a/ykcompile/src/lib.rs b/ykcompile/src/lib.rs index b4eff2669..6e9789599 100644 --- a/ykcompile/src/lib.rs +++ b/ykcompile/src/lib.rs @@ -96,8 +96,9 @@ pub struct TraceCompiler { available_regs: Vec, /// Maps locals to their assigned registers. assigned_regs: HashMap, - /// Stores the destination locals to which we copy RAX to after leaving an inlined call. - leaves: Vec>, + /// Stores the destination local of the outermost function and moves its content into RAX at + /// the end of the trace. + rtn_var: Option, } impl TraceCompiler { @@ -127,6 +128,15 @@ impl TraceCompiler { fn free_register(&mut self, local: &Local) { if let Some(reg) = self.assigned_regs.remove(local) { + if local == &self.rtn_var.as_ref().unwrap().local { + // We currently assume that we only trace a single function which leaves its return + // value in RAX. Since we now inline a function's return variable this won't happen + // automatically anymore. To keep things working, we thus copy the return value of + // the most outer function into RAX at the end of the trace. + dynasm!(self.asm + ; mov rax, Rq(reg) + ); + } self.available_regs.push(reg); } } @@ -217,19 +227,9 @@ impl TraceCompiler { }, } } - // Remember the return destination. - self.leaves.push(dest.as_ref().cloned()); - Ok(()) - } - - fn c_leave(&mut self) -> Result<(), CompileError> { - let dest = self.leaves.pop(); - if let Some(d) = dest { - if let Some(d) = d { - // When we see a leave statement move whatever's left in RAX into the destination - // local. - self.mov_local_local(d.local, Local(0))?; - } + if self.rtn_var.is_none() { + // Remember the return variable of the most outer function. + self.rtn_var = dest.as_ref().cloned(); } Ok(()) } @@ -313,7 +313,7 @@ impl TraceCompiler { }; } Statement::Enter(op, args, dest, off) => self.c_enter(op, args, dest, *off)?, - Statement::Leave => self.c_leave()?, + Statement::Leave => {} Statement::StorageLive(_) => {} Statement::StorageDead(l) => self.free_register(l), Statement::Call(target, args, dest) => self.c_call(target, args, dest)?, @@ -388,7 +388,7 @@ impl TraceCompiler { // Use all the 64-bit registers we can (R15-R8, RDX, RCX). available_regs: vec![15, 14, 13, 12, 11, 10, 9, 8, 2, 1], assigned_regs: HashMap::new(), - leaves: Vec::new(), + rtn_var: None, }; for i in 0..tt.len() { @@ -452,7 +452,7 @@ mod tests { asm: dynasmrt::x64::Assembler::new().unwrap(), available_regs: vec![15, 14, 13, 12, 11, 10, 9, 8, 2, 1], assigned_regs: HashMap::new(), - leaves: Vec::new(), + rtn_var: None, }; for _ in 0..32 { @@ -470,7 +470,7 @@ mod tests { asm: dynasmrt::x64::Assembler::new().unwrap(), available_regs: vec![15, 14, 13, 12, 11, 10, 9, 8, 2, 1], assigned_regs: HashMap::new(), - leaves: Vec::new(), + rtn_var: None, }; let mut seen = HashSet::new(); diff --git a/yktrace/src/tir.rs b/yktrace/src/tir.rs index fb75c5081..7d2186ed1 100644 --- a/yktrace/src/tir.rs +++ b/yktrace/src/tir.rs @@ -117,6 +117,10 @@ impl TirTrace { } }; + // Initialise VarRenamer's accumulator (and thus also set the first offset) to the + // traces most outer number of locals. + rnm.init_acc(body.num_locals); + // When adding statements to the trace, we clone them (rather than referencing the // statements in the SIR) so that we have the freedom to mutate them later. let user_bb_idx_usize = usize::try_from(loc.bb_idx).unwrap(); @@ -158,7 +162,7 @@ impl TirTrace { } => { if let Some(callee_sym) = op.symbol() { // We know the symbol name of the callee at least. - let op = if SIR.bodies.contains_key(callee_sym) { + let op = if let Some(callbody) = SIR.bodies.get(callee_sym) { // We have SIR for the callee, so it will appear inlined in the trace // and we only need to emit Enter/Leave statements. @@ -166,16 +170,16 @@ impl TirTrace { let newdest = dest .as_ref() .map(|(ret_val, _ret_bb)| rnm.rename_place(&ret_val)); - // Rename all `Local`s within the arguments. - let newargs = rnm.rename_args(&args); - // Inform VarRenamer about this function's offset, which is equal to the - // number of variables assigned in the outer body. - rnm.enter(body.num_locals); // FIXME It seems that calls always have a destination despite it being // an `Option`. If this is not always the case, we may want add the // `Local` offset (`var_len`) to this statement so we can assign the // arguments to the correct `Local`s during trace compilation. assert!(newdest.is_some()); + // Rename all `Local`s within the arguments. + let newargs = rnm.rename_args(&args); + // Inform VarRenamer about this function's offset, which is equal to the + // number of variables assigned in the outer body. + rnm.enter(callbody.num_locals, newdest.as_ref().unwrap().clone()); TirOp::Statement(Statement::Enter( op.clone(), newargs, @@ -313,14 +317,22 @@ struct VarRenamer { /// restored again after leaving that call. stack: Vec, /// Current offset used to rename variables. - offset: u32 + offset: u32, + /// Accumulator keeping track of total number of variables used. Needed to use different + /// offsets for consecutive inlined function calls. + acc: Option, + /// Stores the return variables of inlined function calls. Used to replace `$0` during + /// renaming. + returns: Vec } impl VarRenamer { fn new() -> Self { VarRenamer { stack: vec![0], - offset: 0 + offset: 0, + acc: None, + returns: Vec::new() } } @@ -328,18 +340,30 @@ impl VarRenamer { self.offset } - fn enter(&mut self, num_locals: usize) { - // When entering an inlined function call update the current offset by adding the number of - // assigned variables in the outer context. Also add this offset to the stack, so we can - // restore it once we leave the inlined function call again. - self.offset += num_locals as u32; + fn init_acc(&mut self, num_locals: usize) { + if self.acc.is_none() { + self.acc.replace(num_locals as u32); + } + } + + fn enter(&mut self, num_locals: usize, dest: Place) { + // When entering an inlined function call set the offset to the current accumulator. Then + // increment the accumulator by the number of locals in the current function. Also add the + // offset to the stack, so we can restore it once we leave the inlined function call again. + self.offset = self.acc.unwrap(); self.stack.push(self.offset); + match self.acc.as_mut() { + Some(v) => *v += num_locals as u32, + None => {} + } + self.returns.push(dest); } fn leave(&mut self) { // When we leave an inlined function call, we pop the previous offset from the stack, // reverting the offset to what it was before the function was entered. self.stack.pop(); + self.returns.pop(); if let Some(v) = self.stack.last() { self.offset = *v; } else { @@ -380,8 +404,14 @@ impl VarRenamer { fn rename_place(&self, place: &Place) -> Place { if &place.local == &Local(0) { - // Local(0) is used for returning values from calls, so it doesn't need to be renamed. - place.clone() + // Replace the default return variable $0 with the variable in the outer context where + // the return value will end up after leaving the function. This saves us an + // instruction when we compile the trace. + if let Some(v) = self.returns.last() { + v.clone() + } else { + panic!("Expected return value!") + } } else { let mut p = place.clone(); p.local = self.rename_local(&p.local);