feat: multithreaded wasm parser

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

1 files changed, 109 insertions, 12 deletions

diff --git a/crates/parser/src/lib.rs b/crates/parser/src/lib.rs
index 4eda84b..c4f6ff7 100644
--- a/crates/parser/src/lib.rs
+++ b/crates/parser/src/lib.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,6 +174,17 @@ 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();
@@ -146,13 +192,22 @@ impl Parser {
         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())
     }
 }