cargo-show-asm 0.2.58

use anyhow::Context;
use cargo_metadata::{Artifact, Message, MetadataCommand, Package};

#[cfg(feature = "disasm")]
use cargo_show_asm::disasm::dump_disasm;
use cargo_show_asm::{
    asm::Asm,
    dump_function, esafeprintln,
    llvm::Llvm,
    mca::Mca,
    mir::Mir,
    opts::{self, CodeSource, OutputType},
    safeprintln,
};
use std::{
    io::BufReader,
    path::{Path, PathBuf},
    process::{Child, Stdio},
    sync::LazyLock,
};

fn cargo_path() -> &'static Path {
    static CARGO_PATH: LazyLock<PathBuf> =
        LazyLock::new(|| std::env::var_os("CARGO").map_or_else(|| "cargo".into(), PathBuf::from));
    &CARGO_PATH
}

fn rust_path() -> &'static Path {
    static RUSTC_PATH: LazyLock<PathBuf> =
        LazyLock::new(|| std::env::var_os("RUSTC").map_or_else(|| "rustc".into(), PathBuf::from));
    &RUSTC_PATH
}

#[cfg(not(feature = "disasm"))]
macro_rules! no_disasm {
    () => {{
        // Sigh, never type...
        esafeprintln!("This option requires cargo-show-asm to be compiled with \"disasm\" feature");
        std::process::exit(101)
    }};
}

fn spawn_cargo(
    cargo: &opts::Cargo,
    format: &opts::Format,
    syntax: opts::Syntax,
    target_cpu: Option<&str>,
    focus_package: &Package,
    focus_artifact: &opts::Focus,
    force_single_cgu: bool,
) -> std::io::Result<std::process::Child> {
    use std::ffi::OsStr;
    use std::fmt::Write;

    let mut cmd = std::process::Command::new(cargo_path());
    let mut rust_flags = std::env::var("RUSTFLAGS").unwrap_or_default();

    // Cargo flags.
    cmd.arg("rustc")
        // General.
        .args([
            "--message-format=json-render-diagnostics",
            "--color",
            if format.color { "always" } else { "never" },
        ])
        .args(std::iter::repeat_n(
            "-v",
            format.verbosity.saturating_sub(1),
        ))
        // Workspace location.
        .args(
            cargo
                .manifest_path
                .iter()
                .flat_map(|p| ["--manifest-path".as_ref(), p.as_os_str()]),
        )
        .args(["--config", "profile.release.strip=false"])
        // Artifact selectors.
        .args(["--package", &focus_package.name])
        .args(focus_artifact.as_cargo_args())
        // Compile options.
        .args(cargo.config.iter().flat_map(|c| ["--config", c]))
        .args(cargo.dry.then_some("--dry"))
        .args(cargo.frozen.then_some("--frozen"))
        .args(cargo.locked.then_some("--locked"))
        .args(cargo.offline.then_some("--offline"))
        .args(cargo.quiet.then_some("--quiet"))
        .args(cargo.target.iter().flat_map(|t| ["--target", t]))
        .args(cargo.unstable.iter().flat_map(|z| ["-Z", z]))
        .args((syntax.output_type == OutputType::Wasm).then_some("--target=wasm32-unknown-unknown"))
        .args(
            cargo
                .target_dir
                .iter()
                .flat_map(|t| [OsStr::new("--target-dir"), t.as_ref()]),
        )
        .args(
            cargo
                .cli_features
                .no_default_features
                .then_some("--no-default-features"),
        )
        .args(cargo.cli_features.all_features.then_some("--all-features"))
        .args(
            cargo
                .cli_features
                .features
                .iter()
                .flat_map(|feat| ["--features", feat]),
        );
    match &cargo.compile_mode {
        opts::CompileMode::Dev => {}
        opts::CompileMode::Release => {
            cmd.arg("--release");
        }
        opts::CompileMode::Custom(profile) => {
            cmd.args(["--profile", profile]);
        }
    }

    // Cargo flags terminator.
    cmd.arg("--");

    // Rustc flags.
    cmd
        // Start with the user-supplied codegen flags, which we might need to override.
        .args(cargo.codegen.iter().flat_map(|c| ["-C", c]))
        // Next, we care about asm/wasm/llvm-ir/llvm-mac.
        .args(syntax.emit().iter().flat_map(|s| ["--emit", s]))
        .args(syntax.format().iter().flat_map(|s| ["-C", s]));

    if let Some(cpu) = target_cpu {
        write!(rust_flags, " -Ctarget-cpu={cpu}").unwrap();
    }

    {
        // None corresponds to disasm
        if [Some("asm"), None].contains(&syntax.emit()) {
            // Debug info is needed to detect function boundaries in asm (Windows/Mac), and to map asm/wasm
            // output to rust source.
            cmd.arg("-Cdebuginfo=2");
        }
    }

    // current rust does not emit info about generated byproducts, new one will :)
    if force_single_cgu {
        cmd.arg("-Ccodegen-units=1");
    }

    if !rust_flags.is_empty() {
        // `args` from `cargo rustc -- args` are passed only to the final compiler instance.
        // `RUSTFLAGS` envvar is useful for passing flags to all compiler instances.
        cmd.env("RUSTFLAGS", rust_flags.trim_start());
    }

    if format.verbosity >= 3 {
        safeprintln!("Running: {cmd:?}");
    }

    cmd.stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::inherit())
        .spawn()
}

