wasm_encoder/core/
branch_hints.rs

1use crate::{CustomSection, Encode, Section, SectionId};
2use alloc::borrow::Cow;
3use alloc::vec::Vec;
4
5/// Helper structure to encode the `metadata.code.branch_hint` custom section.
6///
7/// This section was defined in the branch-hinting proposal for WebAssembly:
8/// <https://github.com/WebAssembly/branch-hinting>.
9///
10/// # Example
11///
12/// ```
13/// use wasm_encoder::*;
14///
15/// let mut module = Module::new();
16///
17/// let mut types = TypeSection::new();
18/// types.ty().function([], []);
19/// module.section(&types);
20///
21/// let mut funcs = FunctionSection::new();
22/// funcs.function(0);
23/// module.section(&funcs);
24///
25/// let mut code = CodeSection::new();
26/// let mut body = Function::new([]);
27///
28/// body.instructions().i32_const(1);
29/// let if_offset = body.byte_len();
30/// body.instructions()
31///     .if_(BlockType::Empty)
32///     .end()
33///     .end();
34/// code.function(&body);
35///
36/// let mut hints = BranchHints::new();
37/// hints.function_hints(0, [BranchHint {
38///     branch_func_offset: if_offset as u32,
39///     branch_hint_value: 1, // taken
40/// }]);
41/// module.section(&hints);
42/// module.section(&code);
43///
44/// let wasm = module.finish();
45/// let wat = wasmprinter::print_bytes(&wasm).unwrap();
46/// assert_eq!(wat, r#"(module
47///   (type (;0;) (func))
48///   (func (;0;) (type 0)
49///     i32.const 1
50///     (@metadata.code.branch_hint "\01")
51///     if ;; label = @1
52///     end
53///   )
54/// )
55/// "#);
56///
57/// ```
58#[derive(Default, Debug)]
59pub struct BranchHints {
60    bytes: Vec<u8>,
61    num_hints: u32,
62}
63
64/// A single branch hint within a function.
65#[derive(Debug, Clone, Copy)]
66pub struct BranchHint {
67    /// The offset, in bytes from the beginning of the function, to the `if`
68    /// instruction that this is hinting.
69    pub branch_func_offset: u32,
70    /// The value of the hint, 0 for not taken and 1 for taken.
71    pub branch_hint_value: u32,
72}
73
74impl BranchHints {
75    /// Construct an empty encoder for the branch hints custom section.
76    pub fn new() -> Self {
77        Self::default()
78    }
79
80    /// Adds a new set of function hints for the `func` specified.
81    pub fn function_hints<I>(&mut self, func: u32, hints: I)
82    where
83        I: IntoIterator<Item = BranchHint>,
84        I::IntoIter: ExactSizeIterator,
85    {
86        self.num_hints += 1;
87        func.encode(&mut self.bytes);
88        let hints = hints.into_iter();
89        hints.len().encode(&mut self.bytes);
90        for hint in hints {
91            hint.branch_func_offset.encode(&mut self.bytes);
92            1u32.encode(&mut self.bytes);
93            hint.branch_hint_value.encode(&mut self.bytes);
94        }
95    }
96
97    /// Returns if this is an empty section.
98    pub fn is_empty(&self) -> bool {
99        self.num_hints == 0
100    }
101
102    /// Returns the number of functions that have hints registered in this
103    /// sectino.
104    pub fn len(&self) -> u32 {
105        self.num_hints
106    }
107}
108
109impl Encode for BranchHints {
110    fn encode(&self, sink: &mut Vec<u8>) {
111        let mut data = Vec::new();
112        self.num_hints.encode(&mut data);
113        data.extend(&self.bytes);
114
115        CustomSection {
116            name: "metadata.code.branch_hint".into(),
117            data: Cow::Borrowed(&data),
118        }
119        .encode(sink);
120    }
121}
122
123impl Section for BranchHints {
124    fn id(&self) -> u8 {
125        SectionId::Custom.into()
126    }
127}