cargo-treemap 0.1.0

mod apiusage;
mod artifact;
mod attribute;
mod builder;
mod cli;
mod diff;
mod dwarf;
mod model;
mod output;
mod parity;
mod util;

use anyhow::{bail, Context, Result};
use clap::Parser;
use cli::CargoCli;

fn main() -> Result<()> {
    let CargoCli::Treemap(args) = CargoCli::parse();

    if !args.artifact.is_empty() {
        return run_artifacts(&args);
    }

    if args.attribute_features {
        return run_attribute_features(&args);
    }

    eprintln!("cargo-treemap: building (this runs a real cargo build)…");
    let build = builder::build(&args, false, None)?;
    if build.artifacts.is_empty() {
        bail!(
            "no binary/cdylib artifacts were produced.\n\
             Try `--bin <name>`, or `--lib` for a cdylib, or check the package has a binary target."
        );
    }

    // cargo metadata up front: its dep-kind classification lets us exclude build/proc-macro/host
    // rlibs (never linked into the target) from the attribution map, otherwise their symbol
    // names collide with real ones and get misattributed. It also drives parity + api-usage.
    let mut mc = args.manifest.metadata();
    args.features.forward_metadata(&mut mc);
    let md = mc.exec().context("running cargo metadata")?;
    let (selected, _excluded) = args.workspace.partition_packages(&md);

    let linked = parity::linked_crate_names(&md);
    let rlibs: Vec<(String, std::path::PathBuf)> = build
        .rlibs
        .iter()
        .filter(|(name, _)| linked.contains(name))
        .cloned()
        .collect();
    let rlib_map = attribute::RlibMap::build(&rlibs);
    let attributor = attribute::Attributor::new(&rlib_map, &linked, args.split_std);

    let mut nodes = Vec::new();
    let mut format = String::new();
    for path in &build.artifacts {
        let parsed = artifact::parse(path, true)?;
        if format.is_empty() {
            format = parsed.format.to_string();
        }
        let resolver = load_dwarf(args.dwarf, path);
        nodes.push(model::build_tree(&parsed, &attributor, resolver.as_ref()));
    }

    let artifact_names: Vec<String> = nodes.iter().map(|n| n.label.clone()).collect();
    let mut root = model::combine(nodes);

    let (invocation, cwd, generated_unix) = provenance();
    let meta = model::Meta {
        tool: format!("cargo-treemap {}", env!("CARGO_PKG_VERSION")),
        artifacts: artifact_names,
        format,
        target: build.target.clone().unwrap_or_default(),
        profile: build.profile.clone(),
        total_text: root.text,
        total_data: root.data,
        total_gzip: root.gzip,
        default_metric: args.metric.field().to_string(),
        api_accurate: args.api_usage && args.accurate,
        invocation,
        cwd,
        generated_unix,
    };

    // Compare against a git ref: build it in a throwaway worktree, then diff.
    if let Some(git_ref) = &args.compare_ref {
        let base_json = build_ref_baseline(&args, &md, git_ref)?;
        let base = diff::load_base_tree(&base_json)?;
        let _ = std::fs::remove_file(&base_json);
        let d = diff::diff_trees(&base, &root);
        let label = format!("git:{git_ref}");
        return output::emit_diff(&d, &meta, std::path::Path::new(&label), &args);
    }

    // Compare against a saved baseline: diff and emit the delta view instead of the report.
    if let Some(base_path) = &args.compare {
        let base = diff::load_base_tree(base_path)?;
        let d = diff::diff_trees(&base, &root);
        return output::emit_diff(&d, &meta, base_path, &args);
    }

    let wanted: std::collections::HashSet<String> =
        model::crate_bin_totals(&root).into_keys().collect();
    let mut cmetrics = parity::compute(&md, &wanted, &linked);
    if args.compile_time {
        eprintln!("cargo-treemap: --compile-time: full from-scratch timed build (slow)…");
        match parity::compile_times(&args, &md.target_directory) {
            Ok(times) => {
                for (name, ms) in times {
                    if let Some(m) = cmetrics.get_mut(&name) {
                        m.compile_ms = ms;
                    }
                }
            }
            Err(e) => eprintln!("cargo-treemap: compile-time unavailable: {e:#}"),
        }
    }
    model::attach_metrics(&mut root, &cmetrics);
    let crates = parity::crate_rows(&root, &cmetrics);

    let deps = if args.api_usage {
        let used_override = if args.accurate {
            eprintln!("cargo-treemap: --accurate: opt-level-0 v0 build for a truer usage set…");
            let ub = builder::build(&args, true, None)?;
            Some(apiusage::crate_fns(&ub, args.split_std)?)
        } else {
            None
        };
        apiusage::analyze(&md, &selected, &root, used_override.as_ref(), &args)?
    } else {
        Vec::new()
    };

    output::emit(&root, &meta, &deps, &crates, &args)?;
    Ok(())
}

