Merge pull request #3446 from itowlson/dep-wits-functions-in-worlds

Fix WIT extractor ignoring world-level functions

Author: itowlson

crates/dependency-wit/src/lib.rs

@@ -33,6 +33,12 @@ pub async fn extract_wits_into(
     Ok(())
 }
 
+enum ImportKind {
+    Function(spin_serde::KebabId),
+    WholePackage,
+    Interface(spin_serde::KebabId),
+}
+
 pub async fn extract_wits(
     source: impl Iterator<Item = (&DependencyName, &ComponentDependency)>,
     app_root: impl AsRef<Path>,
@@ -49,10 +55,13 @@ pub async fn extract_wits(
 
     // TODO: figure out what to do if we import two itfs from same dep
     for (index, (dependency_name, dependency)) in source.enumerate() {
-        let import_name = match dependency_name {
-            DependencyName::Plain(_) => None,
+        let import_kind = match dependency_name {
+            DependencyName::Plain(name) => ImportKind::Function(name.clone()),
             DependencyName::Package(dependency_package_name) => {
-                dependency_package_name.interface.as_ref()
+                match &dependency_package_name.interface {
+                    Some(itf) => ImportKind::Interface(itf.clone()),
+                    None => ImportKind::WholePackage,
+                }
             }
         };
 
@@ -75,9 +84,18 @@ pub async fn extract_wits(
         let importised = importize(decoded, Some(&impo_world))
             .with_context(|| format!("failed to map importize dependency {dependency_name}"))?;
 
-        let imports = match import_name {
-            None => all_imports(&importised),
-            Some(itf) => one_import(&importised, itf.as_ref())?,
+        let imports = match &import_kind {
+            ImportKind::WholePackage => all_imports(&importised),
+            ImportKind::Interface(itf) => one_import(&importised, itf.as_ref())?,
+            ImportKind::Function(_) => Default::default(),
+        };
+        let func_import = match &import_kind {
+            ImportKind::Function(f) => one_func_import(&importised, f.as_ref())?,
+            _ => Default::default(),
+        };
+        let type_imports = match &import_kind {
+            ImportKind::Function(_) => world_type_imports(&importised),
+            _ => Default::default(),
         };
 
         // Capture WITs for all packages used in the importised thing.
@@ -142,6 +160,36 @@ pub async fn extract_wits(
                 });
             }
         }
+        if let Some(mut func) = func_import {
+            // Remap type IDs in the function to reference the aggregating resolve
+            for param in &mut func.params {
+                if let wit_parser::Type::Id(id) = &mut param.ty {
+                    *id = remap.map_type(*id, Span::default())?;
+                }
+            }
+            if let Some(wit_parser::Type::Id(id)) = &mut func.result {
+                *id = remap.map_type(*id, Span::default())?;
+            }
+
+            // Add world-level type definitions that the function depends on
+            let aggregating_world = aggregating_resolve
+                .worlds
+                .get_mut(aggregating_world_id)
+                .context("aggregated dependency world doesn't exist")?;
+            for (name, type_id) in &type_imports {
+                let mapped_id = remap.map_type(*type_id, Span::default())?;
+                let wk = wit_parser::WorldKey::Name(name.clone());
+                let world_item = wit_parser::WorldItem::Type {
+                    id: mapped_id,
+                    span: Span::default(),
+                };
+                aggregating_world.imports.insert(wk, world_item);
+            }
+
+            let wk = wit_parser::WorldKey::Name(func.name.clone());
+            let world_item = wit_parser::WorldItem::Function(func);
+            aggregating_world.imports.insert(wk, world_item);
+        }
     }
 
     // Text for the root package and world(s)
@@ -321,6 +369,13 @@ fn as_interface(wi: &wit_parser::WorldItem) -> Option<wit_parser::InterfaceId> {
     }
 }
 
+fn as_func(wi: &wit_parser::WorldItem) -> Option<&wit_parser::Function> {
+    match wi {
+        wit_parser::WorldItem::Function(func) => Some(func),
+        _ => None,
+    }
+}
+
 fn one_import(wasm: &DecodedWasm, name: &str) -> anyhow::Result<Vec<wit_parser::InterfaceId>> {
     let id = wasm
         .resolve()
@@ -332,6 +387,47 @@ fn one_import(wasm: &DecodedWasm, name: &str) -> anyhow::Result<Vec<wit_parser::
     Ok(vec![id])
 }
 
+fn world_type_imports(wasm: &DecodedWasm) -> Vec<(String, wit_parser::TypeId)> {
+    wasm.resolve()
+        .worlds
+        .iter()
+        .flat_map(|(_wid, w)| {
+            w.imports.iter().filter_map(|(wk, wi)| {
+                if let wit_parser::WorldItem::Type { id, .. } = wi {
+                    let name = match wk {
+                        wit_parser::WorldKey::Name(n) => n.clone(),
+                        wit_parser::WorldKey::Interface(_) => return None,
+                    };
+                    Some((name, *id))
+                } else {
+                    None
+                }
+            })
+        })
+        .collect()
+}
+
+fn one_func_import(wasm: &DecodedWasm, name: &str) -> anyhow::Result<Option<wit_parser::Function>> {
+    let funcs = wasm
+        .resolve()
+        .worlds
+        .iter()
+        .flat_map(|w| {
+            w.1.imports
+                .iter()
+                .flat_map(|(_wk, wi)| as_func(wi))
+                .filter(|f| f.name == name)
+        })
+        .collect::<Vec<_>>();
+
+    // This shouldn't happen because we are using the compiled Wasm so there should
+    // be only one world in play.  But belt and braces.
+    if funcs.len() > 1 {
+        anyhow::bail!("Dependency exports more than one function named {name}");
+    }
+    Ok(funcs.first().cloned().cloned())
+}
+
 fn read_wasm(wasm_bytes: &[u8]) -> anyhow::Result<DecodedWasm> {
     if wasmparser::Parser::is_component(wasm_bytes) {
         wit_component::decode(wasm_bytes)
@@ -493,4 +589,237 @@ mod test {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn world_level_func_extracted() -> anyhow::Result<()> {
+        let tempdir = tempfile::TempDir::new()?;
+        let dep_file = tempdir.path().join("crimes.wasm");
+
+        let dep_wit = "package my:crimes@1.0.0;\n\nworld crimes {\n  export is-curse: func(s: string) -> bool;\n}";
+        let dep_wasm = generate_dummy_component(dep_wit, "crimes");
+        tokio::fs::write(&dep_file, &dep_wasm).await?;
+
+        let dep_name = DependencyName::Plain("is-curse".to_string().try_into().unwrap());
+        let dep_src = ComponentDependency::Local {
+            path: dep_file,
+            export: None,
+        };
+        let deps = std::iter::once((&dep_name, &dep_src));
+
+        let wit = extract_wits(deps, ".").await?;
+
+        let resolve = parse_wit(&wit).expect("should have emitted valid WIT");
+
+        assert_eq!(1, resolve.packages.len()); // root:component - world-level funcs don't retain their package when importised
+        let (_rc_pkg_id, rc_pkg) = resolve
+            .packages
+            .iter()
+            .find(|(_, p)| p.name.to_string() == "root:component")
+            .expect("should have had `root:component`");
+
+        let root_world_id = rc_pkg
+            .worlds
+            .get("root")
+            .expect("should have had root world");
+        let root_world = resolve
+            .worlds
+            .get(*root_world_id)
+            .expect("should have had root world at that id");
+
+        let func = root_world
+            .imports
+            .iter()
+            .filter_map(|(_, wi)| as_func(wi))
+            .find(|f| f.name == "is-curse")
+            .expect("is-curse function does not appear in root imports");
+
+        assert_eq!(1, func.params.len());
+        assert_eq!(wit_parser::Type::String, func.params.first().unwrap().ty);
+        assert_eq!(wit_parser::Type::Bool, func.result.unwrap());
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn world_level_func_with_record_param() -> anyhow::Result<()> {
+        let tempdir = tempfile::TempDir::new()?;
+        let dep_file = tempdir.path().join("greeter.wasm");
+
+        let dep_wit = r#"package my:greeter@1.0.0;
+
+world greeter {
+  record person {
+    name: string,
+    age: u32,
+  }
+  export greet: func(who: person) -> string;
+}"#;
+        let dep_wasm = generate_dummy_component(dep_wit, "greeter");
+        tokio::fs::write(&dep_file, &dep_wasm).await?;
+
+        let dep_name = DependencyName::Plain("greet".to_string().try_into().unwrap());
+        let dep_src = ComponentDependency::Local {
+            path: dep_file,
+            export: None,
+        };
+        let deps = std::iter::once((&dep_name, &dep_src));
+
+        let wit = extract_wits(deps, ".").await?;
+
+        let resolve = parse_wit(&wit).expect("should have emitted valid WIT");
+
+        let (_rc_pkg_id, rc_pkg) = resolve
+            .packages
+            .iter()
+            .find(|(_, p)| p.name.to_string() == "root:component")
+            .expect("should have had `root:component`");
+
+        let root_world_id = rc_pkg
+            .worlds
+            .get("root")
+            .expect("should have had root world");
+        let root_world = resolve
+            .worlds
+            .get(*root_world_id)
+            .expect("should have had root world at that id");
+
+        let func = root_world
+            .imports
+            .iter()
+            .filter_map(|(_, wi)| as_func(wi))
+            .find(|f| f.name == "greet")
+            .expect("greet function does not appear in root imports");
+
+        assert_eq!(1, func.params.len());
+        // The param should be a user-defined type (record)
+        assert!(
+            matches!(func.params.first().unwrap().ty, wit_parser::Type::Id(_)),
+            "expected record param to be Type::Id"
+        );
+        assert_eq!(wit_parser::Type::String, func.result.unwrap());
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn world_level_func_with_record_result() -> anyhow::Result<()> {
+        let tempdir = tempfile::TempDir::new()?;
+        let dep_file = tempdir.path().join("lookup.wasm");
+
+        let dep_wit = r#"package my:lookup@1.0.0;
+
+world lookup {
+  record info {
+    value: string,
+    found: bool,
+  }
+  export lookup: func(key: string) -> info;
+}"#;
+        let dep_wasm = generate_dummy_component(dep_wit, "lookup");
+        tokio::fs::write(&dep_file, &dep_wasm).await?;
+
+        let dep_name = DependencyName::Plain("lookup".to_string().try_into().unwrap());
+        let dep_src = ComponentDependency::Local {
+            path: dep_file,
+            export: None,
+        };
+        let deps = std::iter::once((&dep_name, &dep_src));
+
+        let wit = extract_wits(deps, ".").await?;
+
+        let resolve = parse_wit(&wit).expect("should have emitted valid WIT");
+
+        let (_rc_pkg_id, rc_pkg) = resolve
+            .packages
+            .iter()
+            .find(|(_, p)| p.name.to_string() == "root:component")
+            .expect("should have had `root:component`");
+
+        let root_world_id = rc_pkg
+            .worlds
+            .get("root")
+            .expect("should have had root world");
+        let root_world = resolve
+            .worlds
+            .get(*root_world_id)
+            .expect("should have had root world at that id");
+
+        let func = root_world
+            .imports
+            .iter()
+            .filter_map(|(_, wi)| as_func(wi))
+            .find(|f| f.name == "lookup")
+            .expect("lookup function does not appear in root imports");
+
+        assert_eq!(1, func.params.len());
+        assert_eq!(wit_parser::Type::String, func.params.first().unwrap().ty);
+        // The result should be a user-defined type (record)
+        assert!(
+            matches!(func.result, Some(wit_parser::Type::Id(_))),
+            "expected record result to be Type::Id"
+        );
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn world_level_func_with_enum_param() -> anyhow::Result<()> {
+        let tempdir = tempfile::TempDir::new()?;
+        let dep_file = tempdir.path().join("color.wasm");
+
+        let dep_wit = r#"package my:colors@1.0.0;
+
+world colors {
+  enum color {
+    red,
+    green,
+    blue,
+  }
+  export color-name: func(c: color) -> string;
+}"#;
+        let dep_wasm = generate_dummy_component(dep_wit, "colors");
+        tokio::fs::write(&dep_file, &dep_wasm).await?;
+
+        let dep_name = DependencyName::Plain("color-name".to_string().try_into().unwrap());
+        let dep_src = ComponentDependency::Local {
+            path: dep_file,
+            export: None,
+        };
+        let deps = std::iter::once((&dep_name, &dep_src));
+
+        let wit = extract_wits(deps, ".").await?;
+
+        let resolve = parse_wit(&wit).expect("should have emitted valid WIT");
+
+        let (_rc_pkg_id, rc_pkg) = resolve
+            .packages
+            .iter()
+            .find(|(_, p)| p.name.to_string() == "root:component")
+            .expect("should have had `root:component`");
+
+        let root_world_id = rc_pkg
+            .worlds
+            .get("root")
+            .expect("should have had root world");
+        let root_world = resolve
+            .worlds
+            .get(*root_world_id)
+            .expect("should have had root world at that id");
+
+        let func = root_world
+            .imports
+            .iter()
+            .filter_map(|(_, wi)| as_func(wi))
+            .find(|f| f.name == "color-name")
+            .expect("color-name function does not appear in root imports");
+
+        assert_eq!(1, func.params.len());
+        assert!(
+            matches!(func.params.first().unwrap().ty, wit_parser::Type::Id(_)),
+            "expected enum param to be Type::Id"
+        );
+        assert_eq!(wit_parser::Type::String, func.result.unwrap());
+
+        Ok(())
+    }
 }