cargo-treemap 0.1.0

//! Parse a compiled artifact (ELF / Mach-O / PE / WASM) into per-symbol sizes.
//!
//! The size of a symbol is the number of bytes it occupies in its section. Two format
//! realities drive this code:
//!   * Mach-O never records symbol sizes, every `st_size` is 0, so within a section we
//!     sort symbols by address and take the delta to the next symbol (or the section end).
//!   * ELF/PE record real sizes; we still fall back to the delta when a size is 0.

use anyhow::{Context, Result};
use object::{
    BinaryFormat, Object, ObjectSection, ObjectSymbol, SectionKind, SymbolKind, SymbolSection,
};
use rayon::prelude::*;
use std::collections::HashMap;
use std::io::Write;
use std::path::Path;

/// Which size bucket a symbol falls into. Read-only-data and data are folded together.
#[derive(Clone, Copy, PartialEq, Eq)]
pub enum SymClass {
    Text,
    Data,
}

pub struct RawSymbol {
    /// Mangled name, exactly as in the symbol table (used for the rlib map lookup).
    pub name: String,
    /// The symbol's address (for DWARF address-to-source resolution).
    pub address: u64,
    pub class: SymClass,
    pub size: u64,
    pub gzip_size: u64,
}

/// Section-level byte totals for the artifact (authoritative rollup; the sum of symbol
/// sizes will be less because of padding, headers and unattributed slack).
pub struct SectionTotals {
    pub text: u64,
    pub data: u64,
    pub text_attributed: u64,
    pub data_attributed: u64,
}

pub struct ParsedArtifact {
    pub name: String,
    pub format: &'static str,
    pub symbols: Vec<RawSymbol>,
    pub totals: SectionTotals,
}

fn classify(kind: SectionKind) -> Option<SymClass> {
    match kind {
        SectionKind::Text => Some(SymClass::Text),
        SectionKind::Data
        | SectionKind::ReadOnlyData
        | SectionKind::ReadOnlyDataWithRel
        | SectionKind::ReadOnlyString => Some(SymClass::Data),
        _ => None,
    }
}

fn format_name(f: BinaryFormat) -> &'static str {
    match f {
        BinaryFormat::Elf => "ELF",
        BinaryFormat::MachO => "Mach-O",
        BinaryFormat::Pe => "PE",
        BinaryFormat::Coff => "COFF",
        BinaryFormat::Wasm => "WASM",
        BinaryFormat::Xcoff => "XCOFF",
        _ => "unknown",
    }
}

fn gz_len(data: &[u8]) -> u64 {
    if data.is_empty() {
        return 0;
    }
    let mut e = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
    if e.write_all(data).is_err() {
        return 0;
    }
    e.finish().map(|v| v.len() as u64).unwrap_or(0)
}

