From cee820e5545c1fb9b423b915cf688414069cc960 Mon Sep 17 00:00:00 2001 From: Henry Date: Sat, 25 Apr 2026 23:17:41 +0200 Subject: chore: add specialized returns, optimize parser Signed-off-by: Henry --- crates/parser/src/macros.rs | 103 ++++++++++++++++- crates/parser/src/module.rs | 1 + crates/parser/src/optimize.rs | 64 +++++------ crates/parser/src/visit.rs | 259 ++++++++++++++++++++---------------------- 4 files changed, 250 insertions(+), 177 deletions(-) (limited to 'crates/parser') diff --git a/crates/parser/src/macros.rs b/crates/parser/src/macros.rs index ebab704..dded6d8 100644 --- a/crates/parser/src/macros.rs +++ b/crates/parser/src/macros.rs @@ -1,12 +1,103 @@ pub(crate) mod visit { macro_rules! validate_then_visit { - ($( @$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*))*) => {$( + ($( @$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*))*) => { + $(validate_then_visit!(@@$proposal $op $({ $($arg: $argty),* })? => $visit ($($ann)*));)* + }; + + // These special-case arms exist so we only clone wasmparser's non-Copy payloads + (@@mvp BrTable { $arg:ident: $argty:ty } => $visit:ident ($($ann:tt)*)) => { + fn $visit(&mut self, $arg: $argty) -> Self::Output { + self.0.$visit($arg.clone()); + let validation = self.0.validator.visitor(self.0.position).$visit($arg); + if let Err(e) = validation { + cold_path(); + self.0.record_error(crate::ParseError::ParseError { message: e.to_string(), offset: self.0.position }); + } + } + }; + + (@@reference_types TypedSelectMulti { $arg:ident: $argty:ty } => $visit:ident ($($ann:tt)*)) => { + fn $visit(&mut self, $arg: $argty) -> Self::Output { + self.0.$visit($arg.clone()); + let validation = self.0.validator.visitor(self.0.position).$visit($arg); + if let Err(e) = validation { + cold_path(); + self.0.record_error(crate::ParseError::ParseError { message: e.to_string(), offset: self.0.position }); + } + } + }; + + (@@exceptions TryTable { $arg:ident: $argty:ty } => $visit:ident ($($ann:tt)*)) => { + fn $visit(&mut self, $arg: $argty) -> Self::Output { + self.0.$visit($arg.clone()); + let validation = self.0.validator.visitor(self.0.position).$visit($arg); + if let Err(e) = validation { + cold_path(); + self.0.record_error(crate::ParseError::ParseError { message: e.to_string(), offset: self.0.position }); + } + } + }; + + (@@stack_switching Resume { cont_type_index: $cont:ty, resume_table: $table:ty } => $visit:ident ($($ann:tt)*)) => { + fn $visit(&mut self, cont_type_index: $cont, resume_table: $table) -> Self::Output { + self.0.$visit(cont_type_index, resume_table.clone()); + let validation = self.0.validator.visitor(self.0.position).$visit(cont_type_index, resume_table); + if let Err(e) = validation { + cold_path(); + self.0.record_error(crate::ParseError::ParseError { message: e.to_string(), offset: self.0.position }); + } + } + }; + + (@@stack_switching ResumeThrow { cont_type_index: $cont:ty, tag_index: $tag:ty, resume_table: $table:ty } => $visit:ident ($($ann:tt)*)) => { + fn $visit(&mut self, cont_type_index: $cont, tag_index: $tag, resume_table: $table) -> Self::Output { + self.0.$visit(cont_type_index, tag_index, resume_table.clone()); + let validation = self.0.validator.visitor(self.0.position).$visit(cont_type_index, tag_index, resume_table); + if let Err(e) = validation { + cold_path(); + self.0.record_error(crate::ParseError::ParseError { message: e.to_string(), offset: self.0.position }); + } + } + }; + + (@@stack_switching ResumeThrowRef { cont_type_index: $cont:ty, resume_table: $table:ty } => $visit:ident ($($ann:tt)*)) => { + fn $visit(&mut self, cont_type_index: $cont, resume_table: $table) -> Self::Output { + self.0.$visit(cont_type_index, resume_table.clone()); + let validation = self.0.validator.visitor(self.0.position).$visit(cont_type_index, resume_table); + if let Err(e) = validation { + cold_path(); + self.0.record_error(crate::ParseError::ParseError { message: e.to_string(), offset: self.0.position }); + } + } + }; + + (@@$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*)) => { fn $visit(&mut self $($(,$arg: $argty)*)?) -> Self::Output { - self.1.$visit($($($arg.clone()),*)?); - self.1.validator_visitor(self.0).$visit($($($arg),*)?)?; - Ok(()) + self.0.$visit($($($arg),*)?); + let validation = self.0.validator.visitor(self.0.position).$visit($($($arg),*)?); + if let Err(e) = validation { + cold_path(); + self.0.record_error(crate::ParseError::ParseError { message: e.to_string(), offset: self.0.position }); + } } - )*}; + }; + } + + macro_rules! validate_then_visit_simd { + ($( @$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*))*) => { + $(validate_then_visit_simd!(@@$proposal $op $({ $($arg: $argty),* })? => $visit ($($ann)*));)* + }; + + (@@$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*)) => { + fn $visit(&mut self $($(,$arg: $argty)*)?) -> Self::Output { + self.0.$visit($($($arg),*)?); + let validation = self.0.validator.simd_visitor(self.0.position).$visit($($($arg),*)?); + if let Err(e) = validation { + cold_path(); + self.0.record_error(crate::ParseError::ParseError { message: e.to_string(), offset: self.0.position }); + } + } + }; } macro_rules! define_operand { @@ -83,7 +174,7 @@ pub(crate) mod visit { pub(crate) use { define_mem_operands, define_mem_operands_simd, define_mem_operands_simd_lane, define_operand, define_operands, - impl_visit_operator, validate_then_visit, + impl_visit_operator, validate_then_visit, validate_then_visit_simd, }; } diff --git a/crates/parser/src/module.rs b/crates/parser/src/module.rs index 548280e..b8dd9e7 100644 --- a/crates/parser/src/module.rs +++ b/crates/parser/src/module.rs @@ -208,6 +208,7 @@ impl ModuleReader { instructions, &mut data, options, + results, self_func, import_mem_count, local_mem_alloc, diff --git a/crates/parser/src/optimize.rs b/crates/parser/src/optimize.rs index a5b793e..5882fea 100644 --- a/crates/parser/src/optimize.rs +++ b/crates/parser/src/optimize.rs @@ -1,7 +1,7 @@ use crate::ParserOptions; use crate::macros::optimize::*; use alloc::vec::Vec; -use tinywasm_types::{BinOp, BinOp128, CmpOp, ConstIdx, Instruction, WasmFunctionData}; +use tinywasm_types::{BinOp, BinOp128, CmpOp, ConstIdx, Instruction, ValueCounts, WasmFunctionData}; pub(crate) struct OptimizeResult { pub(crate) instructions: Vec, @@ -12,12 +12,13 @@ pub(crate) fn optimize_instructions( mut instructions: Vec, function_data: &mut WasmFunctionData, options: &ParserOptions, + function_results: ValueCounts, self_func_addr: u32, imported_memory_count: u32, track_local_memory_usage: bool, ) -> OptimizeResult { let uses_local_memory = if options.optimize_rewrite() { - rewrite(&mut instructions, self_func_addr, imported_memory_count, track_local_memory_usage) + rewrite(&mut instructions, function_results, self_func_addr, imported_memory_count, track_local_memory_usage) } else { track_local_memory_usage && instructions.iter().any(|instr| instr.memory_addr().is_some_and(|mem| mem >= imported_memory_count)) @@ -31,12 +32,20 @@ pub(crate) fn optimize_instructions( fn rewrite( instrs: &mut [Instruction], + function_results: ValueCounts, self_func_addr: u32, imported_memory_count: u32, track_local_memory_usage: bool, ) -> bool { use Instruction::*; let mut uses_local_memory = false; + let return_instr = match function_results { + ValueCounts { c32: 0, c64: 0, c128: 0 } => Some(ReturnVoid), + ValueCounts { c32: 1, c64: 0, c128: 0 } => Some(Return32), + ValueCounts { c32: 0, c64: 1, c128: 0 } => Some(Return64), + ValueCounts { c32: 0, c64: 0, c128: 1 } => Some(Return128), + _ => None, + }; for i in 0..instrs.len() { match instrs[i] { @@ -45,6 +54,7 @@ fn rewrite( LocalCopy128(a, b) if a == b => instrs[i] = Nop, Call(addr) if addr == self_func_addr => instrs[i] = CallSelf, ReturnCall(addr) if addr == self_func_addr => instrs[i] = ReturnCallSelf, + Return if let Some(return_instr) = return_instr => instrs[i] = return_instr, instr @ (I32Add | I32Mul | I32And | I32Or | I32Xor) => { let Some(op) = int_bin_op_32(instr) else { unreachable!() }; rewrite!(instrs, i, [LocalGet32(a), LocalGet32(b)] => BinOpLocalLocal32(op, a, b)); @@ -370,11 +380,7 @@ fn rewrite( ), Jump(ip) => { let target = resolve_jump_target(instrs, ip); - if target == next_non_nop(instrs, i + 1) as u32 { - instrs[i] = Nop; - } else if target != ip { - instrs[i] = Jump(target); - } + canonicalize_jump_like_with_target(instrs, i, target); } JumpIfZero(ip) => { let target = resolve_jump_target(instrs, ip); @@ -430,10 +436,7 @@ fn rewrite( (0, CmpOp::Ne) => JumpIfNonZero64(target), (imm, op) => JumpCmpStackConst64 { target_ip: target, imm, op }, }); - canonicalize_jump_like(instrs, i); - if let JumpIfZero(current) = &mut instrs[i] { - *current = target; - } + canonicalize_jump_like_with_target(instrs, i, target); } JumpIfNonZero(ip) => { let target = resolve_jump_target(instrs, ip); @@ -489,42 +492,27 @@ fn rewrite( (0, CmpOp::Ne) => JumpIfNonZero64(target), (imm, op) => JumpCmpStackConst64 { target_ip: target, imm, op }, }); - canonicalize_jump_like(instrs, i); - if let JumpIfNonZero(current) = &mut instrs[i] { - *current = target; - } + canonicalize_jump_like_with_target(instrs, i, target); } JumpIfZero32(ip) => { let target = resolve_jump_target(instrs, ip); rewrite!(instrs, i, [LocalGet32(local)] => JumpIfLocalZero32 { target_ip: target, local }); - canonicalize_jump_like(instrs, i); - if let JumpIfZero32(current) = &mut instrs[i] { - *current = target; - } + canonicalize_jump_like_with_target(instrs, i, target); } JumpIfNonZero32(ip) => { let target = resolve_jump_target(instrs, ip); rewrite!(instrs, i, [LocalGet32(local)] => JumpIfLocalNonZero32 { target_ip: target, local }); - canonicalize_jump_like(instrs, i); - if let JumpIfNonZero32(current) = &mut instrs[i] { - *current = target; - } + canonicalize_jump_like_with_target(instrs, i, target); } JumpIfZero64(ip) => { let target = resolve_jump_target(instrs, ip); rewrite!(instrs, i, [LocalGet64(local)] => JumpIfLocalZero64 { target_ip: target, local }); - canonicalize_jump_like(instrs, i); - if let JumpIfZero64(current) = &mut instrs[i] { - *current = target; - } + canonicalize_jump_like_with_target(instrs, i, target); } JumpIfNonZero64(ip) => { let target = resolve_jump_target(instrs, ip); rewrite!(instrs, i, [LocalGet64(local)] => JumpIfLocalNonZero64 { target_ip: target, local }); - canonicalize_jump_like(instrs, i); - if let JumpIfNonZero64(current) = &mut instrs[i] { - *current = target; - } + canonicalize_jump_like_with_target(instrs, i, target); } JumpCmpStackConst32 { target_ip, imm: 0, op } => { match op { @@ -573,8 +561,8 @@ fn rewrite( _ => {} } - if track_local_memory_usage { - uses_local_memory |= instrs[i].memory_addr().is_some_and(|mem| mem >= imported_memory_count); + if track_local_memory_usage && !uses_local_memory { + uses_local_memory = instrs[i].memory_addr().is_some_and(|mem| mem >= imported_memory_count); } } @@ -766,11 +754,14 @@ fn inverse_cmp_op(op: CmpOp) -> CmpOp { } } +const PREVIOUS_NON_NOP_BACKTRACK_LIMIT: usize = 32; + fn previous_non_nop(instrs: &[Instruction], read: usize) -> Option<[(usize, Instruction); N]> { let mut out = [(0usize, Instruction::Nop); N]; let mut filled = 0usize; + let start = read.saturating_sub(PREVIOUS_NON_NOP_BACKTRACK_LIMIT); - for idx in (0..read).rev() { + for idx in (start..read).rev() { let instr = instrs[idx]; if matches!(instr, Instruction::MergeBarrier) { return None; @@ -864,7 +855,10 @@ fn canonicalize_jump_like(instrs: &mut [Instruction], idx: usize) { return; }; - let target = resolve_jump_target(instrs, target); + canonicalize_jump_like_with_target(instrs, idx, resolve_jump_target(instrs, target)); +} + +fn canonicalize_jump_like_with_target(instrs: &mut [Instruction], idx: usize, target: u32) { if matches!(instrs[idx], Instruction::Jump(_)) && target == next_non_nop(instrs, idx + 1) as u32 { instrs[idx] = Instruction::Nop; } else { diff --git a/crates/parser/src/visit.rs b/crates/parser/src/visit.rs index 996f728..64821b4 100644 --- a/crates/parser/src/visit.rs +++ b/crates/parser/src/visit.rs @@ -1,6 +1,8 @@ +use core::hint::cold_path; + use crate::{Result, conversion::convert_heaptype, macros::visit::*}; use alloc::string::ToString; -use alloc::{vec, vec::Vec}; +use alloc::vec::Vec; use tinywasm_types::{Instruction, MemoryArg, WasmFunctionData}; use wasmparser::{ FrameKind, FuncValidator, FuncValidatorAllocations, FunctionBody, VisitOperator, VisitSimdOperator, @@ -14,6 +16,13 @@ enum BlockKind { If, } +#[derive(Debug, Clone, Copy)] +enum OperandSize { + S32, + S64, + S128, +} + #[derive(Debug, Clone, Copy, Default)] struct StackBase { s32: u16, @@ -43,10 +52,18 @@ impl FunctionDataBuilder { } } -struct ValidateThenVisit<'a, R: WasmModuleResources>(usize, &'a mut FunctionBuilder); +struct ValidateThenVisit<'a, R: WasmModuleResources>(&'a mut FunctionBuilder); + +fn operand_size(ty: wasmparser::ValType) -> OperandSize { + match ty { + wasmparser::ValType::I32 | wasmparser::ValType::F32 | wasmparser::ValType::Ref(_) => OperandSize::S32, + wasmparser::ValType::I64 | wasmparser::ValType::F64 => OperandSize::S64, + wasmparser::ValType::V128 => OperandSize::S128, + } +} impl<'a, R: WasmModuleResources> VisitOperator<'a> for ValidateThenVisit<'_, R> { - type Output = Result<()>; + type Output = (); wasmparser::for_each_visit_operator!(validate_then_visit); fn simd_visitor(&mut self) -> Option<&mut dyn VisitSimdOperator<'a, Output = Self::Output>> { @@ -55,7 +72,7 @@ impl<'a, R: WasmModuleResources> VisitOperator<'a> for ValidateThenVisit<'_, R> } impl VisitSimdOperator<'_> for ValidateThenVisit<'_, R> { - wasmparser::for_each_visit_simd_operator!(validate_then_visit); + wasmparser::for_each_visit_simd_operator!(validate_then_visit_simd); } pub(crate) fn process_operators_and_validate( @@ -64,28 +81,32 @@ pub(crate) fn process_operators_and_validate( local_addr_map: Vec, ) -> Result<(Vec, WasmFunctionData, FuncValidatorAllocations)> { let mut reader = body.get_operators_reader()?; - let remaining = reader.get_binary_reader().bytes_remaining(); - let mut builder = FunctionBuilder::new(remaining, validator, local_addr_map); + let mut builder = FunctionBuilder::new(validator, local_addr_map); while !reader.eof() { - reader.visit_operator(&mut ValidateThenVisit(reader.original_position(), &mut builder))??; + builder.position = reader.original_position(); + if let Err(e) = reader.visit_operator(&mut ValidateThenVisit(&mut builder)) { + cold_path(); + return Err(crate::ParseError::ParseError { message: e.to_string(), offset: builder.position }); + } } reader.finish()?; - if !builder.errors.is_empty() { - return Err(builder.errors.remove(0)); + if let Some(error) = builder.error { + return Err(error); } Ok((builder.instructions, builder.data.finish(), builder.validator.into_allocations())) } -pub(crate) struct FunctionBuilder { +pub(crate) struct FunctionBuilder { validator: FuncValidator, + position: usize, instructions: Vec, data: FunctionDataBuilder, ctx_stack: Vec, local_addr_map: Vec, - errors: Vec, + error: Option, } impl<'a, R: WasmModuleResources> wasmparser::VisitOperator<'a> for FunctionBuilder { @@ -157,26 +178,20 @@ impl<'a, R: WasmModuleResources> wasmparser::VisitOperator<'a> for FunctionBuild fn visit_global_set(&mut self, global_index: u32) -> Self::Output { if let Some(Some(t)) = self.validator.get_operand_type(0) { - self.instructions.push(match t { - wasmparser::ValType::I32 => Instruction::GlobalSet32(global_index), - wasmparser::ValType::F32 => Instruction::GlobalSet32(global_index), - wasmparser::ValType::I64 => Instruction::GlobalSet64(global_index), - wasmparser::ValType::F64 => Instruction::GlobalSet64(global_index), - wasmparser::ValType::V128 => Instruction::GlobalSet128(global_index), - wasmparser::ValType::Ref(_) => Instruction::GlobalSet32(global_index), + self.instructions.push(match operand_size(t) { + OperandSize::S32 => Instruction::GlobalSet32(global_index), + OperandSize::S64 => Instruction::GlobalSet64(global_index), + OperandSize::S128 => Instruction::GlobalSet128(global_index), }) } } fn visit_drop(&mut self) -> Self::Output { if let Some(Some(t)) = self.validator.get_operand_type(0) { - self.instructions.push(match t { - wasmparser::ValType::I32 => Instruction::Drop32, - wasmparser::ValType::F32 => Instruction::Drop32, - wasmparser::ValType::I64 => Instruction::Drop64, - wasmparser::ValType::F64 => Instruction::Drop64, - wasmparser::ValType::V128 => Instruction::Drop128, - wasmparser::ValType::Ref(_) => Instruction::Drop32, + self.instructions.push(match operand_size(t) { + OperandSize::S32 => Instruction::Drop32, + OperandSize::S64 => Instruction::Drop64, + OperandSize::S128 => Instruction::Drop128, }) } } @@ -190,27 +205,22 @@ impl<'a, R: WasmModuleResources> wasmparser::VisitOperator<'a> for FunctionBuild fn visit_local_get(&mut self, idx: u32) -> Self::Output { let resolved_idx = self.local_addr_map[idx as usize]; - use wasmparser::ValType::*; if let Some(t) = self.validator.get_local_type(idx) { - match t { - I32 | F32 => self.instructions.push(Instruction::LocalGet32(resolved_idx)), - I64 | F64 => self.instructions.push(Instruction::LocalGet64(resolved_idx)), - V128 => self.instructions.push(Instruction::LocalGet128(resolved_idx)), - Ref(_) => self.instructions.push(Instruction::LocalGet32(resolved_idx)), - } + self.instructions.push(match operand_size(t) { + OperandSize::S32 => Instruction::LocalGet32(resolved_idx), + OperandSize::S64 => Instruction::LocalGet64(resolved_idx), + OperandSize::S128 => Instruction::LocalGet128(resolved_idx), + }); } } fn visit_local_set(&mut self, idx: u32) -> Self::Output { let resolved_idx = self.local_addr_map[idx as usize]; if let Some(Some(t)) = self.validator.get_operand_type(0) { - self.instructions.push(match t { - wasmparser::ValType::I32 => Instruction::LocalSet32(resolved_idx), - wasmparser::ValType::F32 => Instruction::LocalSet32(resolved_idx), - wasmparser::ValType::I64 => Instruction::LocalSet64(resolved_idx), - wasmparser::ValType::F64 => Instruction::LocalSet64(resolved_idx), - wasmparser::ValType::V128 => Instruction::LocalSet128(resolved_idx), - wasmparser::ValType::Ref(_) => Instruction::LocalSet32(resolved_idx), + self.instructions.push(match operand_size(t) { + OperandSize::S32 => Instruction::LocalSet32(resolved_idx), + OperandSize::S64 => Instruction::LocalSet64(resolved_idx), + OperandSize::S128 => Instruction::LocalSet128(resolved_idx), }) } } @@ -218,58 +228,36 @@ impl<'a, R: WasmModuleResources> wasmparser::VisitOperator<'a> for FunctionBuild fn visit_local_tee(&mut self, idx: u32) -> Self::Output { let resolved_idx = self.local_addr_map[idx as usize]; if let Some(Some(t)) = self.validator.get_operand_type(0) { + let size = operand_size(t); let last = self.instructions.last(); - let src = match t { - wasmparser::ValType::I32 | wasmparser::ValType::F32 => { - if let Some(Instruction::LocalGet32(src)) = last { Some(*src) } else { None } - } - wasmparser::ValType::I64 | wasmparser::ValType::F64 => { - if let Some(Instruction::LocalGet64(src)) = last { Some(*src) } else { None } - } - wasmparser::ValType::V128 => { - if let Some(Instruction::LocalGet128(src)) = last { - Some(*src) - } else { - None - } - } - wasmparser::ValType::Ref(_) => { - if let Some(Instruction::LocalGet32(src)) = last { - Some(*src) - } else { - None - } - } + let src = match (size, last) { + (OperandSize::S32, Some(Instruction::LocalGet32(src))) => Some(*src), + (OperandSize::S64, Some(Instruction::LocalGet64(src))) => Some(*src), + (OperandSize::S128, Some(Instruction::LocalGet128(src))) => Some(*src), + _ => None, }; if let Some(src) = src { self.instructions.pop(); - match t { - wasmparser::ValType::I32 | wasmparser::ValType::F32 => { + match size { + OperandSize::S32 => { self.instructions.push(Instruction::LocalCopy32(src, resolved_idx)); self.instructions.push(Instruction::LocalGet32(resolved_idx)); } - wasmparser::ValType::I64 | wasmparser::ValType::F64 => { + OperandSize::S64 => { self.instructions.push(Instruction::LocalCopy64(src, resolved_idx)); self.instructions.push(Instruction::LocalGet64(resolved_idx)); } - wasmparser::ValType::V128 => { + OperandSize::S128 => { self.instructions.push(Instruction::LocalCopy128(src, resolved_idx)); self.instructions.push(Instruction::LocalGet128(resolved_idx)); } - wasmparser::ValType::Ref(_) => { - self.instructions.push(Instruction::LocalCopy32(src, resolved_idx)); - self.instructions.push(Instruction::LocalGet32(resolved_idx)); - } } } else { - self.instructions.push(match t { - wasmparser::ValType::I32 => Instruction::LocalTee32(resolved_idx), - wasmparser::ValType::F32 => Instruction::LocalTee32(resolved_idx), - wasmparser::ValType::I64 => Instruction::LocalTee64(resolved_idx), - wasmparser::ValType::F64 => Instruction::LocalTee64(resolved_idx), - wasmparser::ValType::V128 => Instruction::LocalTee128(resolved_idx), - wasmparser::ValType::Ref(_) => Instruction::LocalTee32(resolved_idx), + self.instructions.push(match size { + OperandSize::S32 => Instruction::LocalTee32(resolved_idx), + OperandSize::S64 => Instruction::LocalTee64(resolved_idx), + OperandSize::S128 => Instruction::LocalTee128(resolved_idx), }) } } @@ -368,31 +356,37 @@ impl<'a, R: WasmModuleResources> wasmparser::VisitOperator<'a> for FunctionBuild let branch_table_start = self.data.branch_table_targets.len() as u32; self.instructions.push(Instruction::BranchTable(0, branch_table_start, len)); - let mut seen = alloc::collections::BTreeMap::::new(); struct PadInfo { depth: u32, pad_start: usize, jump_or_ret_ip: usize, is_return: bool, } + let mut seen = Vec::<(u32, usize)>::new(); let mut pads: Vec = Vec::new(); for &depth in target_depths.iter().chain(core::iter::once(&default_depth)) { - if seen.contains_key(&depth) { + if seen.iter().any(|&(seen_depth, _)| seen_depth == depth) { continue; } - seen.insert(depth, pads.len()); + seen.push((depth, pads.len())); let (pad_start, jump_or_ret_ip, is_return) = self.emit_br_table_pad(depth); pads.push(PadInfo { depth, pad_start, jump_or_ret_ip, is_return }); } for &depth in &target_depths { - let pad_idx = seen[&depth]; + let pad_idx = seen + .iter() + .find_map(|&(seen_depth, idx)| (seen_depth == depth).then_some(idx)) + .expect("visit_br_table: missing branch table target"); self.data.branch_table_targets.push(pads[pad_idx].pad_start as u32); } - let default_pad_idx = seen[&default_depth]; + let default_pad_idx = seen + .iter() + .find_map(|&(seen_depth, idx)| (seen_depth == default_depth).then_some(idx)) + .expect("visit_br_table: missing default branch table target"); if let Instruction::BranchTable(default_ip, _, _) = &mut self.instructions[header_ip] { *default_ip = pads[default_pad_idx].pad_start as u32; } @@ -432,13 +426,10 @@ impl<'a, R: WasmModuleResources> wasmparser::VisitOperator<'a> for FunctionBuild } fn visit_typed_select(&mut self, ty: wasmparser::ValType) -> Self::Output { - self.instructions.push(match ty { - wasmparser::ValType::I32 => Instruction::Select32, - wasmparser::ValType::F32 => Instruction::Select32, - wasmparser::ValType::I64 => Instruction::Select64, - wasmparser::ValType::F64 => Instruction::Select64, - wasmparser::ValType::V128 => Instruction::Select128, - wasmparser::ValType::Ref(_) => Instruction::Select32, + self.instructions.push(match operand_size(ty) { + OperandSize::S32 => Instruction::Select32, + OperandSize::S64 => Instruction::Select64, + OperandSize::S128 => Instruction::Select128, }); } @@ -537,35 +528,35 @@ impl wasmparser::VisitSimdOperator<'_> for FunctionBuild } impl FunctionBuilder { - pub(crate) fn validator_visitor( - &mut self, - offset: usize, - ) -> impl VisitOperator<'_, Output = Result<(), wasmparser::BinaryReaderError>> + VisitSimdOperator<'_> { - self.validator.simd_visitor(offset) - } - - pub(crate) fn new(instr_capacity: usize, validator: FuncValidator, local_addr_map: Vec) -> Self { + pub(crate) fn new(validator: FuncValidator, local_addr_map: Vec) -> Self { Self { + position: 0, validator, local_addr_map, - instructions: Vec::with_capacity(instr_capacity), + instructions: Vec::with_capacity(1024), data: FunctionDataBuilder::default(), ctx_stack: Vec::with_capacity(256), - errors: Vec::new(), + error: None, + } + } + + fn record_error(&mut self, error: crate::ParseError) { + if self.error.is_none() { + self.error = Some(error); } } fn stack_base_at_frame(&self, depth: usize) -> StackBase { let Some(frame) = self.validator.get_control_frame(depth) else { return StackBase::default() }; let mut base = StackBase::default(); + let stack_height = self.validator.operand_stack_height() as usize; for i in 0..frame.height { - let depth_from_top = self.validator.operand_stack_height() as usize - 1 - i; + let depth_from_top = stack_height - 1 - i; if let Some(Some(ty)) = self.validator.get_operand_type(depth_from_top) { - match ty { - wasmparser::ValType::I32 | wasmparser::ValType::F32 => base.s32 += 1, - wasmparser::ValType::I64 | wasmparser::ValType::F64 => base.s64 += 1, - wasmparser::ValType::V128 => base.s128 += 1, - wasmparser::ValType::Ref(_) => base.s32 += 1, + match operand_size(ty) { + OperandSize::S32 => base.s32 += 1, + OperandSize::S64 => base.s64 += 1, + OperandSize::S128 => base.s128 += 1, } } } @@ -574,7 +565,7 @@ impl FunctionBuilder { } fn unsupported(&mut self, name: &str) { - self.errors.push(crate::ParseError::UnsupportedOperator(name.to_string())); + self.record_error(crate::ParseError::UnsupportedOperator(name.to_string())); } fn is_unreachable(&self) -> bool { @@ -632,18 +623,38 @@ impl FunctionBuilder { fn label_keep_counts(label_types: &[wasmparser::ValType]) -> (u16, u16, u16) { let (mut c32, mut c64, mut c128) = (0, 0, 0); - for ty in label_types { - match ty { - wasmparser::ValType::I32 | wasmparser::ValType::F32 => c32 += 1, - wasmparser::ValType::I64 | wasmparser::ValType::F64 => c64 += 1, - wasmparser::ValType::V128 => c128 += 1, - wasmparser::ValType::Ref(_) => c32 += 1, + for &ty in label_types { + match operand_size(ty) { + OperandSize::S32 => c32 += 1, + OperandSize::S64 => c64 += 1, + OperandSize::S128 => c128 += 1, } } (c32, c64, c128) } + fn label_keep_counts_for_frame(&self, frame: &wasmparser::Frame) -> (u16, u16, u16) { + match &frame.block_type { + wasmparser::BlockType::Empty => (0, 0, 0), + wasmparser::BlockType::Type(ty) => match frame.kind { + FrameKind::Loop => (0, 0, 0), + _ => Self::label_keep_counts(core::slice::from_ref(ty)), + }, + wasmparser::BlockType::FuncType(idx) => { + let sub_type = self.validator.resources().sub_type_at(*idx); + let func_ty = match sub_type { + Some(st) => st.composite_type.unwrap_func(), + None => return (0, 0, 0), + }; + match frame.kind { + FrameKind::Loop => Self::label_keep_counts(func_ty.params()), + _ => Self::label_keep_counts(func_ty.results()), + } + } + } + } + fn emit_dropkeep_to_label(&mut self, label_depth: u32) { if self.is_unreachable() { return; @@ -654,34 +665,11 @@ impl FunctionBuilder { }; let base = self.stack_base_at_frame(label_depth as usize); - let label_types: Vec<_> = self.label_types_for_frame(frame); - let (c32, c64, c128) = Self::label_keep_counts(&label_types); + let (c32, c64, c128) = self.label_keep_counts_for_frame(frame); self.emit_dropkeep(base, c32, c64, c128); } - fn label_types_for_frame(&self, frame: &wasmparser::Frame) -> Vec { - let ty = &frame.block_type; - match ty { - wasmparser::BlockType::Empty => Vec::new(), - wasmparser::BlockType::Type(ty) => match frame.kind { - FrameKind::Loop => Vec::new(), - _ => vec![*ty], - }, - wasmparser::BlockType::FuncType(idx) => { - let sub_type = self.validator.resources().sub_type_at(*idx); - let func_ty = match sub_type { - Some(st) => st.composite_type.unwrap_func(), - None => return Vec::new(), - }; - match frame.kind { - FrameKind::Loop => func_ty.params().to_vec(), - _ => func_ty.results().to_vec(), - } - } - } - } - fn emit_branch_jump_or_return(&mut self, depth: u32) { if let Some(ctx_idx) = self.get_ctx_idx(depth) { let jump_ip = self.instructions.len(); @@ -702,8 +690,7 @@ impl FunctionBuilder { }; let base = self.stack_base_at_frame(depth as usize); - let label_types: Vec<_> = self.label_types_for_frame(frame); - let (c32, c64, c128) = Self::label_keep_counts(&label_types); + let (c32, c64, c128) = self.label_keep_counts_for_frame(frame); self.emit_dropkeep(base, c32, c64, c128); let jump_ip = self.instructions.len(); -- cgit v1.3.1