foxhunt/trading_engine/src/lockfree/small_batch_ring.rs

#![allow(unsafe_code)] // Intentional unsafe for small batch lock-free operations

//! Optimized lock-free ring buffer for small batch processing
//!
//! This implementation is specifically designed for small batch order processing
//! with minimal atomic operations overhead and cache-optimized memory layout.

use std::alloc::{alloc, dealloc, Layout};
use std::cell::UnsafeCell;
use std::ptr::NonNull;
use std::sync::atomic::{compiler_fence, AtomicU64, Ordering};

/// Small batch optimized ring buffer with reduced atomic operations
///
/// This implementation uses compiler fences instead of expensive memory barriers
/// for single-threaded small batch processing, achieving significant performance
/// improvements for 1-10 order batches.
#[repr(align(64))] // Cache line alignment
/// SmallBatchRing
/// 
/// TODO: Add detailed documentation for this struct
pub struct SmallBatchRing<T> {
    buffer: NonNull<UnsafeCell<T>>,
    capacity: usize,
    mask: usize,

    // Hot cache line - frequently accessed
    head: AtomicU64, // Producer position
    tail: AtomicU64, // Consumer position

    // Cold cache line - metadata
    layout: Layout,
    batch_mode: BatchMode,
}

/// Batch processing mode for optimization
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
/// BatchMode
/// 
/// TODO: Add detailed documentation for this enum
pub enum BatchMode {
    /// Single threaded mode - uses compiler fences only
    SingleThreaded,
    /// Multi-threaded mode - uses full memory barriers
    MultiThreaded,
}

// SAFETY: SmallBatchRing is Send because:
// - Batch processing coordinated through atomic batch_head/write_head indices
// - Memory allocation/deallocation uses safe Layout operations
// - All batch operations maintain memory safety through bounds checking
unsafe impl<T: Send> Send for SmallBatchRing<T> {}

// SAFETY: SmallBatchRing is Sync because:
// - Concurrent batch access coordinated via atomic operations (Relaxed ordering)
// - Small batch design (4-8 elements) minimizes contention window
// - Copy trait requirement ensures no shared ownership issues
// - Memory ordering guarantees prevent data races on batch boundaries
unsafe impl<T: Send> Sync for SmallBatchRing<T> {}

impl<T: Copy> SmallBatchRing<T> {
    /// Create new small batch ring buffer
    pub fn new(capacity: usize, batch_mode: BatchMode) -> Result<Self, &'static str> {
        if capacity == 0 {
            return Err("Capacity cannot be zero");
        }

        if !capacity.is_power_of_two() {
            return Err("Capacity must be power of two");
        }

        let layout =
            Layout::array::<UnsafeCell<T>>(capacity).map_err(|_| "Layout creation failed")?;

        let buffer = unsafe {
            let ptr = alloc(layout);
            if ptr.is_null() {
                return Err("Memory allocation failed");
            }
            NonNull::new_unchecked(ptr.cast::<UnsafeCell<T>>())
        };

