Refactor wasmtime's internal cache slightly (#2057)

Be more generic over the type being serialized, and remove an
intermediate struct which was only used for serialization but isn't
necessary.

crates/environ/src/cache.rs

@@ -1,8 +1,3 @@
-use crate::address_map::{ModuleAddressMap, ValueLabelsRanges};
-use crate::compilation::{Compilation, Relocations, StackMaps, Traps};
-use cranelift_codegen::ir;
-use cranelift_entity::PrimaryMap;
-use cranelift_wasm::DefinedFuncIndex;
 use log::{debug, trace, warn};
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
@@ -26,29 +21,6 @@ struct ModuleCacheEntryInner<'config> {
     cache_config: &'config CacheConfig,
 }
 
-/// Cached compilation data of a Wasm module.
-#[derive(Serialize, Deserialize, Debug, PartialEq, Eq)]
-pub struct ModuleCacheData {
-    compilation: Compilation,
-    relocations: Relocations,
-    address_transforms: ModuleAddressMap,
-    value_ranges: ValueLabelsRanges,
-    stack_slots: PrimaryMap<DefinedFuncIndex, ir::StackSlots>,
-    traps: Traps,
-    stack_maps: StackMaps,
-}
-
-/// A type alias over the module cache data as a tuple.
-pub type ModuleCacheDataTupleType = (
-    Compilation,
-    Relocations,
-    ModuleAddressMap,
-    ValueLabelsRanges,
-    PrimaryMap<DefinedFuncIndex, ir::StackSlots>,
-    Traps,
-    StackMaps,
-);
-
 struct Sha256Hasher(Sha256);
 
 impl<'config> ModuleCacheEntry<'config> {
@@ -68,11 +40,11 @@ impl<'config> ModuleCacheEntry<'config> {
         Self(Some(inner))
     }
 
-    pub fn get_data<T: Hash, E>(
-        &self,
-        state: T,
-        compute: fn(T) -> Result<ModuleCacheDataTupleType, E>,
-    ) -> Result<ModuleCacheData, E> {
+    pub fn get_data<T, U, E>(&self, state: T, compute: fn(T) -> Result<U, E>) -> Result<U, E>
+    where
+        T: Hash,
+        U: Serialize + for<'a> Deserialize<'a>,
+    {
         let mut hasher = Sha256Hasher(Sha256::new());
         state.hash(&mut hasher);
         let hash: [u8; 32] = hasher.0.result().into();
@@ -81,7 +53,7 @@ impl<'config> ModuleCacheEntry<'config> {
 
         let inner = match &self.0 {
             Some(inner) => inner,
-            None => return compute(state).map(ModuleCacheData::from_tuple),
+            None => return compute(state),
         };
 
         if let Some(cached_val) = inner.get_data(&hash) {
@@ -89,7 +61,7 @@ impl<'config> ModuleCacheEntry<'config> {
             inner.cache_config.on_cache_get_async(&mod_cache_path); // call on success
             return Ok(cached_val);
         }
-        let val_to_cache = ModuleCacheData::from_tuple(compute(state)?);
+        let val_to_cache = compute(state)?;
         if inner.update_data(&hash, &val_to_cache).is_some() {
             let mod_cache_path = inner.root_path.join(&hash);
             inner.cache_config.on_cache_update_async(&mod_cache_path); // call on success
@@ -141,7 +113,10 @@ impl<'config> ModuleCacheEntryInner<'config> {
         }
     }
 
-    fn get_data(&self, hash: &str) -> Option<ModuleCacheData> {
+    fn get_data<T>(&self, hash: &str) -> Option<T>
+    where
+        T: for<'a> Deserialize<'a>,
+    {
         let mod_cache_path = self.root_path.join(hash);
         trace!("get_data() for path: {}", mod_cache_path.display());
         let compressed_cache_bytes = fs::read(&mod_cache_path).ok()?;
@@ -153,7 +128,7 @@ impl<'config> ModuleCacheEntryInner<'config> {
             .ok()
     }
 
-    fn update_data(&self, hash: &str, data: &ModuleCacheData) -> Option<()> {
+    fn update_data<T: Serialize>(&self, hash: &str, data: &T) -> Option<()> {
         let mod_cache_path = self.root_path.join(hash);
         trace!("update_data() for path: {}", mod_cache_path.display());
         let serialized_data = bincode::serialize(&data)
@@ -197,32 +172,6 @@ impl<'config> ModuleCacheEntryInner<'config> {
     }
 }
 
-impl ModuleCacheData {
-    pub fn from_tuple(data: ModuleCacheDataTupleType) -> Self {
-        Self {
-            compilation: data.0,
-            relocations: data.1,
-            address_transforms: data.2,
-            value_ranges: data.3,
-            stack_slots: data.4,
-            traps: data.5,
-            stack_maps: data.6,
-        }
-    }
-
-    pub fn into_tuple(self) -> ModuleCacheDataTupleType {
-        (
-            self.compilation,
-            self.relocations,
-            self.address_transforms,
-            self.value_ranges,
-            self.stack_slots,
-            self.traps,
-            self.stack_maps,
-        )
-    }
-}
-
 impl Hasher for Sha256Hasher {
     fn finish(&self) -> u64 {
         panic!("Sha256Hasher doesn't support finish!");

crates/environ/src/cache/tests.rs

@@ -1,6 +1,5 @@
 use super::config::tests::test_prolog;
 use super::*;
-use cranelift_entity::PrimaryMap;
 use std::fs;
 
 // Since cache system is a global thing, each test needs to be run in seperate process.
@@ -66,40 +65,28 @@ fn test_write_read_cache() {
     let entry1 = ModuleCacheEntry::from_inner(ModuleCacheEntryInner::new(compiler1, &cache_config));
     let entry2 = ModuleCacheEntry::from_inner(ModuleCacheEntryInner::new(compiler2, &cache_config));
 
-    entry1.get_data(1, |_| new_module_cache_data()).unwrap();
-    entry1.get_data::<_, i32>(1, |_| panic!()).unwrap();
+    entry1.get_data::<_, i32, i32>(1, |_| Ok(100)).unwrap();
+    entry1.get_data::<_, i32, i32>(1, |_| panic!()).unwrap();
 
-    entry1.get_data(2, |_| new_module_cache_data()).unwrap();
-    entry1.get_data::<_, i32>(1, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32>(2, |_| panic!()).unwrap();
+    entry1.get_data::<_, i32, i32>(2, |_| Ok(100)).unwrap();
+    entry1.get_data::<_, i32, i32>(1, |_| panic!()).unwrap();
+    entry1.get_data::<_, i32, i32>(2, |_| panic!()).unwrap();
 
-    entry1.get_data(3, |_| new_module_cache_data()).unwrap();
-    entry1.get_data::<_, i32>(1, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32>(2, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32>(3, |_| panic!()).unwrap();
+    entry1.get_data::<_, i32, i32>(3, |_| Ok(100)).unwrap();
+    entry1.get_data::<_, i32, i32>(1, |_| panic!()).unwrap();
+    entry1.get_data::<_, i32, i32>(2, |_| panic!()).unwrap();
+    entry1.get_data::<_, i32, i32>(3, |_| panic!()).unwrap();
 
-    entry1.get_data(4, |_| new_module_cache_data()).unwrap();
-    entry1.get_data::<_, i32>(1, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32>(2, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32>(3, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32>(4, |_| panic!()).unwrap();
+    entry1.get_data::<_, i32, i32>(4, |_| Ok(100)).unwrap();
+    entry1.get_data::<_, i32, i32>(1, |_| panic!()).unwrap();
+    entry1.get_data::<_, i32, i32>(2, |_| panic!()).unwrap();
+    entry1.get_data::<_, i32, i32>(3, |_| panic!()).unwrap();
+    entry1.get_data::<_, i32, i32>(4, |_| panic!()).unwrap();
 
-    entry2.get_data(1, |_| new_module_cache_data()).unwrap();
-    entry1.get_data::<_, i32>(1, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32>(2, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32>(3, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32>(4, |_| panic!()).unwrap();
-    entry2.get_data::<_, i32>(1, |_| panic!()).unwrap();
-}
-
-fn new_module_cache_data() -> Result<ModuleCacheDataTupleType, ()> {
-    Ok((
-        Compilation::new(PrimaryMap::new()),
-        PrimaryMap::new(),
-        PrimaryMap::new(),
-        PrimaryMap::new(),
-        PrimaryMap::new(),
-        PrimaryMap::new(),
-        PrimaryMap::new(),
-    ))
+    entry2.get_data::<_, i32, i32>(1, |_| Ok(100)).unwrap();
+    entry1.get_data::<_, i32, i32>(1, |_| panic!()).unwrap();
+    entry1.get_data::<_, i32, i32>(2, |_| panic!()).unwrap();
+    entry1.get_data::<_, i32, i32>(3, |_| panic!()).unwrap();
+    entry1.get_data::<_, i32, i32>(4, |_| panic!()).unwrap();
+    entry2.get_data::<_, i32, i32>(1, |_| panic!()).unwrap();
 }

crates/environ/src/compilation.rs

@@ -1,7 +1,7 @@
 //! A `Compilation` contains the compiled function bodies for a WebAssembly
 //! module.
 
-use crate::cache::ModuleCacheDataTupleType;
+use crate::address_map::{ModuleAddressMap, ValueLabelsRanges};
 use crate::CacheConfig;
 use crate::ModuleTranslation;
 use cranelift_codegen::{binemit, ir, isa, isa::unwind::UnwindInfo};
@@ -184,6 +184,17 @@ pub enum CompileError {
     DebugInfoNotSupported,
 }
 
+/// A type alias for the result of `Compiler::compile_module`
+pub type CompileResult = (
+    Compilation,
+    Relocations,
+    ModuleAddressMap,
+    ValueLabelsRanges,
+    PrimaryMap<DefinedFuncIndex, ir::StackSlots>,
+    Traps,
+    StackMaps,
+);
+
 /// An implementation of a compiler from parsed WebAssembly module to native code.
 pub trait Compiler {
     /// Compile a parsed module with the given `TargetIsa`.
@@ -191,5 +202,5 @@ pub trait Compiler {
         translation: &ModuleTranslation,
         isa: &dyn isa::TargetIsa,
         cache_config: &CacheConfig,
-    ) -> Result<ModuleCacheDataTupleType, CompileError>;
+    ) -> Result<CompileResult, CompileError>;
 }

crates/environ/src/cranelift.rs

@@ -86,11 +86,12 @@
 // assume no valid stack pointer will ever be `usize::max_value() - 32k`.
 
 use crate::address_map::{FunctionAddressMap, InstructionAddressMap};
-use crate::cache::{ModuleCacheDataTupleType, ModuleCacheEntry};
+use crate::cache::ModuleCacheEntry;
 use crate::compilation::{
     Compilation, CompileError, CompiledFunction, Relocation, RelocationTarget, StackMapInformation,
     TrapInformation,
 };
+use crate::compilation::{CompileResult, Compiler};
 use crate::func_environ::{get_func_name, FuncEnvironment};
 use crate::{CacheConfig, FunctionBodyData, ModuleLocal, ModuleTranslation, Tunables};
 use cranelift_codegen::ir::{self, ExternalName};
@@ -283,17 +284,17 @@ fn get_function_address_map<'data>(
 /// optimizing it and then translating to assembly.
 pub struct Cranelift;
 
-impl crate::compilation::Compiler for Cranelift {
+impl Compiler for Cranelift {
     /// Compile the module using Cranelift, producing a compilation result with
     /// associated relocations.
     fn compile_module(
         translation: &ModuleTranslation,
         isa: &dyn isa::TargetIsa,
         cache_config: &CacheConfig,
-    ) -> Result<ModuleCacheDataTupleType, CompileError> {
+    ) -> Result<CompileResult, CompileError> {
         let cache_entry = ModuleCacheEntry::new("cranelift", cache_config);
 
-        let data = cache_entry.get_data(
+        let result = cache_entry.get_data(
             CompileEnv {
                 local: &translation.module.local,
                 module_translation: HashedModuleTranslationState(
@@ -305,11 +306,11 @@ impl crate::compilation::Compiler for Cranelift {
             },
             compile,
         )?;
-        Ok(data.into_tuple())
+        Ok(result)
     }
 }
 
-fn compile(env: CompileEnv<'_>) -> Result<ModuleCacheDataTupleType, CompileError> {
+fn compile(env: CompileEnv<'_>) -> Result<CompileResult, CompileError> {
     let Isa(isa) = env.isa;
     let mut functions = PrimaryMap::with_capacity(env.function_body_inputs.len());
     let mut relocations = PrimaryMap::with_capacity(env.function_body_inputs.len());

crates/environ/src/lightbeam.rs

@@ -1,7 +1,6 @@
 //! Support for compiling with Lightbeam.
 
-use crate::cache::ModuleCacheDataTupleType;
-use crate::compilation::{Compilation, CompileError};
+use crate::compilation::{Compilation, CompileError, CompileResult, Compiler};
 use crate::func_environ::FuncEnvironment;
 use crate::CacheConfig;
 use crate::ModuleTranslation;
@@ -15,14 +14,14 @@ use lightbeam::{CodeGenSession, NullOffsetSink, Sinks};
 /// A compiler that compiles a WebAssembly module with Lightbeam, directly translating the Wasm file.
 pub struct Lightbeam;
 
-impl crate::compilation::Compiler for Lightbeam {
+impl Compiler for Lightbeam {
     /// Compile the module using Lightbeam, producing a compilation result with
     /// associated relocations.
     fn compile_module(
         translation: &ModuleTranslation,
         isa: &dyn isa::TargetIsa,
         _cache_config: &CacheConfig,
-    ) -> Result<ModuleCacheDataTupleType, CompileError> {
+    ) -> Result<CompileResult, CompileError> {
         if translation.tunables.debug_info {
             return Err(CompileError::DebugInfoNotSupported);
         }