optirs_gpu/memory/allocation/

buddy_allocator.rs

// Buddy system allocator for GPU memory management
//
// This module implements a buddy system memory allocator that maintains
// power-of-2 sized blocks in a binary tree structure for efficient
// allocation and deallocation with minimal fragmentation.

use std::collections::{HashMap, VecDeque};
use std::sync::{Arc, Mutex};
use std::time::Instant;

/// Buddy system allocator implementation
pub struct BuddyAllocator {
    /// Base address of the memory pool
    base_ptr: *mut u8,
    /// Total size of the memory pool (must be power of 2)
    total_size: usize,
    /// Minimum block size (must be power of 2)
    min_block_size: usize,
    /// Maximum order (log2 of total_size / min_block_size)
    max_order: usize,
    /// Free lists for each order
    free_lists: Vec<VecDeque<BuddyBlock>>,
    /// Allocation tracking for debugging
    allocated_blocks: HashMap<*mut u8, BuddyBlock>,
    /// Statistics
    stats: BuddyStats,
    /// Configuration
    config: BuddyConfig,
}

/// Buddy block representation
#[derive(Debug, Clone)]
pub struct BuddyBlock {
    /// Block address
    pub ptr: *mut u8,
    /// Block size (always power of 2)
    pub size: usize,
    /// Block order (log2 of size / min_block_size)
    pub order: usize,
    /// Whether block is allocated
    pub is_allocated: bool,
    /// Allocation timestamp
    pub allocated_at: Option<Instant>,
    /// Last access timestamp
    pub last_accessed: Option<Instant>,
    /// Access count
    pub access_count: u64,
}

impl BuddyBlock {
    pub fn new(ptr: *mut u8, size: usize, order: usize) -> Self {
        Self {
            ptr,
            size,
            order,
            is_allocated: false,
            allocated_at: None,
            last_accessed: None,
            access_count: 0,
        }
    }

    pub fn allocate(&mut self) {
        self.is_allocated = true;
        self.allocated_at = Some(Instant::now());
        self.access_count += 1;
    }

    pub fn deallocate(&mut self) {
        self.is_allocated = false;
        self.allocated_at = None;
    }

    pub fn access(&mut self) {
        self.last_accessed = Some(Instant::now());
        self.access_count += 1;
    }

    /// Get buddy address for this block.
    ///
    /// The classic buddy-system XOR trick (`offset ^ size`) only identifies
    /// the true buddy when `offset` is measured relative to a base every
    /// block shares. `base_ptr` (the allocator's arena base, e.g.
    /// `BuddyAllocator::base_ptr`) must be that shared base: XOR-ing the
    /// raw absolute pointer would only coincidentally find the real buddy,
    /// since a real heap allocation's address is not generally a multiple
    /// of the arena's total size.
    pub fn get_buddy_address(&self, base_ptr: *mut u8) -> *mut u8 {
        let relative_offset = (self.ptr as usize).wrapping_sub(base_ptr as usize);
        let buddy_relative = relative_offset ^ self.size;
        (base_ptr as usize).wrapping_add(buddy_relative) as *mut u8
    }

    /// Check if two blocks are buddies, relative to the allocator's
    /// `base_ptr` (see [`Self::get_buddy_address`]).
    pub fn is_buddy_of(&self, other: &BuddyBlock, base_ptr: *mut u8) -> bool {
        if self.size != other.size {
            return false;
        }

        other.ptr == self.get_buddy_address(base_ptr)
    }
}

/// Buddy allocator statistics
#[derive(Debug, Clone, Default)]
pub struct BuddyStats {
    pub total_allocations: u64,
    pub total_deallocations: u64,
    pub successful_allocations: u64,
    pub failed_allocations: u64,
    pub split_operations: u64,
    pub merge_operations: u64,
    pub fragmentation_ratio: f64,
    pub average_allocation_time_ns: f64,
    pub peak_allocated_blocks: usize,
    pub current_allocated_blocks: usize,
    pub internal_fragmentation: f64,
    pub external_fragmentation: f64,
}

