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(())
}
}