Skip to content

Commit

Permalink
Save the number of popped registers.
Browse files Browse the repository at this point in the history
This makes the encoding unambiguous. Fixes #25.

I'm removing the separate UnwindRuleOffsetSpAndPopRegisters type in order
to store the fields more compactly inside UnwindRuleX86_64.
mstange committed Mar 14, 2024
1 parent d3fe8ed commit 6697f09
Showing 4 changed files with 69 additions and 44 deletions.
1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
@@ -24,6 +24,7 @@ arrayvec = "0.7.4"
[dev-dependencies]
object = "0.32"
flate2 = "1.0.28"
itertools = "0.12.1"

[profile.release]
debug = true
20 changes: 20 additions & 0 deletions src/rule_cache.rs
Original file line number Diff line number Diff line change
@@ -122,3 +122,23 @@ impl CacheStats {
self.miss_empty_slot_count + self.miss_wrong_modules_count + self.miss_wrong_address_count
}
}

#[cfg(test)]
mod tests {
use crate::{aarch64::UnwindRuleAarch64, x86_64::UnwindRuleX86_64};

use super::*;

// Ensure that the size of Option<CacheEntry<UnwindRuleX86_64>> doesn't change by accident.
#[test]
fn test_cache_entry_size() {
assert_eq!(
std::mem::size_of::<Option<CacheEntry<UnwindRuleX86_64>>>(),
16
);
assert_eq!(
std::mem::size_of::<Option<CacheEntry<UnwindRuleAarch64>>>(),
24 // <-- larger than we'd like
);
}
}
37 changes: 27 additions & 10 deletions src/x86_64/register_ordering.rs
Original file line number Diff line number Diff line change
@@ -12,7 +12,7 @@ const ENCODE_REGISTERS: [Reg; 8] = [
Reg::R15,
];

