diff options
Diffstat (limited to 'crates/parser')
| -rw-r--r-- | crates/parser/src/visit.rs | 425 |
1 files changed, 339 insertions, 86 deletions
diff --git a/crates/parser/src/visit.rs b/crates/parser/src/visit.rs index cb9a3d2..9a707af 100644 --- a/crates/parser/src/visit.rs +++ b/crates/parser/src/visit.rs @@ -1,13 +1,37 @@ use crate::Result; -use crate::conversion::{convert_heaptype, convert_valtype}; +use crate::conversion::convert_heaptype; use alloc::string::ToString; +use alloc::vec; use alloc::vec::Vec; use tinywasm_types::{Instruction, MemoryArg, WasmFunctionData}; use wasmparser::{ - FuncValidator, FuncValidatorAllocations, FunctionBody, VisitOperator, VisitSimdOperator, WasmModuleResources, + FrameKind, FuncValidator, FuncValidatorAllocations, FunctionBody, VisitOperator, VisitSimdOperator, + WasmModuleResources, }; +#[derive(Debug, Clone, Copy)] +enum BlockKind { + Block, + Loop, + If, +} + +#[derive(Debug, Clone, Copy, Default)] +struct StackBase { + s32: u16, + s64: u16, + s128: u16, + sref: u16, +} + +struct LoweringCtx { + kind: BlockKind, + has_else: bool, + start_ip: usize, + branch_jumps: Vec<usize>, +} + struct ValidateThenVisit<'a, R: WasmModuleResources>(usize, &'a mut FunctionBuilder<R>); macro_rules! validate_then_visit { @@ -112,7 +136,7 @@ pub(crate) struct FunctionBuilder<R: WasmModuleResources> { validator: FuncValidator<R>, instructions: Vec<Instruction>, v128_constants: Vec<i128>, - label_ptrs: Vec<usize>, + ctx_stack: Vec<LoweringCtx>, local_addr_map: Vec<u32>, errors: Vec<crate::ParseError>, } @@ -124,23 +148,157 @@ impl<R: WasmModuleResources> FunctionBuilder<R> { ) -> impl VisitOperator<'_, Output = Result<(), wasmparser::BinaryReaderError>> + VisitSimdOperator<'_> { self.validator.simd_visitor(offset) } -} -impl<R: WasmModuleResources> FunctionBuilder<R> { pub(crate) fn new(instr_capacity: usize, validator: FuncValidator<R>, local_addr_map: Vec<u32>) -> Self { Self { validator, local_addr_map, instructions: Vec::with_capacity(instr_capacity), v128_constants: Vec::new(), - label_ptrs: Vec::with_capacity(256), + ctx_stack: Vec::with_capacity(256), errors: Vec::new(), } } + fn stack_base_at_frame(&self, depth: usize) -> StackBase { + let frame = match self.validator.get_control_frame(depth) { + Some(f) => f, + None => return StackBase::default(), + }; + let height = frame.height; + let current = self.validator.operand_stack_height() as usize; + + let mut base = StackBase::default(); + for i in 0..height { + let depth_from_top = current - 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.sref += 1, + } + } + } + + base + } + fn unsupported(&mut self, name: &str) { self.errors.push(crate::ParseError::UnsupportedOperator(name.to_string())); } + + fn current_ip(&self) -> u32 { + self.instructions.len() as u32 + } + + fn is_unreachable(&self) -> bool { + self.validator.get_control_frame(0).is_none_or(|f| f.unreachable) + } + + fn get_ctx_idx(&self, depth: u32) -> Option<usize> { + let len = self.ctx_stack.len(); + let idx = len.checked_sub(depth as usize + 1)?; + Some(idx) + } + + fn emit_dropkeep(&mut self, base: StackBase, c32: u16, c64: u16, c128: u16, cref: u16) { + let fits_u8 = base.s32 <= u8::MAX as u16 + && c32 <= u8::MAX as u16 + && base.s64 <= u8::MAX as u16 + && c64 <= u8::MAX as u16 + && base.s128 <= u8::MAX as u16 + && c128 <= u8::MAX as u16 + && base.sref <= u8::MAX as u16 + && cref <= u8::MAX as u16; + + if fits_u8 { + self.instructions.push(Instruction::DropKeepSmall { + base32: base.s32 as u8, + keep32: c32 as u8, + base64: base.s64 as u8, + keep64: c64 as u8, + base128: base.s128 as u8, + keep128: c128 as u8, + base_ref: base.sref as u8, + keep_ref: cref as u8, + }); + } else { + self.instructions.push(Instruction::DropKeep32(base.s32, c32)); + self.instructions.push(Instruction::DropKeep64(base.s64, c64)); + self.instructions.push(Instruction::DropKeep128(base.s128, c128)); + self.instructions.push(Instruction::DropKeepRef(base.sref, cref)); + } + } + + fn patch_jump(&mut self, jump_ip: usize, target: u32) { + if let Instruction::Jump(ip) = &mut self.instructions[jump_ip] { + *ip = target; + } + } + + fn patch_jump_if_zero(&mut self, jump_ip: usize, target: u32) { + if let Instruction::JumpIfZero(ip) = &mut self.instructions[jump_ip] { + *ip = target; + } + } + + fn label_keep_counts(label_types: &[wasmparser::ValType]) -> (u16, u16, u16, u16) { + let mut c32: u16 = 0; + let mut c64: u16 = 0; + let mut c128: u16 = 0; + let mut cref: u16 = 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(_) => cref += 1, + } + } + + (c32, c64, c128, cref) + } + + fn emit_dropkeep_to_label(&mut self, label_depth: u32) { + if self.is_unreachable() { + return; + } + + let frame = match self.validator.get_control_frame(label_depth as usize) { + Some(f) => f, + None => return, + }; + + 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, cref) = Self::label_keep_counts(&label_types); + + self.emit_dropkeep(base, c32, c64, c128, cref); + } + + fn label_types_for_frame(&self, frame: &wasmparser::Frame) -> Vec<wasmparser::ValType> { + 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(), + } + } + } + } } macro_rules! impl_visit_operator { @@ -177,7 +335,7 @@ impl<'a, R: WasmModuleResources> wasmparser::VisitOperator<'a> for FunctionBuild define_operands! { // basic instructions - visit_br(Br, u32), visit_br_if(BrIf, u32), visit_global_get(GlobalGet, u32), visit_i32_const(I32Const, i32), visit_i64_const(I64Const, i64), visit_call(Call, u32), visit_return_call(ReturnCall, u32), visit_memory_size(MemorySize, u32), visit_memory_grow(MemoryGrow, u32), visit_unreachable(Unreachable), visit_nop(Nop), visit_return(Return), visit_i32_eqz(I32Eqz), visit_i32_eq(I32Eq), visit_i32_ne(I32Ne), visit_i32_lt_s(I32LtS), visit_i32_lt_u(I32LtU), visit_i32_gt_s(I32GtS), visit_i32_gt_u(I32GtU), visit_i32_le_s(I32LeS), visit_i32_le_u(I32LeU), visit_i32_ge_s(I32GeS), visit_i32_ge_u(I32GeU), visit_i64_eqz(I64Eqz), visit_i64_eq(I64Eq), visit_i64_ne(I64Ne), visit_i64_lt_s(I64LtS), visit_i64_lt_u(I64LtU), visit_i64_gt_s(I64GtS), visit_i64_gt_u(I64GtU), visit_i64_le_s(I64LeS), visit_i64_le_u(I64LeU), visit_i64_ge_s(I64GeS), visit_i64_ge_u(I64GeU), visit_f32_eq(F32Eq), visit_f32_ne(F32Ne), visit_f32_lt(F32Lt), visit_f32_gt(F32Gt), visit_f32_le(F32Le), visit_f32_ge(F32Ge), visit_f64_eq(F64Eq), visit_f64_ne(F64Ne), visit_f64_lt(F64Lt), visit_f64_gt(F64Gt), visit_f64_le(F64Le), visit_f64_ge(F64Ge), visit_i32_clz(I32Clz), visit_i32_ctz(I32Ctz), visit_i32_popcnt(I32Popcnt), visit_i32_add(I32Add), visit_i32_sub(I32Sub), visit_i32_mul(I32Mul), visit_i32_div_s(I32DivS), visit_i32_div_u(I32DivU), visit_i32_rem_s(I32RemS), visit_i32_rem_u(I32RemU), visit_i32_and(I32And), visit_i32_or(I32Or), visit_i32_xor(I32Xor), visit_i32_shl(I32Shl), visit_i32_shr_s(I32ShrS), visit_i32_shr_u(I32ShrU), visit_i32_rotl(I32Rotl), visit_i32_rotr(I32Rotr), visit_i64_clz(I64Clz), visit_i64_ctz(I64Ctz), visit_i64_popcnt(I64Popcnt), visit_i64_add(I64Add), visit_i64_sub(I64Sub), visit_i64_mul(I64Mul), visit_i64_div_s(I64DivS), visit_i64_div_u(I64DivU), visit_i64_rem_s(I64RemS), visit_i64_rem_u(I64RemU), visit_i64_and(I64And), visit_i64_or(I64Or), visit_i64_xor(I64Xor), visit_i64_shl(I64Shl), visit_i64_shr_s(I64ShrS), visit_i64_shr_u(I64ShrU), visit_i64_rotl(I64Rotl), visit_i64_rotr(I64Rotr), visit_f32_abs(F32Abs), visit_f32_neg(F32Neg), visit_f32_ceil(F32Ceil), visit_f32_floor(F32Floor), visit_f32_trunc(F32Trunc), visit_f32_nearest(F32Nearest), visit_f32_sqrt(F32Sqrt), visit_f32_add(F32Add), visit_f32_sub(F32Sub), visit_f32_mul(F32Mul), visit_f32_div(F32Div), visit_f32_min(F32Min), visit_f32_max(F32Max), visit_f32_copysign(F32Copysign), visit_f64_abs(F64Abs), visit_f64_neg(F64Neg), visit_f64_ceil(F64Ceil), visit_f64_floor(F64Floor), visit_f64_trunc(F64Trunc), visit_f64_nearest(F64Nearest), visit_f64_sqrt(F64Sqrt), visit_f64_add(F64Add), visit_f64_sub(F64Sub), visit_f64_mul(F64Mul), visit_f64_div(F64Div), visit_f64_min(F64Min), visit_f64_max(F64Max), visit_f64_copysign(F64Copysign), visit_i32_wrap_i64(I32WrapI64), visit_i32_trunc_f32_s(I32TruncF32S), visit_i32_trunc_f32_u(I32TruncF32U), visit_i32_trunc_f64_s(I32TruncF64S), visit_i32_trunc_f64_u(I32TruncF64U), visit_i64_extend_i32_s(I64ExtendI32S), visit_i64_extend_i32_u(I64ExtendI32U), visit_i64_trunc_f32_s(I64TruncF32S), visit_i64_trunc_f32_u(I64TruncF32U), visit_i64_trunc_f64_s(I64TruncF64S), visit_i64_trunc_f64_u(I64TruncF64U), visit_f32_convert_i32_s(F32ConvertI32S), visit_f32_convert_i32_u(F32ConvertI32U), visit_f32_convert_i64_s(F32ConvertI64S), visit_f32_convert_i64_u(F32ConvertI64U), visit_f32_demote_f64(F32DemoteF64), visit_f64_convert_i32_s(F64ConvertI32S), visit_f64_convert_i32_u(F64ConvertI32U), visit_f64_convert_i64_s(F64ConvertI64S), visit_f64_convert_i64_u(F64ConvertI64U), visit_f64_promote_f32(F64PromoteF32), visit_i32_reinterpret_f32(I32ReinterpretF32), visit_i64_reinterpret_f64(I64ReinterpretF64), visit_f32_reinterpret_i32(F32ReinterpretI32), visit_f64_reinterpret_i64(F64ReinterpretI64), + visit_global_get(GlobalGet, u32), visit_i32_const(I32Const, i32), visit_i64_const(I64Const, i64), visit_call(Call, u32), visit_return_call(ReturnCall, u32), visit_memory_size(MemorySize, u32), visit_memory_grow(MemoryGrow, u32), visit_unreachable(Unreachable), visit_nop(Nop), visit_return(Return), visit_i32_eqz(I32Eqz), visit_i32_eq(I32Eq), visit_i32_ne(I32Ne), visit_i32_lt_s(I32LtS), visit_i32_lt_u(I32LtU), visit_i32_gt_s(I32GtS), visit_i32_gt_u(I32GtU), visit_i32_le_s(I32LeS), visit_i32_le_u(I32LeU), visit_i32_ge_s(I32GeS), visit_i32_ge_u(I32GeU), visit_i64_eqz(I64Eqz), visit_i64_eq(I64Eq), visit_i64_ne(I64Ne), visit_i64_lt_s(I64LtS), visit_i64_lt_u(I64LtU), visit_i64_gt_s(I64GtS), visit_i64_gt_u(I64GtU), visit_i64_le_s(I64LeS), visit_i64_le_u(I64LeU), visit_i64_ge_s(I64GeS), visit_i64_ge_u(I64GeU), visit_f32_eq(F32Eq), visit_f32_ne(F32Ne), visit_f32_lt(F32Lt), visit_f32_gt(F32Gt), visit_f32_le(F32Le), visit_f32_ge(F32Ge), visit_f64_eq(F64Eq), visit_f64_ne(F64Ne), visit_f64_lt(F64Lt), visit_f64_gt(F64Gt), visit_f64_le(F64Le), visit_f64_ge(F64Ge), visit_i32_clz(I32Clz), visit_i32_ctz(I32Ctz), visit_i32_popcnt(I32Popcnt), visit_i32_add(I32Add), visit_i32_sub(I32Sub), visit_i32_mul(I32Mul), visit_i32_div_s(I32DivS), visit_i32_div_u(I32DivU), visit_i32_rem_s(I32RemS), visit_i32_rem_u(I32RemU), visit_i32_and(I32And), visit_i32_or(I32Or), visit_i32_xor(I32Xor), visit_i32_shl(I32Shl), visit_i32_shr_s(I32ShrS), visit_i32_shr_u(I32ShrU), visit_i32_rotl(I32Rotl), visit_i32_rotr(I32Rotr), visit_i64_clz(I64Clz), visit_i64_ctz(I64Ctz), visit_i64_popcnt(I64Popcnt), visit_i64_add(I64Add), visit_i64_sub(I64Sub), visit_i64_mul(I64Mul), visit_i64_div_s(I64DivS), visit_i64_div_u(I64DivU), visit_i64_rem_s(I64RemS), visit_i64_rem_u(I64RemU), visit_i64_and(I64And), visit_i64_or(I64Or), visit_i64_xor(I64Xor), visit_i64_shl(I64Shl), visit_i64_shr_s(I64ShrS), visit_i64_shr_u(I64ShrU), visit_i64_rotl(I64Rotl), visit_i64_rotr(I64Rotr), visit_f32_abs(F32Abs), visit_f32_neg(F32Neg), visit_f32_ceil(F32Ceil), visit_f32_floor(F32Floor), visit_f32_trunc(F32Trunc), visit_f32_nearest(F32Nearest), visit_f32_sqrt(F32Sqrt), visit_f32_add(F32Add), visit_f32_sub(F32Sub), visit_f32_mul(F32Mul), visit_f32_div(F32Div), visit_f32_min(F32Min), visit_f32_max(F32Max), visit_f32_copysign(F32Copysign), visit_f64_abs(F64Abs), visit_f64_neg(F64Neg), visit_f64_ceil(F64Ceil), visit_f64_floor(F64Floor), visit_f64_trunc(F64Trunc), visit_f64_nearest(F64Nearest), visit_f64_sqrt(F64Sqrt), visit_f64_add(F64Add), visit_f64_sub(F64Sub), visit_f64_mul(F64Mul), visit_f64_div(F64Div), visit_f64_min(F64Min), visit_f64_max(F64Max), visit_f64_copysign(F64Copysign), visit_i32_wrap_i64(I32WrapI64), visit_i32_trunc_f32_s(I32TruncF32S), visit_i32_trunc_f32_u(I32TruncF32U), visit_i32_trunc_f64_s(I32TruncF64S), visit_i32_trunc_f64_u(I32TruncF64U), visit_i64_extend_i32_s(I64ExtendI32S), visit_i64_extend_i32_u(I64ExtendI32U), visit_i64_trunc_f32_s(I64TruncF32S), visit_i64_trunc_f32_u(I64TruncF32U), visit_i64_trunc_f64_s(I64TruncF64S), visit_i64_trunc_f64_u(I64TruncF64U), visit_f32_convert_i32_s(F32ConvertI32S), visit_f32_convert_i32_u(F32ConvertI32U), visit_f32_convert_i64_s(F32ConvertI64S), visit_f32_convert_i64_u(F32ConvertI64U), visit_f32_demote_f64(F32DemoteF64), visit_f64_convert_i32_s(F64ConvertI32S), visit_f64_convert_i32_u(F64ConvertI32U), visit_f64_convert_i64_s(F64ConvertI64S), visit_f64_convert_i64_u(F64ConvertI64U), visit_f64_promote_f32(F64PromoteF32), visit_i32_reinterpret_f32(I32ReinterpretF32), visit_i64_reinterpret_f64(I64ReinterpretF64), visit_f32_reinterpret_i32(F32ReinterpretI32), visit_f64_reinterpret_i64(F64ReinterpretI64), // sign_extension visit_i32_extend8_s(I32Extend8S), visit_i32_extend16_s(I32Extend16S), visit_i64_extend8_s(I64Extend8S), visit_i64_extend16_s(I64Extend16S), visit_i64_extend32_s(I64Extend32S), @@ -328,110 +486,205 @@ impl<'a, R: WasmModuleResources> wasmparser::VisitOperator<'a> for FunctionBuild } } - fn visit_block(&mut self, blockty: wasmparser::BlockType) -> Self::Output { - self.label_ptrs.push(self.instructions.len()); - self.instructions.push(match blockty { - wasmparser::BlockType::Empty => Instruction::Block(0), - wasmparser::BlockType::FuncType(idx) => Instruction::BlockWithFuncType(idx, 0), - wasmparser::BlockType::Type(ty) => Instruction::BlockWithType(convert_valtype(&ty), 0), + fn visit_block(&mut self, _blockty: wasmparser::BlockType) -> Self::Output { + let start_ip = self.current_ip() as usize; + self.ctx_stack.push(LoweringCtx { + kind: BlockKind::Block, + has_else: false, + start_ip, + branch_jumps: Vec::new(), }); } - fn visit_loop(&mut self, ty: wasmparser::BlockType) -> Self::Output { - self.label_ptrs.push(self.instructions.len()); - self.instructions.push(match ty { - wasmparser::BlockType::Empty => Instruction::Loop(0), - wasmparser::BlockType::FuncType(idx) => Instruction::LoopWithFuncType(idx, 0), - wasmparser::BlockType::Type(ty) => Instruction::LoopWithType(convert_valtype(&ty), 0), - }); + fn visit_loop(&mut self, _ty: wasmparser::BlockType) -> Self::Output { + let start_ip = self.current_ip() as usize; + self.ctx_stack.push(LoweringCtx { kind: BlockKind::Loop, has_else: false, start_ip, branch_jumps: Vec::new() }); } - fn visit_if(&mut self, ty: wasmparser::BlockType) -> Self::Output { - self.label_ptrs.push(self.instructions.len()); - self.instructions.push(match ty { - wasmparser::BlockType::Empty => Instruction::If(0, 0), - wasmparser::BlockType::FuncType(idx) => Instruction::IfWithFuncType(idx, 0, 0), - wasmparser::BlockType::Type(ty) => Instruction::IfWithType(convert_valtype(&ty), 0, 0), + fn visit_if(&mut self, _ty: wasmparser::BlockType) -> Self::Output { + let cond_jump_ip = self.current_ip() as usize; + self.instructions.push(Instruction::JumpIfZero(0)); + let start_ip = self.current_ip() as usize; + self.ctx_stack.push(LoweringCtx { + kind: BlockKind::If, + has_else: false, + start_ip, + branch_jumps: alloc::vec![cond_jump_ip], }); } fn visit_else(&mut self) -> Self::Output { - self.label_ptrs.push(self.instructions.len()); - self.instructions.push(Instruction::Else(0)); + let Some(cond_jump_ip) = self + .ctx_stack + .last() + .and_then(|ctx| if matches!(ctx.kind, BlockKind::If) { Some(ctx.branch_jumps[0]) } else { None }) + else { + return; + }; + + let jump_ip = self.current_ip() as usize; + self.instructions.push(Instruction::Jump(0)); + + let after_jump_ip = self.current_ip(); + let Some(ctx) = self.ctx_stack.last_mut() else { + return; + }; + ctx.has_else = true; + ctx.branch_jumps.push(jump_ip); + self.patch_jump_if_zero(cond_jump_ip, after_jump_ip); } fn visit_end(&mut self) -> Self::Output { - let Some(label_pointer) = self.label_ptrs.pop() else { - return self.instructions.push(Instruction::Return); - }; + if self.ctx_stack.is_empty() { + self.instructions.push(Instruction::Return); + return; + } - let current_instr_ptr = self.instructions.len(); - match self.instructions.get_mut(label_pointer) { - Some(Instruction::Else(else_instr_end_offset)) => { - *else_instr_end_offset = (current_instr_ptr - label_pointer) - .try_into() - .expect("else_instr_end_offset is too large, tinywasm does not support if blocks that large"); + let ctx = self.ctx_stack.pop().unwrap(); + let end_ip = self.current_ip(); - // since we're ending an else block, we need to end the if block as well - let Some(if_label_pointer) = self.label_ptrs.pop() else { - self.errors.push(crate::ParseError::UnsupportedOperator( - "Expected to end an if block, but there was no if block to end".to_string(), - )); + match ctx.kind { + BlockKind::Block | BlockKind::Loop => { + let target = if matches!(ctx.kind, BlockKind::Loop) { ctx.start_ip as u32 } else { end_ip }; + for &jump_ip in &ctx.branch_jumps { + self.patch_jump(jump_ip, target); + } + } + BlockKind::If => { + let cond_jump_ip = ctx.branch_jumps[0]; + if !ctx.has_else { + self.patch_jump_if_zero(cond_jump_ip, end_ip); + } + for &jump_ip in &ctx.branch_jumps[1..] { + self.patch_jump(jump_ip, end_ip); + } + } + } + } - return; - }; + fn visit_br(&mut self, depth: u32) -> Self::Output { + self.emit_dropkeep_to_label(depth); - let if_instruction = &mut self.instructions[if_label_pointer]; + if let Some(ctx_idx) = self.get_ctx_idx(depth) { + let jump_ip = self.current_ip() as usize; + self.instructions.push(Instruction::Jump(0)); + self.ctx_stack[ctx_idx].branch_jumps.push(jump_ip); + } else { + self.instructions.push(Instruction::Return); + } + } - let (Instruction::If(else_offset, end_offset) - | Instruction::IfWithFuncType(_, else_offset, end_offset) - | Instruction::IfWithType(_, else_offset, end_offset)) = if_instruction - else { - return self.errors.push(crate::ParseError::UnsupportedOperator( - "Expected to end an if block, but the last label was not an if".to_string(), - )); - }; + fn visit_br_if(&mut self, depth: u32) -> Self::Output { + let cond_jump_ip = self.current_ip() as usize; + self.instructions.push(Instruction::JumpIfZero(0)); - *else_offset = (label_pointer - if_label_pointer) - .try_into() - .expect("else_instr_end_offset is too large, tinywasm does not support blocks that large"); + self.emit_dropkeep_to_label(depth); - *end_offset = (current_instr_ptr - if_label_pointer) - .try_into() - .expect("else_instr_end_offset is too large, tinywasm does not support blocks that large"); - } - Some( - Instruction::Block(end_offset) - | Instruction::BlockWithType(_, end_offset) - | Instruction::BlockWithFuncType(_, end_offset) - | Instruction::Loop(end_offset) - | Instruction::LoopWithFuncType(_, end_offset) - | Instruction::LoopWithType(_, end_offset) - | Instruction::If(_, end_offset) - | Instruction::IfWithFuncType(_, _, end_offset) - | Instruction::IfWithType(_, _, end_offset), - ) => { - *end_offset = (current_instr_ptr - label_pointer) - .try_into() - .expect("else_instr_end_offset is too large, tinywasm does not support blocks that large"); - } - _ => { - unreachable!("Expected to end a block, but the last label was not a block") - } - }; + if let Some(ctx_idx) = self.get_ctx_idx(depth) { + let jump_ip = self.current_ip() as usize; + self.instructions.push(Instruction::Jump(0)); + self.ctx_stack[ctx_idx].branch_jumps.push(jump_ip); + } else { + self.instructions.push(Instruction::Return); + } - self.instructions.push(Instruction::EndBlockFrame); + self.patch_jump_if_zero(cond_jump_ip, self.current_ip()); } fn visit_br_table(&mut self, targets: wasmparser::BrTable<'_>) -> Self::Output { - let def = targets.default(); - let instrs = targets + let ts = targets .targets() - .map(|t| t.map(Instruction::BrLabel)) - .collect::<Result<Vec<Instruction>, wasmparser::BinaryReaderError>>() - .expect("visit_br_table: BrTable targets are invalid, this should have been caught by the validator"); + .collect::<Result<Vec<_>, wasmparser::BinaryReaderError>>() + .expect("visit_br_table: BrTable targets are invalid"); - self.instructions.extend(([Instruction::BrTable(def, instrs.len() as u32)].into_iter()).chain(instrs)); + let default_depth = targets.default(); + let len = ts.len() as u32; + let target_depths: Vec<u32> = ts; + + let header_ip = self.current_ip() as usize; + self.instructions.push(Instruction::BranchTable(0, len)); + + let target_table_ip = self.current_ip() as usize; + for _ in 0..len { + self.instructions.push(Instruction::BranchTableTarget(0)); + } + let default_target_ip = self.current_ip() as usize; + self.instructions.push(Instruction::BranchTableTarget(0)); + + let mut seen = alloc::collections::BTreeMap::<u32, usize>::new(); + struct PadInfo { + depth: u32, + pad_start: u32, + jump_or_ret_ip: usize, + is_return: bool, + } + let mut pads: Vec<PadInfo> = Vec::new(); + + for &depth in target_depths.iter().chain(core::iter::once(&default_depth)) { + if seen.contains_key(&depth) { + continue; + } + seen.insert(depth, pads.len()); + + let pad_start = self.current_ip(); + + let frame = if self.is_unreachable() { None } else { self.validator.get_control_frame(depth as usize) }; + let Some(frame) = frame else { + let ip = self.current_ip() as usize; + self.instructions.push(Instruction::Return); + pads.push(PadInfo { depth, pad_start, jump_or_ret_ip: ip, is_return: true }); + continue; + }; + + let base = self.stack_base_at_frame(depth as usize); + let label_types: Vec<_> = self.label_types_for_frame(frame); + let (c32, c64, c128, cref) = Self::label_keep_counts(&label_types); + + self.emit_dropkeep(base, c32, c64, c128, cref); + + let jump_ip = self.current_ip() as usize; + self.instructions.push(Instruction::Jump(0)); + pads.push(PadInfo { depth, pad_start, jump_or_ret_ip: jump_ip, is_return: false }); + } + + for (i, &depth) in target_depths.iter().enumerate() { + let pad_idx = seen[&depth]; + if let Instruction::BranchTableTarget(ip) = &mut self.instructions[target_table_ip + i] { + *ip = pads[pad_idx].pad_start; + } + } + + let default_pad_idx = seen[&default_depth]; + if let Instruction::BranchTableTarget(ip) = &mut self.instructions[default_target_ip] { + *ip = pads[default_pad_idx].pad_start; + } + if let Instruction::BranchTable(default_ip, _) = &mut self.instructions[header_ip] { + *default_ip = pads[default_pad_idx].pad_start; + } + + for pad in &pads { + if pad.is_return { + continue; + } + let Some(frame) = self.validator.get_control_frame(pad.depth as usize) else { + self.instructions[pad.jump_or_ret_ip] = Instruction::Return; + continue; + }; + let Some(ctx_idx) = self.get_ctx_idx(pad.depth) else { + self.instructions[pad.jump_or_ret_ip] = Instruction::Return; + continue; + }; + match frame.kind { + FrameKind::Loop => { + if let Instruction::Jump(target) = &mut self.instructions[pad.jump_or_ret_ip] { + *target = self.ctx_stack[ctx_idx].start_ip as u32; + } + } + _ => { + self.ctx_stack[ctx_idx].branch_jumps.push(pad.jump_or_ret_ip); + } + } + } } fn visit_call_indirect(&mut self, ty: u32, table: u32) -> Self::Output { |