use crate::{instruction_categories as gen, Fee};
pub use config::Config;
pub use error::Error;
use wasmparser::{BlockType, BrTable, VisitOperator, VisitSimdOperator};
mod config;
mod error;
mod optimize;
#[derive(Clone, Copy, Debug)]
pub enum InstrumentationKind {
Pure,
Unreachable,
PreControlFlow,
PostControlFlow,
BetweenControlFlow,
MemoryGrow,
MemoryInit,
MemoryCopy,
MemoryFill,
TableInit,
TableGrow,
TableFill,
TableCopy,
}
#[derive(Debug)]
pub(crate) enum BranchTargetKind {
Loop(usize, usize),
UntakenForward,
Forward,
}
#[derive(Debug)]
pub(crate) struct Frame {
pub(crate) stack_polymorphic: bool,
pub(crate) kind: BranchTargetKind,
}
pub(crate) struct ScheduledInstrumentation {
cost: Fee,
kind: InstrumentationKind,
}
pub struct FunctionState {
pub(crate) offsets: Vec<usize>,
pub(crate) costs: Vec<Fee>,
pub(crate) kinds: Vec<InstrumentationKind>,
pub(crate) frame_stack: Vec<Frame>,
pub(crate) current_frame: Frame,
pub(crate) scheduled_instrumentation: Option<ScheduledInstrumentation>,
}
impl FunctionState {
pub fn new() -> Self {
Self {
offsets: vec![],
costs: vec![],
kinds: vec![],
frame_stack: vec![],
current_frame: Frame {
stack_polymorphic: false,
kind: BranchTargetKind::UntakenForward,
},
scheduled_instrumentation: None,
}
}
}
pub struct Visitor<'s, CostModel> {
pub(crate) offset: usize,
pub(crate) model: &'s mut CostModel,
pub(crate) state: &'s mut FunctionState,
}
impl<'a, 'b, CostModel> Visitor<'a, CostModel> {
fn visit_pure_instruction(&mut self, cost: Fee) {
self.push_instrumentation_before(InstrumentationKind::Pure, cost)
}
fn visit_side_effect_instruction(&mut self, cost: Fee) {
self.push_instrumentation_before(InstrumentationKind::PreControlFlow, cost);
self.push_instrumentation_after(InstrumentationKind::PostControlFlow, Fee::ZERO);
}
fn push_instrumentation_before(&mut self, kind: InstrumentationKind, cost: Fee) {
let kind = if self.state.current_frame.stack_polymorphic {
InstrumentationKind::Unreachable
} else {
kind
};
self.state.offsets.push(self.offset);
self.state.kinds.push(kind);
self.state.costs.push(cost);
}
fn push_instrumentation_after(&mut self, kind: InstrumentationKind, cost: Fee) {
assert!(self
.state
.scheduled_instrumentation
.replace(ScheduledInstrumentation { cost, kind })
.is_none());
}
fn new_frame(&mut self, kind: BranchTargetKind) {
let stack_polymorphic = self.state.current_frame.stack_polymorphic;
self.state.frame_stack.push(std::mem::replace(
&mut self.state.current_frame,
Frame {
stack_polymorphic,
kind,
},
));
}
fn end_frame(&mut self) {
if let Some(frame) = self.state.frame_stack.pop() {
self.state.current_frame = frame;
}
}
fn make_polymorphic(&mut self) {
self.state.current_frame.stack_polymorphic = true;
}
fn root_frame_index(&self) -> usize {
self.state.frame_stack.len()
}
fn frame_index(&self, relative_depth: u32) -> Result<usize, Error> {
usize::try_from(relative_depth).map_err(|_| Error::BranchDepthTooLarge(self.offset))
}
fn adjust_branch_target(&mut self, frame_index: usize) -> Result<(), Error> {
let frame = if let Some(frame_stack_index) = frame_index.checked_sub(1) {
self.state
.frame_stack
.iter_mut()
.nth_back(frame_stack_index)
.ok_or(Error::InvalidBrTarget(self.offset))?
} else {
&mut self.state.current_frame
};
match frame.kind {
BranchTargetKind::Forward => (),
BranchTargetKind::UntakenForward => frame.kind = BranchTargetKind::Forward,
BranchTargetKind::Loop(pre_index, post_index) => {
self.state.kinds[post_index] = InstrumentationKind::PostControlFlow;
self.state.kinds[pre_index] = InstrumentationKind::PreControlFlow;
}
}
Ok(())
}
fn visit_conditional_branch(&mut self, frame_index: usize, cost: Fee) -> Result<(), Error> {
self.visit_side_effect_instruction(cost);
self.adjust_branch_target(frame_index)?;
Ok(())
}
fn visit_unconditional_branch(&mut self, frame_index: usize, cost: Fee) -> Result<(), Error> {
self.visit_conditional_branch(frame_index, cost)?;
self.make_polymorphic();
Ok(())
}
}
macro_rules! trapping_insn {
(fn $visit:ident( $($arg:ident: $ty:ty),* )) => {
fn $visit(&mut self, $($arg: $ty),*) -> Output {
let cost = self.model.$visit($($arg),*);
Ok(self.visit_side_effect_instruction(cost))
}
};
($($_t:ident .
$(atomic.rmw)?
$(atomic.cmpxchg)?
$(load)?
$(store)?
= $($insn:ident)|* ;)*) => {
$($(trapping_insn!(fn $insn(mem: wasmparser::MemArg));)*)*
};
($($_t:ident . $(loadlane)? $(storelane)? = $($insn:ident)|* ;)*) => {
$($(trapping_insn!(fn $insn(mem: wasmparser::MemArg, lane: u8));)*)*
};
($($_t:ident . $(binop)? $(cvtop)? = $($insn:ident)|* ;)*) => {
$($(trapping_insn!(fn $insn());)*)*
};
}
macro_rules! pure_insn {
(fn $visit:ident( $($arg:ident: $ty:ty),* )) => {
fn $visit(&mut self, $($arg: $ty),*) -> Output {
let cost = self.model.$visit($($arg),*);
Ok(self.visit_pure_instruction(cost))
}
};
($($_t:ident .
$(unop)?
$(binop)?
$(cvtop)?
$(relop)?
$(testop)?
$(vbitmask)?
$(vinarrowop)?
$(vrelop)?
$(vternop)?
$(vishiftop)?
$(splat)?
= $($insn:ident)|* ;)*) => {
$($(pure_insn!(fn $insn());)*)*
};
($($_t:ident . const = $($insn:ident, $param:ty)|* ;)*) => {
$($(pure_insn!(fn $insn(val: $param));)*)*
};
($($_t:ident . $(extractlane)? $(replacelane)? = $($insn:ident)|* ;)*) => {
$($(pure_insn!(fn $insn(lane: u8));)*)*
};
($($_t:ident . localsglobals = $($insn:ident)|* ;)*) => {
$($(pure_insn!(fn $insn(index: u32));)*)*
};
}
type Output = Result<(), Error>;
impl<'a, 'b, CostModel: VisitOperator<'b, Output = Fee> + VisitSimdOperator<'b, Output = Fee>>
Visitor<'a, CostModel>
{
gen::atomic_cmpxchg!(trapping_insn);
gen::atomic_rmw!(trapping_insn);
gen::load!(trapping_insn);
gen::store!(trapping_insn);
gen::loadlane!(trapping_insn);
gen::storelane!(trapping_insn);
gen::binop_partial!(trapping_insn);
gen::cvtop_partial!(trapping_insn);
trapping_insn!(fn visit_call(index: u32));
trapping_insn!(fn visit_call_ref(type_index: u32));
trapping_insn!(fn visit_call_indirect(ty_index: u32, table_index: u32));
trapping_insn!(fn visit_memory_atomic_notify(mem: wasmparser::MemArg));
trapping_insn!(fn visit_memory_atomic_wait32(mem: wasmparser::MemArg));
trapping_insn!(fn visit_memory_atomic_wait64(mem: wasmparser::MemArg));
trapping_insn!(fn visit_table_set(table: u32));
trapping_insn!(fn visit_table_get(table: u32));
trapping_insn!(fn visit_ref_as_non_null());
fn visit_unreachable(&mut self) -> Output {
let cost = self.model.visit_unreachable();
self.push_instrumentation_before(InstrumentationKind::PreControlFlow, cost);
self.push_instrumentation_after(InstrumentationKind::Unreachable, Fee::ZERO);
self.make_polymorphic();
Ok(())
}
gen::binop_complete!(pure_insn);
gen::cvtop_complete!(pure_insn);
gen::unop!(pure_insn);
gen::relop!(pure_insn);
gen::vrelop!(pure_insn);
gen::vishiftop!(pure_insn);
gen::vternop!(pure_insn);
gen::vbitmask!(pure_insn);
gen::vinarrowop!(pure_insn);
gen::splat!(pure_insn);
gen::r#const!(pure_insn);
gen::extractlane!(pure_insn);
gen::replacelane!(pure_insn);
gen::testop!(pure_insn);
pure_insn!(fn visit_ref_null(t: wasmparser::HeapType));
pure_insn!(fn visit_ref_func(index: u32));
pure_insn!(fn visit_i8x16_shuffle(pattern: [u8; 16]));
pure_insn!(fn visit_atomic_fence());
pure_insn!(fn visit_select());
pure_insn!(fn visit_typed_select(t: wasmparser::ValType));
pure_insn!(fn visit_drop());
pure_insn!(fn visit_nop());
pure_insn!(fn visit_table_size(table: u32));
pure_insn!(fn visit_memory_size(mem: u32));
pure_insn!(fn visit_global_set(global: u32));
pure_insn!(fn visit_global_get(global: u32));
pure_insn!(fn visit_local_set(local: u32));
pure_insn!(fn visit_local_get(local: u32));
pure_insn!(fn visit_local_tee(local: u32));
fn visit_loop(&mut self, blockty: BlockType) -> Output {
let cost = self.model.visit_loop(blockty);
let instrumentation_kind_index_pre = self.state.kinds.len();
self.push_instrumentation_before(InstrumentationKind::Pure, Fee::ZERO);
let instrumentation_kind_index_post = self.state.kinds.len();
self.push_instrumentation_after(InstrumentationKind::Pure, cost);
self.new_frame(BranchTargetKind::Loop(
instrumentation_kind_index_pre,
instrumentation_kind_index_post,
));
Ok(())
}
fn visit_end(&mut self) -> Output {
let cost = self.model.visit_end();
assert!(
matches!(cost, Fee::ZERO),
"the `end` instruction costs aren’t handled right, set it to Fee::ZERO"
);
match self.state.current_frame.kind {
BranchTargetKind::Forward => self.visit_side_effect_instruction(cost),
BranchTargetKind::Loop(_, _) => self.visit_pure_instruction(cost),
BranchTargetKind::UntakenForward => self.visit_pure_instruction(cost),
}
self.end_frame();
Ok(())
}
fn visit_if(&mut self, blockty: BlockType) -> Output {
let cost = self.model.visit_if(blockty);
self.visit_side_effect_instruction(cost);
self.new_frame(BranchTargetKind::Forward);
Ok(())
}
fn visit_else(&mut self) -> Output {
let cost = self.model.visit_else();
assert!(
matches!(cost, Fee::ZERO),
"the `else` instruction costs aren’t handled right, set it to Fee::ZERO"
);
self.end_frame();
self.new_frame(BranchTargetKind::Forward);
self.visit_side_effect_instruction(cost);
Ok(())
}
fn visit_block(&mut self, blockty: BlockType) -> Output {
let cost = self.model.visit_block(blockty);
self.visit_pure_instruction(cost);
self.new_frame(BranchTargetKind::UntakenForward);
Ok(())
}
fn visit_br(&mut self, relative_depth: u32) -> Output {
let frame_idx = self.frame_index(relative_depth)?;
let cost = self.model.visit_br(relative_depth);
self.visit_unconditional_branch(frame_idx, cost)
}
fn visit_br_if(&mut self, relative_depth: u32) -> Output {
let frame_idx = self.frame_index(relative_depth)?;
let cost = self.model.visit_br_if(relative_depth);
self.visit_conditional_branch(frame_idx, cost)
}
fn visit_br_on_null(&mut self, relative_depth: u32) -> Output {
let frame_idx = self.frame_index(relative_depth)?;
let cost = self.model.visit_br_on_null(relative_depth);
self.visit_conditional_branch(frame_idx, cost)
}
fn visit_br_on_non_null(&mut self, relative_depth: u32) -> Output {
let frame_idx = self.frame_index(relative_depth)?;
let cost = self.model.visit_br_on_non_null(relative_depth);
self.visit_conditional_branch(frame_idx, cost)
}
fn visit_br_table(&mut self, targets: BrTable<'b>) -> Output {
let cost = self.model.visit_br_table(targets.clone());
self.visit_side_effect_instruction(cost);
for target in targets.targets() {
let target = target.map_err(Error::ParseBrTable)?;
self.adjust_branch_target(self.frame_index(target)?)?;
}
self.adjust_branch_target(self.frame_index(targets.default())?)?;
self.make_polymorphic();
Ok(())
}
fn visit_return(&mut self) -> Output {
let cost = self.model.visit_return();
self.visit_unconditional_branch(self.root_frame_index(), cost)
}
fn visit_return_call(&mut self, function_index: u32) -> Output {
let cost = self.model.visit_return_call(function_index);
self.visit_unconditional_branch(self.root_frame_index(), cost)
}
fn visit_return_call_ref(&mut self, type_index: u32) -> Output {
let cost = self.model.visit_return_call_ref(type_index);
self.visit_unconditional_branch(self.root_frame_index(), cost)
}
fn visit_return_call_indirect(&mut self, type_index: u32, table_index: u32) -> Output {
let cost = self
.model
.visit_return_call_indirect(type_index, table_index);
self.visit_unconditional_branch(self.root_frame_index(), cost)
}
fn visit_memory_grow(&mut self, mem: u32) -> Output {
let cost = self.model.visit_memory_grow(mem);
self.push_instrumentation_before(InstrumentationKind::MemoryGrow, cost);
self.push_instrumentation_after(InstrumentationKind::PostControlFlow, Fee::ZERO);
Ok(())
}
fn visit_memory_init(&mut self, data_index: u32, mem: u32) -> Output {
let cost = self.model.visit_memory_init(data_index, mem);
self.push_instrumentation_before(InstrumentationKind::MemoryInit, cost);
self.push_instrumentation_after(InstrumentationKind::PostControlFlow, Fee::ZERO);
Ok(())
}
fn visit_data_drop(&mut self, data_index: u32) -> Output {
let cost = self.model.visit_data_drop(data_index);
Ok(self.visit_pure_instruction(cost))
}
fn visit_memory_copy(&mut self, dst_mem: u32, src_mem: u32) -> Output {
let cost = self.model.visit_memory_copy(dst_mem, src_mem);
self.push_instrumentation_before(InstrumentationKind::MemoryCopy, cost);
self.push_instrumentation_after(InstrumentationKind::PostControlFlow, Fee::ZERO);
Ok(())
}
fn visit_memory_fill(&mut self, mem: u32) -> Output {
let cost = self.model.visit_memory_fill(mem);
self.push_instrumentation_before(InstrumentationKind::MemoryFill, cost);
self.push_instrumentation_after(InstrumentationKind::PostControlFlow, Fee::ZERO);
Ok(())
}
fn visit_table_init(&mut self, elem_index: u32, table: u32) -> Output {
let cost = self.model.visit_table_init(elem_index, table);
self.push_instrumentation_before(InstrumentationKind::TableInit, cost);
self.push_instrumentation_after(InstrumentationKind::PostControlFlow, Fee::ZERO);
Ok(())
}
fn visit_elem_drop(&mut self, elem_index: u32) -> Output {
let cost = self.model.visit_elem_drop(elem_index);
Ok(self.visit_pure_instruction(cost))
}
fn visit_table_copy(&mut self, dst_table: u32, src_table: u32) -> Output {
let cost = self.model.visit_table_copy(dst_table, src_table);
self.push_instrumentation_before(InstrumentationKind::TableCopy, cost);
self.push_instrumentation_after(InstrumentationKind::PostControlFlow, Fee::ZERO);
Ok(())
}
fn visit_table_fill(&mut self, table: u32) -> Output {
let cost = self.model.visit_table_fill(table);
self.push_instrumentation_before(InstrumentationKind::TableFill, cost);
self.push_instrumentation_after(InstrumentationKind::PostControlFlow, Fee::ZERO);
Ok(())
}
fn visit_table_grow(&mut self, table: u32) -> Output {
let cost = self.model.visit_table_grow(table);
self.push_instrumentation_before(InstrumentationKind::TableGrow, cost);
self.push_instrumentation_after(InstrumentationKind::PostControlFlow, Fee::ZERO);
Ok(())
}
}
macro_rules! delegate_one_to_inherent {
(@legacy_exceptions $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*)) => {
fn $visit(&mut self $($(,$arg: $argty)*)?) -> Output {
Err(Error::ExceptionsNotSupported(self.offset))
}
};
(@exceptions $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*)) => {
fn $visit(&mut self $($(,$arg: $argty)*)?) -> Output {
Err(Error::ExceptionsNotSupported(self.offset))
}
};
(@gc $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*)) => {
fn $visit(&mut self $($(,$arg: $argty)*)?) -> Output {
Err(Error::GcNotSupported(self.offset))
}
};
(@memory_control $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*)) => {
fn $visit(&mut self $($(,$arg: $argty)*)?) -> Output {
Err(Error::MemoryControlNotSupported(self.offset))
}
};
(@shared_everything_threads $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*)) => {
fn $visit(&mut self $($(,$arg: $argty)*)?) -> Output {
Err(Error::ThreadsNotSupported(self.offset))
}
};
(@stack_switching $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*)) => {
fn $visit(&mut self $($(,$arg: $argty)*)?) -> Output {
Err(Error::StackSwitchingNotSupported(self.offset))
}
};
(@wide_arithmetic $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*)) => {
fn $visit(&mut self $($(,$arg: $argty)*)?) -> Output {
Err(Error::WideArithmeticNotSupported(self.offset))
}
};
(@$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*)) => {
fn $visit(&mut self $($(,$arg: $argty)*)?) -> Output {
self.$visit($($($arg.clone()),*)?)
}
}
}
macro_rules! delegate_to_inherent {
($( @$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*))*) => {
$(delegate_one_to_inherent!{
@$proposal $op $({ $($arg: $argty),* })? => $visit ($($ann)*)
})*
}
}
#[deny(unconditional_recursion)]
impl<'b, 'a, CostModel: VisitSimdOperator<'a, Output = Fee>> VisitOperator<'a>
for Visitor<'b, CostModel>
{
type Output = Output;
fn simd_visitor(
&mut self,
) -> Option<&mut dyn wasmparser::VisitSimdOperator<'a, Output = Self::Output>> {
Some(self)
}
wasmparser::for_each_visit_operator!(delegate_to_inherent);
}
#[deny(unconditional_recursion)]
impl<'a, 'b, CostModel: VisitSimdOperator<'b, Output = Fee>> VisitSimdOperator<'b>
for Visitor<'a, CostModel>
{
wasmparser::for_each_visit_simd_operator!(delegate_to_inherent);
}
impl<'a, 'b, CostModel: VisitOperator<'b, Output = Fee> + VisitSimdOperator<'b>>
crate::visitors::VisitOperatorWithOffset<'b> for Visitor<'a, CostModel>
{
fn set_offset(&mut self, offset: usize) {
self.offset = offset;
if let Some(scheduled) = self.state.scheduled_instrumentation.take() {
self.state.offsets.push(self.offset);
self.state.kinds.push(scheduled.kind);
self.state.costs.push(scheduled.cost);
}
}
}