Add a safe method for accessing memory and T (#2971)

This is currently a very common operation in host bindings where if wasm
gives a host function a relative pointer you'll want to simulataneously
work with the host state and the wasm memory. These two regions are
distinct and safe to borrow mutably simulataneously but it's not obvious
in the Rust type system that this is so, so add a helper method here to
assist in doing so.

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<StoreContextMut<'a, T>>,
+    ) -> (&'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.
     ///

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)),