diff --git a/crates/wasmtime/src/memory.rs b/crates/wasmtime/src/memory.rs index 4b94dad250..7ea2c96a53 100644 --- a/crates/wasmtime/src/memory.rs +++ b/crates/wasmtime/src/memory.rs @@ -338,6 +338,39 @@ impl Memory { } } + /// Same as [`Memory::data_mut`], but also returns the `T` from the + /// [`StoreContextMut`]. + /// + /// This method can be used when you want to simultaneously work with the + /// `T` in the store as well as the memory behind this [`Memory`]. Using + /// [`Memory::data_mut`] would consider the entire store borrowed, whereas + /// this method allows the Rust compiler to see that the borrow of this + /// memory and the borrow of `T` are disjoint. + /// + /// # Panics + /// + /// Panics if this memory doesn't belong to `store`. + pub fn data_and_store_mut<'a, T: 'a>( + &self, + store: impl Into>, + ) -> (&'a mut [u8], &'a mut T) { + // Note the unsafety here. Our goal is to simultaneously borrow the + // memory and custom data from `store`, and the store it's connected + // to. Rust will not let us do that, however, because we must call two + // separate methods (both of which borrow the whole `store`) and one of + // our borrows is mutable (the custom data). + // + // This operation, however, is safe because these borrows do not overlap + // and in the process of borrowing them mutability doesn't actually + // touch anything. This is akin to mutably borrowing two indices in an + // array, which is safe so long as the indices are separate. + unsafe { + let mut store = store.into(); + let data = &mut *(store.data_mut() as *mut T); + (self.data_mut(store), data) + } + } + /// Returns the base pointer, in the host's address space, that the memory /// is located at. /// diff --git a/crates/wiggle/generate/src/wasmtime.rs b/crates/wiggle/generate/src/wasmtime.rs index 76c82baf20..9c8590f2fb 100644 --- a/crates/wiggle/generate/src/wasmtime.rs +++ b/crates/wiggle/generate/src/wasmtime.rs @@ -118,25 +118,9 @@ fn generate_func( return Err(#rt::wasmtime_crate::Trap::new("missing required memory export")); } }; - // Note the unsafety here. Our goal is to simultaneously borrow the - // memory and custom data from `caller`, and the store it's connected - // to. Rust will not let us do that, however, because we must call two - // separate methods (both of which borrow the whole `caller`) and one of - // our borrows is mutable (the custom data). - // - // This operation, however, is safe because these borrows do not overlap - // and in the process of borrowing them mutability doesn't actually - // touch anything. This is akin to mutably borrowing two indices in an - // array, which is safe so long as the indices are separate. - // - // TODO: depending on how common this is for other users to run into we - // may wish to consider adding a dedicated method for this. For now the - // future of `GuestPtr` may be a bit hazy, so let's just get this - // working from the previous iteration for now. - let (ctx, mem) = unsafe { - let mem = &mut *(mem.data_mut(&mut caller) as *mut [u8]); - (get_cx(caller.data_mut()), #rt::wasmtime::WasmtimeGuestMemory::new(mem)) - }; + let (mem , ctx) = mem.data_and_store_mut(&mut caller); + let ctx = get_cx(ctx); + let mem = #rt::wasmtime::WasmtimeGuestMemory::new(mem); match #abi_func(ctx, &mem #(, #arg_names)*) #await_ { Ok(r) => Ok(<#ret_ty>::from(r)), Err(#rt::Trap::String(err)) => Err(#rt::wasmtime_crate::Trap::new(err)),