mldsa-native-rs 0.0.1-alpha.6

FFI bindings and optional wrapper for the mldsa-native ML-DSA implementation
Documentation
use std::env;
use std::path::PathBuf;

use bindgen::callbacks::ParseCallbacks;

struct Dirs {
    _manifest_dir: PathBuf,
    mldsa_src_dir: PathBuf,
    build_harness_dir: PathBuf,
    build_harness_extra_dir: PathBuf,
}

impl Dirs {
    fn new() -> Self {
        let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());

        let mldsa_src_dir = manifest_dir.join("mldsa-native").join("mldsa");
        let build_harness_dir = manifest_dir.join("mldsa-build-harness");
        let build_harness_extra_dir = manifest_dir.join("mldsa-build-harness-extra");

        Self {
            _manifest_dir: manifest_dir,
            mldsa_src_dir,
            build_harness_dir,
            build_harness_extra_dir,
        }
    }
}

#[cfg(feature = "native")]
mod build_native_support {
    use super::Dirs;
    use super::PathBuf;

    mod native_detect;

    pub fn configure_for_native(builder: &mut cc::Build, dirs: &Dirs) {
        println!("cargo:warning=(INFO) \"native\" feature flag enabled.");
        let detected = native_detect::detect();
        println!("cargo:warning=(INFO) {detected:?}");
        detected.apply(builder);

        builder
            // We provide our own `mld_sys_check_capability()` to dispatch
            // between native and portable implementations at runtime.
            .define("MLD_CONFIG_CUSTOM_CAPABILITY_FUNC", "")
            // Enables native arithmetic backend
            .define("MLD_CONFIG_USE_NATIVE_BACKEND_ARITH", "")
            // Enables native FIPS-202 backend
            .define("MLD_CONFIG_USE_NATIVE_BACKEND_FIPS202", "")
            // Adds the assembly sources as a separate compilation unit
            .file(dirs.build_harness_dir.join("mldsa_native_asm_all.S"));
    }

    #[derive(Debug)]
    struct StripEnumPrefix;

    impl bindgen::callbacks::ParseCallbacks for StripEnumPrefix {
        fn enum_variant_name(
            &self,
            enum_name: Option<&str>,
            original_variant_name: &str,
            _variant_value: bindgen::callbacks::EnumVariantValue,
        ) -> Option<String> {
            if Some("mld_sys_cap") == enum_name && original_variant_name.starts_with("MLD_SYS_CAP_")
            {
                // original_variant_name is the raw C name, e.g. "MLD_SYS_CAP_X86_64_AVX2"
                // We strip the prefix, because bindgen will prepend the `enum_name`.
                Some(
                    original_variant_name
                        .strip_prefix("MLD_SYS_CAP_")
                        .map(|s| s.to_string())?,
                )
            } else {
                None
            }
        }
    }

    pub fn generate_support_bindings(dirs: &Dirs) {
        let header_file = dirs.build_harness_dir.join("detect_capabilities.h");

        let bindings = bindgen::Builder::default()
            .clang_args([
                format!("-I{}", dirs.mldsa_src_dir.to_string_lossy()),
                format!("-I{}", dirs.build_harness_dir.to_string_lossy()),
                format!("-D{}", "MLD_CONFIG_CUSTOM_CAPABILITY_FUNC"),
            ])
            .header(header_file.to_string_lossy())
            // Use ctypes from ::core
            .use_core()
            // Use our custom parsing rules for renaming enums
            .parse_callbacks(Box::new(StripEnumPrefix))
            .constified_enum_module("mld_sys_cap")
            // Tell cargo to invalidate the built crate whenever any of the
            // included header files changed.
            .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
            // Finish the builder and generate the bindings.
            .generate()
            // Unwrap the Result and panic on failure.
            .expect("Unable to generate bindings");

        // Write the bindings to the $OUT_DIR/bindings.rs file.
        let out_path = PathBuf::from(std::env::var("OUT_DIR").unwrap());
        bindings
            .write_to_file(out_path.join("detect_capabilities_bindings.rs"))
            .expect("Couldn't write bindings!");
    }
}

fn compile_c_sources(dirs: &Dirs) {
    let cc_flags = [
        "-Wall",
        "-Wextra",
        "-Werror=unused-result",
        "-Wpedantic",
        "-Werror",
        "-Wmissing-prototypes",
        "-Wshadow",
        "-Wpointer-arith",
        "-Wredundant-decls",
        "-Wconversion",
        "-Wsign-conversion",
        "-Wno-long-long",
        "-Wno-unknown-pragmas",
        "-Wno-unused-command-line-argument",
        "-O3",
        "-fomit-frame-pointer",
        "-std=c99",
        "-pedantic",
        "-MMD",
    ];

    const RANDOMBYTES_INTERNAL_NAME: &str = "_mldsa_native_rs_internal_randombytes";
    println!("cargo:rustc-env=RANDOMBYTES_INTERNAL_NAME={RANDOMBYTES_INTERNAL_NAME}");

    let mut builder = cc::Build::new();

    builder
        .flags(&cc_flags)
        .define("randombytes", RANDOMBYTES_INTERNAL_NAME)
        .define("MLD_CONFIG_NAMESPACE_PREFIX", "mldsa")
        .includes([&dirs.mldsa_src_dir, &dirs.build_harness_dir])
        .file(dirs.build_harness_dir.join("mldsa_native_all.c"));

    #[cfg(feature = "native")]
    build_native_support::configure_for_native(&mut builder, dirs);
    #[cfg(feature = "native")]
    build_native_support::generate_support_bindings(dirs);

    builder.compile("mldsa_native");

    println!(
        "cargo::rerun-if-changed={}",
        dirs.mldsa_src_dir.to_string_lossy()
    );
    println!(
        "cargo::rerun-if-changed={}",
        dirs.build_harness_dir.to_string_lossy()
    );
}