impl BuddyStats {
    pub fn record_allocation(
        &mut self,
        success: bool,
        time_ns: u64,
        size_requested: usize,
        size_allocated: usize,
    ) {
        self.total_allocations += 1;

        if success {
            self.successful_allocations += 1;
            self.current_allocated_blocks += 1;

            if self.current_allocated_blocks > self.peak_allocated_blocks {
                self.peak_allocated_blocks = self.current_allocated_blocks;
            }

            // Update average allocation time
            let total_time = self.average_allocation_time_ns
                * (self.successful_allocations - 1) as f64
                + time_ns as f64;
            self.average_allocation_time_ns = total_time / self.successful_allocations as f64;

            // Update internal fragmentation
            if size_allocated > 0 {
                let waste = size_allocated - size_requested;
                let new_frag = waste as f64 / size_allocated as f64;
                self.internal_fragmentation = (self.internal_fragmentation
                    * (self.successful_allocations - 1) as f64
                    + new_frag)
                    / self.successful_allocations as f64;
            }
        } else {
            self.failed_allocations += 1;
        }
    }

    pub fn record_deallocation(&mut self) {
        self.total_deallocations += 1;
        self.current_allocated_blocks = self.current_allocated_blocks.saturating_sub(1);
    }

    pub fn record_split(&mut self) {
        self.split_operations += 1;
    }

    pub fn record_merge(&mut self) {
        self.merge_operations += 1;
    }

    pub fn get_success_rate(&self) -> f64 {
        if self.total_allocations == 0 {
            0.0
        } else {
            self.successful_allocations as f64 / self.total_allocations as f64
        }
    }

    pub fn get_fragmentation_ratio(&self) -> f64 {
        self.fragmentation_ratio
    }
}

/// Buddy allocator configuration
#[derive(Debug, Clone)]
pub struct BuddyConfig {
    /// Enable coalescing of free blocks
    pub enable_coalescing: bool,
    /// Enable split optimization
    pub enable_split_optimization: bool,
    /// Minimum block size (must be power of 2)
    pub min_block_size: usize,
    /// Maximum allocation size
    pub max_allocation_size: usize,
    /// Enable allocation tracking
    pub enable_tracking: bool,
    /// Enable access pattern analysis
    pub enable_access_analysis: bool,
    /// Defragmentation threshold
    pub defrag_threshold: f64,
    /// Enable automatic defragmentation
    pub auto_defrag: bool,
}

impl Default for BuddyConfig {
    fn default() -> Self {
        Self {
            enable_coalescing: true,
            enable_split_optimization: true,
            min_block_size: 256,
            max_allocation_size: 1024 * 1024 * 1024, // 1GB
            enable_tracking: true,
            enable_access_analysis: false,
            defrag_threshold: 0.5,
            auto_defrag: true,
        }
    }
}

impl BuddyAllocator {
    /// Create a new buddy allocator
    pub fn new(
        base_ptr: *mut u8,
        total_size: usize,
        config: BuddyConfig,
    ) -> Result<Self, BuddyError> {
        // Validate that total_size is power of 2
        if !total_size.is_power_of_two() {
            return Err(BuddyError::InvalidSize(format!(
                "Total size {} is not a power of 2",
                total_size
            )));
        }

        // Validate that min_block_size is power of 2
        if !config.min_block_size.is_power_of_two() {
            return Err(BuddyError::InvalidSize(format!(
                "Minimum block size {} is not a power of 2",
                config.min_block_size
            )));
        }

        // Validate size relationships
        if config.min_block_size > total_size {
            return Err(BuddyError::InvalidSize(format!(
                "Minimum block size {} exceeds total size {}",
                config.min_block_size, total_size
            )));
        }

        let max_order = (total_size / config.min_block_size).trailing_zeros() as usize;
        let mut free_lists = vec![VecDeque::new(); max_order + 1];

        // Initialize with one large free block
        let initial_block = BuddyBlock::new(base_ptr, total_size, max_order);
        free_lists[max_order].push_back(initial_block);

        Ok(Self {
            base_ptr,
            total_size,
            min_block_size: config.min_block_size,
            max_order,
            free_lists,
            allocated_blocks: HashMap::new(),
            stats: BuddyStats::default(),
            config,
        })
    }

