mldsa-native-rs 0.0.1-alpha.6

#![allow(non_camel_case_types)]
#![allow(dead_code)]

mod ffi {
    #![allow(non_upper_case_globals)]
    #![allow(non_camel_case_types)]
    #![allow(non_snake_case)]

    include!(concat!(env!("OUT_DIR"), "/detect_capabilities_bindings.rs"));

    pub use ::core::ffi::c_int;
}
use ffi::*;

// These constants are what mldsa-native expects as return values for `mld_sys_check_capability`
const SUPPORTED: c_int = 1;
const NOT_SUPPORTED: c_int = 0;

#[derive(Debug, Clone, Copy)]
pub struct RuntimeCapabilities {
    pub x86_64_avx2: bool,
    pub aarch64_sha3: bool,
    pub armv81m_mve: bool,
}

impl RuntimeCapabilities {
    pub fn probe() -> Self {
        Self {
            x86_64_avx2: check_capability(mld_sys_cap::X86_64_AVX2),
            aarch64_sha3: check_capability(mld_sys_cap::AARCH64_SHA3),
            armv81m_mve: check_capability(mld_sys_cap::ARMV81M_MVE),
        }
    }
}

pub fn mld_sys_cap_to_str(cap: mld_sys_cap::Type) -> &'static str {
    match cap {
        mld_sys_cap::X86_64_AVX2 => stringify!(mld_sys_cap::X86_64_AVX2),
        mld_sys_cap::AARCH64_SHA3 => stringify!(mld_sys_cap::AARCH64_SHA3),
        mld_sys_cap::ARMV81M_MVE => stringify!(mld_sys_cap::ARMV81M_MVE),
        _ => unreachable!("Unknown mld_sys_cap value: {cap:?}"),
    }
}

#[unsafe(no_mangle)]
#[inline(always)]
pub extern "C" fn mldrs_sys_check_capability(cap: mld_sys_cap::Type) -> c_int {
    let is_supported: bool = check_capability(cap);
    //eprintln!("check_capability({cap:?}) -> {is_supported:?}");
    bool_to_c_int(is_supported)
}

#[inline(always)]
pub fn check_capability(cap: mld_sys_cap::Type) -> bool {
    cfg_if::cfg_if! {
        if #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] {
            x86_backend::check_capability(cap)
        } else if #[cfg(any(target_arch = "aarch64"))] {
            aarch64_backend::check_capability(cap)
        } else {
            let _ = cap;
            false
        }
    }
}

#[inline(always)]
fn bool_to_c_int(v: bool) -> c_int {
    if v { SUPPORTED } else { NOT_SUPPORTED }
}

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub(super) mod x86_backend {
    use super::*;

    // The list of required CPU flags comes from
    // https://github.com/pq-code-package/mldsa-native/blob/13977d0a2b613e7f1ea2c011d4e0c760ff8992ff/integration/liboqs/ML-DSA-44_META.yml
    cpufeatures::new!(cpufeats_avx2_flags, "avx2", "bmi2", "popcnt");

    pub fn check_capability(cap: mld_sys_cap::Type) -> bool {
        let token: cpufeats_avx2_flags::InitToken = cpufeats_avx2_flags::init();
        match cap {
            mld_sys_cap::X86_64_AVX2 => token.get() == true,
            //mld_sys_cap::AARCH64_SHA3 => unreachable!("mld_sys_cap::AARCH64_SHA3 should not be tested on x86"),
            _ => false,
        }
    }
}

#[cfg(any(target_arch = "aarch64"))]
pub(super) mod aarch64_backend {
    use super::*;

    cpufeatures::new!(cpufeats_sha3_flags, "sha3");

    pub fn check_capability(cap: mld_sys_cap::Type) -> bool {
        let token: cpufeats_sha3_flags::InitToken = cpufeats_sha3_flags::init();
        match cap {
            // mld_sys_cap::X86_64_AVX2 => { unreachable!("mld_sys_cap::X86_64_AVX2 should not be tested on aarch64") }
            mld_sys_cap::AARCH64_SHA3 => token.get() == true,
            _ => false,
        }
    }
}

