Support writing to uninitialized memory

README.md

@@ -48,7 +48,7 @@ let transform = AffineTransform2D {
     translate: glam::Vec2::ZERO,
 };
 
-let mut buffer = UniformBuffer::new(Vec::new());
+let mut buffer = UniformBuffer::new(Vec::<u8>::new());
 buffer.write(&transform).unwrap();
 let byte_buffer = buffer.into_inner();
 
@@ -94,7 +94,7 @@ let mut positions = Positions {
     ])
 };
 
-let mut byte_buffer = Vec::new();
+let mut byte_buffer: Vec<u8> = Vec::new();
 
 let mut buffer = StorageBuffer::new(&mut byte_buffer);
 buffer.write(&positions).unwrap();
@@ -118,7 +118,7 @@ Write different data types to dynamic storage buffer
 ```rust
 use encase::{ShaderType, DynamicStorageBuffer};
 
-let mut byte_buffer = Vec::new();
+let mut byte_buffer: Vec<u8> = Vec::new();
 
 let mut buffer = DynamicStorageBuffer::new_with_alignment(&mut byte_buffer, 64);
 let offsets = [
@@ -133,6 +133,30 @@ assert_eq!(offsets, [0, 64, 192]);
 
 ```
 
+Supports writing to uninitialized memory as well.
+
+```rust
+use std::mem::MaybeUninit;
+use encase::{ShaderType, DynamicStorageBuffer};
+
+let mut byte_buffer: Vec<MaybeUninit<u8>> = Vec::new();
+
+let mut buffer = DynamicStorageBuffer::new_with_alignment(&mut byte_buffer, 64);
+let offsets = [
+    buffer.write(&[5.; 10]).unwrap(),
+    buffer.write(&vec![3u32; 20]).unwrap(),
+    buffer.write(&glam::Vec3::ONE).unwrap(),
+];
+
+// SAFETY: Vec<u8> and Vec<MaybeUninit<u8>> share the same layout.
+let byte_buffer: Vec<u8> = unsafe { std::mem::transmute(byte_buffer) };
+
+// write byte_buffer to GPU
+
+assert_eq!(offsets, [0, 64, 192]);
+
+```
+
 [host-shareable types]: https://gpuweb.github.io/gpuweb/wgsl/#host-shareable-types
 [features]: https://docs.rs/crate/encase/latest/features
 [`ShaderType`]: https://docs.rs/encase/latest/encase/trait.ShaderType.html

src/core/rw.rs

@@ -1,4 +1,5 @@
 use super::ShaderType;
+use core::mem::MaybeUninit;
 use thiserror::Error;
 
 #[derive(Clone, Copy, Debug, Error)]
@@ -231,6 +232,21 @@ impl BufferMut for [u8] {
     }
 }
 
