multiboot2: Get a mutable reference to the EFI memory map

Author: YtvwlD

multiboot2/src/lib.rs

@@ -520,6 +520,11 @@ impl<T: AsRef<BootInformationInner> + AsMut<BootInformationInner>> BootInformati
         self.get_tag_mut::<BasicMemoryInfoTag, _>(TagType::BasicMeminfo)
     }
 
+    /// 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)
+    }
+
     fn get_tag_mut<TagT: TagTrait + ?Sized, TagType: Into<TagTypeId>>(
         &mut self,
         typ: TagType,

multiboot2/src/memory_map.rs

@@ -328,6 +328,22 @@ impl EFIMemoryMapTag {
             phantom: PhantomData,
         }
     }
+
+    /// 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 *mut EFIMemoryMapTag;
+        let start_area = (&mut self.descs[0]) as *mut 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 *mut () as u64 + self.size as u64),
+            entry_size: self.desc_size,
+            phantom: PhantomData,
+        }
+    }
 }
 
 impl TagTrait for EFIMemoryMapTag {
@@ -393,3 +409,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)
+        }
+    }
+}