    /// Allocate memory of specified size
    pub fn allocate(&mut self, size: usize) -> Result<*mut u8, BuddyError> {
        let start_time = Instant::now();

        if size == 0 {
            self.stats.record_allocation(false, 0, size, 0);
            return Err(BuddyError::InvalidSize(
                "Cannot allocate zero bytes".to_string(),
            ));
        }

        if size > self.config.max_allocation_size {
            self.stats.record_allocation(false, 0, size, 0);
            return Err(BuddyError::InvalidSize(format!(
                "Allocation size {} exceeds maximum {}",
                size, self.config.max_allocation_size
            )));
        }

        // Calculate required order (round up to next power of 2)
        let required_size = size.max(self.min_block_size).next_power_of_two();
        let required_order = (required_size / self.min_block_size).trailing_zeros() as usize;

        if required_order > self.max_order {
            let elapsed = start_time.elapsed().as_nanos() as u64;
            self.stats.record_allocation(false, elapsed, size, 0);
            return Err(BuddyError::OutOfMemory(format!(
                "Required order {} exceeds maximum order {}",
                required_order, self.max_order
            )));
        }

        // Find available block
        match self.find_free_block(required_order) {
            Some(mut block) => {
                block.allocate();
                let ptr = block.ptr;

                if self.config.enable_tracking {
                    self.allocated_blocks.insert(ptr, block);
                }

                let elapsed = start_time.elapsed().as_nanos() as u64;
                self.stats
                    .record_allocation(true, elapsed, size, required_size);

                Ok(ptr)
            }
            None => {
                let elapsed = start_time.elapsed().as_nanos() as u64;
                self.stats.record_allocation(false, elapsed, size, 0);
                Err(BuddyError::OutOfMemory(
                    "No suitable block available".to_string(),
                ))
            }
        }
    }

    /// Deallocate memory at specified pointer
    pub fn deallocate(&mut self, ptr: *mut u8) -> Result<(), BuddyError> {
        if ptr.is_null() {
            return Err(BuddyError::InvalidPointer(
                "Cannot deallocate null pointer".to_string(),
            ));
        }

        // Remove from allocated blocks
        let block = if self.config.enable_tracking {
            self.allocated_blocks.remove(&ptr).ok_or_else(|| {
                BuddyError::InvalidPointer("Pointer not found in allocated blocks".to_string())
            })?
        } else {
            // If tracking is disabled, we need to reconstruct block info
            // This is less safe but more performance-oriented
            return Err(BuddyError::InvalidPointer(
                "Cannot deallocate without tracking enabled".to_string(),
            ));
        };

        self.stats.record_deallocation();

        // Add back to free list with coalescing
        if self.config.enable_coalescing {
            self.free_with_coalescing(block);
        } else {
            self.free_lists[block.order].push_back(block);
        }

        // Trigger automatic defragmentation if needed
        if self.config.auto_defrag {
            let fragmentation = self.calculate_fragmentation();
            if fragmentation > self.config.defrag_threshold {
                self.defragment();
            }
        }

        Ok(())
    }

    /// Find a free block of at least the specified order
    fn find_free_block(&mut self, min_order: usize) -> Option<BuddyBlock> {
        // Look for exact fit first
        if let Some(exact) = self.free_lists[min_order].pop_front() {
            return Some(exact);
        }

        // Look for larger blocks and split them
        for order in (min_order + 1)..=self.max_order {
            if let Some(large_block) = self.free_lists[order].pop_front() {
                return Some(self.split_block(large_block, min_order));
            }
        }

        None
    }

