Merge pull request #3640 from wasmerio/zero-mem-copy

Zero memory copies on IO

lib/api/src/access.rs

@@ -0,0 +1,229 @@
+use std::mem::MaybeUninit;
+
+use crate::{WasmRef, WasmSlice};
+
+pub(super) enum SliceCow<'a, T> {
+    #[allow(dead_code)]
+    Borrowed(&'a mut [T]),
+    #[allow(dead_code)]
+    Owned(Vec<T>, bool),
+}
+
+impl<'a, T> AsRef<[T]> for SliceCow<'a, T> {
+    fn as_ref(&self) -> &[T] {
+        match self {
+            Self::Borrowed(buf) => buf,
+            Self::Owned(buf, _) => buf,
+        }
+    }
+}
+
+impl<'a, T> AsMut<[T]> for SliceCow<'a, T> {
+    fn as_mut(&mut self) -> &mut [T] {
+        // Note: Zero padding is not required here as its a typed copy which does
+        //       not leak the bytes into the memory
+        // https://stackoverflow.com/questions/61114026/does-stdptrwrite-transfer-the-uninitialized-ness-of-the-bytes-it-writes
+        match self {
+            Self::Borrowed(buf) => buf,
+            Self::Owned(buf, modified) => {
+                *modified = true;
+                buf.as_mut()
+            }
+        }
+    }
+}
+
+/// Provides direct memory access to a piece of memory that
+/// is owned by WASM
+pub struct WasmSliceAccess<'a, T>
+where
+    T: wasmer_types::ValueType,
+{
+    pub(super) slice: WasmSlice<'a, T>,
+    pub(super) buf: SliceCow<'a, T>,
+}
+
+impl<'a, T> AsRef<[T]> for WasmSliceAccess<'a, T>
+where
+    T: wasmer_types::ValueType,
+{
+    fn as_ref(&self) -> &[T] {
+        self.buf.as_ref()
+    }
+}
+
+impl<'a, T> AsMut<[T]> for WasmSliceAccess<'a, T>
+where
+    T: wasmer_types::ValueType,
+{
+    fn as_mut(&mut self) -> &mut [T] {
+        self.buf.as_mut()
+    }
+}
+
+impl<'a, T> WasmSliceAccess<'a, T>
+where
+    T: wasmer_types::ValueType,
+{
+    /// Returns an iterator of all the elements in the slice
+    pub fn iter(&'a self) -> std::slice::Iter<'a, T> {
+        self.as_ref().iter()
+    }
+
+    /// Returns an iterator of all the elements in the slice
+    pub fn iter_mut(&'a mut self) -> std::slice::IterMut<'a, T> {
+        self.buf.as_mut().iter_mut()
+    }
+
+    /// Number of elements in this slice
+    pub fn len(&self) -> usize {
+        self.buf.as_ref().len()
+    }
+
+    /// If the slice is empty
+    pub fn is_empty(&self) -> bool {
+        self.buf.as_ref().is_empty()
+    }
+}
+
+impl<'a> WasmSliceAccess<'a, u8> {
+    /// Writes to the address pointed to by this `WasmPtr` in a memory.
+    #[inline]
+    pub fn copy_from_slice(&mut self, src: &[u8]) {
+        let dst = self.buf.as_mut();
+        dst.copy_from_slice(src);
+    }
+}
+
+impl<'a, T> Drop for WasmSliceAccess<'a, T>
+where
+    T: wasmer_types::ValueType,
+{
+    fn drop(&mut self) {
+        if let SliceCow::Owned(buf, modified) = &self.buf {
+            if *modified {
+                self.slice.write_slice(buf.as_ref()).ok();
+            }
+        }
+    }
+}
+
+pub(super) enum RefCow<'a, T> {
+    #[allow(dead_code)]
+    Borrowed(&'a mut T),
+    #[allow(dead_code)]
+    Owned(T, bool),
+}
+
+impl<'a, T> AsRef<T> for RefCow<'a, T> {
+    fn as_ref(&self) -> &T {
+        match self {
+            Self::Borrowed(val) => *val,
+            Self::Owned(val, _) => val,
+        }
+    }
+}
+
+impl<'a, T> AsMut<T> for RefCow<'a, T> {
+    fn as_mut(&mut self) -> &mut T {
+        // Note: Zero padding is not required here as its a typed copy which does
+        //       not leak the bytes into the memory
+        // https://stackoverflow.com/questions/61114026/does-stdptrwrite-transfer-the-uninitialized-ness-of-the-bytes-it-writes
+        match self {
+            Self::Borrowed(val) => *val,
+            Self::Owned(val, modified) => {
+                *modified = true;
+                val
+            }
+        }
+    }
+}
+
+/// Provides direct memory access to a piece of memory that
+/// is owned by WASM
+pub struct WasmRefAccess<'a, T>
+where
+    T: wasmer_types::ValueType,
+{
+    pub(super) ptr: WasmRef<'a, T>,
+    pub(super) buf: RefCow<'a, T>,
+}
+
+impl<'a, T> AsRef<T> for WasmRefAccess<'a, T>
+where
+    T: wasmer_types::ValueType,
+{
+    fn as_ref(&self) -> &T {
+        self.buf.as_ref()
+    }
+}
+
+impl<'a, T> AsMut<T> for WasmRefAccess<'a, T>
+where
+    T: wasmer_types::ValueType,
+{
+    fn as_mut(&mut self) -> &mut T {
+        self.buf.as_mut()
+    }
+}
+
+impl<'a, T> WasmRefAccess<'a, T>
+where
+    T: wasmer_types::ValueType,
+{
+    /// Reads the address pointed to by this `WasmPtr` in a memory.
+    #[inline]
+    #[allow(clippy::clone_on_copy)]
+    pub fn read(&self) -> T
+    where
+        T: Clone,
+    {
+        self.as_ref().clone()
+    }
+
+    /// Writes to the address pointed to by this `WasmPtr` in a memory.
+    #[inline]
+    pub fn write(&mut self, val: T) {
+        // Note: Zero padding is not required here as its a typed copy which does
+        //       not leak the bytes into the memory
+        // https://stackoverflow.com/questions/61114026/does-stdptrwrite-transfer-the-uninitialized-ness-of-the-bytes-it-writes
+        *(self.as_mut()) = val;
+    }
+}
+
+impl<'a, T> Drop for WasmRefAccess<'a, T>
+where
+    T: wasmer_types::ValueType,
+{
+    fn drop(&mut self) {
+        if let RefCow::Owned(val, modified) = &self.buf {
+            if *modified {
+                self.ptr.write(*val).ok();
+            }
+        }
+    }
+}
+
+impl<'a, T> WasmSliceAccess<'a, T>
+where
+    T: wasmer_types::ValueType,
+{
+    /// Returns a mutable slice that is not yet initialized
+    pub fn as_mut_uninit(&mut self) -> &mut [MaybeUninit<T>] {
+        let ret: &mut [T] = self.buf.as_mut();
+        let ret: &mut [MaybeUninit<T>] = unsafe { std::mem::transmute(ret) };
+        ret
+    }
+}
+
+impl<'a, T> WasmRefAccess<'a, T>
+where
+    T: wasmer_types::ValueType,
+{
+    /// Returns a reference to an unitialized reference to this value
+    pub fn as_mut_uninit(&mut self) -> &mut MaybeUninit<T> {
+        let ret: &mut T = self.buf.as_mut();
+        let ret: &mut MaybeUninit<T> = unsafe { std::mem::transmute(ret) };
+        ret
+    }
+}

lib/api/src/js/mem_access.rs

@@ -1,6 +1,8 @@
+use crate::access::{RefCow, SliceCow, WasmRefAccess};
 use crate::js::externals::memory::MemoryBuffer;
 use crate::js::RuntimeError;
 use crate::js::{Memory32, Memory64, MemoryView, WasmPtr};
+use crate::WasmSliceAccess;
 use std::{
     convert::TryInto,
     fmt,
@@ -121,6 +123,12 @@ impl<'a, T: ValueType> WasmRef<'a, T> {
         let data = unsafe { slice::from_raw_parts(data.as_ptr() as *const _, data.len()) };
         self.buffer.write(self.offset, data)
     }
+
+    /// Gains direct access to the memory of this slice
+    #[inline]
+    pub fn access(self) -> Result<WasmRefAccess<'a, T>, MemoryAccessError> {
+        WasmRefAccess::new(self)
+    }
 }
 
 impl<'a, T: ValueType> fmt::Debug for WasmRef<'a, T> {
@@ -250,6 +258,12 @@ impl<'a, T: ValueType> WasmSlice<'a, T> {
         self.index(idx).write(val)
     }
 
+    /// Gains direct access to the memory of this slice
+    #[inline]
+    pub fn access(self) -> Result<WasmSliceAccess<'a, T>, MemoryAccessError> {
+        WasmSliceAccess::new(self)
+    }
+
     /// Reads the entire slice into the given buffer.
     ///
     /// The length of the buffer must match the length of the slice.
@@ -405,3 +419,29 @@ impl<'a, T: ValueType> DoubleEndedIterator for WasmSliceIter<'a, T> {
 }
 
 impl<'a, T: ValueType> ExactSizeIterator for WasmSliceIter<'a, T> {}
+
+impl<'a, T> WasmSliceAccess<'a, T>
+where
+    T: wasmer_types::ValueType,
+{
+    fn new(slice: WasmSlice<'a, T>) -> Result<Self, MemoryAccessError> {
+        let buf = slice.read_to_vec()?;
+        Ok(Self {
+            slice,
+            buf: SliceCow::Owned(buf, false),
+        })
+    }
+}
+
+impl<'a, T> WasmRefAccess<'a, T>
+where
+    T: wasmer_types::ValueType,
+{
+    fn new(ptr: WasmRef<'a, T>) -> Result<Self, MemoryAccessError> {
+        let val = ptr.read()?;
+        Ok(Self {
+            ptr,
+            buf: RefCow::Owned(val, false),
+        })
+    }
+}

lib/api/src/lib.rs

@@ -440,5 +440,7 @@ mod js;
 #[cfg(feature = "js")]
 pub use js::*;
 
+mod access;
 mod into_bytes;
+pub use access::WasmSliceAccess;
 pub use into_bytes::IntoBytes;