Просмотр исходного кода

builder: cleanup large build function

Philipp Schuster 1 год назад
Родитель
Сommit
41e42cf003
2 измененных файлов с 108 добавлено и 85 удалено
  1. 50 38
      multiboot2-header/src/builder/header.rs
  2. 58 47
      multiboot2/src/builder/information.rs

+ 50 - 38
multiboot2-header/src/builder/header.rs

@@ -159,9 +159,11 @@ impl HeaderBuilder {
     }
 
     /// Constructs the bytes for a valid Multiboot2 header with the given properties.
-    pub fn build(self) -> HeaderBytes {
+    pub fn build(mut self) -> HeaderBytes {
         const ALIGN: usize = 8;
 
+        // PHASE 1/3: Prepare Vector
+
         // We allocate more than necessary so that we can ensure an correct
         // alignment within this data.
         let alloc_len = self.expected_len() + 7;
@@ -179,44 +181,12 @@ impl HeaderBuilder {
         let offset = bytes.as_ptr().align_offset(ALIGN);
         bytes.extend([0].repeat(offset));
 
-        Self::build_add_bytes(
-            &mut bytes,
-            // important that we write the correct expected length into the header!
-            &Multiboot2BasicHeader::new(self.arch, self.expected_len() as u32).struct_as_bytes(),
-            false,
-        );
+        // -----------------------------------------------
+        // PHASE 2/3: Add Tags
+        self.build_add_tags(&mut bytes);
 
-        if let Some(irs) = self.information_request_tag {
-            Self::build_add_bytes(&mut bytes, &irs.build(), false)
-        }
-        if let Some(tag) = self.address_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
-        }
-        if let Some(tag) = self.entry_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
-        }
-        if let Some(tag) = self.console_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
-        }
-        if let Some(tag) = self.framebuffer_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
-        }
-        if let Some(tag) = self.module_align_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
-        }
-        if let Some(tag) = self.efi_bs_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
-        }
-        if let Some(tag) = self.efi_32_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
-        }
-        if let Some(tag) = self.efi_64_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
-        }
-        if let Some(tag) = self.relocatable_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
-        }
-        Self::build_add_bytes(&mut bytes, &EndHeaderTag::new().struct_as_bytes(), true);
+        // -----------------------------------------------
+        // PHASE 3/3: Finalize Vector
 
         // Ensure that the vector has the same length as it's capacity. This is
         // important so that miri doesn't complain that the boxed memory is
@@ -244,6 +214,48 @@ impl HeaderBuilder {
         HeaderBytes { offset, bytes }
     }
 
