cpufeatures: add new! syntax for several target feature sets

cpufeatures/CHANGELOG.md

@@ -5,6 +5,13 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Unreleased
+### Added
+- `new!` syntax for declaring several alternative target feature sets, detected at once and
+  cached in a single atomic variable ([#1544])
+
+[#1544]: https://github.com/RustCrypto/utils/pull/1544
+
 ## 0.3.1 (2026-08-26)
 ### Changed
 - Use compile-time target feature detection under Miri ([#1513])

cpufeatures/README.md

@@ -56,6 +56,50 @@ compiler to completely eliminate fallback code.
 After first call macro caches result and returns it in subsequent
 calls, thus runtime overhead for them is minimal.
 
+## Example: several target feature sets
+
+Backends are often selected from a list of alternatives, in which case naming each set
+lets the macro detect all of them at once and cache the outcome in a single atomic
+variable:
+
+```rust
+#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+pub mod x86_dispatch {
+    // This macro creates `backend` module with a `Features` enum
+    cpufeatures::new!(
+        backend;
+        Avx2: "avx2", "aes";
+        Aes: "aes", "sse4.1";
+        Soft;
+    );
+
+    use backend::Features;
+
+    pub fn run() {
+        // A single relaxed load, regardless of the number of declared sets
+        match backend::get() {
+            Features::Avx2 => println!("AVX2 and AES extensions are supported"),
+            Features::Aes => println!("AES and SSE4.1 extensions are supported"),
+            Features::Soft => println!("no hardware acceleration is available"),
+        }
+
+        // `InitToken` works the same way as for a single target feature set
+        let token: backend::InitToken = backend::init();
+        assert_eq!(token.get(), backend::get());
+    }
+}
+```
+
+Sets are probed in the order in which they are declared and the first fully available one
+is selected, so they should be listed from the most to the least preferred. The trailing
+entry carries no target features and names the variant returned when none of the sets is
+available.
+
+If all target features of the *first* set are enabled via compiler options, it is always
+the selected one, so detection and the atomic load are eliminated entirely. The same happens
+on targets without runtime detection, such as SGX, UEFI and freestanding ones, where the
+selected set follows from the compiler options alone.
+
 ## Supported target architectures
 
 *NOTE: target features with an asterisk are unstable (nightly-only) and subject

cpufeatures/src/aarch64.rs

@@ -20,6 +20,21 @@ macro_rules! __unless_target_features {
     };
 }
 
+// Whether target features can be detected at run time. Detection uses `getauxval` on Linux
+// and Android and `sysctlbyname` on Apple platforms; elsewhere `__detect_target_features!`
+// below is a constant `false`.
+#[macro_export]
+#[doc(hidden)]
+macro_rules! __runtime_detection_available {
+    () => {
+        cfg!(any(
+            target_os = "linux",
+            target_os = "android",
+            target_vendor = "apple"
+        ))
+    };
+}
+
 // Linux runtime detection of target CPU features using `getauxval`.
 #[cfg(any(target_os = "linux", target_os = "android"))]
 #[macro_export]

cpufeatures/src/lib.rs

@@ -31,6 +31,42 @@ mod miri;
 compile_error!("This crate works only on `aarch64`, `loongarch64`, `x86`, and `x86-64` targets.");
 
 /// Create module with CPU feature detection code.
+///
+/// # Single target feature set
+///
+/// The module gets a `get` function returning whether all listed target features are available:
+///
+/// ```rust,ignore
+/// cpufeatures::new!(aes_sha, "aes", "sha");
+///
+/// if aes_sha::get() {
+///     // ...
+/// }
+/// ```
+///
+/// # Multiple target feature sets
+///
+/// Several named sets can be declared instead, separated with `;`. The module then gets a
+/// `Features` enum with one variant per set and `get` returns the first variant whose target
+/// features are all available. The trailing entry carries no target features and names the
+/// variant returned when none of the sets is available.
+///
+/// ```rust,ignore
+/// cpufeatures::new!(
+///     backend;
+///     Avx2: "avx2", "aes";
+///     Aes: "aes", "sse4.1";
+///     Soft;
+/// );
+///
+/// use backend::Features;
+///
+/// match backend::get() {
+///     Features::Avx2 => { /* ... */ }
+///     Features::Aes => { /* ... */ }
+///     Features::Soft => { /* ... */ }
+/// }
+/// ```
 #[macro_export]
 macro_rules! new {
     ($mod_name:ident, $($tf:tt),+ $(,)?) => {
@@ -106,4 +142,184 @@ macro_rules! new {
             }
         }
     };
+    // Every entry has the same shape, a name optionally followed by target features, so that the
+    // matcher never has to decide between starting one more set and finishing the list.
+    // Spelling the trailing entry out as a bare name in its own position instead would be a local
+    // ambiguity, as both alternatives bind an `ident`.
+    ($mod_name:ident; $($variant:ident $(: $($tf:tt),+)?);+ $(;)?) => {
+        mod $mod_name {
+            use core::sync::atomic::{AtomicU8, Ordering::Relaxed};
+
+            /// Target feature set detected at runtime.
+            ///
+            /// Variants are ordered as declared, i.e. the detected one is always the first variant
+            /// whose target features are all available. The last variant declares no target
+            /// features and is the one detected when no other is available.
+            #[derive(Copy, Clone, Debug, Eq, PartialEq)]
+            #[repr(u8)]
+            pub enum Features {
+                $(
+                    #[doc = concat!("The `", stringify!($variant), "` target feature set.")]
+                    $(
+                        #[doc = concat!("\nAvailable target features:", $(" `", $tf, "`",)+)]
+                    )?
+                    $variant,
+                )*
+            }
+
+            // Whether each entry declares target features. `!$tf.is_empty()` is a constant `true`
+            // whose only purpose is to depend on `$tf`, so that the term is emitted exactly for
+            // the entries that have one.
+            const HAS_TARGET_FEATURES: &[bool] = &[$(false $($(|| !$tf.is_empty())+)?),+];
+
+            const _: () = {
+                let len = HAS_TARGET_FEATURES.len();
+
+                assert!(len > 1, "`cpufeatures::new!` expects at least one target feature set");
+                assert!(
+                    !HAS_TARGET_FEATURES[len - 1],
+                    "the last `cpufeatures::new!` entry names the variant detected when none \
+                     of the target feature sets is available and must not declare any itself"
+                );
+
+                let mut i = 0;
+                while i < len - 1 {
+                    assert!(
+                        HAS_TARGET_FEATURES[i],
+                        "only the last `cpufeatures::new!` entry may omit target features"
+                    );
+                    i += 1;
+                }
+            };
+
+            /// Variant detected when none of the target feature sets is available.
+            const FALLBACK: Features = {
+                let all = [$(Features::$variant),+];
+                all[all.len() - 1]
+            };
+
+            // Value stored in `STORAGE` until CPU feature detection has been performed.
+            //
+            // `Features` is `#[repr(u8)]` and does not use explicit discriminants, so its tags are
+            // exactly `0..=FALLBACK` and the value right past the last variant can not collide
+            // with any of them.
+            const UNINIT: u8 = FALLBACK as u8 + 1;
+
+            // Every `Features` tag has to stay below `UNINIT`, otherwise `init_get` could not
+            // tell the uninitialized state apart and the transmutes below would be unsound.
+            const _: () = {
+                $(assert!((Features::$variant as u8) < UNINIT);)*
+            };
+
+            // Set when the detected variant follows from compile-time information alone, which
+            // happens in two cases: the first declared set is enabled at compile time, as it is
+            // probed first and therefore always wins, and the target having no runtime detection
+            // at all, where no set beyond the enabled ones can ever be found. Either way no
+            // storage is involved and `init_inner` never runs.
+            const STATICALLY_DETECTED: Option<Features> = {
+                // The last entry declares no target features, so it is always enabled and
+                // terminates the search below.
+                let enabled = [$(cfg!(all($($(target_feature = $tf,)+)?))),+];
+
+                if enabled[0] || !$crate::__runtime_detection_available!() {
+                    let mut i = 0;
+                    while !enabled[i] {
+                        i += 1;
+                    }
+
+                    Some([$(Features::$variant),+][i])
+                } else {
+                    None
+                }
+            };
+
+            static STORAGE: AtomicU8 = AtomicU8::new(UNINIT);
+
+            /// Initialization token
+            #[derive(Copy, Clone, Debug)]
+            pub struct InitToken(());
+
+            impl InitToken {
+                /// Initialize token, performing CPU feature detection.
+                pub fn init() -> Self {
+                    init()
+                }
+
+                /// Initialize token and return the detected target feature set.
+                pub fn init_get() -> (Self, Features) {
+                    init_get()
+                }
+
+                /// Get initialized value.
+                #[inline(always)]
+                pub fn get(&self) -> Features {
+                    match STATICALLY_DETECTED {
+                        Some(features) => features,
+                        None => {
+                            let val = STORAGE.load(Relaxed);
+
+                            // SAFETY: `InitToken` can only be obtained from `init_get`, which
+                            // stores the tag of a valid `Features` value into `STORAGE` before
+                            // constructing the token, and the tag is never modified afterwards.
+                            unsafe { core::mem::transmute::<u8, Features>(val) }
+                        }
+                    }
+                }
+            }
+
+            #[cold]
+            fn init_inner() -> Features {
+                let res = 'detect: {
+                    $($(
+                        if $crate::__unless_target_features! {
+                            $($tf),+ => { $crate::__detect_target_features!($($tf),+) }
+                        } {
+                            break 'detect Features::$variant;
+                        }
+                    )?)*
+
+                    FALLBACK
+                };
+
+                STORAGE.store(res as u8, Relaxed);
+
+                res
+            }
+
+            /// Get detected target feature set and initialization token, initializing underlying
+            /// storage if needed.
+            #[inline]
+            pub fn init_get() -> (InitToken, Features) {
+                let res = match STATICALLY_DETECTED {
+                    Some(features) => features,
+                    None => {
+                        // Relaxed ordering is fine, as we only have a single atomic variable.
+                        let val = STORAGE.load(Relaxed);
+
+                        if val == UNINIT {
+                            init_inner()
+                        } else {
+                            // SAFETY: `STORAGE` contains either `UNINIT`, which is handled above,
+                            // or the tag of a valid `Features` value written by `init_inner`.
+                            unsafe { core::mem::transmute::<u8, Features>(val) }
+                        }
+                    }
+                };
+
+                (InitToken(()), res)
+            }
+
+            /// Initialize underlying storage if needed and get initialization token.
+            #[inline]
+            pub fn init() -> InitToken {
+                init_get().0
+            }
+
+            /// Initialize underlying storage if needed and get detected target feature set.
+            #[inline]
+            pub fn get() -> Features {
+                init_get().1
+            }
+        }
+    };
 }

cpufeatures/src/loongarch64.rs

@@ -18,6 +18,16 @@ macro_rules! __unless_target_features {
     };
 }
 
+// Whether target features can be detected at run time. Off Linux `__detect_target_features!`
+// below is a constant `false`.
+#[macro_export]
+#[doc(hidden)]
+macro_rules! __runtime_detection_available {
+    () => {
+        cfg!(target_os = "linux")
+    };
+}
+
 // Linux runtime detection of target CPU features using `getauxval`.
 #[cfg(target_os = "linux")]
 #[macro_export]

cpufeatures/src/miri.rs

@@ -3,6 +3,14 @@
 //! Miri is an interpreter, and though it tries to emulate the target CPU
 //! it does not support any target features.
 
+#[macro_export]
+#[doc(hidden)]
+macro_rules! __runtime_detection_available {
+    () => {
+        false
+    };
+}
+
 #[macro_export]
 #[doc(hidden)]
 macro_rules! __unless_target_features {

cpufeatures/src/x86.rs

@@ -28,6 +28,22 @@ macro_rules! __unless_target_features {
     }};
 }
 
+/// Whether target features can be detected at run time.
+///
+/// CPUID is not available on SGX. Freestanding and UEFI targets do not support SIMD
+/// features with default compilation flags.
+#[macro_export]
+#[doc(hidden)]
+macro_rules! __runtime_detection_available {
+    () => {
+        cfg!(not(any(
+            target_env = "sgx",
+            target_os = "none",
+            target_os = "uefi"
+        )))
+    };
+}
+
 /// Use CPUID to detect the presence of all supplied target features.
 #[macro_export]
 #[doc(hidden)]