    /// Split a block down to the target order
    fn split_block(&mut self, mut block: BuddyBlock, target_order: usize) -> BuddyBlock {
        while block.order > target_order {
            self.stats.record_split();

            // Create buddy block
            let buddy_size = block.size / 2;
            let buddy_order = block.order - 1;
            let buddy_ptr = unsafe { block.ptr.add(buddy_size) };

            let buddy_block = BuddyBlock::new(buddy_ptr, buddy_size, buddy_order);

            // Update original block
            block.size = buddy_size;
            block.order = buddy_order;

            // Add buddy to free list
            self.free_lists[buddy_order].push_back(buddy_block);
        }

        block
    }

    /// Free a block with coalescing
    fn free_with_coalescing(&mut self, block: BuddyBlock) {
        let mut current_block = block;
        current_block.deallocate();

        // Try to coalesce with buddy blocks
        while current_block.order < self.max_order {
            let buddy_addr = current_block.get_buddy_address(self.base_ptr);

            // Look for buddy in the same order free list.
            let Some(pos) = self.free_lists[current_block.order]
                .iter()
                .position(|b| b.ptr == buddy_addr)
            else {
                // No buddy found, stop coalescing.
                break;
            };

            // `pos` was just found in this exact list, so removal is
            // guaranteed to succeed; the `else` bails out defensively
            // (stopping coalescing for this block) rather than panicking
            // if that invariant were ever violated.
            let Some(buddy) = self.free_lists[current_block.order].remove(pos) else {
                break;
            };
            self.stats.record_merge();

            // Create coalesced block
            let coalesced_ptr = if current_block.ptr < buddy.ptr {
                current_block.ptr
            } else {
                buddy.ptr
            };

            current_block = BuddyBlock::new(
                coalesced_ptr,
                current_block.size * 2,
                current_block.order + 1,
            );
        }

        // Add final block to appropriate free list
        self.free_lists[current_block.order].push_back(current_block);
    }

    /// Calculate current fragmentation level
    pub fn calculate_fragmentation(&self) -> f64 {
        let mut total_free_space = 0;
        let mut largest_free_block = 0;

        for (order, blocks) in self.free_lists.iter().enumerate() {
            let block_size = self.min_block_size * (1 << order);
            let free_space = blocks.len() * block_size;
            total_free_space += free_space;

            if !blocks.is_empty() && block_size > largest_free_block {
                largest_free_block = block_size;
            }
        }

        if total_free_space == 0 {
            0.0
        } else {
            1.0 - (largest_free_block as f64 / total_free_space as f64)
        }
    }

    /// Perform defragmentation by coalescing free blocks
    pub fn defragment(&mut self) -> usize {
        let mut coalesced_blocks = 0;

        // Go through each order and try to coalesce adjacent blocks
        for order in 0..self.max_order {
            let mut blocks_to_process: Vec<BuddyBlock> = self.free_lists[order].drain(..).collect();
            let mut processed = Vec::new();

            while !blocks_to_process.is_empty() {
                let current = blocks_to_process.remove(0);
                let buddy_addr = current.get_buddy_address(self.base_ptr);

                // Look for buddy in remaining blocks
                if let Some(buddy_pos) = blocks_to_process.iter().position(|b| b.ptr == buddy_addr)
                {
                    let buddy = blocks_to_process.remove(buddy_pos);
                    coalesced_blocks += 1;
                    self.stats.record_merge();

                    // Create coalesced block and add to next order
                    let coalesced_ptr = if current.ptr < buddy.ptr {
                        current.ptr
                    } else {
                        buddy.ptr
                    };

                    let coalesced_block =
                        BuddyBlock::new(coalesced_ptr, current.size * 2, current.order + 1);

                    self.free_lists[order + 1].push_back(coalesced_block);
                } else {
                    processed.push(current);
                }
            }

            // Put back uncoalesced blocks
            self.free_lists[order].extend(processed);
        }

        coalesced_blocks
    }

    /// Get allocator statistics
    pub fn get_stats(&self) -> &BuddyStats {
        &self.stats
    }

