buddy_system_allocator 0.13.0

use super::prev_power_of_two;
use alloc::collections::BTreeSet;
use core::alloc::Layout;
use core::cmp::{max, min};
use core::ops::Range;

#[cfg(feature = "use_spin")]
use core::ops::Deref;
#[cfg(feature = "use_spin")]
use spin::Mutex;

/// A frame allocator that uses buddy system, requiring a global allocator.
///
/// The max order of the allocator is determined by the const generic parameter `ORDER` (`MAX_ORDER = ORDER - 1`).
/// The frame allocator will only be able to allocate ranges of size up to 2<sup>MAX_ORDER</sup>, out of a total
/// range of size at most 2<sup>MAX_ORDER + 1</sup> - 1.
///
/// # Usage
///
/// Create a frame allocator and add some frames to it:
/// ```
/// use buddy_system_allocator::*;
/// // Notice that the max order is `ORDER - 1`.
/// let mut frame = FrameAllocator::<33>::new();
/// assert!(frame.alloc(1).is_none());
///
/// frame.add_frame(0, 3);
/// let num = frame.alloc(1);
/// assert_eq!(num, Some(2));
/// let num = frame.alloc(2);
/// assert_eq!(num, Some(0));
/// ```
pub struct FrameAllocator<const ORDER: usize = 33> {
    // buddy system with max order of `ORDER - 1`
    free_list: [BTreeSet<usize>; ORDER],

    // statistics
    allocated: usize,
    total: usize,
}

impl<const ORDER: usize> FrameAllocator<ORDER> {
    /// Create an empty frame allocator
    pub const fn new() -> Self {
        Self {
            free_list: [const { BTreeSet::new() }; ORDER],
            allocated: 0,
            total: 0,
        }
    }

    /// Add a range of frame number [start, end) to the allocator
    pub fn add_frame(&mut self, start: usize, end: usize) {
        assert!(start <= end);

        let mut total = 0;
        let mut current_start = start;
        let max_block_size = 1 << (ORDER - 1);

        while current_start < end {
            let lowbit = if current_start > 0 {
                current_start & (!current_start + 1)
            } else {
                max_block_size
            };
            let size = min(
                min(lowbit, prev_power_of_two(end - current_start)),
                max_block_size,
            );
            total += size;

            self.free_list[size.trailing_zeros() as usize].insert(current_start);
            current_start += size;
        }

        self.total += total;
    }

    /// Add a range of frames to the allocator.
    pub fn insert(&mut self, range: Range<usize>) {
        self.add_frame(range.start, range.end);
    }

    /// Allocate a range of frames from the allocator, returning the first frame of the allocated
    /// range.
    pub fn alloc(&mut self, count: usize) -> Option<usize> {
        let size = count.next_power_of_two();
        self.alloc_power_of_two(size)
    }

    /// Allocate a range of frames with the given size and alignment from the allocator, returning
    /// the first frame of the allocated range.
    /// The allocated size is the maximum of the next power of two of the given size and the
    /// alignment.
    pub fn alloc_aligned(&mut self, layout: Layout) -> Option<usize> {
        let size = max(layout.size().next_power_of_two(), layout.align());
        self.alloc_power_of_two(size)
    }

    /// Try to allocate a specific range of frames `[start, start + count)` from the allocator.
    ///
    /// The `count` will be rounded up to the next power of two, same as [`alloc`]. The `start`
    /// address must be aligned to this rounded-up size (buddy system invariant).
    ///
    /// Returns `Some(start)` if the range was successfully allocated, or `None` if the range is
    /// unavailable (not managed, already allocated, or misaligned).
    pub fn alloc_at(&mut self, start: usize, count: usize) -> Option<usize> {
        let size = count.next_power_of_two();
        self.alloc_at_power_of_two(start, size)
    }