        Ok(Self {
            buffer,
            capacity,
            mask: capacity - 1,
            head: AtomicU64::new(0),
            tail: AtomicU64::new(0),
            layout,
            batch_mode,
        })
    }

    /// Push batch of items with optimized path
    #[inline(always)]
    pub fn push_batch(&self, items: &[T]) -> Result<usize, usize> {
        if items.is_empty() {
            return Ok(0);
        }

        match self.batch_mode {
            BatchMode::SingleThreaded => self.push_batch_st(items),
            BatchMode::MultiThreaded => self.push_batch_mt(items),
        }
    }

    /// Single-threaded batch push with compiler fences only
    #[inline(always)]
    fn push_batch_st(&self, items: &[T]) -> Result<usize, usize> {
        let head = self.head.load(Ordering::Relaxed);
        let tail = self.tail.load(Ordering::Relaxed);

        let available = self.capacity.saturating_sub((head - tail) as usize);
        let push_count = items.len().min(available);

        if push_count == 0 {
            return Err(0);
        }

        // Write items to buffer
        for (i, &item) in items.iter().take(push_count).enumerate() {
            let index = ((head + i as u64) as usize) & self.mask;
            unsafe {
                (*self.buffer.as_ptr().add(index)).get().write(item);
            }
        }

        // Compiler fence ensures writes complete before head update
        compiler_fence(Ordering::SeqCst);

        // Update head position
        self.head.store(head + push_count as u64, Ordering::Relaxed);

    // Ok variant
        Ok(push_count)
    }

    /// Multi-threaded batch push with full memory barriers
    #[inline(always)]
    fn push_batch_mt(&self, items: &[T]) -> Result<usize, usize> {
        let head = self.head.load(Ordering::Relaxed);
        let tail = self.tail.load(Ordering::Acquire);

        let available = self.capacity.saturating_sub((head - tail) as usize);
        let push_count = items.len().min(available);

        if push_count == 0 {
            return Err(0);
        }

        // Write items to buffer
        for (i, &item) in items.iter().take(push_count).enumerate() {
            let index = ((head + i as u64) as usize) & self.mask;
            unsafe {
                (*self.buffer.as_ptr().add(index)).get().write(item);
            }
        }

        // Release ordering ensures writes are visible before head update
        self.head.store(head + push_count as u64, Ordering::Release);

    // Ok variant
        Ok(push_count)
    }

    /// Pop batch of items with optimized path
    #[inline(always)]
    pub fn pop_batch(&self, output: &mut [T]) -> usize {
        if output.is_empty() {
            return 0;
        }

        match self.batch_mode {
            BatchMode::SingleThreaded => self.pop_batch_st(output),
            BatchMode::MultiThreaded => self.pop_batch_mt(output),
        }
    }

    /// Single-threaded batch pop with compiler fences only
    #[inline(always)]
    fn pop_batch_st(&self, output: &mut [T]) -> usize {
        let tail = self.tail.load(Ordering::Relaxed);
        let head = self.head.load(Ordering::Relaxed);

        let available = (head - tail) as usize;
        let pop_count = output.len().min(available);

        if pop_count == 0 {
            return 0;
        }

        // Read items from buffer
        for i in 0..pop_count {
            let index = ((tail + i as u64) as usize) & self.mask;
            unsafe {
                output[i] = (*self.buffer.as_ptr().add(index)).get().read();
            }
        }

        // Compiler fence ensures reads complete before tail update
        compiler_fence(Ordering::SeqCst);

        // Update tail position
        self.tail.store(tail + pop_count as u64, Ordering::Relaxed);

        pop_count
    }

    /// Multi-threaded batch pop with full memory barriers
    #[inline(always)]
    fn pop_batch_mt(&self, output: &mut [T]) -> usize {
        let tail = self.tail.load(Ordering::Relaxed);
        let head = self.head.load(Ordering::Acquire);

        let available = (head - tail) as usize;
        let pop_count = output.len().min(available);

        if pop_count == 0 {
            return 0;
        }

        // Read items from buffer
        for i in 0..pop_count {
            let index = ((tail + i as u64) as usize) & self.mask;
            unsafe {
                output[i] = (*self.buffer.as_ptr().add(index)).get().read();
            }
        }

        // Release ordering ensures reads complete before tail update
        self.tail.store(tail + pop_count as u64, Ordering::Release);

        pop_count
    }

    /// Try to push single item (optimized for small batches)
    #[inline(always)]
    pub fn try_push(&self, item: T) -> Result<(), T> {
        let items = [item];
        match self.push_batch(&items) {
            Ok(1) => Ok(()),
            _ => Err(item),
        }
    }

    /// Try to pop single item (optimized for small batches)
    #[inline(always)]
    pub fn try_pop(&self) -> Option<T> {
        let mut output = [unsafe { std::mem::zeroed() }];
        (self.pop_batch(&mut output) == 1).then(|| output[0])
    }

    /// Get current buffer utilization
    #[inline]
    pub fn utilization(&self) -> f64 {
        let head = self.head.load(Ordering::Relaxed);
        let tail = self.tail.load(Ordering::Relaxed);
        let used = (head - tail) as usize;
        used as f64 / self.capacity as f64
    }

    /// Get buffer capacity
    #[inline]
    pub const fn capacity(&self) -> usize {
        self.capacity
    }

    /// Get current length
    #[inline]
    pub fn len(&self) -> usize {
        let head = self.head.load(Ordering::Relaxed);
        let tail = self.tail.load(Ordering::Relaxed);
        (head - tail) as usize
    }

    /// Check if empty
    #[inline]
    pub fn is_empty(&self) -> bool {
        let head = self.head.load(Ordering::Relaxed);
        let tail = self.tail.load(Ordering::Relaxed);
        head == tail
    }

    /// Check if full
    #[inline]
    pub fn is_full(&self) -> bool {
        let head = self.head.load(Ordering::Relaxed);
        let tail = self.tail.load(Ordering::Relaxed);
        (head - tail) as usize >= self.capacity
    }

    /// Get batch processing mode
    #[inline]
    pub const fn batch_mode(&self) -> BatchMode {
        self.batch_mode
    }

    /// Switch to single-threaded mode for better performance
    pub fn set_single_threaded(&mut self) {
        self.batch_mode = BatchMode::SingleThreaded;
    }

    /// Switch to multi-threaded mode for safety
    pub fn set_multi_threaded(&mut self) {
        self.batch_mode = BatchMode::MultiThreaded;
    }
}