/// Analyze pre-built binaries given via `--artifact`, without running a build.
fn run_artifacts(args: &cli::TreemapArgs) -> Result<()> {
    let mut parsed = Vec::new();
    for p in &args.artifact {
        eprintln!("cargo-treemap: analyzing {}", p.display());
        parsed.push(artifact::parse(p, true)?);
    }
    // No build ⇒ no rlib map / cargo metadata; recover a crate-name set from the symbols.
    let crate_names = derive_crate_names(&parsed);
    let rlib_map = attribute::RlibMap::build(&[]);
    let attributor = attribute::Attributor::new(&rlib_map, &crate_names, args.split_std);

    let mut nodes = Vec::new();
    let mut format = String::new();
    for (path, p) in args.artifact.iter().zip(&parsed) {
        if format.is_empty() {
            format = p.format.to_string();
        }
        let resolver = load_dwarf(args.dwarf, path);
        nodes.push(model::build_tree(p, &attributor, resolver.as_ref()));
    }
    let artifact_names: Vec<String> = nodes.iter().map(|n| n.label.clone()).collect();
    let root = model::combine(nodes);

    if args.api_usage {
        eprintln!("cargo-treemap: --api-usage is ignored with --artifact (needs a build + metadata)");
    }

    let (invocation, cwd, generated_unix) = provenance();
    let meta = model::Meta {
        tool: format!("cargo-treemap {}", env!("CARGO_PKG_VERSION")),
        artifacts: artifact_names,
        format,
        target: String::new(),
        profile: "prebuilt".to_string(),
        total_text: root.text,
        total_data: root.data,
        total_gzip: root.gzip,
        default_metric: args.metric.field().to_string(),
        api_accurate: false,
        invocation,
        cwd,
        generated_unix,
    };
    output::emit(&root, &meta, &[], &[], args)
}

