Inline function return variables.

This change replaces 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.
More importantly however, this guards us against a future optimisation
in rustc that allows SIR to assign to $0 multiple times and at the
beginning of a block, which could lead to another function overwriting
its value (see rust-lang/rust#72205).

Author: ptersilie

ykcompile/src/lib.rs

@@ -27,6 +27,8 @@ pub enum CompileError {
     /// We ran out of registers.
     /// In the long-run, when we have a proper register allocator, this won't be needed.
     OutOfRegisters,
+    /// We tried to retrieve the register for a local which doesn't have a register assigned to it.
+    UnknownLocal,
     /// Compiling this statement is not yet implemented.
     /// The string inside is a hint as to what kind of statement needs to be implemented.
     Unimplemented(String),
@@ -38,6 +40,7 @@ impl Display for CompileError {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         match self {
             Self::OutOfRegisters => write!(f, "Ran out of registers"),
+            Self::UnknownLocal => write!(f, "Couldn't find assigned register for local."),
             Self::Unimplemented(s) => write!(f, "Unimplemented compilation: {}", s),
             Self::UnknownSymbol(s) => write!(f, "Unknown symbol: {}", s),
         }
@@ -96,8 +99,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 outer most function and moves its content into RAX at
+    /// the end of the trace.
+    rtn_var: Option<Place>,
 }
 
 impl TraceCompiler {
@@ -125,8 +129,26 @@ impl TraceCompiler {
         }
     }
 
+    /// Returns the currently assigned register for a given `Local`. Similar to `local_to_reg` but
+    /// read-only, i.e. this function doesn't assign `Local`'s to registers.
+    fn get_reg(&self, local: &Local) -> Result<u8, CompileError> {
+        match self.assigned_regs.get(local) {
+            Some(u) => Ok(*u),
+            None => Err(CompileError::UnknownLocal),
+        }
+    }
+
     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 +239,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 +325,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 +400,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 +464,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 +482,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

@@ -166,16 +166,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(body.num_locals, newdest.as_ref().unwrap().clone());
                             TirOp::Statement(Statement::Enter(
                                 op.clone(),
                                 newargs,
@@ -313,33 +313,37 @@ struct VarRenamer {
     /// restored again after leaving that call.
     stack: Vec<u32>,
     /// Current offset used to rename variables.
-    offset: u32
+    offset: u32,
+    returns: Vec<Place>
 }
 
 impl VarRenamer {
     fn new() -> Self {
         VarRenamer {
             stack: vec![0],
-            offset: 0
+            offset: 0,
+            returns: Vec::new()
         }
     }
 
     fn offset(&self) -> u32 {
         self.offset
     }
 
-    fn enter(&mut self, num_locals: usize) {
+    fn enter(&mut self, num_locals: usize, dest: Place) {
         // 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;
         self.stack.push(self.offset);
+        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 +384,17 @@ 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. More importantly however, this guards us
+            // against a future optimisation in rustc that allows SIR to assign to $0 multiple
+            // times and at the beginning of a block, which could lead to another function
+            // overwriting its value (see https://github.com/rust-lang/rust/pull/72205).
+            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);