InMemory/MultiSource Improvements

lib/wasix/src/runners/wasi.rs

@@ -482,7 +482,7 @@ mod tests {
 
         let fs = &init.state.fs.root_fs;
 
-        fs.read_dir(&std::path::Path::new("/host")).unwrap();
+        fs.read_dir(std::path::Path::new("/host")).unwrap();
     }
 
     #[cfg(all(feature = "host-fs", feature = "sys"))]
@@ -535,7 +535,7 @@ mod tests {
 
         let fs = &init.state.fs.root_fs;
 
-        fs.read_dir(&std::path::Path::new("/host")).unwrap();
-        fs.read_dir(&std::path::Path::new("/settings")).unwrap();
+        fs.read_dir(std::path::Path::new("/host")).unwrap();
+        fs.read_dir(std::path::Path::new("/settings")).unwrap();
     }
 }

lib/wasix/src/runtime/package_loader/builtin_loader.rs

@@ -704,7 +704,7 @@ mod test {
         let path = dir.path();
 
         let contents = "fail";
-        let correct_hash = WebcHash::sha256(&contents);
+        let correct_hash = WebcHash::sha256(contents);
         let used_hash =
             WebcHash::parse_hex("0000a28ea38a000f3a3328cb7fabe330638d3258affe1a869e3f92986222d997")
                 .unwrap();

lib/wasix/src/runtime/resolver/in_memory_source.rs

@@ -68,9 +68,20 @@ impl InMemorySource {
     }
 
     /// Add a new [`PackageSummary`] to the [`InMemorySource`].
+    ///
+    /// Named packages are also made accessible by their hash.
     pub fn add(&mut self, summary: PackageSummary) {
         match summary.pkg.id.clone() {
             PackageId::Named(ident) => {
+                // Also add the package as a hashed package.
+                let pkg_hash = PackageHash::Sha256(wasmer_config::hash::Sha256Hash(
+                    summary.dist.webc_sha256.as_bytes(),
+                ));
+                self.hash_packages
+                    .entry(pkg_hash)
+                    .or_insert_with(|| summary.clone());
+
+                // Add the named package.
                 let summaries = self
                     .named_packages
                     .entry(ident.full_name.clone())
@@ -107,6 +118,17 @@ impl InMemorySource {
             PackageId::Hash(hash) => self.hash_packages.get(hash),
         }
     }
+
+    pub fn is_empty(&self) -> bool {
+        self.named_packages.is_empty() && self.hash_packages.is_empty()
+    }
+
+    /// Returns the number of packages in the source.
+    pub fn len(&self) -> usize {
+        // Only need to count the hash packages,
+        // as the named packages are also always added as hashed.
+        self.hash_packages.len()
+    }
 }
 
 #[async_trait::async_trait]

lib/wasix/src/runtime/resolver/multi_source.rs

@@ -24,8 +24,14 @@ pub struct MultiSource {
     strategy: MultiSourceStrategy,
 }
 
+impl Default for MultiSource {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 impl MultiSource {
-    pub const fn new() -> Self {
+    pub fn new() -> Self {
         MultiSource {
             sources: Vec::new(),
             strategy: MultiSourceStrategy::default(),
@@ -51,19 +57,47 @@ impl MultiSource {
 impl Source for MultiSource {
     #[tracing::instrument(level = "debug", skip_all, fields(%package))]
     async fn query(&self, package: &PackageSource) -> Result<Vec<PackageSummary>, QueryError> {
+        let mut output = Vec::<PackageSummary>::new();
+
         for source in &self.sources {
             match source.query(package).await {
-                Ok(summaries) => return Ok(summaries),
-                Err(QueryError::Unsupported) if self.strategy.continue_if_unsupported => continue,
-                Err(QueryError::NotFound) if self.strategy.continue_if_not_found => continue,
-                Err(QueryError::NoMatches { .. }) if self.strategy.continue_if_no_matches => {
+                Ok(mut summaries) => {
+                    if self.strategy.merge_results {
+                        // Extend matches, but skip already found versions.
+                        summaries.retain(|new| {
+                            !output.iter().any(|existing| new.pkg.id == existing.pkg.id)
+                        });
+                        output.extend(summaries);
+                    } else {
+                        return Ok(summaries);
+                    }
+                }
+                Err(QueryError::Unsupported)
+                    if self.strategy.continue_if_unsupported || self.strategy.merge_results =>
+                {
+                    continue
+                }
+                Err(QueryError::NotFound)
+                    if self.strategy.continue_if_not_found || self.strategy.merge_results =>
+                {
+                    continue
+                }
+                Err(QueryError::NoMatches { .. })
+                    if self.strategy.continue_if_no_matches || self.strategy.merge_results =>
+                {
                     continue
                 }
                 Err(e) => return Err(e),
             }
         }
 
-        Err(QueryError::NotFound)
+        if !output.is_empty() {
+            output.sort_by(|a, b| a.pkg.id.cmp(&b.pkg.id));
+
+            Ok(output)
+        } else {
+            Err(QueryError::NotFound)
+        }
     }
 }
 
@@ -86,20 +120,66 @@ pub struct MultiSourceStrategy {
     ///
     /// This flag is **disabled** by default.
     pub continue_if_no_matches: bool,
+
+    /// Merge results from all sources into a single result.
+    ///
+    /// True by default.
+    pub merge_results: bool,
 }
 
-impl MultiSourceStrategy {
-    pub const fn default() -> Self {
+impl Default for MultiSourceStrategy {
+    fn default() -> Self {
         MultiSourceStrategy {
             continue_if_unsupported: true,
             continue_if_not_found: true,
             continue_if_no_matches: true,
+            merge_results: true,
         }
     }
 }
 
-impl Default for MultiSourceStrategy {
-    fn default() -> Self {
-        MultiSourceStrategy::default()
+#[cfg(test)]
+mod tests {
+    use wasmer_config::package::PackageId;
+
+    use super::super::{DistributionInfo, InMemorySource, PackageInfo, WebcHash};
+    use super::*;
+
+    /// Test that the `MultiSource` can merge results from multiple sources.
+    #[tokio::test]
+    async fn test_multi_source_merge() {
+        let id1 = PackageId::new_named("ns/pkg", "0.0.1".parse().unwrap());
+        let pkg1 = PackageSummary {
+            pkg: PackageInfo {
+                id: id1.clone(),
+                commands: Vec::new(),
+                entrypoint: None,
+                dependencies: Vec::new(),
+                filesystem: Vec::new(),
+            },
+            dist: DistributionInfo {
+                webc: "https://example.com/ns/pkg/0.0.1".parse().unwrap(),
+                webc_sha256: WebcHash([0u8; 32]),
+            },
+        };
+
+        let id2 = PackageId::new_named("ns/pkg", "0.0.2".parse().unwrap());
+        let mut pkg2 = pkg1.clone();
+        pkg2.pkg.id = id2.clone();
+
+        let mut mem1 = InMemorySource::new();
+        mem1.add(pkg1);
+
+        let mut mem2 = InMemorySource::new();
+        mem2.add(pkg2);
+
+        let mut multi = MultiSource::new();
+        multi.add_source(mem1);
+        multi.add_source(mem2);
+
+        let summaries = multi.query(&"ns/pkg".parse().unwrap()).await.unwrap();
+        assert_eq!(summaries.len(), 2);
+        assert_eq!(summaries[0].pkg.id, id1);
+        assert_eq!(summaries[1].pkg.id, id2);
     }
 }

lib/wasix/tests/runners.rs

@@ -236,9 +236,8 @@ fn client() -> Client {
         }
         builder
     };
-    let client = builder.build().unwrap();
 
-    client
+    builder.build().unwrap()
 }
 
 #[cfg(not(target_os = "windows"))]