rust_mir2_core 0.1.5

Shared Rust MIR extraction model and helpers for rust_mir2
Documentation
use std::collections::BTreeSet;
use std::ffi::OsStr;
use std::path::{Path, PathBuf};
use std::process::Command;

use anyhow::{Context, bail};

use crate::fragments::load_fragments_from_dir;
use crate::manifest::try_resolve_target_project;
use crate::model::{
    ALLOWED_TARGETS_ENV, EXPECTED_PACKAGES_ENV, InputKind, InvocationMode, OUTPUT_DIR_ENV,
    OUTPUT_KIND_ENV, RustMir2Error, TargetProject, WRAPPER_MODE_ENV, WorkspaceMirDump,
};
use crate::render::build_packages_from_targets;
use std::fs;

pub fn detect_invocation_mode(args: &[String], forced_mode: Option<&str>) -> InvocationMode {
    if matches!(forced_mode, Some("workspace-wrapper")) {
        return InvocationMode::WorkspaceWrapper;
    }

    if args.get(1).is_some_and(|arg| {
        let path = Path::new(arg);
        path.file_stem()
            .and_then(OsStr::to_str)
            .is_some_and(|stem| stem == "rustc")
    }) {
        return InvocationMode::WorkspaceWrapper;
    }

    InvocationMode::Library
}

fn wrapper_binary_path() -> anyhow::Result<PathBuf> {
    if let Some(exe) = std::env::var_os("RUST_MIR2_WRAPPER_BIN").map(PathBuf::from) {
        if exe.is_file() {
            return Ok(exe);
        }

        bail!(
            "RUST_MIR2_WRAPPER_BIN points to `{}`, but it is not a file",
            exe.display()
        );
    }

    if let Some(path_env) = std::env::var_os("PATH") {
        for dir in std::env::split_paths(&path_env) {
            let candidate = dir.join("rust_mir2-wrapper");
            if candidate.is_file() {
                return Ok(candidate);
            }
        }
    }

    bail!(
        "failed to find published rust_mir2-wrapper; set RUST_MIR2_WRAPPER_BIN or install rust_mir2_wrapper so `rust_mir2-wrapper` is on PATH"
    );
}

fn run_cargo_check(target_project: &TargetProject, output_dir: &Path) -> anyhow::Result<()> {
    run_cargo_check_with_output_kind(target_project, output_dir, "json")
}

fn run_cargo_check_stream_2r(
    target_project: &TargetProject,
    output_dir: &Path,
) -> anyhow::Result<()> {
    run_cargo_check_with_output_kind(target_project, output_dir, "2r")
}

fn run_cargo_check_with_output_kind(
    target_project: &TargetProject,
    output_dir: &Path,
    output_kind: &str,
) -> anyhow::Result<()> {
    let wrapper = wrapper_binary_path()?;
    let cargo = std::env::var("CARGO").unwrap_or_else(|_| "cargo".to_string());
    let target_dir = std::env::var_os("CARGO_TARGET_DIR")
        .map(PathBuf::from)
        .unwrap_or_else(|| output_dir.join("cargo-target"));
    let mut cmd = Command::new(cargo);
    cmd.current_dir(&target_project.project_root)
        .arg("check")
        .arg("--target-dir")
        .arg(&target_dir);

    match target_project.input_kind {
        InputKind::Workspace => {
            cmd.arg("--workspace");
        }
        InputKind::SingleCrate => {
            cmd.arg("--package").arg(
                target_project
                    .expected_package_names
                    .first()
                    .context("missing package name")?,
            );
        }
    }

    let has_lib = target_project
        .default_targets
        .iter()
        .any(|target| target.target_kind == "lib");
    let bins = target_project
        .default_targets
        .iter()
        .filter(|target| target.target_kind == "bin")
        .map(|target| target.target_name.as_str())
        .collect::<Vec<_>>();

    if output_kind == "2r" {
        if has_lib {
            cmd.arg("--lib");
        }
        for bin in bins {
            cmd.arg("--bin").arg(bin);
        }
    } else {
        cmd.arg("--all-targets");
    }

    cmd.env("RUSTC_WORKSPACE_WRAPPER", &wrapper)
        .env(WRAPPER_MODE_ENV, "workspace-wrapper")
        .env(OUTPUT_DIR_ENV, output_dir)
        .env(
            EXPECTED_PACKAGES_ENV,
            target_project.expected_package_names.join(","),
        )
        .env(
            ALLOWED_TARGETS_ENV,
            target_project
                .default_targets
                .iter()
                .map(|target| target.target_identity.as_str())
                .collect::<Vec<_>>()
                .join(","),
        )
        .env(OUTPUT_KIND_ENV, output_kind);

    let status = cmd.status().context("failed to run cargo check")?;
    if !status.success() {
        bail!("cargo check failed with status {status}");
    }

    Ok(())
}

