diff --git a/src/backend/mod.rs b/src/backend/mod.rs index d02c45d..5c81890 100644 --- a/src/backend/mod.rs +++ b/src/backend/mod.rs @@ -105,47 +105,73 @@ impl<'a> WasmFuncBackend<'a> { block: &WasmBlock<'_>, func: &mut wasm_encoder::Function, ) { - match block { - WasmBlock::Block { body, .. } => { - func.instruction(&wasm_encoder::Instruction::Block( - wasm_encoder::BlockType::Empty, - )); - for sub_block in &body[..] { - self.lower_block(ctx, sub_block, func); + enum Work<'b, 'a> { + Block(&'b WasmBlock<'a>), + Else, + End, + } + let mut work = vec![Work::Block(block)]; + while let Some(w) = work.pop() { + let block = match w { + Work::Block(b) => b, + Work::Else => { + func.instruction(&wasm_encoder::Instruction::Else); + continue; } - func.instruction(&wasm_encoder::Instruction::End); - } - WasmBlock::Loop { body, .. } => { - func.instruction(&wasm_encoder::Instruction::Loop( - wasm_encoder::BlockType::Empty, - )); - for sub_block in &body[..] { - self.lower_block(ctx, sub_block, func); + Work::End => { + func.instruction(&wasm_encoder::Instruction::End); + continue; } - func.instruction(&wasm_encoder::Instruction::End); - } - WasmBlock::Br { target } => { - func.instruction(&wasm_encoder::Instruction::Br(target.index())); - } - WasmBlock::If { - cond, - if_true, - if_false, - } => { - self.lower_value(ctx, *cond, func); - func.instruction(&wasm_encoder::Instruction::If( - wasm_encoder::BlockType::Empty, - )); - for sub_block in &if_true[..] { - self.lower_block(ctx, sub_block, func); + }; + match block { + WasmBlock::Block { body, .. } => { + func.instruction(&wasm_encoder::Instruction::Block( + wasm_encoder::BlockType::Empty, + )); + work.push(Work::End); + work.extend(body.iter().rev().map(Work::Block)); } - if if_false.len() > 0 { - func.instruction(&wasm_encoder::Instruction::Else); - for sub_block in &if_false[..] { - self.lower_block(ctx, sub_block, func); + WasmBlock::Loop { body, .. } => { + func.instruction(&wasm_encoder::Instruction::Loop( + wasm_encoder::BlockType::Empty, + )); + work.push(Work::End); + work.extend(body.iter().rev().map(Work::Block)); + } + WasmBlock::If { + cond, + if_true, + if_false, + } => { + self.lower_value(ctx, *cond, func); + func.instruction(&wasm_encoder::Instruction::If( + wasm_encoder::BlockType::Empty, + )); + work.push(Work::End); + if if_false.len() > 0 { + work.extend(if_false.iter().rev().map(Work::Block)); + work.push(Work::Else); } + work.extend(if_true.iter().rev().map(Work::Block)); } - func.instruction(&wasm_encoder::Instruction::End); + other => self.lower_leaf_block(ctx, other, func), + } + } + } + + /// Emit a `WasmBlock` with no nested blocks. + fn lower_leaf_block( + &self, + ctx: &CompileContext<'_>, + block: &WasmBlock<'_>, + func: &mut wasm_encoder::Function, + ) { + match block { + WasmBlock::Block { .. } | WasmBlock::Loop { .. } | WasmBlock::If { .. } => { + unreachable!("nested blocks are lowered by lower_block") + } + WasmBlock::Br { target } => { + func.instruction(&wasm_encoder::Instruction::Br(target.index())); } WasmBlock::Select { selector, diff --git a/src/backend/stackify.rs b/src/backend/stackify.rs index 7788ebe..681bd9f 100644 --- a/src/backend/stackify.rs +++ b/src/backend/stackify.rs @@ -56,6 +56,33 @@ pub enum WasmBlock<'a> { Unreachable, } +/// Note that we need to drop iteratively, avoiding recursion in the +/// default derived impl, to avoid blowing the stack (the tree can get +/// very deep). +impl<'a> Drop for WasmBlock<'a> { + fn drop(&mut self) { + fn children<'a>(b: &mut WasmBlock<'a>, out: &mut Vec>) { + match b { + WasmBlock::Block { body, .. } | WasmBlock::Loop { body, .. } => { + out.append(body); + } + WasmBlock::If { + if_true, if_false, .. + } => { + out.append(if_true); + out.append(if_false); + } + _ => {} + } + } + let mut work = vec![]; + children(self, &mut work); + while let Some(mut b) = work.pop() { + children(&mut b, &mut work); + } + } +} + /// A Wasm branch target label: number of scopes outward to branch to. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct WasmLabel(u32);