Introduce a StreamingBuffer

src/write/coff.rs

@@ -579,7 +579,7 @@ impl Object {
                     debug_assert!(aux_len >= symbol.name.len());
                     let old_len = buffer.len();
                     buffer.write_bytes(&symbol.name);
-                    buffer.resize(old_len + aux_len, 0);
+                    buffer.resize(old_len + aux_len);
                 }
                 SymbolKind::Section => {
                     debug_assert_eq!(number_of_aux_symbols, 1);

src/write/elf.rs

@@ -178,7 +178,9 @@ impl Object {
             .sections
             .iter()
             .map(|section| {
-                let mut reloc_name = Vec::new();
+                let mut reloc_name = Vec::with_capacity(
+                    if is_rela { ".rela".len() } else { ".rel".len() } + section.name.len(),
+                );
                 if !section.relocations.is_empty() {
                     reloc_name.extend_from_slice(if is_rela {
                         &b".rela"[..]
@@ -416,7 +418,7 @@ impl Object {
 
         // Write section data.
         for (index, comdat) in self.comdats.iter().enumerate() {
-            let mut data = Vec::new();
+            let mut data = Vec::with_capacity(comdat_offsets[index].len);
             data.write_pod(&U32::new(self.endian, elf::GRP_COMDAT));
             for section in &comdat.sections {
                 data.write_pod(&U32::new(
@@ -453,7 +455,7 @@ impl Object {
                 st_size: 0,
             },
         );
-        let mut symtab_shndx = Vec::new();
+        let mut symtab_shndx = Vec::with_capacity(symtab_shndx_len);
         if need_symtab_shndx {
             symtab_shndx.write_pod(&U32::new(endian, 0));
         }

src/write/util.rs

@@ -1,3 +1,4 @@
+use std::io::Write;
 use std::vec::Vec;
 
 use crate::pod::{bytes_of, bytes_of_slice, Pod};
@@ -6,14 +7,18 @@ use crate::pod::{bytes_of, bytes_of_slice, Pod};
 #[allow(clippy::len_without_is_empty)]
 pub trait WritableBuffer {
     /// Returns position/offset for data to be written at.
+    /// Should only be used in debug assertions
     fn len(&self) -> usize;
 
     /// Reserves specified number of bytes in the buffer.
-    fn reserve(&mut self, additional: usize) -> Result<(), ()>;
+    /// Must be called exactly once before writing anything to the buffer.
+    /// Must be given the exact of the buffer after writing everything, calling with a smaller size
+    /// may result in a panic, while calling with a bigger size may result in trailing garbage.
+    fn reserve(&mut self, size: usize) -> Result<(), ()>;
 
-    /// Writes the specified value at the end of the buffer
-    /// until the buffer has the specified length.
-    fn resize(&mut self, new_len: usize, value: u8);
+    /// Writes zero bytes at the end of the buffer until the buffer
+    /// has the specified length.
+    fn resize(&mut self, new_len: usize);
 
     /// Writes the specified slice of bytes at the end of the buffer.
     fn write_bytes(&mut self, val: &[u8]);
@@ -54,22 +59,74 @@ impl WritableBuffer for Vec<u8> {
     }
 
     #[inline]
-    fn reserve(&mut self, additional: usize) -> Result<(), ()> {
-        self.reserve(additional);
+    fn reserve(&mut self, size: usize) -> Result<(), ()> {
+        assert_eq!(self.len(), 0, "Buffer must be empty");
+        self.reserve(size);
         Ok(())
     }
 
     #[inline]
-    fn resize(&mut self, new_len: usize, value: u8) {
-        self.resize(new_len, value);
+    fn resize(&mut self, new_len: usize) {
+        debug_assert!(new_len >= self.len());
+        self.resize(new_len, 0);
     }
 
     #[inline]
     fn write_bytes(&mut self, val: &[u8]) {
+        debug_assert!(self.len() + val.len() <= self.capacity());
         self.extend_from_slice(val)
     }
 }
 
+/// A [`WritableBuffer`] that streams data to a [`Write`]r.
+///
+/// It is advisable to use a buffered writer like [`BufWriter`](std::io::BufWriter) instead of an
+/// unbuffered writer like [`File`](std::fs::File).
+#[derive(Debug)]
+pub struct StreamingBuffer<W> {
+    writer: W,
+    length: usize,
+}
+
+impl<W> StreamingBuffer<W> {
+    /// Create a new `WriteBuf` backed by the given writer.
+    pub fn new(writer: W) -> Self {
+        StreamingBuffer { writer, length: 0 }
+    }
+
+    /// Unwraps this [`WriteBuf`] giving back the original writer.
+    pub fn into_inner(self) -> W {
+        self.writer
+    }
+}
+
+impl<W: Write> WritableBuffer for StreamingBuffer<W> {
+    #[inline]
+    fn len(&self) -> usize {
+        self.length
+    }
+
+    #[inline]
+    fn reserve(&mut self, _additional: usize) -> Result<(), ()> {
+        Ok(())
+    }
+
+    #[inline]
+    fn resize(&mut self, new_len: usize) {
+        debug_assert!(new_len >= self.length);
+        while self.length < new_len {
+            let write_amt = (new_len - self.length - 1) % 1024 + 1;
+            self.write_bytes(&[0; 1024][..write_amt]);
+        }
+    }
+
+    #[inline]
+    fn write_bytes(&mut self, val: &[u8]) {
+        self.writer.write_all(val).unwrap();
+        self.length += val.len();
+    }
+}
+
 /// A trait for mutable byte slices.
 ///
 /// It provides convenience methods for `Pod` types.
@@ -99,7 +156,7 @@ pub(crate) fn align_u64(offset: u64, size: u64) -> u64 {
 
 pub(crate) fn write_align(buffer: &mut dyn WritableBuffer, size: usize) {
     let new_len = align(buffer.len(), size);
-    buffer.resize(new_len, 0);
+    buffer.resize(new_len);
 }
 
 #[cfg(test)]