use anyhow::{bail, ensure, Result}; use wasmtime_environ::{VMOffsets, WasmHeapType, WasmValType}; use super::ControlStackFrame; use crate::{ abi::{scratch, vmctx, ABIOperand, ABIResults, RetArea}, codegen::{CodeGenError, CodeGenPhase, Emission, Prologue}, frame::Frame, isa::reg::RegClass, masm::{MacroAssembler, OperandSize, RegImm, SPOffset, ShiftKind, StackSlot}, reg::{writable, Reg}, regalloc::RegAlloc, stack::{Stack, TypedReg, Val}, }; /// The code generation context. /// The code generation context is made up of three /// essential data structures: /// /// * The register allocator, in charge of keeping the inventory of register /// availability. /// * The value stack, which keeps track of the state of the values /// after each operation. /// * The current function's frame. /// /// These data structures normally require cooperating with each other /// to perform most of the operations needed during the code /// generation process. The code generation context should /// be generally used as the single entry point to access /// the compound functionality provided by its elements. pub(crate) struct CodeGenContext<'a, P: CodeGenPhase> { /// The register allocator. pub regalloc: RegAlloc, /// The value stack. pub stack: Stack, /// The current function's frame. pub frame: Frame

, /// Reachability state. pub reachable: bool, /// A reference to the VMOffsets. pub vmoffsets: &'a VMOffsets, } impl<'a> CodeGenContext<'a, Emission> { /// Prepares arguments for emitting an i32 shift operation. pub fn i32_shift(&mut self, masm: &mut M, kind: ShiftKind) -> Result<()> where M: MacroAssembler, { let top = self .stack .peek() .ok_or_else(|| CodeGenError::missing_values_in_stack())?; if top.is_i32_const() { let val = self .stack .pop_i32_const() .ok_or_else(|| CodeGenError::missing_values_in_stack())?; let typed_reg = self.pop_to_reg(masm, None)?; masm.shift_ir( writable!(typed_reg.reg), val as u64, typed_reg.reg, kind, OperandSize::S32, )?; self.stack.push(typed_reg.into()); } else { masm.shift(self, kind, OperandSize::S32)?; } Ok(()) } /// Prepares arguments for emitting an i64 binary operation. pub fn i64_shift(&mut self, masm: &mut M, kind: ShiftKind) -> Result<()> where M: MacroAssembler, { let top = self .stack .peek() .ok_or_else(|| CodeGenError::missing_values_in_stack())?; if top.is_i64_const() { let val = self .stack .pop_i64_const() .ok_or_else(|| CodeGenError::missing_values_in_stack())?; let typed_reg = self.pop_to_reg(masm, None)?; masm.shift_ir( writable!(typed_reg.reg), val as u64, typed_reg.reg, kind, OperandSize::S64, )?; self.stack.push(typed_reg.into()); } else { masm.shift(self, kind, OperandSize::S64)?; }; Ok(()) } } impl<'a> CodeGenContext<'a, Prologue> { /// Create a new code generation context. pub fn new( regalloc: RegAlloc, stack: Stack, frame: Frame, vmoffsets: &'a VMOffsets, ) -> Self { Self { regalloc, stack, frame, reachable: true, vmoffsets, } } /// Prepares the frame for the [`Emission`] code generation phase. pub fn for_emission(self) -> CodeGenContext<'a, Emission> { CodeGenContext { regalloc: self.regalloc, stack: self.stack, reachable: self.reachable, vmoffsets: self.vmoffsets, frame: self.frame.for_emission(), } } } impl<'a> CodeGenContext<'a, Emission> { /// Request a specific register to the register allocator, /// spilling if not available. pub fn reg(&mut self, named: Reg, masm: &mut M) -> Result { self.regalloc.reg(named, |regalloc| { Self::spill_impl(&mut self.stack, regalloc, &self.frame, masm) }) } /// Allocate a register for the given WebAssembly type. pub fn reg_for_type( &mut self, ty: WasmValType, masm: &mut M, ) -> Result { use WasmValType::*; match ty { I32 | I64 => self.reg_for_class(RegClass::Int, masm), F32 | F64 => self.reg_for_class(RegClass::Float, masm), // All of our supported architectures use the float registers for vector operations. V128 => self.reg_for_class(RegClass::Float, masm), Ref(rt) => match rt.heap_type { WasmHeapType::Func | WasmHeapType::Extern => { self.reg_for_class(RegClass::Int, masm) } _ => bail!(CodeGenError::unsupported_wasm_type()), }, } } /// Request the register allocator to provide the next available /// register of the specified class. pub fn reg_for_class( &mut self, class: RegClass, masm: &mut M, ) -> Result { self.regalloc.reg_for_class(class, &mut |regalloc| { Self::spill_impl(&mut self.stack, regalloc, &self.frame, masm) }) } /// Convenience wrapper around `CodeGenContext::reg_for_class`, to /// request the next available general purpose register. pub fn any_gpr(&mut self, masm: &mut M) -> Result { self.reg_for_class(RegClass::Int, masm) } /// Convenience wrapper around `CodeGenContext::reg_for_class`, to /// request the next available floating point register. pub fn any_fpr(&mut self, masm: &mut M) -> Result { self.reg_for_class(RegClass::Float, masm) } /// Executes the provided function, guaranteeing that the specified set of /// registers, if any, remain unallocatable throughout the function's /// execution. pub fn without<'r, T, M, F>( &mut self, regs: impl IntoIterator + Copy, masm: &mut M, mut f: F, ) -> Result where M: MacroAssembler, F: FnMut(&mut Self, &mut M) -> T, { for r in regs { self.reg(*r, masm)?; } let result = f(self, masm); for r in regs { self.free_reg(*r); } Ok(result) } /// Free the given register. pub fn free_reg(&mut self, reg: impl Into) { let reg: Reg = reg.into(); self.regalloc.free(reg); } /// Loads the stack top value into the next available register, if /// it isn't already one; spilling if there are no registers /// available. Optionally the caller may specify a specific /// destination register. /// When a named register is requested and it's not at the top of the /// stack a move from register to register might happen, in which case /// the source register will be freed. pub fn pop_to_reg( &mut self, masm: &mut M, named: Option, ) -> Result { let typed_reg = if let Some(dst) = named { self.stack.pop_named_reg(dst) } else { self.stack.pop_reg() }; if let Some(dst) = typed_reg { return Ok(dst); } let val = self.stack.pop().expect("a value at stack top"); let reg = if let Some(r) = named { self.reg(r, masm)? } else { self.reg_for_type(val.ty(), masm)? }; if val.is_mem() { let mem = val.unwrap_mem(); let curr_offset = masm.sp_offset()?.as_u32(); let slot_offset = mem.slot.offset.as_u32(); ensure!( curr_offset == slot_offset, CodeGenError::invalid_sp_offset(), ); masm.pop(writable!(reg), val.ty().try_into()?)?; } else { self.move_val_to_reg(&val, reg, masm)?; // Free the source value if it is a register. if val.is_reg() { self.free_reg(val.unwrap_reg()); } } Ok(TypedReg::new(val.ty(), reg)) } /// Pops the value stack top and stores it at the specified address. pub fn pop_to_addr(&mut self, masm: &mut M, addr: M::Address) -> Result<()> { let val = self.stack.pop().expect("a value at stack top"); let ty = val.ty(); let size: OperandSize = ty.try_into()?; match val { Val::Reg(tr) => { masm.store(tr.reg.into(), addr, size)?; self.free_reg(tr.reg); } Val::I32(v) => masm.store(RegImm::i32(v), addr, size)?, Val::I64(v) => masm.store(RegImm::i64(v), addr, size)?, Val::F32(v) => masm.store(RegImm::f32(v.bits()), addr, size)?, Val::F64(v) => masm.store(RegImm::f64(v.bits()), addr, size)?, Val::V128(v) => masm.store(RegImm::v128(v), addr, size)?, Val::Local(local) => { let slot = self.frame.get_wasm_local(local.index); let scratch = scratch!(M); let local_addr = masm.local_address(&slot)?; masm.load(local_addr, writable!(scratch), size)?; masm.store(scratch.into(), addr, size)?; } Val::Memory(_) => { let scratch = scratch!(M, &ty); masm.pop(writable!(scratch), size)?; masm.store(scratch.into(), addr, size)?; } } Ok(()) } /// Move a stack value to the given register. pub fn move_val_to_reg( &self, src: &Val, dst: Reg, masm: &mut M, ) -> Result<()> { let size: OperandSize = src.ty().try_into()?; match src { Val::Reg(tr) => masm.mov(writable!(dst), RegImm::reg(tr.reg), size), Val::I32(imm) => masm.mov(writable!(dst), RegImm::i32(*imm), size), Val::I64(imm) => masm.mov(writable!(dst), RegImm::i64(*imm), size), Val::F32(imm) => masm.mov(writable!(dst), RegImm::f32(imm.bits()), size), Val::F64(imm) => masm.mov(writable!(dst), RegImm::f64(imm.bits()), size), Val::V128(imm) => masm.mov(writable!(dst), RegImm::v128(*imm), size), Val::Local(local) => { let slot = self.frame.get_wasm_local(local.index); let addr = masm.local_address(&slot)?; masm.load(addr, writable!(dst), size) } Val::Memory(mem) => { let addr = masm.address_from_sp(mem.slot.offset)?; masm.load(addr, writable!(dst), size) } } } /// Prepares arguments for emitting a unary operation. /// /// The `emit` function returns the `TypedReg` to put on the value stack. pub fn unop(&mut self, masm: &mut M, size: OperandSize, emit: &mut F) -> Result<()> where F: FnMut(&mut M, Reg, OperandSize) -> Result, M: MacroAssembler, { let typed_reg = self.pop_to_reg(masm, None)?; let dst = emit(masm, typed_reg.reg, size)?; self.stack.push(dst.into()); Ok(()) } /// Prepares arguments for emitting a binary operation. /// /// The `emit` function returns the `TypedReg` to put on the value stack. pub fn binop(&mut self, masm: &mut M, size: OperandSize, emit: F) -> Result<()> where F: FnOnce(&mut M, Reg, Reg, OperandSize) -> Result, M: MacroAssembler, { let src = self.pop_to_reg(masm, None)?; let dst = self.pop_to_reg(masm, None)?; let dst = emit(masm, dst.reg, src.reg.into(), size)?; self.free_reg(src); self.stack.push(dst.into()); Ok(()) } /// Prepares arguments for emitting an f32 or f64 comparison operation. pub fn float_cmp_op(&mut self, masm: &mut M, size: OperandSize, emit: F) -> Result<()> where F: FnOnce(&mut M, Reg, Reg, Reg, OperandSize) -> Result<()>, M: MacroAssembler, { let src2 = self.pop_to_reg(masm, None)?; let src1 = self.pop_to_reg(masm, None)?; let dst = self.any_gpr(masm)?; emit(masm, dst, src1.reg, src2.reg, size)?; self.free_reg(src1); self.free_reg(src2); let dst = match size { // Float comparison operators are defined as // [f64 f64] -> i32 // https://webassembly.github.io/spec/core/appendix/index-instructions.html OperandSize::S32 | OperandSize::S64 => TypedReg::i32(dst), OperandSize::S8 | OperandSize::S16 | OperandSize::S128 => { bail!(CodeGenError::unexpected_operand_size()) } }; self.stack.push(dst.into()); Ok(()) } /// Prepares arguments for emitting an i32 binary operation. /// /// The `emit` function returns the `TypedReg` to put on the value stack. pub fn i32_binop(&mut self, masm: &mut M, mut emit: F) -> Result<()> where F: FnMut(&mut M, Reg, RegImm, OperandSize) -> Result, M: MacroAssembler, { let top = self.stack.peek().expect("value at stack top"); if top.is_i32_const() { let val = self .stack .pop_i32_const() .expect("i32 const value at stack top"); let typed_reg = self.pop_to_reg(masm, None)?; let dst = emit(masm, typed_reg.reg, RegImm::i32(val), OperandSize::S32)?; self.stack.push(dst.into()); } else { self.binop(masm, OperandSize::S32, |masm, dst, src, size| { emit(masm, dst, src.into(), size) })?; } Ok(()) } /// Prepares arguments for emitting an i64 binary operation. /// /// The `emit` function returns the `TypedReg` to put on the value stack. pub fn i64_binop(&mut self, masm: &mut M, emit: F) -> Result<()> where F: FnOnce(&mut M, Reg, RegImm, OperandSize) -> Result, M: MacroAssembler, { let top = self.stack.peek().expect("value at stack top"); if top.is_i64_const() { let val = self .stack .pop_i64_const() .expect("i64 const value at stack top"); let typed_reg = self.pop_to_reg(masm, None)?; let dst = emit(masm, typed_reg.reg, RegImm::i64(val), OperandSize::S64)?; self.stack.push(dst.into()); } else { self.binop(masm, OperandSize::S64, |masm, dst, src, size| { emit(masm, dst, src.into(), size) })?; }; Ok(()) } /// Prepares arguments for emitting a convert operation. pub fn convert_op(&mut self, masm: &mut M, dst_ty: WasmValType, emit: F) -> Result<()> where F: FnOnce(&mut M, Reg, Reg, OperandSize) -> Result<()>, M: MacroAssembler, { let src = self.pop_to_reg(masm, None)?; let dst = self.reg_for_type(dst_ty, masm)?; let dst_size = match dst_ty { WasmValType::I32 => OperandSize::S32, WasmValType::I64 => OperandSize::S64, WasmValType::F32 => OperandSize::S32, WasmValType::F64 => OperandSize::S64, WasmValType::V128 => bail!(CodeGenError::unsupported_wasm_type()), WasmValType::Ref(_) => bail!(CodeGenError::unsupported_wasm_type()), }; emit(masm, dst, src.into(), dst_size)?; self.free_reg(src); self.stack.push(TypedReg::new(dst_ty, dst).into()); Ok(()) } /// Prepares arguments for emitting a convert operation with a temporary /// register. pub fn convert_op_with_tmp_reg( &mut self, masm: &mut M, dst_ty: WasmValType, tmp_reg_class: RegClass, emit: F, ) -> Result<()> where F: FnOnce(&mut M, Reg, Reg, Reg, OperandSize) -> Result<()>, M: MacroAssembler, { let tmp_gpr = self.reg_for_class(tmp_reg_class, masm)?; self.convert_op(masm, dst_ty, |masm, dst, src, dst_size| { emit(masm, dst, src, tmp_gpr, dst_size) })?; self.free_reg(tmp_gpr); Ok(()) } /// Drops the last `n` elements of the stack, calling the provided /// function for each `n` stack value. /// The values are dropped in top-to-bottom order. pub fn drop_last(&mut self, last: usize, mut f: F) -> Result<()> where F: FnMut(&mut RegAlloc, &Val) -> Result<()>, { if last > 0 { let len = self.stack.len(); ensure!(last <= len, CodeGenError::unexpected_value_stack_index(),); let truncate = self.stack.len() - last; let stack_mut = self.stack.inner_mut(); // Invoke the callback in top-to-bottom order. for v in stack_mut[truncate..].into_iter().rev() { f(&mut self.regalloc, v)? } stack_mut.truncate(truncate); } Ok(()) } /// Convenience wrapper around [`Self::spill_callback`]. /// /// This function exists for cases in which triggering an unconditional /// spill is needed, like before entering control flow. pub fn spill(&mut self, masm: &mut M) -> Result<()> { Self::spill_impl(&mut self.stack, &mut self.regalloc, &self.frame, masm) } /// Prepares the compiler to emit an uncoditional jump to the given /// destination branch. This process involves: /// * Balancing the machine /// stack pointer and value stack by popping it to match the destination /// branch. /// * Updating the reachability state. /// * Marking the destination frame as a destination target. pub fn unconditional_jump( &mut self, dest: &mut ControlStackFrame, masm: &mut M, mut f: F, ) -> Result<()> where M: MacroAssembler, F: FnMut(&mut M, &mut Self, &mut ControlStackFrame) -> Result<()>, { let state = dest.stack_state(); let target_offset = state.target_offset; let base_offset = state.base_offset; // Invariant: The SP, must be greater or equal to the target // SP, given that we haven't popped any results by this point // yet. But it may happen in the callback. ensure!( masm.sp_offset()?.as_u32() >= base_offset.as_u32(), CodeGenError::invalid_sp_offset() ); f(masm, self, dest)?; // The following snippet, pops the stack pointer to ensure that it // is correctly placed according to the expectations of the destination // branch. // // This is done in the context of unconditional jumps, as the machine // stack might be left unbalanced at the jump site, due to register // spills. Note that in some cases the stack pointer offset might be // already less than or equal to the original stack pointer offset // registered when entering the destination control stack frame, which // effectively means that when reaching the jump site no extra space was // allocated similar to what would happen in a fall through in which we // assume that the program has allocated and deallocated the right // amount of stack space. // // More generally speaking the current stack pointer will be less than // the original stack pointer offset in cases in which the top value in // the value stack is a memory entry which needs to be popped into the // return location according to the ABI (a register for single value // returns and a memory slot for 1+ returns). This could happen in the // callback invocation above if the callback invokes // `ControlStackFrame::pop_abi_results` (e.g. `br` instruction). // // After an unconditional jump, the compiler will enter in an // unreachable state; instead of immediately truncating the value stack // to the expected length of the destination branch, we let the // reachability analysis code decide what should happen with the length // of the value stack once reachability is actually restored. At that // point, the right stack pointer offset will also be restored, which // should match the contents of the value stack. masm.ensure_sp_for_jump(target_offset)?; dest.set_as_target(); masm.jmp(*dest.label())?; self.reachable = false; Ok(()) } /// Push the ABI representation of the results stack. pub fn push_abi_results( &mut self, results: &ABIResults, masm: &mut M, mut calculate_ret_area: F, ) -> Result<()> where M: MacroAssembler, F: FnMut(&ABIResults, &mut CodeGenContext, &mut M) -> Option, { let area = results .on_stack() .then(|| calculate_ret_area(&results, self, masm).unwrap()); for operand in results.operands().iter() { match operand { ABIOperand::Reg { reg, ty, .. } => { ensure!( self.regalloc.reg_available(*reg), CodeGenError::expected_register_to_be_available(), ); let typed_reg = TypedReg::new(*ty, self.reg(*reg, masm)?); self.stack.push(typed_reg.into()); } ABIOperand::Stack { ty, offset, size } => match area.unwrap() { RetArea::SP(sp_offset) => { let slot = StackSlot::new(SPOffset::from_u32(sp_offset.as_u32() - offset), *size); self.stack.push(Val::mem(*ty, slot)); } // This function is only expected to be called when dealing // with control flow and when calling functions; as a // callee, only [Self::pop_abi_results] is needed when // finalizing the function compilation. _ => bail!(CodeGenError::unexpected_function_call()), }, } } Ok(()) } /// Truncates the value stack to the specified target. /// This function is intended to only be used when restoring the code /// generation's reachability state, when handling an unreachable end or /// else. pub fn truncate_stack_to(&mut self, target: usize) -> Result<()> { if self.stack.len() > target { self.drop_last(self.stack.len() - target, |regalloc, val| match val { Val::Reg(tr) => Ok(regalloc.free(tr.reg)), _ => Ok(()), }) } else { Ok(()) } } /// Load the [VMContext] pointer into the designated pinned register. pub fn load_vmctx(&mut self, masm: &mut M) -> Result<()> where M: MacroAssembler, { let addr = masm.local_address(&self.frame.vmctx_slot())?; masm.load_ptr(addr, writable!(vmctx!(M))) } /// Spill locals and registers to memory. // TODO: optimize the spill range; // At any point in the program, the stack might already contain memory // entries; we could effectively ignore that range; only focusing on the // range that contains spillable values. fn spill_impl( stack: &mut Stack, regalloc: &mut RegAlloc, frame: &Frame, masm: &mut M, ) -> Result<()> { for v in stack.inner_mut() { match v { Val::Reg(r) => { let slot = masm.push(r.reg, r.ty.try_into()?)?; regalloc.free(r.reg); *v = Val::mem(r.ty, slot); } Val::Local(local) => { let slot = frame.get_wasm_local(local.index); let addr = masm.local_address(&slot)?; let scratch = scratch!(M, &slot.ty); masm.load(addr, writable!(scratch), slot.ty.try_into()?)?; let stack_slot = masm.push(scratch, slot.ty.try_into()?)?; *v = Val::mem(slot.ty, stack_slot); } _ => {} } } Ok(()) } /// Prepares for emitting a binary operation where four 64-bit operands are /// used to produce two 64-bit operands, e.g. a 128-bit binop. pub fn binop128(&mut self, masm: &mut M, emit: F) -> Result<()> where F: FnOnce(&mut M, Reg, Reg, Reg, Reg) -> Result<(TypedReg, TypedReg)>, M: MacroAssembler, { let rhs_hi = self.pop_to_reg(masm, None)?; let rhs_lo = self.pop_to_reg(masm, None)?; let lhs_hi = self.pop_to_reg(masm, None)?; let lhs_lo = self.pop_to_reg(masm, None)?; let (lo, hi) = emit(masm, lhs_lo.reg, lhs_hi.reg, rhs_lo.reg, rhs_hi.reg)?; self.free_reg(rhs_hi); self.free_reg(rhs_lo); self.stack.push(lo.into()); self.stack.push(hi.into()); Ok(()) } /// Pops a register from the stack and then immediately frees it. Used to /// discard values from the last operation, for example. pub fn pop_and_free(&mut self, masm: &mut M) -> Result<()> { let reg = self.pop_to_reg(masm, None)?; self.free_reg(reg.reg); Ok(()) } }