use std::{
collections::BTreeMap,
fmt::{self, Debug, Display, Formatter},
};
#[derive(Clone, Default, PartialEq, Eq)]
pub(crate) struct Line {
pub(crate) rva: u32,
pub(crate) len: u32,
pub(crate) num: u32,
pub(crate) file_id: u32,
}
impl Debug for Line {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"Line {{ rva: {:x}, len: {:x}, line: {}, file_id: {} }}",
self.rva, self.len, self.num, self.file_id
)
}
}
#[derive(Clone, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub(crate) struct InlineSite {
pub(crate) inline_origin_id: u32,
pub(crate) call_depth: u32,
pub(crate) call_line_number: u32,
pub(crate) call_file_id: u32,
}
impl Debug for InlineSite {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"InlineSite {{ inline_origin_id: {}, call_depth: {}, call_line_number: {}, call_file_id: {} }}",
self.inline_origin_id, self.call_depth, self.call_line_number, self.call_file_id
)
}
}
#[derive(Clone, Default)]
pub(crate) struct InlineAddressRange {
pub(crate) rva: u32,
pub(crate) len: u32,
}
impl Debug for InlineAddressRange {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"InlineAddressRange {{ rva: {:x}, len: {:x} }}",
self.rva, self.len
)
}
}
#[derive(Clone, Debug, Default)]
pub struct Lines {
pub(crate) lines: Vec<Line>,
pub(crate) inlines: BTreeMap<InlineSite, Vec<InlineAddressRange>>,
pub(crate) are_lines_sorted: bool,
pub(crate) last_line_rva: u32,
}
fn write_inline_record(
site: &InlineSite,
ranges: &[InlineAddressRange],
f: &mut Formatter<'_>,
) -> fmt::Result {
write!(
f,
"INLINE {} {} {} {}",
site.call_depth, site.call_line_number, site.call_file_id, site.inline_origin_id,
)?;
for range in ranges {
write!(f, " {:x} {:x}", range.rva, range.len)?;
}
writeln!(f)
}
fn write_line_record(line: &Line, f: &mut Formatter<'_>) -> fmt::Result {
writeln!(
f,
"{:x} {:x} {} {}",
line.rva, line.len, line.num, line.file_id
)
}
impl Display for Lines {
fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
let mut inlines: Vec<(&_, &_)> = self.inlines.iter().collect();
inlines.sort_by_key(|(site, ranges)| (ranges.first().unwrap().rva, site.call_depth));
for (site, ranges) in inlines {
write_inline_record(site, ranges, f)?;
}
for line in &self.lines {
write_line_record(line, f)?;
}
Ok(())
}
}
impl Lines {
pub(crate) fn new() -> Self {
Self {
lines: Vec::new(),
inlines: BTreeMap::new(),
are_lines_sorted: true,
last_line_rva: 0,
}
}
pub(crate) fn add_line(&mut self, rva: u32, num: u32, file_id: u32) {
self.lines.push(Line {
rva,
num,
len: 0,
file_id,
});
self.are_lines_sorted = self.are_lines_sorted && self.last_line_rva <= rva;
self.last_line_rva = rva;
}
pub(crate) fn add_inline(&mut self, site: InlineSite, address_range: InlineAddressRange) {
self.inlines
.entry(site)
.or_insert_with(|| Vec::with_capacity(1))
.push(address_range);
}
pub fn finalize(&mut self, sym_rva: u32, sym_len: u32) {
self.ensure_order();
self.compute_len(sym_rva, sym_len);
}
fn compute_len(&mut self, sym_rva: u32, sym_len: u32) {
if self.lines.is_empty() {
return;
}
assert!(
self.are_lines_sorted,
"Call ensure_order() before calling compute_len()"
);
let lens: Vec<u32> = self.lines.windows(2).map(|w| w[1].rva - w[0].rva).collect();
let (last, lines) = self.lines.split_last_mut().unwrap();
lines
.iter_mut()
.zip(lens.iter())
.for_each(|(line, len)| line.len = *len);
if let Some(function_end_rva) = sym_rva.checked_add(sym_len) {
if last.rva < function_end_rva {
last.len = function_end_rva - last.rva;
}
}
}
fn ensure_order(&mut self) {
if !self.are_lines_sorted {
self.lines.sort_by_key(|x| x.rva);
self.are_lines_sorted = true;
}
for ranges in self.inlines.values_mut() {
ranges.sort_by_key(|range| range.rva);
ranges.dedup_by(|next, current| {
if current.rva.checked_add(current.len) == Some(next.rva) {
current.len += next.len;
true
} else {
false
}
})
}
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn no_overflow_when_lines_spill_out_of_function() {
let function_sym_len = 0x9;
let mut lines = Lines::new();
lines.add_line(0x10, 100, 0);
lines.add_line(0x18, 102, 0);
lines.add_line(0x14, 101, 0);
lines.add_line(0x1c, 103, 0);
lines.finalize(0x10, function_sym_len);
assert_eq!(
lines.lines,
vec![
Line {
rva: 0x10,
len: 0x4,
num: 100,
file_id: 0
},
Line {
rva: 0x14,
len: 0x4,
num: 101,
file_id: 0
},
Line {
rva: 0x18,
len: 0x4, num: 102,
file_id: 0
},
Line {
rva: 0x1c,
len: 0, num: 103,
file_id: 0
},
]
);
}
}