pub fn analyze_workspace(path: &Path) -> Result<WorkspaceMirDump, RustMir2Error> {
    try_analyze_workspace(path).map_err(Into::into)
}

fn try_analyze_workspace(path: &Path) -> anyhow::Result<WorkspaceMirDump> {
    let target_project = try_resolve_target_project(path)?;
    let temp_dir = tempfile::tempdir().context("failed to create temporary output directory")?;

    run_cargo_check(&target_project, temp_dir.path())?;
    let fragments =
        load_fragments_from_dir(temp_dir.path()).map_err(|error| anyhow::anyhow!(error.message))?;

    if fragments.is_empty() {
        bail!("rust_mir2 produced no target fragments");
    }

    let expected = target_project
        .expected_package_names
        .iter()
        .cloned()
        .collect::<BTreeSet<_>>();
    let actual = fragments
        .iter()
        .map(|fragment| fragment.package_name.clone())
        .collect::<BTreeSet<_>>();

    for expected_name in &expected {
        if !actual.contains(expected_name) {
            bail!("missing fragment for local package `{expected_name}`");
        }
    }

    for actual_name in &actual {
        if !expected.contains(actual_name) {
            bail!("unexpected fragment for non-target package `{actual_name}`");
        }
    }

    Ok(WorkspaceMirDump {
        project_path: target_project.project_root,
        packages: build_packages_from_targets(fragments),
    })
}

pub fn analyze_path(path: &Path) -> Result<WorkspaceMirDump, RustMir2Error> {
    analyze_workspace(path)
}

pub fn export_check_stream_2r(
    path: &Path,
    output_dir: &Path,
) -> Result<Vec<PathBuf>, RustMir2Error> {
    try_export_check_stream_2r(path, output_dir).map_err(Into::into)
}

fn try_export_check_stream_2r(path: &Path, output_dir: &Path) -> anyhow::Result<Vec<PathBuf>> {
    let target_project = try_resolve_target_project(path)?;
    fs::create_dir_all(output_dir)
        .with_context(|| format!("failed to create output dir `{}`", output_dir.display()))?;
    run_cargo_check_stream_2r(&target_project, output_dir)?;

    let mut outputs = target_project
        .default_targets
        .iter()
        .map(|target| {
            target_2r_output_path(output_dir, &target.package_name, &target.target_identity)
        })
        .filter(|path| path.is_file())
        .collect::<Vec<_>>();
    outputs.sort();

    if outputs.is_empty() {
        bail!("rust_mir2 check stream produced no .2r target files");
    }

    if let Some(package_name) = target_project.expected_package_names.first() {
        if let Some(default_bin) = target_project
            .default_targets
            .iter()
            .find(|target| target.target_kind == "bin")
        {
            let source = target_2r_output_path(
                output_dir,
                &default_bin.package_name,
                &default_bin.target_identity,
            );
            if source.is_file() {
                let package_output = output_dir.join(format!("{package_name}.2r"));
                fs::copy(&source, &package_output).with_context(|| {
                    format!(
                        "failed to copy default target `{}` to `{}`",
                        source.display(),
                        package_output.display()
                    )
                })?;
                outputs.push(package_output);
            }
        }
    }

    Ok(outputs)
}

