Repository: bytecodealliance/wasm-tools
Pull request #1631 (2024-06-22 23:38:16)

Support `include` statements in `wasm-parser`
We already had basic support for `include`, but this tests it and fixes it. Thanks!


Patch:
diff --git a/crates/wasm-encoder/src/core/code.rs b/crates/wasm-encoder/src/core/code.rs
index 25de73be..7bee3597 100644
--- a/crates/wasm-encoder/src/core/code.rs
+++ b/crates/wasm-encoder/src/core/code.rs
@@ -274,6 +274,45 @@ impl Function {
         self.bytes.len()
     }
 
+    /// Unwraps and returns the raw byte encoding of this function.
+    ///
+    /// This encoding doesn't include the variable-width size field
+    /// that `encode` will write before the added bytes. As such, its
+    /// length will match the return value of [`byte_len`].
+    ///
+    /// # Use Case
+    ///
+    /// This raw byte form is suitable for later using with
+    /// [`CodeSection::raw`]. Note that it *differs* from what results
+    /// from [`Function::encode`] precisely due to the *lack* of the
+    /// length prefix; [`CodeSection::raw`] will use this. Using
+    /// [`Function::encode`] instead produces bytes that cannot be fed
+    /// into other wasm-encoder types without stripping off the length
+    /// prefix, which is awkward and error-prone.
+    ///
+    /// This method combined with [`CodeSection::raw`] may be useful
+    /// together if one wants to save the result of function encoding
+    /// and use it later: for example, caching the result of some code
+    /// generation process.
+    ///
+    /// For example:
+    ///
+    /// ```
+    /// # use wasm_encoder::{CodeSection, Function, Instruction};
+    /// let mut f = Function::new([]);
+    /// f.instruction(&Instruction::End);
+    /// let bytes = f.into_raw_body();
+    /// // (save `bytes` somewhere for later use)
+    /// let mut code = CodeSection::new();
+    /// code.raw(&bytes[..]);
+    ///
+    /// assert_eq!(2, bytes.len());  // Locals count, then `end`
+    /// assert_eq!(3, code.byte_len()); // Function length byte, function body
+    /// ```
+    pub fn into_raw_body(self) -> Vec<u8> {
+        self.bytes
+    }
+
     /// Parses a single instruction from `reader` and adds it to `self`.
     #[cfg(feature = "wasmparser")]
     pub fn parse(
@@ -3753,4 +3792,23 @@ mod tests {
 
         assert_eq!(f1.bytes, f2.bytes)
     }
+
+    #[test]
+    fn func_raw_bytes() {
+        use super::*;
+
+        let mut f = Function::new([(1, ValType::I32), (1, ValType::F32)]);
+        f.instruction(&Instruction::End);
+        let mut code_from_func = CodeSection::new();
+        code_from_func.function(&f);
+        let bytes = f.into_raw_body();
+        let mut code_from_raw = CodeSection::new();
+        code_from_raw.raw(&bytes[..]);
+
+        let mut c1 = vec![];
+        code_from_func.encode(&mut c1);
+        let mut c2 = vec![];
+        code_from_raw.encode(&mut c2);
+        assert_eq!(c1, c2);
+    }
 }
diff --git a/crates/wit-encoder/src/include.rs b/crates/wit-encoder/src/include.rs
index 747186cc..9f276381 100644
--- a/crates/wit-encoder/src/include.rs
+++ b/crates/wit-encoder/src/include.rs
@@ -6,7 +6,7 @@ use crate::{Ident, Render};
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct Include {
     use_path: Ident,
-    include_names_list: Vec<String>,
+    include_names_list: Vec<(String, String)>,
 }
 
 impl Include {
@@ -16,11 +16,29 @@ impl Include {
             include_names_list: vec![],
         }
     }
+
+    pub fn with(&mut self, id: &str, alias: &str) {
+        self.include_names_list
+            .push((id.to_string(), alias.to_string()));
+    }
 }
 
 impl Render for Include {
     fn render(&self, f: &mut fmt::Formatter<'_>, opts: &crate::RenderOpts) -> fmt::Result {
-        write!(f, "{}include {};\n", opts.spaces(), self.use_path)?;
+        match self.include_names_list.len() {
+            0 => write!(f, "{}include {};\n", opts.spaces(), self.use_path)?,
+            len => {
+                write!(f, "{}include {} with {{ ", opts.spaces(), self.use_path)?;
+                for (i, (id, alias)) in self.include_names_list.iter().enumerate() {
+                    if i == len - 1 {
+                        write!(f, "{id} as {alias}")?;
+                    } else {
+                        write!(f, "{id} as {alias}, ")?;
+                    }
+                }
+                write!(f, " }};\n")?;
+            }
+        }
         Ok(())
     }
 }
