From f7fb3a80c3cab7d7916a3d03c0f8b6b0ef8193d1 Mon Sep 17 00:00:00 2001 From: Chris Fallin Date: Sat, 26 Sep 2026 22:56:07 -0700 Subject: [PATCH] Backend: avoid recursion; use explicit stack with iteration instead. This was overflowing the native stack on some very (very) large function bodies. Note that the main lowering loop needs to be made explicitly iterative as well; the auto-derived `Drop` impl was too recursive as well (!). --- src/backend/mod.rs | 98 ++++++++++++++++++++++++++--------------- src/backend/stackify.rs | 27 ++++++++++++ 2 files changed, 89 insertions(+), 36 deletions(-) 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);