//! Fixed-sized block allocator.
//!
//! Implements a fixed-sized block allocator backed by a linked list allocator. Based on the example fixed-sized block
//! allocator presented here: <https://os.phil-opp.com/allocator-designs/#fixed-size-block-allocator>.
//!
//! ## License
//!
//! Copyright (c) Microsoft Corporation.
//!
//! SPDX-License-Identifier: Apache-2.0
//!

extern crate alloc;
use super::{AllocationStatistics, AllocationStrategy, DEFAULT_ALLOCATION_STRATEGY, PageAllocator};

use crate::{gcd::SpinLockedGcd, tpl_mutex};

use alloc::vec::Vec;
use core::{
    alloc::{AllocError, Allocator, GlobalAlloc, Layout},
    cmp::max,
    debug_assert,
    fmt::{self, Display},
    mem::{align_of, size_of},
    ops::Range,
    ptr::{NonNull, slice_from_raw_parts_mut},
    result::Result,
};
use linked_list_allocator::{align_down_size, align_up_size};
use patina::{
    base::{UEFI_PAGE_SIZE, align_up, page_shift_from_alignment},
    error::EfiError,
    pi::dxe_services::GcdMemoryType,
    uefi_pages_to_size, uefi_size_to_pages, writelncrlf,
};
use r_efi::efi;

/// Type for describing errors that this implementation can produce.
#[derive(Debug, PartialEq)]
pub enum FixedSizeBlockAllocatorError {
    /// Could not satisfy allocation request, and expansion failed.
    ///
    /// Specifies how much additional memory is required to be added to the allocator through
    /// [FixedSizeBlockAllocator::expand()] in order to fulfill the attempted allocation.
    OutOfMemory(usize),
    /// The provided layout was invalid.
    InvalidLayout,
    /// The memory region provided to extend the allocator was invalid.
    InvalidExpansion,
    /// An internal error occurred.
    InternalError,
}

const ALIGNMENT: usize = 0x1000;

const BLOCK_SIZES: &[usize] = &[8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096];

// Returns the index in the block list for the minimum size block that will
// satisfy allocation for the given layout
fn list_index(layout: &Layout) -> Option<usize> {
    let required_block_size = layout.size().max(layout.align());
    BLOCK_SIZES.iter().position(|&s| s >= required_block_size)
}

struct BlockListNode {
    next: Option<&'static mut BlockListNode>,
}

struct AllocatorListNode {
    next: Option<*mut AllocatorListNode>,
    allocator: linked_list_allocator::Heap,
}
struct AllocatorIterator {
    current: Option<*mut AllocatorListNode>,
}

impl AllocatorIterator {
    fn new(start_node: Option<*mut AllocatorListNode>) -> Self {
        AllocatorIterator { current: start_node }
    }
}

impl Iterator for AllocatorIterator {
    type Item = *mut AllocatorListNode;
    fn next(&mut self) -> Option<*mut AllocatorListNode> {
        if let Some(current) = self.current {
            // SAFETY: current is a valid node pointer from the allocator list.
            self.current = unsafe { (*current).next };
            Some(current)
        } else {
            None
        }
    }
}

/// Fixed Size Block Allocator
///
/// Implements an expandable memory allocator using fixed-sized blocks for speed backed by a linked-list allocator
/// implementation when an appropriate sized free block is not available. If more memory is required than can be
/// satisfied by either the block list or the linked-list, more memory is is allocated externally, then passed into
/// the allocator where a new backing linked-list is created.
///
pub struct FixedSizeBlockAllocator {
    /// The memory type this allocator manages.
    memory_type: efi::MemoryType,

    /// The heads of the linked lists for each fixed-size block. Each index corresponds to a block size in
    /// `BLOCK_SIZES`.
    list_heads: [Option<&'static mut BlockListNode>; BLOCK_SIZES.len()],

    /// The linked-list of allocators that this allocator uses to back allocations that are larger than the fixed-size
    /// blocks or if the required fixed-size block list is empty.
    allocators: Option<*mut AllocatorListNode>,

    /// The range of memory that is reserved for this allocator. This is used to stabilize the memory map during an
    /// S4 resume.
    pub(crate) reserved_range: Option<Range<efi::PhysicalAddress>>,

    /// Statistics about the allocator's usage.
    stats: AllocationStatistics,

    /// The page allocation granularity used by this allocator. This is expected to be one of the following:
    /// - `SIZE_4KB` for all allocators except AARCH64 runtime memory allocators
    /// - `SIZE_64KB` for AARCH64 runtime memory allocators
    page_allocation_granularity: usize,
}

impl FixedSizeBlockAllocator {
    /// Creates a new empty FixedSizeBlockAllocator
    pub const fn new(memory_type: efi::MemoryType, page_allocation_granularity: usize) -> Self {
        const EMPTY: Option<&'static mut BlockListNode> = None;
        FixedSizeBlockAllocator {
            memory_type,
            list_heads: [EMPTY; BLOCK_SIZES.len()],
            allocators: None,
            reserved_range: None,
            stats: AllocationStatistics::new(),
            page_allocation_granularity,
        }
    }

    // This routine resets some aspects of allocator state for testing purposes.
    // Note: this does not change the page_change_callback.
    #[cfg(test)]
    pub fn reset(&mut self) {
        const EMPTY: Option<&'static mut BlockListNode> = None;
        self.list_heads = [EMPTY; BLOCK_SIZES.len()];
        self.allocators = None;
        self.reserved_range = None;
        self.stats = AllocationStatistics::new();
    }

    /// Expand the memory available to this allocator with a new contiguous region of memory, setting up a new allocator
    /// node to manage this range. `new_region.len() - size_of::<AllocatorListNode>()` additional memory will be available
    /// to the allocator.
    ///
    /// ## Errors
    ///
    /// Returns [`FixedSizeBlockAllocatorError::InvalidExpansion`] if the new region is not larger than and aligned to
    /// AllocatorListNode.
    pub fn expand(&mut self, new_region: NonNull<[u8]>) -> core::result::Result<(), FixedSizeBlockAllocatorError> {
        // Ensure we're expanding enough to fit a new allocator list node
        if new_region.len() <= size_of::<AllocatorListNode>() {
            debug_assert!(false, "FSB expanded with insufficiently sized memory region.");
            return Err(FixedSizeBlockAllocatorError::InvalidExpansion);
        }

        // Interpret the first part of the provided region as an AllocatorListNode
        let alloc_node_ptr = new_region.as_ptr() as *mut AllocatorListNode;

        if !alloc_node_ptr.is_aligned() {
            debug_assert!(false, "FSB expanded with memory region unaligned to AllocatorListNode.");
            return Err(FixedSizeBlockAllocatorError::InvalidExpansion);
        }

        let heap_region: NonNull<[u8]> = NonNull::slice_from_raw_parts(
            // SAFETY: alloc_node_ptr is validated above and points to a region large enough for an AllocatorListNode.
            NonNull::new(unsafe { alloc_node_ptr.add(1) }).unwrap().cast(),
            new_region.len() - size_of::<AllocatorListNode>(),
        );

        //write the allocator node structure into the start of the range, initialize its heap with the remainder of
        //the range, and add the new allocator to the front of the allocator list.
        let node = AllocatorListNode { next: None, allocator: linked_list_allocator::Heap::empty() };
        // SAFETY: alloc_node_ptr is aligned and points to valid writable memory for an AllocatorListNode.
        unsafe {
            alloc_node_ptr.write(node);
            (*alloc_node_ptr).allocator.init(heap_region.cast::<u8>().as_ptr(), heap_region.len());
            (*alloc_node_ptr).next = self.allocators;
        }

        self.allocators = Some(alloc_node_ptr);

        if self.in_reserved_range(alloc_node_ptr.addr() as efi::PhysicalAddress) {
            self.stats.reserved_used += new_region.len();
        } else {
            self.stats.claimed_pages += uefi_size_to_pages!(new_region.len());
        }

        Ok(())
    }

    // allocates from the linked-list backing allocator if a free block of the
    // appropriate size is not available.
    fn fallback_alloc(&mut self, layout: Layout) -> Result<NonNull<[u8]>, FixedSizeBlockAllocatorError> {
        for node in AllocatorIterator::new(self.allocators) {
            // SAFETY: node is a valid allocator list node pointer from the iterator.
            let allocator = unsafe { &mut (*node).allocator };
            if let Ok(ptr) = allocator.allocate_first_fit(layout) {
                return Ok(NonNull::slice_from_raw_parts(ptr, layout.size()));
            }
        }

        // Determine how much additional memory is required
        //
        // Per the `linked_list_allocator::hole::HoleList::new` documentation, depending on the alignment of the
        // hole_addr pointer, the minimum size for storing required metadata is between 2 * size_of::<usize> and
        //  3 * size_of::<usize>. The size reservation for `additional_mem_required` assumed the largest size.
        let additional_mem_required = layout.pad_to_align().size()
            + Layout::new::<AllocatorListNode>().pad_to_align().size()
            + 3 * size_of::<usize>();
        let additional_mem_required = align_up_size(additional_mem_required, align_of::<AllocatorListNode>());

        Err(FixedSizeBlockAllocatorError::OutOfMemory(additional_mem_required))
    }

