Handle export modifier, hopefully

Signed-off-by: itowlson <ivan.towlson@fermyon.com>

Author: itowlson

Cargo.lock

@@ -14,3 +14,8 @@ tokio = { workspace = true, features = ["rt", "time"] }
 wasmparser = { workspace = true }
 wit-component = { workspace = true }
 wit-parser = { workspace = true }
+
+[dev-dependencies]
+tempfile = { workspace = true }
+wasm-pkg-common = { workspace = true }
+wit-component = { workspace = true, features = ["dummy-module"] }

crates/dependency-wit/Cargo.toml

@@ -35,16 +35,35 @@ 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() {
-        // TODO: map `export`
-        let (wasm_path, _export) = loader
+        let import_name = match dependency_name {
+            DependencyName::Plain(_) => None,
+            DependencyName::Package(dependency_package_name) => {
+                dependency_package_name.interface.as_ref()
+                // match dependency_package_name.interface.as_ref() {
+                //     Some(itf) => Some(itf),
+                //     None => None,
+                // }
+            }
+        };
+
+        let (wasm_path, export) = loader
             .load_component_dependency(dependency_name, dependency)
             .await?;
         let wasm_bytes = tokio::fs::read(&wasm_path).await?;
 
         let decoded = read_wasm(&wasm_bytes)?;
+        let decoded = match export {
+            None => decoded,
+            Some(export) => munge_aliased_export(decoded, &export, dependency_name)?,
+        };
         let impo_world = format!("impo-world{index}");
         let importised = importize(decoded, None, Some(&impo_world))?;
 
+        let imports = match import_name {
+            None => all_imports(&importised),
+            Some(itf) => one_import(&importised, itf.as_ref()),
+        };
+
         // Capture WITs for all packages used in the importised thing.
         // Things like WASI packages may be depended on my multiple packages
         // so we index on the package name to avoid emitting them twice.
@@ -68,16 +87,6 @@ pub async fn extract_wits(
 
         // Now add the imports to the aggregating component import world
 
-        let imports = match dependency_name {
-            DependencyName::Plain(_) => all_imports(&importised),
-            DependencyName::Package(dependency_package_name) => {
-                match dependency_package_name.interface.as_ref() {
-                    Some(itf) => one_import(&importised, itf),
-                    None => all_imports(&importised),
-                }
-            }
-        };
-
         let remap = aggregating_resolve.merge(importised.resolve().clone())?;
         for iid in imports {
             let mapped_iid = remap.map_interface(iid, None)?;
@@ -113,6 +122,122 @@ pub async fn extract_wits(
     Ok(buf)
 }
 
+fn munge_aliased_export(
+    decoded: DecodedWasm,
+    export: &str,
+    new_name: &DependencyName,
+) -> anyhow::Result<DecodedWasm> {
+    // TODO: I am not sure how `export` is meant to work if you are
+    // depping on a package rather than an itf
+
+    let export_qname = spin_serde::DependencyPackageName::try_from(export.to_string())?;
+    let Some(export_itf_name) = export_qname.interface.as_ref() else {
+        anyhow::bail!("the export name should be a qualified interface name - missing interface");
+    };
+    let export_pkg_name = wit_parser::PackageName {
+        namespace: export_qname.package.namespace().to_string(),
+        name: export_qname.package.name().to_string(),
+        version: export_qname.version,
+    };
+
+    let DependencyName::Package(new_name) = new_name else {
+        anyhow::bail!("the dependency name should be a qualified interface name - not qualified");
+    };
+    let Some(new_itf_name) = new_name.interface.as_ref() else {
+        anyhow::bail!(
+            "the dependency name should be a qualified interface name - missing interface"
+        );
+    };
+    let new_pkg_name = wit_parser::PackageName {
+        namespace: new_name.package.namespace().to_string(),
+        name: new_name.package.name().to_string(),
+        version: new_name.version.clone(),
+    };
+
+    let (mut resolve, decode_id) = match decoded {
+        DecodedWasm::WitPackage(resolve, id) => (resolve, WorldOrPackageId::Package(id)),
+        DecodedWasm::Component(resolve, id) => (resolve, WorldOrPackageId::World(id)),
+    };
+
+    // Two scenarios:
+    // 1. The new name is in a package that is already in the Resolve
+    //    1a. The package already contains an interface with the right name
+    //    1b. The package does not already contain an interface with the right name
+    // 2. The new name is in a package that is NOT already in the Resolve
+
+    let existing_pkg = resolve
+        .packages
+        .iter()
+        .find(|(_pkg_id, pkg)| pkg.name == new_pkg_name);
+
+    // We address the first level by creating the new-name package if it doesn't exist
+    let (inserting_into_pkg_id, inserting_into_pkg) = match existing_pkg {
+        Some(tup) => tup,
+        None => {
+            // insert the needed package
+            let package_wit = format!("package {new_pkg_name};");
+            let pkg_id = resolve
+                .push_str(std::env::current_dir().unwrap(), &package_wit)
+                .context("failed at setting up fake pkg")?;
+            let pkg = resolve.packages.get(pkg_id).unwrap();
+            (pkg_id, pkg)
+        }
+    };
+
+    // Second level asks if the package already contains the interface
+    let existing_itf = inserting_into_pkg.interfaces.get(new_itf_name.as_ref());
+    if existing_itf.is_some() {
+        // no rename is needed, but we might need to do some extra work to make sure
+        // that the export, rather than the import, gets included in the aggregated world
+        return Ok(decode_id.make_decoded_wasm(resolve));
+    }
+
+    // It does not: we need to slurp the EXPORTED itf into the `inserting_into`
+    // package under the NEW (importing) interface name
+    let Some(export_pkg_id) = resolve.package_names.get(&export_pkg_name) else {
+        anyhow::bail!("export is from a package that doesn't exist");
+    };
+    let Some(export_pkg) = resolve.packages.get(*export_pkg_id) else {
+        anyhow::bail!("export pkg id doesn't point to a package wtf");
+    };
+    let Some(export_itf_id) = export_pkg.interfaces.get(export_itf_name.as_ref()) else {
+        anyhow::bail!("export pkg doesn't contain export itf");
+    };
+    let Some(export_itf) = resolve.interfaces.get(*export_itf_id) else {
+        anyhow::bail!("export pkg doesn't contain export itf part 2");
+    };
+
+    let mut export_itf = export_itf.clone();
+    export_itf.package = Some(inserting_into_pkg_id);
+    export_itf.name = Some(new_itf_name.to_string());
+    let export_itf_id_2 = resolve.interfaces.alloc(export_itf);
+
+    // OKAY TIME TO ADD THIS UNDER THE WRONG NAME TO THE THINGY
+    // oh man there is some nonsense about worlds as well
+    let inserting_into_pkg_mut = resolve.packages.get_mut(inserting_into_pkg_id).unwrap(); // SHENANIGANS to get around a "mutable borrow at the same time as immutable borrow" woe
+    inserting_into_pkg_mut
+        .interfaces
+        .insert(new_itf_name.to_string(), export_itf_id_2);
+
+    let thingy = decode_id.make_decoded_wasm(resolve);
+
+    Ok(thingy)
+}
+
+enum WorldOrPackageId {
+    Package(wit_parser::PackageId),
+    World(wit_parser::WorldId),
+}
+
+impl WorldOrPackageId {
+    pub fn make_decoded_wasm(&self, resolve: wit_parser::Resolve) -> DecodedWasm {
+        match self {
+            Self::Package(id) => DecodedWasm::WitPackage(resolve, *id),
+            Self::World(id) => DecodedWasm::Component(resolve, *id),
+        }
+    }
+}
+
 fn all_imports(wasm: &DecodedWasm) -> Vec<wit_parser::InterfaceId> {
     wasm.resolve()
         .worlds
@@ -129,7 +254,7 @@ fn as_interface(wi: &wit_parser::WorldItem) -> Option<wit_parser::InterfaceId> {
     }
 }
 
-fn one_import(wasm: &DecodedWasm, name: &spin_serde::KebabId) -> Vec<wit_parser::InterfaceId> {
+fn one_import(wasm: &DecodedWasm, name: &str) -> Vec<wit_parser::InterfaceId> {
     let id = wasm
         .resolve()
         .interfaces
@@ -182,28 +307,131 @@ fn importize(
     Ok(DecodedWasm::Component(resolve, world_id))
 }
 
-// fn all_imports(wasm: &DecodedWasm) -> Vec<(wit_parser::PackageName, String)> {
-//     let mut itfs = vec![];
-
-//     for (_pid, pp) in &wasm.resolve().packages {
-//         for (_w, wid) in &pp.worlds {
-//             if let Some(world) = wasm.resolve().worlds.get(*wid) {
-//                 for (_wk, witem) in &world.imports {
-//                     if let wit_parser::WorldItem::Interface { id, .. } = witem {
-//                         if let Some(itf) = wasm.resolve().interfaces.get(*id) {
-//                             if let Some(itfp) = itf.package.as_ref() {
-//                                 if let Some(ppp) = wasm.resolve().packages.get(*itfp) {
-//                                     if let Some(itfname) = itf.name.as_ref() {
-//                                         itfs.push((ppp.name.clone(), itfname.clone()));
-//                                     }
-//                                 }
-//                             }
-//                         }
-//                     }
-//                 }
-//             }
-//         }
-//     }
-
-//     itfs
-// }
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    fn parse_wit(wit: &str) -> anyhow::Result<wit_parser::Resolve> {
+        let mut resolve = wit_parser::Resolve::new();
+        resolve.push_str("dummy.wit", wit)?;
+        Ok(resolve)
+    }
+
+    fn generate_dummy_component(wit: &str, world: &str) -> Vec<u8> {
+        let mut resolve = wit_parser::Resolve::default();
+        let package_id = resolve.push_str("test", wit).expect("should parse WIT");
+        let world_id = resolve
+            .select_world(&[package_id], Some(world))
+            .expect("should select world");
+
+        let mut wasm = wit_component::dummy_module(
+            &resolve,
+            world_id,
+            wit_parser::ManglingAndAbi::Legacy(wit_parser::LiftLowerAbi::Sync),
+        );
+        wit_component::embed_component_metadata(
+            &mut wasm,
+            &resolve,
+            world_id,
+            wit_component::StringEncoding::UTF8,
+        )
+        .expect("should embed component metadata");
+
+        let mut encoder = wit_component::ComponentEncoder::default()
+            .validate(true)
+            .module(&wasm)
+            .expect("should set module");
+        encoder.encode().expect("should encode component")
+    }
+
+    #[tokio::test]
+    async fn if_no_dependencies_then_empty_valid_wit() -> anyhow::Result<()> {
+        let wit = extract_wits(std::iter::empty(), ".").await?;
+
+        let resolve = parse_wit(&wit).expect("should have emitted valid WIT");
+
+        assert_eq!(1, resolve.packages.len());
+        assert_eq!(
+            "root:component",
+            resolve.packages.iter().next().unwrap().1.name.to_string()
+        );
+
+        assert_eq!(0, resolve.interfaces.len());
+
+        assert_eq!(1, resolve.worlds.len());
+
+        let world = resolve.worlds.iter().next().unwrap().1;
+        assert_eq!("root", world.name);
+        assert_eq!(0, world.imports.len());
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn single_dep_wit_extracted() -> anyhow::Result<()> {
+        let tempdir = tempfile::TempDir::new()?;
+        let dep_file = tempdir.path().join("regex.wasm");
+
+        let dep_wit = "package my:regex@1.0.0;\n\ninterface regex {\n  matches: func(s: string) -> bool;\n}\nworld matcher {\n  export regex;\n}";
+        let dep_wasm = generate_dummy_component(dep_wit, "matcher");
+        tokio::fs::write(&dep_file, &dep_wasm).await?;
+
+        let dep_name =
+            DependencyName::Package("my:regex/regex@1.0.0".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!(2, resolve.packages.len()); // root:component and my:regex
+        let (_rc_pkg_id, rc_pkg) = resolve
+            .packages
+            .iter()
+            .find(|(_, p)| p.name.to_string() == "root:component")
+            .expect("should have had `root:component`");
+        let (_mr_pkg_id, _mr_pkg) = resolve
+            .packages
+            .iter()
+            .find(|(_, p)| p.name.to_string() == "my:regex@1.0.0")
+            .expect("should have had `my:regex`");
+
+        assert_eq!(1, resolve.interfaces.len());
+        assert_eq!(
+            "regex",
+            resolve
+                .interfaces
+                .iter()
+                .next()
+                .unwrap()
+                .1
+                .name
+                .as_ref()
+                .unwrap()
+        );
+        let regex_itf_id = resolve.interfaces.iter().next().unwrap().0;
+
+        assert_eq!(2, rc_pkg.worlds.len()); // root and synthetic "impo*" wart
+        let root_world_id = rc_pkg
+            .worlds
+            .iter()
+            .find(|w| w.0 == "root")
+            .expect("should have had `root` world")
+            .1;
+
+        let world = resolve.worlds.get(*root_world_id).unwrap();
+        assert_eq!(1, world.imports.len());
+        let expected_import = wit_parser::WorldItem::Interface {
+            id: regex_itf_id,
+            stability: wit_parser::Stability::Unknown,
+        };
+        let import = world.imports.values().next().unwrap();
+        assert_eq!(&expected_import, import);
+
+        Ok(())
+    }
+}