diff --git a/src/frontend.rs b/src/frontend.rs index 999a6b7..94442e8 100644 --- a/src/frontend.rs +++ b/src/frontend.rs @@ -110,29 +110,62 @@ struct FunctionBodyBuilder<'a, 'b> { #[derive(Debug)] enum Frame { Block { + start_depth: usize, out: BlockId, params: Vec, results: Vec, }, Loop { - top: BlockId, + start_depth: usize, + header: BlockId, out: BlockId, params: Vec, results: Vec, }, If { + start_depth: usize, out: BlockId, el: BlockId, + param_values: Vec, params: Vec, results: Vec, }, Else { + start_depth: usize, out: BlockId, params: Vec, results: Vec, }, } +impl Frame { + fn start_depth(&self) -> usize { + match self { + Frame::Block { start_depth, .. } + | Frame::Loop { start_depth, .. } + | Frame::If { start_depth, .. } + | Frame::Else { start_depth, .. } => *start_depth, + } + } + + fn br_args(&self) -> &[Type] { + match self { + Frame::Block { results, .. } + | Frame::If { results, .. } + | Frame::Else { results, .. } => &results[..], + Frame::Loop { params, .. } => ¶ms[..], + } + } + + fn br_target(&self) -> BlockId { + match self { + Frame::Block { out, .. } => *out, + Frame::Loop { header, .. } => *header, + Frame::If { out, .. } | Frame::Else { out, .. } => *out, + } + } +} + impl<'a, 'b> FunctionBodyBuilder<'a, 'b> { fn new(module: &'b Module<'a>, my_sig: SignatureId, body: &'b mut FunctionBody<'a>) -> Self { body.blocks.push(Block::default()); @@ -146,8 +179,12 @@ impl<'a, 'b> FunctionBodyBuilder<'a, 'b> { } } + fn add_block_params(&mut self, block: BlockId, tys: &[Type]) { + self.body.blocks[block].params.extend_from_slice(tys); + } + fn handle_op(&mut self, op: Operator<'a>) -> Result<()> { - match op { + match &op { Operator::Unreachable | Operator::Call { .. } | Operator::LocalGet { .. } @@ -158,27 +195,171 @@ impl<'a, 'b> FunctionBodyBuilder<'a, 'b> { None => { self.emit(Operator::Return)?; } - Some(Frame::Block { out, .. }) | Some(Frame::Loop { out, .. }) => { - // No need to manipulate stack: assuming the input - // Wasm was validated properly, the `results` - // values must be on the top of the stack now, and - // they will remain so once we exit the block. - self.emit_branch(out); + Some(Frame::Block { + start_depth, + out, + results, + .. + }) + | Some(Frame::Loop { + start_depth, + out, + results, + .. + }) => { + // Generate a branch to the out-block with + // blockparams for the results. + let result_values = self.op_stack.split_off(results.len()); + self.emit_branch(out, &result_values[..]); + assert_eq!(self.op_stack.len(), start_depth); self.cur_block = out; + self.push_block_params(&results[..]); + } + Some(Frame::If { + start_depth, + out, + el, + param_values, + results, + .. + }) => { + // Generate a branch to the out-block with + // blockparams for the results. + let result_values = self.op_stack.split_off(results.len()); + self.emit_branch(out, &result_values[..]); + // No `else`, so we need to generate a trivial + // branch in the else-block. If the if-block-type + // has results, they must be exactly the params. + let else_result_values = param_values; + assert_eq!(else_result_values.len(), results.len()); + self.emit_branch(el, &else_result_values[..]); + assert_eq!(self.op_stack.len(), start_depth); + self.cur_block = out; + self.push_block_params(&results[..]); + } + Some(Frame::Else { + out, + results, + start_depth, + .. + }) => { + // Generate a branch to the out-block with + // blockparams for the results. + let result_values = self.op_stack.split_off(results.len()); + assert_eq!(self.op_stack.len(), start_depth); + self.emit_branch(out, &result_values[..]); + self.cur_block = out; + self.push_block_params(&results[..]); } - _ => bail!("unsupported block type"), }, Operator::Block { ty } => { - let (params, results) = self.block_params_and_results(ty); + let (params, results) = self.block_params_and_results(*ty); let out = self.create_block(); + self.add_block_params(out, &results[..]); + let start_depth = self.op_stack.len() - params.len(); self.ctrl_stack.push(Frame::Block { + start_depth, out, params, results, }); } + Operator::Loop { ty } => { + let (params, results) = self.block_params_and_results(*ty); + let header = self.create_block(); + self.add_block_params(header, ¶ms[..]); + let initial_args = self.op_stack.split_off(params.len()); + let start_depth = self.op_stack.len(); + self.emit_branch(header, &initial_args[..]); + self.cur_block = header; + self.push_block_params(¶ms[..]); + let out = self.create_block(); + self.ctrl_stack.push(Frame::Loop { + start_depth, + header, + out, + params, + results, + }); + } + + Operator::If { ty } => { + let (params, results) = self.block_params_and_results(*ty); + let if_true = self.create_block(); + let if_false = self.create_block(); + let join = self.create_block(); + self.add_block_params(join, &results[..]); + let cond = self.op_stack.pop().unwrap(); + let param_values = self.op_stack[self.op_stack.len() - params.len()..].to_vec(); + let start_depth = self.op_stack.len(); + self.ctrl_stack.push(Frame::If { + start_depth, + out: join, + el: if_false, + param_values, + params, + results, + }); + self.cur_block = if_true; + self.emit_cond_branch(cond, if_true, &[], if_false, &[]); + } + + Operator::Else => { + if let Frame::If { + start_depth, + out, + el, + param_values, + params, + results, + } = self.ctrl_stack.pop().unwrap() + { + let if_results = self.op_stack.split_off(results.len()); + self.emit_branch(out, &if_results[..]); + self.op_stack.extend(param_values); + self.ctrl_stack.push(Frame::Else { + start_depth, + out, + params, + results, + }); + self.cur_block = el; + } else { + bail!("Else without If on top of frame stack"); + } + } + + Operator::Br { relative_depth } | Operator::BrIf { relative_depth } => { + let cond = match &op { + Operator::Br { .. } => None, + Operator::BrIf { .. } => Some(self.op_stack.pop().unwrap()), + _ => unreachable!(), + }; + // Pop skipped-over frames. + let _ = self.ctrl_stack.split_off(*relative_depth as usize); + // Get the frame we're branching to. + let frame = self.ctrl_stack.pop().unwrap(); + // Get the args off the stack. + let args = self.op_stack.split_off(frame.br_args().len()); + // Truncate the result stack down to the expected height. + self.op_stack.truncate(frame.start_depth()); + // Finally, generate the branch itself. + match cond { + None => { + self.emit_branch(frame.br_target(), &args[..]); + } + Some(cond) => { + let cont = self.create_block(); + self.emit_cond_branch(cond, frame.br_target(), &args[..], cont, &[]); + self.cur_block = cont; + } + } + } + + Operator::BrTable { .. } => {} + _ => bail!("Unsupported operator: {:?}", op), } @@ -204,11 +385,58 @@ impl<'a, 'b> FunctionBodyBuilder<'a, 'b> { } } - fn emit_branch(&mut self, target: BlockId) { + fn emit_branch(&mut self, target: BlockId, args: &[ValueId]) { let block = self.cur_block; + let args = args.iter().map(|&val| Operand::Value(val)).collect(); + let target = BlockTarget { + block: target, + args, + }; self.body.blocks[block].terminator = Terminator::Br { target }; } + fn emit_cond_branch( + &mut self, + cond: ValueId, + if_true: BlockId, + if_true_args: &[ValueId], + if_false: BlockId, + if_false_args: &[ValueId], + ) { + let block = self.cur_block; + let if_true_args = if_true_args + .iter() + .map(|&val| Operand::Value(val)) + .collect(); + let if_false_args = if_false_args + .iter() + .map(|&val| Operand::Value(val)) + .collect(); + self.body.blocks[block].terminator = Terminator::CondBr { + cond: Operand::Value(cond), + if_true: BlockTarget { + block: if_true, + args: if_true_args, + }, + if_false: BlockTarget { + block: if_false, + args: if_false_args, + }, + }; + } + + fn push_block_params(&mut self, tys: &[Type]) { + assert_eq!(tys, self.body.blocks[self.cur_block].params); + for (i, &ty) in tys.iter().enumerate() { + let value_id = self.body.values.len() as ValueId; + self.body.values.push(ValueDef { + kind: ValueKind::BlockParam(self.cur_block, i), + ty, + }); + self.op_stack.push(value_id); + } + } + fn emit(&mut self, op: Operator<'a>) -> Result<()> { let block = self.cur_block; let inst = self.body.blocks[block].insts.len() as InstId; diff --git a/src/ir.rs b/src/ir.rs index f783fd4..eac9af7 100644 --- a/src/ir.rs +++ b/src/ir.rs @@ -68,20 +68,26 @@ pub enum Operand<'a> { Sub(Box>), } +#[derive(Clone, Debug)] +pub struct BlockTarget<'a> { + pub block: BlockId, + pub args: Vec>, +} + #[derive(Clone, Debug)] pub enum Terminator<'a> { Br { - target: BlockId, + target: BlockTarget<'a>, }, CondBr { cond: Operand<'a>, - if_true: BlockId, - if_false: BlockId, + if_true: BlockTarget<'a>, + if_false: BlockTarget<'a>, }, Select { value: Operand<'a>, - targets: Vec, - default: BlockId, + targets: Vec>, + default: BlockTarget<'a>, }, Return { values: Vec>,