    /// Allocates and returns a pointer to a memory buffer for the given layout.
    ///
    ///
    /// Memory allocated by this routine should be deallocated with
    /// [`Self::dealloc`]
    ///
    /// ## Errors
    ///
    /// Returns [`FixedSizeBlockAllocatorError::OutOfMemory`] when the allocator doesn't have enough memory.
    /// Returns [`FixedSizeBlockAllocatorError::InvalidLayout`] when the layout provided is invalid.
    /// Returns [`FixedSizeBlockAllocatorError::InternalError`] when an internal error occurs.
    pub fn alloc(&mut self, layout: Layout) -> Result<NonNull<[u8]>, FixedSizeBlockAllocatorError> {
        self.stats.pool_allocation_calls += 1;

        match list_index(&layout) {
            Some(index) => {
                let head = self.list_heads.get_mut(index).ok_or(FixedSizeBlockAllocatorError::InternalError)?;
                match head.take() {
                    Some(node) => {
                        let head = self.list_heads.get_mut(index).ok_or(FixedSizeBlockAllocatorError::InternalError)?;
                        *head = node.next.take();
                        let ptr: NonNull<u8> = NonNull::from(node).cast();
                        Ok(NonNull::slice_from_raw_parts(ptr, layout.size()))
                    }
                    None => {
                        // no block exists in list => allocate new block
                        let block_size = *BLOCK_SIZES.get(index).ok_or(FixedSizeBlockAllocatorError::InternalError)?;
                        // only works if all block sizes are a power of 2
                        let block_align = block_size;
                        let layout = match Layout::from_size_align(block_size, block_align) {
                            Ok(layout) => layout,
                            Err(_) => return Err(FixedSizeBlockAllocatorError::InvalidLayout),
                        };
                        self.fallback_alloc(layout)
                    }
                }
            }
            None => self.fallback_alloc(layout),
        }
    }

    // deallocates back to the linked-list backing allocator if the size of
    // layout being freed is too big to be tracked as a fixed-size free block.
    fn fallback_dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
        for node in AllocatorIterator::new(self.allocators) {
            // SAFETY: node is produced by AllocatorIterator and points to a valid AllocatorListNode.
            let allocator = unsafe { &mut (*node).allocator };
            if (allocator.bottom() <= ptr.as_ptr()) && (ptr.as_ptr() < allocator.top()) {
                // SAFETY: ptr was allocated by this allocator for the given layout.
                unsafe { allocator.deallocate(ptr, layout) };
            }
        }
    }

    /// Deallocates a buffer allocated by [`Self::alloc`].
    ///
    /// ## Safety
    ///
    /// Caller must ensure that `ptr` was created by a call to [`Self::alloc`] with the same `layout`.
    pub unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
        self.stats.pool_free_calls += 1;
        match list_index(&layout) {
            Some(index) => {
                let head = self.list_heads.get_mut(index).expect("list_index guarantees valid index");
                let new_node = BlockListNode { next: head.take() };
                let block_size = *BLOCK_SIZES.get(index).expect("list_index guarantees valid index");
                // verify that block has size and alignment required for storing node
                if size_of::<BlockListNode>() > block_size || align_of::<BlockListNode>() > block_size {
                    // Should never reach this statement under normal operation since BlockListNode is a single pointer and all block sizes are >= 8 bytes,
                    // Failure indicates corruption of the allocator's internal state.
                    panic!("FSB deallocating block too small to store BlockListNode.");
                }
                let new_node_ptr = ptr.as_ptr() as *mut BlockListNode;
                // SAFETY: new_node_ptr points to memory returned by alloc for this layout.
                unsafe {
                    new_node_ptr.write(new_node);
                    let head = self.list_heads.get_mut(index).expect("list_index guarantees valid index");
                    *head = Some(&mut *new_node_ptr);
                }
            }
            None => {
                self.fallback_dealloc(ptr, layout);
            }
        }
    }

    /// Returns whether the provided address is in the FSB's reserved range
    pub fn in_reserved_range(&self, address: efi::PhysicalAddress) -> bool {
        match &self.reserved_range {
            Some(reserved_range) => reserved_range.contains(&address),
            _ => false,
        }
    }

    /// Sets the reserved memory range (bin range) for this allocator.
    ///
    /// ## Errors
    ///
    /// Returns [`EfiError::AlreadyStarted`] if a reserved range has already been set.
    pub fn set_reserved_range(&mut self, range: Range<efi::PhysicalAddress>) -> Result<(), EfiError> {
        if self.reserved_range.is_some() {
            Err(EfiError::AlreadyStarted)?;
        }

        let size = (range.end - range.start) as usize;
        self.reserved_range = Some(range);
        self.stats.reserved_size = size;
        self.stats.reserved_used = 0;
        self.stats.claimed_pages += uefi_size_to_pages!(size);

        Ok(())
    }

    /// Indicates whether the given pointer falls within a memory region managed by this allocator.
    ///
    /// Note: `true` does not indicate that the pointer corresponds to an active allocation - it may be in either
    /// allocated or freed memory. `true` just means that the pointer falls within a memory region that this allocator
    /// manages.
    pub fn contains(&self, ptr: *mut u8) -> bool {
        AllocatorIterator::new(self.allocators).any(|node| {
            // SAFETY: node is produced by AllocatorIterator and points to a valid AllocatorListNode.
            let allocator = unsafe { &mut (*node).allocator };
            (allocator.bottom() <= ptr) && (ptr < allocator.top())
        })
    }

    /// Tracks page allocations for record keeping
    pub fn notify_page_allocation(&mut self, allocation: NonNull<[u8]>) {
        if self.in_reserved_range(allocation.addr().get() as efi::PhysicalAddress) {
            self.stats.reserved_used += allocation.len();
        } else {
            self.stats.claimed_pages += uefi_size_to_pages!(allocation.len());
        }
    }

    /// Tracks page freeing for record keeping
    pub fn notify_pages_freed(&mut self, address: efi::PhysicalAddress, pages: usize) {
        if self.in_reserved_range(address) {
            self.stats.reserved_used = self.stats.reserved_used.saturating_sub(pages * ALIGNMENT);
        } else {
            self.stats.claimed_pages = self.stats.claimed_pages.saturating_sub(pages);
        }
    }

    /// Get the ranges of the memory owned by this allocator
    ///
    /// Returns an iterator of ranges of the memory owned by this allocator.
    /// If the allocator does not own any memory, it will return an empty iterator.
    pub(crate) fn get_memory_ranges(&self) -> impl Iterator<Item = Range<usize>> {
        AllocatorIterator::new(self.allocators).map(|node| {
            // SAFETY: node is produced by AllocatorIterator and points to a valid AllocatorListNode.
            let allocator = unsafe { &(*node).allocator };
            allocator.bottom() as usize..allocator.top() as usize
        })
    }

    /// Returns the memory type for this allocator
    #[inline(always)]
    pub fn memory_type(&self) -> efi::MemoryType {
        self.memory_type
    }

    /// Returns a reference to the allocation stats for this allocator.
    pub fn stats(&self) -> &AllocationStatistics {
        &self.stats
    }
}

impl Display for FixedSizeBlockAllocator {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writelncrlf!(f, "Memory Type: {:x?}", self.memory_type())?;
        writelncrlf!(f, "Allocation Ranges:")?;
        for node in AllocatorIterator::new(self.allocators) {
            // SAFETY: node is produced by AllocatorIterator and points to a valid AllocatorListNode.
            let allocator = unsafe { &mut (*node).allocator };
            writelncrlf!(
                f,
                "  PhysRange: {:#x}-{:#x}, Size: {:#x}, Used: {:#x} Free: {:#x}",
                align_down_size(allocator.bottom() as usize, 0x1000), //account for AllocatorListNode
                allocator.top() as usize,
                align_up_size(allocator.size(), 0x1000), //account for AllocatorListNode
                allocator.used(),
                allocator.free(),
            )?;
        }
        writelncrlf!(f, "Bucket Range: {:x?}", self.reserved_range)?;
        writelncrlf!(f, "Allocation Stats:")?;
        writelncrlf!(f, "  pool_allocation_calls: {}", self.stats.pool_allocation_calls)?;
        writelncrlf!(f, "  pool_free_calls: {}", self.stats.pool_free_calls)?;
        writelncrlf!(f, "  page_allocation_calls: {}", self.stats.page_allocation_calls)?;
        writelncrlf!(f, "  page_free_calls: {}", self.stats.page_free_calls)?;
        writelncrlf!(f, "  reserved_size: {}", self.stats.reserved_size)?;
        writelncrlf!(f, "  reserved_used: {}", self.stats.reserved_used)?;
        writelncrlf!(f, "  claimed_pages: {}", self.stats.claimed_pages)?;
        Ok(())
    }
}

/// Spin Locked Fixed Size Block Allocator
///
/// A wrapper for [`FixedSizeBlockAllocator`] that allocates additional memory as needed from a GCD
/// and provides Sync/Send via means of a spin mutex.
///
/// Note: [SpinLockedFixedSizeBlockAllocator::alloc()] and [SpinLockedFixedSizeBlockAllocator::allocate()] will call
/// alloc() twice when additional memory is required.
pub struct SpinLockedFixedSizeBlockAllocator {
    /// The GCD instance that this allocator uses to allocate additional memory as needed.
    gcd: &'static SpinLockedGcd,

