From 6f685937830a1fa0c29b5804b9b7b6ae98ff555c Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Wed, 27 Mar 2024 18:08:05 -0700
Subject: [PATCH] Fix compilation cache for components

This commit fixes a mistake in #8181 which meant that the caching for
components was no longer working. The mistake is fixed in this commit as
well as a new test being added too.
---
 crates/wasmtime/src/compile/code_builder.rs | 12 ++++----
 crates/wasmtime/src/engine/serialization.rs | 32 +++++++++++++++++++++
 2 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/crates/wasmtime/src/compile/code_builder.rs b/crates/wasmtime/src/compile/code_builder.rs
index fb026d2a1ac3..9393278edf1f 100644
--- a/crates/wasmtime/src/compile/code_builder.rs
+++ b/crates/wasmtime/src/compile/code_builder.rs
@@ -179,11 +179,13 @@ impl<'a> CodeBuilder<'a> {
                         // Implementation of how to serialize artifacts
                         |(_engine, _wasm, _), (code, _info_and_types)| Some(code.mmap().to_vec()),
                         // Cache hit, deserialize the provided artifacts
-                        |(engine, _wasm, _), serialized_bytes| {
-                            let code = engine
-                                .0
-                                .load_code_bytes(&serialized_bytes, ObjectKind::Module)
-                                .ok()?;
+                        |(engine, wasm, _), serialized_bytes| {
+                            let kind = if wasmparser::Parser::is_component(&wasm) {
+                                ObjectKind::Component
+                            } else {
+                                ObjectKind::Module
+                            };
+                            let code = engine.0.load_code_bytes(&serialized_bytes, kind).ok()?;
                             Some((code, None))
                         },
                     )?;
diff --git a/crates/wasmtime/src/engine/serialization.rs b/crates/wasmtime/src/engine/serialization.rs
index b2d5915b406e..3f1473fc1fa8 100644
--- a/crates/wasmtime/src/engine/serialization.rs
+++ b/crates/wasmtime/src/engine/serialization.rs
@@ -812,4 +812,36 @@ Caused by:
 
         Ok(())
     }
+
+    #[test]
+    #[cfg_attr(miri, ignore)]
+    #[cfg(feature = "component-model")]
+    fn components_are_cached() -> Result<()> {
+        use crate::component::Component;
+
+        let td = TempDir::new()?;
+        let config_path = td.path().join("config.toml");
+        std::fs::write(
+            &config_path,
+            &format!(
+                "
+                    [cache]
+                    enabled = true
+                    directory = '{}'
+                ",
+                td.path().join("cache").display()
+            ),
+        )?;
+        let mut cfg = Config::new();
+        cfg.cache_config_load(&config_path)?;
+        let engine = Engine::new(&cfg)?;
+        Component::new(&engine, "(component (core module (func)))")?;
+        assert_eq!(engine.config().cache_config.cache_hits(), 0);
+        assert_eq!(engine.config().cache_config.cache_misses(), 1);
+        Component::new(&engine, "(component (core module (func)))")?;
+        assert_eq!(engine.config().cache_config.cache_hits(), 1);
+        assert_eq!(engine.config().cache_config.cache_misses(), 1);
+
+        Ok(())
+    }
 }