[wast] Branch hinting proposal implementation (#1394)

* wast: parsing and encoding branch hint annotations.

* Implement parsing/printing function branch hints

* Adding a test for branch hinting

* Touch up some branch hint parsing

* Remove unnecessary inclusion in `Custom` enum since this isn't parsed
  there.
* Refactor parsing slightly.
* Fix encoding to use function indices instead of function offsets.
* Refactor Rust idioms in encoding and remove some unnecessary
  abstractions.

---------

Co-authored-by: Alex Crichton <alex@alexcrichton.com>

Author: jpages

crates/wasmparser/src/readers/core.rs

@@ -1,3 +1,4 @@
+mod branch_hinting;
 mod code;
 mod coredumps;
 mod custom;
@@ -18,6 +19,7 @@ mod tables;
 mod tags;
 mod types;
 
+pub use self::branch_hinting::*;
 pub use self::code::*;
 pub use self::coredumps::*;
 pub use self::custom::*;

crates/wasmparser/src/readers/core/branch_hinting.rs

@@ -0,0 +1,59 @@
+use crate::{BinaryReader, FromReader, Result, SectionLimited};
+
+/// A reader for the `metadata.code.branch_hint` custom section.
+pub type BranchHintSectionReader<'a> = SectionLimited<'a, BranchHintFunction<'a>>;
+
+/// Branch hints for a single function.
+///
+/// Produced from [`BranchHintSectionReader`].
+#[derive(Debug, Clone)]
+pub struct BranchHintFunction<'a> {
+    /// The function that these branch hints apply to.
+    pub func: u32,
+    /// The branch hints available for this function.
+    pub hints: SectionLimited<'a, BranchHint>,
+}
+
+impl<'a> FromReader<'a> for BranchHintFunction<'a> {
+    fn from_reader(reader: &mut BinaryReader<'a>) -> Result<Self> {
+        let func = reader.read_var_u32()?;
+        // FIXME(#188) ideally wouldn't have to do skips here
+        let hints = reader.skip(|reader| {
+            let items_count = reader.read_var_u32()?;
+            for _ in 0..items_count {
+                reader.read::<BranchHint>()?;
+            }
+            Ok(())
+        })?;
+        Ok(BranchHintFunction {
+            func,
+            hints: SectionLimited::new(hints.remaining_buffer(), hints.original_position())?,
+        })
+    }
+}
+
+/// A hint for a single branch.
+#[derive(Debug, Copy, Clone)]
+pub struct BranchHint {
+    /// The byte offset, from the start of the function's body, of where the
+    /// hinted instruction lives.
+    pub func_offset: u32,
+    /// Whether or not the branch is hinted to be taken or not.
+    pub taken: bool,
+}
+
+impl<'a> FromReader<'a> for BranchHint {
+    fn from_reader(reader: &mut BinaryReader<'a>) -> Result<Self> {
+        let func_offset = reader.read_var_u32()?;
+        match reader.read_u8()? {
+            1 => {}
+            n => reader.invalid_leading_byte(n, "invalid branch hint byte")?,
+        }
+        let taken = match reader.read_u8()? {
+            0 => false,
+            1 => true,
+            n => reader.invalid_leading_byte(n, "invalid branch hint taken byte")?,
+        };
+        Ok(BranchHint { func_offset, taken })
+    }
+}

crates/wasmprinter/src/lib.rs

@@ -18,6 +18,7 @@ use wasmparser::*;
 const MAX_LOCALS: u32 = 50000;
 const MAX_NESTING_TO_PRINT: u32 = 50;
 const MAX_WASM_FUNCTIONS: u32 = 1_000_000;
+const MAX_WASM_FUNCTION_SIZE: u32 = 128 * 1024;
 
 mod operator;
 
@@ -52,6 +53,7 @@ pub struct Printer {
     nesting: u32,
     line: usize,
     group_lines: Vec<usize>,
+    code_section_hints: Vec<(u32, Vec<(usize, BranchHint)>)>,
 }
 
 #[derive(Default)]
