From 6c70b81ff55fc8dff574dccd0408fccff7d404bc Mon Sep 17 00:00:00 2001
From: Pat Hickey <phickey@fastly.com>
Date: Thu, 21 Oct 2021 12:04:20 -0700
Subject: [PATCH] review feedback

---
 crates/runtime/src/instance/allocator.rs | 21 +++++++++------------
 crates/runtime/src/memory.rs             |  3 +++
 crates/wasmtime/src/memory.rs            | 10 +++++++++-
 3 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/crates/runtime/src/instance/allocator.rs b/crates/runtime/src/instance/allocator.rs
index fead139467..abcb94b11f 100644
--- a/crates/runtime/src/instance/allocator.rs
+++ b/crates/runtime/src/instance/allocator.rs
@@ -97,10 +97,10 @@ impl StorePtr {
         self.0.clone()
     }
     /// Use the StorePtr as a mut ref to the Store.
-    // XXX should this be an unsafe fn? is it always safe at a use site?
-    pub(crate) fn get(&mut self) -> Option<&mut dyn Store> {
+    /// Safety: must not be used outside the original lifetime of the borrow.
+    pub(crate) unsafe fn get(&mut self) -> Option<&mut dyn Store> {
         match self.0 {
-            Some(ptr) => Some(unsafe { &mut *ptr }),
+            Some(ptr) => Some(&mut *ptr),
             None => None,
         }
     }
@@ -630,12 +630,11 @@ impl OnDemandInstanceAllocator {
             PrimaryMap::with_capacity(module.table_plans.len() - num_imports);
         for table in &module.table_plans.values().as_slice()[num_imports..] {
             tables.push(
-                Table::new_dynamic(
-                    table,
+                Table::new_dynamic(table, unsafe {
                     store
                         .get()
-                        .expect("if module has table plans, store is not empty"),
-                )
+                        .expect("if module has table plans, store is not empty")
+                })
                 .map_err(InstantiationError::Resource)?,
             );
         }
@@ -656,13 +655,11 @@ impl OnDemandInstanceAllocator {
             PrimaryMap::with_capacity(module.memory_plans.len() - num_imports);
         for plan in &module.memory_plans.values().as_slice()[num_imports..] {
             memories.push(
-                Memory::new_dynamic(
-                    plan,
-                    creator,
+                Memory::new_dynamic(plan, creator, unsafe {
                     store
                         .get()
-                        .expect("if module has memory plans, store is not empty"),
-                )
+                        .expect("if module has memory plans, store is not empty")
+                })
                 .map_err(InstantiationError::Resource)?,
             );
         }
diff --git a/crates/runtime/src/memory.rs b/crates/runtime/src/memory.rs
index 7e2180e5d3..5a69d2136e 100644
--- a/crates/runtime/src/memory.rs
+++ b/crates/runtime/src/memory.rs
@@ -378,6 +378,9 @@ impl Memory {
     ///
     /// Generally, prefer using `InstanceHandle::memory_grow`, which encapsulates
     /// this unsafety.
+    ///
+    /// Ensure that the provided Store is not used to get access any Memory
+    /// which lives inside it.
     pub unsafe fn grow(
         &mut self,
         delta_pages: u64,
diff --git a/crates/wasmtime/src/memory.rs b/crates/wasmtime/src/memory.rs
index f541c2f81b..3f22f3dd2f 100644
--- a/crates/wasmtime/src/memory.rs
+++ b/crates/wasmtime/src/memory.rs
@@ -223,8 +223,12 @@ impl Memory {
         Memory::_new(store.as_context_mut().0, ty)
     }
 
-    /// Async variant of [`Memory::new`]. You must use this variant with Stores which have a
+    /// Async variant of [`Memory::new`]. You must use this variant with [`Store`]s which have a
     /// [`ResourceLimiterAsync`].
+    ///
+    /// # Panics
+    ///
+    /// This function will panic when used with a non-async [`Store`].
     #[cfg(feature = "async")]
     pub async fn new_async<T>(
         mut store: impl AsContextMut<Data = T>,
@@ -491,6 +495,10 @@ impl Memory {
     }
 
     /// Async variant of [`Memory::grow`]. Required when using a [`ResourceLimiterAsync`].
+    ///
+    /// # Panics
+    ///
+    /// This function will panic when used with a non-async [`Store`].
     #[cfg(feature = "async")]
     pub async fn grow_async<T>(
         &self,