rustytorch/crates/training/rtx-preprocessing/src/loaders/memory_mapped.rs

//! Memory-mapped file loader with lazy loading and prefetching support
use crate::{PreprocessingError, Result};
use dashmap::DashMap;
use memmap2::{Mmap, MmapOptions};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::Instant;

use rtx_tensor::{Device, Tensor};

/// Configuration for memory-mapped file loading
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct LoaderConfig {
    /// Page size for memory mapping (in bytes)
    pub page_size: usize,
    /// Prefetching strategy to use
    pub prefetch_strategy: PrefetchStrategy,
    /// Whether to use shared memory segments
    pub use_shared_memory: bool,
    /// Whether to enable lazy loading with page faults
    pub enable_lazy_loading: bool,
}

impl Default for LoaderConfig {
    fn default() -> Self {
        Self {
            page_size: 4096,
            prefetch_strategy: PrefetchStrategy::None,
            use_shared_memory: false,
            enable_lazy_loading: false,
        }
    }
}

/// Prefetching strategies for optimizing memory access patterns
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum PrefetchStrategy {
    /// No prefetching
    None,
    /// Sequential prefetching with configurable window size
    Sequential { window_size: usize },
    /// Random access optimized prefetching
    Random { cache_size: usize },
    /// Adaptive prefetching based on access patterns
    Adaptive { max_window_size: usize },
}

/// Statistics for memory-mapped file loader performance
#[derive(Debug, Clone, Default)]
pub struct LoaderStatistics {
    pub total_reads: u64,
    pub cache_hits: u64,
    pub cache_misses: u64,
    pub bytes_read: u64,
    pub pages_loaded: usize,
}

/// Page information for lazy loading
#[derive(Debug, Clone)]
struct PageInfo {
    page_id: usize,
    offset: usize,
    size: usize,
    loaded_at: Option<Instant>,
    access_count: u64,
}

/// Global shared memory registry for cross-process sharing
static SHARED_MEMORY_REGISTRY: std::sync::LazyLock<DashMap<PathBuf, Arc<SharedMemorySegment>>> =
    std::sync::LazyLock::new(DashMap::new);

/// Shared memory segment that can be accessed across processes
#[derive(Debug)]
struct SharedMemorySegment {
    mmap: Mmap,
    file_size: usize,
    path: PathBuf,
    reference_count: Arc<Mutex<usize>>,
}

/// Memory-mapped file loader with advanced features
pub struct MemoryMappedFileLoader {
    config: LoaderConfig,
    mmap: Option<Mmap>,
    shared_segment: Option<Arc<SharedMemorySegment>>,
    file_path: Option<PathBuf>,
    file_size: usize,
    pages: RwLock<HashMap<usize, PageInfo>>,
    statistics: Arc<RwLock<LoaderStatistics>>,
    access_pattern_history: Arc<RwLock<Vec<(usize, Instant)>>>,
}

impl MemoryMappedFileLoader {
    /// Create a new memory-mapped file loader with the given configuration
    pub fn new(config: LoaderConfig) -> Self {
        Self {
            config,
            mmap: None,
            shared_segment: None,
            file_path: None,
            file_size: 0,
            pages: RwLock::new(HashMap::new()),
            statistics: Arc::new(RwLock::new(LoaderStatistics::default())),
            access_pattern_history: Arc::new(RwLock::new(Vec::new())),
        }
    }

    /// Get the loader configuration
    pub fn config(&self) -> &LoaderConfig {
        &self.config
    }

    /// Check if a file is currently loaded
    pub fn is_loaded(&self) -> bool {
        self.mmap.is_some() || self.shared_segment.is_some()
    }

