uefi: introduce traits MemoryMap and MemoryMapMut

Author: phip1611

uefi/src/table/boot.rs

@@ -11,7 +11,7 @@ use crate::{Char16, Error, Event, Guid, Handle, Result, Status, StatusExt};
 use core::cell::UnsafeCell;
 use core::ffi::c_void;
 use core::mem::{self, MaybeUninit};
-use core::ops::{Deref, DerefMut};
+use core::ops::{Deref, DerefMut, Index, IndexMut};
 use core::ptr::NonNull;
 use core::sync::atomic::{AtomicPtr, Ordering};
 use core::{ptr, slice};
@@ -1804,6 +1804,95 @@ impl MemoryMapMeta {
     }
 }
 
+/// An accessory to the UEFI memory map that can be either iterated or indexed
+/// like an array.
+///
+/// A [`MemoryMap`] is always associated with the unique [`MemoryMapKey`]
+/// bundled with the ma.
+///
+/// To iterate over the entries, call [`MemoryMap::entries`].
+///
+/// ## UEFI pitfalls
+/// **Please note** that when working with memory maps, the `entry_size` is
+/// usually larger than `size_of::<MemoryDescriptor` [[0]]. So to be safe,
+/// always use `entry_size` as step-size when interfacing with the memory map on
+/// a low level.
+///
+/// [0]: https://github.com/tianocore/edk2/blob/7142e648416ff5d3eac6c6d607874805f5de0ca8/MdeModulePkg/Core/PiSmmCore/Page.c#L1059
+pub trait MemoryMap: Index<usize, Output = MemoryDescriptor> {
+    // TODO also require IntoIterator?! :)
+
+    /// Returns the associated [`MemoryMapMeta`].
+    #[must_use]
+    fn meta(&self) -> MemoryMapMeta;
+
+    /// Returns the associated [`MemoryMapKey`].
+    #[must_use]
+    fn key(&self) -> MemoryMapKey;
+
+    /// Returns the number of keys in the map.
+    #[must_use]
+    fn len(&self) -> usize;
+
+    /// Returns a reference to the [`MemoryDescriptor`] at the given index, if
+    /// present.
+    #[must_use]
+    fn get(&self, index: usize) -> Option<&MemoryDescriptor> {
+        if index >= self.len() {
+            None
+        } else {
+            let offset = index * self.meta().desc_size;
+            unsafe {
+                self.buffer()
+                    .as_ptr()
+                    .add(offset)
+                    .cast::<MemoryDescriptor>()
+                    .as_ref()
+            }
+        }
+    }
+
+    /// Returns a reference to the underlying memory.
+    fn buffer(&self) -> &[u8];
+
+    /// Returns an Iterator of type [`MemoryMapIter`].
+    fn entries(&self) -> MemoryMapIter<'_>;
+}
+
+/// Extension to [`MemoryMap`] that adds mutable operations. This also includes
+/// the ability to sort the memory map.
+pub trait MemoryMapMut: MemoryMap + IndexMut<usize> {
+    /// Returns a mutable reference to the [`MemoryDescriptor`] at the given
+    /// index, if present.
+    #[must_use]
+    fn get_mut(&mut self, index: usize) -> Option<&mut MemoryDescriptor> {
+        if index >= self.len() {
+            None
+        } else {
+            let offset = index * self.meta().desc_size;
+            unsafe {
+                self.buffer_mut()
+                    .as_mut_ptr()
+                    .add(offset)
+                    .cast::<MemoryDescriptor>()
+                    .as_mut()
+            }
+        }
+    }
+
+    /// Sorts the memory map by physical address in place. This operation is
+    /// optional and should be invoked only once.
+    #[must_use]
+    fn sort(&mut self);
+
+    /// Returns a reference to the underlying memory.
+    ///
+    /// # Safety
+    ///
+    /// This is unsafe as there is a potential to create invalid entries.
+    unsafe fn buffer_mut(&self) -> &mut [u8];
+}
+
 /// An accessory to the memory map that can be either iterated or
 /// indexed like an array.
 ///