impl<T> Drop for SmallBatchRing<T> {
    fn drop(&mut self) {
        unsafe {
            dealloc(self.buffer.as_ptr().cast::<u8>(), self.layout);
        }
    }
}

impl<T> std::fmt::Debug for SmallBatchRing<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let head = self.head.load(Ordering::Relaxed);
        let tail = self.tail.load(Ordering::Relaxed);
        let len = (head - tail) as usize;
        f.debug_struct("SmallBatchRing")
            .field("capacity", &self.capacity)
            .field("head", &head)
            .field("tail", &tail)
            .field("len", &len)
            .field("batch_mode", &self.batch_mode)
            .finish()
    }
}

/// Cache-optimized structure-of-arrays layout for small batch orders
#[repr(align(64))]
/// SmallBatchOrdersSoA
/// 
/// TODO: Add detailed documentation for this struct
pub struct SmallBatchOrdersSoA {
    /// Order IDs (cache line 1)
    pub order_ids: [u64; 8],

    /// Prices (cache line 2)
    pub prices: [f64; 8],

    /// Quantities (cache line 3)
    pub quantities: [f64; 8],

    /// Timestamps (cache line 4)
    pub timestamps: [u64; 8],

    /// Sides and order types (packed into cache line 5)
    pub sides: [u8; 8], // 0 = Buy, 1 = Sell
    /// Order Types
    pub order_types: [u8; 8], // 0 = Market, 1 = Limit, etc.
    /// Symbols
    pub symbols: [u64; 6],    // Symbol hashes (remaining space)

    /// Batch size
    pub count: usize,
}

impl SmallBatchOrdersSoA {
    /// Create new empty structure-of-arrays
    #[must_use]
    pub const fn new() -> Self {
        Self {
            order_ids: [0; 8],
            prices: [0.0; 8],
            quantities: [0.0; 8],
            timestamps: [0; 8],
            sides: [0; 8],
            order_types: [0; 8],
            symbols: [0; 6],
            count: 0,
        }
    }

    /// Add order to structure-of-arrays layout
    #[inline(always)]
    pub fn add_order(
        &mut self,
        order_id: u64,
        symbol_hash: u64,
        side: u8,
        order_type: u8,
        quantity: f64,
        price: f64,
        timestamp: u64,
    ) -> bool {
        if self.count >= 8 {
            return false;
        }

        let idx = self.count;
        self.order_ids[idx] = order_id;
        self.prices[idx] = price;
        self.quantities[idx] = quantity;
        self.timestamps[idx] = timestamp;
        self.sides[idx] = side;
        self.order_types[idx] = order_type;

        if idx < 6 {
            self.symbols[idx] = symbol_hash;
        }

        self.count += 1;
        true
    }