    /// Open and memory-map a file
    pub fn open(&mut self, path: &Path) -> Result<()> {
        // Check if file exists
        if !path.exists() {
            return Err(PreprocessingError::invalid_input(format!(
                "File does not exist: {}",
                path.display()
            )));
        }

        let file = File::open(path)
            .map_err(|e| PreprocessingError::invalid_input(format!("Failed to open file: {e}")))?;

        let file_size = file
            .metadata()
            .map_err(|e| {
                PreprocessingError::invalid_input(format!("Failed to get file metadata: {e}"))
            })?
            .len() as usize;

        if self.config.use_shared_memory {
            // Check if this file is already in shared memory
            let path_buf = path.to_path_buf();
            if let Some(shared_segment) = SHARED_MEMORY_REGISTRY.get(&path_buf) {
                // Increment reference count
                let mut ref_count = shared_segment.reference_count.lock().unwrap();
                *ref_count += 1;
                self.shared_segment = Some(shared_segment.clone());
                self.file_size = shared_segment.file_size;
                self.file_path = Some(path_buf);
                return Ok(());
            }

            // Create new shared segment
            let mmap = unsafe {
                MmapOptions::new().map(&file).map_err(|e| {
                    PreprocessingError::invalid_input(format!("Failed to memory-map file: {e}"))
                })?
            };

            let shared_segment = Arc::new(SharedMemorySegment {
                mmap,
                file_size,
                path: path_buf.clone(),
                reference_count: Arc::new(Mutex::new(1)),
            });

            SHARED_MEMORY_REGISTRY.insert(path_buf.clone(), shared_segment.clone());
            self.shared_segment = Some(shared_segment);
            self.file_path = Some(path_buf);
        } else {
            // Create private memory mapping
            let mmap = unsafe {
                MmapOptions::new().map(&file).map_err(|e| {
                    PreprocessingError::invalid_input(format!("Failed to memory-map file: {e}"))
                })?
            };
            self.mmap = Some(mmap);
            self.file_path = Some(path.to_path_buf());
        }

        self.file_size = file_size;

        // Initialize lazy loading pages if enabled
        if self.config.enable_lazy_loading {
            self.initialize_pages();
        }

        Ok(())
    }

    /// Initialize page structure for lazy loading
    fn initialize_pages(&self) {
        if self.file_size == 0 {
            return;
        }

        let mut pages = self.pages.write();
        let page_count = self.file_size.div_ceil(self.config.page_size);

        for page_id in 0..page_count {
            let offset = page_id * self.config.page_size;
            let size = std::cmp::min(self.config.page_size, self.file_size - offset);

            pages.insert(
                page_id,
                PageInfo {
                    page_id,
                    offset,
                    size,
                    loaded_at: None,
                    access_count: 0,
                },
            );
        }
    }

    /// Get the size of the currently loaded file
    pub fn file_size(&self) -> usize {
        self.file_size
    }

    /// Get the number of loaded pages
    pub fn loaded_pages_count(&self) -> usize {
        let pages = self.pages.read();
        pages
            .values()
            .filter(|page| page.loaded_at.is_some())
            .count()
    }

    /// Read a chunk of data from the file
    pub fn read_chunk(&self, offset: usize, size: usize, device: &Device) -> Result<Tensor> {
        if !self.is_loaded() {
            return Err(PreprocessingError::invalid_input("File not loaded"));
        }

        if offset + size > self.file_size {
            return Err(PreprocessingError::invalid_input(format!(
                "Read beyond file size: {} + {} > {}",
                offset, size, self.file_size
            )));
        }

        // Update statistics
        {
            let mut stats = self.statistics.write();
            stats.total_reads += 1;
            stats.bytes_read += size as u64;
        }

        // Record access pattern for adaptive prefetching
        {
            let mut history = self.access_pattern_history.write();
            history.push((offset, Instant::now()));
            // Keep history bounded to prevent memory growth
            if history.len() > 1000 {
                history.remove(0);
            }
        }

        // Handle lazy loading if enabled
        if self.config.enable_lazy_loading {
            self.handle_lazy_loading(offset, size)?;
        }

        // Handle prefetching
        self.handle_prefetching(offset, size);

        // Get the actual data
        let data = self.get_raw_data(offset, size)?;

        // Convert to f32 and create tensor
        let float_data: Vec<f32> = data
            .chunks_exact(4)
            .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
            .collect();

        let tensor_size = float_data.len();
        Tensor::from_slice(&float_data, &[tensor_size], device).map_err(|e| {
            PreprocessingError::invalid_input(format!("Failed to create tensor: {e:?}"))
        })
    }

    /// Handle lazy loading for the requested range
    fn handle_lazy_loading(&self, offset: usize, size: usize) -> Result<()> {
        let start_page = offset / self.config.page_size;
        let end_page = (offset + size).div_ceil(self.config.page_size);

        let mut pages = self.pages.write();
        let now = Instant::now();

        for page_id in start_page..end_page {
            if let Some(page) = pages.get_mut(&page_id) {
                if page.loaded_at.is_none() {
                    // Simulate page fault and loading
                    page.loaded_at = Some(now);
                    page.access_count += 1;

                    // Update statistics
                    let mut stats = self.statistics.write();
                    stats.cache_misses += 1;
                    stats.pages_loaded += 1;
                } else {
                    page.access_count += 1;

                    // Update statistics
                    let mut stats = self.statistics.write();
                    stats.cache_hits += 1;
                }
            }
        }

        Ok(())
    }

