wasmparser: improve visit_operator performance (#2228)

* 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.

* Don't require handling OperatorsReaderAllocations

Make it a non-default API that other various locations can optionally
use but aren't required to do so.

* More test fixes

* Even more test fixes

* Simplify some adapter definitions

* Simplify definition of `OperatorFactory`

* Minor doc updates

* Fix doc link

---------

Co-authored-by: Alex Crichton <[email protected]>

Author: keithw

crates/wasmparser/benches/benchmark.rs

@@ -2,10 +2,14 @@ use anyhow::Result;
 use criterion::{Criterion, criterion_group, criterion_main};
 use once_cell::unsync::Lazy;
 use std::fs;
+use std::mem;
 use std::path::Path;
 use std::path::PathBuf;
 use wasmparser::VisitSimdOperator;
-use wasmparser::{DataKind, ElementKind, Parser, Payload, Validator, VisitOperator, WasmFeatures};
+use wasmparser::{
+    BinaryReader, DataKind, ElementKind, OperatorsReader, Parser, Payload, Validator,
+    VisitOperator, WasmFeatures,
+};
 
 /// A benchmark input.
 pub struct BenchmarkInput {
@@ -80,6 +84,17 @@ 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();
+    let mut read_expr = |reader: BinaryReader<'_>| -> Result<_> {
+        let mut ops = OperatorsReader::new_with_allocs(reader, mem::take(&mut allocs));
+
+        while !ops.eof() {
+            ops.visit_operator(&mut NopVisit)?;
+        }
+        ops.finish()?;
+        allocs = ops.into_allocations();
+        Ok(())
+    };
     for item in Parser::new(0).parse_all(wasm) {
         match item? {
             TypeSection(s) => {
@@ -114,9 +129,7 @@ fn read_all_wasm(wasm: &[u8]) -> Result<()> {
             }
             GlobalSection(s) => {
                 for item in s {
-                    for op in item?.init_expr.get_operators_reader() {
-                        op?;
-                    }
+                    read_expr(item?.init_expr.get_binary_reader())?;
                 }
             }
             ExportSection(s) => {
@@ -128,9 +141,7 @@ 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() {
-                            op?;
-                        }
+                        read_expr(offset_expr.get_binary_reader())?;
                     }
                     match item.items {
                         wasmparser::ElementItems::Functions(r) => {
@@ -150,9 +161,7 @@ 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() {
-                            op?;
-                        }
+                        read_expr(offset_expr.get_binary_reader())?;
                     }
                 }
             }
@@ -161,11 +170,7 @@ fn read_all_wasm(wasm: &[u8]) -> Result<()> {
                 for item in locals.by_ref() {
                     let _ = item?;
                 }
-                let mut ops = locals.into_operators_reader();
-                while !ops.eof() {
-                    ops.visit_operator(&mut NopVisit)?;
-                }
-                ops.finish()?;
+                read_expr(locals.into_binary_reader_for_operators())?;
             }
 
             // 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.