Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Zero memory copies on IO #3640

Merged
merged 14 commits into from
Mar 4, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
229 changes: 229 additions & 0 deletions lib/api/src/access.rs
Original file line number Diff line number Diff line change
@@ -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
}
}
40 changes: 40 additions & 0 deletions lib/api/src/js/mem_access.rs
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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> {
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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),
})
}
}
2 changes: 2 additions & 0 deletions lib/api/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Loading