use std::collections::{BTreeMap, BTreeSet};
use crate::{ControlFlow, Instruction};
#[derive(Debug)]
pub struct BasicBlock {
pub instructions: Vec<Instruction>,
pub next: NextBranch,
}
#[derive(Debug)]
pub enum NextBranch {
None,
Goto(usize),
Cond {
t: usize,
f: usize,
},
}
impl NextBranch {
pub fn iter(&self) -> impl Iterator<Item = usize> {
match self {
NextBranch::None => [None, None].into_iter().flatten(),
NextBranch::Goto(t) => [Some(*t), None].into_iter().flatten(),
NextBranch::Cond { t, f } => [Some(*t), Some(*f)].into_iter().flatten(),
}
}
}
pub fn basic_blocks(bytecode: &[u16], entries: &[usize]) -> BTreeMap<usize, BasicBlock> {
let mut bbs = BTreeMap::new();
let mut search_next = BTreeSet::from([0]);
for e in entries {
search_next.insert(*e);
}
while let Some(start_addr) = search_next.pop_first() {
let bb = decode_bb(bytecode, start_addr, &search_next);
for next in bb.next.iter() {
if !bbs.contains_key(&next) {
search_next.insert(next);
}
}
bbs.insert(start_addr, bb);
}
bbs
}
fn decode_bb(bytecode: &[u16], entry_point: usize, avoid: &BTreeSet<usize>) -> BasicBlock {
let mut instructions = Vec::new();
let next;
let mut cursor = entry_point;
loop {
let inst = crate::decode::decode_one(&mut &bytecode[cursor..]).unwrap();
let cf = inst.control_flow();
let len = inst.len();
instructions.push(inst);
next = match cf {
ControlFlow::FallThrough => {
cursor += len;
if !avoid.contains(&cursor) {
continue;
}
NextBranch::Goto(cursor)
}
ControlFlow::GoTo(t) => NextBranch::Goto((cursor as i32 + t) as usize),
ControlFlow::Branch(t) => NextBranch::Cond {
t: (cursor as i32 + t as i32) as usize,
f: cursor + len,
},
ControlFlow::Terminate => NextBranch::None,
};
break;
}
BasicBlock { instructions, next }
}