Skip to main content

cpufetch_rs/cpu/
flags.rs

1//! CPU feature flag detection and representation.
2//!
3//! This module provides functionality for detecting and representing CPU feature flags
4//! across different architectures. It uses dynamic feature detection where available
5//! and falls back to static detection where necessary.
6
7use bitflags::bitflags;
8use serde::{Deserialize, Serialize};
9
10/// Error types specific to CPU feature detection
11#[derive(Debug, thiserror::Error)]
12pub enum FeatureError {
13    #[error("Feature detection not supported on this architecture")]
14    UnsupportedArch,
15    #[error("Failed to detect feature {0}")]
16    DetectionFailed(String),
17}
18
19bitflags! {
20    /// CPU features for x86/x86_64 architectures
21    #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
22    pub struct X86Features: u64 {
23        const SSE = 1 << 0;
24        const SSE2 = 1 << 1;
25        const SSE3 = 1 << 2;
26        const SSSE3 = 1 << 3;
27        const SSE4_1 = 1 << 4;
28        const SSE4_2 = 1 << 5;
29        const AVX = 1 << 6;
30        const AVX2 = 1 << 7;
31        const FMA = 1 << 8;
32        const BMI1 = 1 << 9;
33        const BMI2 = 1 << 10;
34        const F16C = 1 << 11;
35        const POPCNT = 1 << 12;
36        const AES = 1 << 13;
37        const AVX512F = 1 << 14;
38        const AVX512BW = 1 << 15;
39        const AVX512CD = 1 << 16;
40        const AVX512DQ = 1 << 17;
41        const AVX512VL = 1 << 18;
42    }
43}
44
45bitflags! {
46    /// CPU features for ARM architectures
47    #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
48    pub struct ArmFeatures: u64 {
49        const NEON = 1 << 0;
50        const AES = 1 << 1;
51        const PMULL = 1 << 2;
52        const SHA1 = 1 << 3;
53        const SHA2 = 1 << 4;
54        const CRC32 = 1 << 5;
55        const ATOMICS = 1 << 6;
56        const FP = 1 << 7;
57        const ASIMD = 1 << 8;
58        const FPHP = 1 << 9;
59        const ASIMDHP = 1 << 10;
60        const ASIMDDP = 1 << 11;
61        const ASIMDFHM = 1 << 12;
62    }
63}
64
65/// Detect CPU features for the current architecture
66///
67/// # Errors
68///
69/// Returns `FeatureError::UnsupportedArch` on non-x86 targets.
70#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
71pub fn detect_features() -> Result<X86Features, FeatureError> {
72    let mut features = X86Features::empty();
73
74    // Using is_x86_feature_detected! for runtime detection
75    if cfg!(target_arch = "x86") || cfg!(target_arch = "x86_64") {
76        if std::is_x86_feature_detected!("sse") {
77            features |= X86Features::SSE;
78        }
79        if std::is_x86_feature_detected!("sse2") {
80            features |= X86Features::SSE2;
81        }
82        if std::is_x86_feature_detected!("sse3") {
83            features |= X86Features::SSE3;
84        }
85        if std::is_x86_feature_detected!("ssse3") {
86            features |= X86Features::SSSE3;
87        }
88        if std::is_x86_feature_detected!("sse4.1") {
89            features |= X86Features::SSE4_1;
90        }
91        if std::is_x86_feature_detected!("sse4.2") {
92            features |= X86Features::SSE4_2;
93        }
94        if std::is_x86_feature_detected!("avx") {
95            features |= X86Features::AVX;
96        }
97        if std::is_x86_feature_detected!("avx2") {
98            features |= X86Features::AVX2;
99        }
100        if std::is_x86_feature_detected!("fma") {
101            features |= X86Features::FMA;
102        }
103        if std::is_x86_feature_detected!("bmi1") {
104            features |= X86Features::BMI1;
105        }
106        if std::is_x86_feature_detected!("bmi2") {
107            features |= X86Features::BMI2;
108        }
109        if std::is_x86_feature_detected!("f16c") {
110            features |= X86Features::F16C;
111        }
112        if std::is_x86_feature_detected!("popcnt") {
113            features |= X86Features::POPCNT;
114        }
115        if std::is_x86_feature_detected!("aes") {
116            features |= X86Features::AES;
117        }
118        if std::is_x86_feature_detected!("avx512f") {
119            features |= X86Features::AVX512F;
120        }
121        if std::is_x86_feature_detected!("avx512bw") {
122            features |= X86Features::AVX512BW;
123        }
124        if std::is_x86_feature_detected!("avx512cd") {
125            features |= X86Features::AVX512CD;
126        }
127        if std::is_x86_feature_detected!("avx512dq") {
128            features |= X86Features::AVX512DQ;
129        }
130        if std::is_x86_feature_detected!("avx512vl") {
131            features |= X86Features::AVX512VL;
132        }
133    }
134
135    Ok(features)
136}
137
138/// Detect CPU features for ARM architectures
139///
140/// # Errors
141///
142/// Returns `FeatureError` if feature detection is not supported.
143#[cfg(target_arch = "aarch64")]
144pub fn detect_features() -> Result<ArmFeatures, FeatureError> {
145    let mut features = ArmFeatures::empty();
146
147    // Using target_feature detection for ARM
148    if cfg!(target_arch = "aarch64") {
149        if std::arch::is_aarch64_feature_detected!("neon") {
150            features |= ArmFeatures::NEON;
151        }
152        if std::arch::is_aarch64_feature_detected!("aes") {
153            features |= ArmFeatures::AES;
154        }
155        if std::arch::is_aarch64_feature_detected!("pmull") {
156            features |= ArmFeatures::PMULL;
157        }
158        if std::arch::is_aarch64_feature_detected!("sha2") {
159            features |= ArmFeatures::SHA2;
160        }
161        if std::arch::is_aarch64_feature_detected!("crc") {
162            features |= ArmFeatures::CRC32;
163        }
164        if std::arch::is_aarch64_feature_detected!("lse") {
165            features |= ArmFeatures::ATOMICS;
166        }
167        if std::arch::is_aarch64_feature_detected!("fp") {
168            features |= ArmFeatures::FP;
169        }
170        if std::arch::is_aarch64_feature_detected!("asimd") {
171            features |= ArmFeatures::ASIMD;
172        }
173        // Note: Some features might not be available for detection in all environments
174    }
175
176    Ok(features)
177}
178
179/// Detect CPU features for unsupported architectures
180#[cfg(not(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64")))]
181pub fn detect_features() -> Result<(), FeatureError> {
182    Err(FeatureError::UnsupportedArch)
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    #[test]
190    fn test_x86_features_flags() {
191        let features = X86Features::SSE | X86Features::SSE2;
192        assert!(features.contains(X86Features::SSE));
193        assert!(features.contains(X86Features::SSE2));
194        assert!(!features.contains(X86Features::AVX));
195    }
196
197    #[test]
198    fn test_arm_features_flags() {
199        let features = ArmFeatures::NEON | ArmFeatures::AES;
200        assert!(features.contains(ArmFeatures::NEON));
201        assert!(features.contains(ArmFeatures::AES));
202        assert!(!features.contains(ArmFeatures::SHA2));
203    }
204
205    // Note: We can't reliably test actual feature detection in unit tests
206    // as it depends on the CPU capabilities of the test machine
207}