    /// Clear all orders
    #[inline(always)]
    pub fn clear(&mut self) {
        self.count = 0;
        // No need to zero arrays for performance
    }

    /// Get SIMD-friendly price slice
    #[inline]
    #[must_use]
    pub fn prices_simd(&self) -> &[f64] {
        &self.prices[..self.count]
    }

    /// Get SIMD-friendly quantity slice
    #[inline]
    #[must_use]
    pub fn quantities_simd(&self) -> &[f64] {
        &self.quantities[..self.count]
    }

    /// Calculate total notional using SIMD if available
    #[cfg(target_arch = "x86_64")]
    #[must_use]
    pub fn calculate_total_notional_simd(&self) -> f64 {
        if self.count == 0 {
            return 0.0;
        }

        if std::arch::is_x86_feature_detected!("avx2") && self.count >= 4 {
            unsafe { self.calculate_total_notional_avx2() }
        } else {
            self.calculate_total_notional_scalar()
        }
    }

    /// AVX2 implementation for notional calculation
    #[cfg(target_arch = "x86_64")]
    #[target_feature(enable = "avx2")]
    unsafe fn calculate_total_notional_avx2(&self) -> f64 {
        use std::arch::x86_64::{
            _mm256_add_pd, _mm256_castpd256_pd128, _mm256_extractf128_pd, _mm256_hadd_pd,
            _mm256_loadu_pd, _mm256_mul_pd, _mm256_setzero_pd, _mm_add_pd, _mm_cvtsd_f64,
        };

        let mut sum_vec = _mm256_setzero_pd();
        let mut i = 0;

        // Process 4 orders at a time
        while i + 4 <= self.count {
            let prices_vec = _mm256_loadu_pd(&self.prices[i]);
            let quantities_vec = _mm256_loadu_pd(&self.quantities[i]);

            let notional_vec = _mm256_mul_pd(prices_vec, quantities_vec);
            sum_vec = _mm256_add_pd(sum_vec, notional_vec);

            i += 4;
        }

        // Sum vector components
        let sum_high_low = _mm256_hadd_pd(sum_vec, sum_vec);
        let sum_128 = _mm256_extractf128_pd(sum_high_low, 1);
        let sum_64 = _mm_add_pd(_mm256_castpd256_pd128(sum_high_low), sum_128);
        let mut total = _mm_cvtsd_f64(sum_64);

        // Add remaining scalar elements
        for j in i..self.count {
            total += self.prices[j] * self.quantities[j];
        }

        total
    }

    /// Scalar fallback for notional calculation
    #[must_use]
    pub fn calculate_total_notional_scalar(&self) -> f64 {
        self.prices[..self.count]
            .iter()
            .zip(&self.quantities[..self.count])
            .map(|(&price, &quantity)| price * quantity)
            .sum()
    }
}

impl Default for SmallBatchOrdersSoA {
    fn default() -> Self {
        Self::new()
    }
}