    /// Get current memory usage
    pub fn get_memory_usage(&self) -> MemoryUsage {
        let mut total_allocated = 0;
        let mut total_free = 0;

        // Calculate allocated memory
        for block in self.allocated_blocks.values() {
            total_allocated += block.size;
        }

        // Calculate free memory
        for (order, blocks) in self.free_lists.iter().enumerate() {
            let block_size = self.min_block_size * (1 << order);
            total_free += blocks.len() * block_size;
        }

        MemoryUsage {
            total_size: self.total_size,
            allocated_size: total_allocated,
            free_size: total_free,
            fragmentation_ratio: self.calculate_fragmentation(),
            allocated_blocks: self.allocated_blocks.len(),
            free_blocks: self.free_lists.iter().map(|l| l.len()).sum(),
        }
    }

    /// Get detailed statistics about free blocks
    pub fn get_free_block_stats(&self) -> Vec<FreeBlockStats> {
        self.free_lists
            .iter()
            .enumerate()
            .map(|(order, blocks)| {
                let block_size = self.min_block_size * (1 << order);
                FreeBlockStats {
                    order,
                    block_size,
                    block_count: blocks.len(),
                    total_size: blocks.len() * block_size,
                }
            })
            .collect()
    }

    /// Check allocator consistency
    pub fn validate_consistency(&self) -> Result<(), BuddyError> {
        let mut total_free = 0;
        let mut total_allocated = 0;

        // Check free blocks
        for (order, blocks) in self.free_lists.iter().enumerate() {
            let expected_size = self.min_block_size * (1 << order);

            for block in blocks {
                if block.size != expected_size {
                    return Err(BuddyError::CorruptedState(format!(
                        "Free block at order {} has incorrect size: expected {}, got {}",
                        order, expected_size, block.size
                    )));
                }

                if block.is_allocated {
                    return Err(BuddyError::CorruptedState(
                        "Free block marked as allocated".to_string(),
                    ));
                }

                total_free += block.size;
            }
        }

        // Check allocated blocks
        for block in self.allocated_blocks.values() {
            if !block.is_allocated {
                return Err(BuddyError::CorruptedState(
                    "Allocated block marked as free".to_string(),
                ));
            }

            total_allocated += block.size;
        }

        // Check total memory accounting
        if total_free + total_allocated != self.total_size {
            return Err(BuddyError::CorruptedState(format!(
                "Memory accounting error: total_free ({}) + total_allocated ({}) != total_size ({})",
                total_free, total_allocated, self.total_size
            )));
        }

        Ok(())
    }

    /// Reset allocator to initial state
    pub fn reset(&mut self) {
        self.free_lists = vec![VecDeque::new(); self.max_order + 1];
        self.allocated_blocks.clear();
        self.stats = BuddyStats::default();

        // Add initial large block
        let initial_block = BuddyBlock::new(self.base_ptr, self.total_size, self.max_order);
        self.free_lists[self.max_order].push_back(initial_block);
    }

    /// Access a previously allocated block (for statistics)
    pub fn access_block(&mut self, ptr: *mut u8) -> Result<(), BuddyError> {
        if self.config.enable_access_analysis {
            if let Some(block) = self.allocated_blocks.get_mut(&ptr) {
                block.access();
                Ok(())
            } else {
                Err(BuddyError::InvalidPointer("Block not found".to_string()))
            }
        } else {
            Ok(()) // Silently succeed if access analysis is disabled
        }
    }

    /// Get allocation info for a pointer
    pub fn get_allocation_info(&self, ptr: *mut u8) -> Option<AllocationInfo> {
        self.allocated_blocks.get(&ptr).map(|block| AllocationInfo {
            ptr: block.ptr,
            size: block.size,
            order: block.order,
            allocated_at: block.allocated_at,
            last_accessed: block.last_accessed,
            access_count: block.access_count,
        })
    }
}

// Safety: BuddyAllocator manages GPU memory pointers. While *mut u8 is not Send/Sync by default,
// it's safe to share BuddyAllocator across threads when protected by Arc<Mutex<>> because:
// 1. The pointers point to GPU memory managed by the GPU driver
// 2. The Mutex provides exclusive access for all mutable operations
// 3. No thread-local state is maintained
unsafe impl Send for BuddyAllocator {}
unsafe impl Sync for BuddyAllocator {}

