Skip to content

Commit

Permalink
Clean-up Instruction::update_branch_offset extension method (#1230)
Browse files Browse the repository at this point in the history 
* add ComparatorExt::from_cmp_branch_instruction

* clean-up Instruction::update_branch_offset
  • Loading branch information
@Robbepop
Robbepop authored Oct 7, 2024
1 parent 66075b9 commit 8f76eef
Show file tree
Hide file tree
Showing 2 changed files with 162 additions and 88 deletions.
51 changes: 50 additions & 1 deletion crates/wasmi/src/engine/translator/comparator.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,12 @@ use crate::ir::{self, BranchOffset16, Comparator, Const16, Instruction, Reg};

/// Extensional functionality for [`Comparator`].
pub trait ComparatorExt: Sized {
/// Creates a [`Comparator`] from an [`Instruction`].
/// Creates a [`Comparator`] from a comparison [`Instruction`].
fn from_cmp_instruction(instr: Instruction) -> Option<Self>;

/// Creates a [`Comparator`] from a fused compare+branch [`Instruction`].
fn from_cmp_branch_instruction(instr: Instruction) -> Option<Self>;

/// Returns the negated version of `self` if possible.
///
/// # Note
Expand Down Expand Up @@ -64,6 +67,52 @@ impl ComparatorExt for Comparator {
Some(cmp)
}

fn from_cmp_branch_instruction(instr: Instruction) -> Option<Self> {
use Instruction as I;
let cmp = match instr {
I::BranchI32And { .. } | I::BranchI32AndImm { .. } => Self::I32And,
I::BranchI32Or { .. } | I::BranchI32OrImm { .. } => Self::I32Or,
I::BranchI32Xor { .. } | I::BranchI32XorImm { .. } => Self::I32Xor,
I::BranchI32AndEqz { .. } | I::BranchI32AndEqzImm { .. } => Self::I32AndEqz,
I::BranchI32OrEqz { .. } | I::BranchI32OrEqzImm { .. } => Self::I32OrEqz,
I::BranchI32XorEqz { .. } | I::BranchI32XorEqzImm { .. } => Self::I32XorEqz,
I::BranchI32Eq { .. } | I::BranchI32EqImm { .. } => Self::I32Eq,
I::BranchI32Ne { .. } | I::BranchI32NeImm { .. } => Self::I32Ne,
I::BranchI32LtS { .. } | I::BranchI32LtSImm { .. } => Self::I32LtS,
I::BranchI32LtU { .. } | I::BranchI32LtUImm { .. } => Self::I32LtU,
I::BranchI32LeS { .. } | I::BranchI32LeSImm { .. } => Self::I32LeS,
I::BranchI32LeU { .. } | I::BranchI32LeUImm { .. } => Self::I32LeU,
I::BranchI32GtS { .. } | I::BranchI32GtSImm { .. } => Self::I32GtS,
I::BranchI32GtU { .. } | I::BranchI32GtUImm { .. } => Self::I32GtU,
I::BranchI32GeS { .. } | I::BranchI32GeSImm { .. } => Self::I32GeS,
I::BranchI32GeU { .. } | I::BranchI32GeUImm { .. } => Self::I32GeU,
I::BranchI64Eq { .. } | I::BranchI64EqImm { .. } => Self::I64Eq,
I::BranchI64Ne { .. } | I::BranchI64NeImm { .. } => Self::I64Ne,
I::BranchI64LtS { .. } | I::BranchI64LtSImm { .. } => Self::I64LtS,
I::BranchI64LtU { .. } | I::BranchI64LtUImm { .. } => Self::I64LtU,
I::BranchI64LeS { .. } | I::BranchI64LeSImm { .. } => Self::I64LeS,
I::BranchI64LeU { .. } | I::BranchI64LeUImm { .. } => Self::I64LeU,
I::BranchI64GtS { .. } | I::BranchI64GtSImm { .. } => Self::I64GtS,
I::BranchI64GtU { .. } | I::BranchI64GtUImm { .. } => Self::I64GtU,
I::BranchI64GeS { .. } | I::BranchI64GeSImm { .. } => Self::I64GeS,
I::BranchI64GeU { .. } | I::BranchI64GeUImm { .. } => Self::I64GeU,
I::BranchF32Eq { .. } => Self::F32Eq,
I::BranchF32Ne { .. } => Self::F32Ne,
I::BranchF32Lt { .. } => Self::F32Lt,
I::BranchF32Le { .. } => Self::F32Le,
I::BranchF32Gt { .. } => Self::F32Gt,
I::BranchF32Ge { .. } => Self::F32Ge,
I::BranchF64Eq { .. } => Self::F64Eq,
I::BranchF64Ne { .. } => Self::F64Ne,
I::BranchF64Lt { .. } => Self::F64Lt,
I::BranchF64Le { .. } => Self::F64Le,
I::BranchF64Gt { .. } => Self::F64Gt,
I::BranchF64Ge { .. } => Self::F64Ge,
_ => return None,
};
Some(cmp)
}

fn negate(self) -> Option<Self> {
let negated = match self {
Self::I32And => Self::I32AndEqz,
Expand Down
199 changes: 112 additions & 87 deletions crates/wasmi/src/engine/translator/instr_encoder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1310,105 +1310,130 @@ trait UpdateBranchOffset {
impl UpdateBranchOffset for Instruction {
#[rustfmt::skip]
fn update_branch_offset(&mut self, stack: &mut ValueStack, new_offset: BranchOffset) -> Result<(), Error> {
/// Initializes the 16-bit offset of `instr` if possible.
/// Updates the [`BranchOffset16`] to `new_offset` if possible.
///
/// If `new_offset` cannot be encoded as 16-bit offset `self` is replaced with a fallback instruction.
macro_rules! init_offset {
($lhs:expr, $rhs:expr, $offset:expr, $new_offset:expr, $cmp:expr) => {{
if let Err(_) = $offset.init($new_offset) {
let params = stack.alloc_const(ComparatorAndOffset::new($cmp, $new_offset))?;
*self = Instruction::branch_cmp_fallback(*$lhs, *$rhs, params);
}
Ok(())
}}
}

macro_rules! init_offset_imm {
($ty:ty, $lhs:expr, $rhs:expr, $offset:expr, $new_offset:expr, $cmp:expr) => {{
if let Err(_) = $offset.init($new_offset) {
let rhs = stack.alloc_const(<$ty>::from(*$rhs))?;
let params = stack.alloc_const(ComparatorAndOffset::new($cmp, $new_offset))?;
*self = Instruction::branch_cmp_fallback(*$lhs, rhs, params);
/// Otherwise returns `Some` fallback `Instruction` that replaces the outer `self`.
fn init_offset_imm<T>(
stack: &mut ValueStack,
lhs: Reg,
rhs: Const16<T>,
offset: &mut BranchOffset16,
new_offset: BranchOffset,
cmp: Comparator,
) -> Result<Option<Instruction>, Error>
where
T: From<Const16<T>> + Into<UntypedVal>,
{
match offset.init(new_offset) {
Ok(_) => Ok(None),
Err(_) => {
let rhs = stack.alloc_const(<T>::from(rhs))?;
let params = stack.alloc_const(ComparatorAndOffset::new(cmp, new_offset))?;
Ok(Some(Instruction::branch_cmp_fallback(lhs, rhs, params)))
}
Ok(())
}};
}
}

use Instruction as I;
use Comparator as Cmp;
match self {
Instruction::Branch { offset } |
Instruction::BranchTableTarget { offset, .. } |
Instruction::BranchTableTargetNonOverlapping { offset, .. } => {
offset.init(new_offset);
Ok(())
return Ok(())
}
_ => {}
};
let Some(comparator) = Comparator::from_cmp_branch_instruction(*self) else {
panic!("expected a Wasmi branch+cmp instruction but found: {:?}", *self)
};
let update = match self {
I::BranchI32And { lhs, rhs, offset } |
I::BranchI32Or { lhs, rhs, offset } |
I::BranchI32Xor { lhs, rhs, offset } |
I::BranchI32AndEqz { lhs, rhs, offset } |
I::BranchI32OrEqz { lhs, rhs, offset } |
I::BranchI32XorEqz { lhs, rhs, offset } |
I::BranchI32Eq { lhs, rhs, offset } |
I::BranchI32Ne { lhs, rhs, offset } |
I::BranchI32LtS { lhs, rhs, offset } |
I::BranchI32LtU { lhs, rhs, offset } |
I::BranchI32LeS { lhs, rhs, offset } |
I::BranchI32LeU { lhs, rhs, offset } |
I::BranchI32GtS { lhs, rhs, offset } |
I::BranchI32GtU { lhs, rhs, offset } |
I::BranchI32GeS { lhs, rhs, offset } |
I::BranchI32GeU { lhs, rhs, offset } |
I::BranchI64Eq { lhs, rhs, offset } |
I::BranchI64Ne { lhs, rhs, offset } |
I::BranchI64LtS { lhs, rhs, offset } |
I::BranchI64LtU { lhs, rhs, offset } |
I::BranchI64LeS { lhs, rhs, offset } |
I::BranchI64LeU { lhs, rhs, offset } |
I::BranchI64GtS { lhs, rhs, offset } |
I::BranchI64GtU { lhs, rhs, offset } |
I::BranchI64GeS { lhs, rhs, offset } |
I::BranchI64GeU { lhs, rhs, offset } |
I::BranchF32Eq { lhs, rhs, offset } |
I::BranchF32Ne { lhs, rhs, offset } |
I::BranchF32Lt { lhs, rhs, offset } |
I::BranchF32Le { lhs, rhs, offset } |
I::BranchF32Gt { lhs, rhs, offset } |
I::BranchF32Ge { lhs, rhs, offset } |
I::BranchF64Eq { lhs, rhs, offset } |
I::BranchF64Ne { lhs, rhs, offset } |
I::BranchF64Lt { lhs, rhs, offset } |
I::BranchF64Le { lhs, rhs, offset } |
I::BranchF64Gt { lhs, rhs, offset } |
I::BranchF64Ge { lhs, rhs, offset } => {
match offset.init(new_offset) {
Ok(_) => None,
Err(_) => {
let params = stack.alloc_const(ComparatorAndOffset::new(comparator, new_offset))?;
Some(Instruction::branch_cmp_fallback(*lhs, *rhs, params))
}
}
}
I::BranchI32AndImm { lhs, rhs, offset } |
I::BranchI32OrImm { lhs, rhs, offset } |
I::BranchI32XorImm { lhs, rhs, offset } |
I::BranchI32AndEqzImm { lhs, rhs, offset } |
I::BranchI32OrEqzImm { lhs, rhs, offset } |
I::BranchI32XorEqzImm { lhs, rhs, offset } |
I::BranchI32EqImm { lhs, rhs, offset } |
I::BranchI32NeImm { lhs, rhs, offset } |
I::BranchI32LtSImm { lhs, rhs, offset } |
I::BranchI32LeSImm { lhs, rhs, offset } |
I::BranchI32GtSImm { lhs, rhs, offset } |
I::BranchI32GeSImm { lhs, rhs, offset } => {
init_offset_imm::<i32>(stack, *lhs, *rhs, offset, new_offset, comparator)?
}
I::BranchI32LtUImm { lhs, rhs, offset } |
I::BranchI32LeUImm { lhs, rhs, offset } |
I::BranchI32GtUImm { lhs, rhs, offset } |
I::BranchI32GeUImm { lhs, rhs, offset } => {
init_offset_imm::<u32>(stack, *lhs, *rhs, offset, new_offset, comparator)?
}
I::BranchI32And { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32And),
I::BranchI32Or { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32Or),
I::BranchI32Xor { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32Xor),
I::BranchI32AndEqz { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32AndEqz),
I::BranchI32OrEqz { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32OrEqz),
I::BranchI32XorEqz { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32XorEqz),
I::BranchI32Eq { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32Eq),
I::BranchI32Ne { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32Ne),
I::BranchI32LtS { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32LtS),
I::BranchI32LtU { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32LtU),
I::BranchI32LeS { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32LeS),
I::BranchI32LeU { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32LeU),
I::BranchI32GtS { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32GtS),
I::BranchI32GtU { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32GtU),
I::BranchI32GeS { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32GeS),
I::BranchI32GeU { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32GeU),
I::BranchI64Eq { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64Eq),
I::BranchI64Ne { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64Ne),
I::BranchI64LtS { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64LtS),
I::BranchI64LtU { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64LtU),
I::BranchI64LeS { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64LeS),
I::BranchI64LeU { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64LeU),
I::BranchI64GtS { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64GtS),
I::BranchI64GtU { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64GtU),
I::BranchI64GeS { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64GeS),
I::BranchI64GeU { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64GeU),
I::BranchF32Eq { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F32Eq),
I::BranchF32Ne { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F32Ne),
I::BranchF32Lt { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F32Lt),
I::BranchF32Le { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F32Le),
I::BranchF32Gt { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F32Gt),
I::BranchF32Ge { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F32Ge),
I::BranchF64Eq { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F64Eq),
I::BranchF64Ne { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F64Ne),
I::BranchF64Lt { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F64Lt),
I::BranchF64Le { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F64Le),
I::BranchF64Gt { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F64Gt),
I::BranchF64Ge { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F64Ge),
I::BranchI32AndImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32And),
I::BranchI32OrImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32Or),
I::BranchI32XorImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32Xor),
I::BranchI32AndEqzImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32AndEqz),
I::BranchI32OrEqzImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32OrEqz),
I::BranchI32XorEqzImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32XorEqz),
I::BranchI32EqImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32Eq),
I::BranchI32NeImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32Ne),
I::BranchI32LtSImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32LtS),
I::BranchI32LeSImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32LeS),
I::BranchI32GtSImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32GtS),
I::BranchI32GeSImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32GeS),
I::BranchI32LtUImm { lhs, rhs, offset } => init_offset_imm!(u32, lhs, rhs, offset, new_offset, Cmp::I32LtU),
I::BranchI32LeUImm { lhs, rhs, offset } => init_offset_imm!(u32, lhs, rhs, offset, new_offset, Cmp::I32LeU),
I::BranchI32GtUImm { lhs, rhs, offset } => init_offset_imm!(u32, lhs, rhs, offset, new_offset, Cmp::I32GtU),
I::BranchI32GeUImm { lhs, rhs, offset } => init_offset_imm!(u32, lhs, rhs, offset, new_offset, Cmp::I32GeU),
I::BranchI64EqImm { lhs, rhs, offset } => init_offset_imm!(i64, lhs, rhs, offset, new_offset, Cmp::I64Eq),
I::BranchI64NeImm { lhs, rhs, offset } => init_offset_imm!(i64, lhs, rhs, offset, new_offset, Cmp::I64Ne),
I::BranchI64LtSImm { lhs, rhs, offset } => init_offset_imm!(i64, lhs, rhs, offset, new_offset, Cmp::I64LtS),
I::BranchI64LeSImm { lhs, rhs, offset } => init_offset_imm!(i64, lhs, rhs, offset, new_offset, Cmp::I64LeS),
I::BranchI64GtSImm { lhs, rhs, offset } => init_offset_imm!(i64, lhs, rhs, offset, new_offset, Cmp::I64GtS),
I::BranchI64GeSImm { lhs, rhs, offset } => init_offset_imm!(i64, lhs, rhs, offset, new_offset, Cmp::I64GeS),
I::BranchI64LtUImm { lhs, rhs, offset } => init_offset_imm!(u64, lhs, rhs, offset, new_offset, Cmp::I64LtU),
I::BranchI64LeUImm { lhs, rhs, offset } => init_offset_imm!(u64, lhs, rhs, offset, new_offset, Cmp::I64LeU),
I::BranchI64GtUImm { lhs, rhs, offset } => init_offset_imm!(u64, lhs, rhs, offset, new_offset, Cmp::I64GtU),
I::BranchI64GeUImm { lhs, rhs, offset } => init_offset_imm!(u64, lhs, rhs, offset, new_offset, Cmp::I64GeU),
_ => panic!("tried to update branch offset of a non-branch instruction: {self:?}"),
I::BranchI64EqImm { lhs, rhs, offset } |
I::BranchI64NeImm { lhs, rhs, offset } |
I::BranchI64LtSImm { lhs, rhs, offset } |
I::BranchI64LeSImm { lhs, rhs, offset } |
I::BranchI64GtSImm { lhs, rhs, offset } |
I::BranchI64GeSImm { lhs, rhs, offset } => {
init_offset_imm::<i64>(stack, *lhs, *rhs, offset, new_offset, comparator)?
}
I::BranchI64LtUImm { lhs, rhs, offset } |
I::BranchI64LeUImm { lhs, rhs, offset } |
I::BranchI64GtUImm { lhs, rhs, offset } |
I::BranchI64GeUImm { lhs, rhs, offset } => {
init_offset_imm::<u64>(stack, *lhs, *rhs, offset, new_offset, comparator)?
}
_ => panic!("expected a Wasmi branch+cmp instruction but found: {:?}", *self),
};
if let Some(update) = update {
*self = update;
}
Ok(())
}
}

Expand Down

0 comments on commit 8f76eef

Please sign in to comment.