From a22561f691ffb5cc16c220a0b3dfdd1885113dad Mon Sep 17 00:00:00 2001
From: Nick Fitzgerald <fitzgen@gmail.com>
Date: Mon, 25 Sep 2023 12:20:35 -0700
Subject: [PATCH] Wasmtime: Add `hash_key` methods for `Func` and `Table`

---
 crates/wasmtime/src/externals/table.rs | 53 +++++++++++++++++++++++++
 crates/wasmtime/src/func.rs            | 54 ++++++++++++++++++++++++++
 2 files changed, 107 insertions(+)

diff --git a/crates/wasmtime/src/externals/table.rs b/crates/wasmtime/src/externals/table.rs
index 3505cb5dab24..3404324cd3b2 100644
--- a/crates/wasmtime/src/externals/table.rs
+++ b/crates/wasmtime/src/externals/table.rs
@@ -346,4 +346,57 @@ impl Table {
             vmctx: export.vmctx,
         }
     }
+
+    /// Get a stable hash key for this table.
+    ///
+    /// Even if the same underlying table definition is added to the
+    /// `StoreData` multiple times and becomes multiple `wasmtime::Table`s,
+    /// this hash key will be consistent across all of these tables.
+    #[allow(dead_code)] // Not used yet, but added for consistency.
+    pub(crate) fn hash_key(&self, store: &StoreOpaque) -> impl std::hash::Hash + Eq {
+        store[self.0].definition as usize
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{Instance, Module, Store};
+
+    #[test]
+    fn hash_key_is_stable_across_duplicate_store_data_entries() -> Result<()> {
+        let mut store = Store::<()>::default();
+        let module = Module::new(
+            store.engine(),
+            r#"
+                (module
+                    (table (export "t") 1 1 externref)
+                )
+            "#,
+        )?;
+        let instance = Instance::new(&mut store, &module, &[])?;
+
+        // Each time we `get_table`, we call `Table::from_wasmtime` which adds
+        // a new entry to `StoreData`, so `t1` and `t2` will have different
+        // indices into `StoreData`.
+        let t1 = instance.get_table(&mut store, "t").unwrap();
+        let t2 = instance.get_table(&mut store, "t").unwrap();
+
+        // That said, they really point to the same table.
+        assert!(t1.get(&mut store, 0).unwrap().unwrap_externref().is_none());
+        assert!(t2.get(&mut store, 0).unwrap().unwrap_externref().is_none());
+        t1.set(&mut store, 0, Val::ExternRef(Some(ExternRef::new(42))))?;
+        assert!(t1.get(&mut store, 0).unwrap().unwrap_externref().is_some());
+        assert!(t2.get(&mut store, 0).unwrap().unwrap_externref().is_some());
+
+        // And therefore their hash keys are the same.
+        assert!(t1.hash_key(&store.as_context().0) == t2.hash_key(&store.as_context().0));
+
+        // But the hash keys are different from different tables.
+        let instance2 = Instance::new(&mut store, &module, &[])?;
+        let t3 = instance2.get_table(&mut store, "t").unwrap();
+        assert!(t1.hash_key(&store.as_context().0) != t3.hash_key(&store.as_context().0));
+
+        Ok(())
+    }
 }
diff --git a/crates/wasmtime/src/func.rs b/crates/wasmtime/src/func.rs
index 3ec071c72980..134a8321a59d 100644
--- a/crates/wasmtime/src/func.rs
+++ b/crates/wasmtime/src/func.rs
@@ -1339,6 +1339,16 @@ impl Func {
         // (unsafely), which should be safe since we just did the type check above.
         unsafe { Ok(TypedFunc::new_unchecked(*self)) }
     }
+
+    /// Get a stable hash key for this function.
+    ///
+    /// Even if the same underlying function is added to the `StoreData`
+    /// multiple times and becomes multiple `wasmtime::Func`s, this hash key
+    /// will be consistent across all of these functions.
+    #[allow(dead_code)] // Not used yet, but added for consistency.
+    pub(crate) fn hash_key(&self, store: &mut StoreOpaque) -> impl std::hash::Hash + Eq {
+        self.caller_checked_func_ref(store).as_ptr() as usize
+    }
 }
 
 /// Prepares for entrance into WebAssembly.
@@ -2391,3 +2401,47 @@ mod rooted {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{Instance, Module, Store};
+
+    #[test]
+    fn hash_key_is_stable_across_duplicate_store_data_entries() -> Result<()> {
+        let mut store = Store::<()>::default();
+        let module = Module::new(
+            store.engine(),
+            r#"
+                (module
+                    (func (export "f")
+                        nop
+                    )
+                )
+            "#,
+        )?;
+        let instance = Instance::new(&mut store, &module, &[])?;
+
+        // Each time we `get_func`, we call `Func::from_wasmtime` which adds a
+        // new entry to `StoreData`, so `f1` and `f2` will have different
+        // indices into `StoreData`.
+        let f1 = instance.get_func(&mut store, "f").unwrap();
+        let f2 = instance.get_func(&mut store, "f").unwrap();
+
+        // But their hash keys are the same.
+        assert!(
+            f1.hash_key(&mut store.as_context_mut().0)
+                == f2.hash_key(&mut store.as_context_mut().0)
+        );
+
+        // But the hash keys are different from different funcs.
+        let instance2 = Instance::new(&mut store, &module, &[])?;
+        let f3 = instance2.get_func(&mut store, "f").unwrap();
+        assert!(
+            f1.hash_key(&mut store.as_context_mut().0)
+                != f3.hash_key(&mut store.as_context_mut().0)
+        );
+
+        Ok(())
+    }
+}