/// Memory usage information
#[derive(Debug, Clone)]
pub struct MemoryUsage {
    pub total_size: usize,
    pub allocated_size: usize,
    pub free_size: usize,
    pub fragmentation_ratio: f64,
    pub allocated_blocks: usize,
    pub free_blocks: usize,
}

/// Free block statistics
#[derive(Debug, Clone)]
pub struct FreeBlockStats {
    pub order: usize,
    pub block_size: usize,
    pub block_count: usize,
    pub total_size: usize,
}

/// Allocation information
#[derive(Debug, Clone)]
pub struct AllocationInfo {
    pub ptr: *mut u8,
    pub size: usize,
    pub order: usize,
    pub allocated_at: Option<Instant>,
    pub last_accessed: Option<Instant>,
    pub access_count: u64,
}

/// Buddy allocator errors
#[derive(Debug, Clone)]
pub enum BuddyError {
    InvalidSize(String),
    OutOfMemory(String),
    InvalidPointer(String),
    CorruptedState(String),
}

impl std::fmt::Display for BuddyError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            BuddyError::InvalidSize(msg) => write!(f, "Invalid size: {}", msg),
            BuddyError::OutOfMemory(msg) => write!(f, "Out of memory: {}", msg),
            BuddyError::InvalidPointer(msg) => write!(f, "Invalid pointer: {}", msg),
            BuddyError::CorruptedState(msg) => write!(f, "Corrupted state: {}", msg),
        }
    }
}

impl std::error::Error for BuddyError {}

/// Thread-safe buddy allocator wrapper
pub struct ThreadSafeBuddyAllocator {
    allocator: Arc<Mutex<BuddyAllocator>>,
}

impl ThreadSafeBuddyAllocator {
    pub fn new(
        base_ptr: *mut u8,
        total_size: usize,
        config: BuddyConfig,
    ) -> Result<Self, BuddyError> {
        let allocator = BuddyAllocator::new(base_ptr, total_size, config)?;
        Ok(Self {
            allocator: Arc::new(Mutex::new(allocator)),
        })
    }

    pub fn allocate(&self, size: usize) -> Result<*mut u8, BuddyError> {
        let mut allocator = self.allocator.lock().unwrap_or_else(|e| e.into_inner());
        allocator.allocate(size)
    }

    pub fn deallocate(&self, ptr: *mut u8) -> Result<(), BuddyError> {
        let mut allocator = self.allocator.lock().unwrap_or_else(|e| e.into_inner());
        allocator.deallocate(ptr)
    }

    pub fn get_stats(&self) -> BuddyStats {
        let allocator = self.allocator.lock().unwrap_or_else(|e| e.into_inner());
        allocator.get_stats().clone()
    }

    pub fn get_memory_usage(&self) -> MemoryUsage {
        let allocator = self.allocator.lock().unwrap_or_else(|e| e.into_inner());
        allocator.get_memory_usage()
    }

