use crate::ir::id::{FunctionID, GlobalID, LocalID};
use crate::ir::module::module_functions::{FuncKind, LocalFunction};
use crate::ir::module::module_globals::Global;
use crate::ir::module::{Iter, Module};
use crate::ir::types::{DataType, FuncInstrMode, InstrumentationMode, Location};
use crate::iterator::iterator_trait::{IteratingInstrumenter, Iterator};
use crate::module_builder::AddLocal;
use crate::opcode::{Inject, InjectAt, Instrumenter, MacroOpcode, Opcode};
use crate::subiterator::module_subiterator::ModuleSubIterator;
use std::collections::HashMap;
use wasmparser::Operator;
pub struct ModuleIterator<'a, 'b> {
pub module: &'a mut Module<'b>,
mod_iterator: ModuleSubIterator,
}
#[allow(dead_code)]
impl<'a, 'b> ModuleIterator<'a, 'b> {
pub fn new(module: &'a mut Module<'b>, skip_funcs: &Vec<FunctionID>) -> Self {
let mut metadata = HashMap::new();
for func in module.functions.iter() {
match &func.kind {
FuncKind::Import(_) => {}
FuncKind::Local(LocalFunction { func_id, body, .. }) => {
metadata.insert(*func_id, body.num_instructions);
}
}
}
let num_funcs = module.num_local_functions;
ModuleIterator {
module,
mod_iterator: ModuleSubIterator::new(num_funcs, metadata, skip_funcs.to_owned()),
}
}
pub fn curr_op_owned(&self) -> Option<Operator<'b>> {
if self.mod_iterator.end() {
None
} else if let (
Location::Module {
func_idx,
instr_idx,
..
},
..,
) = self.mod_iterator.curr_loc()
{
match &self.module.functions.get(func_idx as FunctionID).kind {
FuncKind::Import(_) => panic!("Cannot get an instruction to an imported function"),
FuncKind::Local(l) => Some(l.body.instructions[instr_idx].op.clone()),
}
} else {
panic!("Should have gotten Module Location!")
}
}
}
impl<'a, 'b> Inject<'b> for ModuleIterator<'a, 'b> {
fn inject(&mut self, instr: Operator<'b>) {
if let (
Location::Module {
func_idx,
instr_idx,
..
},
..,
) = self.curr_loc()
{
match self.module.functions.get_mut(func_idx as FunctionID).kind {
FuncKind::Import(_) => panic!("Cannot get an instruction to an imported function"),
FuncKind::Local(ref mut l) => l.add_instr(instr, instr_idx),
}
} else {
panic!("Should have gotten Module Location!")
}
}
}
impl<'a> InjectAt<'a> for ModuleIterator<'_, 'a> {
fn inject_at(&mut self, idx: usize, mode: InstrumentationMode, instr: Operator<'a>) {
if let (Location::Module { func_idx, .. }, ..) = self.curr_loc() {
let loc = Location::Module {
func_idx,
instr_idx: idx,
};
self.set_instrument_mode_at(mode, loc);
self.add_instr_at(loc, instr);
} else {
panic!("Should have gotten Module Location!")
}
}
}
impl<'a> Opcode<'a> for ModuleIterator<'_, 'a> {}
impl<'a> MacroOpcode<'a> for ModuleIterator<'_, 'a> {}
impl<'a> Instrumenter<'a> for ModuleIterator<'_, 'a> {
fn curr_instrument_mode(&self) -> &Option<InstrumentationMode> {
if let (
Location::Module {
func_idx,
instr_idx,
..
},
..,
) = self.mod_iterator.curr_loc()
{
match &self.module.functions.get(func_idx as FunctionID).kind {
FuncKind::Import(_) => panic!("Cannot get an instruction to an imported function"),
FuncKind::Local(l) => &l.body.instructions[instr_idx].instr_flag.current_mode,
}
} else {
panic!("Should have gotten Module Location and not Module Location!")
}
}
fn set_instrument_mode_at(&mut self, mode: InstrumentationMode, loc: Location) {
if let Location::Module {
func_idx,
instr_idx,
..
} = loc
{
match self.module.functions.get_mut(func_idx as FunctionID).kind {
FuncKind::Import(_) => panic!("Cannot add an instruction to an imported function"),
FuncKind::Local(ref mut l) => {
l.body.instructions[instr_idx].instr_flag.current_mode = Some(mode)
}
}
} else {
panic!("Should have gotten module location!")
}
}
fn curr_func_instrument_mode(&self) -> &Option<FuncInstrMode> {
if let (Location::Module { func_idx, .. }, ..) = self.mod_iterator.curr_loc() {
match &self.module.functions.get(func_idx as FunctionID).kind {
FuncKind::Import(_) => panic!("Cannot get an instruction to an imported function"),
FuncKind::Local(l) => &l.instr_flag.current_mode,
}
} else {
panic!("Should have gotten Module Location and not Module Location!")
}
}
fn set_func_instrument_mode(&mut self, mode: FuncInstrMode) {
if let (Location::Module { func_idx, .. }, ..) = self.mod_iterator.curr_loc() {
match self.module.functions.get_mut(func_idx as FunctionID).kind {
FuncKind::Import(_) => panic!("Cannot get an instruction to an imported function"),
FuncKind::Local(ref mut l) => l.instr_flag.current_mode = Some(mode),
}
} else {
panic!("Should have gotten Module Location and not Module Location!")
}
}
fn clear_instr_at(&mut self, loc: Location, mode: InstrumentationMode) {
if let Location::Module {
func_idx,
instr_idx,
..
} = loc
{
match self.module.functions.get_mut(func_idx as FunctionID).kind {
FuncKind::Import(_) => panic!("Cannot add an instruction to an imported function"),
FuncKind::Local(ref mut l) => {
l.clear_instr_at(instr_idx, mode);
}
}
} else {
panic!("Should have gotten Module Location!")
}
}
fn add_instr_at(&mut self, loc: Location, instr: Operator<'a>) {
if let Location::Module {
func_idx,
instr_idx,
..
} = loc
{
match self.module.functions.get_mut(func_idx as FunctionID).kind {
FuncKind::Import(_) => panic!("Cannot add an instruction to an imported function"),
FuncKind::Local(ref mut l) => {
l.add_instr(instr, instr_idx);
}
}
} else {
panic!("Should have gotten Module Location!")
}
}
fn empty_alternate_at(&mut self, loc: Location) -> &mut Self {
if let Location::Module {
func_idx,
instr_idx,
..
} = loc
{
match self.module.functions.get_mut(func_idx as FunctionID).kind {
FuncKind::Import(_) => panic!("Cannot instrument an imported function"),
FuncKind::Local(ref mut l) => {
l.body.instructions[instr_idx].instr_flag.alternate = Some(vec![])
}
}
} else {
panic!("Should have gotten Module Location and not Module Location!")
}
self
}
fn empty_block_alt_at(&mut self, loc: Location) -> &mut Self {
if let Location::Module {
func_idx,
instr_idx,
..
} = loc
{
match self.module.functions.get_mut(func_idx as FunctionID).kind {
FuncKind::Import(_) => panic!("Cannot instrument an imported function"),
FuncKind::Local(ref mut l) => {
l.body.instructions[instr_idx].instr_flag.block_alt = Some(vec![]);
l.instr_flag.has_special_instr |= true;
}
}
} else {
panic!("Should have gotten Module Location and not Module Location!")
}
self
}
fn get_injected_val(&self, idx: usize) -> &Operator {
if let (
Location::Module {
func_idx,
instr_idx,
..
},
..,
) = self.mod_iterator.curr_loc()
{
match &self.module.functions.get(func_idx as FunctionID).kind {
FuncKind::Import(_) => panic!("Cannot get an instruction to an imported function"),
FuncKind::Local(l) => l.body.instructions[instr_idx].instr_flag.get_instr(idx),
}
} else {
panic!("Should have gotten Component Location and not Module Location!")
}
}
}
impl<'a> IteratingInstrumenter<'a> for ModuleIterator<'_, 'a> {
fn set_instrument_mode(&mut self, mode: InstrumentationMode) {
self.set_instrument_mode_at(mode, self.curr_loc().0);
}
fn add_global(&mut self, global: Global) -> GlobalID {
self.module.globals.add(global)
}
}
impl AddLocal for ModuleIterator<'_, '_> {
fn add_local(&mut self, val_type: DataType) -> LocalID {
let curr_loc = self.curr_loc();
if let (Location::Module { func_idx, .. }, ..) = curr_loc {
self.module.functions.add_local(func_idx, val_type)
} else {
panic!("Should have gotten Module Location!")
}
}
}
impl<'a> Iterator for ModuleIterator<'_, 'a> {
fn reset(&mut self) {
self.mod_iterator.reset();
}
fn next(&mut self) -> Option<&Operator> {
match self.mod_iterator.next() {
false => None,
true => self.curr_op(),
}
}
fn curr_loc(&self) -> (Location, bool) {
self.mod_iterator.curr_loc()
}
fn curr_op(&self) -> Option<&Operator<'a>> {
if self.mod_iterator.end() {
None
} else if let (
Location::Module {
func_idx,
instr_idx,
..
},
..,
) = self.mod_iterator.curr_loc()
{
match &self.module.functions.get(func_idx as FunctionID).kind {
FuncKind::Import(_) => panic!("Cannot get an instruction to an imported function"),
FuncKind::Local(l) => Some(&l.body.instructions[instr_idx].op),
}
} else {
panic!("Should have gotten Module Location!")
}
}
}