From 7b36db9c12f1ae82806602be48bede1fac05dda4 Mon Sep 17 00:00:00 2001
From: Alex Zaslavsky <azaslavsky@users.noreply.github.com>
Date: Fri, 31 May 2024 17:49:46 -0700
Subject: [PATCH] Fix pretty printing signature and output

---
 crates/wit-component/src/printing.rs     | 21 +++++++++++++++++----
 crates/wit-component/tests/components.rs |  2 +-
 crates/wit-component/tests/interfaces.rs |  2 +-
 crates/wit-component/tests/merge.rs      |  2 +-
 src/bin/wasm-tools/component.rs          |  4 ++--
 5 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/crates/wit-component/src/printing.rs b/crates/wit-component/src/printing.rs
index cb0488ac92..28406a5023 100644
--- a/crates/wit-component/src/printing.rs
+++ b/crates/wit-component/src/printing.rs
@@ -51,13 +51,14 @@ impl WitPrinter {
     }
 
     /// Print a set of one or more WIT packages into a string.
-    pub fn print(&mut self, resolve: &Resolve, pkg_ids: Vec<PackageId>) -> Result<String> {
+    pub fn print(&mut self, resolve: &Resolve, pkg_ids: &[PackageId]) -> Result<String> {
+        let has_multiple_packages = pkg_ids.len() > 1;
         for (i, pkg_id) in pkg_ids.into_iter().enumerate() {
             if i > 0 {
                 self.output.push_str("\n\n");
             }
 
-            let pkg = &resolve.packages[pkg_id];
+            let pkg = &resolve.packages[pkg_id.clone()];
             self.print_docs(&pkg.docs);
             self.output.push_str("package ");
             self.print_name(&pkg.name.namespace);
@@ -66,8 +67,15 @@ impl WitPrinter {
             if let Some(version) = &pkg.name.version {
                 self.output.push_str(&format!("@{version}"));
             }
-            self.print_semicolon();
-            self.output.push_str("\n\n");
+
+            if has_multiple_packages {
+                self.output.push_str("{");
+                self.output.indent += 1
+            } else {
+                self.print_semicolon();
+                self.output.push_str("\n\n");
+            }
+
             for (name, id) in pkg.interfaces.iter() {
                 self.print_docs(&resolve.interfaces[*id].docs);
                 self.print_stability(&resolve.interfaces[*id].stability);
@@ -87,6 +95,11 @@ impl WitPrinter {
                 self.print_world(resolve, *id)?;
                 writeln!(&mut self.output, "}}")?;
             }
+
+            if has_multiple_packages {
+                self.output.push_str("}");
+                self.output.indent -= 1
+            }
         }
 
         Ok(std::mem::take(&mut self.output).into())
diff --git a/crates/wit-component/tests/components.rs b/crates/wit-component/tests/components.rs
index a0a61a0c06..8de9056674 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, vec![pkg])
+            .print(&resolve, &[pkg])
             .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 be1bce46ca..30f7f389aa 100644
--- a/crates/wit-component/tests/interfaces.rs
+++ b/crates/wit-component/tests/interfaces.rs
@@ -101,7 +101,7 @@ fn assert_print(
     path: &Path,
     is_dir: bool,
 ) -> Result<()> {
-    let output = WitPrinter::default().print(resolve, pkg_ids.clone())?;
+    let output = WitPrinter::default().print(resolve, &pkg_ids)?;
     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 7620232282..f9a70859bb 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, vec![id])?;
+                    let output = WitPrinter::default().print(&into, &[id])?;
                     assert_output(&expected, &output)?;
                 }
             }
diff --git a/src/bin/wasm-tools/component.rs b/src/bin/wasm-tools/component.rs
index bd9426c7b0..b6fb37d8d5 100644
--- a/src/bin/wasm-tools/component.rs
+++ b/src/bin/wasm-tools/component.rs
@@ -637,7 +637,7 @@ impl WitOpts {
                 }
 
                 for (id, pkg) in resolve.packages.iter() {
-                    let output = printer.print(resolve, vec![id])?;
+                    let output = printer.print(resolve, &[id])?;
                     let out_dir = if main.contains(&id) {
                         dir.clone()
                     } else {
@@ -664,7 +664,7 @@ impl WitOpts {
                 }
             }
             None => {
-                let output = printer.print(resolve, main)?;
+                let output = printer.print(resolve, &main)?;
                 self.output.output(Output::Wat(&output))?;
             }
         }