    pub fn defragment(&self) -> usize {
        let mut allocator = self.allocator.lock().unwrap_or_else(|e| e.into_inner());
        allocator.defragment()
    }
}

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

    #[test]
    fn test_buddy_allocator_creation() {
        let size = 1024 * 1024; // 1MB
        let config = BuddyConfig::default();

        // Simulate memory allocation
        let memory = vec![0u8; size];
        let ptr = memory.as_ptr() as *mut u8;

        let allocator = BuddyAllocator::new(ptr, size, config);
        assert!(allocator.is_ok());
    }

    #[test]
    fn test_basic_allocation() {
        let size = 1024 * 1024;
        let config = BuddyConfig::default();
        let memory = vec![0u8; size];
        let ptr = memory.as_ptr() as *mut u8;

        let mut allocator = BuddyAllocator::new(ptr, size, config).expect("unwrap failed");

        // Allocate some memory
        let alloc1 = allocator.allocate(1024);
        assert!(alloc1.is_ok());

        let alloc2 = allocator.allocate(2048);
        assert!(alloc2.is_ok());

        // Check stats
        let stats = allocator.get_stats();
        assert_eq!(stats.total_allocations, 2);
        assert_eq!(stats.successful_allocations, 2);
    }

    #[test]
    fn test_deallocation() {
        let size = 1024 * 1024;
        let config = BuddyConfig::default();
        let memory = vec![0u8; size];
        let ptr = memory.as_ptr() as *mut u8;

        let mut allocator = BuddyAllocator::new(ptr, size, config).expect("unwrap failed");

        let alloc_ptr = allocator.allocate(1024).expect("unwrap failed");
        let dealloc_result = allocator.deallocate(alloc_ptr);
        assert!(dealloc_result.is_ok());

        let stats = allocator.get_stats();
        assert_eq!(stats.total_deallocations, 1);
    }

    #[test]
    fn test_coalescing() {
        let size = 1024 * 1024;
        let config = BuddyConfig::default();
        let memory = vec![0u8; size];
        let ptr = memory.as_ptr() as *mut u8;

        let mut allocator = BuddyAllocator::new(ptr, size, config).expect("unwrap failed");

        // Allocate a larger block that will be split
        let large_ptr = allocator.allocate(4096).expect("unwrap failed");

        // Free it - this will add it back to the free list
        allocator.deallocate(large_ptr).expect("unwrap failed");

        // Now allocate two smaller blocks that are buddies
        // The allocator will split the 4096 block into two 2048 blocks
        let ptr1 = allocator.allocate(2048).expect("unwrap failed");
        let ptr2 = allocator.allocate(2048).expect("unwrap failed");

        // Deallocate them - they should coalesce back into the 4096 block
        allocator.deallocate(ptr1).expect("unwrap failed");
        allocator.deallocate(ptr2).expect("unwrap failed");

        let stats = allocator.get_stats();
        // `memory` is a real heap allocation, so `ptr` is an arbitrary
        // (non-zero, not power-of-two-aligned-to-`size`) address -- this is
        // a regression test for the buddy-address computation being
        // relative to the arena's `base_ptr` rather than the raw absolute
        // pointer (see `BuddyBlock::get_buddy_address`): with the wrong
        // (absolute) formula, this real address essentially never finds its
        // true buddy and coalescing silently never happens.
        assert!(
            stats.merge_operations > 0,
            "two adjacent same-size blocks with a real (non-zero) base pointer must coalesce"
        );
    }

    #[test]
    fn test_fragmentation_calculation() {
        let size = 1024 * 1024;
        let config = BuddyConfig::default();
        let memory = vec![0u8; size];
        let ptr = memory.as_ptr() as *mut u8;

        let allocator = BuddyAllocator::new(ptr, size, config).expect("unwrap failed");
        let fragmentation = allocator.calculate_fragmentation();

        // With one large free block, fragmentation should be minimal
        assert!(fragmentation < 0.1);
    }

    #[test]
    fn test_memory_usage() {
        let size = 1024 * 1024;
        let config = BuddyConfig::default();
        let memory = vec![0u8; size];
        let ptr = memory.as_ptr() as *mut u8;

        let mut allocator = BuddyAllocator::new(ptr, size, config).expect("unwrap failed");

        let usage_before = allocator.get_memory_usage();
        assert_eq!(usage_before.total_size, size);
        assert_eq!(usage_before.allocated_size, 0);

        allocator.allocate(1024).expect("unwrap failed");

        let usage_after = allocator.get_memory_usage();
        assert!(usage_after.allocated_size > 0);
    }

    #[test]
    fn test_thread_safe_allocator() {
        let size = 1024 * 1024;
        let config = BuddyConfig::default();
        let memory = vec![0u8; size];
        let ptr = memory.as_ptr() as *mut u8;

        let allocator = ThreadSafeBuddyAllocator::new(ptr, size, config).expect("unwrap failed");

        let alloc_result = allocator.allocate(1024);
        assert!(alloc_result.is_ok());

        let stats = allocator.get_stats();
        assert_eq!(stats.total_allocations, 1);
    }
}