zkvmc_core/
addr2frame.rs

1use std::rc::Rc;
2
3use addr2line::{fallible_iterator::FallibleIterator, Context, LookupResult};
4use elf::{abi::STT_FUNC, endian::LittleEndian, ElfBytes};
5use eyre::Result;
6use gimli::{EndianRcSlice, RunTimeEndian};
7use object::File;
8use rustc_demangle::demangle;
9use rustc_hash::FxHashMap;
10
11type GimliReader = gimli::EndianRcSlice<gimli::RunTimeEndian>;
12
13fn demangle_name(name: String) -> String {
14    if let Some(index) = name.rfind("::") {
15        let truncated = &name[0..index];
16        truncated.to_string()
17    } else {
18        name
19    }
20}
21
22fn build_symbol_table(elf_data: &[u8]) -> Result<FxHashMap<u64, String>> {
23    let mut symbol_table = FxHashMap::default();
24
25    let elf = ElfBytes::<LittleEndian>::minimal_parse(elf_data)?;
26    if let Some((symtab, strtab)) = elf.symbol_table()? {
27        for sym in symtab {
28            if sym.st_symtype() == STT_FUNC {
29                let name = strtab.get(sym.st_name as usize)?;
30                symbol_table.insert(sym.st_value, demangle(name).to_string());
31            }
32        }
33    }
34
35    Ok(symbol_table)
36}
37
38fn lookup_pc_in_dwarf(pc: u32, ctx: &Context<GimliReader>) -> Vec<Frame> {
39    let frames = match ctx.find_frames(pc as u64) {
40        LookupResult::Output(result) => result.unwrap(),
41        LookupResult::Load {
42            load: _,
43            continuation: _,
44        } => unimplemented!(),
45    };
46    frames
47        .filter_map(|frame| Ok(decode_frame(frame)))
48        .collect()
49        .unwrap()
50}
51
52fn decode_frame(fr: addr2line::Frame<EndianRcSlice<RunTimeEndian>>) -> Option<Frame> {
53    Some(Frame {
54        name: fr.function.as_ref()?.demangle().ok()?.to_string(),
55        lineno: fr.location.as_ref()?.line? as i64,
56        filename: fr.location.as_ref()?.file?.to_string(),
57    })
58}
59
60fn load_dwarf<'data, O: object::Object<'data>>(file: &O) -> Result<gimli::Dwarf<GimliReader>> {
61    let endian = if file.is_little_endian() {
62        gimli::RunTimeEndian::Little
63    } else {
64        gimli::RunTimeEndian::Big
65    };
66
67    fn load_section<'data, O, Endian>(
68        id: gimli::SectionId,
69        file: &O,
70        endian: Endian,
71    ) -> std::result::Result<gimli::EndianRcSlice<Endian>, gimli::Error>
72    where
73        O: object::Object<'data>,
74        Endian: gimli::Endianity,
75    {
76        use object::ObjectSection as _;
77
78        let data = file
79            .section_by_name(id.name())
80            .and_then(|section| section.uncompressed_data().ok())
81            .unwrap_or(std::borrow::Cow::Borrowed(&[]));
82        Ok(gimli::EndianRcSlice::new(Rc::from(&*data), endian))
83    }
84
85    Ok(gimli::Dwarf::load(|id| load_section(id, file, endian))?)
86}
87
88/// Represents a frame.
89#[derive(Clone, Debug, PartialEq, Eq)]
90pub struct Frame {
91    /// Function name
92    pub name: String,
93
94    /// Line number
95    pub lineno: i64,
96
97    /// Filename where this function is defined
98    pub filename: String,
99}
100
101pub trait LookupPc {
102    fn lookup_pc(&self, pc: u32) -> Option<Frame>;
103}
104
105pub struct Addr2Frame {
106    ctx: Context<GimliReader>,
107    symbol_table: FxHashMap<u64, String>,
108}
109
110impl Addr2Frame {
111    pub fn new(input: &[u8]) -> Result<Self> {
112        let file = File::parse(input)?;
113        let dwarf = load_dwarf(&file)?;
114        let ctx = Context::from_dwarf(dwarf)?;
115        let symbol_table = build_symbol_table(input)?;
116        Ok(Self { ctx, symbol_table })
117    }
118}
119
120impl LookupPc for Addr2Frame {
121    fn lookup_pc(&self, pc: u32) -> Option<Frame> {
122        let dwarf_frames = lookup_pc_in_dwarf(pc, &self.ctx);
123        let symbol = self.symbol_table.get(&(pc as u64)).cloned();
124
125        if !dwarf_frames.is_empty() {
126            Some(dwarf_frames.last().unwrap().clone())
127        } else {
128            symbol.map(|symbol| Frame {
129                name: demangle_name(symbol).replace('&', ""),
130                lineno: 0,
131                filename: "unknown".into(),
132            })
133        }
134    }
135}
zkvmc_core/addr2frame.rs

zkvmc_core/
addr2frame.rs
