mldsa-native-rs 0.0.1-alpha.6

//! build_native_detect.rs
//!
//! Architecture and feature detection for compiling the mldsa-native /
//! mlkem-native C sources from a Cargo build script. It reproduces the
//! define/flag *logic* of `mldsa-native/test/mk/auto.mk`, but the detection
//! itself is Rust/Cargo-native and target-focused, so it works correctly with
//! Cargo's cross-compilation.
//!
//! Design:
//! - **Architecture** comes from `CARGO_CFG_TARGET_ARCH` / `CARGO_CFG_TARGET_ENDIAN`
//!   (the *target*, set by Cargo for build scripts), normalized to auto.mk's
//!   ARCH names. This is what makes cross-compilation Just Work.
//! - **Feature availability** comes in two flavors:
//!   * `detect_target_features()` is suitable for native builds, where
//!     the building host is also the target host.
//!     Target features are derived from `CARGO_CFG_TARGET_FEATURE` — the set of
//!     target features rustc is actually compiling with.
//!     Matching the C side to this is what avoids C-vs-Rust ABI mismatches on FFI.
//!     This *replaces* auto.mk's `/proc/cpuinfo`/`sysctl` host probing (which inspects the build
//!     machine and breaks under cross-compilation).
//!   * `force_target_features()` reports the features expected for the target
//!     platform to emit the optimized native code; in this case we can
//!     rely on runtime CPU feature detection to ensure the optimized native
//!     code is executed only when supported in the execution environment.
//! - **Compiler confirmation** in two stages: first we ask
//!   the cc crate whether the compiler *accepts* the flag
//!   (`cc::Build::is_flag_supported`); then we compile a tiny inline-asm
//!   snippet under that flag (compile-only, via a throwaway `cc::Build` driven
//!   by `try_compile_intermediates`) to confirm the instructions actually
//!   *assemble*. Either failing skips the flag with a `cargo:warning`. This is
//!   auto.mk's `MK_COMPILER_SUPPORTS_*` probe, demoted from "detection" to
//!   "check".
//!
//! Define/flag mapping (identical to auto.mk):
//!   x86_64      -> -DMLD_FORCE_X86_64,   +`-mavx2` if avx2,  +`-mbmi2` if bmi2
//!                  (sse2 is inspected but, as in auto.mk, never adds a flag)
//!   aarch64     -> -DMLD_FORCE_AARCH64,  +`-march=armv8.4-a+sha3` if sha3
//!   aarch64_be  -> -DMLD_FORCE_AARCH64_EB
//!   riscv64     -> -DMLD_FORCE_RISCV64,  +`-march=rv64gcv` if the `v` extension
//!   riscv32     -> -DMLD_FORCE_RISCV32
//!   powerpc64le -> -DMLD_FORCE_PPC64LE
//!

use std::env;

use force_target_features as target_features;
//use detect_target_features as target_features;

/// Architecture names normalized to match auto.mk's `ARCH` values.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Arch {
    X86_64,
    Aarch64,
    Aarch64Be,
    Riscv64,
    Riscv32,
    Ppc64le,
    Arm,
    Unknown,
}

impl Arch {
    /// The `MLD_FORCE_*` define associated with this architecture, if any.
    pub fn force_define(self) -> Option<&'static str> {
        match self {
            Arch::X86_64 => Some("MLD_FORCE_X86_64"),
            Arch::Aarch64 => Some("MLD_FORCE_AARCH64"),
            Arch::Aarch64Be => Some("MLD_FORCE_AARCH64_EB"),
            Arch::Riscv64 => Some("MLD_FORCE_RISCV64"),
            Arch::Riscv32 => Some("MLD_FORCE_RISCV32"),
            Arch::Ppc64le => Some("MLD_FORCE_PPC64LE"),
            Arch::Arm | Arch::Unknown => None,
        }
    }
}

/// A preprocessor define: `(name, value)`. `value == None` means a bare
/// `-Dname` (matches how auto.mk writes `-DMLD_FORCE_*`).
pub type Define = (String, Option<String>);

/// The outcome of detection: everything needed to feed a compiler.
#[derive(Debug, Clone, Default)]
pub struct Detection {
    #[expect(dead_code)] // only used by Debug
    /// Resolved target architecture.
    pub arch: Option<Arch>,
    /// `-D` defines to pass to the compiler.
    pub defines: Vec<Define>,
    /// Compiler flags to pass, e.g. `-mavx2`, `-march=armv8.4-a+sha3`.
    pub flags: Vec<String>,
}