pub fn target_2r_output_path(
    output_dir: &Path,
    package_name: &str,
    target_identity: &str,
) -> PathBuf {
    output_dir.join(format!(
        "{}__{}.2r",
        sanitize_output_part(package_name),
        sanitize_output_part(target_identity)
    ))
}

fn sanitize_output_part(part: &str) -> String {
    let mut out = String::with_capacity(part.len());
    let mut previous_was_underscore = false;
    for ch in part.chars() {
        let keep = ch.is_ascii_alphanumeric() || ch == '-' || ch == '_';
        if keep {
            out.push(ch);
            previous_was_underscore = false;
        } else if !previous_was_underscore {
            out.push('_');
            previous_was_underscore = true;
        }
    }
    let trimmed = out.trim_matches('_');
    if trimmed.is_empty() {
        "unknown".to_string()
    } else {
        trimmed.to_string()
    }
}

pub fn export_emit_mir(path: &Path, output_dir: &Path) -> Result<Vec<PathBuf>, RustMir2Error> {
    try_export_emit_mir(path, output_dir).map_err(Into::into)
}

fn try_export_emit_mir(path: &Path, output_dir: &Path) -> anyhow::Result<Vec<PathBuf>> {
    let target_project = try_resolve_target_project(path)?;
    let package_name = target_project
        .expected_package_names
        .first()
        .cloned()
        .context("missing package name")?;
    let cargo = std::env::var("CARGO").unwrap_or_else(|_| "cargo".to_string());
    let target_dir = output_dir.join("cargo-target");

    fs::create_dir_all(output_dir)
        .with_context(|| format!("failed to create output dir `{}`", output_dir.display()))?;
    let mut cmd = Command::new(cargo);
    cmd.current_dir(&target_project.project_root)
        .arg("rustc")
        .arg("--package")
        .arg(&package_name)
        .arg("--bin")
        .arg(&package_name)
        .arg("--target-dir")
        .arg(&target_dir)
        .arg("--")
        .arg("--emit=mir");

    let status = cmd
        .status()
        .context("failed to run cargo rustc --emit=mir")?;
    if !status.success() {
        bail!("cargo rustc --emit=mir failed with status {status}");
    }

    let deps_dir = target_dir.join("debug").join("deps");
    let mut copied = Vec::new();
    for entry in fs::read_dir(&deps_dir)
        .with_context(|| format!("failed to read `{}`", deps_dir.display()))?
    {
        let entry = entry?;
        let path = entry.path();
        if path.extension().and_then(|value| value.to_str()) != Some("mir") {
            continue;
        }
        let target = output_dir.join(path.file_name().context("mir file had no file name")?);
        fs::copy(&path, &target).with_context(|| {
            format!(
                "failed to copy `{}` to `{}`",
                path.display(),
                target.display()
            )
        })?;
        copied.push(target);
    }

    if copied.is_empty() {
        bail!("cargo rustc --emit=mir produced no .mir files");
    }

    copied.sort();
    Ok(copied)
}

#[cfg(test)]
mod tests {
    use super::wrapper_binary_path;
    use std::sync::Mutex;

    static ENV_LOCK: Mutex<()> = Mutex::new(());

    #[test]
    fn wrapper_binary_path_falls_back_to_path() {
        let _guard = ENV_LOCK.lock().unwrap();
        unsafe {
            std::env::remove_var("RUST_MIR2_WRAPPER_BIN");
        }

        let path = wrapper_binary_path().expect("PATH wrapper lookup should succeed");
        assert_eq!(
            path.file_name().and_then(|name| name.to_str()),
            Some("rust_mir2-wrapper")
        );
    }

    #[test]
    fn wrapper_binary_path_rejects_missing_file() {
        let _guard = ENV_LOCK.lock().unwrap();
        unsafe {
            std::env::set_var(
                "RUST_MIR2_WRAPPER_BIN",
                "/definitely/not/a/rust_mir2-wrapper",
            );
        }

        let error = wrapper_binary_path().expect_err("missing wrapper file should fail");
        assert!(error.to_string().contains("but it is not a file"));
    }
}