    /// The handle associated with this allocator. It is used to track ownership of the memory allocated
    /// by this allocator in the GCD.
    handle: efi::Handle,

    /// The inner allocator that is protected by a TPL mutex.
    inner: tpl_mutex::TplMutex<FixedSizeBlockAllocator>,

    /// The minimum amount of memory to allocate when expanding the allocator.
    min_expansion: usize,
}

impl SpinLockedFixedSizeBlockAllocator {
    /// Creates a new empty FixedSizeBlockAllocator that will request memory from `gcd` as needed to satisfy
    /// requests.
    pub const fn new(
        gcd: &'static SpinLockedGcd,
        allocator_handle: efi::Handle,
        memory_type: efi::MemoryType,
        page_allocation_granularity: usize,
        min_expansion: usize,
    ) -> Self {
        SpinLockedFixedSizeBlockAllocator {
            gcd,
            handle: allocator_handle,
            inner: tpl_mutex::TplMutex::new(
                efi::TPL_HIGH_LEVEL,
                FixedSizeBlockAllocator::new(memory_type, page_allocation_granularity),
                "FsbLock",
            ),
            min_expansion,
        }
    }

    // This routine resets some aspects of allocator state for testing purposes.
    // Note: this does not reset the GCD nor change the page_change_callback.
    #[cfg(test)]
    pub fn reset(&self) {
        self.lock().reset();
    }

    /// Locks the allocator
    ///
    /// This can be used to do several actions on the allocator atomically.
    pub fn lock(&self) -> tpl_mutex::TplGuard<'_, FixedSizeBlockAllocator> {
        self.inner.lock()
    }

    /// Indicates whether the given pointer falls within a memory region managed by this allocator.
    ///
    /// See [`FixedSizeBlockAllocator::contains()`]
    pub fn contains(&self, ptr: NonNull<u8>) -> bool {
        self.lock().contains(ptr.as_ptr())
    }

    /// Attempts to allocate the given number of pages according to the given allocation strategy.
    /// Valid allocation strategies are:
    /// - BottomUp(None): Allocate the block of pages from the lowest available free memory.
    /// - BottomUp(Some(address)): Allocate the block of pages from the lowest available free memory. Fail if memory
    ///   cannot be found below `address`.
    /// - TopDown(None): Allocate the block of pages from the highest available free memory.
    /// - TopDown(Some(address)): Allocate the block of pages from the highest available free memory. Fail if memory
    ///   cannot be found above `address`.
    /// - Address(address): Allocate the block of pages at exactly the given address (or fail).
    ///
    /// If an address is specified as part of a strategy, it must be page-aligned.
    pub fn allocate_pages(
        &self,
        allocation_strategy: AllocationStrategy,
        pages: usize,
        alignment: usize,
    ) -> Result<NonNull<[u8]>, EfiError> {
        // Record this call in the FSB's stats
        self.lock().stats.page_allocation_calls += 1;
        let granularity = self.lock().page_allocation_granularity;

        // Granularity and alignment both are powers of two, so we can use the max of the two
        let required_alignment = max(granularity, alignment);

        // Ensure that the requested number of pages is a multiple of the granularity
        let required_pages = align_up(pages, uefi_size_to_pages!(granularity))?;

        let align_shift = page_shift_from_alignment(required_alignment)?;

        if let AllocationStrategy::Address(address) = allocation_strategy {
            // validate allocation strategy addresses for direct address allocation is properly aligned.
            // for BottomUp and TopDown strategies, the address parameter doesn't have to be page-aligned, but
            // the resulting allocation will be page-aligned.
            if address % required_alignment != 0 {
                return Err(EfiError::InvalidParameter);
            }
        }

        // Page allocations and pool allocations are disjoint; page allocations are allocated directly from the GCD and are
        // freed straight back to GCD. As such, a tracking allocator structure is not required.
        let start_address = self
            .allocate_from_gcd(allocation_strategy, align_shift, uefi_pages_to_size!(required_pages))
            .map_err(|err| match err {
                EfiError::InvalidParameter | EfiError::NotFound => err,
                _ => EfiError::OutOfResources,
            })?;

        let allocation = slice_from_raw_parts_mut(start_address as *mut u8, uefi_pages_to_size!(required_pages));
        let allocation = NonNull::new(allocation).ok_or(EfiError::OutOfResources)?;

        // Notify the FSB that additional pages were allocated for record keeping
        self.lock().notify_page_allocation(allocation);

        Ok(allocation)
    }

    /// Frees the block of pages at the given address of the given size.
    ///
    /// ## Safety
    /// Caller must ensure that the given address corresponds to a valid block of pages that was allocated with
    /// [Self::allocate_pages]
    pub unsafe fn free_pages(&self, address: usize, pages: usize) -> Result<(), EfiError> {
        self.lock().stats.page_free_calls += 1;

        let granularity = self.lock().page_allocation_granularity;

        // Ensure that the requested number of pages is a multiple of the granularity
        let required_pages = align_up(pages, uefi_size_to_pages!(granularity))?;

        if !address.is_multiple_of(granularity) {
            return Err(EfiError::InvalidParameter);
        }

        let descriptor =
            self.gcd.get_existent_memory_descriptor_for_address(address as efi::PhysicalAddress).map_err(|err| {
                match err {
                    EfiError::NotFound => err,
                    _ => EfiError::InvalidParameter,
                }
            })?;

        if descriptor.image_handle != self.handle {
            Err(EfiError::NotFound)?;
        }

        if self.lock().in_reserved_range(address as efi::PhysicalAddress) {
            self.gcd.free_memory_space_preserving_ownership(address, uefi_pages_to_size!(required_pages)).map_err(
                |err| match err {
                    EfiError::NotFound => err,
                    _ => EfiError::InvalidParameter,
                },
            )?;
        } else {
            self.gcd.free_memory_space(address, uefi_pages_to_size!(required_pages)).map_err(|err| match err {
                EfiError::NotFound => err,
                _ => EfiError::InvalidParameter,
            })?;
        }

        // Notify the FSB that pages were freed for record keeping
        self.lock().notify_pages_freed(address as efi::PhysicalAddress, required_pages);

        Ok(())
    }

    /// Sets the reserved memory range (bin range) for this allocator.
    ///
    /// See [`FixedSizeBlockAllocator::set_reserved_range()`] for details on the accounting model.
    ///
    /// ## Errors
    ///
    /// Returns [`EfiError::AlreadyStarted`] if a reserved range has already been set.
    pub fn set_reserved_range(&self, range: Range<efi::PhysicalAddress>) -> Result<(), EfiError> {
        self.lock().set_reserved_range(range)
    }

    /// Attempts to allocate from the GCD, preferring the reserved (bin) range if one exists.
    ///
    /// For strategies that do not exclude the bin range, this first tries to allocate within the
    /// bin range so that special-type pages land in their designated bin. Specifically, bin
    /// preference is attempted for:
    /// - `TopDown(None)` / `BottomUp(None)`: fully unconstrained strategies.
    /// - `TopDown(Some(max))` / `BottomUp(Some(max))`: constrained strategies where `max` is at
    ///   or above the bin range end, meaning the bin is reachable.
    ///
    /// `Address(addr)` strategies are never redirected because the caller requires an exact address.
    ///
    /// If the bin is full or no bin exists, the allocation falls through to the original strategy.
    fn allocate_from_gcd(
        &self,
        strategy: AllocationStrategy,
        align_shift: usize,
        size: usize,
    ) -> Result<usize, EfiError> {
        let reserved_range = self.lock().reserved_range.clone();

        // Determine whether to attempt bin-preference allocation.
        let try_bin = match strategy {
            AllocationStrategy::TopDown(None) | AllocationStrategy::BottomUp(None) => true,
            AllocationStrategy::TopDown(Some(max)) | AllocationStrategy::BottomUp(Some(max)) => {
                if let Some(ref reserved) = reserved_range { max as u64 >= reserved.start + size as u64 } else { false }
            }
            _ => false,
        };

        if try_bin
            && let Some(ref reserved) = reserved_range
            && let Ok(addr) = self.gcd.allocate_memory_space(
                AllocationStrategy::TopDown(Some(reserved.end as usize)),
                GcdMemoryType::SystemMemory,
                align_shift,
                size,
                self.handle,
                None,
            )
        {
            if addr >= reserved.start as usize {
                return Ok(addr);
            }
            // Landed below the bin, free and fall through.
            let _ = self.gcd.free_memory_space(addr, size);
        }

        // Normal allocation path.
        self.gcd.allocate_memory_space(strategy, GcdMemoryType::SystemMemory, align_shift, size, self.handle, None)
    }

    /// Returns an iterator of the ranges of memory owned by this allocator
    /// Returns an empty iterator if the allocator does not own any memory.
    pub fn get_memory_ranges(&self) -> alloc::vec::IntoIter<Range<usize>> {
        let ranges: Vec<_> = self.lock().get_memory_ranges().collect();
        ranges.into_iter()
    }

    /// Returns the allocator handle associated with this allocator.
    pub fn handle(&self) -> efi::Handle {
        self.handle
    }

    /// Returns the reserved memory range, if any.
    #[coverage(off)]
    pub fn reserved_range(&self) -> Option<Range<efi::PhysicalAddress>> {
        self.inner.lock().reserved_range.clone()
    }

    /// Returns the memory type for this allocator.
    #[allow(dead_code)]
    #[coverage(off)]
    pub fn memory_type(&self) -> efi::MemoryType {
        self.inner.lock().memory_type()
    }

    /// Returns allocation statistics for this allocator.
    #[allow(dead_code)]
    pub fn stats(&self) -> AllocationStatistics {
        *self.inner.lock().stats()
    }
}