/// This module is dedicated to portably implement the "call_once" logic
/// to initialize the CPU feature detection once so it is accessible from
/// the C side.
mod init_once {
    pub use ::core::ffi::c_int;

    #[derive(Copy, Clone, Debug)]
    struct InitError();

    type InitResult<T> = core::result::Result<T, InitError>;

    /// This internal function wraps the extern C function to use Rust-y Result semantics
    #[inline(always)]
    fn mldrs_init_sys_check_capability_wrapper() -> InitResult<()> {
        unsafe extern "C" {
            /// The actual initialization is done by this function implemented on the C-side.
            ///
            /// Returns 1 on SUCCESS, 0 on FAILURE
            fn mldrs_init_sys_check_capability() -> c_int;
        }

        let result = unsafe { mldrs_init_sys_check_capability() };

        if result != 1 {
            Err(InitError())
        } else {
            Ok(())
        }
    }

    /// This function ensures the initialization is executed only once,
    /// even when concurrent threads call it.
    ///
    /// This is just a Rust wrapper to ensure protable "call_once" logic,
    /// the actual initialization is done by the C side.
    #[unsafe(no_mangle)]
    pub extern "C" fn mldrs_call_once_init_sys_check_capability() {
        use std::sync::OnceLock;

        static MLDRS_INIT_SYS_CHECK_CAPABILITY_RESULT: OnceLock<InitResult<()>> = OnceLock::new();

        let result = *MLDRS_INIT_SYS_CHECK_CAPABILITY_RESULT
            .get_or_init(|| mldrs_init_sys_check_capability_wrapper());

        if result.is_err() {
            unreachable! {"mldrs_sys_check_capability() failed!"}
        }
    }

    #[cfg(test)]
    mod test {
        use super::super::ffi;
        use super::super::*;
        use super::*;

        // Test entrypoint provided by the C-side: it wraps the inlined
        // `mld_sys_check_capability()` used by mldsa-native.
        use ffi::mldrs_test_mld_sys_check_capability;

        #[test]
        fn mldrs_init_sys_check_capability_fails_if_called_twice() {
            // We do not know if some other test already triggered the initialization,
            // so we cannot assume the next line is the first call: we cannot assert
            // it should succeed or fail.
            let _ = mldrs_init_sys_check_capability_wrapper();

            // At this point the initialization was triggered at least once, so we can assert subsequent inits fail.
            assert!(
                mldrs_init_sys_check_capability_wrapper().is_err(),
                "Should fail the second time"
            );
        }

        fn assert_capability_result(cap: mld_sys_cap::Type, cap_name: &'static str) {
            let got: c_int = unsafe { mldrs_test_mld_sys_check_capability(cap) };

            println!("Testing `mld_sys_check_capability({cap_name})` returned {got:?}");

            assert!(
                got == SUPPORTED || got == NOT_SUPPORTED,
                "unexpected capability result for {cap_name}: {got}"
            );

            let rs_got: c_int = mldrs_sys_check_capability(cap);
            assert_eq!(got, rs_got, "mismatched capability result for {cap_name}");
        }

        macro_rules! capability_tests {
            (
                $(
                    $test_name:ident => $cap:path
                ),+ $(,)?
            ) => {
                $(
                    #[test]
                    fn $test_name() {
                        println!("");
                        assert_capability_result($cap, stringify!($cap));
                        print!("... ");
                    }
                )+
            };
        }

        capability_tests! {
            test_mld_sys_cap_x86_64_avx2 => mld_sys_cap::X86_64_AVX2,
            test_mld_sys_cap_aarch64_sha3 => mld_sys_cap::AARCH64_SHA3,
            test_mld_sys_cap_armv81m_mve => mld_sys_cap::ARMV81M_MVE,
        }

        #[test]
        fn test_all_mld_sys_cap() {
            use ffi::MLDRS_SYS_CAP_MAX_VALUE;

            println!("\nTesting all capabilities");
            let max = MLDRS_SYS_CAP_MAX_VALUE;
            for cap in 0..=max {
                let cap_name = mld_sys_cap_to_str(cap);
                assert_capability_result(cap, cap_name);
            }
            print!("DONE\t... ");
        }
    }
}