cpufetch-rs 0.0.5

//! ARM64 architecture-specific CPU detection.
//!
//! On macOS, Apple Silicon chips are identified via `hw.cpufamily` sysctl and
//! distinguished by P-core / E-core counts.  On Linux and other platforms a
//! generic ARM fallback is returned.

use crate::cpu::info::Frequency;
use crate::cpu::{ArmFeatures, CpuError, CpuInfo, Vendor, Version};

/// Detect CPU information for ARM64 systems.
///
/// # Errors
///
/// Returns `CpuError` if CPU detection fails.
pub fn detect_cpu() -> Result<CpuInfo, CpuError> {
    // On macOS, attempt Apple Silicon identification first.
    #[cfg(all(target_os = "macos", feature = "macos"))]
    if let Some(info) = apple_silicon::detect() {
        return Ok(info);
    }

    // Generic ARM fallback (Linux, bare-metal, etc.)
    Ok(CpuInfo {
        vendor: Vendor::ARM,
        brand_string: String::from("ARM Processor"),
        version: Version {
            family: 0,
            model: 0,
            stepping: 0,
        },
        physical_cores: u32::try_from(num_cpus::get_physical()).unwrap_or(0),
        logical_cores: u32::try_from(num_cpus::get()).unwrap_or(0),
        frequency: Frequency::default(),
        cache_sizes: [None; 4],
        features: detect_arm_features(),
        microarch: None,
        hypervisor: None,
        peak_flops: None,
        p_cores: None,
        e_cores: None,
    })
}

// ── ARM feature detection ────────────────────────────────────────────────────

fn detect_arm_features() -> ArmFeatures {
    let mut features = ArmFeatures::empty();

    #[cfg(target_arch = "aarch64")]
    {
        // NEON / FP / ASIMD are mandatory on all AArch64 — set unconditionally.
        features |= ArmFeatures::NEON | ArmFeatures::FP | ArmFeatures::ASIMD;

        // On Apple Silicon (macOS aarch64), all M-series chips include AES, PMULL,
        // SHA2, CRC32, and LSE atomics as baseline features — set unconditionally.
        // On Linux, use runtime detection via the OS auxiliary vector.
        #[cfg(target_os = "macos")]
        {
            features |= ArmFeatures::AES | ArmFeatures::PMULL | ArmFeatures::SHA2;
            features |= ArmFeatures::CRC32 | ArmFeatures::ATOMICS;
        }

        // On Linux, parse /proc/cpuinfo "Features" line for runtime detection.
        // This is more portable than std::arch::is_aarch64_feature_detected!(),
        // which fails to compile with some cross-compilation toolchains.
        #[cfg(target_os = "linux")]
        if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
            if let Some(feat_line) = cpuinfo
                .lines()
                .find(|l| l.starts_with("Features"))
                .and_then(|l| l.split_once(':'))
                .map(|(_, v)| v)
            {
                let has = |name: &str| feat_line.split_whitespace().any(|f| f.eq_ignore_ascii_case(name));
                if has("aes") {
                    features |= ArmFeatures::AES;
                }
                if has("pmull") {
                    features |= ArmFeatures::PMULL;
                }
                if has("sha2") {
                    features |= ArmFeatures::SHA2;
                }
                if has("crc32") {
                    features |= ArmFeatures::CRC32;
                }
                if has("atomics") {
                    features |= ArmFeatures::ATOMICS;
                }
            }
        }
    }

    features
}

// ── Apple Silicon detection (macOS only) ─────────────────────────────────────

#[cfg(all(target_os = "macos", feature = "macos"))]
mod apple_silicon {
    use super::detect_arm_features;
    use crate::cpu::info::Frequency;
    use crate::cpu::uarch::Microarch;
    use crate::cpu::{CpuInfo, Vendor, Version};

    /// Read a sysctl key as a `u32`, reinterpreting signed bits correctly.
    ///
    /// Many macOS `hw.*` keys are `CTLTYPE_INT` (signed 32-bit) but contain values
    /// that look like large unsigned constants (e.g. CPU family IDs such as
    /// `0xDA33_D83D`).  Casting `i32 as u32` preserves the bit pattern.
    #[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
    fn sysctl_u32(name: &str) -> Option<u32> {
        use sysctl::{Ctl, CtlValue, Sysctl};
        let ctl = Ctl::new(name).ok()?;
        match ctl.value().ok()? {
            CtlValue::Int(i) => Some(i as u32),
            CtlValue::Uint(u) => Some(u),
            CtlValue::Long(l) => Some(l as u32),
            CtlValue::Ulong(u) => Some(u as u32),
            _ => None,
        }
    }

    /// Read a sysctl key as a `u64` for cache sizes (bytes).
    #[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
    fn sysctl_u64(name: &str) -> Option<u64> {
        use sysctl::{Ctl, CtlValue, Sysctl};
        let ctl = Ctl::new(name).ok()?;
        match ctl.value().ok()? {
            CtlValue::Int(i) => Some(i as u64),
            CtlValue::Uint(u) => Some(u64::from(u)),
            CtlValue::Long(l) => Some(l as u64),
            CtlValue::Ulong(u) | CtlValue::U64(u) => Some(u),
            CtlValue::S64(s) => Some(s as u64),
            _ => None,
        }
    }