// SAFETY: SpinLockedFixedSizeBlockAllocator serializes access and delegates to the inner allocator.
unsafe impl GlobalAlloc for SpinLockedFixedSizeBlockAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        match self.allocate(layout) {
            Ok(alloc) => alloc.as_ptr() as *mut u8,
            Err(_) => core::ptr::null_mut(),
        }
    }
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        if let Some(ptr) = NonNull::new(ptr) {
            // SAFETY: ptr came from alloc with the same layout.
            unsafe { self.deallocate(ptr, layout) }
        }
    }
}

// SAFETY: SpinLockedFixedSizeBlockAllocator serializes access and delegates to the inner allocator.
unsafe impl Allocator for SpinLockedFixedSizeBlockAllocator {
    fn allocate(&self, layout: Layout) -> core::result::Result<NonNull<[u8]>, AllocError> {
        let allocation = self.lock().alloc(layout);
        match allocation {
            Ok(alloc) => Ok(alloc),
            Err(FixedSizeBlockAllocatorError::OutOfMemory(additional_mem_required)) => {
                // Compile-time check to ensure ALIGNMENT is compatible with the alignment requirements
                // of `expand()` and `page_shift_from_alignment()`
                const _: () = assert!(ALIGNMENT.is_multiple_of(align_of::<AllocatorListNode>()));
                const _: () = assert!(ALIGNMENT.is_multiple_of(UEFI_PAGE_SIZE) && ALIGNMENT > 0);

                // As a matter of policy, allocate at least the minimum expansion amount of memory and ensure the
                // size is aligned to `ALIGNMENT`.
                let mut allocation_size = max(additional_mem_required, self.min_expansion);
                let required_alignment = self.lock().page_allocation_granularity;

                // Ensure that the requested number of pages is a multiple of the granularity
                let required_pages =
                    align_up(uefi_size_to_pages!(allocation_size), uefi_size_to_pages!(required_alignment)).map_err(
                        |_| {
                            debug_assert!(false);
                            AllocError
                        },
                    )?;

                allocation_size = uefi_pages_to_size!(required_pages);

                // Allocate additional memory through the GCD, returning AllocError
                // if the GCD returns an error
                let start_address: usize = self
                    .allocate_from_gcd(
                        DEFAULT_ALLOCATION_STRATEGY,
                        page_shift_from_alignment(required_alignment).map_err(|_| {
                            debug_assert!(false);
                            AllocError
                        })?,
                        allocation_size,
                    )
                    .map_err(|err| {
                        log::error!(
                            "Allocator Expansion via GCD failed: [{err:?}], {{ Bytes: {allocation_size:#x}, Alignment: {required_alignment:#x}, Page Count: {required_pages:#x} }}",
                        );
                        AllocError
                    })?;

                // Expand the FSB using the allocated memory region
                let allocated_ptr = NonNull::new(start_address as *mut u8).ok_or_else(|| {
                    debug_assert!(false);
                    AllocError
                })?;
                if self.lock().expand(NonNull::slice_from_raw_parts(allocated_ptr, allocation_size)).is_err() {
                    debug_assert!(false);
                    return Err(AllocError);
                }

                // Try the allocation one more time
                match self.lock().alloc(layout) {
                    Ok(alloc) => Ok(alloc),
                    Err(_) => {
                        debug_assert!(false);
                        Err(AllocError)
                    }
                }
            }
            Err(_) => {
                debug_assert!(false);
                Err(AllocError)
            }
        }
    }
    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
        // SAFETY: ptr came from allocate with the same layout.
        unsafe { self.lock().dealloc(ptr, layout) }
    }
}

impl Display for SpinLockedFixedSizeBlockAllocator {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.lock().fmt(f)
    }
}

// SAFETY: SpinLockedFixedSizeBlockAllocator protects internal state with a spin lock.
unsafe impl Sync for SpinLockedFixedSizeBlockAllocator {}
// SAFETY: SpinLockedFixedSizeBlockAllocator protects internal state with a spin lock.
unsafe impl Send for SpinLockedFixedSizeBlockAllocator {}

impl PageAllocator for SpinLockedFixedSizeBlockAllocator {
    fn allocate_pages(
        &self,
        allocation_strategy: AllocationStrategy,
        pages: usize,
        alignment: usize,
    ) -> Result<NonNull<[u8]>, EfiError> {
        Self::allocate_pages(self, allocation_strategy, pages, alignment)
    }

    /// Frees the block of pages at the given address of the given size.
    ///
    /// ## Safety
    ///
    /// Caller must ensure that the given address corresponds to a valid block of pages that was allocated with
    /// [Self::allocate_pages].
    unsafe fn free_pages(&self, address: usize, pages: usize) -> Result<(), EfiError> {
        // SAFETY: address/pages must refer to a valid allocation owned by this allocator
        // per the free_pages safety contract.
        unsafe { Self::free_pages(self, address, pages) }
    }

    fn set_reserved_range(&self, range: Range<efi::PhysicalAddress>) -> Result<(), EfiError> {
        Self::set_reserved_range(self, range)
    }

    fn get_memory_ranges(&self) -> alloc::vec::IntoIter<Range<usize>> {
        Self::get_memory_ranges(self)
    }

    fn contains(&self, ptr: NonNull<u8>) -> bool {
        Self::contains(self, ptr)
    }

    fn handle(&self) -> efi::Handle {
        Self::handle(self)
    }

    fn reserved_range(&self) -> Option<Range<efi::PhysicalAddress>> {
        Self::reserved_range(self)
    }

    fn stats(&self) -> AllocationStatistics {
        Self::stats(self)
    }

    #[cfg(test)]
    fn reset(&self) {
        Self::reset(self)
    }
}

#[cfg(test)]
#[coverage(off)]
mod tests {
    extern crate std;
    use crate::{
        allocator::{
            DEFAULT_ALLOCATION_STRATEGY, DEFAULT_PAGE_ALLOCATION_GRANULARITY, HIGH_TRAFFIC_ALLOC_MIN_EXPANSION,
        },
        gcd, test_support,
    };
    use alloc::vec::Vec;
    use core::{alloc::GlobalAlloc, ffi::c_void, panic};
    use std::alloc::System;

    use patina::{
        base::{SIZE_64KB, UEFI_PAGE_SHIFT, UEFI_PAGE_SIZE},
        uefi_pages_to_size,
    };

    use super::*;

    fn init_gcd(gcd: &SpinLockedGcd, size: usize) -> u64 {
        // SAFETY: Resetting the test GCD is safe in a test.
        unsafe { gcd.reset() };

        gcd.init(48, 16);
        let layout = Layout::from_size_align(size, UEFI_PAGE_SIZE).unwrap();
        // SAFETY: System allocator is used with a valid layout for test memory backing.
        let base = unsafe { System.alloc(layout) as u64 };
        // SAFETY: init_memory_blocks is used with a valid backing allocation for tests.
        unsafe {
            gcd.init_memory_blocks(GcdMemoryType::SystemMemory, base as usize, size, efi::MEMORY_WB, efi::MEMORY_WB)
                .unwrap();
        }
        base
    }

    // this runs each test twice, once with 4KB page allocation granularity and once with 64KB page allocation
    // granularity. This is to ensure that the allocator works correctly with both page allocation granularities.
    fn with_granularity_modulation<F: Fn(usize) + std::panic::RefUnwindSafe>(f: F) {
        f(DEFAULT_PAGE_ALLOCATION_GRANULARITY);
        f(SIZE_64KB);
    }

    fn with_locked_state<F: Fn() + std::panic::RefUnwindSafe>(f: F) {
        test_support::with_global_lock(|| {
            test_support::init_test_logger();
            f();
        })
        .unwrap();
    }

    #[test]
    fn test_get_memory_ranges_returns_allocated_region() {
        with_granularity_modulation(|granularity| {
            with_locked_state(|| {
                static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

                init_gcd(&GCD, 0x400000);

                let mut fsb = FixedSizeBlockAllocator::new(efi::BOOT_SERVICES_DATA, granularity);

                assert_eq!(fsb.get_memory_ranges().count(), 0);

                let allocation_size = DEFAULT_PAGE_ALLOCATION_GRANULARITY;
                let allocated_address = GCD
                    .allocate_memory_space(
                        DEFAULT_ALLOCATION_STRATEGY,
                        GcdMemoryType::SystemMemory,
                        page_shift_from_alignment(granularity).unwrap(),
                        allocation_size,
                        DUMMY_HANDLE,
                        None,
                    )
                    .unwrap();

                fsb.expand(NonNull::slice_from_raw_parts(
                    NonNull::new(allocated_address as *mut u8).unwrap(),
                    allocation_size,
                ))
                .unwrap();

                let ranges: Vec<_> = fsb.get_memory_ranges().collect();
                assert_eq!(ranges.len(), 1);

                let expected_start = allocated_address + size_of::<AllocatorListNode>();
                let expected_end = expected_start + allocation_size - size_of::<AllocatorListNode>();
                assert_eq!(ranges[0], expected_start..expected_end);
            });
        });
    }
    const DUMMY_HANDLE: *mut c_void = 0xDEADBEEF as *mut c_void;

