wasmparser: improve visit_operator performance

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.