    /// Handle prefetching based on the configured strategy
    fn handle_prefetching(&self, offset: usize, size: usize) {
        match &self.config.prefetch_strategy {
            PrefetchStrategy::None => {
                // No prefetching
            }
            PrefetchStrategy::Sequential { window_size } => {
                self.prefetch_sequential(offset, size, *window_size);
            }
            PrefetchStrategy::Random { cache_size } => {
                self.prefetch_random(offset, size, *cache_size);
            }
            PrefetchStrategy::Adaptive { max_window_size } => {
                self.prefetch_adaptive(offset, size, *max_window_size);
            }
        }
    }

    /// Sequential prefetching implementation
    fn prefetch_sequential(&self, offset: usize, _size: usize, window_size: usize) {
        if !self.config.enable_lazy_loading {
            return;
        }

        let start_page = offset / self.config.page_size;
        let prefetch_start = start_page + 1;
        let prefetch_end = std::cmp::min(
            prefetch_start + window_size,
            self.file_size.div_ceil(self.config.page_size),
        );

        let mut pages = self.pages.write();
        let now = Instant::now();

        for page_id in prefetch_start..prefetch_end {
            if let Some(page) = pages.get_mut(&page_id)
                && page.loaded_at.is_none()
            {
                page.loaded_at = Some(now);
            }
        }
    }

    /// Random access prefetching implementation
    fn prefetch_random(&self, _offset: usize, _size: usize, _cache_size: usize) {
        // Implement random access prefetching based on cache size
        // This would involve keeping frequently accessed pages in memory
    }

    /// Adaptive prefetching implementation
    fn prefetch_adaptive(&self, offset: usize, size: usize, max_window_size: usize) {
        // Analyze access pattern history to predict next accesses
        let history = self.access_pattern_history.read();
        if history.len() < 3 {
            // Fall back to sequential with small window
            self.prefetch_sequential(offset, size, 1);
            return;
        }

        // Simple pattern detection: check if accesses are sequential
        let recent_accesses: Vec<usize> = history
            .iter()
            .rev()
            .take(5)
            .map(|(offset, _)| *offset)
            .collect();

        let mut is_sequential = true;
        let mut stride = 0;

        if recent_accesses.len() >= 2 {
            stride = recent_accesses[0].saturating_sub(recent_accesses[1]);

            for i in 1..recent_accesses.len() - 1 {
                let current_stride = recent_accesses[i].saturating_sub(recent_accesses[i + 1]);
                if current_stride != stride {
                    is_sequential = false;
                    break;
                }
            }
        }

        if is_sequential && stride > 0 {
            // Predict next accesses based on detected stride
            let predicted_window =
                std::cmp::min(stride / self.config.page_size + 1, max_window_size);
            self.prefetch_sequential(offset, size, predicted_window);
        } else {
            // Fall back to small sequential window
            self.prefetch_sequential(offset, size, 1);
        }
    }

    /// Get raw data from the memory-mapped file
    fn get_raw_data(&self, offset: usize, size: usize) -> Result<&[u8]> {
        if let Some(ref mmap) = self.mmap {
            Ok(&mmap[offset..offset + size])
        } else if let Some(ref shared_segment) = self.shared_segment {
            Ok(&shared_segment.mmap[offset..offset + size])
        } else {
            Err(PreprocessingError::invalid_input(
                "No memory mapping available",
            ))
        }
    }

    /// Check if this loader shares memory with another loader
    pub fn shares_memory_with(&self, other: &Self) -> bool {
        match (&self.shared_segment, &other.shared_segment) {
            (Some(self_segment), Some(other_segment)) => Arc::ptr_eq(self_segment, other_segment),
            _ => false,
        }
    }

    /// Get performance statistics
    pub fn statistics(&self) -> LoaderStatistics {
        self.statistics.read().clone()
    }
}

impl Drop for MemoryMappedFileLoader {
    fn drop(&mut self) {
        if let Some(ref shared_segment) = self.shared_segment
            && let Some(path) = &self.file_path
        {
            let mut ref_count = shared_segment.reference_count.lock().unwrap();
            *ref_count -= 1;

            if *ref_count == 0 {
                SHARED_MEMORY_REGISTRY.remove(path);
            }
        }
    }
}

// Add Send and Sync bounds for thread safety
unsafe impl Send for MemoryMappedFileLoader {}
unsafe impl Sync for MemoryMappedFileLoader {}