Factor out functionality to write_section_raw

marxin · Sep 4, 2024 · a7b21df · a7b21df
1 parent 2903c06
commit a7b21df
Showing 1 changed file with 30 additions and 35 deletions.
diff --git a/wild_lib/src/elf_writer.rs b/wild_lib/src/elf_writer.rs
@@ -52,6 +52,7 @@ use crate::threading::prelude::*;
 use ahash::AHashMap;
 use anyhow::anyhow;
 use anyhow::bail;
+use anyhow::ensure;
 use anyhow::Context;
 use linker_utils::elf::rel_type_to_string;
 use linker_utils::elf::shf;
@@ -1142,14 +1143,15 @@ impl<'data> ObjectLayout<'data> {
                 SectionSlot::Loaded(sec) => {
                     self.write_section(layout, sec, buffers, table_writer)?
                 }
-                SectionSlot::LoadedDebugInfo(part_id, section_id, section_size) => self
-                    .write_debug_section(
+                SectionSlot::LoadedDebugInfo(part_id, section_id, section_size) => {
+                    self.write_section_raw(
                         layout,
                         part_id,
                         *section_id,
                         *section_size as usize,
                         buffers,
-                    )?,
+                    )?;
+                }
                 SectionSlot::EhFrameData(section_index) => {
                     self.write_eh_frame_data(*section_index, layout, table_writer)?;
                 }
@@ -1203,21 +1205,15 @@ impl<'data> ObjectLayout<'data> {
             .output_sections
             .has_data_in_file(sec.output_section_id())
         {
-            let section_buffer = buffers.get_mut(sec.output_part_id());
-            let allocation_size = sec.capacity() as usize;
-            if section_buffer.len() < allocation_size {
-                bail!(
-                    "Insufficient space allocated to section `{}`. Tried to take {} bytes, but only {} remain",
-                    self.object.section_display_name(sec.index),
-                    allocation_size, section_buffer.len()
-                );
-            }
-            let out = slice_take_prefix_mut(section_buffer, allocation_size);
-            // Cut off any padding so that our output buffer is the size of our input buffer.
             let object_section = self.object.section(sec.index)?;
-            let section_size = self.object.section_size(object_section)?;
-            let out = &mut out[..section_size as usize];
-            self.object.copy_section_data(object_section, out)?;
+            let section_size = self.object.section_size(object_section)? as usize;
+            let out = self.write_section_raw(
+                layout,
+                &sec.output_part_id(),
+                sec.index,
+                section_size,
+                buffers,
+            )?;
             self.apply_relocations(out, sec, layout, table_writer)
                 .with_context(|| {
                     format!(
@@ -1238,35 +1234,34 @@ impl<'data> ObjectLayout<'data> {
         Ok(())
     }
 
-    fn write_debug_section(
+    fn write_section_raw(
         &self,
         layout: &Layout,
         part_id: &PartId,
         section_index: SectionIndex,
         section_size: usize,
-        buffers: &mut OutputSectionPartMap<&mut [u8]>,
-    ) -> Result {
-        if layout
+        buffers: &'data mut OutputSectionPartMap<&mut [u8]>,
+    ) -> Result<&'data mut [u8]> {
+        ensure!(layout
             .output_sections
-            .has_data_in_file(part_id.output_section_id())
-        {
-            let section_buffer = buffers.get_mut(*part_id);
-            let allocation_size = section_size;
-            if section_buffer.len() < allocation_size {
-                bail!(
+            .has_data_in_file(part_id.output_section_id()));
+
+        let section_buffer = buffers.get_mut(*part_id);
+        let allocation_size = section_size;
+        if section_buffer.len() < allocation_size {
+            bail!(
                     "Insufficient space allocated to section `{}`. Tried to take {} bytes, but only {} remain",
                     self.object.section_display_name(section_index),
                     allocation_size, section_buffer.len()
                 );
-            }
-            let out = slice_take_prefix_mut(section_buffer, allocation_size);
-            // Cut off any padding so that our output buffer is the size of our input buffer.
-            let object_section = self.object.section(section_index)?;
-            let section_size = self.object.section_size(object_section)?;
-            let out = &mut out[..section_size as usize];
-            self.object.copy_section_data(object_section, out)?;
         }
-        Ok(())
+        let out = slice_take_prefix_mut(section_buffer, allocation_size);
+        // Cut off any padding so that our output buffer is the size of our input buffer.
+        let object_section = self.object.section(section_index)?;
+        let section_size = self.object.section_size(object_section)?;
+        let out: &'data mut [u8] = &mut out[..section_size as usize];
+        self.object.copy_section_data(object_section, out)?;
+        Ok(out)
     }
 
     /// Writes debug symbols.