diff options
Diffstat (limited to 'crates/parser/src/module.rs')
| -rw-r--r-- | crates/parser/src/module.rs | 258 |
1 files changed, 217 insertions, 41 deletions
diff --git a/crates/parser/src/module.rs b/crates/parser/src/module.rs index b8dd9e7..02a3297 100644 --- a/crates/parser/src/module.rs +++ b/crates/parser/src/module.rs @@ -2,21 +2,59 @@ use crate::log::debug; use crate::{ParseError, ParserOptions, Result, conversion, optimize}; use alloc::sync::Arc; use alloc::{format, string::ToString, vec::Vec}; +use core::marker::PhantomData; +use core::ops::Range; use tinywasm_types::*; use wasmparser::{FuncValidatorAllocations, Payload, Validator}; -pub(crate) type Code = (Vec<Instruction>, WasmFunctionData, ValueCounts); +pub(crate) struct FunctionCode { + pub instructions: Vec<Instruction>, + pub data: WasmFunctionData, + pub locals: ValueCounts, + pub uses_local_memory: bool, +} + +pub(crate) fn imported_func_count(imports: &[Import]) -> usize { + imports.iter().filter(|i| matches!(&i.kind, ImportKind::Function(_))).count() +} + +pub(crate) fn imported_memory_count(imports: &[Import]) -> u32 { + imports.iter().filter(|i| matches!(&i.kind, ImportKind::Memory(_))).count() as u32 +} + +pub(crate) fn optimize_function_code( + mut code: FunctionCode, + options: &ParserOptions, + function_results: ValueCounts, + self_func_addr: u32, + imported_memory_count: u32, +) -> FunctionCode { + let optimized = optimize::optimize_instructions( + code.instructions, + &mut code.data, + options, + function_results, + self_func_addr, + imported_memory_count, + ); + + code.instructions = optimized.instructions; + code.uses_local_memory = optimized.uses_local_memory; + code +} #[derive(Default)] -pub(crate) struct ModuleReader { +pub(crate) struct ModuleReader<'a> { func_validator_allocations: Option<FuncValidatorAllocations>, + has_code_section: bool, + marker: PhantomData<&'a [u8]>, pub(crate) version: Option<u16>, pub(crate) start_func: Option<u32>, pub(crate) func_types: Vec<Arc<FuncType>>, pub(crate) code_type_addrs: Vec<u32>, pub(crate) exports: Vec<Export>, - pub(crate) code: Vec<Code>, + pub(crate) code: Vec<FunctionCode>, pub(crate) globals: Vec<Global>, pub(crate) table_types: Vec<TableType>, pub(crate) memory_types: Vec<MemoryType>, @@ -24,9 +62,18 @@ pub(crate) struct ModuleReader { pub(crate) data: Vec<Data>, pub(crate) elements: Vec<Element>, pub(crate) end_reached: bool, + + #[cfg(parallel_parser)] + pending_functions: Option<Vec<crate::parallel::PendingFunction<'a>>>, } -impl ModuleReader { +impl<'a> ModuleReader<'a> { + fn function_results(&self, ordinal: usize) -> ValueCounts { + let ty_idx = self.code_type_addrs[ordinal]; + let ty = self.func_types.get(ty_idx as usize).expect("No func type for func, this is a bug"); + ValueCounts::from_iter(ty.results()) + } + pub(crate) fn process_payload(&mut self, payload: Payload<'_>, validator: &mut Validator) -> Result<()> { match payload { Payload::Version { num, encoding, range } => { @@ -55,7 +102,6 @@ impl ModuleReader { self.func_types = reader.into_iter().map(|t| conversion::convert_module_type(t?)).collect::<Result<Vec<_>>>()?; } - Payload::GlobalSection(reader) => { if !self.globals.is_empty() { return Err(ParseError::DuplicateSection("Global section".into())); @@ -69,6 +115,7 @@ impl ModuleReader { if !self.table_types.is_empty() { return Err(ParseError::DuplicateSection("Table section".into())); } + debug!("Found table section"); validator.table_section(&reader)?; self.table_types = conversion::convert_module_tables(reader)?; @@ -112,22 +159,6 @@ impl ModuleReader { validator.function_section(&reader)?; self.code_type_addrs = reader.into_iter().map(|f| Ok(f?)).collect::<Result<Vec<_>>>()?; } - Payload::CodeSectionStart { count, range, .. } => { - debug!("Found code section ({count} functions)"); - if !self.code.is_empty() { - return Err(ParseError::DuplicateSection("Code section".into())); - } - self.code.reserve(count as usize); - validator.code_section_start(&range)?; - } - Payload::CodeSectionEntry(function) => { - debug!("Found code section entry"); - let v = validator.code_section_entry(&function)?; - let func_validator = v.into_validator(self.func_validator_allocations.take().unwrap_or_default()); - let (code, allocations) = conversion::convert_module_code(function, func_validator)?; - self.code.push(code); - self.func_validator_allocations = Some(allocations); - } Payload::ImportSection(reader) => { if !self.imports.is_empty() { return Err(ParseError::DuplicateSection("Import section".into())); @@ -160,9 +191,161 @@ impl ModuleReader { debug!("Found custom section"); debug!("Skipping custom section: {:?}", _reader.name()); } + Payload::CodeSectionStart { .. } | Payload::CodeSectionEntry(_) => { + return Err(ParseError::Other("code section payload handled separately".into())); + } Payload::UnknownSection { .. } => return Err(ParseError::UnsupportedSection("Unknown section".into())), section => return Err(ParseError::UnsupportedSection(format!("Unsupported section: {section:?}"))), + } + + Ok(()) + } + + pub(crate) fn begin_code_section( + &mut self, + count: u32, + range: Range<usize>, + size: u32, + validator: &mut Validator, + options: &ParserOptions, + ) -> Result<bool> { + debug!("Found code section ({count} functions)"); + if self.has_code_section { + return Err(ParseError::DuplicateSection("Code section".into())); + } + + self.has_code_section = true; + self.code.reserve(count as usize); + validator.code_section_start(&range)?; + + #[cfg(parallel_parser)] + { + let defer = crate::parallel::should_use_parallel(options, count as usize, size as usize); + if defer { + debug!("Queuing {count} functions from {size} byte code section"); + self.pending_functions = Some(Vec::with_capacity(count as usize)); + } + Ok(defer) + } + + #[cfg(not(parallel_parser))] + { + let _ = (size, options); + Ok(false) + } + } + + pub(crate) fn process_inline_code_section_entry( + &mut self, + function: wasmparser::FunctionBody<'_>, + validator: &mut Validator, + options: &ParserOptions, + ) -> Result<()> { + debug!("Found code section entry"); + let ordinal = self.code.len(); + let func_to_validate = validator.code_section_entry(&function)?; + let func_validator = + func_to_validate.into_validator(self.func_validator_allocations.take().unwrap_or_default()); + let (code, allocations) = conversion::convert_module_code(function, func_validator)?; + let code = optimize_function_code( + code, + options, + self.function_results(ordinal), + (imported_func_count(&self.imports) + ordinal) as u32, + imported_memory_count(&self.imports), + ); + self.code.push(code); + self.func_validator_allocations = Some(allocations); + Ok(()) + } + + pub(crate) fn process_borrowed_code_section_entry( + &mut self, + function: wasmparser::FunctionBody<'a>, + validator: &mut Validator, + options: &ParserOptions, + ) -> Result<()> { + debug!("Found code section entry"); + + #[cfg(parallel_parser)] + if let Some(pending) = self.pending_functions.as_mut() { + let func_to_validate = validator.code_section_entry(&function)?; + let ordinal = self.code.len() + pending.len(); + let ty_idx = self.code_type_addrs[ordinal]; + pending.push(crate::parallel::PendingFunction { + ordinal, + ty_idx, + func_to_validate, + body: crate::parallel::FunctionBodyInput::Borrowed(function), + }); + return Ok(()); + } + + self.process_inline_code_section_entry(function, validator, options) + } + + #[cfg(parallel_parser)] + pub(crate) fn queue_owned_code_section( + &mut self, + count: u32, + body_offset: usize, + section_bytes: Arc<[u8]>, + validator: &mut Validator, + ) -> Result<()> { + let code_len = self.code.len(); + let pending = self + .pending_functions + .as_mut() + .ok_or_else(|| ParseError::Other("owned code section queued without pending storage".into()))?; + + let mut reader = wasmparser::BinaryReader::new(§ion_bytes, body_offset); + for _ in 0..count { + let body_reader = reader.read_reader()?; + let body_range = body_reader.range(); + let function = wasmparser::FunctionBody::new(body_reader); + let func_to_validate = validator.code_section_entry(&function)?; + let ordinal = code_len + pending.len(); + let ty_idx = self.code_type_addrs[ordinal]; + pending.push(crate::parallel::PendingFunction { + ordinal, + ty_idx, + func_to_validate, + body: crate::parallel::FunctionBodyInput::Owned(crate::parallel::OwnedFunctionBody { + section_bytes: section_bytes.clone(), + body_range: (body_range.start - body_offset)..(body_range.end - body_offset), + body_offset: body_range.start, + }), + }); + } + + if reader.bytes_remaining() != 0 { + return Err(ParseError::ParseError { + message: "trailing bytes at end of section".into(), + offset: reader.original_position(), + }); + } + + Ok(()) + } + + #[cfg(parallel_parser)] + pub(crate) fn process_pending_functions(&mut self, options: &ParserOptions) -> Result<()> { + let Some(pending) = self.pending_functions.take().filter(|pending| !pending.is_empty()) else { + return Ok(()); }; + + self.code.extend(crate::parallel::process_pending( + pending, + options, + &self.func_types, + imported_func_count(&self.imports), + imported_memory_count(&self.imports), + )?); + Ok(()) + } + + #[cfg(not(parallel_parser))] + pub(crate) fn process_pending_functions(&mut self, _options: &ParserOptions) -> Result<()> { Ok(()) } @@ -175,8 +358,7 @@ impl ModuleReader { return Err(ParseError::Other("Code and code type address count mismatch".to_string())); } - let imported_func_count = self.imports.iter().filter(|i| matches!(&i.kind, ImportKind::Function(_))).count(); - let import_mem_count = self.imports.iter().filter(|i| matches!(&i.kind, ImportKind::Memory(_))).count() as u32; + let import_mem_count = imported_memory_count(&self.imports); let has_local_mem_export = self.exports.iter().any(|export| export.kind == ExternalKind::Memory && export.index >= import_mem_count); let has_active_data_segment_on_local_memory = self.data.iter().any(|data| match &data.kind { @@ -193,33 +375,27 @@ impl ModuleReader { } else { LocalMemoryAllocation::Skip }; + let mut funcs = Vec::with_capacity(self.code.len()); - for (func_idx, ((instructions, mut data, locals), ty_idx)) in - self.code.into_iter().zip(self.code_type_addrs).enumerate() - { + for (code, ty_idx) in self.code.into_iter().zip(self.code_type_addrs) { let ty = self.func_types.get(ty_idx as usize).expect("No func type for func, this is a bug").clone(); let params = ValueCounts::from_iter(ty.params()); let results = ValueCounts::from_iter(ty.results()); - let self_func = (imported_func_count + func_idx) as u32; - let local_mem_alloc = - optimize_local_memory_allocation && local_memory_allocation != LocalMemoryAllocation::Eager; - let optimized = optimize::optimize_instructions( - instructions, - &mut data, - options, - results, - self_func, - import_mem_count, - local_mem_alloc, - ); - - if optimized.uses_local_memory { + if code.uses_local_memory { local_memory_allocation = LocalMemoryAllocation::Eager; } funcs.push( - WasmFunction { instructions: optimized.instructions.into(), data, locals, params, results, ty }.into(), + WasmFunction { + instructions: code.instructions.into(), + data: code.data, + locals: code.locals, + params, + results, + ty, + } + .into(), ); } |