@@ -242,6 +244,15 @@ impl Printer {
                     drop(self.register_component_names(state, reader));
                 }
 
+                Payload::CustomSection(c) if c.name() == "metadata.code.branch_hint" => {
+                    drop(
+                        self.register_branch_hint_section(BranchHintSectionReader::new(
+                            c.data(),
+                            c.data_offset(),
+                        )?),
+                    );
+                }
+
                 Payload::End(_) => break,
                 _ => {}
             }
@@ -1109,10 +1120,20 @@ impl Printer {
                 .print_core_functype_idx(state, ty, Some(func_idx))?
                 .unwrap_or(0);
 
+            // Hints are stored on `self` in reverse order of function index so
+            // check the last one and see if it matches this function.
+            let hints = match self.code_section_hints.last() {
+                Some((f, _)) if *f == func_idx => {
+                    let (_, hints) = self.code_section_hints.pop().unwrap();
+                    hints
+                }
+                _ => Vec::new(),
+            };
+
             if self.print_skeleton {
                 self.result.push_str(" ...");
             } else {
-                self.print_func_body(state, func_idx, params, &mut body)?;
+                self.print_func_body(state, func_idx, params, &mut body, &hints)?;
             }
 
             self.end_group();
@@ -1128,10 +1149,12 @@ impl Printer {
         func_idx: u32,
         params: u32,
         body: &mut BinaryReader<'_>,
+        mut branch_hints: &[(usize, BranchHint)],
     ) -> Result<()> {
         let mut first = true;
         let mut local_idx = 0;
         let mut locals = NamedLocalPrinter::new("local");
+        let func_start = body.original_position();
         for _ in 0..body.read_var_u32()? {
             let offset = body.original_position();
             let cnt = body.read_var_u32()?;
@@ -1163,7 +1186,21 @@ impl Printer {
         let mut buf = String::new();
         let mut op_printer = operator::PrintOperator::new(self, state);
         while !body.eof() {
-            // TODO
+            // Branch hints are stored in increasing order of their body offset
+            // so print them whenever their instruction comes up.
+            if let Some(((hint_offset, hint), rest)) = branch_hints.split_first() {
+                if hint.func_offset == (body.original_position() - func_start) as u32 {
+                    branch_hints = rest;
+                    op_printer.printer.newline(*hint_offset);
+                    let desc = if hint.taken { "\"\\01\"" } else { "\"\\00\"" };
+                    write!(
+                        op_printer.printer.result,
+                        "(@metadata.code.branch_hint {})",
+                        desc
+                    )?;
+                }
+            }
+
             let offset = body.original_position();
             mem::swap(&mut buf, &mut op_printer.printer.result);
             let op_kind = body.visit_operator(&mut op_printer)??;
@@ -2688,6 +2725,26 @@ impl Printer {
         }
         Ok(())
     }
+
+    fn register_branch_hint_section(&mut self, section: BranchHintSectionReader<'_>) -> Result<()> {
+        self.code_section_hints.clear();
+        for func in section {
+            let func = func?;
+            if self.code_section_hints.len() >= MAX_WASM_FUNCTIONS as usize {
+                bail!("found too many hints");
+            }
+            if func.hints.count() >= MAX_WASM_FUNCTION_SIZE {
+                bail!("found too many hints");
+            }
+            let hints = func
+                .hints
+                .into_iter_with_offsets()
+                .collect::<wasmparser::Result<Vec<_>>>()?;
+            self.code_section_hints.push((func.func, hints));
+        }
+        self.code_section_hints.reverse();
+        Ok(())
+    }
 }
 
 struct NamedLocalPrinter {

crates/wast/src/component/component.rs

@@ -111,6 +111,7 @@ impl<'a> Parse<'a> for Component<'a> {
         let _r = parser.register_annotation("custom");
         let _r = parser.register_annotation("producers");
         let _r = parser.register_annotation("name");
+        let _r = parser.register_annotation("metadata.code.branch_hint");
 
         let span = parser.parse::<kw::component>()?.0;
         let id = parser.parse()?;