diff --git a/crates/wit-encoder/src/world.rs b/crates/wit-encoder/src/world.rs
index 5f837678..7203524b 100644
--- a/crates/wit-encoder/src/world.rs
+++ b/crates/wit-encoder/src/world.rs
@@ -52,8 +52,8 @@ impl World {
     pub fn function_export(&mut self, value: StandaloneFunc) {
         self.item(WorldItem::function_export(value));
     }
-    pub fn include(&mut self, value: impl Into<Ident>) {
-        self.item(WorldItem::include(value));
+    pub fn include(&mut self, include: Include) {
+        self.item(WorldItem::Include(include));
     }
 
     /// Set the documentation


Test patch:
diff --git a/crates/wit-encoder/tests/include-aliases.rs b/crates/wit-encoder/tests/include-aliases.rs
new file mode 100644
index 00000000..140ce44d
--- /dev/null
+++ b/crates/wit-encoder/tests/include-aliases.rs
@@ -0,0 +1,68 @@
+use pretty_assertions::assert_eq;
+use wit_encoder::{Include, Package, PackageName, StandaloneFunc, World};
+
+const PACKAGE: &str = indoc::indoc! {"
+    package foo:foo;
+
+    world foo {
+      import a: func();
+    }
+
+    world bar {
+      import a: func();
+      import b: func();
+    }
+
+    world baz {
+      import a: func();
+      import b: func();
+      import c: func();
+    }
+
+    world quux {
+      include foo with { a as b };
+      include bar with { a as b, b as c };
+      include baz with { a as b, b as c, c as d };
+    }
+"};
+
+#[test]
+fn concrete_types() {
+    let mut package = Package::new(PackageName::new("foo", "foo", None));
+
+    let mut world = World::new("foo");
+    world.function_import(StandaloneFunc::new("a"));
+    package.world(world);
+
+    let mut world = World::new("bar");
+    world.function_import(StandaloneFunc::new("a"));
+    world.function_import(StandaloneFunc::new("b"));
+    package.world(world);
+
+    let mut world = World::new("baz");
+    world.function_import(StandaloneFunc::new("a"));
+    world.function_import(StandaloneFunc::new("b"));
+    world.function_import(StandaloneFunc::new("c"));
+    package.world(world);
+
+    let mut world = World::new("quux");
+
+    let mut include = Include::new("foo");
+    include.with("a", "b");
+    world.include(include);
+
+    let mut include = Include::new("bar");
+    include.with("a", "b");
+    include.with("b", "c");
+    world.include(include);
+
+    let mut include = Include::new("baz");
+    include.with("a", "b");
+    include.with("b", "c");
+    include.with("c", "d");
+    world.include(include);
+
+    package.world(world);
+
+    assert_eq!(package.to_string(), PACKAGE);
+}
diff --git a/crates/wit-encoder/tests/include-reps.rs b/crates/wit-encoder/tests/include-reps.rs
index 5341e053..a05fb3ac 100644
--- a/crates/wit-encoder/tests/include-reps.rs
+++ b/crates/wit-encoder/tests/include-reps.rs
@@ -1,5 +1,5 @@
 use pretty_assertions::assert_eq;
-use wit_encoder::{Interface, Package, PackageName, World};
+use wit_encoder::{Include, Interface, Package, PackageName, World};
 
 const PACKAGE: &str = indoc::indoc! {"
     package foo:foo;
@@ -33,9 +33,9 @@ fn concrete_types() {
     package.world(world);
 
     let mut world = World::new("foo");
-    world.include("bar");
-    world.include("bar");
-    world.include("bar");
+    world.include(Include::new("bar"));
+    world.include(Include::new("bar"));
+    world.include(Include::new("bar"));
     package.world(world);
 
     assert_eq!(package.to_string(), PACKAGE);
diff --git a/crates/wit-encoder/tests/kebab-name-include-with.rs b/crates/wit-encoder/tests/kebab-name-include-with.rs
new file mode 100644
index 00000000..f8afb2ab
--- /dev/null
+++ b/crates/wit-encoder/tests/kebab-name-include-with.rs
@@ -0,0 +1,41 @@
+use pretty_assertions::assert_eq;
+use wit_encoder::{Include, Package, PackageName, StandaloneFunc, World};
+
+const PACKAGE: &str = indoc::indoc! {"
+    package foo:foo;
+
+    world foo {
+      import a: func();
+    }
+
+    world bar {
+      import a: func();
+    }
+
+    world baz {
+      include bar;
+      include foo with { a as b };
+    }
+"};
+
+#[test]
+fn concrete_types() {
+    let mut package = Package::new(PackageName::new("foo", "foo", None));
+
+    let mut world = World::new("foo");
+    world.function_import(StandaloneFunc::new("a"));
+    package.world(world);
+
+    let mut world = World::new("bar");
+    world.function_import(StandaloneFunc::new("a"));
+    package.world(world);
+
+    let mut world = World::new("baz");
+    world.include(Include::new("bar"));
+    let mut include = Include::new("foo");
+    include.with("a", "b");
+    world.include(include);
+    package.world(world);
+
+    assert_eq!(package.to_string(), PACKAGE);
+}