/// Build once per Cargo feature (toggled on against a no-default-features baseline) and report
/// the marginal binary-size cost of each feature.
fn run_attribute_features(args: &cli::TreemapArgs) -> Result<()> {
    let mut mc = args.manifest.metadata();
    args.features.forward_metadata(&mut mc);
    let md = mc.exec().context("running cargo metadata")?;
    let (selected, _excluded) = args.workspace.partition_packages(&md);

    let mut features: Vec<String> = Vec::new();
    for pkg in &selected {
        for f in pkg.features.keys() {
            if f != "default" && !features.contains(f) {
                features.push(f.clone());
            }
        }
    }
    features.sort();
    if features.is_empty() {
        bail!("the selected package(s) declare no Cargo features to attribute");
    }

    eprintln!("cargo-treemap: feature attribution: baseline build (--no-default-features)…");
    let base = builder::build(args, false, Some((true, Vec::new())))?;
    if base.artifacts.is_empty() {
        bail!("baseline build produced no artifact; the package may need a feature to build a binary");
    }
    let (bt, bd) = total_size(&base)?;

    #[derive(serde::Serialize)]
    struct FeatDelta {
        feature: String,
        text_delta: i64,
        data_delta: i64,
    }
    let mut rows: Vec<FeatDelta> = Vec::new();
    for f in &features {
        eprintln!("cargo-treemap: feature attribution: building with `{f}`…");
        match builder::build(args, false, Some((true, vec![f.clone()]))).and_then(|b| total_size(&b)) {
            Ok((t, d)) => rows.push(FeatDelta {
                feature: f.clone(),
                text_delta: t as i64 - bt as i64,
                data_delta: d as i64 - bd as i64,
            }),
            Err(e) => eprintln!("cargo-treemap: feature `{f}` unavailable, skipping: {e:#}"),
        }
    }
    rows.sort_by(|a, b| {
        (b.text_delta + b.data_delta)
            .abs()
            .cmp(&(a.text_delta + a.data_delta).abs())
    });

    if matches!(args.format, cli::Format::Json) {
        let out = serde_json::json!({
            "baseline": { "text": bt, "data": bd },
            "features": rows,
        });
        println!("{}", serde_json::to_string_pretty(&out)?);
    } else {
        println!("\n  Feature size attribution (marginal cost vs --no-default-features)\n");
        println!(
            "  baseline: .text {}, .data {}\n",
            util::human_bytes(bt),
            util::human_bytes(bd)
        );
        println!("  {:>11}  {:>11}  feature", "+.text", "+.data");
        println!("  {}", "-".repeat(40));
        for r in &rows {
            println!(
                "  {:>11}  {:>11}  {}",
                util::signed_bytes(r.text_delta),
                util::signed_bytes(r.data_delta),
                r.feature
            );
        }
        println!("\n  Marginal cost of enabling each feature alone; features can overlap, so these don't sum.\n");
    }
    Ok(())
}

fn total_size(build: &builder::BuildOutput) -> Result<(u64, u64)> {
    let (mut text, mut data) = (0u64, 0u64);
    for path in &build.artifacts {
        let parsed = artifact::parse(path, false)?;
        text += parsed.totals.text;
        data += parsed.totals.data;
    }
    Ok((text, data))
}

/// Build a git ref in a throwaway worktree and return the path to its saved JSON report.
fn build_ref_baseline(
    args: &cli::TreemapArgs,
    md: &cargo_metadata::Metadata,
    git_ref: &str,
) -> Result<std::path::PathBuf> {
    let ws = md.workspace_root.as_std_path();
    let repo_root = git_toplevel(ws)?;
    let rel = ws.strip_prefix(&repo_root).unwrap_or(std::path::Path::new(""));

    let pid = std::process::id();
    let tmp = std::env::temp_dir();
    let wt = tmp.join(format!("cargo-treemap-ref-{pid}"));
    let out_json = tmp.join(format!("cargo-treemap-ref-{pid}.json"));
    let _ = std::fs::remove_dir_all(&wt);

    eprintln!("cargo-treemap: --compare-ref: checking out `{git_ref}` in a worktree…");
    let wt_str = wt.to_str().context("worktree path is not valid UTF-8")?;
    run_git(&repo_root, &["worktree", "add", "--detach", "--force", wt_str, git_ref])
        .with_context(|| format!("creating a git worktree for `{git_ref}`"))?;

    let build = (|| -> Result<()> {
        let wt_manifest = wt.join(rel).join("Cargo.toml");
        let exe = std::env::current_exe().context("locating own binary")?;
        eprintln!("cargo-treemap: --compare-ref: building `{git_ref}` (a full build of the ref)…");
        let mut cmd = std::process::Command::new(exe);
        cmd.arg("treemap")
            .arg("--manifest-path")
            .arg(&wt_manifest)
            .arg("--format")
            .arg("json")
            .arg("-o")
            .arg(&out_json);
        forward_build_flags(&mut cmd, args);
        cmd.stdout(std::process::Stdio::inherit())
            .stderr(std::process::Stdio::inherit());
        if !cmd.status().context("running the baseline analysis")?.success() {
            bail!("analyzing ref `{git_ref}` failed");
        }
        Ok(())
    })();

    // Always tear the worktree down, even if the build failed.
    let _ = run_git(&repo_root, &["worktree", "remove", "--force", wt_str]);
    let _ = std::fs::remove_dir_all(&wt);
    build?;
    Ok(out_json)
}

