cargo-treemap 0.1.0

//! Attribute each symbol to the crate that owns it.
//!
//! Priority: (1) exact match from the dependency/std rlib symbol tables, authoritative,
//! because it says which crate's rlib actually emitted the symbol (this is what correctly
//! blames a monomorphized generic on the instantiating crate); (2) the first identifier in
//! the demangled name that names a known crate, a heuristic that inlining and
//! monomorphization can fool; (3) a `[runtime]`/`[unknown]` bucket for symbols that name no
//! crate at all (compiler-outlined functions, `_main`, EH machinery).

use crate::dwarf::DwarfResolver;
use crate::util::std_crates;
use object::{Object, ObjectSymbol};
use std::collections::{HashMap, HashSet};
use std::path::Path;

#[derive(Clone, Copy, PartialEq, Eq)]
pub enum Confidence {
    Exact,
    Heuristic,
    Unknown,
}

impl Confidence {
    pub fn as_str(self) -> &'static str {
        match self {
            Confidence::Exact => "exact",
            Confidence::Heuristic => "heuristic",
            Confidence::Unknown => "unknown",
        }
    }
}

/// Maps a mangled symbol name to the crate whose rlib defined it.
pub struct RlibMap {
    map: HashMap<String, String>,
}

impl RlibMap {
    pub fn build(rlibs: &[(String, std::path::PathBuf)]) -> Self {
        let mut map = HashMap::new();
        for (crate_name, path) in rlibs {
            if let Err(e) = index_rlib(&mut map, crate_name, path) {
                eprintln!(
                    "cargo-treemap: warning: could not index {}: {e:#}",
                    path.display()
                );
            }
        }
        RlibMap { map }
    }

    fn get(&self, sym: &str) -> Option<&str> {
        self.map.get(sym).map(String::as_str)
    }
}

fn index_rlib(
    map: &mut HashMap<String, String>,
    crate_name: &str,
    path: &Path,
) -> anyhow::Result<()> {
    let file = std::fs::File::open(path)?;
    let mmap = unsafe { memmap2::Mmap::map(&file) }?;
    let archive = object::read::archive::ArchiveFile::parse(&*mmap)?;
    for member in archive.members() {
        let Ok(member) = member else { continue };
        let Ok(data) = member.data(&*mmap) else {
            continue;
        };
        let Ok(obj) = object::File::parse(data) else {
            continue;
        };
        for sym in obj.symbols() {
            if !sym.is_definition() {
                continue;
            }
            if let Ok(name) = sym.name() {
                if !name.is_empty() {
                    map.entry(name.to_string())
                        .or_insert_with(|| crate_name.to_string());
                }
            }
        }
    }
    Ok(())
}

pub struct Attributor<'a> {
    rlib: &'a RlibMap,
    crates: &'a HashSet<String>,
    std_set: HashSet<&'static str>,
    split_std: bool,
}

pub struct Attributed {
    pub crate_name: String,
    pub confidence: Confidence,
    pub demangled: String,
}

impl<'a> Attributor<'a> {
    pub fn new(rlib: &'a RlibMap, crates: &'a HashSet<String>, split_std: bool) -> Self {
        Attributor {
            rlib,
            crates,
            std_set: std_crates(),
            split_std,
        }
    }

    fn canon(&self, name: &str) -> String {
        // Collapse into `std` only when the crate isn't a real direct dependency; otherwise a
        // user's own `object`/`gimli`/etc. would be mislabeled as std.
        if !self.split_std && self.std_set.contains(name) && !self.crates.contains(name) {
            "std".to_string()
        } else {
            name.to_string()
        }
    }

    /// First identifier token anywhere in the demangled name that names a known crate.
    /// Walking every token (not just segment heads) is what lets `<str as core::fmt::
    /// Display>::fmt` resolve to `core` rather than the primitive `str`.
    fn first_known_ident(&self, s: &str) -> Option<String> {
        let b = s.as_bytes();
        let mut i = 0;
        while i < b.len() {
            let c = b[i];
            if c.is_ascii_alphabetic() || c == b'_' {
                let start = i;
                while i < b.len() && (b[i].is_ascii_alphanumeric() || b[i] == b'_') {
                    i += 1;
                }
                let tok = &s[start..i];
                if self.crates.contains(tok) || self.std_set.contains(tok) {
                    return Some(tok.to_string());
                }
            } else {
                i += 1;
            }
        }
        None
    }

    pub fn attribute(&self, mangled: &str, address: u64, dwarf: Option<&DwarfResolver>) -> Attributed {
        // `{:#}` drops the trailing legacy hash / v0 disambiguator suffix.
        let demangled = format!("{:#}", rustc_demangle::demangle(mangled));

        // DWARF source-file attribution is the most accurate, it places monomorphized code
        // and anonymous rodata by where they actually come from.
        if let Some(dw) = dwarf {
            if let Some(c) = dw.crate_for_addr(address) {
                return Attributed {
                    crate_name: self.canon(&c),
                    confidence: Confidence::Exact,
                    demangled,
                };
            }
        }

        if let Some(c) = self.rlib.get(mangled) {
            return Attributed {
                crate_name: self.canon(c),
                confidence: Confidence::Exact,
                demangled,
            };
        }
        if let Some(c) = self.first_known_ident(&demangled) {
            return Attributed {
                crate_name: self.canon(&c),
                confidence: Confidence::Heuristic,
                demangled,
            };
        }
        // Names no crate: a compiler/runtime symbol (outlined fn, `_main`, EH) if it has no
        // path, else a genuinely-unknown Rust path.
        let bucket = if demangled.contains("::") {
            "[unknown]"
        } else {
            "[runtime]"
        };
        Attributed {
            crate_name: bucket.to_string(),
            confidence: Confidence::Unknown,
            demangled,
        }
    }
}

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

    fn crates(names: &[&str]) -> HashSet<String> {
        names.iter().map(|s| s.to_string()).collect()
    }

    #[test]
    fn attribution_priority_and_buckets() {
        let rlib = RlibMap::build(&[]); // no rlibs; forces the heuristic path
        let set = crates(&["regex_automata", "serde"]);
        let attr = Attributor::new(&rlib, &set, false);

        // first known crate ident in the path wins
        assert_eq!(attr.attribute("regex_automata::dfa::search", 0, None).crate_name, "regex_automata");
        // impl symbol: the trait crate (core) is found deep in the path, collapsed to std
        assert_eq!(attr.attribute("<str as core::fmt::Display>::fmt", 0, None).crate_name, "std");
        // a rust path naming no known crate -> [unknown]
        assert_eq!(attr.attribute("mystery::thing", 0, None).crate_name, "[unknown]");
        // no path at all (e.g. a C/runtime symbol) -> [runtime]
        assert_eq!(attr.attribute("_main", 0, None).crate_name, "[runtime]");
    }

    #[test]
    fn split_std_keeps_core_separate() {
        let rlib = RlibMap::build(&[]);
        let set = crates(&[]);
        let attr = Attributor::new(&rlib, &set, true); // --split-std
        assert_eq!(attr.attribute("core::ptr::drop_in_place", 0, None).crate_name, "core");
    }

    #[test]
    fn real_dep_named_like_std_vendored_is_not_collapsed() {
        // `object` is std-vendored, but if it's a real dependency it must stay itself.
        let rlib = RlibMap::build(&[]);
        let set = crates(&["object"]);
        let attr = Attributor::new(&rlib, &set, false);
        assert_eq!(attr.attribute("object::read::parse", 0, None).crate_name, "object");
    }
}