Refactor: AlignedMemory::resize()

src/aligned_memory.rs

@@ -9,6 +9,7 @@ pub struct AlignedMemory<const ALIGN: usize> {
     max_len: usize,
     align_offset: usize,
     mem: Vec<u8>,
+    zero_up_to_max_len: bool,
 }
 
 impl<const ALIGN: usize> AlignedMemory<ALIGN> {
@@ -19,7 +20,6 @@ impl<const ALIGN: usize> AlignedMemory<ALIGN> {
         mem.resize(align_offset, 0);
         (mem, align_offset)
     }
-
     fn get_mem_zeroed(max_len: usize) -> (Vec<u8>, usize) {
         // use calloc() to get zeroed memory from the OS instead of using
         // malloc() + memset(), see
@@ -29,34 +29,47 @@ impl<const ALIGN: usize> AlignedMemory<ALIGN> {
         mem.resize(align_offset + max_len, 0);
         (mem, align_offset)
     }
-
-    /// Return a new AlignedMemory type
-    pub fn new(max_len: usize) -> Self {
+    /// Returns a filled AlignedMemory by copying the given slice
+    pub fn from_slice(data: &[u8]) -> Self {
+        let max_len = data.len();
+        let (mut mem, align_offset) = Self::get_mem(max_len);
+        mem.extend_from_slice(data);
+        Self {
+            max_len,
+            align_offset,
+            mem,
+            zero_up_to_max_len: false,
+        }
+    }
+    /// Returns a new empty AlignedMemory with uninitialized preallocated memory
+    pub fn with_capacity(max_len: usize) -> Self {
         let (mem, align_offset) = Self::get_mem(max_len);
         Self {
             max_len,
             align_offset,
             mem,
+            zero_up_to_max_len: false,
         }
     }
-    /// Return a pre-filled AlignedMemory type
-    pub fn new_with_size(len: usize) -> Self {
-        let (mem, align_offset) = Self::get_mem_zeroed(len);
+    /// Returns a new empty AlignedMemory with zero initialized preallocated memory
+    pub fn with_capacity_zeroed(max_len: usize) -> Self {
+        let (mut mem, align_offset) = Self::get_mem_zeroed(max_len);
+        mem.truncate(align_offset);
         Self {
-            max_len: len,
+            max_len,
             align_offset,
             mem,
+            zero_up_to_max_len: true,
         }
     }
-    /// Return a pre-filled AlignedMemory type
-    pub fn new_with_data(data: &[u8]) -> Self {
-        let max_len = data.len();
-        let (mut mem, align_offset) = Self::get_mem(max_len);
-        mem.extend_from_slice(data);
+    /// Returns a new filled AlignedMemory with zero initialized preallocated memory
+    pub fn zero_filled(max_len: usize) -> Self {
+        let (mem, align_offset) = Self::get_mem_zeroed(max_len);
         Self {
             max_len,
             align_offset,
             mem,
+            zero_up_to_max_len: true,
         }
     }
     /// Calculate memory size
@@ -87,15 +100,22 @@ impl<const ALIGN: usize> AlignedMemory<ALIGN> {
         let end = self.mem.len();
         &mut self.mem[start..end]
     }
-    /// resize memory with value starting at the write_index
-    pub fn resize(&mut self, num: usize, value: u8) -> std::io::Result<()> {
+    /// Grows memory with `value` repeated `num` times starting at the `write_index`
+    pub fn fill_write(&mut self, num: usize, value: u8) -> std::io::Result<()> {
         if self.mem.len() + num > self.align_offset + self.max_len {
             return Err(std::io::Error::new(
                 std::io::ErrorKind::InvalidInput,
                 "aligned memory resize failed",
             ));
         }
-        self.mem.resize(self.mem.len() + num, value);
+        if self.zero_up_to_max_len && value == 0 {
+            // Safe because everything up to `max_len` is zeroed and no shrinking is allowed
+            unsafe {
+                self.mem.set_len(self.mem.len() + num);
+            }
+        } else {
+            self.mem.resize(self.mem.len() + num, value);
+        }
         Ok(())
     }
 }
@@ -105,7 +125,7 @@ impl<const ALIGN: usize> AlignedMemory<ALIGN> {
 // aligned.
 impl<const ALIGN: usize> Clone for AlignedMemory<ALIGN> {
     fn clone(&self) -> Self {
-        AlignedMemory::new_with_data(self.as_slice())
+        AlignedMemory::from_slice(self.as_slice())
     }
 }
 
@@ -127,7 +147,7 @@ impl<const ALIGN: usize> std::io::Write for AlignedMemory<ALIGN> {
 
 impl<const ALIGN: usize, T: AsRef<[u8]>> From<T> for AlignedMemory<ALIGN> {
     fn from(bytes: T) -> Self {
-        AlignedMemory::new_with_data(bytes.as_ref())
+        AlignedMemory::from_slice(bytes.as_ref())
     }
 }
 
@@ -145,7 +165,7 @@ mod tests {
     use std::io::Write;
 
     fn do_test<const ALIGN: usize>() {
-        let mut aligned_memory = AlignedMemory::<ALIGN>::new(10);
+        let mut aligned_memory = AlignedMemory::<ALIGN>::with_capacity(10);
 
         assert_eq!(aligned_memory.write(&[42u8; 1]).unwrap(), 1);
         assert_eq!(aligned_memory.write(&[42u8; 9]).unwrap(), 9);
@@ -157,16 +177,16 @@ mod tests {
         aligned_memory.as_slice_mut().copy_from_slice(&[84u8; 10]);
         assert_eq!(aligned_memory.as_slice(), &[84u8; 10]);
 
-        let mut aligned_memory = AlignedMemory::<ALIGN>::new(10);
-        aligned_memory.resize(5, 0).unwrap();
-        aligned_memory.resize(2, 1).unwrap();
+        let mut aligned_memory = AlignedMemory::<ALIGN>::with_capacity_zeroed(10);
+        aligned_memory.fill_write(5, 0).unwrap();
+        aligned_memory.fill_write(2, 1).unwrap();
         assert_eq!(aligned_memory.write(&[2u8; 3]).unwrap(), 3);
         assert_eq!(aligned_memory.as_slice(), &[0, 0, 0, 0, 0, 1, 1, 2, 2, 2]);
-        aligned_memory.resize(1, 3).unwrap_err();
+        aligned_memory.fill_write(1, 3).unwrap_err();
         aligned_memory.write(&[4u8; 1]).unwrap_err();
         assert_eq!(aligned_memory.as_slice(), &[0, 0, 0, 0, 0, 1, 1, 2, 2, 2]);
 
-        let aligned_memory = AlignedMemory::<ALIGN>::new_with_size(10);
+        let aligned_memory = AlignedMemory::<ALIGN>::zero_filled(10);
         assert_eq!(aligned_memory.len(), 10);
         assert_eq!(aligned_memory.as_slice(), &[0u8; 10]);
     }

src/call_frames.rs

@@ -33,7 +33,7 @@ impl<'a> CallFrames<'a> {
     /// New call frame, depth indicates maximum call depth
     pub fn new(config: &'a Config) -> Self {
         let stack_len = config.stack_size();
-        let stack = AlignedMemory::new_with_size(stack_len);
+        let stack = AlignedMemory::zero_filled(stack_len);
 
         let mut frames = CallFrames {
             config,

src/elf.rs

@@ -366,7 +366,7 @@ impl<E: UserDefinedError, I: InstructionMeter> Executable<E, I> {
         syscall_registry: SyscallRegistry,
         bpf_functions: BTreeMap<u32, (usize, String)>,
     ) -> Self {
-        let elf_bytes = AlignedMemory::new_with_data(text_bytes);
+        let elf_bytes = AlignedMemory::from_slice(text_bytes);
         let enable_symbol_and_section_labels = config.enable_symbol_and_section_labels;
         Self {
             config,
@@ -402,7 +402,7 @@ impl<E: UserDefinedError, I: InstructionMeter> Executable<E, I> {
             let bytes = if is_memory_aligned(bytes, HOST_ALIGN) {
                 bytes
             } else {
-                aligned = AlignedMemory::<{ HOST_ALIGN }>::new_with_data(bytes);
+                aligned = AlignedMemory::<{ HOST_ALIGN }>::from_slice(bytes);
                 aligned.as_slice()
             };
             Self::load_with_parser(&NewParser::parse(bytes)?, config, bytes, syscall_registry)
@@ -422,7 +422,7 @@ impl<E: UserDefinedError, I: InstructionMeter> Executable<E, I> {
         bytes: &[u8],
         syscall_registry: SyscallRegistry,
     ) -> Result<Self, ElfError> {
-        let mut elf_bytes = AlignedMemory::new_with_data(bytes);
+        let mut elf_bytes = AlignedMemory::from_slice(bytes);
 
         Self::validate(&mut config, elf, elf_bytes.as_slice())?;
 
@@ -2143,6 +2143,6 @@ mod test {
             Executable::jit_compile(&mut executable).unwrap();
         }
 
-        assert_eq!(18640, executable.mem_size());
+        assert_eq!(18656, executable.mem_size());
     }
 }