    /// Read cache sizes from macOS sysctl.
    ///
    /// Returns `[L1i, L1d, L2, L3]` in KB. Uses P-core (perflevel0) values
    /// since those are the primary/larger caches.
    fn detect_cache_sizes() -> [Option<u32>; 4] {
        let bytes_to_kb = |bytes: u64| -> u32 {
            #[allow(clippy::cast_possible_truncation)]
            let kb = (bytes / 1024) as u32;
            kb
        };

        [
            sysctl_u64("hw.perflevel0.l1icachesize").map(bytes_to_kb),
            sysctl_u64("hw.perflevel0.l1dcachesize").map(bytes_to_kb),
            sysctl_u64("hw.perflevel0.l2cachesize").map(bytes_to_kb),
            // macOS doesn't expose L3 via perflevel sysctl; it may not exist
            // as a distinct level on Apple Silicon (the SLC is not reported here).
            None,
        ]
    }

    /// Known maximum P-core frequencies (MHz) for Apple Silicon chips.
    ///
    /// Apple does not expose CPU frequency via sysctl on Apple Silicon.
    /// These values come from official Apple specs and Geekbench/AnandTech measurements.
    fn lookup_frequency(generation: &str, variant: &str) -> Option<f64> {
        match (generation, variant) {
            // M1 family — all share the same Firestorm P-core at 3228 MHz
            ("M1", _) => Some(3228.0),
            // M2 family
            ("M2", " Max" | " Ultra") => Some(3680.0),
            ("M2", _) => Some(3490.0),
            // M3 family
            ("M3", _) => Some(4056.0),
            // M4 family
            ("M4", " Max") => Some(4608.0),
            ("M4", " Pro") => Some(4512.0),
            ("M4", _) => Some(4408.0),
            _ => None,
        }
    }

    /// Map `hw.cpufamily` → (generation label, Microarch).
    ///
    /// Values are constants from `<mach/machine.h>` (macOS 15 / Sequoia):
    /// ```text
    /// CPUFAMILY_ARM_FIRESTORM_ICESTORM  0x1b58_8bb3  (M1 / A14)
    /// CPUFAMILY_ARM_BLIZZARD_AVALANCHE  0xda33_d83d  (M2 / A15)
    /// CPUFAMILY_ARM_EVEREST_SAWTOOTH    0x8765_edea  (M3 / A16/A17 Pro)
    /// CPUFAMILY_ARM_COLL                0x17d5_b93a  (M4 / A18 Pro)
    /// ```
    fn classify_family(family: u32) -> Option<(&'static str, Microarch)> {
        match family {
            0x1b58_8bb3 => Some(("M1", Microarch::AppleM1)),
            0xda33_d83d => Some(("M2", Microarch::AppleM2)),
            0x8765_edea => Some(("M3", Microarch::AppleM3)),
            0x17d5_b93a => Some(("M4", Microarch::AppleM4)),
            _ => None,
        }
    }

