memkit 0.2.0-beta.1

//! Huge page support for large allocations.
//!
//! Uses huge pages (2MB) to reduce TLB misses and improve performance
//! for large arenas.

#[cfg(target_os = "linux")]
use std::fs::OpenOptions;
#[cfg(target_os = "linux")]
use std::io::Write;
#[cfg(target_os = "linux")]
use std::os::fd::AsFd;
use std::ptr::NonNull;

/// Huge page allocator configuration.
#[derive(Debug, Clone, Copy)]
pub struct HugePageConfig {
    /// Enable huge page allocation.
    pub enabled: bool,
    /// Minimum size to use huge pages.
    pub min_size: usize,
    /// Prefer huge pages but fallback to normal if unavailable.
    pub fallback: bool,
}

impl Default for HugePageConfig {
    fn default() -> Self {
        Self {
            enabled: cfg!(target_os = "linux"), // Only enable on Linux by default
            min_size: 2 * 1024 * 1024, // 2MB
            fallback: true,
        }
    }
}

/// Allocate memory with huge pages if available.
pub fn alloc_huge(size: usize, config: HugePageConfig) -> Option<*mut u8> {
    if !config.enabled || size < config.min_size {
        return alloc_normal(size);
    }

    #[cfg(target_os = "linux")]
    {
        // Try MAP_HUGETLB on Linux
        alloc_linux_huge(size, config.fallback)
    }

    #[cfg(not(target_os = "linux"))]
    {
        // No huge page support on other platforms
        if config.fallback {
            alloc_normal(size)
        } else {
            None
        }
    }
}

/// Deallocate huge page memory.
pub unsafe fn dealloc_huge(ptr: *mut u8, size: usize, config: HugePageConfig) {
    if !config.enabled || size < config.min_size {
        dealloc_normal(ptr, size);
        return;
    }

    #[cfg(target_os = "linux")]
    {
        // Check if this was a huge page allocation
        // For simplicity, we always use normal dealloc
        // In a production implementation, we'd track which allocations used huge pages
        dealloc_normal(ptr, size);
    }

    #[cfg(not(target_os = "linux"))]
    {
        dealloc_normal(ptr, size);
    }
}

/// Normal allocation fallback.
fn alloc_normal(size: usize) -> Option<*mut u8> {
    use std::alloc::{alloc, Layout};
    
    // Align to cache line
    let aligned_size = (size + 63) & !63;
    let layout = Layout::from_size_align(aligned_size, 64).ok()?;
    
    let ptr = unsafe { alloc(layout) };
    NonNull::new(ptr).map(|p| p.as_ptr())
}

/// Normal deallocation.
unsafe fn dealloc_normal(ptr: *mut u8, size: usize) {
    use std::alloc::{dealloc, Layout};
    
    let aligned_size = (size + 63) & !63;
    let layout = Layout::from_size_align_unchecked(aligned_size, 64);
    dealloc(ptr, layout);
}

#[cfg(target_os = "linux")]
fn alloc_linux_huge(size: usize, fallback: bool) -> Option<*mut u8> {
    use std::alloc::Layout;
    use std::ptr;
    
    // Try to reserve huge pages first
    if let Ok(_) = reserve_huge_pages(size) {
        // Use mmap with MAP_HUGETLB
        let ptr = unsafe {
            libc::mmap(
                ptr::null_mut(),
                size,
                libc::PROT_READ | libc::PROT_WRITE,
                libc::MAP_PRIVATE | libc::MAP_ANONYMOUS | libc::MAP_HUGETLB,
                -1,
                0,
            )
        };
        
        if ptr != libc::MAP_FAILED {
            return Some(ptr as *mut u8);
        }
    }
    
    // Fallback to normal allocation
    if fallback {
        alloc_normal(size)
    } else {
        None
    }
}

#[cfg(target_os = "linux")]
fn reserve_huge_pages(size: usize) -> Result<(), std::io::Error> {
    use std::path::Path;
    
    let huge_pages_needed = (size + 2 * 1024 * 1024 - 1) / (2 * 1024 * 1024);
    
    // Try to reserve from /proc/sys/vm/nr_hugepages
    if Path::new("/proc/sys/vm/nr_hugepages").exists() {
        let mut file = OpenOptions::new()
            .write(true)
            .open("/proc/sys/vm/nr_hugepages")?;
            
        // Note: This requires root privileges in most configurations
        // For production use, this should be configured at system level
        writeln!(file, "{}", huge_pages_needed)?;
    }
    
    Ok(())
}

/// Check if huge pages are available on the system.
pub fn huge_pages_available() -> bool {
    #[cfg(target_os = "linux")]
    {
        std::path::Path::new("/proc/sys/vm/nr_hugepages").exists()
    }
    
    #[cfg(not(target_os = "linux"))]
    {
        false
    }
}

/// Get the huge page size on the system.
pub fn huge_page_size() -> Option<usize> {
    #[cfg(target_os = "linux")]
    {
        if let Ok(content) = std::fs::read_to_string("/proc/meminfo") {
            for line in content.lines() {
                if line.starts_with("Hugepagesize:") {
                    if let Some(kb_str) = line.split_whitespace().nth(1) {
                        if let Ok(kb) = kb_str.parse::<usize>() {
                            return Some(kb * 1024);
                        }
                    }
                }
            }
        }
        None
    }
    
    #[cfg(not(target_os = "linux"))]
    {
        None
    }
}

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

    #[test]
    fn test_huge_page_config() {
        let config = HugePageConfig::default();
        assert!(config.enabled == cfg!(target_os = "linux"));
        assert_eq!(config.min_size, 2 * 1024 * 1024);
        assert!(config.fallback);
    }

    #[test]
    fn test_normal_allocation() {
        let size = 1024;
        let ptr = alloc_normal(size);
        assert!(!ptr.unwrap().is_null());
        
        unsafe {
            dealloc_normal(ptr.unwrap(), size);
        }
    }

    #[test]
    fn test_huge_page_info() {
        // These should not panic
        let _available = huge_pages_available();
        let _size = huge_page_size();
    }

    #[test]
    fn test_huge_page_fallback() {
        let config = HugePageConfig {
            enabled: true,
            min_size: 2 * 1024 * 1024,
            fallback: true,
        };
        
        // Small allocation should use normal path
        let ptr = alloc_huge(1024, config);
        assert!(ptr.is_some());
        
        unsafe {
            dealloc_huge(ptr.unwrap(), 1024, config);
        }
        
        // Large allocation should try huge pages
        let ptr = alloc_huge(4 * 1024 * 1024, config);
        assert!(ptr.is_some());
        
        unsafe {
            dealloc_huge(ptr.unwrap(), 4 * 1024 * 1024, config);
        }
    }
}