+impl BufferMut for [MaybeUninit<u8>] {
+    #[inline]
+    fn capacity(&self) -> usize {
+        self.len()
+    }
+
+    #[inline]
+    fn write<const N: usize>(&mut self, offset: usize, val: &[u8; N]) {
+        use crate::utils::SliceExt;
+        // SAFETY: &[u8; N] and &[MaybeUninit<u8>; N] have the same layout
+        let val: &[MaybeUninit<u8>; N] = unsafe { core::mem::transmute(val) };
+        *self.array_mut(offset) = *val;
+    }
+}
+
 impl<const LEN: usize> BufferMut for [u8; LEN] {
     #[inline]
     fn capacity(&self) -> usize {
@@ -243,6 +259,18 @@ impl<const LEN: usize> BufferMut for [u8; LEN] {
     }
 }
 
+impl<const LEN: usize> BufferMut for [MaybeUninit<u8>; LEN] {
+    #[inline]
+    fn capacity(&self) -> usize {
+        <[MaybeUninit<u8>] as BufferMut>::capacity(self)
+    }
+
+    #[inline]
+    fn write<const N: usize>(&mut self, offset: usize, val: &[u8; N]) {
+        <[MaybeUninit<u8>] as BufferMut>::write(self, offset, val)
+    }
+}
+
 impl BufferMut for Vec<u8> {
     #[inline]
     fn capacity(&self) -> usize {
@@ -257,7 +285,25 @@ impl BufferMut for Vec<u8> {
     #[inline]
     fn try_enlarge(&mut self, wanted: usize) -> core::result::Result<(), EnlargeError> {
         use crate::utils::ByteVecExt;
-        self.try_extend_zeroed(wanted).map_err(EnlargeError::from)
+        self.try_extend(wanted).map_err(EnlargeError::from)
+    }
+}
+
+impl BufferMut for Vec<MaybeUninit<u8>> {
+    #[inline]
+    fn capacity(&self) -> usize {
+        self.capacity()
+    }
+
+    #[inline]
+    fn write<const N: usize>(&mut self, offset: usize, val: &[u8; N]) {
+        <[MaybeUninit<u8>] as BufferMut>::write(self, offset, val)
+    }
+
+    #[inline]
+    fn try_enlarge(&mut self, wanted: usize) -> core::result::Result<(), EnlargeError> {
+        use crate::utils::ByteVecExt;
+        self.try_extend(wanted).map_err(EnlargeError::from)
     }
 }
 

src/utils.rs

@@ -1,3 +1,5 @@
+use core::mem::MaybeUninit;
+
 #[track_caller]
 pub const fn consume_zsts<const N: usize>(_: [(); N]) {}
 
@@ -63,17 +65,11 @@ macro_rules! array_mut_to_2d_array_mut {
 
 pub(crate) trait ByteVecExt {
     /// Tries to extend `self` with `0`s up to `new_len`, using memset.
-    fn try_extend_zeroed(
-        &mut self,
-        new_len: usize,
-    ) -> Result<(), std::collections::TryReserveError>;
+    fn try_extend(&mut self, new_len: usize) -> Result<(), std::collections::TryReserveError>;
 }
 
 impl ByteVecExt for Vec<u8> {
-    fn try_extend_zeroed(
-        &mut self,
-        new_len: usize,
-    ) -> Result<(), std::collections::TryReserveError> {
+    fn try_extend(&mut self, new_len: usize) -> Result<(), std::collections::TryReserveError> {
         let additional = new_len.saturating_sub(self.len());
         if additional > 0 {
             self.try_reserve(additional)?;
@@ -92,6 +88,25 @@ impl ByteVecExt for Vec<u8> {
     }
 }
 
+impl<T> ByteVecExt for Vec<MaybeUninit<T>> {
+    fn try_extend(&mut self, new_len: usize) -> Result<(), std::collections::TryReserveError> {
+        let additional = new_len.saturating_sub(self.len());
+        if additional > 0 {
+            self.try_reserve(additional)?;
+
+            // It's OK to not initialize the extended elements as MaybeUninit allows
+            // uninitialized memory.
+
+            // SAFETY
+            // 1. new_len is less than or equal to Vec::capacity() since we reserved at least `additional` elements
+            // 2. The elements at old_len..new_len are initialized since we wrote `additional` bytes
+            // 3. MaybeUninit
+            unsafe { self.set_len(new_len) }
+        }
+        Ok(())
+    }
+}
+
 pub(crate) trait SliceExt<T> {
     /// Returns a "window" (shared reference to an array of length `N`) into this slice.
     ///
@@ -139,29 +154,29 @@ mod byte_vec_ext {
     use crate::utils::ByteVecExt;
 
     #[test]
-    fn try_extend_zeroed() {
-        let mut vec = Vec::new();
+    fn try_extend() {
+        let mut vec: Vec<u8> = Vec::new();
 
-        vec.try_extend_zeroed(10).unwrap();
+        vec.try_extend(10).unwrap();
 
         assert_eq!(vec.len(), 10);
         assert!(vec.iter().all(|val| *val == 0));
     }
 
     #[test]
-    fn try_extend_zeroed_noop() {
+    fn try_extend_noop() {
         let mut vec = vec![0; 12];
 
-        vec.try_extend_zeroed(10).unwrap();
+        vec.try_extend(10).unwrap();
 
         assert_eq!(vec.len(), 12);
     }
 
     #[test]
-    fn try_extend_zeroed_err() {
+    fn try_extend_err() {
         let mut vec = vec![0; 12];
 
-        assert!(vec.try_extend_zeroed(usize::MAX).is_err());
+        assert!(vec.try_extend(usize::MAX).is_err());
     }
 }