Transition more encoding in wast to using wasm-encoder (#1803)

* wast: Use wasm-encoder for import sections

* wast: Use wasm-encoder for the table section

* wast: Use wasm-encoder for the memory section

* wast: Use wasm-encoder for the tag section

* wast: Use wasm-encoder for the global section

* wast: Use wasm-encoder for the element section

* wast: Use wasm-encoder for the data section

* wast: Use wasm-encoder for the func section

* wast: Use wasm-encoder for the name section

* wast: Use wasm-encoder for more custom sections

* wast: Use wasm-encoder for encoding branch hints

This additionally entailed adding support to wasm-encoder for emitting
branch hints.

* Add example for `wasm_encoder::BranchHints`

* Fix rebase conflict

Author: alexcrichton

Cargo.lock

@@ -30,3 +30,4 @@ wasmparser = { optional = true, workspace = true }
 anyhow = { workspace = true }
 tempfile = "3.2.0"
 wasmparser = { path = "../wasmparser" }
+wasmprinter = { path = "../wasmprinter" }

crates/wasm-encoder/Cargo.toml

@@ -1,3 +1,4 @@
+mod branch_hints;
 mod code;
 mod custom;
 mod data;
@@ -16,6 +17,7 @@ mod tables;
 mod tags;
 mod types;
 
+pub use branch_hints::*;
 pub use code::*;
 pub use custom::*;
 pub use data::*;

crates/wasm-encoder/src/core.rs

@@ -0,0 +1,125 @@
+use crate::{CustomSection, Encode, Section, SectionId};
+use std::borrow::Cow;
+
+/// Helper structure to encode the `metadata.code.branch_hint` custom section.
+///
+/// This section was defined in the branch-hinting proposal for WebAssembly:
+/// <https://github.com/WebAssembly/branch-hinting>.
+///
+/// # Example
+///
+/// ```
+/// use wasm_encoder::*;
+///
+/// let mut module = Module::new();
+///
+/// let mut types = TypeSection::new();
+/// types.ty().function([], []);
+/// module.section(&types);
+///
+/// let mut funcs = FunctionSection::new();
+/// funcs.function(0);
+/// module.section(&funcs);
+///
+/// let mut code = CodeSection::new();
+/// let mut body = Function::new([]);
+///
+/// body.instruction(&Instruction::I32Const(1));
+/// let if_offset = body.byte_len();
+/// body.instruction(&Instruction::If(BlockType::Empty));
+/// body.instruction(&Instruction::End);
+/// body.instruction(&Instruction::End);
+/// code.function(&body);
+///
+/// let mut hints = BranchHints::new();
+/// hints.function_hints(0, [BranchHint {
+///     branch_func_offset: if_offset as u32,
+///     branch_hint_value: 1, // taken
+/// }]);
+/// module.section(&hints);
+/// module.section(&code);
+///
+/// let wasm = module.finish();
+/// let wat = wasmprinter::print_bytes(&wasm).unwrap();
+/// assert_eq!(wat, r#"(module
+///   (type (;0;) (func))
+///   (func (;0;) (type 0)
+///     i32.const 1
+///     (@metadata.code.branch_hint "\01")
+///     if ;; label = @1
+///     end
+///   )
+/// )
+/// "#);
+///
+/// ```
+#[derive(Default, Debug)]
+pub struct BranchHints {
+    bytes: Vec<u8>,
+    num_hints: u32,
+}
+
+/// A single branch hint within a function.
+#[derive(Debug, Clone, Copy)]
+pub struct BranchHint {
+    /// The offset, in bytes from the beginning of the function, to the `if`
+    /// instruction that this is hinting.
+    pub branch_func_offset: u32,
+    /// The value of the hint, 0 for not taken and 1 for taken.
+    pub branch_hint_value: u32,
+}
+
+impl BranchHints {
+    /// Construct an empty encoder for the branch hints custom section.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Adds a new set of function hints for the `func` specified.
+    pub fn function_hints<I>(&mut self, func: u32, hints: I)
+    where
+        I: IntoIterator<Item = BranchHint>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        self.num_hints += 1;
+        func.encode(&mut self.bytes);
+        let hints = hints.into_iter();
+        hints.len().encode(&mut self.bytes);
+        for hint in hints {
+            hint.branch_func_offset.encode(&mut self.bytes);
+            1u32.encode(&mut self.bytes);
+            hint.branch_hint_value.encode(&mut self.bytes);
+        }
+    }
+
+    /// Returns if this is an empty section.
+    pub fn is_empty(&self) -> bool {
+        self.num_hints == 0
+    }
+
+    /// Returns the number of functions that have hints registered in this
+    /// sectino.
+    pub fn len(&self) -> u32 {
+        self.num_hints
+    }
+}
+
+impl Encode for BranchHints {
+    fn encode(&self, sink: &mut Vec<u8>) {
+        let mut data = Vec::new();
+        self.num_hints.encode(&mut data);
+        data.extend(&self.bytes);
+
+        CustomSection {
+            name: "metadata.code.branch_hint".into(),
+            data: Cow::Borrowed(&data),
+        }
+        .encode(sink);
+    }
+}
+
+impl Section for BranchHints {
+    fn id(&self) -> u8 {
+        SectionId::Custom.into()
+    }
+}

crates/wasm-encoder/src/core/branch_hints.rs

@@ -1,13 +1,13 @@
 use crate::component::*;
 use crate::core;
 use crate::core::EncodeOptions;
