Merge #67

67: Replace return variable with its destination. r=vext01 a=ptersilie

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).

While we are here, also fix a bug in the variable renaming code.
Currently, functions that are called consecutively (i.e. they are not                                                                                                            
nested) use the same offset to rename their variables (the offset of                                                                                                             
their outer context), which leads to them sharing the same variable                                                                                                              
names. By using an accumulator which is continuously increased to                                                                                                                
calculate the offset, we make sure consecutive functions have increasing                                                                                                         
variable names, even when after leaving a function the offset is                                                                                                                 
temporarily reset to that of the outer context.

Co-authored-by: Lukas Diekmann <[email protected]>

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);