    #[test]
    fn allocate_deallocate_test() {
        with_granularity_modulation(|granularity| {
            with_locked_state(|| {
                // Create a static GCD for test.
                static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

                // Allocate some space on the heap with the global allocator (std) to be used by expand().
                init_gcd(&GCD, 0x400000);

                let fsb = SpinLockedFixedSizeBlockAllocator::new(
                    &GCD,
                    DUMMY_HANDLE,
                    efi::RUNTIME_SERVICES_DATA,
                    granularity,
                    DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                );

                let layout = Layout::from_size_align(0x8, 0x8).unwrap();
                let allocation = fsb.allocate(layout).unwrap().cast::<u8>();

                // SAFETY: Allocation was returned by fsb for this layout.
                unsafe { fsb.deallocate(allocation, layout) };

                let layout = Layout::from_size_align(0x20, 0x20).unwrap();
                let allocation = fsb.allocate(layout).unwrap().cast::<u8>();

                // SAFETY: Allocation was returned by fsb for this layout.
                unsafe { fsb.deallocate(allocation, layout) };
            });
        });
    }

    #[test]
    fn test_list_index() {
        let layout = Layout::from_size_align(8, 1).unwrap();
        assert_eq!(list_index(&layout), Some(0));

        let layout = Layout::from_size_align(12, 8).unwrap();
        assert_eq!(list_index(&layout), Some(1));

        let layout = Layout::from_size_align(8, 32).unwrap();
        assert_eq!(list_index(&layout), Some(2));

        let layout = Layout::from_size_align(4096, 32).unwrap();
        assert_eq!(list_index(&layout), Some(9));

        let layout = Layout::from_size_align(1, 4096).unwrap();
        assert_eq!(list_index(&layout), Some(9));

        let layout = Layout::from_size_align(8192, 1).unwrap();
        assert_eq!(list_index(&layout), None);
    }

    #[test]
    fn test_construct_empty_fixed_size_block_allocator() {
        with_locked_state(|| {
            let fsb = FixedSizeBlockAllocator::new(efi::BOOT_SERVICES_DATA, DEFAULT_PAGE_ALLOCATION_GRANULARITY);
            assert!(fsb.list_heads.iter().all(|x| x.is_none()));
            assert!(fsb.allocators.is_none());
        });
    }

    #[test]
    fn test_expand() {
        with_granularity_modulation(|granularity| {
            with_locked_state(|| {
                // Create a static GCD
                static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

                // Allocate some space on the heap with the global allocator (std) to be used by expand().
                let base = init_gcd(&GCD, 0x4000000);

                //verify no allocators exist before expand.
                let mut fsb = FixedSizeBlockAllocator::new(efi::RUNTIME_SERVICES_DATA, granularity);
                assert!(fsb.allocators.is_none());

                let allocation_size = DEFAULT_PAGE_ALLOCATION_GRANULARITY;

                // Allocate one page to expand by
                let allocated_address = GCD
                    .allocate_memory_space(
                        DEFAULT_ALLOCATION_STRATEGY,
                        GcdMemoryType::SystemMemory,
                        UEFI_PAGE_SHIFT,
                        allocation_size,
                        DUMMY_HANDLE,
                        None,
                    )
                    .unwrap();

                fsb.expand(NonNull::slice_from_raw_parts(
                    NonNull::new(allocated_address as *mut u8).unwrap(),
                    allocation_size,
                ))
                .unwrap();

                assert!(fsb.allocators.is_some());
                // SAFETY: fsb.allocators points to valid list nodes after expand.
                unsafe {
                    assert!((*fsb.allocators.unwrap()).next.is_none());
                    assert!((*fsb.allocators.unwrap()).allocator.bottom() as usize > base as usize);
                    assert_eq!(
                        (*fsb.allocators.unwrap()).allocator.free(),
                        allocation_size - size_of::<AllocatorListNode>()
                    );
                }

                //expand by larger than the minimum expansion.
                let allocation_size = DEFAULT_PAGE_ALLOCATION_GRANULARITY + 0x1000;

                // Allocate one page to expand by
                let allocated_address = GCD
                    .allocate_memory_space(
                        DEFAULT_ALLOCATION_STRATEGY,
                        GcdMemoryType::SystemMemory,
                        UEFI_PAGE_SHIFT,
                        allocation_size,
                        DUMMY_HANDLE,
                        None,
                    )
                    .unwrap();

                fsb.expand(NonNull::slice_from_raw_parts(
                    NonNull::new(allocated_address as *mut u8).unwrap(),
                    allocation_size,
                ))
                .unwrap();
                assert!(fsb.allocators.is_some());
                // SAFETY: fsb.allocators points to valid list nodes after expand.
                unsafe {
                    assert!((*fsb.allocators.unwrap()).next.is_some());
                    assert!((*(*fsb.allocators.unwrap()).next.unwrap()).next.is_none());
                    assert!((*fsb.allocators.unwrap()).allocator.bottom() as usize > base as usize);
                    assert_eq!(
                        (*fsb.allocators.unwrap()).allocator.free(),
                        //expected free: size + a page to hold allocator node - size of allocator node.
                        allocation_size - size_of::<AllocatorListNode>()
                    );
                }
            });
        });
    }

    #[test]
    fn test_allocation_iterator() {
        with_locked_state(|| {
            // Create a static GCD
            static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

            // Allocate some space on the heap with the global allocator (std) to be used by expand().
            init_gcd(&GCD, 0x800000);

            let mut fsb = FixedSizeBlockAllocator::new(efi::BOOT_SERVICES_DATA, DEFAULT_PAGE_ALLOCATION_GRANULARITY);

            const NUM_ALLOCATIONS: usize = 5;

            const ALLOCATION_SIZE: usize = DEFAULT_PAGE_ALLOCATION_GRANULARITY;
            for _ in 0..NUM_ALLOCATIONS {
                fsb.expand(NonNull::slice_from_raw_parts(
                    NonNull::new(
                        GCD.allocate_memory_space(
                            DEFAULT_ALLOCATION_STRATEGY,
                            GcdMemoryType::SystemMemory,
                            UEFI_PAGE_SHIFT,
                            ALLOCATION_SIZE,
                            DUMMY_HANDLE,
                            None,
                        )
                        .unwrap() as *mut u8,
                    )
                    .unwrap(),
                    ALLOCATION_SIZE,
                ))
                .unwrap();
            }

            assert_eq!(NUM_ALLOCATIONS, AllocatorIterator::new(fsb.allocators).count());
            let expected_free = ALLOCATION_SIZE - size_of::<AllocatorListNode>();
            assert!(AllocatorIterator::new(fsb.allocators).all(|node| {
                // SAFETY: node pointers come from AllocatorIterator over fsb.allocators.
                unsafe { (*node).allocator.free() == expected_free }
            }));
        });
    }

    #[test]
    fn test_fallback_alloc() {
        with_granularity_modulation(|granularity| {
            with_locked_state(|| {
                // Create a static GCD
                static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

                // Allocate some space on the heap with the global allocator (std) to be used by expand().
                let _ = init_gcd(&GCD, 0x400000);

                let mut fsb = FixedSizeBlockAllocator::new(efi::RUNTIME_SERVICES_DATA, granularity);

                // Test fallback_alloc with size < size_of::<AllocatorListNode>()
                let allocation_size = size_of::<AllocatorListNode>() / 2;
                let layout = Layout::from_size_align(allocation_size, 0x10).unwrap();
                match fsb.fallback_alloc(layout) {
                    Err(FixedSizeBlockAllocatorError::OutOfMemory(mem_req)) => {
                        assert!(
                            mem_req
                                >= layout.pad_to_align().size()
                                    + Layout::new::<AllocatorListNode>().pad_to_align().size(),
                            "fallback_alloc should request enough memory to fit aligned layout and an aligned AllocatorListNode"
                        );
                    }
                    _ => {
                        panic!(
                            "fallback_alloc with no allocators should return FixedSizeBlockAllocatorError::OutOfMemory"
                        )
                    }
                }
                assert!(fsb.allocators.is_none());

                // Test fallback_alloc with size > size_of::<AllocatorListNode>(), but unaligned to AllocatorListNode
                let allocation_size = size_of::<AllocatorListNode>() + align_of::<AllocatorListNode>() / 2;
                let layout = Layout::from_size_align(allocation_size, 0x10).unwrap();
                match fsb.fallback_alloc(layout) {
                    Err(FixedSizeBlockAllocatorError::OutOfMemory(mem_req)) => {
                        assert!(
                            mem_req
                                >= layout.pad_to_align().size()
                                    + Layout::new::<AllocatorListNode>().pad_to_align().size(),
                            "fallback_alloc should request enough memory to fit aligned layout and an aligned AllocatorListNode"
                        );
                    }
                    _ => {
                        panic!(
                            "fallback_alloc with no allocators should return FixedSizeBlockAllocatorError::OutOfMemory"
                        )
                    }
                }
                assert!(fsb.allocators.is_none());
            });
        });
    }

