feat: add static buffer as an alternative to allocating memory (#939)

* feat: add static buffer as an alternative to allocating memory

* fix: zeroing static buffer gradually

crates/wasmi/src/memory/buffer.rs

@@ -1,6 +1,7 @@
+use core::mem::ManuallyDrop;
 use std::{vec, vec::Vec};
 
-/// A `Vec`-based byte buffer implementation.
+/// A byte buffer implementation.
 ///
 /// # Note
 ///
@@ -9,9 +10,16 @@ use std::{vec, vec::Vec};
 /// solution fitting any platform.
 #[derive(Debug)]
 pub struct ByteBuffer {
-    bytes: Vec<u8>,
+    ptr: *mut u8,
+    len: usize,
+    capacity: usize,
+    is_static: bool,
 }
 
+// Safety: `ByteBuffer` is essentially an enum of `Vec<u8>` or `&'static mut [u8]`.
+// They both are `Send` so this is sound.
+unsafe impl Send for ByteBuffer {}
+
 impl ByteBuffer {
     /// Creates a new byte buffer with the given initial length.
     ///
@@ -20,33 +28,153 @@ impl ByteBuffer {
     /// - If the initial length is 0.
     /// - If the initial length exceeds the maximum supported limit.
     pub fn new(initial_len: usize) -> Self {
+        let mut vec = ManuallyDrop::new(vec![0x00_u8; initial_len]);
+        let (ptr, len, capacity) = (vec.as_mut_ptr(), vec.len(), vec.capacity());
         Self {
-            bytes: vec![0x00_u8; initial_len],
+            ptr,
+            len,
+            capacity,
+            is_static: false,
+        }
+    }
+
+    /// Creates a new byte buffer with the given initial length.
+    ///
+    /// # Errors
+    ///
+    /// - If the initial length is 0.
+    /// - If the initial length exceeds the maximum supported limit.
+    pub fn new_static(buf: &'static mut [u8], initial_len: usize) -> Self {
+        assert!(initial_len <= buf.len());
+        buf[..initial_len].fill(0x00_u8);
+        Self {
+            ptr: buf.as_mut_ptr(),
+            len: initial_len,
+            capacity: buf.len(),
+            is_static: true,
         }
     }
 
     /// Grows the byte buffer to the given `new_size`.
     ///
     /// # Panics
     ///
-    /// If the current size of the [`ByteBuffer`] is larger than `new_size`.
+    /// - If the current size of the [`ByteBuffer`] is larger than `new_size`.
+    /// - If backed by static buffer and `new_size` is larger than it's capacity.
     pub fn grow(&mut self, new_size: usize) {
         assert!(new_size >= self.len());
-        self.bytes.resize(new_size, 0x00_u8);
+        if self.is_static {
+            if self.capacity < new_size {
+                panic!("Cannot grow static byte buffer more then it's capacity")
+            }
+            let len = self.len();
+            self.len = new_size;
+            self.data_mut()[len..new_size].fill(0x00_u8);
+        } else {
+            // Safety: those parts have been obtained from `Vec`.
+            let vec = unsafe { Vec::from_raw_parts(self.ptr, self.len, self.capacity) };
+            let mut vec = ManuallyDrop::new(vec);
+            vec.resize(new_size, 0x00_u8);
+            let (ptr, len, capacity) = (vec.as_mut_ptr(), vec.len(), vec.capacity());
+            self.ptr = ptr;
+            self.len = len;
+            self.capacity = capacity;
+        }
     }
 
     /// Returns the length of the byte buffer in bytes.
     pub fn len(&self) -> usize {
-        self.bytes.len()
+        self.len
     }
 
     /// Returns a shared slice to the bytes underlying to the byte buffer.
     pub fn data(&self) -> &[u8] {
-        &self.bytes[..]
+        // Safety: either is backed by a `Vec` or a static buffer, ptr[0..len] is valid.
+        unsafe { core::slice::from_raw_parts(self.ptr, self.len) }
     }
 
     /// Returns an exclusive slice to the bytes underlying to the byte buffer.
     pub fn data_mut(&mut self) -> &mut [u8] {
-        &mut self.bytes[..]
+        // Safety: either is backed by a `Vec` or a static buffer, ptr[0..len] is valid.
+        unsafe { core::slice::from_raw_parts_mut(self.ptr, self.len) }
+    }
+}
+
+impl Drop for ByteBuffer {
+    fn drop(&mut self) {
+        if !self.is_static {
+            // Safety: those parts have been obtained from `Vec`.
+            unsafe { Vec::from_raw_parts(self.ptr, self.len, self.capacity) };
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    #[test]
+    fn test_basic_allocation_deallocation() {
+        let buffer = ByteBuffer::new(10);
+        assert_eq!(buffer.len(), 10);
+        // Dropping the buffer should not cause UB.
+    }
+
+    #[test]
+    fn test_basic_data_manipulation() {
+        let mut buffer = ByteBuffer::new(10);
+        assert_eq!(buffer.len(), 10);
+        let data = buffer.data(); // test we can read the data
+        assert_eq!(data, &[0; 10]);
+        let data = buffer.data_mut(); // test we can take a mutable reference to the data
+        data[4] = 4; // test we can write to the data and it is not UB
+        let data = buffer.data(); // test we can take a new reference to the data
+        assert_eq!(data, &[0, 0, 0, 0, 4, 0, 0, 0, 0, 0]); // test we can read the data
+                                                           // test drop is okay
+    }
+
+    #[test]
+    fn test_static_buffer_initialization() {
+        static mut BUF: [u8; 10] = [7; 10];
+        let buf = unsafe { &mut *core::ptr::addr_of_mut!(BUF) };
+        let mut buffer = ByteBuffer::new_static(buf, 5);
+        assert_eq!(buffer.len(), 5);
+        // Modifying the static buffer through ByteBuffer and checking its content.
+        let data = buffer.data_mut();
+        data[0] = 1;
+        unsafe {
+            assert_eq!(BUF[0], 1);
+        }
+    }
+
+    #[test]
+    fn test_growing_buffer() {
+        let mut buffer = ByteBuffer::new(5);
+        buffer.grow(10);
+        assert_eq!(buffer.len(), 10);
+        assert_eq!(buffer.data(), &[0; 10]);
+    }
+
+    #[test]
+    fn test_growing_static() {
+        static mut BUF: [u8; 10] = [7; 10];
+        let buf = unsafe { &mut *core::ptr::addr_of_mut!(BUF) };
+        let mut buffer = ByteBuffer::new_static(buf, 5);
+        assert_eq!(buffer.len(), 5);
+        assert_eq!(buffer.data(), &[0; 5]);
+        buffer.grow(8);
+        assert_eq!(buffer.len(), 8);
+        assert_eq!(buffer.data(), &[0; 8]);
+        buffer.grow(10);
+        assert_eq!(buffer.len(), 10);
+        assert_eq!(buffer.data(), &[0; 10]);
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_static_buffer_overflow() {
+        static mut BUF: [u8; 5] = [7; 5];
+        let buf = unsafe { &mut *core::ptr::addr_of_mut!(BUF) };
+        let mut buffer = ByteBuffer::new_static(buf, 5);
+        buffer.grow(10); // This should panic.
     }
 }

crates/wasmi/src/memory/error.rs

@@ -22,6 +22,8 @@ pub enum MemoryError {
     },
     /// Tried to create too many memories
     TooManyMemories,
+    /// Tried to create memory with invalid static buffer size
+    InvalidStaticBufferSize,
 }
 
 impl Display for MemoryError {
@@ -45,6 +47,9 @@ impl Display for MemoryError {
             Self::TooManyMemories => {
                 write!(f, "too many memories")
             }
+            Self::InvalidStaticBufferSize => {
+                write!(f, "tried to use too small static buffer")
+            }
         }
     }
 }

crates/wasmi/src/memory/mod.rs

@@ -165,6 +165,45 @@ impl MemoryEntity {
         }
     }
 
+    /// Creates a new memory entity with the given memory type.
+    pub fn new_static(
+        memory_type: MemoryType,
+        limiter: &mut ResourceLimiterRef<'_>,
+        buf: &'static mut [u8],
+    ) -> Result<Self, MemoryError> {
+        let initial_pages = memory_type.initial_pages();
+        let initial_len = initial_pages.to_bytes();
+        let maximum_pages = memory_type.maximum_pages().unwrap_or_else(Pages::max);
+        let maximum_len = maximum_pages.to_bytes();
+
+        if let Some(limiter) = limiter.as_resource_limiter() {
+            if !limiter.memory_growing(0, initial_len.unwrap_or(usize::MAX), maximum_len)? {
+                // Here there's no meaningful way to map Ok(false) to
+                // INVALID_GROWTH_ERRCODE, so we just translate it to an
+                // appropriate Err(...)
+                return Err(MemoryError::OutOfBoundsAllocation);
+            }
+        }
+
+        if let Some(initial_len) = initial_len {
+            if buf.len() < initial_len {
+                return Err(MemoryError::InvalidStaticBufferSize);
+            }
+            let memory = Self {
+                bytes: ByteBuffer::new_static(buf, initial_len),
+                memory_type,
+                current_pages: initial_pages,
+            };
+            Ok(memory)
+        } else {
+            let err = MemoryError::OutOfBoundsAllocation;
+            if let Some(limiter) = limiter.as_resource_limiter() {
+                limiter.memory_grow_failed(&err)
+            }
+            Err(err)
+        }
+    }
+
     /// Returns the memory type of the linear memory.
     pub fn ty(&self) -> MemoryType {
         self.memory_type
@@ -339,6 +378,27 @@ impl Memory {
         Ok(memory)
     }
 
+    /// Creates a new linear memory to the store.
+    ///
+    /// # Errors
+    ///
+    /// If more than [`u32::MAX`] much linear memory is allocated.
+    /// - If static buffer is invalid
+    pub fn new_static(
+        mut ctx: impl AsContextMut,
+        ty: MemoryType,
+        buf: &'static mut [u8],
+    ) -> Result<Self, MemoryError> {
+        let (inner, mut resource_limiter) = ctx
+            .as_context_mut()
+            .store
+            .store_inner_and_resource_limiter_ref();
+
+        let entity = MemoryEntity::new_static(ty, &mut resource_limiter, buf)?;
+        let memory = inner.alloc_memory(entity);
+        Ok(memory)
+    }
+
     /// Returns the memory type of the linear memory.
     ///
     /// # Panics