/// Parse a compiled artifact. `want_gzip` computes the per-symbol gzip metric (the expensive
/// part); pass false when only sizes/names are needed.
pub fn parse(path: &Path, want_gzip: bool) -> Result<ParsedArtifact> {
    let file = std::fs::File::open(path).with_context(|| format!("open {}", path.display()))?;
    let mmap =
        unsafe { memmap2::Mmap::map(&file) }.with_context(|| format!("mmap {}", path.display()))?;
    let obj = object::File::parse(&*mmap)
        .with_context(|| format!("parse object file {}", path.display()))?;
    let is_macho = obj.format() == BinaryFormat::MachO;

    // Group defined symbols by their section index: (address, reported_size, name).
    let mut by_section: HashMap<usize, Vec<(u64, u64, String)>> = HashMap::new();
    for sym in obj.symbols() {
        if !sym.is_definition() {
            continue;
        }
        if !matches!(
            sym.kind(),
            SymbolKind::Text | SymbolKind::Data | SymbolKind::Unknown
        ) {
            continue;
        }
        let SymbolSection::Section(idx) = sym.section() else {
            continue;
        };
        let name = match sym.name() {
            Ok(n) if !n.is_empty() => n.to_string(),
            _ => continue,
        };
        by_section
            .entry(idx.0)
            .or_default()
            .push((sym.address(), sym.size(), name));
    }

    let mut symbols = Vec::new();
    let mut totals = SectionTotals {
        text: 0,
        data: 0,
        text_attributed: 0,
        data_attributed: 0,
    };

    for section in obj.sections() {
        let Some(class) = classify(section.kind()) else {
            continue;
        };
        let ssize = section.size();
        match class {
            SymClass::Text => totals.text += ssize,
            SymClass::Data => totals.data += ssize,
        }

        let Some(mut syms) = by_section.remove(&section.index().0) else {
            continue;
        };
        let saddr = section.address();
        let sec_end = saddr.saturating_add(ssize);
        let sdata = section.data().unwrap_or(&[]);
        // Address asc, then reported size desc, so a same-address alias group leads with its
        // authoritative sized member.
        syms.sort_by(|a, b| a.0.cmp(&b.0).then(b.1.cmp(&a.1)));

        // Phase 1 (serial): resolve each symbol's byte extent. The high-water mark clamps every
        // extent to start at already-attributed bytes, so interior/overlapping labels can't
        // double-count and the per-section sum never exceeds the section size.
        let mut extents: Vec<(String, u64, u64, usize, usize)> = Vec::new();
        let mut prev_addr: Option<u64> = None;
        let mut covered_end = saddr;
        for i in 0..syms.len() {
            let (addr, reported, ref name) = syms[i];
            // A symbol tagged to this section but outside its byte range can't be sized against
            // it, notably wasm data symbols, whose addresses are linear-memory offsets.
            if addr < saddr || addr >= sec_end {
                continue;
            }
            if prev_addr == Some(addr) {
                continue;
            }
            prev_addr = Some(addr);

            // Next distinct address in the section, bounding the delta for size-0 symbols.
            let mut next = sec_end;
            for sym in &syms[i + 1..] {
                if sym.0 > addr {
                    next = sym.0.min(sec_end);
                    break;
                }
            }
            // Mach-O reports every size as 0, so it always takes the delta path.
            let raw_end = if reported > 0 && !is_macho {
                addr.saturating_add(reported).min(sec_end)
            } else {
                next
            };
            let start = addr.max(covered_end);
            let size = raw_end.saturating_sub(start);
            covered_end = covered_end.max(raw_end);
            if size == 0 {
                continue;
            }
            let off = (start - saddr) as usize;
            let end = off.saturating_add(size as usize).min(sdata.len());
            extents.push((name.clone(), addr, size, off, end));
        }

        let section_attr: u64 = extents.iter().map(|e| e.2).sum();
        match class {
            SymClass::Text => totals.text_attributed += section_attr,
            SymClass::Data => totals.data_attributed += section_attr,
        }

        // Phase 2 (parallel): gzip each extent's bytes (the dominant CPU cost). `sdata` is a
        // read-only mmap slice, so parallel reads are sound.
        let gzips: Vec<u64> = if want_gzip {
            extents
                .par_iter()
                .map(|&(_, _, _, off, end)| if off < end { gz_len(&sdata[off..end]) } else { 0 })
                .collect()
        } else {
            vec![0; extents.len()]
        };

        // Phase 3 (serial): emit symbols in address order.
        for ((name, addr, size, _, _), gzip_size) in extents.into_iter().zip(gzips) {
            symbols.push(RawSymbol {
                name,
                address: addr,
                class,
                size,
                gzip_size,
            });
        }
    }

    let name = path
        .file_name()
        .map(|s| s.to_string_lossy().into_owned())
        .unwrap_or_else(|| "artifact".to_string());

    Ok(ParsedArtifact {
        name,
        format: format_name(obj.format()),
        symbols,
        totals,
    })
}