    #[test]
    fn test_alloc() {
        with_granularity_modulation(|granularity| {
            with_locked_state(|| {
                // Create a static GCD
                static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

                // Allocate some space on the heap with the global allocator (std) to be used by expand().
                let base = init_gcd(&GCD, 0x400000);

                let fsb = SpinLockedFixedSizeBlockAllocator::new(
                    &GCD,
                    1 as _,
                    efi::RUNTIME_SERVICES_DATA,
                    granularity,
                    DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                );

                let layout = Layout::from_size_align(0x1000, 0x10).unwrap();
                // SAFETY: fsb is initialized and used with a valid layout in tests.
                let allocation = unsafe { fsb.alloc(layout) };
                assert!(fsb.lock().allocators.is_some());
                assert!((allocation as u64) > base);
                assert!((allocation as u64) < base + 0x400000);
            });
        });
    }

    #[test]
    fn test_allocate() {
        with_granularity_modulation(|granularity| {
            with_locked_state(|| {
                // Create a static GCD
                static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

                // Allocate some space on the heap with the global allocator (std) to be used by expand().
                let base = init_gcd(&GCD, 0x400000);

                let fsb = SpinLockedFixedSizeBlockAllocator::new(
                    &GCD,
                    1 as _,
                    efi::RUNTIME_SERVICES_DATA,
                    granularity,
                    DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                );

                let layout = Layout::from_size_align(0x1000, 0x10).unwrap();
                let allocation = fsb.allocate(layout).unwrap().as_ptr() as *mut u8;
                assert!(fsb.lock().allocators.is_some());
                assert!((allocation as u64) > base);
                assert!((allocation as u64) < base + 0x400000);
            });
        });
    }

    #[test]
    fn test_fallback_dealloc() {
        with_granularity_modulation(|granularity| {
            with_locked_state(|| {
                // Create a static GCD
                static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

                // Allocate some space on the heap with the global allocator (std) to be used by expand().
                init_gcd(&GCD, 0x400000);

                let mut fsb = FixedSizeBlockAllocator::new(efi::RUNTIME_SERVICES_DATA, granularity);

                let layout = Layout::from_size_align(0x8, 0x8).unwrap();

                // Expand the FSB by the minimum expansion to fit the allocation
                let expansion_size = DEFAULT_PAGE_ALLOCATION_GRANULARITY;
                fsb.expand(NonNull::slice_from_raw_parts(
                    NonNull::new(
                        GCD.allocate_memory_space(
                            DEFAULT_ALLOCATION_STRATEGY,
                            GcdMemoryType::SystemMemory,
                            UEFI_PAGE_SHIFT,
                            expansion_size,
                            DUMMY_HANDLE,
                            None,
                        )
                        .unwrap() as *mut u8,
                    )
                    .unwrap(),
                    expansion_size,
                ))
                .unwrap();

                let allocation = fsb.fallback_alloc(layout).unwrap();

                // Finally, we can test fallback_dealloc
                fsb.fallback_dealloc(allocation.cast(), layout);
                let expected_free = DEFAULT_PAGE_ALLOCATION_GRANULARITY - size_of::<AllocatorListNode>();
                // SAFETY: fsb.allocators points to a valid allocator after expand.
                unsafe {
                    assert_eq!((*fsb.allocators.unwrap()).allocator.free(), expected_free);
                }
            });
        });
    }

    #[test]
    fn test_dealloc() {
        with_granularity_modulation(|granularity| {
            with_locked_state(|| {
                // Create a static GCD
                static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

                // Allocate some space on the heap with the global allocator (std) to be used by expand().
                init_gcd(&GCD, 0x400000);

                let fsb = SpinLockedFixedSizeBlockAllocator::new(
                    &GCD,
                    1 as _,
                    efi::RUNTIME_SERVICES_DATA,
                    granularity,
                    DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                );

                let layout = Layout::from_size_align(0x8, 0x8).unwrap();
                // SAFETY: fsb is initialized and used with a valid layout in tests.
                let allocation = unsafe { fsb.alloc(layout) };

                // SAFETY: Allocation was returned by fsb for this layout.
                unsafe { fsb.dealloc(allocation, layout) };
                let free_block_ptr = fsb.lock().list_heads[list_index(&layout).unwrap()].take().unwrap()
                    as *mut BlockListNode as *mut u8;
                assert_eq!(free_block_ptr, allocation);

                let layout = Layout::from_size_align(0x20, 0x20).unwrap();
                // SAFETY: fsb is initialized and used with a valid layout in tests.
                let allocation = unsafe { fsb.alloc(layout) };

                // SAFETY: Allocation was returned by fsb for this layout.
                unsafe { fsb.dealloc(allocation, layout) };
                let free_block_ptr = fsb.lock().list_heads[list_index(&layout).unwrap()].take().unwrap()
                    as *mut BlockListNode as *mut u8;
                assert_eq!(free_block_ptr, allocation);
            });
        });
    }

    #[test]
    fn test_deallocate() {
        with_granularity_modulation(|granularity| {
            with_locked_state(|| {
                // Create a static GCD
                static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

                // Allocate some space on the heap with the global allocator (std) to be used by expand().
                init_gcd(&GCD, 0x400000);

                let fsb = SpinLockedFixedSizeBlockAllocator::new(
                    &GCD,
                    1 as _,
                    efi::RUNTIME_SERVICES_DATA,
                    granularity,
                    DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                );

                let layout = Layout::from_size_align(0x8, 0x8).unwrap();
                let allocation = fsb.allocate(layout).unwrap().cast::<u8>();
                let allocation_ptr = allocation.as_ptr();

                // SAFETY: Allocation was returned by fsb for this layout.
                unsafe { fsb.deallocate(allocation, layout) };
                let free_block_ptr = fsb.lock().list_heads[list_index(&layout).unwrap()].take().unwrap()
                    as *mut BlockListNode as *mut u8;
                assert_eq!(free_block_ptr, allocation_ptr);

                let layout = Layout::from_size_align(0x20, 0x20).unwrap();
                let allocation = fsb.allocate(layout).unwrap().cast::<u8>();
                let allocation_ptr = allocation.as_ptr();

                // SAFETY: Allocation was returned by fsb for this layout.
                unsafe { fsb.deallocate(allocation, layout) };
                let free_block_ptr = fsb.lock().list_heads[list_index(&layout).unwrap()].take().unwrap()
                    as *mut BlockListNode as *mut u8;
                assert_eq!(free_block_ptr, allocation_ptr);
            });
        });
    }

    #[test]
    fn test_contains() {
        with_locked_state(|| {
            // Create a static GCD
            static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

            // Allocate some space on the heap with the global allocator (std) to be used by expand().
            init_gcd(&GCD, 0x400000);

            let fsb = SpinLockedFixedSizeBlockAllocator::new(
                &GCD,
                1 as _,
                efi::BOOT_SERVICES_DATA,
                DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                DEFAULT_PAGE_ALLOCATION_GRANULARITY,
            );

            let layout = Layout::from_size_align(0x8, 0x8).unwrap();
            let allocation = fsb.allocate(layout).unwrap().cast::<u8>();
            assert!(fsb.contains(allocation));
        });
    }

    #[test]
    fn test_allocate_pages() {
        with_granularity_modulation(|granularity| {
            with_locked_state(|| {
                // Create a static GCD
                static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

                // Allocate some space on the heap with the global allocator (std) to back the test GCD.
                let address = init_gcd(&GCD, 0x1000000);

                let fsb = SpinLockedFixedSizeBlockAllocator::new(
                    &GCD,
                    1 as _,
                    efi::RUNTIME_SERVICES_DATA,
                    granularity,
                    DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                );

                let pages = 4;

                let allocation =
                    fsb.allocate_pages(gcd::AllocateType::BottomUp(None), pages, UEFI_PAGE_SIZE).unwrap().cast::<u8>();

                assert!(allocation.as_ptr() as u64 >= address);
                assert!((allocation.as_ptr() as u64) < address + 0x1000000);

                // SAFETY: free_pages uses a valid test allocation pointer and page count.
                unsafe {
                    match fsb.free_pages(0, pages) {
                        Err(EfiError::NotFound) => {}
                        _ => panic!("Expected NOT_FOUND"),
                    };
                };

                // SAFETY: allocation and page count come from allocate_pages in this test.
                unsafe {
                    fsb.free_pages(allocation.as_ptr() as usize, pages).unwrap();
                };
            });
        });
    }

    #[test]
    fn test_allocate_at_address() {
        with_granularity_modulation(|granularity| {
            with_locked_state(|| {
                // Create a static GCD
                static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

                // Allocate some space on the heap with the global allocator (std) to back the test GCD.
                let address = init_gcd(&GCD, 0x1000000);

                let fsb = SpinLockedFixedSizeBlockAllocator::new(
                    &GCD,
                    1 as _,
                    efi::RUNTIME_SERVICES_DATA,
                    granularity,
                    DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                );

                let target_address = (address + 0x400000 - max(8_u64 * ALIGNMENT as u64, granularity as u64))
                    & (!(granularity as u64 - 1_u64));
                let pages = 4;

                let allocation = fsb
                    .allocate_pages(gcd::AllocateType::Address(target_address as usize), pages, UEFI_PAGE_SIZE)
                    .unwrap()
                    .cast::<u8>();

                assert_eq!(allocation.as_ptr() as u64, target_address);

                // SAFETY: free_pages uses a valid test allocation pointer and page count.
                unsafe {
                    fsb.free_pages(allocation.as_ptr() as usize, pages).unwrap();
                };
            });
        });
    }

