wasmparser: improve visit_operator performance

- The BinaryReader can visit_operator by itself again,
  with no need for OperatorsReader or extra state, when
  the visitor implements the FrameStack trait. Validation
  uses this to avoid duplicating the frame stack.

- OperatorsReader keeps an internal FrameStack and can be created
  from its allocations (and turned back into them), avoiding
  the need to allocate a stack of frames.

Author: keithw

crates/wasm-encoder/src/reencode.rs

@@ -1594,7 +1594,7 @@ pub mod utils {
         reencoder: &mut T,
         const_expr: wasmparser::ConstExpr,
     ) -> Result<crate::ConstExpr, Error<T::Error>> {
-        let mut ops = const_expr.get_operators_reader();
+        let mut ops = const_expr.get_operators_reader(Default::default());
         let mut bytes = Vec::new();
 
         while !ops.is_end_then_eof() {
@@ -1740,7 +1740,7 @@ pub mod utils {
         func: wasmparser::FunctionBody<'_>,
     ) -> Result<(), Error<T::Error>> {
         let mut f = reencoder.new_function_with_parsed_locals(&func)?;
-        let mut reader = func.get_operators_reader()?;
+        let mut reader = func.get_operators_reader(Default::default())?;
         while !reader.eof() {
             f.instruction(&reencoder.parse_instruction(&mut reader)?);
         }

crates/wasm-mutate/src/mutators/codemotion.rs

@@ -75,7 +75,7 @@ impl CodemotionMutator {
         for fidx in (function_to_mutate..function_count).chain(0..function_to_mutate) {
             config.consume_fuel(1)?;
             let reader = all_readers[fidx as usize].clone();
-            let operatorreader = reader.get_operators_reader()?;
+            let operatorreader = reader.get_operators_reader(Default::default())?;
 
             let operators = operatorreader
                 .into_iter_with_offsets()

crates/wasm-mutate/src/mutators/modify_const_exprs.rs

@@ -62,7 +62,7 @@ impl<'cfg, 'wasm> Reencode for InitTranslator<'cfg, 'wasm> {
             return reencode::utils::const_expr(self, e);
         }
 
-        let mut reader = e.get_operators_reader();
+        let mut reader = e.get_operators_reader(Default::default());
 
         if !self.config.reduce {
             // FIXME: implement non-reducing mutations for constant expressions.

crates/wasm-mutate/src/mutators/peephole.rs

@@ -122,7 +122,7 @@ impl PeepholeMutator {
             }
 
             let reader = readers[function_to_mutate as usize].clone();
-            let operatorreader = reader.get_operators_reader()?;
+            let operatorreader = reader.get_operators_reader(Default::default())?;
             let mut localsreader = reader.get_locals_reader()?;
             let operators = operatorreader
                 .into_iter_with_offsets()
@@ -490,7 +490,7 @@ macro_rules! match_code_mutation {
                     });
                 }
                 Payload::CodeSectionEntry(reader) => {
-                    let operatorsreader = reader.get_operators_reader().unwrap();
+                    let operatorsreader = reader.get_operators_reader(Default::default()).unwrap();
                     let range = operatorsreader.get_binary_reader().range();
                     let operators = operatorsreader
                         .into_iter_with_offsets()

crates/wasm-mutate/src/mutators/peephole/dfg.rs

@@ -1314,7 +1314,7 @@ mod tests {
             match payload {
                 wasmparser::Payload::CodeSectionEntry(reader) => {
                     let operators = reader
-                        .get_operators_reader()
+                        .get_operators_reader(Default::default())
                         .unwrap()
                         .into_iter_with_offsets()
                         .collect::<wasmparser::Result<Vec<OperatorAndByteOffset>>>()
@@ -1370,7 +1370,7 @@ mod tests {
             match payload {
                 wasmparser::Payload::CodeSectionEntry(reader) => {
                     let operators = reader
-                        .get_operators_reader()
+                        .get_operators_reader(Default::default())
                         .unwrap()
                         .into_iter_with_offsets()
                         .collect::<wasmparser::Result<Vec<OperatorAndByteOffset>>>()
@@ -1452,7 +1452,7 @@ mod tests {
             match payload {
                 wasmparser::Payload::CodeSectionEntry(reader) => {
                     let operators = reader
-                        .get_operators_reader()
+                        .get_operators_reader(Default::default())
                         .unwrap()
                         .into_iter_with_offsets()
                         .collect::<wasmparser::Result<Vec<OperatorAndByteOffset>>>()
@@ -1509,7 +1509,7 @@ mod tests {
             match payload {
                 wasmparser::Payload::CodeSectionEntry(reader) => {
                     let operators = reader
-                        .get_operators_reader()
+                        .get_operators_reader(Default::default())
                         .unwrap()
                         .into_iter_with_offsets()
                         .collect::<wasmparser::Result<Vec<OperatorAndByteOffset>>>()

crates/wasmparser/benches/benchmark.rs

@@ -80,6 +80,7 @@ fn collect_test_files(path: &Path, list: &mut Vec<BenchmarkInput>) -> Result<()>
 /// so that we can report better errors in case of failures.
 fn read_all_wasm(wasm: &[u8]) -> Result<()> {
     use Payload::*;
+    let mut allocs = wasmparser::OperatorsReaderAllocations::default();
     for item in Parser::new(0).parse_all(wasm) {
         match item? {
             TypeSection(s) => {
@@ -114,9 +115,14 @@ fn read_all_wasm(wasm: &[u8]) -> Result<()> {
             }
             GlobalSection(s) => {
                 for item in s {
-                    for op in item?.init_expr.get_operators_reader() {
+                    let mut ops = item?
+                        .init_expr
+                        .get_operators_reader(std::mem::take(&mut allocs))
+                        .into_iter();
+                    for op in ops.by_ref() {
                         op?;
                     }
+                    allocs = ops.into_allocations();
                 }
             }
             ExportSection(s) => {
@@ -128,9 +134,13 @@ fn read_all_wasm(wasm: &[u8]) -> Result<()> {
                 for item in s {
                     let item = item?;
                     if let ElementKind::Active { offset_expr, .. } = item.kind {
-                        for op in offset_expr.get_operators_reader() {
+                        let mut ops = offset_expr
+                            .get_operators_reader(std::mem::take(&mut allocs))
+                            .into_iter();
+                        for op in ops.by_ref() {
                             op?;
                         }
+                        allocs = ops.into_allocations();
                     }
                     match item.items {
                         wasmparser::ElementItems::Functions(r) => {
@@ -150,9 +160,13 @@ fn read_all_wasm(wasm: &[u8]) -> Result<()> {
                 for item in s {
                     let item = item?;
                     if let DataKind::Active { offset_expr, .. } = item.kind {
-                        for op in offset_expr.get_operators_reader() {
+                        let mut ops = offset_expr
+                            .get_operators_reader(std::mem::take(&mut allocs))
+                            .into_iter();
+                        for op in ops.by_ref() {
                             op?;
                         }
+                        allocs = ops.into_allocations();
                     }
                 }
             }
@@ -161,11 +175,12 @@ fn read_all_wasm(wasm: &[u8]) -> Result<()> {
                 for item in locals.by_ref() {
                     let _ = item?;
                 }
-                let mut ops = locals.into_operators_reader();
+                let mut ops = locals.into_operators_reader(std::mem::take(&mut allocs));
                 while !ops.eof() {
                     ops.visit_operator(&mut NopVisit)?;
                 }
                 ops.finish()?;
+                allocs = ops.into_allocations();
             }
 
             // Component sections

crates/wasmparser/src/arity.rs

@@ -15,8 +15,8 @@
 
 use crate::prelude::*;
 use crate::{
-    BinaryReaderError, BlockType, BrTable, CompositeInnerType, ContType, FrameKind, FuncType,
-    Operator, OperatorsReader, RefType, Result, ResumeTable, SubType, TryTable, ValType,
+    BlockType, BrTable, CompositeInnerType, ContType, FrameKind, FuncType, Operator, RefType,
+    ResumeTable, SubType, TryTable, ValType,
 };
 
 /// To compute the arity (param and result counts) of "variable-arity"
@@ -70,15 +70,6 @@ pub trait ModuleArity {
     }
 }
 
-impl OperatorsReader<'_> {
-    /// Read the next operator and compute its arity (param and result counts)
-    pub fn operator_arity(&self, module: &impl ModuleArity) -> Result<(u32, u32)> {
-        self.clone().read()?.operator_arity(module).ok_or_else(|| {
-            BinaryReaderError::new("operator arity is unknown", self.original_position())
-        })
-    }
-}
-
 impl Operator<'_> {
     /// Compute the arity (param and result counts) of the operator, given
     /// an impl ModuleArity, which stores the necessary module state.