Add Sorted Iterator for the UEFI Memory Map (Issue#661)

uefi-test-runner/src/boot/memory.rs

@@ -91,17 +91,28 @@ fn memory_map(bt: &BootServices) {
     // We will use vectors for convenience.
     let mut buffer = vec![0_u8; buf_sz];
 
-    let (_key, desc_iter) = bt
+    let mut memory_map = bt
         .memory_map(&mut buffer)
         .expect("Failed to retrieve UEFI memory map");
 
+    memory_map.sort();
+
     // Collect the descriptors into a vector
-    let descriptors = desc_iter.copied().collect::<Vec<_>>();
+    let descriptors = memory_map.entries().copied().collect::<Vec<_>>();
 
     // Ensured we have at least one entry.
     // Real memory maps usually have dozens of entries.
     assert!(!descriptors.is_empty(), "Memory map is empty");
 
+    let mut curr_value = descriptors[0];
+
+    for value in descriptors.iter().skip(1) {
+        if value.phys_start <= curr_value.phys_start {
+            panic!("memory map sorting failed");
+        }
+        curr_value = *value;
+    }
+
     // This is pretty much a sanity test to ensure returned memory isn't filled with random values.
     let first_desc = descriptors[0];
 

uefi/src/table/boot.rs

@@ -425,10 +425,7 @@ impl BootServices {
     ///
     /// * [`uefi::Status::BUFFER_TOO_SMALL`]
     /// * [`uefi::Status::INVALID_PARAMETER`]
-    pub fn memory_map<'buf>(
-        &self,
-        buffer: &'buf mut [u8],
-    ) -> Result<(MemoryMapKey, MemoryMapIter<'buf>)> {
+    pub fn memory_map<'buf>(&self, buffer: &'buf mut [u8]) -> Result<MemoryMap<'buf>> {
         let mut map_size = buffer.len();
         MemoryDescriptor::assert_aligned(buffer);
         let map_buffer = buffer.as_mut_ptr().cast::<MemoryDescriptor>();
@@ -453,13 +450,13 @@ impl BootServices {
         }
         .into_with_val(move || {
             let len = map_size / entry_size;
-            let iter = MemoryMapIter {
-                buffer,
+
+            MemoryMap {
+                key: map_key,
+                buf: buffer,
                 entry_size,
-                index: 0,
                 len,
-            };
-            (map_key, iter)
+            }
         })
     }
 
@@ -1993,6 +1990,108 @@ pub struct MemoryMapSize {
     pub map_size: usize,
 }
 
+/// An iterator of [`MemoryDescriptor`] that is always associated with the
+/// unique [`MemoryMapKey`] contained in the struct.
+///
+/// To iterate over the entries, call [`MemoryMap::entries`]. To get a sorted
+/// map, you manually have to call [`MemoryMap::sort`] first.
+pub struct MemoryMap<'buf> {
+    key: MemoryMapKey,
+    buf: &'buf mut [u8],
+    entry_size: usize,
+    len: usize,
+}
+
+impl<'buf> MemoryMap<'buf> {
+    #[must_use]
+    /// Returns the unique [`MemoryMapKey`] associated with the memory map.
+    pub fn key(&self) -> MemoryMapKey {
+        self.key
+    }
+
+    /// Sorts the memory map by physical address in place.
+    /// This operation is optional and should be invoked only once.
+    pub fn sort(&mut self) {
+        unsafe {
+            self.qsort(0, self.len - 1);
+        }
+    }
+
+    /// Hoare partition scheme for quicksort.
+    /// Must be called with `low` and `high` being indices within bounds.
+    unsafe fn qsort(&mut self, low: usize, high: usize) {
+        if low >= high {
+            return;
+        }
+
+        let p = self.partition(low, high);
+        self.qsort(low, p);
+        self.qsort(p + 1, high);
+    }
+
+    unsafe fn partition(&mut self, low: usize, high: usize) -> usize {
+        let pivot = self.get_element_phys_addr(low + (high - low) / 2);
+
+        let mut left_index = low.wrapping_sub(1);
+        let mut right_index = high.wrapping_add(1);
+
+        loop {
+            while {
+                left_index = left_index.wrapping_add(1);
+
+                self.get_element_phys_addr(left_index) < pivot
+            } {}
+
+            while {
+                right_index = right_index.wrapping_sub(1);
+
+                self.get_element_phys_addr(right_index) > pivot
+            } {}
+
+            if left_index >= right_index {
+                return right_index;
+            }
+
+            self.swap(left_index, right_index);
+        }
+    }
+
+    /// Indices must be smaller than len.
+    unsafe fn swap(&mut self, index1: usize, index2: usize) {
+        if index1 == index2 {
+            return;
+        }
+
+        let base = self.buf.as_mut_ptr();
+
+        unsafe {
+            ptr::swap_nonoverlapping(
+                base.add(index1 * self.entry_size),
+                base.add(index2 * self.entry_size),
+                self.entry_size,
+            );
+        }
+    }
+
+    fn get_element_phys_addr(&self, index: usize) -> PhysicalAddress {
+        let offset = index.checked_mul(self.entry_size).unwrap();
+        let elem = unsafe { &*self.buf.as_ptr().add(offset).cast::<MemoryDescriptor>() };
+        elem.phys_start
+    }
+
+    /// Returns an iterator over the contained memory map. To get a sorted map,
+    /// call [`MemoryMap::sort`] first.
+    #[must_use]
+    pub fn entries(&self) -> MemoryMapIter {
+        MemoryMapIter {
+            buffer: self.buf,
+            entry_size: self.entry_size,
+            index: 0,
+            len: self.len,
+        }
+    }
+}
+
 /// An iterator of [`MemoryDescriptor`]. The underlying memory map is always
 /// associated with a unique [`MemoryMapKey`].
 #[derive(Debug, Clone)]