    #[test]
    fn test_allocate_below_address_bottom_up() {
        with_granularity_modulation(|granularity| {
            with_locked_state(|| {
                // Create a static GCD
                static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

                // Allocate some space on the heap with the global allocator (std) to back the test GCD.
                let address = init_gcd(&GCD, 0x1000000);

                let fsb = SpinLockedFixedSizeBlockAllocator::new(
                    &GCD,
                    1 as _,
                    efi::RUNTIME_SERVICES_DATA,
                    granularity,
                    DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                );

                let target_address = address + 0x400000 - 8 * (ALIGNMENT as u64);
                let pages = 4;

                let allocation = fsb
                    .allocate_pages(gcd::AllocateType::BottomUp(Some(target_address as usize)), pages, UEFI_PAGE_SIZE)
                    .unwrap()
                    .cast::<u8>();
                assert!((allocation.as_ptr() as u64) < target_address);

                // SAFETY: free_pages uses a valid test allocation pointer and page count.
                unsafe {
                    fsb.free_pages(allocation.as_ptr() as usize, pages).unwrap();
                };
            });
        });
    }

    #[test]
    fn test_allocate_below_address_top_down() {
        with_granularity_modulation(|granularity| {
            with_locked_state(|| {
                // Create a static GCD
                static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

                // Allocate some space on the heap with the global allocator (std) to back the test GCD.
                let address = init_gcd(&GCD, 0x1000000);

                let fsb = SpinLockedFixedSizeBlockAllocator::new(
                    &GCD,
                    1 as _,
                    efi::RUNTIME_SERVICES_DATA,
                    granularity,
                    DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                );

                let target_address = address + 0x400000 - 8 * (ALIGNMENT as u64);
                let pages = 4;

                let allocation = fsb
                    .allocate_pages(gcd::AllocateType::TopDown(Some(target_address as usize)), pages, UEFI_PAGE_SIZE)
                    .unwrap()
                    .cast::<u8>();
                assert!((allocation.as_ptr() as usize + uefi_pages_to_size!(pages)) <= target_address as usize);

                // SAFETY: allocation and page count come from allocate_pages in this test.
                unsafe {
                    fsb.free_pages(allocation.as_ptr() as usize, pages).unwrap();
                };
            });
        });
    }

    #[test]
    fn test_allocator_commands_with_invalid_parameters() {
        with_locked_state(|| {
            // Create a static GCD
            static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

            // Allocate some space on the heap with the global allocator (std) to be used by expand().
            let _ = init_gcd(&GCD, 0x400000);

            // Test commands with bad handle.
            let fsb = SpinLockedFixedSizeBlockAllocator::new(
                &GCD,
                0 as _,
                efi::BOOT_SERVICES_DATA,
                DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                DEFAULT_PAGE_ALLOCATION_GRANULARITY,
            );
            match fsb.allocate_pages(AllocationStrategy::Address(0x1000), 5, UEFI_PAGE_SIZE) {
                Err(EfiError::InvalidParameter) => {}
                _ => panic!("Expected INVALID_PARAMETER"),
            }

            let fsb = SpinLockedFixedSizeBlockAllocator::new(
                &GCD,
                1 as _,
                efi::BOOT_SERVICES_DATA,
                DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                DEFAULT_PAGE_ALLOCATION_GRANULARITY,
            );

            let allocation_strategy = AllocationStrategy::Address(0x1000);
            match fsb.allocate_pages(allocation_strategy, 5, UEFI_PAGE_SIZE) {
                Err(EfiError::NotFound) => {}
                _ => panic!("Expected NOT_FOUND"),
            }
            // Test invalid alignment
            let allocation_strategy = AllocationStrategy::Address(0x1001);
            match fsb.allocate_pages(allocation_strategy, 5, UEFI_PAGE_SIZE) {
                Err(EfiError::InvalidParameter) => {}
                _ => panic!("Expected INVALID_PARAMETER"),
            }

            // SAFETY: Invalid parameters are intentionally tested (unit test).
            unsafe {
                match fsb.free_pages(0x1001, 5) {
                    Err(EfiError::InvalidParameter) => {}
                    _ => panic!("Expected INVALID_PARAMETER"),
                }
            }
        });
    }

    #[test]
    fn validate_fsb_display_impl_does_not_panic() {
        with_locked_state(|| {
            // Create a static GCD
            static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

            // Allocate some space on the heap with the global allocator (std) to be used by expand().
            let _ = init_gcd(&GCD, 0x400000);

            let fsb = SpinLockedFixedSizeBlockAllocator::new(
                &GCD,
                1 as _,
                efi::BOOT_SERVICES_DATA,
                DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                DEFAULT_PAGE_ALLOCATION_GRANULARITY,
            );
            fsb.allocate_pages(DEFAULT_ALLOCATION_STRATEGY, 5, UEFI_PAGE_SIZE).unwrap();

            let layout = Layout::from_size_align(0x1000, 0x10).unwrap();
            let _ = fsb.allocate(layout); // Triggers expansion + allocation

            // Call format on the inner FixedSizeBlockAllocator
            let _ = std::format!("{}", fsb.lock());
        });
    }

