Hardware Optimization in Rust: Cache-Friendly Layouts, Data-Oriented Design, and SIMD with std::arch

// SLOW: each shape is a heap allocation; vtable dispatch on every call
trait Shape { fn area(&self) -> f32; }
struct Circle  { r: f32 }
struct Rect    { w: f32, h: f32 }

impl Shape for Circle { fn area(&self) -> f32 { std::f32::consts::PI * self.r * self.r } }
impl Shape for Rect   { fn area(&self) -> f32 { self.w * self.h } }

let shapes: Vec<Box<dyn Shape>> = vec![
    Box::new(Circle { r: 1.0 }),
    Box::new(Rect { w: 2.0, h: 3.0 }),
];
let total: f32 = shapes.iter().map(|s| s.area()).sum();

// FAST: all variants in one enum, stored in a flat Vec — zero indirection
#[derive(Copy, Clone)]
enum ShapeFlat {
    Circle { r: f32 },
    Rect   { w: f32, h: f32 },
}

impl ShapeFlat {
    fn area(self) -> f32 {
        match self {
            ShapeFlat::Circle { r } => std::f32::consts::PI * r * r,
            ShapeFlat::Rect { w, h } => w * h,
        }
    }
}

let shapes: Vec<ShapeFlat> = vec![
    ShapeFlat::Circle { r: 1.0 },
    ShapeFlat::Rect { w: 2.0, h: 3.0 },
];
let total: f32 = shapes.iter().copied().map(ShapeFlat::area).sum();
// All data is contiguous — CPU prefetcher loves this

use std::arch::x86_64::*;

/// Sum a slice of f32s, using AVX2 if available, scalar otherwise.
pub fn sum_f32(data: &[f32]) -> f32 {
    // Runtime CPU feature detection — safe to call anywhere
    if is_x86_feature_detected!("avx2") {
        // SAFETY: we just checked that AVX2 is available
        unsafe { sum_f32_avx2(data) }
    } else {
        data.iter().sum()
    }
}

/// AVX2 implementation: processes 8 f32s per iteration.
///
/// # Safety
/// Caller must guarantee that AVX2 is available on this CPU.
#[target_feature(enable = "avx2")]
unsafe fn sum_f32_avx2(data: &[f32]) -> f32 {
    let mut acc = _mm256_setzero_ps(); // [0.0 × 8]

    let chunks = data.len() / 8;
    for i in 0..chunks {
        // Load 8 f32s from memory (unaligned load is fine on modern CPUs)
        let v = _mm256_loadu_ps(data.as_ptr().add(i * 8));
        // Add 8 floats in parallel: acc += v
        acc = _mm256_add_ps(acc, v);
    }

    // Horizontal reduction: sum all 8 lanes of 'acc' into one f32
    // Step 1: fold upper 4 lanes into lower 4
    let low  = _mm256_castps256_ps128(acc);
    let high = _mm256_extractf128_ps(acc, 1);
    let sum4 = _mm_add_ps(low, high);
    // Step 2: fold upper 2 lanes into lower 2
    let shuf = _mm_movehdup_ps(sum4);
    let sum2 = _mm_add_ps(sum4, shuf);
    // Step 3: fold upper 1 lane into lower 1
    let shuf = _mm_movehl_ps(shuf, sum2);
    let sum1 = _mm_add_ss(sum2, shuf);
    let result = _mm_cvtss_f32(sum1);

    // Handle remainder elements the scalar way
    let remainder: f32 = data[chunks * 8..].iter().sum();
    result + remainder
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "sse2")]
pub unsafe fn find_byte_sse2(haystack: &[u8], needle: u8) -> Option<usize> {
    use std::arch::x86_64::*;

    // Broadcast 'needle' to all 16 lanes of an SSE2 register
    let target = _mm_set1_epi8(needle as i8);

    let mut i = 0;
    while i + 16 <= haystack.len() {
        // Load 16 bytes from haystack
        let chunk = _mm_loadu_si128(haystack.as_ptr().add(i) as *const __m128i);
        // Compare all 16 bytes simultaneously: result lane = 0xFF if equal, 0x00 if not
        let cmp = _mm_cmpeq_epi8(chunk, target);
        // Pack comparison results into a 16-bit bitmask (1 bit per byte)
        let mask = _mm_movemask_epi8(cmp) as u32;
        if mask != 0 {
            // trailing_zeros gives the position of the first match within the chunk
            return Some(i + mask.trailing_zeros() as usize);
        }
        i += 16;
    }

    // Scalar fallback for remaining bytes
    haystack[i..].iter().position(|&b| b == needle).map(|pos| i + pos)
}