@@ -2014,12 +2113,16 @@ impl<'buf> Iterator for MemoryMapIter<'buf> {
 
     fn next(&mut self) -> Option<Self::Item> {
         if self.index < self.len {
-            let ptr = self.buffer.as_ptr() as usize + self.entry_size * self.index;
+            let descriptor = unsafe {
+                &*self
+                    .buffer
+                    .as_ptr()
+                    .add(self.entry_size * self.index)
+                    .cast::<MemoryDescriptor>()
+            };
 
             self.index += 1;
 
-            let descriptor = unsafe { &*(ptr as *const MemoryDescriptor) };
-
             Some(descriptor)
         } else {
             None
@@ -2197,3 +2300,93 @@ pub enum InterfaceType: i32 => {
 #[derive(Debug, Clone, Copy)]
 #[repr(transparent)]
 pub struct ProtocolSearchKey(NonNull<c_void>);
+
+#[cfg(test)]
+mod tests {
+    use core::mem::size_of;
+
+    use crate::table::boot::{MemoryAttribute, MemoryMap, MemoryMapKey, MemoryType};
+
+    use super::{MemoryDescriptor, MemoryMapIter};
+
+    #[test]
+    fn mem_map_sorting() {
+        // Doesn't matter what type it is.
+        const TY: MemoryType = MemoryType::RESERVED;
+
+        const BASE: MemoryDescriptor = MemoryDescriptor {
+            ty: TY,
+            phys_start: 0,
+            virt_start: 0,
+            page_count: 0,
+            att: MemoryAttribute::empty(),
+        };
+
+        let mut buffer = [
+            MemoryDescriptor {
+                phys_start: 2000,
+                ..BASE
+            },
+            MemoryDescriptor {
+                phys_start: 3000,
+                ..BASE
+            },
+            BASE,
+            MemoryDescriptor {
+                phys_start: 1000,
+                ..BASE
+            },
+        ];
+
+        let desc_count = buffer.len();
+
+        let byte_buffer = {
+            let size = desc_count * size_of::<MemoryDescriptor>();
+            unsafe { core::slice::from_raw_parts_mut(buffer.as_mut_ptr() as *mut u8, size) }
+        };
+
+        let mut mem_map = MemoryMap {
+            // Key doesn't matter
+            key: MemoryMapKey(0),
+            len: desc_count,
+            buf: byte_buffer,
+            entry_size: size_of::<MemoryDescriptor>(),
+        };
+
+        mem_map.sort();
+
+        if !is_sorted(&mem_map.entries()) {
+            panic!("mem_map is not sorted: {}", mem_map);
+        }
+    }
+
+    // Added for debug purposes on test failure
+    impl core::fmt::Display for MemoryMap<'_> {
+        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+            writeln!(f)?;
+            for desc in self.entries() {
+                writeln!(f, "{:?}", desc)?;
+            }
+            Ok(())
+        }
+    }
+
+    fn is_sorted(iter: &MemoryMapIter) -> bool {
+        let mut iter = iter.clone();
+        let mut curr_start;
+
+        if let Some(val) = iter.next() {
+            curr_start = val.phys_start;
+        } else {
+            return true;
+        }
+
+        for desc in iter {
+            if desc.phys_start <= curr_start {
+                return false;
+            }
+            curr_start = desc.phys_start
+        }
+        true
+    }
+}

uefi/src/table/system.rs

@@ -7,7 +7,7 @@ use core::{ptr, slice};
 use crate::proto::console::text;
 use crate::{CStr16, Char16, Handle, Result, Status};
 
-use super::boot::{BootServices, MemoryDescriptor, MemoryMapIter, MemoryType};
+use super::boot::{BootServices, MemoryDescriptor, MemoryMap, MemoryType};
 use super::runtime::{ResetType, RuntimeServices};
 use super::{cfg, Header, Revision};
 
@@ -158,20 +158,20 @@ impl SystemTable<Boot> {
     unsafe fn get_memory_map_and_exit_boot_services(
         &self,
         buf: &'static mut [u8],
-    ) -> Result<MemoryMapIter<'static>> {
+    ) -> Result<MemoryMap<'static>> {
         let boot_services = self.boot_services();
 
         // Get the memory map.
-        let (memory_map_key, memory_map_iter) = boot_services.memory_map(buf)?;
+        let memory_map = boot_services.memory_map(buf)?;
 
         // Try to exit boot services using the memory map key. Note that after
         // the first call to `exit_boot_services`, there are restrictions on
         // what boot services functions can be called. In UEFI 2.8 and earlier,
         // only `get_memory_map` and `exit_boot_services` are allowed. Starting
         // in UEFI 2.9 other memory allocation functions may also be called.
         boot_services
-            .exit_boot_services(boot_services.image_handle(), memory_map_key)
-            .map(move |()| memory_map_iter)
+            .exit_boot_services(boot_services.image_handle(), memory_map.key())
+            .map(move |()| memory_map)
     }
 
     /// Exit the UEFI boot services.
@@ -212,7 +212,7 @@ impl SystemTable<Boot> {
     /// [`Logger::disable`]: crate::logger::Logger::disable
     /// [`uefi_services::init`]: https://docs.rs/uefi-services/latest/uefi_services/fn.init.html
     #[must_use]
-    pub fn exit_boot_services(self) -> (SystemTable<Runtime>, MemoryMapIter<'static>) {
+    pub fn exit_boot_services(self) -> (SystemTable<Runtime>, MemoryMap<'static>) {
         let boot_services = self.boot_services();
 
         // Reboot the device.