refactor: expose fn for embedding component metadata (#1233)

The logic for embedding component metadata as a custom section was
used only in the course of `wasm-tools component embed`, which made it
difficult to use from outside wasm-tools.

This commit refeactors the code that actually writes out the metadata
to a function that can be re-used and is exposed by wit-component.

Signed-off-by: Victor Adossi <[email protected]>

Author: vados-cosmonic

crates/wit-component/src/lib.rs

@@ -2,11 +2,12 @@
 
 #![deny(missing_docs)]
 
-use anyhow::{bail, Context, Result};
 use std::str::FromStr;
-use std::{fmt::Display, path::Path};
-use wasm_encoder::CanonicalOption;
-use wit_parser::{PackageId, Resolve, UnresolvedPackage};
+use std::{borrow::Cow, fmt::Display, path::Path};
+
+use anyhow::{bail, Context, Result};
+use wasm_encoder::{CanonicalOption, Encode, Section};
+use wit_parser::{PackageId, Resolve, UnresolvedPackage, WorldId};
 
 mod decoding;
 mod encoding;
@@ -178,3 +179,94 @@ pub fn is_wasm_binary_or_wat(bytes: impl AsRef<[u8]>) -> bool {
 
     false
 }
+
+/// Embed component metadata in a buffer of bytes that contains a Wasm module
+pub fn embed_component_metadata(
+    bytes: &mut Vec<u8>,
+    wit_resolver: &Resolve,
+    world: WorldId,
+    encoding: StringEncoding,
+) -> Result<()> {
+    let encoded = metadata::encode(&wit_resolver, world, encoding, None)?;
+
+    let section = wasm_encoder::CustomSection {
+        name: "component-type".into(),
+        data: Cow::Borrowed(&encoded),
+    };
+    bytes.push(section.id());
+    section.encode(bytes);
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use std::path::Path;
+
+    use anyhow::Result;
+    use wasmparser::Payload;
+    use wit_parser::{Resolve, UnresolvedPackage};
+
+    use super::{embed_component_metadata, StringEncoding};
+
+    const MODULE_WAT: &str = r#"
+(module
+  (type (;0;) (func))
+  (func (;0;) (type 0)
+    nop
+  )
+)
+"#;
+
+    const COMPONENT_WIT: &str = r#"
+package test:foo;
+world test-world {}
+"#;
+
+    #[test]
+    fn component_metadata_embedding_works() -> Result<()> {
+        let mut bytes = wat::parse_str(MODULE_WAT)?;
+
+        // Get original len & custom section count
+        let original_len = bytes.len();
+        let payloads = wasmparser::Parser::new(0).parse_all(&bytes);
+        let original_custom_section_count = payloads.fold(0, |acc, payload| {
+            if let Ok(Payload::CustomSection { .. }) = payload {
+                acc + 1
+            } else {
+                acc
+            }
+        });
+
+        // Parse pre-canned WIT to build resolver
+        let mut resolver = Resolve::default();
+        let pkg = UnresolvedPackage::parse(&Path::new("in-code.wit"), COMPONENT_WIT)?;
+        let pkg_id = resolver.push(pkg)?;
+        let world = resolver.select_world(pkg_id, Some("test-world").into())?;
+
+        // Embed component metadata
+        embed_component_metadata(&mut bytes, &resolver, world, StringEncoding::UTF8)?;
+
+        // Re-retrieve custom section count, and search for the component-type custom section along the way
+        let mut found_component_section = false;
+        let new_custom_section_count =
+            wasmparser::Parser::new(0)
+                .parse_all(&bytes)
+                .fold(0, |acc, payload| {
+                    if let Ok(Payload::CustomSection(reader)) = payload {
+                        if reader.name() == "component-type" {
+                            found_component_section = true;
+                        }
+                        acc + 1
+                    } else {
+                        acc
+                    }
+                });
+
+        assert!(original_len < bytes.len());
+        assert_eq!(original_custom_section_count + 1, new_custom_section_count);
+        assert!(found_component_section);
+
+        Ok(())
+    }
+}

src/bin/wasm-tools/component.rs

@@ -1,16 +1,16 @@
 //! The WebAssembly component tool command line interface.
 
-use anyhow::{bail, Context, Result};
-use clap::Parser;
-use std::borrow::Cow;
 use std::collections::HashMap;
 use std::io::Read;
 use std::path::{Path, PathBuf};
-use wasm_encoder::{Encode, Section};
+
+use anyhow::{bail, Context, Result};
+use clap::Parser;
+
 use wasm_tools::Output;
 use wit_component::{
-    is_wasm_binary_or_wat, parse_wit_from_path, ComponentEncoder, DecodedWasm, Linker,
-    StringEncoding, WitPrinter,
+    embed_component_metadata, is_wasm_binary_or_wat, parse_wit_from_path, ComponentEncoder,
+    DecodedWasm, Linker, StringEncoding, WitPrinter,
 };
 use wit_parser::{Resolve, UnresolvedPackage};
 
@@ -216,21 +216,15 @@ impl EmbedOpts {
         let (resolve, id) = parse_wit_from_path(self.wit)?;
         let world = resolve.select_world(id, self.world.as_deref())?;
 
-        let encoded = wit_component::metadata::encode(
+        let mut wasm = wasm.unwrap_or_else(|| wit_component::dummy_module(&resolve, world));
+
+        embed_component_metadata(
+            &mut wasm,
             &resolve,
             world,
             self.encoding.unwrap_or(StringEncoding::UTF8),
-            None,
         )?;
 
-        let section = wasm_encoder::CustomSection {
-            name: "component-type".into(),
-            data: Cow::Borrowed(&encoded),
-        };
-        let mut wasm = wasm.unwrap_or_else(|| wit_component::dummy_module(&resolve, world));
-        wasm.push(section.id());
-        section.encode(&mut wasm);
-
         self.io.output(Output::Wasm {
             bytes: &wasm,
             wat: self.wat,