+    /// Helper method that adds all the tags to the given vector.
+    fn build_add_tags(&mut self, bytes: &mut Vec<u8>) {
+        Self::build_add_bytes(
+            bytes,
+            // important that we write the correct expected length into the header!
+            &Multiboot2BasicHeader::new(self.arch, self.expected_len() as u32).struct_as_bytes(),
+            false,
+        );
+
+        if let Some(irs) = self.information_request_tag.take() {
+            Self::build_add_bytes(bytes, &irs.build(), false)
+        }
+        if let Some(tag) = self.address_tag.as_ref() {
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
+        }
+        if let Some(tag) = self.entry_tag.as_ref() {
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
+        }
+        if let Some(tag) = self.console_tag.as_ref() {
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
+        }
+        if let Some(tag) = self.framebuffer_tag.as_ref() {
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
+        }
+        if let Some(tag) = self.module_align_tag.as_ref() {
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
+        }
+        if let Some(tag) = self.efi_bs_tag.as_ref() {
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
+        }
+        if let Some(tag) = self.efi_32_tag.as_ref() {
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
+        }
+        if let Some(tag) = self.efi_64_tag.as_ref() {
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
+        }
+        if let Some(tag) = self.relocatable_tag.as_ref() {
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
+        }
+        Self::build_add_bytes(bytes, &EndHeaderTag::new().struct_as_bytes(), true);
+    }
+
     // clippy thinks this can be a const fn but the compiler denies it
     #[allow(clippy::missing_const_for_fn)]
     pub fn information_request_tag(

+ 58 - 47
multiboot2/src/builder/information.rs

@@ -182,6 +182,8 @@ impl InformationBuilder {
     pub fn build(self) -> BootInformationBytes {
         const ALIGN: usize = 8;
 
+        // PHASE 1/3: Prepare Vector
+
         // We allocate more than necessary so that we can ensure an correct
         // alignment within this data.
         let alloc_len = self.expected_len() + 7;
@@ -199,90 +201,99 @@ impl InformationBuilder {
         let offset = bytes.as_ptr().align_offset(ALIGN);
         bytes.extend([0].repeat(offset));
 
+        // -----------------------------------------------
+        // PHASE 2/3: Add Tags
+        self.build_add_tags(&mut bytes);
+
+        // -----------------------------------------------
+        // PHASE 3/3: Finalize Vector
+
+        // Ensure that the vector has the same length as it's capacity. This is
+        // important so that miri doesn't complain that the boxed memory is
+        // smaller than the original allocation.
+        bytes.extend([0].repeat(bytes.capacity() - bytes.len()));
+
+        assert_eq!(
+            alloc_ptr,
+            bytes.as_ptr(),
+            "Vector was reallocated. Alignment of MBI probably broken!"
+        );
+        assert_eq!(
+            bytes[0..offset].iter().sum::<u8>(),
+            0,
+            "The offset to alignment area should be zero."
+        );
+
+        // Construct a box from a vec without `into_boxed_slice`. The latter
+        // calls `shrink` on the allocator, which might reallocate this memory.
+        // We don't want that!
+        let bytes = unsafe { Box::from_raw(bytes.leak()) };
+
+        assert_eq!(bytes.len(), alloc_len);
+
+        BootInformationBytes { offset, bytes }
+    }
+
+    /// Helper method that adds all the tags to the given vector.
+    fn build_add_tags(&self, bytes: &mut Vec<u8>) {
         Self::build_add_bytes(
-            &mut bytes,
+            bytes,
             // important that we write the correct expected length into the header!
             &BootInformationHeader::new(self.expected_len() as u32).struct_as_bytes(),
             false,
         );
-
         if let Some(tag) = self.basic_memory_info_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.boot_loader_name_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.command_line_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.efisdt32_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.efisdt64_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.efi_boot_services_not_exited_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.efi_image_handle32.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.efi_image_handle64.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.efi_memory_map_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.elf_sections_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.framebuffer_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.image_load_addr.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.memory_map_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
-        for tag in self.module_tags {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+        for tag in &self.module_tags {
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.rsdp_v1_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.rsdp_v2_tag.as_ref() {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
-        for tag in self.smbios_tags {
-            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
+        for tag in &self.smbios_tags {
+            Self::build_add_bytes(bytes, &tag.struct_as_bytes(), false)
         }
-        Self::build_add_bytes(&mut bytes, &EndTag::default().struct_as_bytes(), true);
-
-        // Ensure that the vector has the same length as it's capacity. This is
-        // important so that miri doesn't complain that the boxed memory is
-        // smaller than the original allocation.
-        bytes.extend([0].repeat(bytes.capacity() - bytes.len()));
-
-        assert_eq!(
-            alloc_ptr,
-            bytes.as_ptr(),
-            "Vector was reallocated. Alignment of MBI probably broken!"
-        );
-        assert_eq!(
-            bytes[0..offset].iter().sum::<u8>(),
-            0,
-            "The offset to alignment area should be zero."
-        );
-
-        // Construct a box from a vec without `into_boxed_slice`. The latter
-        // calls `shrink` on the allocator, which might reallocate this memory.
-        // We don't want that!
-        let bytes = unsafe { Box::from_raw(bytes.leak()) };
-
-        assert_eq!(bytes.len(), alloc_len);
-
-        BootInformationBytes { offset, bytes }
+        Self::build_add_bytes(bytes, &EndTag::default().struct_as_bytes(), true);
     }
 
     pub fn basic_memory_info_tag(&mut self, basic_memory_info_tag: BasicMemoryInfoTag) {