feat: multithreaded wasm parser

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

Author: explodingcamera

Cargo.lock

@@ -16,6 +16,10 @@ log={workspace=true, optional=true}
 tinywasm-types={version="0.9.0-alpha.0", path="../types", default-features=false}
 
 [features]
-default=["std", "log"]
+default=["std", "log", "parallel"]
 log=["dep:log"]
 std=["tinywasm-types/std", "wasmparser/std"]
+parallel=["std"]
+
+[dev-dependencies]
+wat.workspace=true

crates/parser/Cargo.toml

@@ -0,0 +1,16 @@
+use std::env;
+
+fn main() {
+    println!("cargo::rustc-check-cfg=cfg(parallel_parser)");
+
+    if env::var("CARGO_FEATURE_PARALLEL").is_err() {
+        return;
+    }
+
+    let arch = env::var("CARGO_CFG_TARGET_ARCH").unwrap_or_default();
+    let os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default();
+
+    if !arch.starts_with("wasm") && os != "unknown" && os != "none" {
+        println!("cargo:rustc-cfg=parallel_parser");
+    }
+}

crates/parser/build.rs

@@ -1,6 +1,6 @@
 use alloc::sync::Arc;
 
-use crate::{Result, module::Code, visit::process_operators_and_validate};
+use crate::{Result, module::FunctionCode, visit::process_operators_and_validate};
 use alloc::{boxed::Box, format, string::ToString, vec::Vec};
 use tinywasm_types::*;
 use wasmparser::{FuncValidator, FuncValidatorAllocations, OperatorsReader, ValidatorResources};