fn sysroot() -> anyhow::Result<PathBuf> {
    let output = std::process::Command::new(rust_path())
        .arg("--print=sysroot")
        .stdin(Stdio::null())
        .stderr(Stdio::inherit())
        .stdout(Stdio::piped())
        .output()?;
    if !output.status.success() {
        anyhow::bail!(
            "Failed to get sysroot. '{:?} --print=sysroot' exited with {}",
            rust_path(),
            output.status,
        );
    }
    // `rustc` prints a trailing newline.
    Ok(PathBuf::from(
        std::str::from_utf8(&output.stdout)?.trim_end(),
    ))
}

#[allow(clippy::too_many_lines)]
fn main() -> anyhow::Result<()> {
    let opts = opts::options().run();
    owo_colors::set_override(opts.format.color);

    let cargo = match opts.code_source {
        CodeSource::FromCargo { ref cargo } => cargo,
        CodeSource::File { ref file } => {
            if opts.format.verbosity > 1 {
                esafeprintln!("Processing a given single file");
            }
            match file.extension() {
                Some(ext) if ext == "s" => {
                    let nope = PathBuf::new();
                    let mut asm = Asm::new(&nope, &nope);
                    let mut format = opts.format;
                    // For standalone file we don't know the matching
                    // system root so don't even try to dump it
                    format.rust = false;
                    dump_function(&mut asm, opts.to_dump, file, &format)?;
                }
                _ => {
                    #[cfg(feature = "disasm")]
                    {
                        dump_disasm(opts.to_dump, file, &opts.format, opts.syntax.output_style)?;
                    }
                    #[cfg(not(feature = "disasm"))]
                    {
                        no_disasm!()
                    }
                }
            }
            return Ok(());
        }
    };

    let sysroot = sysroot()?;
    if opts.format.verbosity > 1 {
        esafeprintln!("Found sysroot: {}", sysroot.display());
    }

    let unstable = cargo
        .unstable
        .iter()
        .flat_map(|x| ["-Z".to_owned(), x.clone()])
        .collect::<Vec<_>>();

    let mut metadata = MetadataCommand::new();
    metadata.cargo_path(cargo_path());
    if let Some(path) = &cargo.manifest_path {
        metadata.manifest_path(path);
    }
    let metadata = metadata.other_options(unstable).no_deps().exec()?;

    #[cfg(feature = "_unstable")]
    let target_dir_to_build_dir = metadata
        .build_directory
        .as_ref()
        .map(|build| (metadata.target_directory.as_std_path(), build.as_std_path()));
    #[cfg(not(feature = "_unstable"))]
    let target_dir_to_build_dir = None;

    let focus_package = match opts
        .select_fragment
        .package
        .as_ref()
        .map(|n| n.as_str().trim())
    {
        Some(name) => metadata
            .packages
            .iter()
            .find(|p| *p.name == name || p.id.repr.as_str() == name)
            .with_context(|| format!("Package '{name}' is not found"))?,
        None if metadata.packages.len() == 1 => &metadata.packages[0],
        None => {
            esafeprintln!(
                "{:?} refers to multiple packages, you need to specify which one to use",
                cargo.manifest_path
            );
            for package in &metadata.packages {
                esafeprintln!("\t-p {}", package.name);
            }
            anyhow::bail!("Multiple packages found")
        }
    };

    let focus_artifact = match opts.select_fragment.focus {
        Some(ref focus) => focus.clone(),
        None => {
            let candidates = focus_package
                .targets
                .iter()
                .filter_map(|t| opts::Focus::try_from(t).ok())
                .collect::<Vec<_>>();
            match candidates.as_slice() {
                [] => anyhow::bail!("No targets found"),
                [c] => c.clone(),
                xs => {
                    esafeprintln!(
                        "{} defines multiple targets, you need to specify which one to use:",
                        focus_package.name
                    );
                    for focus in xs {
                        esafeprintln!("\t{}", focus.as_cargo_args().collect::<Vec<_>>().join(" "));
                    }
                    anyhow::bail!("Multiple targets found")
                }
            }
        }
    };

    // Pending on this https://github.com/rust-lang/rust/pull/122597

    #[cfg(feature = "disasm")]
    let force_single_cgu = opts.syntax.output_type != OutputType::Disasm;

    #[cfg(not(feature = "disasm"))]
    let force_single_cgu = true;

    let cargo_child = spawn_cargo(
        cargo,
        &opts.format,
        opts.syntax,
        opts.target_cpu.as_deref(),
        focus_package,
        &focus_artifact,
        force_single_cgu,
    )?;

    let asm_path = cargo_to_asm_path(cargo_child, &focus_artifact, &opts, target_dir_to_build_dir)?;

    if opts.format.verbosity > 3 {
        safeprintln!("goal: {:?}", opts.to_dump);
    }

    match opts.syntax.output_type {
        OutputType::Asm | OutputType::Wasm => {
            let mut asm = Asm::new(metadata.workspace_root.as_std_path(), &sysroot);
            dump_function(&mut asm, opts.to_dump, &asm_path, &opts.format)
        }
        OutputType::Llvm | OutputType::LlvmInput => {
            dump_function(&mut Llvm, opts.to_dump, &asm_path, &opts.format)
        }
        OutputType::Mir => dump_function(&mut Mir, opts.to_dump, &asm_path, &opts.format),
        OutputType::Mca => {
            let mut mca = Mca::new(
                &opts.mca_arg,
                cargo.target.as_deref(),
                opts.target_cpu.as_deref(),
            );
            dump_function(&mut mca, opts.to_dump, &asm_path, &opts.format)
        }
        #[cfg(not(feature = "disasm"))]
        OutputType::Disasm => no_disasm!(),

        #[cfg(feature = "disasm")]
        OutputType::Disasm => dump_disasm(
            opts.to_dump,
            &asm_path,
            &opts.format,
            opts.syntax.output_style,
        ),
    }
}

