Replace return variable with its destination.

ykcompile/src/lib.rs

@@ -96,8 +96,9 @@ pub struct TraceCompiler {
     available_regs: Vec<u8>,
     /// Maps locals to their assigned registers.
     assigned_regs: HashMap<Local, u8>,
-    /// Stores the destination locals to which we copy RAX to after leaving an inlined call.
-    leaves: Vec<Option<Place>>,
+    /// Stores the destination local of the outermost function and moves its content into RAX at
+    /// the end of the trace.
+    rtn_var: Option<Place>,
 }
 
 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();

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,24 +162,24 @@ 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.
 
                             // Rename the destination if there is one.
                             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,33 +317,53 @@ struct VarRenamer {
     /// restored again after leaving that call.
     stack: Vec<u32>,
     /// 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<u32>,
+    /// Stores the return variables of inlined function calls. Used to replace `$0` during
+    /// renaming.
+    returns: Vec<Place>
 }
 
 impl VarRenamer {
     fn new() -> Self {
         VarRenamer {
             stack: vec![0],
-            offset: 0
+            offset: 0,
+            acc: None,
+            returns: Vec::new()
         }
     }
 
     fn offset(&self) -> u32 {
         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);