@@ -163,7 +163,7 @@ pub(crate) fn convert_module_export(export: wasmparser::Export<'_>) -> Result<Ex
 pub(crate) fn convert_module_code(
     func: wasmparser::FunctionBody<'_>,
     mut validator: FuncValidator<ValidatorResources>,
-) -> Result<(Code, FuncValidatorAllocations)> {
+) -> Result<(FunctionCode, FuncValidatorAllocations)> {
     let locals_reader = func.get_locals_reader()?;
     let count = locals_reader.get_count();
     let pos = locals_reader.original_position();
@@ -200,7 +200,7 @@ pub(crate) fn convert_module_code(
     }
 
     let (body, data, allocations) = process_operators_and_validate(validator, func, local_addr_map)?;
-    Ok(((body, data, local_counts), allocations))
+    Ok((FunctionCode { instructions: body, data, locals: local_counts, uses_local_memory: false }, allocations))
 }
 
 pub(crate) fn convert_module_type(ty: wasmparser::RecGroup) -> Result<Arc<FuncType>> {

crates/parser/src/conversion.rs

@@ -35,6 +35,10 @@ mod macros;
 mod module;
 mod optimize;
 mod visit;
+
+#[cfg(parallel_parser)]
+mod parallel;
+
 pub use error::*;
 use module::ModuleReader;
 use wasmparser::{Validator, WasmFeatures};
@@ -51,11 +55,27 @@ pub struct ParserOptions {
     pub optimize_rewrite: bool,
     /// Whether to remove `Nop` and `MergeBarrier` instructions after rewriting.
     pub optimize_remove_nop: bool,
+
+    #[cfg(parallel_parser)]
+    /// Number of threads to use for parallel parsing.
+    ///
+    /// Requires the `parallel` feature. Ignored when the feature is disabled.
+    ///
+    /// - `None`: auto-detect based on available parallelism
+    /// - `Some(1)`: force single-threaded
+    /// - `Some(n)`: use up to `n` workers
+    pub parser_threads: Option<usize>,
 }
 
 impl Default for ParserOptions {
     fn default() -> Self {
-        Self { optimize_local_memory_allocation: true, optimize_rewrite: true, optimize_remove_nop: true }
+        Self {
+            optimize_local_memory_allocation: true,
+            optimize_rewrite: true,
+            optimize_remove_nop: true,
+            #[cfg(parallel_parser)]
+            parser_threads: None,
+        }
     }
 }
 
@@ -92,6 +112,21 @@ impl ParserOptions {
     pub const fn optimize_remove_nop(&self) -> bool {
         self.optimize_remove_nop
     }
+
+    #[cfg(parallel_parser)]
+    /// Set the number of threads for parallel parsing.
+    ///
+    /// Requires the `parallel` feature to have any effect.
+    pub const fn with_parser_threads(mut self, threads: usize) -> Self {
+        self.parser_threads = Some(threads);
+        self
+    }
+
+    #[cfg(parallel_parser)]
+    /// Returns the configured parser thread count, or `None` for auto-detect.
+    pub const fn parser_threads(&self) -> Option<usize> {
+        self.parser_threads
+    }
 }
 
 /// A WebAssembly parser
@@ -139,20 +174,40 @@ impl Parser {
         Validator::new_with_features(features)
     }
 
+    #[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));
+        let read_bytes = stream
+            .read(&mut buffer[len..])
+            .map_err(|e| ParseError::Other(alloc::format!("Error reading from stream: {e}")))?;
+        buffer.truncate(len + read_bytes);
+        Ok(read_bytes)
+    }
+
     /// Parse a [`Module`] from bytes
     pub fn parse_module_bytes(&self, wasm: impl AsRef<[u8]>) -> Result<Module> {
         let wasm = wasm.as_ref();
         let mut validator = Self::create_validator(self.options.clone());
         let mut reader = ModuleReader::default();
 
         for payload in wasmparser::Parser::new(0).parse_all(wasm) {
-            reader.process_payload(payload?, &mut validator)?;
+            match payload? {
+                wasmparser::Payload::CodeSectionStart { count, range, size } => {
+                    reader.begin_code_section(count, range, size, &mut validator, &self.options)?;
+                }
+                wasmparser::Payload::CodeSectionEntry(function) => {
+                    reader.process_borrowed_code_section_entry(function, &mut validator, &self.options)?;
+                }
+                payload => reader.process_payload(payload, &mut validator)?,
+            }
         }
 
         if !reader.end_reached {
             return Err(ParseError::EndNotReached);
         }
 
+        reader.process_pending_functions(&self.options)?;
         reader.into_module(&self.options)
     }
 
@@ -176,18 +231,60 @@ impl Parser {
         loop {
             match parser.parse(&buffer, eof)? {
                 wasmparser::Chunk::NeedMoreData(hint) => {
-                    let len = buffer.len();
-                    buffer.extend((0..hint).map(|_| 0u8));
-                    let read_bytes = stream
-                        .read(&mut buffer[len..])
-                        .map_err(|e| ParseError::Other(alloc::format!("Error reading from stream: {e}")))?;
-                    buffer.truncate(len + read_bytes);
+                    let read_bytes = Self::read_more(&mut stream, &mut buffer, hint as usize)?;
                     eof = read_bytes == 0;
                 }
                 wasmparser::Chunk::Parsed { consumed, payload } => {
-                    reader.process_payload(payload, &mut validator)?;
-                    buffer.drain(..consumed);
+                    #[cfg(parallel_parser)]
+                    let mut deferred_code_section = None;
+
+                    match payload {
+                        wasmparser::Payload::CodeSectionStart { count, range, size } => {
+                            let defer =
+                                reader.begin_code_section(count, range.clone(), size, &mut validator, &self.options)?;
+
+                            #[cfg(parallel_parser)]
+                            if defer {
+                                deferred_code_section = Some((count, range.end - size as usize, size as usize));
+                            }
+
+                            #[cfg(not(parallel_parser))]
+                            let _ = defer;
+
+                            buffer.drain(..consumed);
+                        }
+                        wasmparser::Payload::CodeSectionEntry(function) => {
+                            reader.process_inline_code_section_entry(function, &mut validator, &self.options)?;
+                            buffer.drain(..consumed);
+                        }
+                        payload => {
+                            reader.process_payload(payload, &mut validator)?;
+                            buffer.drain(..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();
+                            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(),
+                                });
+                            }
+                        }
+
+                        let section_bytes = alloc::sync::Arc::<[u8]>::from(buffer[..section_size].to_vec());
+                        reader.queue_owned_code_section(count, body_offset, section_bytes, &mut validator)?;
+                        parser.skip_section();
+                        buffer.drain(..section_size);
+                        continue;
+                    }
+
                     if eof || reader.end_reached {
+                        reader.process_pending_functions(&self.options)?;
                         return reader.into_module(&self.options);
                     }
                 }
@@ -196,10 +293,10 @@ impl Parser {
     }
 }
 
-impl TryFrom<ModuleReader> for Module {
+impl TryFrom<ModuleReader<'_>> for Module {
     type Error = ParseError;
 
-    fn try_from(reader: ModuleReader) -> Result<Self> {
+    fn try_from(reader: ModuleReader<'_>) -> Result<Self> {
         reader.into_module(&ParserOptions::default())
     }
 }