Skip to content

Commit 8f76eef

Browse files
RobbepopRobbepop
authored
Clean-up Instruction::update_branch_offset extension method (#1230)
* add ComparatorExt::from_cmp_branch_instruction * clean-up Instruction::update_branch_offset
1 parent 66075b9 commit 8f76eef

File tree

2 files changed

+162
-88
lines changed

2 files changed

+162
-88
lines changed

‎crates/wasmi/src/engine/translator/comparator.rs‎

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,12 @@ use crate::ir::{self, BranchOffset16, Comparator, Const16, Instruction, Reg};
22

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

8+
/// Creates a [`Comparator`] from a fused compare+branch [`Instruction`].
9+
fn from_cmp_branch_instruction(instr: Instruction) -> Option<Self>;
10+
811
/// Returns the negated version of `self` if possible.
912
///
1013
/// # Note
@@ -64,6 +67,52 @@ impl ComparatorExt for Comparator {
6467
Some(cmp)
6568
}
6669

70+
fn from_cmp_branch_instruction(instr: Instruction) -> Option<Self> {
71+
use Instruction as I;
72+
let cmp = match instr {
73+
I::BranchI32And { .. } | I::BranchI32AndImm { .. } => Self::I32And,
74+
I::BranchI32Or { .. } | I::BranchI32OrImm { .. } => Self::I32Or,
75+
I::BranchI32Xor { .. } | I::BranchI32XorImm { .. } => Self::I32Xor,
76+
I::BranchI32AndEqz { .. } | I::BranchI32AndEqzImm { .. } => Self::I32AndEqz,
77+
I::BranchI32OrEqz { .. } | I::BranchI32OrEqzImm { .. } => Self::I32OrEqz,
78+
I::BranchI32XorEqz { .. } | I::BranchI32XorEqzImm { .. } => Self::I32XorEqz,
79+
I::BranchI32Eq { .. } | I::BranchI32EqImm { .. } => Self::I32Eq,
80+
I::BranchI32Ne { .. } | I::BranchI32NeImm { .. } => Self::I32Ne,
81+
I::BranchI32LtS { .. } | I::BranchI32LtSImm { .. } => Self::I32LtS,
82+
I::BranchI32LtU { .. } | I::BranchI32LtUImm { .. } => Self::I32LtU,
83+
I::BranchI32LeS { .. } | I::BranchI32LeSImm { .. } => Self::I32LeS,
84+
I::BranchI32LeU { .. } | I::BranchI32LeUImm { .. } => Self::I32LeU,
85+
I::BranchI32GtS { .. } | I::BranchI32GtSImm { .. } => Self::I32GtS,
86+
I::BranchI32GtU { .. } | I::BranchI32GtUImm { .. } => Self::I32GtU,
87+
I::BranchI32GeS { .. } | I::BranchI32GeSImm { .. } => Self::I32GeS,
88+
I::BranchI32GeU { .. } | I::BranchI32GeUImm { .. } => Self::I32GeU,
89+
I::BranchI64Eq { .. } | I::BranchI64EqImm { .. } => Self::I64Eq,
90+
I::BranchI64Ne { .. } | I::BranchI64NeImm { .. } => Self::I64Ne,
91+
I::BranchI64LtS { .. } | I::BranchI64LtSImm { .. } => Self::I64LtS,
92+
I::BranchI64LtU { .. } | I::BranchI64LtUImm { .. } => Self::I64LtU,
93+
I::BranchI64LeS { .. } | I::BranchI64LeSImm { .. } => Self::I64LeS,
94+
I::BranchI64LeU { .. } | I::BranchI64LeUImm { .. } => Self::I64LeU,
95+
I::BranchI64GtS { .. } | I::BranchI64GtSImm { .. } => Self::I64GtS,
96+
I::BranchI64GtU { .. } | I::BranchI64GtUImm { .. } => Self::I64GtU,
97+
I::BranchI64GeS { .. } | I::BranchI64GeSImm { .. } => Self::I64GeS,
98+
I::BranchI64GeU { .. } | I::BranchI64GeUImm { .. } => Self::I64GeU,
99+
I::BranchF32Eq { .. } => Self::F32Eq,
100+
I::BranchF32Ne { .. } => Self::F32Ne,
101+
I::BranchF32Lt { .. } => Self::F32Lt,
102+
I::BranchF32Le { .. } => Self::F32Le,
103+
I::BranchF32Gt { .. } => Self::F32Gt,
104+
I::BranchF32Ge { .. } => Self::F32Ge,
105+
I::BranchF64Eq { .. } => Self::F64Eq,
106+
I::BranchF64Ne { .. } => Self::F64Ne,
107+
I::BranchF64Lt { .. } => Self::F64Lt,
108+
I::BranchF64Le { .. } => Self::F64Le,
109+
I::BranchF64Gt { .. } => Self::F64Gt,
110+
I::BranchF64Ge { .. } => Self::F64Ge,
111+
_ => return None,
112+
};
113+
Some(cmp)
114+
}
115+
67116
fn negate(self) -> Option<Self> {
68117
let negated = match self {
69118
Self::I32And => Self::I32AndEqz,

‎crates/wasmi/src/engine/translator/instr_encoder.rs‎

Lines changed: 112 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -1310,105 +1310,130 @@ trait UpdateBranchOffset {
13101310
impl UpdateBranchOffset for Instruction {
13111311
#[rustfmt::skip]
13121312
fn update_branch_offset(&mut self, stack: &mut ValueStack, new_offset: BranchOffset) -> Result<(), Error> {
1313-
/// Initializes the 16-bit offset of `instr` if possible.
1313+
/// Updates the [`BranchOffset16`] to `new_offset` if possible.
13141314
///
1315-
/// If `new_offset` cannot be encoded as 16-bit offset `self` is replaced with a fallback instruction.
1316-
macro_rules! init_offset {
1317-
($lhs:expr, $rhs:expr, $offset:expr, $new_offset:expr, $cmp:expr) => {{
1318-
if let Err(_) = $offset.init($new_offset) {
1319-
let params = stack.alloc_const(ComparatorAndOffset::new($cmp, $new_offset))?;
1320-
*self = Instruction::branch_cmp_fallback(*$lhs, *$rhs, params);
1321-
}
1322-
Ok(())
1323-
}}
1324-
}
1325-
1326-
macro_rules! init_offset_imm {
1327-
($ty:ty, $lhs:expr, $rhs:expr, $offset:expr, $new_offset:expr, $cmp:expr) => {{
1328-
if let Err(_) = $offset.init($new_offset) {
1329-
let rhs = stack.alloc_const(<$ty>::from(*$rhs))?;
1330-
let params = stack.alloc_const(ComparatorAndOffset::new($cmp, $new_offset))?;
1331-
*self = Instruction::branch_cmp_fallback(*$lhs, rhs, params);
1315+
/// Otherwise returns `Some` fallback `Instruction` that replaces the outer `self`.
1316+
fn init_offset_imm<T>(
1317+
stack: &mut ValueStack,
1318+
lhs: Reg,
1319+
rhs: Const16<T>,
1320+
offset: &mut BranchOffset16,
1321+
new_offset: BranchOffset,
1322+
cmp: Comparator,
1323+
) -> Result<Option<Instruction>, Error>
1324+
where
1325+
T: From<Const16<T>> + Into<UntypedVal>,
1326+
{
1327+
match offset.init(new_offset) {
1328+
Ok(_) => Ok(None),
1329+
Err(_) => {
1330+
let rhs = stack.alloc_const(<T>::from(rhs))?;
1331+
let params = stack.alloc_const(ComparatorAndOffset::new(cmp, new_offset))?;
1332+
Ok(Some(Instruction::branch_cmp_fallback(lhs, rhs, params)))
13321333
}
1333-
Ok(())
1334-
}};
1334+
}
13351335
}
13361336

13371337
use Instruction as I;
1338-
use Comparator as Cmp;
13391338
match self {
13401339
Instruction::Branch { offset } |
13411340
Instruction::BranchTableTarget { offset, .. } |
13421341
Instruction::BranchTableTargetNonOverlapping { offset, .. } => {
13431342
offset.init(new_offset);
1344-
Ok(())
1343+
return Ok(())
1344+
}
1345+
_ => {}
1346+
};
1347+
let Some(comparator) = Comparator::from_cmp_branch_instruction(*self) else {
1348+
panic!("expected a Wasmi branch+cmp instruction but found: {:?}", *self)
1349+
};
1350+
let update = match self {
1351+
I::BranchI32And { lhs, rhs, offset } |
1352+
I::BranchI32Or { lhs, rhs, offset } |
1353+
I::BranchI32Xor { lhs, rhs, offset } |
1354+
I::BranchI32AndEqz { lhs, rhs, offset } |
1355+
I::BranchI32OrEqz { lhs, rhs, offset } |
1356+
I::BranchI32XorEqz { lhs, rhs, offset } |
1357+
I::BranchI32Eq { lhs, rhs, offset } |
1358+
I::BranchI32Ne { lhs, rhs, offset } |
1359+
I::BranchI32LtS { lhs, rhs, offset } |
1360+
I::BranchI32LtU { lhs, rhs, offset } |
1361+
I::BranchI32LeS { lhs, rhs, offset } |
1362+
I::BranchI32LeU { lhs, rhs, offset } |
1363+
I::BranchI32GtS { lhs, rhs, offset } |
1364+
I::BranchI32GtU { lhs, rhs, offset } |
1365+
I::BranchI32GeS { lhs, rhs, offset } |
1366+
I::BranchI32GeU { lhs, rhs, offset } |
1367+
I::BranchI64Eq { lhs, rhs, offset } |
1368+
I::BranchI64Ne { lhs, rhs, offset } |
1369+
I::BranchI64LtS { lhs, rhs, offset } |
1370+
I::BranchI64LtU { lhs, rhs, offset } |
1371+
I::BranchI64LeS { lhs, rhs, offset } |
1372+
I::BranchI64LeU { lhs, rhs, offset } |
1373+
I::BranchI64GtS { lhs, rhs, offset } |
1374+
I::BranchI64GtU { lhs, rhs, offset } |
1375+
I::BranchI64GeS { lhs, rhs, offset } |
1376+
I::BranchI64GeU { lhs, rhs, offset } |
1377+
I::BranchF32Eq { lhs, rhs, offset } |
1378+
I::BranchF32Ne { lhs, rhs, offset } |
1379+
I::BranchF32Lt { lhs, rhs, offset } |
1380+
I::BranchF32Le { lhs, rhs, offset } |
1381+
I::BranchF32Gt { lhs, rhs, offset } |
1382+
I::BranchF32Ge { lhs, rhs, offset } |
1383+
I::BranchF64Eq { lhs, rhs, offset } |
1384+
I::BranchF64Ne { lhs, rhs, offset } |
1385+
I::BranchF64Lt { lhs, rhs, offset } |
1386+
I::BranchF64Le { lhs, rhs, offset } |
1387+
I::BranchF64Gt { lhs, rhs, offset } |
1388+
I::BranchF64Ge { lhs, rhs, offset } => {
1389+
match offset.init(new_offset) {
1390+
Ok(_) => None,
1391+
Err(_) => {
1392+
let params = stack.alloc_const(ComparatorAndOffset::new(comparator, new_offset))?;
1393+
Some(Instruction::branch_cmp_fallback(*lhs, *rhs, params))
1394+
}
1395+
}
1396+
}
1397+
I::BranchI32AndImm { lhs, rhs, offset } |
1398+
I::BranchI32OrImm { lhs, rhs, offset } |
1399+
I::BranchI32XorImm { lhs, rhs, offset } |
1400+
I::BranchI32AndEqzImm { lhs, rhs, offset } |
1401+
I::BranchI32OrEqzImm { lhs, rhs, offset } |
1402+
I::BranchI32XorEqzImm { lhs, rhs, offset } |
1403+
I::BranchI32EqImm { lhs, rhs, offset } |
1404+
I::BranchI32NeImm { lhs, rhs, offset } |
1405+
I::BranchI32LtSImm { lhs, rhs, offset } |
1406+
I::BranchI32LeSImm { lhs, rhs, offset } |
1407+
I::BranchI32GtSImm { lhs, rhs, offset } |
1408+
I::BranchI32GeSImm { lhs, rhs, offset } => {
1409+
init_offset_imm::<i32>(stack, *lhs, *rhs, offset, new_offset, comparator)?
1410+
}
1411+
I::BranchI32LtUImm { lhs, rhs, offset } |
1412+
I::BranchI32LeUImm { lhs, rhs, offset } |
1413+
I::BranchI32GtUImm { lhs, rhs, offset } |
1414+
I::BranchI32GeUImm { lhs, rhs, offset } => {
1415+
init_offset_imm::<u32>(stack, *lhs, *rhs, offset, new_offset, comparator)?
13451416
}
1346-
I::BranchI32And { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32And),
1347-
I::BranchI32Or { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32Or),
1348-
I::BranchI32Xor { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32Xor),
1349-
I::BranchI32AndEqz { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32AndEqz),
1350-
I::BranchI32OrEqz { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32OrEqz),
1351-
I::BranchI32XorEqz { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32XorEqz),
1352-
I::BranchI32Eq { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32Eq),
1353-
I::BranchI32Ne { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32Ne),
1354-
I::BranchI32LtS { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32LtS),
1355-
I::BranchI32LtU { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32LtU),
1356-
I::BranchI32LeS { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32LeS),
1357-
I::BranchI32LeU { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32LeU),
1358-
I::BranchI32GtS { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32GtS),
1359-
I::BranchI32GtU { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32GtU),
1360-
I::BranchI32GeS { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32GeS),
1361-
I::BranchI32GeU { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I32GeU),
1362-
I::BranchI64Eq { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64Eq),
1363-
I::BranchI64Ne { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64Ne),
1364-
I::BranchI64LtS { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64LtS),
1365-
I::BranchI64LtU { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64LtU),
1366-
I::BranchI64LeS { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64LeS),
1367-
I::BranchI64LeU { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64LeU),
1368-
I::BranchI64GtS { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64GtS),
1369-
I::BranchI64GtU { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64GtU),
1370-
I::BranchI64GeS { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64GeS),
1371-
I::BranchI64GeU { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::I64GeU),
1372-
I::BranchF32Eq { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F32Eq),
1373-
I::BranchF32Ne { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F32Ne),
1374-
I::BranchF32Lt { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F32Lt),
1375-
I::BranchF32Le { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F32Le),
1376-
I::BranchF32Gt { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F32Gt),
1377-
I::BranchF32Ge { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F32Ge),
1378-
I::BranchF64Eq { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F64Eq),
1379-
I::BranchF64Ne { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F64Ne),
1380-
I::BranchF64Lt { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F64Lt),
1381-
I::BranchF64Le { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F64Le),
1382-
I::BranchF64Gt { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F64Gt),
1383-
I::BranchF64Ge { lhs, rhs, offset } => init_offset!(lhs, rhs, offset, new_offset, Cmp::F64Ge),
1384-
I::BranchI32AndImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32And),
1385-
I::BranchI32OrImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32Or),
1386-
I::BranchI32XorImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32Xor),
1387-
I::BranchI32AndEqzImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32AndEqz),
1388-
I::BranchI32OrEqzImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32OrEqz),
1389-
I::BranchI32XorEqzImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32XorEqz),
1390-
I::BranchI32EqImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32Eq),
1391-
I::BranchI32NeImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32Ne),
1392-
I::BranchI32LtSImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32LtS),
1393-
I::BranchI32LeSImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32LeS),
1394-
I::BranchI32GtSImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32GtS),
1395-
I::BranchI32GeSImm { lhs, rhs, offset } => init_offset_imm!(i32, lhs, rhs, offset, new_offset, Cmp::I32GeS),
1396-
I::BranchI32LtUImm { lhs, rhs, offset } => init_offset_imm!(u32, lhs, rhs, offset, new_offset, Cmp::I32LtU),
1397-
I::BranchI32LeUImm { lhs, rhs, offset } => init_offset_imm!(u32, lhs, rhs, offset, new_offset, Cmp::I32LeU),
1398-
I::BranchI32GtUImm { lhs, rhs, offset } => init_offset_imm!(u32, lhs, rhs, offset, new_offset, Cmp::I32GtU),
1399-
I::BranchI32GeUImm { lhs, rhs, offset } => init_offset_imm!(u32, lhs, rhs, offset, new_offset, Cmp::I32GeU),
1400-
I::BranchI64EqImm { lhs, rhs, offset } => init_offset_imm!(i64, lhs, rhs, offset, new_offset, Cmp::I64Eq),
1401-
I::BranchI64NeImm { lhs, rhs, offset } => init_offset_imm!(i64, lhs, rhs, offset, new_offset, Cmp::I64Ne),
1402-
I::BranchI64LtSImm { lhs, rhs, offset } => init_offset_imm!(i64, lhs, rhs, offset, new_offset, Cmp::I64LtS),
1403-
I::BranchI64LeSImm { lhs, rhs, offset } => init_offset_imm!(i64, lhs, rhs, offset, new_offset, Cmp::I64LeS),
1404-
I::BranchI64GtSImm { lhs, rhs, offset } => init_offset_imm!(i64, lhs, rhs, offset, new_offset, Cmp::I64GtS),
1405-
I::BranchI64GeSImm { lhs, rhs, offset } => init_offset_imm!(i64, lhs, rhs, offset, new_offset, Cmp::I64GeS),
1406-
I::BranchI64LtUImm { lhs, rhs, offset } => init_offset_imm!(u64, lhs, rhs, offset, new_offset, Cmp::I64LtU),
1407-
I::BranchI64LeUImm { lhs, rhs, offset } => init_offset_imm!(u64, lhs, rhs, offset, new_offset, Cmp::I64LeU),
1408-
I::BranchI64GtUImm { lhs, rhs, offset } => init_offset_imm!(u64, lhs, rhs, offset, new_offset, Cmp::I64GtU),
1409-
I::BranchI64GeUImm { lhs, rhs, offset } => init_offset_imm!(u64, lhs, rhs, offset, new_offset, Cmp::I64GeU),
1410-
_ => panic!("tried to update branch offset of a non-branch instruction: {self:?}"),
1417+
I::BranchI64EqImm { lhs, rhs, offset } |
1418+
I::BranchI64NeImm { lhs, rhs, offset } |
1419+
I::BranchI64LtSImm { lhs, rhs, offset } |
1420+
I::BranchI64LeSImm { lhs, rhs, offset } |
1421+
I::BranchI64GtSImm { lhs, rhs, offset } |
1422+
I::BranchI64GeSImm { lhs, rhs, offset } => {
1423+
init_offset_imm::<i64>(stack, *lhs, *rhs, offset, new_offset, comparator)?
1424+
}
1425+
I::BranchI64LtUImm { lhs, rhs, offset } |
1426+
I::BranchI64LeUImm { lhs, rhs, offset } |
1427+
I::BranchI64GtUImm { lhs, rhs, offset } |
1428+
I::BranchI64GeUImm { lhs, rhs, offset } => {
1429+
init_offset_imm::<u64>(stack, *lhs, *rhs, offset, new_offset, comparator)?
1430+
}
1431+
_ => panic!("expected a Wasmi branch+cmp instruction but found: {:?}", *self),
1432+
};
1433+
if let Some(update) = update {
1434+
*self = update;
14111435
}
1436+
Ok(())
14121437
}
14131438
}
14141439

0 commit comments

Comments
 (0)