impl std::fmt::Debug for SmallBatchOrdersSoA {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SmallBatchOrdersSoA")
            .field("count", &self.count)
            .field("order_ids", &&self.order_ids[..self.count])
            .field("prices", &&self.prices[..self.count])
            .field("quantities", &&self.quantities[..self.count])
            .field("timestamps", &&self.timestamps[..self.count])
            .finish_non_exhaustive()
    }
}

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

    #[test]
    fn test_small_batch_ring_creation() {
        let ring = SmallBatchRing::<u64>::new(8, BatchMode::SingleThreaded)
            .expect("Failed to create ring");

        assert_eq!(ring.capacity(), 8);
        assert_eq!(ring.len(), 0);
        assert!(ring.is_empty());
        assert!(!ring.is_full());
        assert_eq!(ring.batch_mode(), BatchMode::SingleThreaded);
    }

    #[test]
    fn test_batch_operations() {
        let ring = SmallBatchRing::<u32>::new(16, BatchMode::SingleThreaded)
            .expect("Failed to create ring");

        // Test batch push
        let items = [1, 2, 3, 4, 5];
        let pushed = ring.push_batch(&items).expect("Failed to push batch");
        assert_eq!(pushed, 5);
        assert_eq!(ring.len(), 5);

        // Test batch pop
        let mut output = [0u32; 3];
        let popped = ring.pop_batch(&mut output);
        assert_eq!(popped, 3);
        assert_eq!(output, [1, 2, 3]);
        assert_eq!(ring.len(), 2);

        // Test remaining items
        let mut remaining = [0u32; 5];
        let remaining_count = ring.pop_batch(&mut remaining);
        assert_eq!(remaining_count, 2);
        assert_eq!(remaining[0], 4);
        assert_eq!(remaining[1], 5);
        assert!(ring.is_empty());
    }

    #[test]
    fn test_single_vs_multi_threaded_mode() {
        let mut ring =
            SmallBatchRing::<u64>::new(8, BatchMode::MultiThreaded).expect("Failed to create ring");

        assert_eq!(ring.batch_mode(), BatchMode::MultiThreaded);

        ring.set_single_threaded();
        assert_eq!(ring.batch_mode(), BatchMode::SingleThreaded);

        ring.set_multi_threaded();
        assert_eq!(ring.batch_mode(), BatchMode::MultiThreaded);
    }

    #[test]
    fn test_structure_of_arrays() {
        let mut soa = SmallBatchOrdersSoA::new();

        // Add some orders
        assert!(soa.add_order(1, 0x123, 0, 1, 1.5, 50000.0, 1000));
        assert!(soa.add_order(2, 0x456, 1, 0, 2.0, 3000.0, 2000));
        assert_eq!(soa.count, 2);

        // Test SIMD-friendly access
        let prices = soa.prices_simd();
        assert_eq!(prices.len(), 2);
        assert_eq!(prices[0], 50000.0);
        assert_eq!(prices[1], 3000.0);

        // Test notional calculation
        let total_notional = soa.calculate_total_notional_scalar();
        let expected = 1.5 * 50000.0 + 2.0 * 3000.0; // 75000 + 6000 = 81000
        assert!((total_notional - expected).abs() < 1e-6);
    }

    #[test]
    fn test_performance_characteristics() {
        let ring = SmallBatchRing::<u64>::new(1024, BatchMode::SingleThreaded)
            .expect("Failed to create ring");

        const NUM_BATCHES: usize = 1000;
        const BATCH_SIZE: usize = 8;

        let start = std::time::Instant::now();

        for batch_id in 0..NUM_BATCHES {
            // Create batch
            let mut items = [0u64; BATCH_SIZE];
            for i in 0..BATCH_SIZE {
                items[i] = (batch_id * BATCH_SIZE + i) as u64;
            }

            // Push batch
            ring.push_batch(&items).expect("Failed to push batch");

            // Pop batch
            let mut output = [0u64; BATCH_SIZE];
            let popped = ring.pop_batch(&mut output);
            assert_eq!(popped, BATCH_SIZE);
        }

        let duration = start.elapsed();
        let ops_per_sec = (NUM_BATCHES * BATCH_SIZE * 2) as f64 / duration.as_secs_f64();
        let avg_latency_ns = duration.as_nanos() / (NUM_BATCHES * 2) as u128;

        println!("Small batch ring performance:");
        println!("  Operations per second: {:.0}", ops_per_sec);
        println!("  Average latency per batch: {}ns", avg_latency_ns);

        // Should be significantly faster than regular lock-free operations
        assert!(
            avg_latency_ns < 500,
            "Latency too high: {}ns",
            avg_latency_ns
        );
    }
}