fn cargo_to_asm_path(
    mut cargo: Child,
    focus_artifact: &opts::Focus,
    opts: &crate::opts::Options,
    target_dir_to_build_dir: Option<(&Path, &Path)>,
) -> anyhow::Result<PathBuf> {
    let mut result_artifact = None;
    let mut success = false;
    for msg in Message::parse_stream(BufReader::new(cargo.stdout.take().unwrap())) {
        match msg? {
            Message::CompilerArtifact(artifact) if focus_artifact.matches_artifact(&artifact) => {
                result_artifact = Some(artifact);
            }
            Message::BuildFinished(fin) => {
                success = fin.success;
                break;
            }
            _ => {}
        }
    }
    // add some spacing between cargo's output and ours
    esafeprintln!();
    if !success {
        let status = cargo.wait().context("cargo process failed")?;
        esafeprintln!("Cargo failed with {status}");
        std::process::exit(101);
    }
    let artifact = result_artifact.context("No artifact found")?;

    if opts.format.verbosity > 1 {
        esafeprintln!("Artifact files: {:?}", artifact.filenames);
    }

    let asm_path = match opts.syntax.ext() {
        Some(expect_ext) => {
            locate_asm_path_via_artifact(&artifact, expect_ext, target_dir_to_build_dir)?
        }
        None => {
            if let Some(executable) = artifact.executable {
                executable.into()
            } else if let Some(rlib) = artifact
                .filenames
                .iter()
                .find(|f| f.extension() == Some("rlib"))
            {
                rlib.into()
            } else {
                anyhow::bail!(
                    "Looking for an executable or an rlib file to work on, got {:?} instead.",
                    artifact
                );
            }
        }
    };
    if opts.format.verbosity > 1 {
        esafeprintln!("Working with file: {}", asm_path.display());
    }
    Ok(asm_path)
}

