Add PEI memory bin support [Rebase & FF]

docs/src/SUMMARY.md

@@ -68,6 +68,7 @@
   - [Event, Timer, and Task Priority](dxe_core/events.md)
   - [Image Loading and Execution](dxe_core/images.md)
   - [Memory Management](dxe_core/memory_management.md)
+  - [Memory Bins](dxe_core/memory_bins.md)
   - [Protocol Database](dxe_core/protocol_database.md)
   - [Synchronization](dxe_core/synchronization.md)
   - [Testing](dxe_core/testing.md)

patina_dxe_core/src/allocator/uefi_allocator.rs

@@ -70,23 +70,10 @@ where
         self.memory_type
     }
 
-    /// Reserves a range of memory to be used by this allocator of the given size in pages.
-    ///
-    /// The caller specifies a maximum number of pages this allocator is expected to require, and as long as the number
-    /// of pages actually used by the allocator is less than that amount, then all the allocations for this allocator
-    /// will be in a single contiguous block. This capability can be used to ensure that the memory map presented to the
-    /// OS is stable from boot-to-boot despite small boot-to-boot variations in actual page usage.
-    ///
-    /// For best memory stability, this routine should be called only during the initialization of the memory subsystem;
-    /// calling it after other allocations/frees have occurred will not cause allocation errors, but may cause the
-    /// memory map to vary from boot-to-boot.
-    ///
-    /// This routine will return Err(efi::Status::ALREADY_STARTED) if it is called more than once.
-    ///
-    pub fn reserve_memory_pages(&self, pages: usize) -> Result<(), EfiError> {
-        self.allocator.reserve_memory_pages(pages)
+    /// Sets the reserved memory range (bin range) for this allocator.
+    pub fn set_reserved_range(&self, range: Range<efi::PhysicalAddress>) {
+        self.allocator.set_reserved_range(range);
     }
-
     /// Returns an iterator over the memory ranges managed by this allocator.
     /// Returns an empty iterator if the allocator has no memory ranges.
     pub(crate) fn get_memory_ranges(&self) -> impl Iterator<Item = Range<efi::PhysicalAddress>> {
@@ -286,8 +273,9 @@ mod tests {
     use std::alloc::{GlobalAlloc, System};
 
     use patina::{
-        base::{SIZE_4KB, SIZE_64KB, UEFI_PAGE_SIZE},
+        base::{SIZE_4KB, SIZE_64KB, UEFI_PAGE_SIZE, align_up, page_shift_from_alignment},
         pi::dxe_services,
+        uefi_pages_to_size, uefi_size_to_pages,
     };
 
     use crate::{
@@ -346,7 +334,7 @@ mod tests {
             let fsb = SpinLockedFixedSizeBlockAllocator::new(
                 &GCD,
                 1 as _,
-                NonNull::from_ref(GCD.memory_type_info(efi::BOOT_SERVICES_DATA)),
+                efi::BOOT_SERVICES_DATA,
                 DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                 HIGH_TRAFFIC_ALLOC_MIN_EXPANSION,
             );
@@ -366,7 +354,7 @@ mod tests {
                 let fsb = SpinLockedFixedSizeBlockAllocator::new(
                     &GCD,
                     1 as _,
-                    NonNull::from_ref(GCD.memory_type_info(efi::RUNTIME_SERVICES_DATA)),
+                    efi::RUNTIME_SERVICES_DATA,
                     granularity,
                     LOW_TRAFFIC_RUNTIME_ALLOC_MIN_EXPANSION,
                 );
@@ -408,7 +396,7 @@ mod tests {
                 let fsb = SpinLockedFixedSizeBlockAllocator::new(
                     &GCD,
                     1 as _,
-                    NonNull::from_ref(GCD.memory_type_info(efi::RUNTIME_SERVICES_DATA)),
+                    efi::RUNTIME_SERVICES_DATA,
                     granularity,
                     LOW_TRAFFIC_RUNTIME_ALLOC_MIN_EXPANSION,
                 );
@@ -456,7 +444,7 @@ mod tests {
                 let fsb = SpinLockedFixedSizeBlockAllocator::new(
                     &GCD,
                     1 as _,
-                    NonNull::from_ref(GCD.memory_type_info(efi::RUNTIME_SERVICES_DATA)),
+                    efi::RUNTIME_SERVICES_DATA,
                     granularity,
                     LOW_TRAFFIC_RUNTIME_ALLOC_MIN_EXPANSION,
                 );
@@ -499,7 +487,7 @@ mod tests {
                 let fsb = SpinLockedFixedSizeBlockAllocator::new(
                     &GCD,
                     1 as _,
-                    NonNull::from_ref(GCD.memory_type_info(efi::RUNTIME_SERVICES_DATA)),
+                    efi::RUNTIME_SERVICES_DATA,
                     granularity,
                     LOW_TRAFFIC_RUNTIME_ALLOC_MIN_EXPANSION,
                 );
@@ -541,7 +529,7 @@ mod tests {
             let bs_fsb = SpinLockedFixedSizeBlockAllocator::new(
                 &GCD,
                 1 as _,
-                NonNull::from_ref(GCD.memory_type_info(efi::BOOT_SERVICES_DATA)),
+                efi::BOOT_SERVICES_DATA,
                 DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                 HIGH_TRAFFIC_ALLOC_MIN_EXPANSION,
             );
@@ -550,7 +538,7 @@ mod tests {
             let bc_fsb = SpinLockedFixedSizeBlockAllocator::new(
                 &GCD,
                 2 as _,
-                NonNull::from_ref(GCD.memory_type_info(efi::BOOT_SERVICES_CODE)),
+                efi::BOOT_SERVICES_CODE,
                 DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                 LOW_TRAFFIC_ALLOC_MIN_EXPANSION,
             );
@@ -583,7 +571,7 @@ mod tests {
                 let fsb = SpinLockedFixedSizeBlockAllocator::new(
                     &GCD,
                     1 as _,
-                    NonNull::from_ref(GCD.memory_type_info(efi::RUNTIME_SERVICES_DATA)),
+                    efi::RUNTIME_SERVICES_DATA,
                     granularity,
                     LOW_TRAFFIC_RUNTIME_ALLOC_MIN_EXPANSION,
                 );
@@ -617,7 +605,7 @@ mod tests {
             let fsb = SpinLockedFixedSizeBlockAllocator::new(
                 &GCD,
                 1 as _,
-                NonNull::from_ref(GCD.memory_type_info(efi::BOOT_SERVICES_DATA)),
+                efi::BOOT_SERVICES_DATA,
                 DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                 HIGH_TRAFFIC_ALLOC_MIN_EXPANSION,
             );
@@ -641,7 +629,7 @@ mod tests {
             let fsb = SpinLockedFixedSizeBlockAllocator::new(
                 &GCD,
                 1 as _,
-                NonNull::from_ref(GCD.memory_type_info(efi::BOOT_SERVICES_DATA)),
+                efi::BOOT_SERVICES_DATA,
                 DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                 HIGH_TRAFFIC_ALLOC_MIN_EXPANSION,
             );
@@ -669,6 +657,31 @@ mod tests {
         });
     }
 
+    /// Allocates a GCD region with ownership preservation and sets it as the allocator's reserved range.
+    fn setup_reserved_range(
+        gcd: &SpinLockedGcd,
+        allocator: &UefiAllocator<SpinLockedFixedSizeBlockAllocator>,
+        pages: usize,
+        granularity: usize,
+    ) -> Range<efi::PhysicalAddress> {
+        let required_pages = align_up(pages, uefi_size_to_pages!(granularity)).unwrap();
+        let size = uefi_pages_to_size!(required_pages);
+        let addr = gcd
+            .allocate_memory_space(
+                DEFAULT_ALLOCATION_STRATEGY,
+                dxe_services::GcdMemoryType::SystemMemory,
+                page_shift_from_alignment(granularity).unwrap(),
+                size,
+                allocator.handle(),
+                None,
+            )
+            .unwrap();
+        gcd.free_memory_space_preserving_ownership(addr, size).unwrap();
+        let range = addr as efi::PhysicalAddress..(addr + size) as efi::PhysicalAddress;
+        allocator.set_reserved_range(range.clone());
+        range
+    }
+
     #[test]
     fn reserve_memory_pages_reserves_the_pages() {
         with_granularity_modulation(|granularity| {
@@ -681,17 +694,17 @@ mod tests {
                 let reserved_fsb = SpinLockedFixedSizeBlockAllocator::new(
                     &GCD,
                     1 as _,
-                    NonNull::from_ref(GCD.memory_type_info(efi::RUNTIME_SERVICES_DATA)),
+                    efi::RUNTIME_SERVICES_DATA,
                     granularity,
                     LOW_TRAFFIC_RUNTIME_ALLOC_MIN_EXPANSION,
                 );
                 let reserved_allocator = UefiAllocator::new(reserved_fsb, efi::RUNTIME_SERVICES_DATA);
-                reserved_allocator.reserve_memory_pages(0x100).unwrap();
+                setup_reserved_range(&GCD, &reserved_allocator, 0x100, granularity);
 
                 let unreserved_fsb = SpinLockedFixedSizeBlockAllocator::new(
                     &GCD,
                     2 as _,
-                    NonNull::from_ref(GCD.memory_type_info(efi::LOADER_DATA)),
+                    efi::LOADER_DATA,
                     DEFAULT_PAGE_ALLOCATION_GRANULARITY,
                     LOW_TRAFFIC_ALLOC_MIN_EXPANSION,
                 );
@@ -783,4 +796,46 @@ mod tests {
             });
         });
     }
+
+    #[test]
+    fn allocate_max_address_prefers_reserved_range() {
+        with_granularity_modulation(|granularity| {
+            with_locked_state(|| {
+                static GCD: SpinLockedGcd = SpinLockedGcd::new(None);
+
+                let base = init_gcd(&GCD, 0x400000);
+                let gcd_end = base + 0x400000;
+
+                let reserved_fsb = SpinLockedFixedSizeBlockAllocator::new(
+                    &GCD,
+                    1 as _,
+                    efi::RUNTIME_SERVICES_DATA,
+                    granularity,
+                    LOW_TRAFFIC_RUNTIME_ALLOC_MIN_EXPANSION,
+                );
+                let reserved_allocator = UefiAllocator::new(reserved_fsb, efi::RUNTIME_SERVICES_DATA);
+                let reserved_range = setup_reserved_range(&GCD, &reserved_allocator, 0x100, granularity);
+
+                // TopDown(Some(max)) where max is above the reserved range should land in the bin.
+                let page = reserved_allocator
+                    .allocate_pages(AllocationStrategy::TopDown(Some(gcd_end as usize)), 1, UEFI_PAGE_SIZE)
+                    .unwrap();
+                let page_addr = page.as_ptr() as *mut u8 as u64;
+                assert!(
+                    reserved_range.contains(&page_addr),
+                    "TopDown(Some(gcd_end)) should land in the bin: addr={page_addr:#x}, range={reserved_range:#x?}",
+                );
+
+                // TopDown(Some(max)) where max is below the reserved range must not try the bin.
+                let page_below = reserved_allocator
+                    .allocate_pages(AllocationStrategy::TopDown(Some(reserved_range.start as usize)), 1, UEFI_PAGE_SIZE)
+                    .unwrap();
+                let page_below_addr = page_below.as_ptr() as *mut u8 as u64;
+                assert!(
+                    !reserved_range.contains(&page_below_addr),
+                    "TopDown(Some(below_bin)) should not land in the bin: addr={page_below_addr:#x}, range={reserved_range:#x?}",
+                );
+            });
+        });
+    }
 }

patina_dxe_core/src/gcd.rs

@@ -21,7 +21,7 @@ use patina::{
     error::EfiError,
     pi::{
         dxe_services::{GcdIoType, GcdMemoryType, MemorySpaceDescriptor},
-        hob::{self, Hob, HobList, PhaseHandoffInformationTable},
+        hob::{self, Hob, HobList, MEMORY_TYPE_INFO_HOB_GUID, PhaseHandoffInformationTable},
     },
 };
 use patina_internal_cpu::paging::{PatinaPageTable, create_cpu_paging};
@@ -557,6 +557,12 @@ pub fn add_hob_resource_descriptors_to_gcd(hob_list: &HobList) {
             None => continue, // Not a resource descriptor HOB or unsupported version for this build
         };
 
+        // Skip the PEI memory bin region to avoid a conflict. It will overlap the system
+        // memory region it was allocated from.
+        if res_desc.owner == MEMORY_TYPE_INFO_HOB_GUID {
+            continue;
+        }
+
         let mem_range = res_desc.physical_start
             ..res_desc.physical_start.checked_add(res_desc.resource_length).expect("Invalid resource descriptor hob");
 

patina_dxe_core/src/gcd/spin_locked_gcd.rs

@@ -17,7 +17,7 @@ use patina::{
     guids::{self, CACHE_ATTRIBUTE_CHANGE_EVENT_GROUP},
     pi::{
         dxe_services::{self, GcdMemoryType, MemorySpaceDescriptor},
-        hob::{self, EFiMemoryTypeInformation},
+        hob,
     },
     uefi_pages_to_size, uefi_size_to_pages,
 };
@@ -1971,7 +1971,6 @@ pub struct SpinLockedGcd {
     memory: tpl_mutex::TplMutex<GCD>,
     io: tpl_mutex::TplMutex<IoGCD>,
     memory_change_callback: Option<MapChangeCallback>,
-    memory_type_info_table: [EFiMemoryTypeInformation; 17],
     page_table: tpl_mutex::TplMutex<Option<Box<dyn PatinaPageTable>>>,
     /// Contains the current memory protection policy
     pub(crate) memory_protection_policy: MemoryProtectionPolicy,
@@ -2007,25 +2006,6 @@ impl SpinLockedGcd {
                 "GcdIoLock",
             ),
             memory_change_callback,
-            memory_type_info_table: [
-                EFiMemoryTypeInformation { memory_type: efi::RESERVED_MEMORY_TYPE, number_of_pages: 0 },
-                EFiMemoryTypeInformation { memory_type: efi::LOADER_CODE, number_of_pages: 0 },
-                EFiMemoryTypeInformation { memory_type: efi::LOADER_DATA, number_of_pages: 0 },
-                EFiMemoryTypeInformation { memory_type: efi::BOOT_SERVICES_CODE, number_of_pages: 0 },
-                EFiMemoryTypeInformation { memory_type: efi::BOOT_SERVICES_DATA, number_of_pages: 0 },
-                EFiMemoryTypeInformation { memory_type: efi::RUNTIME_SERVICES_CODE, number_of_pages: 0 },
-                EFiMemoryTypeInformation { memory_type: efi::RUNTIME_SERVICES_DATA, number_of_pages: 0 },
-                EFiMemoryTypeInformation { memory_type: efi::CONVENTIONAL_MEMORY, number_of_pages: 0 },
-                EFiMemoryTypeInformation { memory_type: efi::UNUSABLE_MEMORY, number_of_pages: 0 },
-                EFiMemoryTypeInformation { memory_type: efi::ACPI_RECLAIM_MEMORY, number_of_pages: 0 },
-                EFiMemoryTypeInformation { memory_type: efi::ACPI_MEMORY_NVS, number_of_pages: 0 },
-                EFiMemoryTypeInformation { memory_type: efi::MEMORY_MAPPED_IO, number_of_pages: 0 },
-                EFiMemoryTypeInformation { memory_type: efi::MEMORY_MAPPED_IO_PORT_SPACE, number_of_pages: 0 },
-                EFiMemoryTypeInformation { memory_type: efi::PAL_CODE, number_of_pages: 0 },
-                EFiMemoryTypeInformation { memory_type: efi::PERSISTENT_MEMORY, number_of_pages: 0 },
-                EFiMemoryTypeInformation { memory_type: efi::UNACCEPTED_MEMORY_TYPE, number_of_pages: 0 },
-                EFiMemoryTypeInformation { memory_type: 16 /*EfiMaxMemoryType*/, number_of_pages: 0 },
-            ],
             page_table: tpl_mutex::TplMutex::new(efi::TPL_HIGH_LEVEL, None, "GcdPageTableLock"),
             memory_protection_policy: MemoryProtectionPolicy::new(),
             last_efi_memory_map_key: tpl_mutex::TplMutex::new(efi::TPL_HIGH_LEVEL, None, "LastEfiMemoryMapKeyLock"),
@@ -2055,16 +2035,6 @@ impl SpinLockedGcd {
         self.memory.lock().prioritize_32_bit_memory = value;
     }
 
-    /// Returns a reference to the memory type information table.
-    pub const fn memory_type_info_table(&self) -> &[EFiMemoryTypeInformation; 17] {
-        &self.memory_type_info_table
-    }
-
-    /// Returns a pointer to the memory type information for the given memory type.
-    pub const fn memory_type_info(&self, memory_type: u32) -> &EFiMemoryTypeInformation {
-        &self.memory_type_info_table[memory_type as usize]
-    }
-
     fn set_paging_attributes(&self, base_address: usize, len: usize, attributes: u64) -> Result<(), EfiError> {
         if let Some(page_table) = &mut *self.page_table.lock() {
             // only apply page table attributes to the page table, not our virtual GCD attributes

patina_dxe_core/src/lib.rs

@@ -84,6 +84,7 @@ mod events;
 mod filesystems;
 mod gcd;
 mod memory_attributes_protocol;
+mod memory_bin;
 mod memory_manager;
 mod misc_boot_services;
 mod pecoff;
@@ -428,6 +429,7 @@ impl<P: PlatformInfo> Core<P> {
 
         //make sure that well-known handles exist.
         PROTOCOL_DB.init_protocol_db();
+
         // Initialize full allocation support.
         allocator::init_memory_support(&hob_list);