multiboot2: Get a mutable reference to the EFI memory map

Author: YtvwlD

multiboot2/src/lib.rs

@@ -345,6 +345,11 @@ impl BootInformation {
         }
     }
 
+    /// Search for the EFI Memory map tag, return a mutable reference.
+    pub fn efi_memory_map_tag_mut(&mut self) -> Option<&mut EFIMemoryMapTag> {
+        self.get_tag_mut::<EFIMemoryMapTag, _>(TagType::EfiMmap)
+    }
+
     /// Search for the EFI 32-bit image handle pointer.
     pub fn efi_32_ih(&self) -> Option<&EFIImageHandle32> {
         self.get_tag::<EFIImageHandle32, _>(TagType::Efi32Ih)

multiboot2/src/memory_map.rs

@@ -314,6 +314,23 @@ impl EFIMemoryMapTag {
             phantom: PhantomData,
         }
     }
+
+    #[cfg(feature = "builder")]
+    /// Return an iterator over ALL marked memory areas, mutably.
+    ///
+    /// This differs from `MemoryMapTag` as for UEFI, the OS needs some non-
+    /// available memory areas for tables and such.
+    pub fn memory_areas_mut(&mut self) -> EFIMemoryAreaIterMut {
+        let self_ptr = self as *const EFIMemoryMapTag;
+        let start_area = (&self.descs[0]) as *const EFIMemoryDesc;
+        EFIMemoryAreaIterMut {
+            current_area: start_area as u64,
+            // NOTE: `last_area` is only a bound, it doesn't necessarily point exactly to the last element
+            last_area: (self_ptr as *const () as u64 + self.size as u64),
+            entry_size: self.desc_size,
+            phantom: PhantomData,
+        }
+    }
 }
 
 impl TagTrait for EFIMemoryMapTag {
@@ -532,3 +549,25 @@ impl<'a> Iterator for EFIMemoryAreaIter<'a> {
         }
     }
 }
+
+/// An iterator over ALL EFI memory areas, mutably.
+#[derive(Clone, Debug)]
+pub struct EFIMemoryAreaIterMut<'a> {
+    current_area: u64,
+    last_area: u64,
+    entry_size: u32,
+    phantom: PhantomData<&'a mut EFIMemoryDesc>,
+}
+
+impl<'a> Iterator for EFIMemoryAreaIterMut<'a> {
+    type Item = &'a mut EFIMemoryDesc;
+    fn next(&mut self) -> Option<&'a mut EFIMemoryDesc> {
+        if self.current_area > self.last_area {
+            None
+        } else {
+            let area = unsafe { &mut *(self.current_area as *mut EFIMemoryDesc) };
+            self.current_area += self.entry_size as u64;
+            Some(area)
+        }
+    }
+}