    /// Determine the chip variant (Pro / Max / Ultra) from P-core and E-core counts.
    ///
    /// `hw.perflevel0.physicalcpu` = P-core count (performance, highest frequency)
    /// `hw.perflevel1.physicalcpu` = E-core count (efficiency, lower frequency)
    #[allow(clippy::match_same_arms)]
    fn chip_variant(generation: &str, p_cores: u32, e_cores: u32) -> &'static str {
        let total = p_cores + e_cores;
        match generation {
            "M1" => match (p_cores, e_cores, total) {
                (4, 4, 8) => "",        // M1 base
                (6, 2, 8) => " Pro",    // M1 Pro (8-core config)
                (8, 2, 10) => " Pro",   // M1 Pro (10-core config)
                (8, 4, 12) => " Max",   // M1 Max
                (_, _, 20) => " Ultra", // M1 Ultra
                _ => "",
            },
            "M2" => match (p_cores, e_cores, total) {
                (4, 4, 8) => "",        // M2 base
                (6, 4, 10) => " Pro",   // M2 Pro (10-core config)
                (8, 4, 12) => " Pro",   // M2 Pro (12-core config)
                (_, _, 16) => " Max",   // M2 Max
                (_, _, 24) => " Ultra", // M2 Ultra
                _ => "",
            },
            "M3" => match (p_cores, e_cores, total) {
                (4, 4, 8) => "",        // M3 base
                (5, 6, 11) => " Pro",   // M3 Pro (11-core config)
                (6, 6, 12) => " Pro",   // M3 Pro (12-core config)
                (12, 4, 16) => " Max",  // M3 Max
                (_, _, 32) => " Ultra", // M3 Ultra
                _ => "",
            },
            "M4" => match (p_cores, e_cores, total) {
                (4, 6, 10) => "",            // M4 base
                (10, 4, 14) => " Pro",       // M4 Pro 14-core
                (12, 4, 16) => " Max",       // M4 Max (16-core, if/when released)
                (_, _, 20) => " Max",        // M4 Max larger config
                (_, _, 28 | 40) => " Ultra", // M4 Ultra (speculative)
                _ => "",
            },
            _ => "",
        }
    }

    /// Perform Apple Silicon detection and return a populated `CpuInfo`.
    ///
    /// Returns `None` if the CPU family is unrecognised (non-Apple ARM hardware).
    pub fn detect() -> Option<CpuInfo> {
        let family = sysctl_u32("hw.cpufamily")?;
        let (generation, microarch) = classify_family(family)?;

        // P-cores are perflevel 0 (fastest), E-cores are perflevel 1.
        let p_cores = sysctl_u32("hw.perflevel0.physicalcpu").unwrap_or(0);
        let e_cores = sysctl_u32("hw.perflevel1.physicalcpu").unwrap_or(0);

        let variant = chip_variant(generation, p_cores, e_cores);
        let brand_string = format!("Apple {generation}{variant}");

        let physical_cores = u32::try_from(num_cpus::get_physical()).unwrap_or(0);
        let logical_cores = u32::try_from(num_cpus::get()).unwrap_or(0);

        let features = detect_arm_features();
        let cache_sizes = detect_cache_sizes();

        // Apple Silicon frequency from lookup table (not available via sysctl)
        let max_freq = lookup_frequency(generation, variant);
        let frequency = Frequency {
            base: None,
            max: max_freq,
            current: None,
        };

        // Peak FLOPS: NEON is 128-bit = 2 DP ops/cycle.
        // Apple Silicon has FMA so multiply-add counts as 2 FLOP/cycle.
        // Use P-core count and max frequency for peak calculation.
        let peak_flops = max_freq.map(|mhz| {
            let clock_ghz = mhz / 1000.0;
            let cores = f64::from(p_cores);
            let neon_dp_width = 2.0; // 128-bit NEON = 2 doubles
            let fma_factor = 2.0; // Apple Silicon always has FMA
            cores * clock_ghz * neon_dp_width * fma_factor
        });

        Some(CpuInfo {
            vendor: Vendor::Apple,
            brand_string,
            version: Version {
                family: 0,
                model: 0,
                stepping: 0,
            },
            physical_cores,
            logical_cores,
            frequency,
            cache_sizes,
            features,
            microarch: Some(microarch),
            hypervisor: None,
            peak_flops,
            p_cores: Some(p_cores),
            e_cores: Some(e_cores),
        })
    }

    #[cfg(test)]
    mod tests {
        use super::*;
        use crate::cpu::uarch::Microarch;

        #[test]
        fn test_known_families() {
            // M1 — CPUFAMILY_ARM_FIRESTORM_ICESTORM
            assert!(matches!(classify_family(0x1b58_8bb3), Some(("M1", Microarch::AppleM1))));
            // M2 — CPUFAMILY_ARM_BLIZZARD_AVALANCHE
            assert!(matches!(classify_family(0xda33_d83d), Some(("M2", Microarch::AppleM2))));
            // M3 — CPUFAMILY_ARM_EVEREST_SAWTOOTH
            assert!(matches!(classify_family(0x8765_edea), Some(("M3", Microarch::AppleM3))));
            // M4 — CPUFAMILY_ARM_COLL (0x17D5_B93A — the user's chip)
            assert!(matches!(classify_family(0x17d5_b93a), Some(("M4", Microarch::AppleM4))));
            // Unknown family → None (no panic)
            assert!(classify_family(0xdead_beef).is_none());
        }

        #[test]
        fn test_m4_pro_variant() {
            // M4 Pro 14-core: 10 P-cores + 4 E-cores
            assert_eq!(chip_variant("M4", 10, 4), " Pro");
            // M4 base: 4 P-cores + 6 E-cores
            assert_eq!(chip_variant("M4", 4, 6), "");
        }

        #[test]
        fn test_m1_variants() {
            assert_eq!(chip_variant("M1", 4, 4), ""); // M1
            assert_eq!(chip_variant("M1", 8, 2), " Pro"); // M1 Pro 10-core
            assert_eq!(chip_variant("M1", 8, 4), " Max"); // M1 Max
        }
    }
}

// ── Unit tests ───────────────────────────────────────────────────────────────

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

    #[test]
    fn test_detect_cpu_runs() {
        let info = detect_cpu().unwrap();
        // On macOS Apple Silicon the vendor will be Apple; on other ARM it is ARM.
        assert!(
            matches!(info.vendor, Vendor::ARM | Vendor::Apple),
            "Expected ARM or Apple vendor, got {:?}",
            info.vendor
        );
        assert!(!info.brand_string.is_empty());
        assert!(info.logical_cores > 0);
        assert!(info.physical_cores > 0);
    }

    #[test]
    #[cfg(all(target_os = "macos", feature = "macos"))]
    fn test_apple_silicon_detected() {
        let info = detect_cpu().unwrap();
        assert_eq!(info.vendor, Vendor::Apple, "On macOS aarch64, vendor must be Apple");
        assert!(info.microarch.is_some(), "Apple Silicon microarch must be detected");
        println!(
            "Detected: {} ({:?}) — {}P + {}E cores",
            info.brand_string, info.microarch, info.physical_cores, info.logical_cores,
        );
    }
}