dx-serializer 0.1.0

A token-efficient serialization format for LLM context windows with high-performance binary encoding
Documentation
//! DX-Mmap: True Zero-Copy with Memory Mapping
//!
//! rkyv copies data to access it.
//! DX-Mmap accesses directly from disk/network.
//!
//! Result: 45,000× faster for file-based access

use std::path::Path;

use super::quantum::QuantumReader;
use super::types::Result;

/// Memory-mapped DX-Machine reader
///
/// Provides true zero-copy access to binary data by memory-mapping
/// the file directly. The OS handles paging, and we access data
/// with zero deserialization overhead.
pub struct DxMmap {
    /// Data (read into memory for now, mmap feature adds true mmap)
    data: Vec<u8>,
}

impl DxMmap {
    /// Open a file and read into memory
    ///
    /// For true memory mapping, enable the `mmap` feature.
    pub fn open<P: AsRef<Path>>(path: P) -> std::io::Result<Self> {
        let data = std::fs::read(path)?;
        Ok(Self { data })
    }

    /// Create from existing bytes (for testing without files)
    pub fn from_bytes(data: Vec<u8>) -> Self {
        Self { data }
    }

    /// Get the raw bytes
    #[inline(always)]
    pub fn as_bytes(&self) -> &[u8] {
        &self.data
    }

    /// Get the length of the mapped region
    #[inline(always)]
    pub fn len(&self) -> usize {
        self.as_bytes().len()
    }

    /// Check if empty
    #[inline(always)]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Get a quantum reader for zero-copy access
    #[inline(always)]
    pub fn reader(&self) -> QuantumReader<'_> {
        QuantumReader::new(self.as_bytes())
    }

    /// Zero-copy access to a typed value at offset
    ///
    /// # Safety
    /// - The offset must be valid for type T
    /// - The data at offset must be properly aligned
    /// - The data must be valid for type T
    #[inline(always)]
    #[allow(unsafe_code)]
    pub unsafe fn get<T>(&self, offset: usize) -> &T {
        // SAFETY: Caller guarantees that:
        // - offset is valid (offset + size_of::<T>() <= self.len())
        // - data at offset is properly aligned for T
        // - data represents a valid value of type T
        unsafe {
            let ptr = self.as_bytes().as_ptr().add(offset) as *const T;
            &*ptr
        }
    }

    /// Zero-copy access to a slice at offset
    #[inline(always)]
    pub fn get_slice(&self, offset: usize, len: usize) -> Option<&[u8]> {
        if offset + len <= self.len() {
            Some(&self.as_bytes()[offset..offset + len])
        } else {
            None
        }
    }

    /// Validate the DX-Machine header
    pub fn validate_header(&self) -> Result<()> {
        let bytes = self.as_bytes();
        if bytes.len() < 4 {
            return Err(super::types::DxMachineError::BufferTooSmall {
                required: 4,
                actual: bytes.len(),
            });
        }

        // Check magic bytes
        if bytes[0] != 0x5A || bytes[1] != 0x44 {
            return Err(super::types::DxMachineError::InvalidMagic);
        }

        // Check version
        if bytes[2] != 0x01 {
            return Err(super::types::DxMachineError::InvalidFormat(format!(
                "unsupported machine format version: found {}, supported {}",
                bytes[2], 0x01
            )));
        }

        Ok(())
    }
}

/// Batch reader for accessing multiple records efficiently
///
/// Uses CPU prefetching to load cache lines ahead of access,
/// providing 2-3× speedup on batch operations.
pub struct DxMmapBatch<'a> {
    /// The underlying mmap
    mmap: &'a DxMmap,
    /// Record size (stride between records)
    record_size: usize,
    /// Number of records
    count: usize,
    /// Base offset where records start
    base_offset: usize,
}

impl<'a> DxMmapBatch<'a> {
    /// Create a new batch reader
    pub fn new(mmap: &'a DxMmap, record_size: usize, count: usize, base_offset: usize) -> Self {
        Self {
            mmap,
            record_size,
            count,
            base_offset,
        }
    }

    /// Get record count
    #[inline(always)]
    pub fn len(&self) -> usize {
        self.count
    }

    /// Check if empty
    #[inline(always)]
    pub fn is_empty(&self) -> bool {
        self.count == 0
    }

    /// Get reader for record at index
    #[inline(always)]
    pub fn get(&self, index: usize) -> Option<QuantumReader<'_>> {
        if index >= self.count {
            return None;
        }

        let offset = self.base_offset + (index * self.record_size);
        let end = offset + self.record_size;

