jvmrs 0.1.2

//! Unified Memory Allocation Interface
//!
//! This module provides a unified interface for memory allocation across
//! different allocation strategies (TLAB, Arena, Direct) in JVMRS.
//! It eliminates code duplication and provides consistent allocation
//! behavior throughout the system.

use crate::error::MemoryError;

/// Memory type information for allocation
#[derive(Debug, Clone)]
pub enum MemoryType {
    /// Regular object allocation
    Object(String),
    /// Primitive array
    PrimitiveArray(PrimitiveArrayType),
    /// Object array
    ObjectArray(String),
    /// Multi-dimensional array
    MultiArray(String, Vec<i32>),
}

/// Primitive array types
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PrimitiveArrayType {
    Int, Byte, Short, Char, Long, Float, Double, Boolean,
}

/// Statistics for memory allocation
#[derive(Debug, Clone)]
pub struct MemoryStats {
    pub total_allocated: usize,
    pub allocation_count: u64,
    pub deallocation_count: u64,
    pub current_usage: usize,
}

/// Unified memory allocation trait
pub trait MemoryAllocator {
    /// Allocate a new object
    fn allocate_object(&mut self, class_name: &str) -> Result<u32, MemoryError>;
    
    /// Allocate an array
    fn allocate_array(&mut self, array_type: MemoryType, length: i32) -> Result<u32, MemoryError>;
    
    /// Get memory statistics
    fn get_stats(&self) -> MemoryStats;
    
    /// Check if allocation is possible
    fn can_allocate(&self, size: usize) -> bool;
    
    /// Cleanup unused memory
    fn cleanup(&mut self) -> Result<(), MemoryError>;
}

/// Arena-based allocator
pub struct ArenaAllocator {
    arenas: Vec<MemoryArena>,
    current_arena: usize,
    stats: MemoryStats,
}

/// Memory arena for batch allocations
pub struct MemoryArena {
    base: *mut u8,
    size: usize,
    used: usize,
    objects: Vec<u32>,
}

impl ArenaAllocator {
    /// Create a new arena allocator
    pub fn new(arena_size: usize) -> Self {
        // Create a dummy arena for now to avoid complex memory management
        let arenas = vec![
            MemoryArena {
                base: std::ptr::null_mut(), // Use null for now
                size: arena_size,
                used: 0,
                objects: vec![],
            }
        ];

        Self {
            arenas,
            current_arena: 0,
            stats: MemoryStats {
                total_allocated: arena_size,
                allocation_count: 0,
                deallocation_count: 0,
                current_usage: 0,
            },
        }
    }

    /// Add a new arena
    fn add_arena(&mut self, size: usize) -> Result<(), MemoryError> {
        let layout = std::alloc::Layout::from_size_align(size, 8)
            .map_err(|_| MemoryError::InvalidArrayLength(size))?;
        
        let base = unsafe { std::alloc::alloc(layout) };
        if base.is_null() {
            return Err(MemoryError::OutOfMemory);
        }

        self.arenas.push(MemoryArena {
            base,
            size,
            used: 0,
            objects: vec![],
        });

        self.stats.total_allocated += size;
        self.current_arena = self.arenas.len() - 1;
        Ok(())
    }
}

impl MemoryAllocator for ArenaAllocator {
    fn allocate_object(&mut self, class_name: &str) -> Result<u32, MemoryError> {
        let object_size = 16; // Base object size
        if !self.can_allocate(object_size) {
            self.cleanup()?;
        }

        // Calculate address before borrowing arena
        let base_address = self.arenas.len() * 1024;
        
        let arena = &mut self.arenas[self.current_arena];
        if arena.used + object_size > arena.size {
            // Current arena is full, try to create new one
            let new_size = arena.size * 2;
            self.add_arena(new_size)?;
            return self.allocate_object(class_name);
        }

        // For now, just return sequential addresses
        let address = (base_address + arena.used) as u32;
        arena.used += object_size;
        arena.objects.push(address);
        self.stats.allocation_count += 1;
        self.stats.current_usage += object_size;
        
        Ok(address)
    }