pub fn decode(encoded_ordering: u16) -> ArrayVec<Reg, 8> {
pub fn decode(count: u8, encoded_ordering: u16) -> ArrayVec<Reg, 8> {
let mut regs: ArrayVec<Reg, 8> = ENCODE_REGISTERS.into();
let mut r = encoded_ordering;
let mut n: u16 = 8;
@@ -24,15 +24,16 @@ pub fn decode(encoded_ordering: u16) -> ArrayVec<Reg, 8> {
r /= n;
n -= 1;
}
regs.truncate(8 - n as usize);
regs.truncate(count as usize);
regs
}

pub fn encode(registers: &[Reg]) -> Option<u16> {
pub fn encode(registers: &[Reg]) -> Option<(u8, u16)> {
if registers.len() > ENCODE_REGISTERS.len() {
return None;
}

let count = registers.len() as u8;
let mut r: u16 = 0;
let mut reg_order: ArrayVec<Reg, 8> = ENCODE_REGISTERS.into();

@@ -45,23 +46,39 @@ pub fn encode(registers: &[Reg]) -> Option<u16> {
r += index as u16 * scale;
scale *= 8 - i as u16;
}
Some(r)
Some((count, r))
}

#[cfg(test)]
mod test {
use super::*;

#[test]
fn register_compression() {
fn unhandled_orderings() {
use super::Reg::*;

assert_eq!(encode(&[RAX]), None, "RAX is a volatile register, i.e. not a callee-save register, so it does not need to be restored during epilogs and is not covered by the encoding.");
assert_eq!(encode(&[RSI, RSI]), None, "Valid register orderings only contain each register (at most) once, so there is no encoding for a sequence with repeated registers.");
assert_eq!(
decode(encode(&[RSI, R12, R15, R14, RBX]).unwrap()).as_slice(),
&[RSI, R12, R15, R14, RBX],
"This particular register ordering should roundtrip successfully"
);
}

#[test]
fn roundtrip_all() {
// Test all possible register orderings.
// That is, for all permutations of length 0 to 8 of the ENCODE_REGISTERS array, check that
// the register ordering rountrips successfully through encoding and decoding.
use itertools::Itertools;
for permutation in (0..=8).flat_map(|k| ENCODE_REGISTERS.iter().cloned().permutations(k)) {
let permutation = permutation.as_slice();
let encoding = encode(permutation);
if let Some((count, encoded)) = encoding {
assert_eq!(
decode(count, encoded).as_slice(),
permutation,
"Register permutation should roundtrip correctly",
);
} else {
panic!("Register permutation failed to encode: {permutation:?}");
}
}
}
}
55 changes: 21 additions & 34 deletions src/x86_64/unwind_rule.rs
Original file line number Diff line number Diff line change
@@ -31,7 +31,14 @@ pub enum UnwindRuleX86_64 {
///
/// The registers are stored in a separate compressed ordering to facilitate restoring register
/// values if desired. If not for this we could simply store the total offset.
OffsetSpAndPopRegisters(UnwindRuleOffsetSpAndPopRegisters),
OffsetSpAndPopRegisters {
/// The additional stack pointer offset to undo before popping the registers, divided by 8 bytes.
sp_offset_by_8: u16,
/// The number of registers to pop from the stack.
register_count: u8,
/// An encoded ordering of the callee-save registers to pop from the stack, see register_ordering.
encoded_registers_to_pop: u16,
},
}

pub enum OffsetOrPop {
@@ -69,40 +76,16 @@ impl UnwindRuleX86_64 {
if regs.is_empty() && sp_offset_by_8 == 0 {
Some(Self::JustReturn)
} else {
let rule = UnwindRuleOffsetSpAndPopRegisters::try_new(sp_offset_by_8, &regs)?;
Some(Self::OffsetSpAndPopRegisters(rule))
let (register_count, encoded_registers_to_pop) = register_ordering::encode(&regs)?;
Some(Self::OffsetSpAndPopRegisters {
sp_offset_by_8,
register_count,
encoded_registers_to_pop,
})
}
}
}

#[derive(Default, Clone, Copy, Debug, PartialEq, Eq)]
pub struct UnwindRuleOffsetSpAndPopRegisters {
/// The additional stack pointer offset to undo before popping the registers, divided by 8 bytes.
pub sp_offset_by_8: u16,
/// An encoded ordering of the callee-save registers to pop from the stack, see register_ordering.
pub encoded_registers_to_pop: u16,
}

impl UnwindRuleOffsetSpAndPopRegisters {
pub fn try_new(sp_offset_by_8: u16, registers_to_pop: &[Reg]) -> Option<Self> {
let encoded_registers_to_pop = register_ordering::encode(registers_to_pop)?;
Some(Self {
sp_offset_by_8,
encoded_registers_to_pop,
})
}

/// Return the initial stack pointer offset, in bytes.
pub fn sp_offset(&self) -> u64 {
self.sp_offset_by_8 as u64 * 8
}

/// Return the ordered sequence of registers to pop.
pub fn registers_to_pop(&self) -> ArrayVec<Reg, 8> {
register_ordering::decode(self.encoded_registers_to_pop)
}
}

impl UnwindRule for UnwindRuleX86_64 {
type UnwindRegs = UnwindRegsX86_64;

@@ -236,12 +219,16 @@ impl UnwindRule for UnwindRuleX86_64 {

(new_sp, new_bp)
}
UnwindRuleX86_64::OffsetSpAndPopRegisters(r) => {
UnwindRuleX86_64::OffsetSpAndPopRegisters {
sp_offset_by_8,
register_count,
encoded_registers_to_pop,
} => {
let sp = regs.sp();
let mut sp = sp
.checked_add(r.sp_offset())
.checked_add(sp_offset_by_8 as u64 * 8)
.ok_or(Error::IntegerOverflow)?;
for reg in r.registers_to_pop() {
for reg in register_ordering::decode(register_count, encoded_registers_to_pop) {
let value = read_stack(sp).map_err(|_| Error::CouldNotReadStack(sp))?;
sp = sp.checked_add(8).ok_or(Error::IntegerOverflow)?;
regs.set(reg, value);

0 comments on commit 6697f09

Please sign in to comment.