use crate::ir::component::Component;
use crate::ir::id::{FunctionID, GlobalID, LocalID, ModuleID};
use crate::ir::module::module_functions::{FuncKind, LocalFunction};
use crate::ir::module::module_globals::Global;
use crate::ir::module::Iter;
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::component_subiterator::ComponentSubIterator;
use std::collections::HashMap;
use std::iter::Iterator as StdIter;
use wasmparser::Operator;
pub struct ComponentIterator<'a, 'b> {
pub comp: &'a mut Component<'b>,
comp_iterator: ComponentSubIterator,
}
fn print_metadata(metadata: &HashMap<ModuleID, HashMap<FunctionID, usize>>) {
for c in metadata.keys() {
println!("Module: {:?}", c);
for (m, i) in metadata.get(c).unwrap().iter() {
println!("Function: {:?} Instr: {:?}", m, i);
}
}
}
#[allow(dead_code)]
impl<'a, 'b> ComponentIterator<'a, 'b> {
pub fn new(
comp: &'a mut Component<'b>,
skip_funcs: HashMap<ModuleID, Vec<FunctionID>>,
) -> Self {
let mut metadata = HashMap::new();
for (mod_idx, m) in comp.modules.iter().enumerate() {
let mut mod_metadata = HashMap::new();
for func in m.functions.iter() {
match &func.kind {
FuncKind::Import(_) => {}
FuncKind::Local(LocalFunction { func_id, body, .. }) => {
mod_metadata.insert(*func_id, body.num_instructions);
}
}
}
metadata.insert(ModuleID(mod_idx as u32), mod_metadata);
}
print_metadata(&metadata);
let num_modules = comp.num_modules;
ComponentIterator {
comp,
comp_iterator: ComponentSubIterator::new(
ModuleID(0),
num_modules,
metadata,
skip_funcs,
),
}
}
pub fn curr_module(&self) -> ModuleID {
if let (
Location::Component {
mod_idx,
func_idx: _func_idx,
instr_idx: _instr_idx,
..
},
..,
) = self.curr_loc()
{
mod_idx
} else {
panic!("Should have gotten component location");
}
}
pub fn curr_op_owned(&self) -> Option<Operator<'b>> {
if self.comp_iterator.end() {
None
} else if let (
Location::Component {
mod_idx,
func_idx,
instr_idx,
..
},
..,
) = self.comp_iterator.curr_loc()
{
match &self.comp.modules[*mod_idx as usize]
.functions
.get(func_idx)
.kind
{
FuncKind::Import(_) => None,
FuncKind::Local(l) => Some(l.body.instructions[instr_idx].op.clone()),
}
} else {
panic!("Should have gotten Component Location!")
}
}
}
impl<'a, 'b> Inject<'b> for ComponentIterator<'a, 'b> {
fn inject(&mut self, instr: Operator<'b>) {
if let (
Location::Component {
mod_idx,
func_idx,
instr_idx,
..
},
..,
) = self.curr_loc()
{
match self.comp.modules[*mod_idx as usize]
.functions
.get_mut(func_idx)
.kind
{
FuncKind::Import(_) => panic!("Can't inject into an imported function!"),
FuncKind::Local(ref mut l) => l.add_instr(instr, instr_idx),
}
} else {
panic!("Should have gotten component location!")
}
}
}
impl<'a, 'b> InjectAt<'b> for ComponentIterator<'a, 'b> {
fn inject_at(&mut self, idx: usize, mode: InstrumentationMode, instr: Operator<'b>) {
if let (
Location::Component {
mod_idx, func_idx, ..
},
..,
) = self.curr_loc()
{
let loc = Location::Component {
mod_idx,
func_idx,
instr_idx: idx,
};
self.set_instrument_mode_at(mode, loc);
self.add_instr_at(loc, instr);
} else {
panic!("Should have gotten Component Location!")
}
}
}
impl<'a, 'b> Opcode<'b> for ComponentIterator<'a, 'b> {}
impl<'a, 'b> MacroOpcode<'b> for ComponentIterator<'a, 'b> {}
impl<'a, 'b> Instrumenter<'b> for ComponentIterator<'a, 'b> {
fn curr_instrument_mode(&self) -> &Option<InstrumentationMode> {
if let (
Location::Component {
mod_idx,
func_idx,
instr_idx,
..
},
..,
) = self.comp_iterator.curr_loc()
{
match &self.comp.modules[*mod_idx as usize]
.functions
.get(func_idx)
.kind
{
FuncKind::Import(_) => {
panic!("Can't get instrumentation from an imported function!")
}
FuncKind::Local(l) => &l.body.instructions[instr_idx].instr_flag.current_mode,
}
} else {
panic!("Should have gotten Component Location and not Module Location!")
}
}
fn set_instrument_mode_at(&mut self, mode: InstrumentationMode, loc: Location) {
if let Location::Component {
mod_idx,
func_idx,
instr_idx,
..
} = loc
{
match self.comp.modules[*mod_idx as usize]
.functions
.get_mut(func_idx)
.kind
{
FuncKind::Import(_) => panic!("Can't instrument into an imported function!"),
FuncKind::Local(ref mut l) => {
l.body.instructions[instr_idx].instr_flag.current_mode = Some(mode)
}
}
} else {
panic!("Should have gotten component location!")
}
}
fn curr_func_instrument_mode(&self) -> &Option<FuncInstrMode> {
if let (
Location::Component {
mod_idx, func_idx, ..
},
..,
) = self.comp_iterator.curr_loc()
{
match &self.comp.modules[*mod_idx as usize]
.functions
.get(func_idx)
.kind
{
FuncKind::Import(_) => {
panic!("Can't get instrumentation from an imported function!")
}
FuncKind::Local(l) => &l.instr_flag.current_mode,
}
} else {
panic!("Should have gotten Component Location and not Module Location!")
}
}
fn set_func_instrument_mode(&mut self, mode: FuncInstrMode) {
if let (
Location::Component {
mod_idx, func_idx, ..
},
..,
) = self.curr_loc()
{
match self.comp.modules[*mod_idx as usize]
.functions
.get_mut(func_idx)
.kind
{
FuncKind::Import(_) => panic!("Can't instrument into an imported function!"),
FuncKind::Local(ref mut l) => l.instr_flag.current_mode = Some(mode),
}
} else {
panic!("Should have gotten component location!")
}
}
fn clear_instr_at(&mut self, loc: Location, mode: InstrumentationMode) {
if let Location::Component {
mod_idx,
func_idx,
instr_idx,
..
} = loc
{
match self.comp.modules[*mod_idx as usize]
.functions
.get_mut(func_idx)
.kind
{
FuncKind::Import(_) => panic!("Can't instrument into an imported function!"),
FuncKind::Local(ref mut l) => l.clear_instr_at(instr_idx, mode),
}
} else {
panic!("Should have gotten component location!")
}
}
fn add_instr_at(&mut self, loc: Location, instr: Operator<'b>) {
if let Location::Component {
mod_idx,
func_idx,
instr_idx,
..
} = loc
{
match self.comp.modules[*mod_idx as usize]
.functions
.get_mut(func_idx)
.kind
{
FuncKind::Import(_) => panic!("Can't instrument into an imported function!"),
FuncKind::Local(ref mut l) => {
l.add_instr(instr, instr_idx);
}
}
} else {
panic!("Should have gotten Component Location and not Module Location!")
}
}
fn empty_alternate_at(&mut self, loc: Location) -> &mut Self {
if let Location::Component {
mod_idx,
func_idx,
instr_idx,
..
} = loc
{
match self.comp.modules[*mod_idx as usize]
.functions
.get_mut(func_idx)
.kind
{
FuncKind::Import(_) => panic!("Can't instrument into an imported function!"),
FuncKind::Local(ref mut l) => {
l.body.instructions[instr_idx].instr_flag.alternate = Some(vec![]);
}
}
} else {
panic!("Should have gotten Component Location and not Module Location!")
}
self
}
fn empty_block_alt_at(&mut self, loc: Location) -> &mut Self {
if let Location::Component {
mod_idx,
func_idx,
instr_idx,
..
} = loc
{
match self.comp.modules[*mod_idx as usize]
.functions
.get_mut(func_idx)
.kind
{
FuncKind::Import(_) => panic!("Can't instrument into 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 Component Location and not Module Location!")
}
self
}
fn get_injected_val(&self, idx: usize) -> &Operator {
if let (
Location::Component {
mod_idx,
func_idx,
instr_idx,
..
},
..,
) = self.comp_iterator.curr_loc()
{
match &self.comp.modules[*mod_idx as usize]
.functions
.get(func_idx)
.kind
{
FuncKind::Import(_) => panic!("Can't inject into 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, 'b> IteratingInstrumenter<'b> for ComponentIterator<'a, 'b> {
fn set_instrument_mode(&mut self, mode: InstrumentationMode) {
if let (
Location::Component {
mod_idx: _mod_idx,
func_idx: _func_idx,
instr_idx: _instr_idx,
..
},
..,
) = self.curr_loc()
{
self.set_instrument_mode_at(mode, self.curr_loc().0);
} else {
panic!("Should have gotten component location!")
}
}
fn add_global(&mut self, global: Global) -> GlobalID {
let curr_mod = *self.curr_module() as usize;
self.comp.modules[curr_mod].globals.add(global)
}
}
impl<'a, 'b> Iterator for ComponentIterator<'a, 'b> {
fn reset(&mut self) {
self.comp_iterator.reset();
}
fn next(&mut self) -> Option<&Operator> {
match self.comp_iterator.next() {
false => None,
true => self.curr_op(),
}
}
fn curr_loc(&self) -> (Location, bool) {
self.comp_iterator.curr_loc()
}
fn curr_op(&self) -> Option<&Operator> {
if self.comp_iterator.end() {
None
} else if let (
Location::Component {
mod_idx,
func_idx,
instr_idx,
..
},
..,
) = self.comp_iterator.curr_loc()
{
match &self.comp.modules[*mod_idx as usize]
.functions
.get(func_idx)
.kind
{
FuncKind::Import(_) => panic!("Can't inject into an imported function!"),
FuncKind::Local(l) => Some(&l.body.instructions[instr_idx].op),
}
} else {
panic!("Should have gotten Component Location and not Module Location!")
}
}
}
impl AddLocal for ComponentIterator<'_, '_> {
fn add_local(&mut self, val_type: DataType) -> LocalID {
let curr_loc = self.curr_loc().0;
if let Location::Component {
mod_idx, func_idx, ..
} = curr_loc
{
{
self.comp.modules[*mod_idx as usize]
.functions
.add_local(func_idx, val_type)
}
} else {
panic!("Should have gotten Component Location and not Module Location!")
}
}
}