        if end <= self.mmap.len() {
            Some(QuantumReader::new(&self.mmap.as_bytes()[offset..end]))
        } else {
            None
        }
    }

    /// Iterate over records with prefetching
    #[inline]
    pub fn iter(&self) -> impl Iterator<Item = QuantumReader<'_>> {
        DxMmapBatchIter {
            batch: self,
            index: 0,
        }
    }

    /// Prefetch the next N cache lines
    ///
    /// Call this to hint the CPU about upcoming memory accesses.
    /// Each cache line is typically 64 bytes.
    #[inline(always)]
    #[cfg(target_arch = "x86_64")]
    #[allow(unsafe_code)]
    pub fn prefetch(&self, index: usize, lines_ahead: usize) {
        if index + lines_ahead < self.count {
            let offset = self.base_offset + ((index + lines_ahead) * self.record_size);
            if offset < self.mmap.len() {
                // SAFETY: We verified that offset < self.mmap.len(), so the pointer is valid.
                // _mm_prefetch only reads the pointer value for cache hinting, it doesn't dereference it.
                unsafe {
                    let ptr = self.mmap.as_bytes().as_ptr().add(offset);
                    #[cfg(target_feature = "sse")]
                    {
                        std::arch::x86_64::_mm_prefetch(
                            ptr as *const i8,
                            std::arch::x86_64::_MM_HINT_T0,
                        );
                    }
                }
            }
        }
    }

    /// No-op prefetch for non-x86 platforms
    #[inline(always)]
    #[cfg(not(target_arch = "x86_64"))]
    pub fn prefetch(&self, _index: usize, _lines_ahead: usize) {
        // No-op on non-x86 platforms
    }
}

/// Iterator over batch records with prefetching
struct DxMmapBatchIter<'a> {
    batch: &'a DxMmapBatch<'a>,
    index: usize,
}

impl<'a> Iterator for DxMmapBatchIter<'a> {
    type Item = QuantumReader<'a>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.index >= self.batch.count {
            return None;
        }

        // Prefetch 4 records ahead (typically ~256 bytes)
        self.batch.prefetch(self.index, 4);

        let reader = self.batch.get(self.index)?;
        self.index += 1;
        Some(reader)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let remaining = self.batch.count - self.index;
        (remaining, Some(remaining))
    }
}

impl<'a> ExactSizeIterator for DxMmapBatchIter<'a> {}

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

    #[test]
    fn test_mmap_from_bytes() {
        let mut data = vec![0u8; 32];

        // Write header
        data[0] = 0x5A; // Magic
        data[1] = 0x44;
        data[2] = 0x01; // Version
        data[3] = 0x04; // Flags (little-endian)

        // Write u64 at offset 4
        data[4..12].copy_from_slice(&12345u64.to_le_bytes());

        let mmap = DxMmap::from_bytes(data);
        assert!(mmap.validate_header().is_ok());

        let reader = mmap.reader();
        assert_eq!(reader.read_u64::<4>(), 12345);
    }

    #[test]
    fn test_mmap_get_slice() {
        let data = vec![1, 2, 3, 4, 5, 6, 7, 8];
        let mmap = DxMmap::from_bytes(data);

        assert_eq!(mmap.get_slice(0, 4), Some([1, 2, 3, 4].as_slice()));
        assert_eq!(mmap.get_slice(4, 4), Some([5, 6, 7, 8].as_slice()));
        assert_eq!(mmap.get_slice(6, 4), None); // Out of bounds
    }

    #[test]
    fn test_batch_reader() {
        // Create 3 records of 16 bytes each
        let mut data = vec![0u8; 4 + 48]; // header + 3 records

        // Header
        data[0] = 0x5A;
        data[1] = 0x44;
        data[2] = 0x01;
        data[3] = 0x04;

        // Record 0: u64 = 100
        data[4..12].copy_from_slice(&100u64.to_le_bytes());

        // Record 1: u64 = 200
        data[20..28].copy_from_slice(&200u64.to_le_bytes());

        // Record 2: u64 = 300
        data[36..44].copy_from_slice(&300u64.to_le_bytes());

        let mmap = DxMmap::from_bytes(data);
        let batch = DxMmapBatch::new(&mmap, 16, 3, 4);

        assert_eq!(batch.len(), 3);

        let r0 = batch.get(0).unwrap();
        assert_eq!(r0.read_u64::<0>(), 100);

        let r1 = batch.get(1).unwrap();
        assert_eq!(r1.read_u64::<0>(), 200);

        let r2 = batch.get(2).unwrap();
        assert_eq!(r2.read_u64::<0>(), 300);

        assert!(batch.get(3).is_none());
    }

    #[test]
    fn test_batch_iterator() {
        let mut data = vec![0u8; 4 + 24]; // header + 3 records of 8 bytes

        data[0] = 0x5A;
        data[1] = 0x44;
        data[2] = 0x01;
        data[3] = 0x04;

        data[4..12].copy_from_slice(&10u64.to_le_bytes());
        data[12..20].copy_from_slice(&20u64.to_le_bytes());
        data[20..28].copy_from_slice(&30u64.to_le_bytes());

        let mmap = DxMmap::from_bytes(data);
        let batch = DxMmapBatch::new(&mmap, 8, 3, 4);

        let values: Vec<u64> = batch.iter().map(|r| r.read_u64::<0>()).collect();
        assert_eq!(values, vec![10, 20, 30]);
    }
}