Add importize method to Resolve (#1784)

* Add importize method to Resolve

A new method `importize` is added to the Resolve. This is to allow to
mutate the Resolve to state where it would resemble what a consuming
component would expect to see during composition.

Signed-off-by: karthik2804 <[email protected]>

* Updates to importize

* Update the CLI to have `--importize` and `--importize-world`
* Rewrite the test to use these flags and have multiple tests in one
  file, each with a smaller world.
* Update the implementation to preserve allow-listed imports instead of
  removing all imports and recreating what needs to be preserved.

* Enable fuzz-testing of `importize`

---------

Signed-off-by: karthik2804 <[email protected]>
Co-authored-by: Alex Crichton <[email protected]>

crates/wit-parser/src/resolve.rs

@@ -931,6 +931,101 @@ package {name} is defined in two different locations:\n\
         Some(self.id_of_name(interface.package.unwrap(), interface.name.as_ref()?))
     }
 
+    /// Convert a world to an "importized" version where the world is updated
+    /// in-place to reflect what it would look like to be imported.
+    ///
+    /// This is a transformation which is used as part of the process of
+    /// importing a component today. For example when a component depends on
+    /// another component this is useful for generating WIT which can be use to
+    /// represent the component being imported. The general idea is that this
+    /// function will update the `world_id` specified such it imports the
+    /// functionality that it previously exported. The world will be left with
+    /// no exports.
+    ///
+    /// This world is then suitable for merging into other worlds or generating
+    /// bindings in a context that is importing the original world. This
+    /// is intended to be used as part of language tooling when depending on
+    /// other components.
+    pub fn importize(&mut self, world_id: WorldId) -> Result<()> {
+        // Collect the set of interfaces which are depended on by exports. Also
+        // all imported types are assumed to stay so collect any interfaces
+        // they depend on.
+        let mut live_through_exports = IndexSet::default();
+        for (_, export) in self.worlds[world_id].exports.iter() {
+            if let WorldItem::Interface { id, .. } = export {
+                self.collect_interface_deps(*id, &mut live_through_exports);
+            }
+        }
+        for (_, import) in self.worlds[world_id].imports.iter() {
+            if let WorldItem::Type(ty) = import {
+                if let Some(dep) = self.type_interface_dep(*ty) {
+                    self.collect_interface_deps(dep, &mut live_through_exports);
+                }
+            }
+        }
+
+        // Rename the world to avoid having it get confused with the original
+        // name of the world. Add `-importized` to it for now. Precisely how
+        // this new world is created may want to be updated over time if this
+        // becomes problematic.
+        let world = &mut self.worlds[world_id];
+        let pkg = &mut self.packages[world.package.unwrap()];
+        pkg.worlds.shift_remove(&world.name);
+        world.name.push_str("-importized");
+        pkg.worlds.insert(world.name.clone(), world_id);
+
+        // Trim all unnecessary imports first.
+        world.imports.retain(|name, item| match (name, item) {
+            // Remove imports which can't be used by import such as:
+            //
+            // * `import foo: interface { .. }`
+            // * `import foo: func();`
+            (WorldKey::Name(_), WorldItem::Interface { .. } | WorldItem::Function(_)) => false,
+
+            // Coarsely say that all top-level types are required to avoid
+            // calculating precise liveness of them right now.
+            (WorldKey::Name(_), WorldItem::Type(_)) => true,
+
+            // Only retain interfaces if they're needed somehow transitively
+            // for the exports.
+            (WorldKey::Interface(id), _) => live_through_exports.contains(id),
+        });
+
+        // After all unnecessary imports are gone remove all exports and move
+        // them all to imports, failing if there's an overlap.
+        for (name, export) in mem::take(&mut world.exports) {
+            match (name.clone(), world.imports.insert(name, export)) {
+                // no previous item? this insertion was ok
+                (_, None) => {}
+
+                // cannot overwrite an import with an export
+                (WorldKey::Name(name), _) => {
+                    bail!("world export `{name}` conflicts with import of same name");
+                }
+
+                // interface overlap is ok and is always allowed.
+                (WorldKey::Interface(id), Some(WorldItem::Interface { id: other, .. })) => {
+                    assert_eq!(id, other);
+                }
+
+                (WorldKey::Interface(_), _) => unreachable!(),
+            }
+        }
+
+        #[cfg(debug_assertions)]
+        self.assert_valid();
+        Ok(())
+    }
+
+    fn collect_interface_deps(&self, interface: InterfaceId, deps: &mut IndexSet<InterfaceId>) {
+        if !deps.insert(interface) {
+            return;
+        }
+        for dep in self.interface_direct_deps(interface) {
+            self.collect_interface_deps(dep, deps);
+        }
+    }
+
     /// Returns the ID of the specified `name` within the `pkg`.
     pub fn id_of_name(&self, pkg: PackageId, name: &str) -> String {
         let package = &self.packages[pkg];

fuzz/src/roundtrip_wit.rs

@@ -34,15 +34,18 @@ pub fn run(u: &mut Unstructured<'_>) -> Result<()> {
     // to avoid timing out this fuzzer with asan enabled.
     let mut decoded_worlds = Vec::new();
     for (id, world) in resolve.worlds.iter().take(20) {
-        log::debug!("testing world {}", world.name);
+        log::debug!("embedding world {} as in a dummy module", world.name);
         let mut dummy = wit_component::dummy_module(&resolve, id);
         wit_component::embed_component_metadata(&mut dummy, &resolve, id, StringEncoding::UTF8)
             .unwrap();
         write_file("dummy.wasm", &dummy);
 
+        // Decode what was just created and record it later for testing merging
+        // worlds together.
         let (_, decoded) = wit_component::metadata::decode(&dummy).unwrap();
         decoded_worlds.push(decoded.resolve);
 
+        log::debug!("... componentizing the world into a binary component");
         let wasm = wit_component::ComponentEncoder::default()
             .module(&dummy)
             .unwrap()
@@ -55,7 +58,15 @@ pub fn run(u: &mut Unstructured<'_>) -> Result<()> {
         .validate_all(&wasm)
         .unwrap();
 
+        log::debug!("... decoding the component itself");
         wit_component::decode(&wasm).unwrap();
+
+        // Test out importizing the world and then assert the world is still
+        // valid.
+        log::debug!("... importizing this world");
+        let mut resolve2 = resolve.clone();
+        let _ = resolve2.importize(id);
+        resolve.assert_valid();
     }
 
     if decoded_worlds.len() < 2 {

src/bin/wasm-tools/component.rs

@@ -4,6 +4,7 @@ use anyhow::{bail, Context, Result};
 use clap::Parser;
 use std::collections::HashMap;
 use std::io::Read;
+use std::mem;
 use std::path::{Path, PathBuf};
 use wasm_encoder::reencode::{Error, Reencode, ReencodeComponent, RoundtripReencoder};
 use wasm_encoder::ModuleType;
@@ -499,6 +500,30 @@ pub struct WitOpts {
     )]
     json: bool,
 
+    /// Generates WIT to import the component specified to this command.
+    ///
+    /// This flags requires that the input is a binary component, not a
+    /// wasm-encoded WIT package. This will then generate a WIT world and output
+    /// that. The returned world will have imports corresponding to the exports
+    /// of the component which is input.
+    ///
+    /// This is similar to `--importize-world`, but is used with components.
+    #[clap(long, conflicts_with = "importize_world")]
+    importize: bool,
+
+    /// Generates a WIT world to import a component which corresponds to the
+    /// selected world.
+    ///
+    /// This flag is used to indicate that the input is a WIT package and the
+    /// world passed here is the name of a WIT `world` within the package. The
+    /// output of the command will be the same WIT world but one that's
+    /// importing the selected world. This effectively moves the world's exports
+    /// to imports.
+    ///
+    /// This is similar to `--importize`, but is used with WIT packages.
+    #[clap(long, conflicts_with = "importize", value_name = "WORLD")]
+    importize_world: Option<String>,
+
     /// Features to enable when parsing the `wit` option.
     ///
     /// This flag enables the `@unstable` feature in WIT documents where the
@@ -521,7 +546,13 @@ impl WitOpts {
 
     /// Executes the application.
     fn run(self) -> Result<()> {
-        let decoded = self.decode_input()?;
+        let mut decoded = self.decode_input()?;
+
+        if self.importize {
+            self.importize(&mut decoded, None)?;
+        } else if self.importize_world.is_some() {
+            self.importize(&mut decoded, self.importize_world.as_deref())?;
+        }
 
         // Now that the WIT document has been decoded, it's time to emit it.
         // This interprets all of the output options and performs such a task.
@@ -605,6 +636,30 @@ impl WitOpts {
         }
     }
 
+    fn importize(&self, decoded: &mut DecodedWasm, world: Option<&str>) -> Result<()> {
+        let (resolve, world_id) = match (&mut *decoded, world) {
+            (DecodedWasm::Component(resolve, world), None) => (resolve, *world),
+            (DecodedWasm::Component(..), Some(_)) => {
+                bail!(
+                    "the `--importize-world` flag is not compatible with a \
+                     component input, use `--importize` instead"
+                );
+            }
+            (DecodedWasm::WitPackage(resolve, id), world) => {
+                let world = resolve.select_world(*id, world)?;
+                (resolve, world)
+            }
+        };
+        // let pkg = decoded.package();
+        // let world_id = decoded.resolve().select_world(main, None)?;
+        resolve
+            .importize(world_id)
+            .context("failed to move world exports to imports")?;
+        let resolve = mem::take(resolve);
+        *decoded = DecodedWasm::Component(resolve, world_id);
+        Ok(())
+    }
+
     fn emit_wasm(&self, decoded: &DecodedWasm) -> Result<()> {
         assert!(self.wasm || self.wat);
         assert!(self.out_dir.is_none());

tests/cli.rs

@@ -149,27 +149,32 @@ fn execute(cmd: &mut Command, stdin: Option<&[u8]>, should_fail: bool) -> Result
 
     let output = p.wait_with_output()?;
 
-    if !output.status.success() {
-        if !should_fail {
-            bail!(
-                "{cmd:?} failed:
-                status: {}
-                stdout: {}
-                stderr: {}",
-                output.status,
-                String::from_utf8_lossy(&output.stdout),
-                String::from_utf8_lossy(&output.stderr)
-            );
+    let mut failure = None;
+    match output.status.code() {
+        Some(0) => {
+            if should_fail {
+                failure = Some("succeeded instead of failed");
+            }
         }
-    } else if should_fail {
+        Some(1) | Some(2) => {
+            if !should_fail {
+                failure = Some("failed");
+            }
+        }
+        _ => failure = Some("unknown exit code"),
+    }
+    if let Some(msg) = failure {
         bail!(
-            "{cmd:?} succeeded instead of failed
-                stdout: {}
-                stderr: {}",
+            "{cmd:?} {msg}:
+             status: {}
+             stdout: {}
+             stderr: {}",
+            output.status,
             String::from_utf8_lossy(&output.stdout),
             String::from_utf8_lossy(&output.stderr)
         );
     }
+
     Ok(output)
 }
 

tests/cli/importize.wit

@@ -0,0 +1,80 @@
+// RUN[simple]: component wit --importize-world simple %
+// RUN[simple-component]: component embed --dummy --world simple % | \
+//                        component wit --importize
+// RUN[with-deps]: component wit --importize-world with-deps %
+// RUN[simple-toplevel]: component wit --importize-world simple-toplevel %
+// RUN[toplevel-deps]: component wit --importize-world toplevel-deps %
+// FAIL[fail1]: component wit --importize-world fail1 %
+// RUN[trim-imports]: component wit --importize-world trim-imports %
+// RUN[tricky-import]: component wit --importize-world tricky-import %
+
+package importize:importize;
+
+interface t {
+    resource r;
+}
+interface bar {
+    use t.{r};
+    record foo {
+      x: string
+    }
+    importize: func(name: r);
+}
+
+interface qux {
+   use bar.{foo};
+   blah: func(boo: foo);
+}
+
+interface something-else-dep {
+    type t = u32;
+}
+
+world simple {
+    export t;
+}
+
+world with-deps {
+    export qux;
+}
+
+world simple-toplevel {
+    export foo: func();
+    export something: interface {
+        foo: func();
+    }
+}
+
+world toplevel-deps {
+    type s = u32;
+    export bar: func() -> s;
+    export something-else: interface {
+        use something-else-dep.{t};
+        bar: func() -> t;
+    }
+}
+
+world fail1 {
+  type foo = u32;
+  export foo: func() -> foo;
+}
+
+interface a {}
+interface b {}
+
+world trim-imports {
+  import a;
+  import foo: func();
+  import bar: interface {}
+  type t = u32;
+  export b;
+}
+
+interface with-dep {
+  type t = u32;
+}
+
+world tricky-import {
+  use with-dep.{t};
+  export f: func() -> t;
+}

tests/cli/importize.wit.fail1.stderr

@@ -0,0 +1,4 @@
+error: failed to move world exports to imports
+
+Caused by:
+    0: world export `foo` conflicts with import of same name

tests/cli/importize.wit.simple-component.stdout

@@ -0,0 +1,13 @@
+package root:root;
+
+world root-importized {
+  import importize:importize/t;
+}
+package importize:importize {
+  interface t {
+    resource r;
+  }
+  world simple {
+    export t;
+  }
+}