fn locate_asm_path_via_artifact(
    artifact: &Artifact,
    expect_ext: &str,
    target_dir_to_build_dir: Option<(&Path, &Path)>,
) -> anyhow::Result<PathBuf> {
    // When build-dir is enabled, deps are in the build dir rather than the target dir.
    let build_dir = |path: &Path| {
        target_dir_to_build_dir
            .and_then(|(target, build)| Some(build.join(path.strip_prefix(target).ok()?)))
            .unwrap_or_else(|| path.into())
    };
    let read_dir = |path: &Path| {
        build_dir(path)
            .read_dir()
            .ok()
            .into_iter()
            .flat_map(|iter| iter.filter_map(|entry| entry.ok()))
    };

    // In Cargo's legacy build-dir layout
    // For lib, test, bench, lib-type example, `filenames` hint the file stem of the asm file.
    // We could locate asm files precisely.
    //
    // `filenames`:
    // [..]/target/debug/deps/libfoo-01234567.rmeta         # lib by-product
    // [..]/target/debug/deps/foo-01234567                  # test & bench
    // [..]/target/debug/deps/example/libfoo-01234567.rmeta # lib-type example by-product
    // Asm files:
    // [..]/target/debug/deps/foo-01234567.s
    // [..]/target/debug/deps/example/foo-01234567.s
    if let Some(path) = artifact
        .filenames
        .iter()
        .filter(|path| {
            matches!(
                path.parent().unwrap().file_name(),
                Some("deps" | "examples")
            )
        })
        .find_map(|path| {
            let path = build_dir(path.as_ref()).with_extension(expect_ext);
            if path.exists() {
                return Some(path);
            }
            let path = path.with_file_name(path.file_name()?.to_str()?.strip_prefix("lib")?);
            if path.exists() {
                return Some(path);
            }
            None
        })
    {
        return Ok(path);
    }

    // In Cargo's new build-dir layout, the compiler rmeta output is in the same directory as the
    // asm file so we can just check the current directory.
    //
    // `filenames`:
    // [..]/build-dir/release/build/foo/01234567/out/libfoo-01234567.rmeta" # lib by-product
    // Asm files:
    // [..]/build-dir/release/build/foo/01234567/out/foo-01234567.s
    if let Some(path) = artifact
        .filenames
        .iter()
        .filter(|path| matches!(path.extension(), Some("rmeta")))
        .find_map(|path| {
            for entry in read_dir(build_dir(path.parent().unwrap().as_ref()).as_ref()) {
                let asm_file = entry.path().with_extension(expect_ext);
                if asm_file.exists() {
                    return Some(asm_file);
                }
            }
            None
        })
    {
        return Ok(path);
    }

    // then there's rlib with filenames as following:
    // `filenames`:
    // [..]/target/debug/libfoo.a              <+
    // [..]/target/debug/libfoo.rlib            | <+ Hard linked.
    // Asm files:                               |  | Or same contents at least
    // [..]/target/debug/libfoo-01234567.a     <+  |
    // [..]/target/debug/libfoo-01234567.rlib     <+
    // [..]/target/debug/foo-01234567.s

    if let Some(rlib_path) = artifact
        .filenames
        .iter()
        .find(|f| f.extension() == Some("rlib"))
    {
        let deps_dir = rlib_path.with_file_name("deps");

        for entry in read_dir(deps_dir.as_ref()) {
            let maybe_origin = entry.path();
            if same_contents(rlib_path.as_ref(), &maybe_origin) {
                let name = maybe_origin
                    .file_name()
                    .unwrap()
                    .to_str()
                    .unwrap()
                    .strip_prefix("lib")
                    .unwrap();
                let asm_file = maybe_origin.with_file_name(name).with_extension(expect_ext);
                if asm_file.exists() {
                    return Ok(asm_file);
                }
            }
        }
    }

    // for cdylib we have
    // [..]/target/debug/deps/xx.d
    // [..]/target/debug/deps/libxx.so <+ Hard linked/same contents
    // [..]/target/debug/deps/xx.s      | <- asm file
    // [..]/target/debug/libxx.d        |
    // [..]/target/debug/libxx.so      <+ <- artifact
    //
    // on windows it's xx.dll / xx.s, on MacOS it's libxx.dylib / xx.s...
    //    if artifact.target.kind.iter().any(|k| k == "cdylib") {
    //
    // In Cargo's new build-dir layout the asm file is in the same directory as the dylib, so we
    // search for the dylib in the build-dir, then search for the asm in that directory.
    if let Some(cdylib_path) = artifact.filenames.iter().find(|f| {
        f.extension()
            .is_some_and(|e| ["so", "dylib", "dll"].contains(&e))
    }) {
        let deps_dir = cdylib_path.with_file_name("deps");
        for entry in read_dir(deps_dir.as_ref()) {
            let maybe_origin = entry.path();
            if same_contents(cdylib_path.as_ref(), &maybe_origin) {
                let Some(name) = maybe_origin.file_name() else {
                    continue;
                };
                let Some(name) = name.to_str() else { continue };
                let name = name.strip_prefix("lib").unwrap_or(name);
                // on windows this is xx.dll -> xx.s, no lib....
                let asm_file = maybe_origin.with_file_name(name).with_extension(expect_ext);
                if asm_file.exists() {
                    return Ok(asm_file);
                }
            }
        }

        // New Cargo build-dir layout
        let build = cdylib_path.with_file_name("build");

        for maybe_origin in read_dir_recursive(build.as_ref(), &build_dir) {
            if same_contents(cdylib_path.as_ref(), &maybe_origin) {
                // We found the dylib in the build-dir, so the asm file should be in the same
                // directory with it. So search for the asm file.
                for entry in read_dir(maybe_origin.parent().unwrap()) {
                    let asm_file = entry.path().with_extension(expect_ext);
                    if asm_file.exists() {
                        return Ok(asm_file);
                    }
                }
            }
        }
    }

    // For bin or bin-type example artifacts, `filenames` provide hard-linked paths
    // without extra-filename.
    // We scan all possible original artifacts by checking hard links,
    // in order to retrieve the correct extra-filename, and then locate asm files.
    //
    // In Cargo's legacy build-dir layout
    // `filenames`, also `executable`:
    // [..]/target/debug/foobin                    <+
    // [..]/target/debug/examples/fooexample        | <+ Hard linked.
    // Origins:                                     |  |
    // [..]/target/debug/deps/foobin-01234567      <+  |
    // [..]/target/debug/examples/fooexample-01234567 <+
    // Asm files:
    // [..]/target/debug/deps/foobin-01234567.s
    // [..]/target/debug/examples/fooexample-01234567.s
    //
    // In Cargo's new build-dir layout
    // `filenames`, also `executable`:
    // [..]/target/debug/foobin                                   <+
    // [..]/target/debug/examples/fooexample                       | <+ Hard linked.
    // Origins:                                                    |  |
    // [..]/build-dir/debug/build/foobin/01234567/out/foobin      <+  |
    // [..]/build-dir/debug/build/fooexample/01234567/out/fooexample <+
    // Asm files:
    // [..]/build-dir/debug/deps/foobin-01234567.s
    // [..]/build-dir/debug/examples/fooexample-01234567.s
    if let Some(exe_path) = &artifact.executable {
        let parent = exe_path.parent().unwrap();
        let deps_dir = if parent.file_name() == Some("examples") {
            parent.to_owned()
        } else {
            exe_path.with_file_name("deps")
        };

        for entry in read_dir(deps_dir.as_ref()) {
            let maybe_origin = entry.path();
            if same_contents(exe_path.as_ref(), &maybe_origin) {
                let asm_file = maybe_origin.with_extension(expect_ext);
                if asm_file.exists() {
                    return Ok(asm_file);
                }
            }
        }

        let build = if parent.file_name() == Some("examples") {
            exe_path.parent().unwrap().with_file_name("build")
        } else {
            exe_path.with_file_name("build")
        };

        for maybe_origin in read_dir_recursive(build.as_ref(), &build_dir) {
            if same_contents(exe_path.as_ref(), &maybe_origin) {
                let asm_file = maybe_origin.with_extension(expect_ext);
                if asm_file.exists() {
                    return Ok(asm_file);
                }
            }
        }
    }

    anyhow::bail!(
        "Cannot locate the path to the asm file\nArtifact paths: {}",
        artifact
            .filenames
            .iter()
            .chain(artifact.executable.as_ref())
            .map(|f| f.as_str())
            .collect::<Vec<_>>()
            .join(", ")
    );
}

fn read_dir_recursive(path: &Path, build_dir: &impl Fn(&Path) -> PathBuf) -> Vec<PathBuf> {
    build_dir(path)
        .read_dir()
        .ok()
        .into_iter()
        .flat_map(|iter| iter.filter_map(|entry| entry.ok()))
        .flat_map(|entry| {
            if entry.path().is_dir() {
                read_dir_recursive(&entry.path(), build_dir).into_iter()
            } else {
                vec![entry.path()].into_iter()
            }
        })
        .collect()
}

/// Check if files have the same contents
///
/// Before that we check if they are the same according to hardlinks (sorry filesystems with no
/// hardlinks, you are not real), and before reading files - we check if sizes are the same - files
/// are likely identical.
fn same_contents(a: &Path, b: &Path) -> bool {
    same_file::is_same_file(a, b).unwrap_or(false)
        || (std::fs::metadata(a)
            .ok()
            .zip(std::fs::metadata(b).ok())
            .is_some_and(|(a, b)| a.len() == b.len())
            && std::fs::read(a)
                .ok()
                .zip(std::fs::read(b).ok())
                .is_some_and(|(a, b)| a == b))
}