    fn allocate_array(&mut self, array_type: MemoryType, length: i32) -> Result<u32, MemoryError> {
        if length < 0 {
            return Err(MemoryError::InvalidArrayLength(length as usize));
        }

        let element_size = match array_type {
            MemoryType::PrimitiveArray(PrimitiveArrayType::Int) => 4,
            MemoryType::PrimitiveArray(PrimitiveArrayType::Byte) => 1,
            MemoryType::PrimitiveArray(PrimitiveArrayType::Short) => 2,
            MemoryType::PrimitiveArray(PrimitiveArrayType::Char) => 2,
            MemoryType::PrimitiveArray(PrimitiveArrayType::Long) => 8,
            MemoryType::PrimitiveArray(PrimitiveArrayType::Float) => 4,
            MemoryType::PrimitiveArray(PrimitiveArrayType::Double) => 8,
            MemoryType::PrimitiveArray(PrimitiveArrayType::Boolean) => 1,
            _ => return Err(MemoryError::InvalidArrayOperation("Unsupported array type".to_string())),
        };

        let total_size = element_size * length as usize + 16; // Array header

        if !self.can_allocate(total_size) {
            self.cleanup()?;
        }

        // Calculate address before borrowing arena
        let base_address = self.arenas.len() * 1024;
        
        let arena = &mut self.arenas[self.current_arena];
        if arena.used + total_size > arena.size {
            // Current arena is full, try to create new one
            let new_size = arena.size * 2;
            self.add_arena(new_size)?;
            return self.allocate_array(array_type, length);
        }

        // For now, just return sequential addresses
        let address = (base_address + arena.used) as u32;
        arena.used += total_size;
        arena.objects.push(address);
        self.stats.allocation_count += 1;
        self.stats.current_usage += total_size;
        
        Ok(address)
    }

    fn get_stats(&self) -> MemoryStats {
        self.stats.clone()
    }

    fn can_allocate(&self, size: usize) -> bool {
        if let Some(arena) = self.arenas.get(self.current_arena) {
            arena.used + size <= arena.size
        } else {
            false
        }
    }

    fn cleanup(&mut self) -> Result<(), MemoryError> {
        // Simple arena cleanup - just mark current arena for reuse
        if let Some(arena) = self.arenas.get_mut(self.current_arena) {
            arena.used = 0;
            arena.objects.clear();
            self.stats.deallocation_count += 1;
            self.stats.current_usage = 0;
        }
        Ok(())
    }
}

impl Drop for MemoryArena {
    fn drop(&mut self) {
        if !self.base.is_null() {
            unsafe {
                std::alloc::dealloc(
                    self.base,
                    std::alloc::Layout::from_size_align(self.size, 8).unwrap_or_else(|_| {
                        // Fallback layout if it fails
                        std::alloc::Layout::from_size_align_unchecked(self.size, 8)
                    }),
                );
            }
        }
    }
}

/// TLAB-based allocator
pub struct TlabAllocator {
    tlab_size: usize,
    tlab: Option<Tlab>,
    stats: MemoryStats,
}

impl TlabAllocator {
    pub fn new(tlab_size: usize) -> Self {
        Self {
            tlab_size,
            tlab: None,
            stats: MemoryStats {
                total_allocated: 0,
                allocation_count: 0,
                deallocation_count: 0,
                current_usage: 0,
            },
        }
    }

    fn get_or_create_tlab(&mut self) -> &mut Tlab {
        if self.tlab.is_none() {
            let layout = std::alloc::Layout::from_size_align(self.tlab_size, 8)
                .unwrap_or_else(|_| unsafe { std::alloc::Layout::from_size_align_unchecked(self.tlab_size, 8) });
            
            let base = unsafe { std::alloc::alloc(layout) };
            if !base.is_null() {
                self.tlab = Some(Tlab {
                    base,
                    size: self.tlab_size,
                    used: 0,
                });
                self.stats.total_allocated += self.tlab_size;
            }
        }
        self.tlab.as_mut().unwrap_or_else(|| {
            // Fallback if allocation failed - use panic as this should never happen in production
            panic!("Critical: Failed to allocate TLAB");
        })
    }
}

