Make wasmtime_environ::Module serializable (#2005)

* Define WasmType/WasmFuncType in the Cranelift
* Make `Module` serializable

crates/environ/Cargo.toml

@@ -19,7 +19,7 @@ cranelift-frontend = { path = "../../cranelift/frontend", version = "0.65.0" }
 cranelift-wasm = { path = "../../cranelift/wasm", version = "0.65.0", features = ["enable-serde"] }
 wasmparser = "0.58.0"
 lightbeam = { path = "../lightbeam", optional = true, version = "0.18.0" }
-indexmap = "1.0.2"
+indexmap = { version = "1.0.2", features = ["serde-1"] }
 rayon = { version = "1.2.1", optional = true }
 thiserror = "1.0.4"
 directories = "2.0.1"

crates/environ/src/module.rs

@@ -11,6 +11,7 @@ use cranelift_wasm::{
 };
 use indexmap::IndexMap;
 use more_asserts::assert_ge;
+use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::sync::{
     atomic::{AtomicUsize, Ordering::SeqCst},
@@ -18,7 +19,7 @@ use std::sync::{
 };
 
 /// A WebAssembly table initializer.
-#[derive(Clone, Debug, Hash)]
+#[derive(Clone, Debug, Hash, Serialize, Deserialize)]
 pub struct TableElements {
     /// The index of a table to initialize.
     pub table_index: TableIndex,
@@ -31,7 +32,7 @@ pub struct TableElements {
 }
 
 /// An index of an entity.
-#[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
+#[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
 pub enum EntityIndex {
     /// Function index.
     Function(FuncIndex),
@@ -44,7 +45,7 @@ pub enum EntityIndex {
 }
 
 /// Implemenation styles for WebAssembly linear memory.
-#[derive(Debug, Clone, Hash)]
+#[derive(Debug, Clone, Hash, Serialize, Deserialize)]
 pub enum MemoryStyle {
     /// The actual memory can be resized and moved.
     Dynamic,
@@ -80,7 +81,7 @@ impl MemoryStyle {
 
 /// A WebAssembly linear memory description along with our chosen style for
 /// implementing it.
-#[derive(Debug, Clone, Hash)]
+#[derive(Debug, Clone, Hash, Serialize, Deserialize)]
 pub struct MemoryPlan {
     /// The WebAssembly linear memory description.
     pub memory: Memory,
@@ -103,7 +104,7 @@ impl MemoryPlan {
 }
 
 /// Implemenation styles for WebAssembly tables.
-#[derive(Debug, Clone, Hash)]
+#[derive(Debug, Clone, Hash, Serialize, Deserialize)]
 pub enum TableStyle {
     /// Signatures are stored in the table and checked in the caller.
     CallerChecksSignature,
@@ -118,7 +119,7 @@ impl TableStyle {
 
 /// A WebAssembly table description along with our chosen style for
 /// implementing it.
-#[derive(Debug, Clone, Hash)]
+#[derive(Debug, Clone, Hash, Serialize, Deserialize)]
 pub struct TablePlan {
     /// The WebAssembly table description.
     pub table: cranelift_wasm::Table,
@@ -136,9 +137,10 @@ impl TablePlan {
 
 /// A translated WebAssembly module, excluding the function bodies and
 /// memory initializers.
-#[derive(Debug)]
+#[derive(Debug, Serialize, Deserialize)]
 pub struct Module {
     /// A unique identifier (within this process) for this module.
+    #[serde(skip_serializing, default = "Module::next_id")]
     pub id: usize,
 
     /// The name of this wasm module, often found in the wasm file.
@@ -166,6 +168,7 @@ pub struct Module {
     pub passive_elements: HashMap<ElemIndex, Box<[FuncIndex]>>,
 
     /// WebAssembly passive data segments.
+    #[serde(with = "passive_data_serde")]
     pub passive_data: HashMap<DataIndex, Arc<[u8]>>,
 
     /// WebAssembly table initializers.
@@ -178,7 +181,7 @@ pub struct Module {
 /// This is stored within a `Module` and it implements `Hash`, unlike `Module`,
 /// and is used as part of the cache key when we load compiled modules from the
 /// global cache.
-#[derive(Debug, Hash)]
+#[derive(Debug, Hash, Serialize, Deserialize)]
 pub struct ModuleLocal {
     /// Unprocessed signatures exactly as provided by `declare_signature()`.
     pub signatures: PrimaryMap<SignatureIndex, (WasmFuncType, ir::Signature)>,
@@ -211,10 +214,8 @@ pub struct ModuleLocal {
 impl Module {
     /// Allocates the module data structures.
     pub fn new() -> Self {
-        static NEXT_ID: AtomicUsize = AtomicUsize::new(0);
-
         Self {
-            id: NEXT_ID.fetch_add(1, SeqCst),
+            id: Self::next_id(),
             name: None,
             imports: Vec::new(),
             exports: IndexMap::new(),
@@ -241,6 +242,11 @@ impl Module {
     pub fn get_passive_element(&self, index: ElemIndex) -> Option<&[FuncIndex]> {
         self.passive_elements.get(&index).map(|es| &**es)
     }
+
+    fn next_id() -> usize {
+        static NEXT_ID: AtomicUsize = AtomicUsize::new(0);
+        NEXT_ID.fetch_add(1, SeqCst)
+    }
 }
 
 impl ModuleLocal {
@@ -344,3 +350,49 @@ impl ModuleLocal {
         &self.signatures[self.functions[func_index]].0
     }
 }
+
+mod passive_data_serde {
+    use super::{Arc, DataIndex, HashMap};
+    use serde::{de::MapAccess, de::Visitor, ser::SerializeMap, Deserializer, Serializer};
+    use std::fmt;
+
+    pub(super) fn serialize<S>(
+        data: &HashMap<DataIndex, Arc<[u8]>>,
+        ser: S,
+    ) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        let mut map = ser.serialize_map(Some(data.len()))?;
+        for (k, v) in data {
+            map.serialize_entry(k, v.as_ref())?;
+        }
+        map.end()
+    }
+
+    struct PassiveDataVisitor;
+    impl<'de> Visitor<'de> for PassiveDataVisitor {
+        type Value = HashMap<DataIndex, Arc<[u8]>>;
+
+        fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+            formatter.write_str("a passive_data map")
+        }
+        fn visit_map<M>(self, mut access: M) -> Result<Self::Value, M::Error>
+        where
+            M: MapAccess<'de>,
+        {
+            let mut map = HashMap::with_capacity(access.size_hint().unwrap_or(0));
+            while let Some((key, value)) = access.next_entry::<_, Vec<u8>>()? {
+                map.insert(key, value.into());
+            }
+            Ok(map)
+        }
+    }
+
+    pub(super) fn deserialize<'de, D>(de: D) -> Result<HashMap<DataIndex, Arc<[u8]>>, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        de.deserialize_map(PassiveDataVisitor)
+    }
+}

crates/environ/src/module_environ.rs

@@ -110,14 +110,10 @@ impl<'data> cranelift_wasm::ModuleEnvironment<'data> for ModuleEnvironment<'data
         Ok(())
     }
 
-    fn declare_signature(&mut self, wasm: &WasmFuncType, sig: ir::Signature) -> WasmResult<()> {
+    fn declare_signature(&mut self, wasm: WasmFuncType, sig: ir::Signature) -> WasmResult<()> {
         let sig = translate_signature(sig, self.pointer_type());
         // TODO: Deduplicate signatures.
-        self.result
-            .module
-            .local
-            .signatures
-            .push((wasm.clone(), sig));
+        self.result.module.local.signatures.push((wasm, sig));
         Ok(())
     }
 

crates/wasmtime/src/externals.rs

@@ -202,13 +202,11 @@ impl Global {
         // The original export is coming from wasmtime_runtime itself we should
         // support all the types coming out of it, so assert such here.
         GlobalType::from_wasmtime_global(&self.wasmtime_export.global)
-            .expect("core wasm global type should be supported")
     }
 
     /// Returns the value type of this `global`.
     pub fn val_type(&self) -> ValType {
         ValType::from_wasm_type(&self.wasmtime_export.global.wasm_ty)
-            .expect("core wasm type should be supported")
     }
 
     /// Returns the underlying mutability of this `global`.

crates/wasmtime/src/func.rs

@@ -545,7 +545,7 @@ impl Func {
         // This is only called with `Export::Function`, and since it's coming
         // from wasmtime_runtime itself we should support all the types coming
         // out of it, so assert such here.
-        FuncType::from_wasm_func_type(&wft).expect("core wasm signature should be supported")
+        FuncType::from_wasm_func_type(&wft)
     }
 
     /// Returns the number of parameters that this function takes.

crates/wasmtime/src/types.rs

@@ -123,16 +123,15 @@ impl ValType {
         }
     }
 
-    pub(crate) fn from_wasm_type(ty: &wasm::WasmType) -> Option<Self> {
+    pub(crate) fn from_wasm_type(ty: &wasm::WasmType) -> Self {
         match ty {
-            wasm::WasmType::I32 => Some(Self::I32),
-            wasm::WasmType::I64 => Some(Self::I64),
-            wasm::WasmType::F32 => Some(Self::F32),
-            wasm::WasmType::F64 => Some(Self::F64),
-            wasm::WasmType::V128 => Some(Self::V128),
-            wasm::WasmType::FuncRef => Some(Self::FuncRef),
-            wasm::WasmType::ExternRef => Some(Self::ExternRef),
-            wasm::WasmType::Func | wasm::WasmType::EmptyBlockType => None,
+            wasm::WasmType::I32 => Self::I32,
+            wasm::WasmType::I64 => Self::I64,
+            wasm::WasmType::F32 => Self::F32,
+            wasm::WasmType::F64 => Self::F64,
+            wasm::WasmType::V128 => Self::V128,
+            wasm::WasmType::FuncRef => Self::FuncRef,
+            wasm::WasmType::ExternRef => Self::ExternRef,
         }
     }
 }
@@ -279,21 +278,21 @@ impl FuncType {
     /// Returns `None` if any types in the signature can't be converted to the
     /// types in this crate, but that should very rarely happen and largely only
     /// indicate a bug in our cranelift integration.
-    pub(crate) fn from_wasm_func_type(signature: &wasm::WasmFuncType) -> Option<FuncType> {
+    pub(crate) fn from_wasm_func_type(signature: &wasm::WasmFuncType) -> FuncType {
         let params = signature
             .params
             .iter()
             .map(|p| ValType::from_wasm_type(p))
-            .collect::<Option<Vec<_>>>()?;
+            .collect::<Vec<_>>();
         let results = signature
             .returns
             .iter()
             .map(|r| ValType::from_wasm_type(r))
-            .collect::<Option<Vec<_>>>()?;
-        Some(FuncType {
+            .collect::<Vec<_>>();
+        FuncType {
             params: params.into_boxed_slice(),
             results: results.into_boxed_slice(),
-        })
+        }
     }
 }
 
@@ -332,14 +331,14 @@ impl GlobalType {
 
     /// Returns `None` if the wasmtime global has a type that we can't
     /// represent, but that should only very rarely happen and indicate a bug.
-    pub(crate) fn from_wasmtime_global(global: &wasm::Global) -> Option<GlobalType> {
-        let ty = ValType::from_wasm_type(&global.wasm_ty)?;
+    pub(crate) fn from_wasmtime_global(global: &wasm::Global) -> GlobalType {
+        let ty = ValType::from_wasm_type(&global.wasm_ty);
         let mutability = if global.mutability {
             Mutability::Var
         } else {
             Mutability::Const
         };
-        Some(GlobalType::new(ty, mutability))
+        GlobalType::new(ty, mutability)
     }
 }
 
@@ -451,14 +450,10 @@ impl<'module> EntityType<'module> {
     /// Convert this `EntityType` to an `ExternType`.
     pub(crate) fn extern_type(&self) -> ExternType {
         match self {
-            EntityType::Function(sig) => FuncType::from_wasm_func_type(sig)
-                .expect("core wasm function type should be supported")
-                .into(),
+            EntityType::Function(sig) => FuncType::from_wasm_func_type(sig).into(),
             EntityType::Table(table) => TableType::from_wasmtime_table(table).into(),
             EntityType::Memory(memory) => MemoryType::from_wasmtime_memory(memory).into(),
-            EntityType::Global(global) => GlobalType::from_wasmtime_global(global)
-                .expect("core wasm global type should be supported")
-                .into(),
+            EntityType::Global(global) => GlobalType::from_wasmtime_global(global).into(),
         }
     }
 }