From 753b05d8354d0633705f16fe9855ae62174b3297 Mon Sep 17 00:00:00 2001
From: Jeron Aldaron Lau <aldaronlau@gmail.com>
Date: Sat, 17 Sep 2022 16:51:02 -0500
Subject: [PATCH] Add `Memory::data_and_store_mut` (#448)

* Add `Memory::data_and_store_mut`

* Fix doc comments

* Fix comments for `data*` methods on `Memory`
---
 wasmi_v1/src/memory/mod.rs | 18 ++++++++++++++++--
 wasmi_v1/src/store.rs      | 18 ++++++++++++++++++
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/wasmi_v1/src/memory/mod.rs b/wasmi_v1/src/memory/mod.rs
index 45be0f9ae0..6c9c5b01b2 100644
--- a/wasmi_v1/src/memory/mod.rs
+++ b/wasmi_v1/src/memory/mod.rs
@@ -311,7 +311,7 @@ impl Memory {
             .grow(additional)
     }
 
-    /// Returns a shared slice to the bytes underlying to the byte buffer.
+    /// Returns a shared slice to the bytes underlying the [`Memory`].
     ///
     /// # Panics
     ///
@@ -320,7 +320,7 @@ impl Memory {
         ctx.into().store.resolve_memory(*self).data()
     }
 
-    /// Returns an exclusive slice to the bytes underlying to the byte buffer.
+    /// Returns an exclusive slice to the bytes underlying the [`Memory`].
     ///
     /// # Panics
     ///
@@ -329,6 +329,20 @@ impl Memory {
         ctx.into().store.resolve_memory_mut(*self).data_mut()
     }
 
+    /// Returns an exclusive slice to the bytes underlying the [`Memory`], and an exclusive
+    /// reference to the user provided state.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `ctx` does not own this [`Memory`].
+    pub fn data_and_store_mut<'a, T: 'a>(
+        &self,
+        ctx: impl Into<StoreContextMut<'a, T>>,
+    ) -> (&'a mut [u8], &'a mut T) {
+        let (memory, store) = ctx.into().store.resolve_memory_and_state_mut(*self);
+        (memory.data_mut(), store)
+    }
+
     /// Reads `n` bytes from `memory[offset..offset+n]` into `buffer`
     /// where `n` is the length of `buffer`.
     ///
diff --git a/wasmi_v1/src/store.rs b/wasmi_v1/src/store.rs
index 78aafbbf45..a541db24af 100644
--- a/wasmi_v1/src/store.rs
+++ b/wasmi_v1/src/store.rs
@@ -299,6 +299,24 @@ impl<T> Store<T> {
             .unwrap_or_else(|| panic!("failed to resolve stored linear memory: {:?}", entity_index))
     }
 
+    /// Returns an exclusive reference to the associated entity of the linear memory and an
+    /// exclusive reference to the user provided state.
+    ///
+    /// # Panics
+    ///
+    /// - If the linear memory does not originate from this store.
+    /// - If the linear memory cannot be resolved to its entity.
+    pub(super) fn resolve_memory_and_state_mut(
+        &mut self,
+        memory: Memory,
+    ) -> (&mut MemoryEntity, &mut T) {
+        let entity_index = self.unwrap_index(memory.into_inner());
+        let memory_entity = self.memories.get_mut(entity_index).unwrap_or_else(|| {
+            panic!("failed to resolve stored linear memory: {:?}", entity_index)
+        });
+        (memory_entity, &mut self.user_state)
+    }
+
     /// Returns a shared reference to the associated entity of the Wasm or host function.
     ///
     /// # Panics