zkvmc-core 0.0.1

use std::rc::Rc;

use addr2line::{fallible_iterator::FallibleIterator, Context, LookupResult};
use elf::{abi::STT_FUNC, endian::LittleEndian, ElfBytes};
use eyre::Result;
use gimli::{EndianRcSlice, RunTimeEndian};
use object::File;
use rustc_demangle::demangle;
use rustc_hash::FxHashMap;

type GimliReader = gimli::EndianRcSlice<gimli::RunTimeEndian>;

fn demangle_name(name: String) -> String {
    if let Some(index) = name.rfind("::") {
        let truncated = &name[0..index];
        truncated.to_string()
    } else {
        name
    }
}

fn build_symbol_table(elf_data: &[u8]) -> Result<FxHashMap<u64, String>> {
    let mut symbol_table = FxHashMap::default();

    let elf = ElfBytes::<LittleEndian>::minimal_parse(elf_data)?;
    if let Some((symtab, strtab)) = elf.symbol_table()? {
        for sym in symtab {
            if sym.st_symtype() == STT_FUNC {
                let name = strtab.get(sym.st_name as usize)?;
                symbol_table.insert(sym.st_value, demangle(name).to_string());
            }
        }
    }

    Ok(symbol_table)
}

fn lookup_pc_in_dwarf(pc: u32, ctx: &Context<GimliReader>) -> Vec<Frame> {
    let frames = match ctx.find_frames(pc as u64) {
        LookupResult::Output(result) => result.unwrap(),
        LookupResult::Load {
            load: _,
            continuation: _,
        } => unimplemented!(),
    };
    frames
        .filter_map(|frame| Ok(decode_frame(frame)))
        .collect()
        .unwrap()
}

fn decode_frame(fr: addr2line::Frame<EndianRcSlice<RunTimeEndian>>) -> Option<Frame> {
    Some(Frame {
        name: fr.function.as_ref()?.demangle().ok()?.to_string(),
        lineno: fr.location.as_ref()?.line? as i64,
        filename: fr.location.as_ref()?.file?.to_string(),
    })
}

fn load_dwarf<'data, O: object::Object<'data>>(file: &O) -> Result<gimli::Dwarf<GimliReader>> {
    let endian = if file.is_little_endian() {
        gimli::RunTimeEndian::Little
    } else {
        gimli::RunTimeEndian::Big
    };

    fn load_section<'data, O, Endian>(
        id: gimli::SectionId,
        file: &O,
        endian: Endian,
    ) -> std::result::Result<gimli::EndianRcSlice<Endian>, gimli::Error>
    where
        O: object::Object<'data>,
        Endian: gimli::Endianity,
    {
        use object::ObjectSection as _;

        let data = file
            .section_by_name(id.name())
            .and_then(|section| section.uncompressed_data().ok())
            .unwrap_or(std::borrow::Cow::Borrowed(&[]));
        Ok(gimli::EndianRcSlice::new(Rc::from(&*data), endian))
    }

    Ok(gimli::Dwarf::load(|id| load_section(id, file, endian))?)
}

/// Represents a frame.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Frame {
    /// Function name
    pub name: String,

    /// Line number
    pub lineno: i64,

    /// Filename where this function is defined
    pub filename: String,
}

pub trait LookupPc {
    fn lookup_pc(&self, pc: u32) -> Option<Frame>;
}

pub struct Addr2Frame {
    ctx: Context<GimliReader>,
    symbol_table: FxHashMap<u64, String>,
}

impl Addr2Frame {
    pub fn new(input: &[u8]) -> Result<Self> {
        let file = File::parse(input)?;
        let dwarf = load_dwarf(&file)?;
        let ctx = Context::from_dwarf(dwarf)?;
        let symbol_table = build_symbol_table(input)?;
        Ok(Self { ctx, symbol_table })
    }
}

impl LookupPc for Addr2Frame {
    fn lookup_pc(&self, pc: u32) -> Option<Frame> {
        let dwarf_frames = lookup_pc_in_dwarf(pc, &self.ctx);
        let symbol = self.symbol_table.get(&(pc as u64)).cloned();

        if !dwarf_frames.is_empty() {
            Some(dwarf_frames.last().unwrap().clone())
        } else {
            symbol.map(|symbol| Frame {
                name: demangle_name(symbol).replace('&', ""),
                lineno: 0,
                filename: "unknown".into(),
            })
        }
    }
}