    /// Allocate a range of frames of the given size from the allocator. The size must be a power of
    /// two. The allocated range will have alignment equal to the size.
    fn alloc_power_of_two(&mut self, size: usize) -> Option<usize> {
        let class = size.trailing_zeros() as usize;
        for i in class..self.free_list.len() {
            // Find the first non-empty size class
            if !self.free_list[i].is_empty() {
                // Split buffers
                for j in (class + 1..i + 1).rev() {
                    if let Some(block_ref) = self.free_list[j].iter().next() {
                        let block = *block_ref;
                        self.free_list[j - 1].insert(block + (1 << (j - 1)));
                        self.free_list[j - 1].insert(block);
                        self.free_list[j].remove(&block);
                    } else {
                        return None;
                    }
                }

                let result = self.free_list[class].iter().next();
                if let Some(result_ref) = result {
                    let result = *result_ref;
                    self.free_list[class].remove(&result);
                    self.allocated += size;
                    return Some(result);
                } else {
                    return None;
                }
            }
        }
        None
    }

    /// Allocate a specific range of frames at the given start address with the given power-of-two
    /// size. The start address must be aligned to the size.
    fn alloc_at_power_of_two(&mut self, start: usize, size: usize) -> Option<usize> {
        let class = size.trailing_zeros() as usize;

        if start & (size - 1) != 0 {
            return None;
        }

        for i in class..self.free_list.len() {
            let block_start = start & !((1 << i) - 1);
            if self.free_list[i].remove(&block_start) {
                let mut current_start = block_start;
                for j in (class..i).rev() {
                    let mid = current_start + (1 << j);
                    if start >= mid {
                        self.free_list[j].insert(current_start);
                        current_start = mid;
                    } else {
                        self.free_list[j].insert(mid);
                    }
                }
                self.allocated += size;
                return Some(start);
            }
        }

        None
    }

    /// Deallocate a range of frames [frame, frame+count) from the frame allocator.
    ///
    /// The range should be exactly the same when it was allocated, as in heap allocator
    pub fn dealloc(&mut self, start_frame: usize, count: usize) {
        let size = count.next_power_of_two();
        self.dealloc_power_of_two(start_frame, size)
    }

    /// Deallocate a range of frames which was previously allocated by [`alloc_aligned`].
    ///
    /// The layout must be exactly the same as when it was allocated.
    pub fn dealloc_aligned(&mut self, start_frame: usize, layout: Layout) {
        let size = max(layout.size().next_power_of_two(), layout.align());
        self.dealloc_power_of_two(start_frame, size)
    }

    /// Deallocate a range of frames with the given size from the allocator. The size must be a
    /// power of two.
    fn dealloc_power_of_two(&mut self, start_frame: usize, size: usize) {
        let class = size.trailing_zeros() as usize;

        // Merge free buddy lists
        let mut current_ptr = start_frame;
        let mut current_class = class;
        while current_class < self.free_list.len() - 1 {
            let buddy = current_ptr ^ (1 << current_class);
            if self.free_list[current_class].remove(&buddy) {
                // Free buddy found
                current_ptr = min(current_ptr, buddy);
                current_class += 1;
            } else {
                break;
            }
        }

        self.free_list[current_class].insert(current_ptr);
        self.allocated -= size;
    }
}

impl<const ORDER: usize> Default for FrameAllocator<ORDER> {
    fn default() -> Self {
        Self::new()
    }
}

/// A locked version of `FrameAllocator`
///
/// # Usage
///
/// Create a locked frame allocator and add frames to it:
/// ```
/// use buddy_system_allocator::*;
/// // Notice that the max order is `ORDER - 1`.
/// let mut frame = LockedFrameAllocator::<33>::new();
/// assert!(frame.lock().alloc(1).is_none());
///
/// frame.lock().add_frame(0, 3);
/// let num = frame.lock().alloc(1);
/// assert_eq!(num, Some(2));
/// let num = frame.lock().alloc(2);
/// assert_eq!(num, Some(0));
/// ```
#[cfg(feature = "use_spin")]
#[derive(Default)]
pub struct LockedFrameAllocator<const ORDER: usize = 33>(Mutex<FrameAllocator<ORDER>>);

#[cfg(feature = "use_spin")]
impl<const ORDER: usize> LockedFrameAllocator<ORDER> {
    /// Creates an empty frame allocator.
    pub const fn new() -> Self {
        Self(Mutex::new(FrameAllocator::new()))
    }
}

#[cfg(feature = "use_spin")]
impl<const ORDER: usize> Deref for LockedFrameAllocator<ORDER> {
    type Target = Mutex<FrameAllocator<ORDER>>;

    fn deref(&self) -> &Mutex<FrameAllocator<ORDER>> {
        &self.0
    }
}