-use crate::token::{Id, Index, NameAnnotation, Span};
+use crate::token::{Id, Index, NameAnnotation};
 use wasm_encoder::{
     CanonicalFunctionSection, ComponentAliasSection, ComponentCoreTypeEncoder,
     ComponentDefinedTypeEncoder, ComponentExportSection, ComponentImportSection,
     ComponentInstanceSection, ComponentNameSection, ComponentSection, ComponentSectionId,
     ComponentStartSection, ComponentTypeEncoder, ComponentTypeSection, CoreTypeSection,
-    InstanceSection, NameMap, NestedComponentSection, RawSection, SectionId, SubType,
+    InstanceSection, NameMap, NestedComponentSection, RawSection, SubType,
 };
 
 pub fn encode(component: &Component<'_>, options: &EncodeOptions) -> Vec<u8> {
@@ -177,19 +177,12 @@ impl<'a> Encoder<'a> {
     fn encode_custom(&mut self, custom: &Custom) {
         // Flush any in-progress section before encoding the customs section
         self.flush(None);
-        self.component.section(custom);
+        self.component.section(&custom.to_section());
     }
 
     fn encode_producers(&mut self, custom: &core::Producers) {
-        use crate::encode::Encode;
-
-        let mut data = Vec::new();
-        custom.encode(&mut data);
-        self.encode_custom(&Custom {
-            name: "producers",
-            span: Span::from_offset(0),
-            data: vec![&data],
-        })
+        self.flush(None);
+        self.component.section(&custom.to_section());
     }
 
     fn encode_core_module(&mut self, module: &CoreModule<'a>, options: &EncodeOptions) {
@@ -549,32 +542,16 @@ fn get_name<'a>(id: &Option<Id<'a>>, name: &Option<NameAnnotation<'a>>) -> Optio
     })
 }
 
-// This implementation is much like `wasm_encoder::CustomSection`, except
-// that it extends via a list of slices instead of a single slice.
-impl wasm_encoder::Encode for Custom<'_> {
-    fn encode(&self, sink: &mut Vec<u8>) {
-        let mut buf = [0u8; 5];
-        let encoded_name_len =
-            leb128::write::unsigned(&mut &mut buf[..], u64::try_from(self.name.len()).unwrap())
-                .unwrap();
-        let data_len = self.data.iter().fold(0, |acc, s| acc + s.len());
-
-        // name length
-        (encoded_name_len + self.name.len() + data_len).encode(sink);
-
-        // name
-        self.name.encode(sink);
-
-        // data
-        for s in &self.data {
-            sink.extend(*s);
+impl Custom<'_> {
+    fn to_section(&self) -> wasm_encoder::CustomSection<'_> {
+        let mut ret = Vec::new();
+        for list in self.data.iter() {
+            ret.extend_from_slice(list);
+        }
+        wasm_encoder::CustomSection {
+            name: self.name.into(),
+            data: ret.into(),
         }
-    }
-}
-
-impl wasm_encoder::ComponentSection for Custom<'_> {
-    fn id(&self) -> u8 {
-        SectionId::Custom.into()
     }
 }
 
@@ -629,63 +606,6 @@ impl From<core::HeapType<'_>> for wasm_encoder::HeapType {
     }
 }
 
-impl From<&core::ItemKind<'_>> for wasm_encoder::EntityType {
-    fn from(kind: &core::ItemKind) -> Self {
-        match kind {
-            core::ItemKind::Func(t) => Self::Function(t.into()),
-            core::ItemKind::Table(t) => Self::Table((*t).into()),
-            core::ItemKind::Memory(t) => Self::Memory((*t).into()),
-            core::ItemKind::Global(t) => Self::Global((*t).into()),
-            core::ItemKind::Tag(t) => Self::Tag(t.into()),
-        }
-    }
-}
-
-impl From<core::TableType<'_>> for wasm_encoder::TableType {
-    fn from(ty: core::TableType) -> Self {
-        Self {
-            element_type: ty.elem.into(),
-            minimum: ty.limits.min,
-            maximum: ty.limits.max,
-            table64: ty.limits.is64,
-            shared: ty.shared,
-        }
-    }
-}
-
-impl From<core::MemoryType> for wasm_encoder::MemoryType {
-    fn from(ty: core::MemoryType) -> Self {
-        Self {
-            minimum: ty.limits.min,
-            maximum: ty.limits.max,
-            memory64: ty.limits.is64,
-            shared: ty.shared,
-            page_size_log2: ty.page_size_log2,
-        }
-    }
-}
-
-impl From<core::GlobalType<'_>> for wasm_encoder::GlobalType {
-    fn from(ty: core::GlobalType) -> Self {
-        Self {
-            val_type: ty.ty.into(),
-            mutable: ty.mutable,
-            shared: ty.shared,
-        }
-    }
-}
-
-impl From<&core::TagType<'_>> for wasm_encoder::TagType {
-    fn from(ty: &core::TagType) -> Self {
-        match ty {
-            core::TagType::Exception(r) => Self {
-                kind: wasm_encoder::TagKind::Exception,
-                func_type_idx: r.into(),
-            },
-        }
-    }
-}
-
 impl<T: std::fmt::Debug> From<&core::TypeUse<'_, T>> for u32 {
     fn from(u: &core::TypeUse<'_, T>) -> Self {
         match &u.index {
@@ -898,10 +818,10 @@ impl From<&ModuleType<'_>> for wasm_encoder::ModuleType {
                     _ => unreachable!("only outer type aliases are supported"),
                 },
                 ModuleTypeDecl::Import(i) => {
-                    encoded.import(i.module, i.field, (&i.item.kind).into());
+                    encoded.import(i.module, i.field, i.item.to_entity_type());
                 }
                 ModuleTypeDecl::Export(name, item) => {
-                    encoded.export(name, (&item.kind).into());
+                    encoded.export(name, item.to_entity_type());
                 }
             }
         }