impl MemoryAllocator for TlabAllocator {
    fn allocate_object(&mut self, _class_name: &str) -> Result<u32, MemoryError> {
        let object_size = 16;
        
        loop {
            let tlab = self.get_or_create_tlab();
            
            if tlab.used + object_size <= tlab.size {
                // TLAB has enough space
                let address = tlab.base as u32 + tlab.used as u32;
                tlab.used += object_size;
                self.stats.allocation_count += 1;
                self.stats.current_usage += object_size;
                return Ok(address);
            }
            
            // TLAB is full, cleanup and try again
            self.cleanup()?;
        }
    }

    fn allocate_array(&mut self, _array_type: MemoryType, length: i32) -> Result<u32, MemoryError> {
        if length < 0 {
            return Err(MemoryError::InvalidArrayLength(length as usize));
        }

        let element_size = match _array_type {
            MemoryType::PrimitiveArray(PrimitiveArrayType::Int) => 4,
            MemoryType::PrimitiveArray(PrimitiveArrayType::Byte) => 1,
            MemoryType::PrimitiveArray(PrimitiveArrayType::Short) => 2,
            MemoryType::PrimitiveArray(PrimitiveArrayType::Char) => 2,
            MemoryType::PrimitiveArray(PrimitiveArrayType::Long) => 8,
            MemoryType::PrimitiveArray(PrimitiveArrayType::Float) => 4,
            MemoryType::PrimitiveArray(PrimitiveArrayType::Double) => 8,
            MemoryType::PrimitiveArray(PrimitiveArrayType::Boolean) => 1,
            _ => return Err(MemoryError::InvalidArrayOperation("Unsupported array type".to_string())),
        };

        let total_size = element_size * length as usize + 16;
        
        loop {
            let tlab = self.get_or_create_tlab();
            
            if tlab.used + total_size <= tlab.size {
                // TLAB has enough space
                let address = tlab.base as u32 + tlab.used as u32;
                tlab.used += total_size;
                self.stats.allocation_count += 1;
                self.stats.current_usage += total_size;
                return Ok(address);
            }
            
            // TLAB is full, cleanup and try again
            self.cleanup()?;
        }
    }

    fn get_stats(&self) -> MemoryStats {
        self.stats.clone()
    }

    fn can_allocate(&self, size: usize) -> bool {
        if let Some(tlab) = &self.tlab {
            tlab.used + size <= tlab.size
        } else {
            true // Can always allocate first TLAB
        }
    }

    fn cleanup(&mut self) -> Result<(), MemoryError> {
        if let Some(tlab) = &mut self.tlab {
            tlab.used = 0;
            self.stats.deallocation_count += 1;
            self.stats.current_usage = 0;
        }
        Ok(())
    }
}

/// Thread-Local Allocation Buffer (TLAB)
struct Tlab {
    base: *mut u8,
    size: usize,
    used: usize,
}

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

    #[test]
    fn test_arena_allocator() {
        let mut allocator = ArenaAllocator::new(1024);
        
        // Allocate object
        let obj1 = allocator.allocate_object("Test").unwrap();
        let obj2 = allocator.allocate_object("Test").unwrap();
        
        assert_ne!(obj1, obj2);
        assert_eq!(allocator.get_stats().allocation_count, 2);
    }

    #[test]
    fn test_tlab_allocator() {
        let mut allocator = TlabAllocator::new(512);
        
        // Allocate object
        let obj1 = allocator.allocate_object("Test").unwrap();
        let obj2 = allocator.allocate_object("Test").unwrap();
        
        assert_ne!(obj1, obj2);
        assert_eq!(allocator.get_stats().allocation_count, 2);
    }

    #[test]
    fn test_array_allocation() {
        let mut allocator = ArenaAllocator::new(1024);
        
        // Allocate int array
        let arr = allocator.allocate_array(
            MemoryType::PrimitiveArray(PrimitiveArrayType::Int),
            10
        ).unwrap();
        
        assert_ne!(arr, 0);
        assert_eq!(allocator.get_stats().allocation_count, 1);
    }
}