/// Wrap the entire doc comment in a "```text" block.
///
/// Any existing "```" delimiters are removed; the content of those blocks
/// remains and its formatting is otherwise preserved.
#[derive(Debug)]
struct CommentFormattingEscaper;

impl ParseCallbacks for CommentFormattingEscaper {
    fn process_comment(&self, comment: &str) -> Option<String> {
        // get rid of lines containing triple-backtick delimiters, since we're about to add our own
        let comment = comment
            .lines()
            .filter(|line| !line.contains("```"))
            .collect::<Vec<_>>()
            .join("\n");

        Some(format!("```text\n{comment}\n```"))
    }
}

fn generate_bindings(dirs: &Dirs) {
    let wrapper_h_file = dirs.build_harness_extra_dir.join("wrapper.h");

    let bindings = bindgen::Builder::default()
        .clang_args([
            format!("-I{}", dirs.mldsa_src_dir.to_string_lossy()),
            format!("-I{}", dirs.build_harness_dir.to_string_lossy()),
        ])
        .use_core()
        .header(wrapper_h_file.to_string_lossy())
        // Tell cargo to invalidate the built crate whenever any of the
        // included header files changed.
        .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
        // Render all doc comments as preformatted plain text.
        .parse_callbacks(Box::new(CommentFormattingEscaper))
        // Finish the builder and generate the bindings.
        .generate()
        // Unwrap the Result and panic on failure.
        .expect("Unable to generate bindings");

    // Write the bindings to the $OUT_DIR/bindings.rs file.
    let out_path = PathBuf::from(std::env::var("OUT_DIR").unwrap());
    bindings
        .write_to_file(out_path.join("bindings.rs"))
        .expect("Couldn't write bindings!");
}

#[cfg(feature = "built_info")]
mod built_support {

    /// Return `Some(true)` if the working tree is dirty, `Some(false)`
    /// if clean, and `None` if we could not determine it (not a repo,
    /// or git2 error such as a shallow CI clone).
    /// Mirrors how `built` itself may end up with `None`.
    fn repo_is_dirty(repo: &git2::Repository) -> Option<bool> {
        // Exclude untracked and ignored files from the "dirty"
        // judgement to match the common definition (tracked-content
        // changes).
        let mut opts = git2::StatusOptions::new();
        opts.include_untracked(false).include_ignored(false);

        let statuses = repo.statuses(Some(&mut opts)).ok()?;
        let dirty = !statuses.is_empty();

        Some(dirty)
    }

    /// Policy (a deliberate cost/accuracy tradeoff):
    ///
    ///  * Not in a git repo  -> emit no git-related rerun directive.
    ///    `built` writes None for the git fields; nothing to keep
    ///    fresh.
    ///    (Cargo still re-runs the script if a crate source file
    ///    changes, via its default package scan, so non-git metadata
    ///    stays correct.)
    ///
    ///  * In a repo          -> watch `.git/HEAD`.
    ///    Re-runs on commit / checkout / branch switch, refreshing
    ///    GIT_COMMIT_HASH and GIT_HEAD_REF.
    ///    We do NOT force an every-build rerun here, preserving
    ///    build caching.
    ///
    ///  * In a repo, dirty   -> force an unconditional rerun (via a
    ///    path that never exists, which Cargo always treats as
    ///    "changed").
    ///    Once dirty, we re-observe on every build so GIT_DIRTY
    ///    tracks further edits and the eventual return to clean.
    ///
    /// KNOWN LIMITATION (accepted): a clean -> dirty transition
    /// caused by editing a file that is NOT one of THIS crate's
    /// build inputs (e.g., only README.md, or a sibling crate) will
    /// not trigger a rerun, so GIT_DIRTY can read stale-clean until
    /// some build input or HEAD changes.
    /// Edits to this crate's compiled sources DO trigger Cargo's
    /// normal rebuild, which re-observes dirtiness.
    /// In other words, GIT_DIRTY is fresh relative to this crate's
    /// build inputs + HEAD, not relative to the entire work tree.
    pub fn configure_built_rerun() {
        let manifest_dir = std::env::var("CARGO_MANIFEST_DIR")
            .expect("CARGO_MANIFEST_DIR is always set for build scripts");

        // Discover the enclosing repo (walking up from `manifest_dir`).
        // `None` if not in a repo or git2 errors.
        if let Some(repo) = git2::Repository::discover(manifest_dir).ok() {
            // `.path()` is the resolved git dir (e.g., `.../.git/`, or
            // the real dir a worktree/submodule gitfile points at).
            let git_dir = repo.path().to_path_buf();

            let git_head = git_dir.join("HEAD");
            if git_head.exists() {
                // Ask Cargo to watch HEAD so commits/checkouts
                // refresh the commit hash without disabling caching.
                println!("cargo:rerun-if-changed={}", git_head.display());
            } else {
                // Fallback: watch the git dir itself so ref changes
                // still trigger a rerun, even under reftable / worktree
                // / bare layouts where HEAD isn't a loose file at
                // path()/HEAD.
                println!("cargo:rerun-if-changed={}", git_dir.display());
            }

            if Some(true) == repo_is_dirty(&repo) {
                // Dirty: force re-run every build to keep GIT_DIRTY live.
                println!("cargo:rerun-if-changed=__force_rerun_while_dirty__");
            }
        } else {
            // Nothin to emit, this is intentionally empty
        }
    }
}

fn main() {
    let dirs = Dirs::new();

    compile_c_sources(&dirs);
    generate_bindings(&dirs);

    #[cfg(feature = "built_info")]
    {
        built::write_built_file().expect("Failed to acquire build-time information");
        built_support::configure_built_rerun();
    }
}