perf: improve parser performance

Signed-off-by: Henry <mail@henrygressmann.de>

3 files changed, 120 insertions, 111 deletions

diff --git a/crates/parser/src/conversion.rs b/crates/parser/src/conversion.rs
index cfc0e84..586cac5 100644
--- a/crates/parser/src/conversion.rs
+++ b/crates/parser/src/conversion.rs
@@ -180,8 +180,7 @@ pub(crate) fn convert_module_type(ty: wasmparser::RecGroup) -> Result<Arc<FuncTy
             ty.composite_type
         )));
     };
-    let params: Vec<_> = ty.params().iter().map(convert_valtype).collect();
-    let params = params.into_iter().collect::<Result<Vec<_>>>()?;
+    let params = ty.params().iter().map(convert_valtype).collect::<Result<Vec<_>>>()?;
     let results = ty.results().iter().map(convert_valtype).collect::<Result<Vec<_>>>()?;
     Ok(FuncType::new(&params, &results).into())
 }
@@ -209,17 +208,20 @@ pub(crate) fn convert_valtype(valtype: &wasmparser::ValType) -> Result<WasmType>
 }
 
 pub(crate) fn process_const_operators(ops: OperatorsReader<'_>) -> Result<Box<[ConstInstruction]>> {
-    let ops = ops.into_iter().collect::<wasmparser::Result<Vec<_>>>()?;
-    // In practice, the len can never be something other than 2,
-    // but we'll keep this here since it's part of the spec
-    // Invalid modules will be rejected by the validator anyway (there are also tests for this in the testsuite)
-    debug_assert!(ops.len() >= 2);
-    debug_assert!(matches!(ops[ops.len() - 1], wasmparser::Operator::End));
+    let mut out = Vec::new();
+    let mut operator_count = 0;
+    let mut end_reached = false;
+
+    for op in ops {
+        let op = op?;
+        operator_count += 1;
+        if matches!(op, wasmparser::Operator::End) {
+            end_reached = true;
+            break;
+        }
 
-    let mut out = Vec::with_capacity(ops.len().saturating_sub(1));
-    for op in ops.iter().take(ops.len() - 1) {
         let instr = match op {
-            wasmparser::Operator::RefNull { hty } => match convert_heaptype(*hty)? {
+            wasmparser::Operator::RefNull { hty } => match convert_heaptype(hty)? {
                 WasmType::RefFunc => ConstInstruction::RefFunc(None),
                 WasmType::RefExtern => ConstInstruction::RefExtern(None),
                 other => {
@@ -228,13 +230,13 @@ pub(crate) fn process_const_operators(ops: OperatorsReader<'_>) -> Result<Box<[C
                     )));
                 }
             },
-            wasmparser::Operator::RefFunc { function_index } => ConstInstruction::RefFunc(Some(*function_index)),
-            wasmparser::Operator::I32Const { value } => ConstInstruction::I32Const(*value),
-            wasmparser::Operator::I64Const { value } => ConstInstruction::I64Const(*value),
+            wasmparser::Operator::RefFunc { function_index } => ConstInstruction::RefFunc(Some(function_index)),
+            wasmparser::Operator::I32Const { value } => ConstInstruction::I32Const(value),
+            wasmparser::Operator::I64Const { value } => ConstInstruction::I64Const(value),
             wasmparser::Operator::F32Const { value } => ConstInstruction::F32Const(f32::from_bits(value.bits())),
             wasmparser::Operator::F64Const { value } => ConstInstruction::F64Const(f64::from_bits(value.bits())),
             wasmparser::Operator::V128Const { value } => ConstInstruction::V128Const(*value.bytes()),
-            wasmparser::Operator::GlobalGet { global_index } => ConstInstruction::GlobalGet(*global_index),
+            wasmparser::Operator::GlobalGet { global_index } => ConstInstruction::GlobalGet(global_index),
             wasmparser::Operator::I32Add => ConstInstruction::I32Add,
             wasmparser::Operator::I32Sub => ConstInstruction::I32Sub,
             wasmparser::Operator::I32Mul => ConstInstruction::I32Mul,
@@ -250,6 +252,10 @@ pub(crate) fn process_const_operators(ops: OperatorsReader<'_>) -> Result<Box<[C
         out.push(instr);
     }
 
+    if operator_count < 2 || !end_reached {
+        return Err(crate::ParseError::Other("constant expression did not end correctly".into()));
+    }
+
     Ok(out.into_boxed_slice())
 }
 
diff --git a/crates/parser/src/lib.rs b/crates/parser/src/lib.rs
index cc305a5..7936675 100644
--- a/crates/parser/src/lib.rs
+++ b/crates/parser/src/lib.rs
@@ -177,7 +177,7 @@ impl Parser {
     #[cfg(feature = "std")]
     fn read_more(stream: &mut impl std::io::Read, buffer: &mut alloc::vec::Vec<u8>, hint: usize) -> Result<usize> {
         let len = buffer.len();
-        buffer.extend((0..hint).map(|_| 0u8));
+        buffer.resize(len + hint, 0);
         let read_bytes = stream
             .read(&mut buffer[len..])
             .map_err(|e| ParseError::Other(alloc::format!("Error reading from stream: {e}")))?;
@@ -227,10 +227,16 @@ impl Parser {
         let mut buffer = alloc::vec::Vec::new();
         let mut parser = wasmparser::Parser::new(0);
         let mut eof = false;
+        let mut buffer_offset = 0;
 
         loop {
-            match parser.parse(&buffer, eof)? {
+            match parser.parse(&buffer[buffer_offset..], eof)? {
                 wasmparser::Chunk::NeedMoreData(hint) => {
+                    if buffer_offset != 0 {
+                        buffer.copy_within(buffer_offset.., 0);
+                        buffer.truncate(buffer.len() - buffer_offset);
+                        buffer_offset = 0;
+                    }
                     let read_bytes = Self::read_more(&mut stream, &mut buffer, hint as usize)?;
                     eof = read_bytes == 0;
                 }
@@ -263,30 +269,31 @@ impl Parser {
                             reader.process_payload(payload, validator.as_mut())?;
                         }
                     }
-                    buffer.drain(..consumed);
+                    buffer_offset += consumed;
 
                     #[cfg(parallel_parser)]
                     if let Some((count, body_offset, section_size)) = deferred_code_section {
-                        while buffer.len() < section_size {
-                            let remaining = section_size - buffer.len();
+                        while buffer.len() - buffer_offset < section_size {
+                            let remaining = section_size - (buffer.len() - buffer_offset);
                             let read_bytes = Self::read_more(&mut stream, &mut buffer, remaining)?;
                             if read_bytes == 0 {
                                 return Err(ParseError::ParseError {
                                     message: "unexpected end-of-file".into(),
-                                    offset: body_offset + buffer.len(),
+                                    offset: body_offset + buffer.len() - buffer_offset,
                                 });
                             }
                         }
 
-                        let section_bytes = alloc::sync::Arc::<[u8]>::from(buffer[..section_size].to_vec());
+                        let section_end = buffer_offset + section_size;
+                        let section_bytes = alloc::sync::Arc::<[u8]>::from(buffer[buffer_offset..section_end].to_vec());
                         reader.queue_owned_code_section(count, body_offset, section_bytes, validator.as_mut())?;
                         parser.skip_section();
-                        buffer.drain(..section_size);
+                        buffer_offset = section_end;
                         continue;
                     }
 
                     if reader.end_reached {
-                        if !buffer.is_empty() {
+                        if buffer_offset != buffer.len() {
                             return Err(ParseError::Other("trailing bytes after end of module".into()));
                         }
 
diff --git a/crates/parser/src/parallel.rs b/crates/parser/src/parallel.rs
index 137eb2f..88d1035 100644
--- a/crates/parser/src/parallel.rs
+++ b/crates/parser/src/parallel.rs
@@ -4,7 +4,7 @@ use alloc::sync::Arc;
 use alloc::vec::Vec;
 use core::ops::Range;
 use tinywasm_types::ValueCounts;
-use wasmparser::{FuncValidatorAllocations, ValidatorResources};
+use wasmparser::{FuncValidatorAllocations, OperatorsReaderAllocations, ValidatorResources};
 
 pub(crate) enum FunctionBodyInput<'a> {
     Borrowed(wasmparser::FunctionBody<'a>),
@@ -22,25 +22,29 @@ pub(crate) struct OwnedFunctionBody {
 pub(crate) struct PendingFunction<'a> {
     pub ordinal: usize,
     pub results: ValueCounts,
-    pub ty_idx: u32,
     pub func_to_validate: Option<wasmparser::FuncToValidate<ValidatorResources>>,
+    pub ty_idx: u32,
     pub body: FunctionBodyInput<'a>,
 }
 
-pub(crate) const MIN_FUNCTIONS: usize = 8;
-const MIN_CODE_SECTION_BYTES: usize = 32 * 1024;
-const MIN_FUNCTION_BODY_BYTES: usize = 4;
+const MIN_FUNCTIONS: usize = 8;
+const MIN_CODE_SECTION_BYTES: usize = 16 * 1024;
+const MIN_FUNCTION_BODY_BYTES: usize = 16;
+const MIN_FUNCTIONS_PER_WORKER: usize = 4;
+const MAX_WORKERS: usize = 12;
 
-pub(crate) fn should_parallelize_function(body_len: usize) -> bool {
-    body_len >= MIN_FUNCTION_BODY_BYTES
+fn body_len(body: &FunctionBodyInput<'_>) -> usize {
+    match body {
+        FunctionBodyInput::Borrowed(func) => func.as_bytes().len(),
+        FunctionBodyInput::Owned(body) => body.body_range.len(),
+    }
 }
 
 pub(crate) fn should_use_parallel(options: &ParserOptions, num_functions: usize, code_section_bytes: usize) -> bool {
-    if num_functions < MIN_FUNCTIONS || code_section_bytes < MIN_CODE_SECTION_BYTES {
-        return false;
-    }
-
-    worker_count(options, num_functions) > 1
+    num_functions >= MIN_FUNCTIONS
+        && code_section_bytes >= MIN_CODE_SECTION_BYTES
+        && code_section_bytes / num_functions >= MIN_FUNCTION_BODY_BYTES
+        && worker_count(options, num_functions) > 1
 }
 
 fn worker_count(options: &ParserOptions, num_functions: usize) -> usize {
@@ -48,15 +52,7 @@ fn worker_count(options: &ParserOptions, num_functions: usize) -> usize {
         .parser_threads()
         .unwrap_or_else(|| std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1))
         .max(1);
-
-    requested.min(num_functions).max(1)
-}
-
-fn body_len(body: &FunctionBodyInput<'_>) -> usize {
-    match body {
-        FunctionBodyInput::Borrowed(func) => func.as_bytes().len(),
-        FunctionBodyInput::Owned(body) => body.body_range.len(),
-    }
+    requested.min(MAX_WORKERS).min(num_functions.div_ceil(MIN_FUNCTIONS_PER_WORKER)).max(1)
 }
 
 fn process_function_job(
@@ -65,16 +61,18 @@ fn process_function_job(
     options: &ParserOptions,
     imported_func_count: usize,
     imported_memory_count: u32,
-) -> Result<(usize, FunctionCode)> {
-    let validator = job.func_to_validate.map(|func| func.into_validator(FuncValidatorAllocations::default()));
-    let (code, _, _) = match job.body {
+    validator_allocs: Option<FuncValidatorAllocations>,
+    reader_allocs: OperatorsReaderAllocations,
+) -> Result<(FunctionCode, Option<FuncValidatorAllocations>, OperatorsReaderAllocations)> {
+    let validator = job.func_to_validate.map(|func| func.into_validator(validator_allocs.unwrap_or_default()));
+    let (code, validator_allocs, reader_allocs) = match job.body {
         FunctionBodyInput::Borrowed(func) => {
-            conversion::convert_module_code(func, validator, Default::default(), metadata, job.ty_idx)?
+            conversion::convert_module_code(func, validator, reader_allocs, metadata, job.ty_idx)?
         }
         FunctionBodyInput::Owned(body) => {
             let reader = wasmparser::BinaryReader::new(&body.section_bytes[body.body_range], body.body_offset);
             let func = wasmparser::FunctionBody::new(reader);
-            conversion::convert_module_code(func, validator, Default::default(), metadata, job.ty_idx)?
+            conversion::convert_module_code(func, validator, reader_allocs, metadata, job.ty_idx)?
         }
     };
 
@@ -86,79 +84,77 @@ fn process_function_job(
         imported_memory_count,
     )?;
 
-    Ok((job.ordinal, code))
+    Ok((code, validator_allocs, reader_allocs))
 }
 
-pub(crate) fn process_pending(
-    pending: Vec<PendingFunction<'_>>,
+fn process_chunk<'a>(
+    jobs: impl IntoIterator<Item = PendingFunction<'a>>,
     metadata: &crate::visit::ModuleMetadata,
     options: &ParserOptions,
     imported_func_count: usize,
     imported_memory_count: u32,
 ) -> Result<Vec<FunctionCode>> {
-    let (small_jobs, large_jobs): (Vec<_>, Vec<_>) =
-        pending.into_iter().partition(|job| !should_parallelize_function(body_len(&job.body)));
+    let mut validator_allocs = None;
+    let mut reader_allocs = OperatorsReaderAllocations::default();
+    let jobs = jobs.into_iter();
+    let mut codes = Vec::with_capacity(jobs.size_hint().0);
 
-    let mut codes = small_jobs
-        .into_iter()
-        .map(|job| process_function_job(job, metadata, options, imported_func_count, imported_memory_count))
-        .collect::<Result<Vec<_>>>()?;
+    for job in jobs {
+        let (code, next_validator_allocs, next_reader_allocs) = process_function_job(
+            job,
+            metadata,
+            options,
+            imported_func_count,
+            imported_memory_count,
+            validator_allocs,
+            reader_allocs,
+        )?;
+        codes.push(code);
+        validator_allocs = next_validator_allocs;
+        reader_allocs = next_reader_allocs;
+    }
 
-    let num_workers = worker_count(options, large_jobs.len());
+    Ok(codes)
+}
+
+pub(crate) fn process_pending(
+    pending: Vec<PendingFunction<'_>>,
+    metadata: &crate::visit::ModuleMetadata,
+    options: &ParserOptions,
+    imported_func_count: usize,
+    imported_memory_count: u32,
+) -> Result<Vec<FunctionCode>> {
+    let num_workers = worker_count(options, pending.len());
     if num_workers == 1 {
-        codes.extend(
-            large_jobs
-                .into_iter()
-                .map(|job| process_function_job(job, metadata, options, imported_func_count, imported_memory_count))
-                .collect::<Result<Vec<_>>>()?,
-        );
-    } else {
-        let chunk_size = large_jobs.len().div_ceil(num_workers);
-        let chunks = {
-            let mut chunks = Vec::with_capacity(num_workers);
-            let mut iter = large_jobs.into_iter();
-            while let Some(first) = iter.next() {
-                let mut chunk = alloc::vec![first];
-                for _ in 1..chunk_size {
-                    match iter.next() {
-                        Some(job) => chunk.push(job),
-                        None => break,
-                    }
-                }
-                chunks.push(chunk);
+        return process_chunk(pending, metadata, options, imported_func_count, imported_memory_count);
+    }
+    let code_count = pending.len();
+    let chunk_size = pending.len().div_ceil(num_workers);
+    let chunk_bytes = pending.iter().map(|job| body_len(&job.body)).sum::<usize>().div_ceil(num_workers);
+    std::thread::scope(|scope| {
+        let mut jobs = pending.into_iter();
+        let mut handles = Vec::with_capacity(num_workers);
+        while let Some(first) = jobs.next() {
+            let mut chunk = Vec::with_capacity(chunk_size);
+            let mut bytes = body_len(&first.body);
+            chunk.push(first);
+            while bytes < chunk_bytes
+                && let Some(job) = jobs.next()
+            {
+                bytes += body_len(&job.body);
+                chunk.push(job);
             }
-            chunks
-        };
-
-        let results: Vec<Result<(usize, FunctionCode)>> = std::thread::scope(|s| {
-            let handles: Vec<_> = chunks
-                .into_iter()
-                .map(|chunk| {
-                    s.spawn(move || {
-                        chunk
-                            .into_iter()
-                            .map(|job| {
-                                process_function_job(job, metadata, options, imported_func_count, imported_memory_count)
-                            })
-                            .collect::<Vec<_>>()
-                    })
-                })
-                .collect();
-
             handles
-                .into_iter()
-                .flat_map(|handle| match handle.join() {
-                    Ok(results) => results,
-                    Err(_) => alloc::vec![Err(ParseError::Other("worker thread panicked".into()))],
-                })
-                .collect()
-        });
-
-        for result in results {
-            codes.push(result?);
+                .push(scope.spawn(move || {
+                    process_chunk(chunk, metadata, options, imported_func_count, imported_memory_count)
+                }));
         }
-    }
 
-    codes.sort_by_key(|(ordinal, _)| *ordinal);
-    Ok(codes.into_iter().map(|(_, code)| code).collect())
+        let mut codes = Vec::with_capacity(code_count);
+        for handle in handles {
+            let chunk = handle.join().map_err(|_| ParseError::Other("worker thread panicked".into()))??;
+            codes.extend(chunk);
+        }
+        Ok(codes)
+    })
 }