fn git_toplevel(dir: &std::path::Path) -> Result<std::path::PathBuf> {
    let out = std::process::Command::new("git")
        .arg("-C")
        .arg(dir)
        .args(["rev-parse", "--show-toplevel"])
        .output()
        .context("running git (is it installed?)")?;
    if !out.status.success() {
        bail!("{} is not inside a git repository", dir.display());
    }
    Ok(std::path::PathBuf::from(
        String::from_utf8_lossy(&out.stdout).trim(),
    ))
}

fn run_git(repo_root: &std::path::Path, args: &[&str]) -> Result<()> {
    let status = std::process::Command::new("git")
        .arg("-C")
        .arg(repo_root)
        .args(args)
        .status()
        .context("running git")?;
    if !status.success() {
        bail!("`git {}` failed", args.join(" "));
    }
    Ok(())
}

/// Forward build-selection flags (plus `--split-std`) to the compare-ref self-invocation.
fn forward_build_flags(cmd: &mut std::process::Command, args: &cli::TreemapArgs) {
    if let Some(p) = &args.profile {
        cmd.arg("--profile").arg(p);
    } else if args.release {
        cmd.arg("--release");
    }
    if args.features.all_features {
        cmd.arg("--all-features");
    }
    if args.features.no_default_features {
        cmd.arg("--no-default-features");
    }
    if !args.features.features.is_empty() {
        cmd.arg("--features").arg(args.features.features.join(","));
    }
    builder::push_build_flags(cmd, args);
    if args.split_std {
        cmd.arg("--split-std"); // a cargo-treemap flag, not a cargo one
    }
}

/// Capture what produced this report: the (re-runnable) command line, the working directory,
/// and a wall-clock timestamp (unix seconds; the browser formats it for display).
fn provenance() -> (String, String, u64) {
    let args: Vec<String> = std::env::args().collect();
    let invocation = if args.get(1).map(|s| s == "treemap").unwrap_or(false) {
        format!("cargo treemap {}", shell_join(&args[2..]))
    } else {
        shell_join(&args)
    };
    let cwd = std::env::current_dir()
        .map(|p| p.display().to_string())
        .unwrap_or_default();
    let generated_unix = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0);
    (invocation, cwd, generated_unix)
}

/// Shell-quote args so the recorded command line round-trips (paths with spaces, etc.).
fn shell_join(args: &[String]) -> String {
    args.iter()
        .map(|a| {
            let safe = !a.is_empty()
                && a.bytes()
                    .all(|b| b.is_ascii_alphanumeric() || b"-_=/.:,+@%".contains(&b));
            if safe {
                a.clone()
            } else {
                format!("'{}'", a.replace('\'', "'\\''"))
            }
        })
        .collect::<Vec<_>>()
        .join(" ")
}

fn load_dwarf(enabled: bool, path: &std::path::Path) -> Option<dwarf::DwarfResolver> {
    if !enabled {
        return None;
    }
    match dwarf::DwarfResolver::load(path) {
        Ok(r) => Some(r),
        Err(e) => {
            eprintln!(
                "cargo-treemap: --dwarf: no usable debug info for {} ({e:#}); using name-based attribution",
                path.display()
            );
            None
        }
    }
}

fn derive_crate_names(parsed: &[artifact::ParsedArtifact]) -> std::collections::HashSet<String> {
    use util::{is_primitive_or_keyword, leading_ident, split_path, std_crates};
    // Exclude std/std-vendored crates so canon() still collapses them into `std` (matching
    // build mode); everything else that heads a path is treated as a crate.
    let std = std_crates();
    let mut set = std::collections::HashSet::new();
    for a in parsed {
        for sym in &a.symbols {
            let demangled = format!("{:#}", rustc_demangle::demangle(&sym.name));
            let segs = split_path(&demangled);
            if segs.len() > 1 {
                let id = leading_ident(&segs[0]);
                if id.len() > 1 && !is_primitive_or_keyword(id) && !std.contains(id) {
                    set.insert(id.to_string());
                }
            }
        }
    }
    set
}