    #[test]
    fn test_allocation_stats() {
        with_locked_state(|| {
            // Create a static GCD
            static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

            const TEST_MIN_EXPANSION_SIZE: usize = HIGH_TRAFFIC_ALLOC_MIN_EXPANSION;

            // Allocate some space on the heap with the global allocator (std) to be used by expand().
            let _ = init_gcd(&GCD, 0x1000000);

            // Make a fixed-sized-block allocator
            let fsb = SpinLockedFixedSizeBlockAllocator::new(
                &GCD,
                1 as _,
                efi::BOOT_SERVICES_DATA,
                DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                HIGH_TRAFFIC_ALLOC_MIN_EXPANSION,
            );

            let stats = fsb.stats();
            assert_eq!(stats.pool_allocation_calls, 0);
            assert_eq!(stats.pool_free_calls, 0);
            assert_eq!(stats.page_allocation_calls, 0);
            assert_eq!(stats.page_free_calls, 0);
            assert_eq!(stats.reserved_size, 0);
            assert_eq!(stats.reserved_used, 0);
            assert_eq!(stats.claimed_pages, 0);

            //test alloc/deallocate and stats
            // SAFETY: fsb is initialized and used with a valid layout for testing.
            let ptr = unsafe {
                fsb.alloc(
                    Layout::from_size_align(TEST_MIN_EXPANSION_SIZE - size_of::<AllocatorListNode>(), 0x8).unwrap(),
                )
            };

            let stats = fsb.stats();
            //an additional allocation call will be made as the first one will fail due to a lack of memory
            assert_eq!(stats.pool_allocation_calls, 2);
            assert_eq!(stats.pool_free_calls, 0);
            assert_eq!(stats.page_allocation_calls, 0);
            assert_eq!(stats.page_free_calls, 0);
            assert_eq!(stats.reserved_size, 0);
            assert_eq!(stats.reserved_used, 0);
            let initial_claimed = stats.claimed_pages;

            // SAFETY: Allocation was returned by fsb for this layout.
            unsafe {
                fsb.dealloc(ptr, Layout::from_size_align(0x100, 0x8).unwrap());
            }

            let stats = fsb.stats();
            //an additional allocation call will be made as the first one will fail due to a lack of memory
            assert_eq!(stats.pool_allocation_calls, 2);
            assert_eq!(stats.pool_free_calls, 1);
            assert_eq!(stats.page_allocation_calls, 0);
            assert_eq!(stats.page_free_calls, 0);
            assert_eq!(stats.reserved_size, 0);
            assert_eq!(stats.reserved_used, 0);
            assert_eq!(stats.claimed_pages, initial_claimed);

            //test alloc/deallocate and stats blowing the bucket
            // SAFETY: fsb is initialized and used with a valid layout in tests.
            let ptr = unsafe { fsb.alloc(Layout::from_size_align(TEST_MIN_EXPANSION_SIZE * 3, 0x8).unwrap()) };

            //after this allocate, the basic memory map of the FSB should look like:
            //1MB range as a result of previous pool allocation expand - available for pool allocation.
            //    Claims first 1MB of 2MB reserved region.
            //1MB free but owned by the allocator (not pool) as a result of 2MB reservation.
            //3MB+1 page range as a result of 3MB allocation + 1 page to hold allocator node.

            let stats = fsb.stats();
            //an additional allocation call will be made as the first one will fail due to a lack of memory
            assert_eq!(stats.pool_allocation_calls, 4);
            assert_eq!(stats.pool_free_calls, 1);
            assert_eq!(stats.page_allocation_calls, 0);
            assert_eq!(stats.page_free_calls, 0);
            assert_eq!(stats.reserved_size, 0);
            assert_eq!(stats.reserved_used, 0);
            let claimed_after_3mb = stats.claimed_pages;

            // SAFETY: Allocation was returned by fsb for this layout.
            unsafe {
                fsb.dealloc(ptr, Layout::from_size_align(TEST_MIN_EXPANSION_SIZE * 3, 0x8).unwrap());
            }

            //after this free, the basic memory map of the FSB should look like:
            //1MB range as a result of previous pool allocation expand - available for pool allocation.
            //    Claims first 1MB of 2MB reserved region.
            //1MB free but owned by the allocator (not pool) as a result of 2MB reservation.
            //3MB+1 page range as a result of 3MB allocation + 1 page to hold allocator node - available for pool allocation.

            let stats = fsb.stats();
            assert_eq!(stats.pool_allocation_calls, 4);
            assert_eq!(stats.pool_free_calls, 2);
            assert_eq!(stats.page_allocation_calls, 0);
            assert_eq!(stats.page_free_calls, 0);
            assert_eq!(stats.reserved_size, 0);
            assert_eq!(stats.reserved_used, 0);
            assert_eq!(stats.claimed_pages, claimed_after_3mb);

            // test that a small page allocation fits in the 1MB free reserved region.
            let ptr = fsb.allocate_pages(DEFAULT_ALLOCATION_STRATEGY, 0x4, UEFI_PAGE_SIZE).unwrap().as_ptr();

            //after this allocate_pages, the basic memory map of the FSB should look like:
            //1MB range as a result of previous pool allocation expand - available for pool allocation.
            //    Claims first 1MB of 2MB reserved region.
            //16K allocated.
            //1MB-16k free but owned by the allocator (not pool) as a result of 2MB reservation.
            //3MB+1 page range as a result of 3MB allocation + 1 page to hold allocator node - available for pool allocation.

            let stats = fsb.stats();
            assert_eq!(stats.pool_allocation_calls, 4);
            assert_eq!(stats.pool_free_calls, 2);
            assert_eq!(stats.page_allocation_calls, 1);
            assert_eq!(stats.page_free_calls, 0);
            assert_eq!(stats.reserved_size, 0);
            assert_eq!(stats.reserved_used, 0);
            assert_eq!(stats.claimed_pages, claimed_after_3mb + 0x4);

            // SAFETY: free_pages uses a valid test allocation pointer and page count.
            unsafe {
                fsb.free_pages(ptr as *mut u8 as usize, 0x4).unwrap();
            }

            //after this free, the basic memory map of the FSB should look like:
            //1MB range as a result of previous pool allocation expand - available for pool allocation.
            //    Claims first 1MB of 2MB reserved region.
            //1MB free but owned by the allocator (not pool) as a result of 2MB reservation.
            //3MB+1 page range as a result of 3MB allocation + 1 page to hold allocator node - available for pool allocation.

            let stats = fsb.stats();
            assert_eq!(stats.pool_allocation_calls, 4);
            assert_eq!(stats.pool_free_calls, 2);
            assert_eq!(stats.page_allocation_calls, 1);
            assert_eq!(stats.page_free_calls, 1);
            assert_eq!(stats.reserved_size, 0);
            assert_eq!(stats.reserved_used, 0);
            assert_eq!(stats.claimed_pages, claimed_after_3mb);

            //test that a lage page allocation results in more claimed pages.
            let ptr = fsb.allocate_pages(DEFAULT_ALLOCATION_STRATEGY, 0x104, UEFI_PAGE_SIZE).unwrap().as_ptr();

            //after this allocate_pages, the basic memory map of the FSB should look like:
            //1MB range as a result of previous pool allocation expand - available for pool allocation.
            //    Claims first 1MB of 2MB reserved region.
            //1MB free but owned by the allocator (not pool) as a result of 2MB reservation.
            //3MB+1 page range as a result of 3MB allocation + 1 page to hold allocator node - available for pool allocation.
            //104 pages (1MB+16K) page as a result of allocation.

            let stats = fsb.stats();
            assert_eq!(stats.pool_allocation_calls, 4);
            assert_eq!(stats.pool_free_calls, 2);
            assert_eq!(stats.page_allocation_calls, 2);
            assert_eq!(stats.page_free_calls, 1);
            assert_eq!(stats.reserved_size, 0);
            assert_eq!(stats.reserved_used, 0);
            assert_eq!(stats.claimed_pages, claimed_after_3mb + 0x104);

            // test that a small page allocation fits in the 1MB free reserved region.
            let ptr1 = fsb.allocate_pages(DEFAULT_ALLOCATION_STRATEGY, 0x4, UEFI_PAGE_SIZE).unwrap().as_ptr();

            //after this allocate_pages, the basic memory map of the FSB should look like:
            //1MB range as a result of previous pool allocation expand - available for pool allocation.
            //    Claims first 1MB of 2MB reserved region.
            //16K allocated.
            //1MB-16k free but owned by the allocator (not pool) as a result of 2MB reservation.
            //3MB+1 page range as a result of 3MB allocation + 1 page to hold allocator node - available for pool allocation.
            //104 pages (1MB+16K) page as a result of allocation.

            let stats = fsb.stats();
            assert_eq!(stats.pool_allocation_calls, 4);
            assert_eq!(stats.pool_free_calls, 2);
            assert_eq!(stats.page_allocation_calls, 3);
            assert_eq!(stats.page_free_calls, 1);
            assert_eq!(stats.reserved_size, 0);
            assert_eq!(stats.reserved_used, 0);
            assert_eq!(stats.claimed_pages, claimed_after_3mb + 0x104 + 0x4);

            // SAFETY: free_pages uses a valid test allocation pointer and page count.
            unsafe {
                fsb.free_pages(ptr1 as *mut u8 as usize, 0x4).unwrap();
            }
            // SAFETY: free_pages uses a valid test allocation pointer and page count.
            unsafe {
                fsb.free_pages(ptr as *mut u8 as usize, 0x104).unwrap();
            }

            //after this free, the basic memory map of the FSB should look like:
            //1MB range as a result of previous pool allocation expand - available for pool allocation.
            //    Claims first 1MB of 2MB reserved region.
            //1MB free but owned by the allocator (not pool) as a result of 2MB reservation.
            //3MB+1 page range as a result of 3MB allocation + 1 page to hold allocator node - available for pool allocation.

            let stats = fsb.stats();
            assert_eq!(stats.pool_allocation_calls, 4);
            assert_eq!(stats.pool_free_calls, 2);
            assert_eq!(stats.page_allocation_calls, 3);
            assert_eq!(stats.page_free_calls, 3);
            assert_eq!(stats.reserved_size, 0);
            assert_eq!(stats.reserved_used, 0);
            assert_eq!(stats.claimed_pages, claimed_after_3mb);
        });
    }

    #[test]
    fn test_get_memory_ranges() {
        with_granularity_modulation(|granularity| {
            with_locked_state(|| {
                // Create a static GCD
                static GCD: SpinLockedGcd = SpinLockedGcd::new(None);

                // Allocate some space on the heap with the global allocator (std) to be used by expand().
                let base = init_gcd(&GCD, 0x400000);

                let mut fsb = FixedSizeBlockAllocator::new(efi::RUNTIME_SERVICES_DATA, granularity);

                const NUM_ALLOCATIONS: usize = 3;

                // Expand the FSB multiple times
                for _ in 0..NUM_ALLOCATIONS {
                    fsb.expand(NonNull::slice_from_raw_parts(
                        NonNull::new(
                            GCD.allocate_memory_space(
                                DEFAULT_ALLOCATION_STRATEGY,
                                GcdMemoryType::SystemMemory,
                                UEFI_PAGE_SHIFT,
                                DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                                DUMMY_HANDLE,
                                None,
                            )
                            .unwrap() as *mut u8,
                        )
                        .unwrap(),
                        DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                    ))
                    .unwrap();
                }

                // Collect the memory ranges reported by the allocator
                let memory_ranges: Vec<_> = fsb.get_memory_ranges().collect();

                // Verify that the reported ranges match the expected ranges
                assert_eq!(memory_ranges.len(), NUM_ALLOCATIONS);
                for range in &memory_ranges {
                    assert!(range.start >= base as usize);
                    assert!(range.end <= (base + 0x400000) as usize);
                    assert!(range.start < range.end);
                }

                // Ensure that the ranges do not overlap
                for i in 0..memory_ranges.len() {
                    for j in i + 1..memory_ranges.len() {
                        assert!(
                            memory_ranges[i].end <= memory_ranges[j].start
                                || memory_ranges[j].end <= memory_ranges[i].start
                        );
                    }
                }
            });
        });
    }

    #[test]
    fn test_page_shift_from_alignment() {
        #[derive(Debug)]
        struct TestConfig {
            alignment: usize,
            expected: Result<usize, EfiError>,
        }

        let configs = [
            TestConfig { alignment: 0x1000, expected: Ok(12) },
            TestConfig { alignment: 0x2000, expected: Ok(13) },
            TestConfig { alignment: 0x400000, expected: Ok(22) },
            TestConfig { alignment: 0x6000, expected: Err(EfiError::InvalidParameter) },
            TestConfig { alignment: 0x800, expected: Err(EfiError::InvalidParameter) },
            TestConfig { alignment: 0, expected: Err(EfiError::InvalidParameter) },
        ];

        for config in configs {
            let result = page_shift_from_alignment(config.alignment);
            assert_eq!(result, config.expected, "Test config: {config:?}");
        }
    }
}