impl Detection {
    /// Apply all detected defines and flags to a `cc::Build`.
    pub fn apply(&self, build: &mut cc::Build) {
        for (name, value) in &self.defines {
            build.define(name, value.as_deref());
        }
        for flag in &self.flags {
            build.flag(flag);
        }
    }
}

/// Detect architecture + features for the *target* and return defines/flags.
pub fn detect() -> Detection {
    let arch = detect_arch();
    println!("cargo:warning=(INFO) Architecture: `{arch:?}`");

    let features = target_features(arch);
    println!("cargo:warning=(INFO) {features:?}");

    let mut defines: Vec<Define> = Vec::new();
    let mut flags: Vec<String> = Vec::new();

    if let Some(def) = arch.force_define() {
        defines.push((def.to_string(), None));
    }

    match arch {
        Arch::X86_64 => {
            // AVX2: target-enabled AND compiler accepts the flag AND can emit it.
            if features.has("avx2")
                && compiler_can_emit(
                    "MK_COMPILER_SUPPORTS_AVX2",
                    "-mavx2",
                    Some(
                        r#"int main() { __asm__("vpxor %%ymm0, %%ymm1, %%ymm2" ::: "ymm0", "ymm1", "ymm2"); return 0; }"#,
                    ),
                )
            {
                flags.push("-mavx2".into());
            }
            // BMI2.
            if features.has("bmi2")
                && compiler_can_emit(
                    "MK_COMPILER_SUPPORTS_BMI2",
                    "-mbmi2",
                    Some(
                        r#"int main() { __asm__("pdep %%eax, %%ebx, %%ecx" ::: "eax", "ebx", "ecx"); return 0; }"#,
                    ),
                )
            {
                flags.push("-mbmi2".into());
            }
            // SSE2 is inspected in auto.mk but never turns into a flag. We honor
            // the override var for parity but add nothing.
            let _ = env_override("MK_COMPILER_SUPPORTS_SSE2");
        }
        Arch::Aarch64 => {
            if features.has("sha3")
                && compiler_can_emit(
                    "MK_COMPILER_SUPPORTS_SHA3",
                    "-march=armv8.4-a+sha3",
                    Some(
                        r#"int main() { __asm__("eor3 v0.16b, v1.16b, v2.16b, v3.16b" ::: "v0", "v1", "v2", "v3"); return 0; }"#,
                    ),
                )
            {
                flags.push("-march=armv8.4-a+sha3".into());
            }
        }
        Arch::Riscv64 => {
            // The RISC-V vector extension is the `v` target feature in rustc.
            if features.has("v")
                && compiler_can_emit(
                    "MK_COMPILER_SUPPORTS_RVV",
                    "-march=rv64gcv",
                    Some(r#"int main() { __asm__("vadd.vv v0, v1, v2"); return 0; }"#),
                )
            {
                flags.push("-march=rv64gcv".into());
            }
        }
        // Force define only; no optional feature flags in auto.mk.
        Arch::Aarch64Be | Arch::Riscv32 | Arch::Ppc64le | Arch::Arm | Arch::Unknown => {}
    }

    Detection {
        arch: Some(arch),
        defines,
        flags,
    }
}

/// Map Cargo's *target* to auto.mk's `ARCH`. Cargo always provides the target
/// arch to build scripts, so this transparently covers cross-compilation
/// (no CROSS_PREFIX parsing needed).
pub fn detect_arch() -> Arch {
    let target_arch = env::var("CARGO_CFG_TARGET_ARCH")
        .expect("CARGO_CFG_TARGET_ARCH not set (run as a Cargo build script)");
    let target_endian = env::var("CARGO_CFG_TARGET_ENDIAN").unwrap_or_default();
    match target_arch.as_str() {
        "x86_64" => Arch::X86_64,
        "aarch64" => {
            if target_endian == "big" {
                Arch::Aarch64Be
            } else {
                Arch::Aarch64
            }
        }
        "riscv64" => Arch::Riscv64,
        "riscv32" => Arch::Riscv32,
        "powerpc64" if target_endian == "little" => Arch::Ppc64le,
        "arm" => Arch::Arm,
        _ => Arch::Unknown,
    }
}

/// The set of target features rustc is compiling the target with, taken from
/// `CARGO_CFG_TARGET_FEATURE` (comma-separated). This is the Rust-native
/// replacement for auto.mk's host-CPU probing and reflects what the Rust side
/// is built with, keeping the C side ABI-compatible.
#[derive(Debug)]
struct TargetFeatures {
    set: Vec<String>,
}

impl TargetFeatures {
    const EMPTY: Self = Self { set: Vec::new() };

    fn new<I, S>(features: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        Self {
            set: features.into_iter().map(Into::into).collect(),
        }
    }

    fn has(&self, feat: &str) -> bool {
        self.set.iter().any(|f| f == feat)
    }
}

#[allow(dead_code)]
fn detect_target_features(_arch: Arch) -> TargetFeatures {
    // Not present when no features are enabled; default to empty.
    let raw = env::var("CARGO_CFG_TARGET_FEATURE").unwrap_or_default();
    let set = raw
        .split(',')
        .map(str::trim)
        .filter(|s| !s.is_empty())
        .map(str::to_string)
        .collect();
    TargetFeatures { set }
}

#[allow(dead_code)]
fn force_target_features(arch: Arch) -> TargetFeatures {
    match arch {
        Arch::X86_64 => TargetFeatures::new(["avx2", "bmi2", "popcnt"]),
        Arch::Aarch64 => TargetFeatures::new(["sha3"]),
        Arch::Riscv64 => TargetFeatures::new(["v"]),
        // Force define only; no optional feature flags in auto.mk.
        Arch::Aarch64Be | Arch::Riscv32 | Arch::Ppc64le | Arch::Arm | Arch::Unknown => {
            TargetFeatures::EMPTY
        }
    }
}

/// Returns Some(true/false) if `MK_COMPILER_SUPPORTS_<FEATURE>` is set.
fn env_override(name: &str) -> Option<bool> {
    println!("cargo:rerun-if-env-changed={}", name);
    env::var(name).ok().map(|v| v.trim() == "1")
}

/// Confirmation check (not detection), in up to two stages:
///   1. Does the compiler *accept* the flag? (`cc::Build::is_flag_supported`,
///      cheap and cached.)
///   2. If `asm_probe` is `Some`, can it actually *emit* the instructions? We
///      compile that tiny inline-asm snippet under the flag via
///      `compiler_can_assemble` (compile-only, no linking), inheriting the same
///      target/cross configuration the real build uses. Pass `None` to check
///      flag acceptance only and skip this stage.
///
/// Honors the `MK_COMPILER_SUPPORTS_*` override, which short-circuits both
/// stages. On any negative/error result, warns and returns false so the caller
/// skips the flag rather than failing the build.
fn compiler_can_emit(override_var: &str, flag: &str, asm_probe: Option<&str>) -> bool {
    if let Some(v) = env_override(override_var) {
        return v;
    }

    // Stage 1: flag acceptance.
    match cc::Build::new().is_flag_supported(flag) {
        Ok(true) => {}
        Ok(false) => {
            println!(
                "cargo:warning=build_native_detect: C compiler does not accept `{}`; skipping it",
                flag
            );
            return false;
        }
        Err(e) => {
            println!(
                "cargo:warning=build_native_detect: could not probe flag `{}` ({}); skipping it",
                flag, e
            );
            return false;
        }
    }

    // Stage 2 (optional): actually emit the instructions.
    if let Some(probe) = asm_probe {
        if !compiler_can_assemble(flag, probe) {
            println!(
                "cargo:warning=build_native_detect: C compiler accepts `{}` but cannot assemble \
                 the corresponding instructions; skipping it",
                flag
            );
            return false;
        }
    }

    true
}

/// Compile `asm_probe` (a small C program with inline asm) with `flag` using a
/// throwaway `cc::Build`, compile-only. Uses `try_compile_intermediates` so it
/// returns a `Result` instead of panicking, compiles to object files without
/// linking (linking would break cross-compilation when the C compiler can't
/// link), and inherits the same configuration the real build uses. Cargo
/// metadata and warnings are silenced so this probe doesn't pollute output.
fn compiler_can_assemble(flag: &str, asm_probe: &str) -> bool {
    let out_dir = env::var("OUT_DIR").unwrap_or_else(|_| ".".into());
    let src_path = format!("{}/mld_emit_probe.c", out_dir);

    if std::fs::write(&src_path, asm_probe).is_err() {
        return false;
    }

    cc::Build::new()
        .file(&src_path)
        .flag(flag)
        .cargo_metadata(false)
        .cargo_warnings(false)
        .try_compile_intermediates()
        .is_ok()
}