From d4bdb5bc294caea4f851940639e63039f6431716 Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Thu, 11 Jul 2024 14:04:14 -0500
Subject: [PATCH] Improve the output of `wasm-tools component wit` (#1657)

* Improve the `--out-dir` flag of `component wit`

* Lift files in `deps` outside of their folders to be `deps/*.wit`
  instead of `deps/*/main.wit` since this structure is now supported.
* Handle multi-package input documents by placing each package in the
  main folder.
* Name the main wit file after the package name.
* Add some tests for these cases.

* Print all WIT packages by default

This commit changes the output of `wasm-tools component wit` to by
default output all WIT packages using the multi-package syntax
implemented in #1577. This means that a complete view of the WIT will be
seen when running this command instead of just the top-level world.

* Fix build of fuzzer

* Sanitize output to be more platform agnostic
---
 crates/wit-component/src/printing.rs          | 19 +++++----
 crates/wit-component/tests/components.rs      |  2 +-
 crates/wit-component/tests/interfaces.rs      |  2 +-
 crates/wit-component/tests/merge.rs           |  2 +-
 fuzz/src/roundtrip_wit.rs                     |  2 +-
 src/bin/wasm-tools/component.rs               | 39 +++++++++----------
 src/lib.rs                                    | 17 ++++----
 tests/cli.rs                                  | 16 +++++++-
 ...core-wasm-wit-multiple-packages.wit.stdout | 17 ++++++--
 tests/cli/print-core-wasm-wit.wit.stdout      | 17 ++++++--
 .../wit-directory-output-in-deps-folder.wit   | 14 +++++++
 ...directory-output-in-deps-folder.wit.stdout |  3 ++
 tests/cli/wit-directory-output-valid.wit      | 14 +++++++
 .../cli/wit-directory-output-valid.wit.stdout | 16 ++++++++
 tests/cli/wit-directory-output.wit            | 14 +++++++
 tests/cli/wit-directory-output.wit.stdout     |  2 +
 16 files changed, 148 insertions(+), 48 deletions(-)
 create mode 100644 tests/cli/wit-directory-output-in-deps-folder.wit
 create mode 100644 tests/cli/wit-directory-output-in-deps-folder.wit.stdout
 create mode 100644 tests/cli/wit-directory-output-valid.wit
 create mode 100644 tests/cli/wit-directory-output-valid.wit.stdout
 create mode 100644 tests/cli/wit-directory-output.wit
 create mode 100644 tests/cli/wit-directory-output.wit.stdout

diff --git a/crates/wit-component/src/printing.rs b/crates/wit-component/src/printing.rs
index 28406a5023..5c2ce5813f 100644
--- a/crates/wit-component/src/printing.rs
+++ b/crates/wit-component/src/printing.rs
@@ -51,8 +51,13 @@ impl WitPrinter {
     }
 
     /// Print a set of one or more WIT packages into a string.
-    pub fn print(&mut self, resolve: &Resolve, pkg_ids: &[PackageId]) -> Result<String> {
-        let has_multiple_packages = pkg_ids.len() > 1;
+    pub fn print(
+        &mut self,
+        resolve: &Resolve,
+        pkg_ids: &[PackageId],
+        force_print_package_in_curlies: bool,
+    ) -> Result<String> {
+        let print_package_in_curlies = force_print_package_in_curlies || pkg_ids.len() > 1;
         for (i, pkg_id) in pkg_ids.into_iter().enumerate() {
             if i > 0 {
                 self.output.push_str("\n\n");
@@ -68,9 +73,8 @@ impl WitPrinter {
                 self.output.push_str(&format!("@{version}"));
             }
 
-            if has_multiple_packages {
-                self.output.push_str("{");
-                self.output.indent += 1
+            if print_package_in_curlies {
+                self.output.push_str(" {\n");
             } else {
                 self.print_semicolon();
                 self.output.push_str("\n\n");
@@ -96,9 +100,8 @@ impl WitPrinter {
                 writeln!(&mut self.output, "}}")?;
             }
 
-            if has_multiple_packages {
-                self.output.push_str("}");
-                self.output.indent -= 1
+            if print_package_in_curlies {
+                self.output.push_str("}\n");
             }
         }
 
diff --git a/crates/wit-component/tests/components.rs b/crates/wit-component/tests/components.rs
index de2ff79c67..9d00e2f764 100644
--- a/crates/wit-component/tests/components.rs
+++ b/crates/wit-component/tests/components.rs
@@ -171,7 +171,7 @@ fn run_test(path: &Path) -> Result<()> {
                 }
             };
         let wit = WitPrinter::default()
-            .print(&resolve, &[pkg])
+            .print(&resolve, &[pkg], false)
             .context("failed to print WIT")?;
         assert_output(&wit, &component_wit_path)?;
 
diff --git a/crates/wit-component/tests/interfaces.rs b/crates/wit-component/tests/interfaces.rs
index 87c7ac2d4f..55a4cadd88 100644
--- a/crates/wit-component/tests/interfaces.rs
+++ b/crates/wit-component/tests/interfaces.rs
@@ -96,7 +96,7 @@ fn run_test(path: &Path, is_dir: bool) -> Result<()> {
 }
 
 fn assert_print(resolve: &Resolve, pkg_ids: &[PackageId], path: &Path, is_dir: bool) -> Result<()> {
-    let output = WitPrinter::default().print(resolve, &pkg_ids)?;
+    let output = WitPrinter::default().print(resolve, &pkg_ids, false)?;
     for pkg_id in pkg_ids {
         let pkg = &resolve.packages[*pkg_id];
         let expected = if is_dir {
diff --git a/crates/wit-component/tests/merge.rs b/crates/wit-component/tests/merge.rs
index f9a70859bb..3dd25da957 100644
--- a/crates/wit-component/tests/merge.rs
+++ b/crates/wit-component/tests/merge.rs
@@ -46,7 +46,7 @@ fn merging() -> Result<()> {
                         .join("merge")
                         .join(&pkg.name.name)
                         .with_extension("wit");
-                    let output = WitPrinter::default().print(&into, &[id])?;
+                    let output = WitPrinter::default().print(&into, &[id], false)?;
                     assert_output(&expected, &output)?;
                 }
             }
diff --git a/fuzz/src/roundtrip_wit.rs b/fuzz/src/roundtrip_wit.rs
index b15cd8c415..5cc6678265 100644
--- a/fuzz/src/roundtrip_wit.rs
+++ b/fuzz/src/roundtrip_wit.rs
@@ -86,7 +86,7 @@ fn roundtrip_through_printing(file: &str, resolve: &Resolve, wasm: &[u8]) {
     for (id, pkg) in resolve.packages.iter() {
         let mut map = SourceMap::new();
         let pkg_name = &pkg.name;
-        let doc = WitPrinter::default().print(resolve, &[id]).unwrap();
+        let doc = WitPrinter::default().print(resolve, &[id], false).unwrap();
         write_file(&format!("{file}-{pkg_name}.wit"), &doc);
         map.push(format!("{pkg_name}.wit").as_ref(), doc);
         let unresolved = map.parse().unwrap();
diff --git a/src/bin/wasm-tools/component.rs b/src/bin/wasm-tools/component.rs
index d546f82883..1b15b2f5a4 100644
--- a/src/bin/wasm-tools/component.rs
+++ b/src/bin/wasm-tools/component.rs
@@ -523,9 +523,7 @@ impl WitOpts {
         // This interprets all of the output options and performs such a task.
         if self.json {
             self.emit_json(&decoded)?;
-            return Ok(());
-        }
-        if self.wasm || self.wat {
+        } else if self.wasm || self.wat {
             self.emit_wasm(&decoded)?;
         } else {
             self.emit_wit(&decoded)?;
@@ -618,7 +616,6 @@ impl WitOpts {
         assert!(!self.wasm && !self.wat);
 
         let resolve = decoded.resolve();
-        let main = decoded.packages();
 
         let mut printer = WitPrinter::default();
         printer.emit_docs(!self.no_docs);
@@ -641,28 +638,31 @@ impl WitOpts {
                     *cnt += 1;
                 }
 
+                let main = decoded.packages();
                 for (id, pkg) in resolve.packages.iter() {
-                    let output = printer.print(resolve, &[id])?;
-                    let out_dir = if main.contains(&id) {
+                    let is_main = main.contains(&id);
+                    let output = printer.print(resolve, &[id], is_main)?;
+                    let out_dir = if is_main {
                         dir.clone()
                     } else {
-                        let dir = dir.join("deps");
-                        let packages_with_same_name = &names[&pkg.name.name];
-                        if packages_with_same_name.len() == 1 {
-                            dir.join(&pkg.name.name)
+                        dir.join("deps")
+                    };
+                    let packages_with_same_name = &names[&pkg.name.name];
+                    let stem = if packages_with_same_name.len() == 1 {
+                        pkg.name.name.clone()
+                    } else {
+                        let packages_with_same_namespace =
+                            packages_with_same_name[&pkg.name.namespace];
+                        if packages_with_same_namespace == 1 {
+                            format!("{}:{}", pkg.name.namespace, pkg.name.name)
                         } else {
-                            let packages_with_same_namespace =
-                                packages_with_same_name[&pkg.name.namespace];
-                            if packages_with_same_namespace == 1 {
-                                dir.join(format!("{}:{}", pkg.name.namespace, pkg.name.name))
-                            } else {
-                                dir.join(pkg.name.to_string())
-                            }
+                            pkg.name.to_string()
                         }
                     };
                     std::fs::create_dir_all(&out_dir)
                         .with_context(|| format!("failed to create directory: {out_dir:?}"))?;
-                    let path = out_dir.join("main.wit");
+                    let filename = format!("{stem}.wit");
+                    let path = out_dir.join(&filename);
                     std::fs::write(&path, &output)
                         .with_context(|| format!("failed to write file: {path:?}"))?;
                     println!("Writing: {}", path.display());
@@ -672,8 +672,7 @@ impl WitOpts {
                 self.output.output(
                     &self.general,
                     Output::Wit {
-                        resolve: &resolve,
-                        ids: &main,
+                        wit: &decoded,
                         printer,
                     },
                 )?;
diff --git a/src/lib.rs b/src/lib.rs
index b9536eff81..70210b8356 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -143,8 +143,7 @@ pub struct OutputArg {
 pub enum Output<'a> {
     #[cfg(feature = "component")]
     Wit {
-        resolve: &'a wit_parser::Resolve,
-        ids: &'a [wit_parser::PackageId],
+        wit: &'a wit_component::DecodedWasm,
         printer: wit_component::WitPrinter,
     },
     Wasm(&'a [u8]),
@@ -233,12 +232,14 @@ impl OutputArg {
             }
             Output::Json(s) => self.output_str(s),
             #[cfg(feature = "component")]
-            Output::Wit {
-                resolve,
-                ids,
-                mut printer,
-            } => {
-                let output = printer.print(resolve, ids)?;
+            Output::Wit { wit, mut printer } => {
+                let resolve = wit.resolve();
+                let ids = resolve
+                    .packages
+                    .iter()
+                    .map(|(id, _)| id)
+                    .collect::<Vec<_>>();
+                let output = printer.print(resolve, &ids, false)?;
                 self.output_str(&output)
             }
         }
diff --git a/tests/cli.rs b/tests/cli.rs
index 3c8d2cfd23..2e9d16645a 100644
--- a/tests/cli.rs
+++ b/tests/cli.rs
@@ -32,6 +32,7 @@ use std::env;
 use std::io::Write;
 use std::path::{Path, PathBuf};
 use std::process::{Command, Output, Stdio};
+use tempfile::TempDir;
 
 fn main() {
     let mut tests = Vec::new();
@@ -76,6 +77,7 @@ fn run_test(test: &Path, bless: bool) -> Result<()> {
 
     let mut cmd = wasm_tools_exe();
     let mut stdin = None;
+    let tempdir = TempDir::new()?;
     for arg in line.split_whitespace() {
         if arg == "|" {
             let output = execute(&mut cmd, stdin.as_deref(), false)?;
@@ -83,6 +85,8 @@ fn run_test(test: &Path, bless: bool) -> Result<()> {
             cmd = wasm_tools_exe();
         } else if arg == "%" {
             cmd.arg(test);
+        } else if arg == "%tmpdir" {
+            cmd.arg(tempdir.path());
         } else {
             cmd.arg(arg);
         }
@@ -94,12 +98,14 @@ fn run_test(test: &Path, bless: bool) -> Result<()> {
         bless,
         &output.stdout,
         &test.with_extension(&format!("{extension}.stdout")),
+        &tempdir,
     )
     .context("failed to check stdout expectation (auto-update with BLESS=1)")?;
     assert_output(
         bless,
         &output.stderr,
         &test.with_extension(&format!("{extension}.stderr")),
+        &tempdir,
     )
     .context("failed to check stderr expectation (auto-update with BLESS=1)")?;
     Ok(())
@@ -145,7 +151,14 @@ fn execute(cmd: &mut Command, stdin: Option<&[u8]>, should_fail: bool) -> Result
     Ok(output)
 }
 
-fn assert_output(bless: bool, output: &[u8], path: &Path) -> Result<()> {
+fn assert_output(bless: bool, output: &[u8], path: &Path, tempdir: &TempDir) -> Result<()> {
+    let tempdir = tempdir.path().to_str().unwrap();
+    // sanitize the output to be consistent across platforms and handle per-test
+    // differences such as `%tmpdir`.
+    let output = String::from_utf8_lossy(output)
+        .replace(tempdir, "%tmpdir")
+        .replace("\\", "/");
+
     if bless {
         if output.is_empty() {
             drop(std::fs::remove_file(path));
@@ -162,7 +175,6 @@ fn assert_output(bless: bool, output: &[u8], path: &Path) -> Result<()> {
             Ok(())
         }
     } else {
-        let output = std::str::from_utf8(output)?;
         let contents = std::fs::read_to_string(path)
             .with_context(|| format!("failed to read {path:?}"))?
             .replace("\r\n", "\n");
diff --git a/tests/cli/print-core-wasm-wit-multiple-packages.wit.stdout b/tests/cli/print-core-wasm-wit-multiple-packages.wit.stdout
index 75feb032d8..57c65fce90 100644
--- a/tests/cli/print-core-wasm-wit-multiple-packages.wit.stdout
+++ b/tests/cli/print-core-wasm-wit-multiple-packages.wit.stdout
@@ -1,5 +1,16 @@
-package root:root;
+package root:root {
+  world root {
+    import bar:bar/my-interface;
+  }
+}
+
+
+package bar:bar {
+  interface my-interface {
+    foo: func();
+  }
 
-world root {
-  import bar:bar/my-interface;
+  world my-world {
+    import my-interface;
+  }
 }
diff --git a/tests/cli/print-core-wasm-wit.wit.stdout b/tests/cli/print-core-wasm-wit.wit.stdout
index e560f82ea0..c1eb1f6319 100644
--- a/tests/cli/print-core-wasm-wit.wit.stdout
+++ b/tests/cli/print-core-wasm-wit.wit.stdout
@@ -1,5 +1,16 @@
-package root:root;
+package root:root {
+  world root {
+    import foo:foo/my-interface;
+  }
+}
+
+
+package foo:foo {
+  interface my-interface {
+    foo: func();
+  }
 
-world root {
-  import foo:foo/my-interface;
+  world my-world {
+    import my-interface;
+  }
 }
diff --git a/tests/cli/wit-directory-output-in-deps-folder.wit b/tests/cli/wit-directory-output-in-deps-folder.wit
new file mode 100644
index 0000000000..58ff346e05
--- /dev/null
+++ b/tests/cli/wit-directory-output-in-deps-folder.wit
@@ -0,0 +1,14 @@
+// RUN: component embed --dummy % --world a:b/c | component wit --out-dir %tmpdir
+
+package a:b {
+  interface foo {}
+
+  world c {
+    import foo;
+    import a:c/foo;
+  }
+}
+
+package a:c {
+  interface foo {}
+}
diff --git a/tests/cli/wit-directory-output-in-deps-folder.wit.stdout b/tests/cli/wit-directory-output-in-deps-folder.wit.stdout
new file mode 100644
index 0000000000..8c2fc622ef
--- /dev/null
+++ b/tests/cli/wit-directory-output-in-deps-folder.wit.stdout
@@ -0,0 +1,3 @@
+Writing: %tmpdir/root.wit
+Writing: %tmpdir/deps/b.wit
+Writing: %tmpdir/deps/c.wit
diff --git a/tests/cli/wit-directory-output-valid.wit b/tests/cli/wit-directory-output-valid.wit
new file mode 100644
index 0000000000..01e60dca1b
--- /dev/null
+++ b/tests/cli/wit-directory-output-valid.wit
@@ -0,0 +1,14 @@
+// RUN: component wit % --out-dir %tmpdir | component wit %tmpdir
+
+package a:b {
+  interface foo {}
+
+  world c {
+    import foo;
+    import a:c/foo;
+  }
+}
+
+package a:c {
+  interface foo {}
+}
diff --git a/tests/cli/wit-directory-output-valid.wit.stdout b/tests/cli/wit-directory-output-valid.wit.stdout
new file mode 100644
index 0000000000..33576bc7a2
--- /dev/null
+++ b/tests/cli/wit-directory-output-valid.wit.stdout
@@ -0,0 +1,16 @@
+package a:c {
+  interface foo {
+  }
+
+}
+
+
+package a:b {
+  interface foo {
+  }
+
+  world c {
+    import foo;
+    import a:c/foo;
+  }
+}
diff --git a/tests/cli/wit-directory-output.wit b/tests/cli/wit-directory-output.wit
new file mode 100644
index 0000000000..88d6f83d7d
--- /dev/null
+++ b/tests/cli/wit-directory-output.wit
@@ -0,0 +1,14 @@
+// RUN: component wit % --out-dir %tmpdir
+
+package a:b {
+  interface foo {}
+
+  world c {
+    import foo;
+    import a:c/foo;
+  }
+}
+
+package a:c {
+  interface foo {}
+}
diff --git a/tests/cli/wit-directory-output.wit.stdout b/tests/cli/wit-directory-output.wit.stdout
new file mode 100644
index 0000000000..9524cea580
--- /dev/null
+++ b/tests/cli/wit-directory-output.wit.stdout
@@ -0,0 +1,2 @@
+Writing: %tmpdir/c.wit
+Writing: %tmpdir/b.wit