use std::{
collections::{HashMap, VecDeque},
fmt,
};
use crate::{
analysis::ssa::{PhiNode, PhiOperand, SsaInstruction, SsaOp, SsaVarId},
utils::BitSet,
};
#[derive(Debug, Clone, Copy, Default)]
pub struct ReplaceResult {
pub replaced: usize,
pub skipped: usize,
}
impl ReplaceResult {
#[must_use]
pub const fn is_complete(&self) -> bool {
self.skipped == 0
}
}
impl std::ops::Add for ReplaceResult {
type Output = Self;
fn add(self, rhs: Self) -> Self {
Self {
replaced: self.replaced + rhs.replaced,
skipped: self.skipped + rhs.skipped,
}
}
}
#[derive(Debug, Clone)]
pub struct SsaBlock {
id: usize,
phi_nodes: Vec<PhiNode>,
instructions: Vec<SsaInstruction>,
}
impl SsaBlock {
#[must_use]
pub fn new(id: usize) -> Self {
Self {
id,
phi_nodes: Vec::new(),
instructions: Vec::new(),
}
}
#[must_use]
pub fn with_capacity(id: usize, phi_capacity: usize, instr_capacity: usize) -> Self {
Self {
id,
phi_nodes: Vec::with_capacity(phi_capacity),
instructions: Vec::with_capacity(instr_capacity),
}
}
#[must_use]
pub const fn id(&self) -> usize {
self.id
}
pub fn set_id(&mut self, id: usize) {
self.id = id;
}
#[must_use]
pub fn phi_nodes(&self) -> &[PhiNode] {
&self.phi_nodes
}
pub fn phi_nodes_mut(&mut self) -> &mut Vec<PhiNode> {
&mut self.phi_nodes
}
#[must_use]
pub fn instructions(&self) -> &[SsaInstruction] {
&self.instructions
}
pub fn instructions_mut(&mut self) -> &mut Vec<SsaInstruction> {
&mut self.instructions
}
#[must_use]
pub fn phi_count(&self) -> usize {
self.phi_nodes.len()
}
#[must_use]
pub fn instruction_count(&self) -> usize {
self.instructions.len()
}
#[must_use]
pub fn has_no_phis(&self) -> bool {
self.phi_nodes.is_empty()
}
#[must_use]
pub fn has_no_instructions(&self) -> bool {
self.instructions.is_empty()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.phi_nodes.is_empty() && self.instructions.is_empty()
}
pub fn clear(&mut self) {
self.phi_nodes.clear();
self.instructions.clear();
}
pub fn add_phi(&mut self, phi: PhiNode) {
self.phi_nodes.push(phi);
}
pub fn add_instruction(&mut self, instr: SsaInstruction) {
self.instructions.push(instr);
}
#[must_use]
pub fn phi(&self, index: usize) -> Option<&PhiNode> {
self.phi_nodes.get(index)
}
pub fn phi_mut(&mut self, index: usize) -> Option<&mut PhiNode> {
self.phi_nodes.get_mut(index)
}
#[must_use]
pub fn instruction(&self, index: usize) -> Option<&SsaInstruction> {
self.instructions.get(index)
}
pub fn instruction_mut(&mut self, index: usize) -> Option<&mut SsaInstruction> {
self.instructions.get_mut(index)
}
#[must_use]
pub fn terminator(&self) -> Option<&SsaInstruction> {
self.instructions.last()
}
#[must_use]
pub fn terminator_op(&self) -> Option<&SsaOp> {
self.instructions.last().map(SsaInstruction::op)
}
#[must_use]
pub fn successors(&self) -> Vec<usize> {
self.terminator_op()
.map_or_else(Vec::new, super::SsaOp::successors)
}
pub fn redirect_target(&mut self, old_target: usize, new_target: usize) -> bool {
if let Some(terminator) = self.instructions.last_mut() {
return terminator.op_mut().redirect_target(old_target, new_target);
}
false
}
pub fn set_target(&mut self, target: usize) {
if let Some(terminator) = self.instructions.last_mut() {
match terminator.op_mut() {
SsaOp::Jump { target: t } | SsaOp::Leave { target: t } => {
*t = target;
}
SsaOp::Branch {
true_target,
false_target,
..
}
| SsaOp::BranchCmp {
true_target,
false_target,
..
} => {
*true_target = target;
*false_target = target;
}
SsaOp::Switch {
targets, default, ..
} => {
*default = target;
for t in targets.iter_mut() {
*t = target;
}
}
_ => {
*terminator = SsaInstruction::synthetic(SsaOp::Jump { target });
}
}
} else {
self.instructions
.push(SsaInstruction::synthetic(SsaOp::Jump { target }));
}
}
pub fn replace_uses(&mut self, old_var: SsaVarId, new_var: SsaVarId) -> ReplaceResult {
let mut replaced = 0;
let mut skipped = 0;
for instr in &mut self.instructions {
let op = instr.op_mut();
if let Some(dest) = op.dest() {
if dest == new_var {
if op.uses().contains(&old_var) {
skipped += 1;
}
continue;
}
}
replaced += op.replace_uses(old_var, new_var);
}
ReplaceResult { replaced, skipped }
}
pub(crate) fn replace_uses_including_phis(
&mut self,
old_var: SsaVarId,
new_var: SsaVarId,
) -> ReplaceResult {
let mut result = self.replace_uses(old_var, new_var);
for phi in &mut self.phi_nodes {
for operand in phi.operands_mut() {
if operand.value() == old_var {
*operand = PhiOperand::new(new_var, operand.predecessor());
result.replaced += 1;
}
}
}
result
}
#[must_use]
pub fn find_phi_defining(&self, var: SsaVarId) -> Option<&PhiNode> {
self.phi_nodes.iter().find(|phi| phi.result() == var)
}
#[must_use]
pub fn is_trampoline(&self) -> Option<usize> {
if !self.phi_nodes.is_empty() {
return None;
}
if self.instructions.len() != 1 {
return None;
}
match self.instructions[0].op() {
SsaOp::Jump { target } | SsaOp::Leave { target } => Some(*target),
_ => None,
}
}
pub fn defined_variables(&self) -> impl Iterator<Item = SsaVarId> + '_ {
let phi_defs = self.phi_nodes.iter().map(PhiNode::result);
let instr_defs = self.instructions.iter().filter_map(SsaInstruction::def);
phi_defs.chain(instr_defs)
}
pub fn used_variables(&self) -> impl Iterator<Item = SsaVarId> + '_ {
let phi_uses = self.phi_nodes.iter().flat_map(PhiNode::used_variables);
let instr_uses = self.instructions.iter().flat_map(SsaInstruction::uses);
phi_uses.chain(instr_uses)
}
pub fn sort_instructions_topologically(&mut self) -> bool {
if self.instructions.len() <= 1 {
return true;
}
let mut terminators: Vec<(usize, SsaInstruction)> = Vec::new();
let mut non_terminator_indices: Vec<usize> = Vec::new();
for (idx, instr) in self.instructions.iter().enumerate() {
if instr.is_terminator() {
terminators.push((idx, instr.clone()));
} else {
non_terminator_indices.push(idx);
}
}
if non_terminator_indices.is_empty() {
return true;
}
let mut def_index: HashMap<SsaVarId, usize> = HashMap::new();
for &idx in &non_terminator_indices {
if let Some(dest) = self.instructions[idx].def() {
def_index.insert(dest, idx);
}
}
let max_phi_var = self
.phi_nodes
.iter()
.map(|phi| phi.result().index())
.max()
.map_or(0, |m| m + 1);
let mut phi_defs = BitSet::new(max_phi_var);
for phi in &self.phi_nodes {
phi_defs.insert(phi.result().index());
}
let idx_to_pos: HashMap<usize, usize> = non_terminator_indices
.iter()
.enumerate()
.map(|(pos, &idx)| (idx, pos))
.collect();
let n = non_terminator_indices.len();
let mut deps: Vec<BitSet> = (0..n).map(|_| BitSet::new(n)).collect();
let mut rdeps: Vec<BitSet> = (0..n).map(|_| BitSet::new(n)).collect();
let mut prev_side_effect_pos: Option<usize> = None;
for (pos, &idx) in non_terminator_indices.iter().enumerate() {
let instr = &self.instructions[idx];
for used in &instr.uses() {
if used.index() < phi_defs.len() && phi_defs.contains(used.index()) {
continue;
}
if let Some(&dep_idx) = def_index.get(used) {
if dep_idx != idx {
if let Some(&dep_pos) = idx_to_pos.get(&dep_idx) {
deps[pos].insert(dep_pos);
rdeps[dep_pos].insert(pos);
}
}
}
}
if !instr.op().is_pure() {
if let Some(prev_pos) = prev_side_effect_pos {
deps[pos].insert(prev_pos);
rdeps[prev_pos].insert(pos);
}
prev_side_effect_pos = Some(pos);
}
}
let mut in_degree: Vec<usize> = deps.iter().map(BitSet::count).collect();
let mut ready: VecDeque<usize> = VecDeque::new();
for (pos, °) in in_degree.iter().enumerate() {
if deg == 0 {
ready.push_back(pos);
}
}
let mut sorted_positions: Vec<usize> = Vec::with_capacity(n);
while let Some(pos) = ready.pop_front() {
sorted_positions.push(pos);
for dep_pos in rdeps[pos].iter() {
in_degree[dep_pos] -= 1;
if in_degree[dep_pos] == 0 {
ready.push_back(dep_pos);
}
}
}
if sorted_positions.len() != n {
return false;
}
let mut temp: Vec<Option<SsaInstruction>> = self.instructions.drain(..).map(Some).collect();
for pos in sorted_positions {
let original_idx = non_terminator_indices[pos];
if let Some(instr) = temp[original_idx].take() {
self.instructions.push(instr);
}
}
terminators.sort_by_key(|(idx, _)| *idx);
for (_, instr) in terminators {
self.instructions.push(instr);
}
true
}
}
impl fmt::Display for SsaBlock {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
writeln!(f, "B{}:", self.id)?;
for phi in &self.phi_nodes {
writeln!(f, " {phi}")?;
}
for instr in &self.instructions {
writeln!(f, " {instr}")?;
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
analysis::{
ssa::{PhiNode, PhiOperand, SsaInstruction, SsaOp, SsaVarId, VariableOrigin},
SsaFunctionBuilder,
},
assembly::{FlowType, Instruction, InstructionCategory, Operand, StackBehavior},
};
fn make_test_cil_instruction(mnemonic: &'static str, pops: u8, pushes: u8) -> Instruction {
Instruction {
rva: 0x1000,
offset: 0,
size: 1,
opcode: 0x00,
prefix: 0,
mnemonic,
category: InstructionCategory::Arithmetic,
flow_type: FlowType::Sequential,
operand: Operand::None,
stack_behavior: StackBehavior {
pops,
pushes,
net_effect: i8::try_from(i16::from(pushes) - i16::from(pops)).unwrap_or(0),
},
branch_targets: vec![],
}
}
#[test]
fn test_ssa_block_creation() {
let block = SsaBlock::new(5);
assert_eq!(block.id(), 5);
assert!(block.is_empty());
assert!(block.has_no_phis());
assert!(block.has_no_instructions());
}
#[test]
fn test_ssa_block_with_capacity() {
let block = SsaBlock::with_capacity(0, 2, 10);
assert_eq!(block.id(), 0);
assert!(block.is_empty());
}
#[test]
fn test_ssa_block_add_phi() {
let mut block = SsaBlock::new(0);
let result = SsaVarId::from_index(0);
let v1 = SsaVarId::from_index(1);
let v2 = SsaVarId::from_index(2);
let mut phi = PhiNode::new(result, VariableOrigin::Local(0));
phi.add_operand(PhiOperand::new(v1, 0));
phi.add_operand(PhiOperand::new(v2, 1));
block.add_phi(phi);
assert!(!block.has_no_phis());
assert_eq!(block.phi_count(), 1);
assert!(block.phi(0).is_some());
assert_eq!(block.phi(0).unwrap().result(), result);
}
#[test]
fn test_ssa_block_add_instruction() {
let mut block = SsaBlock::new(0);
let v0 = SsaVarId::from_index(0);
let v1 = SsaVarId::from_index(1);
let v2 = SsaVarId::from_index(2);
let cil = make_test_cil_instruction("add", 2, 1);
let instr = SsaInstruction::new(
cil,
SsaOp::Add {
dest: v2,
left: v0,
right: v1,
},
);
block.add_instruction(instr);
assert!(!block.has_no_instructions());
assert_eq!(block.instruction_count(), 1);
assert!(block.instruction(0).is_some());
}
#[test]
fn test_ssa_block_phi_access() {
let mut block = SsaBlock::new(0);
let r1 = SsaVarId::from_index(0);
let r2 = SsaVarId::from_index(1);
block.add_phi(PhiNode::new(r1, VariableOrigin::Local(0)));
block.add_phi(PhiNode::new(r2, VariableOrigin::Local(1)));
assert_eq!(block.phi_count(), 2);
assert!(block.phi(0).is_some());
assert!(block.phi(1).is_some());
assert!(block.phi(2).is_none());
}
#[test]
fn test_ssa_block_instruction_access() {
let mut block = SsaBlock::new(0);
let cil1 = make_test_cil_instruction("nop", 0, 0);
let cil2 = make_test_cil_instruction("ret", 0, 0);
block.add_instruction(SsaInstruction::new(cil1, SsaOp::Nop));
block.add_instruction(SsaInstruction::new(cil2, SsaOp::Return { value: None }));
assert_eq!(block.instruction_count(), 2);
assert!(block.instruction(0).is_some());
assert!(block.instruction(1).is_some());
assert!(block.instruction(2).is_none());
}
#[test]
fn test_ssa_block_find_phi_defining() {
let mut block = SsaBlock::new(0);
let r1 = SsaVarId::from_index(0);
let r2 = SsaVarId::from_index(1);
let other = SsaVarId::from_index(2);
block.add_phi(PhiNode::new(r1, VariableOrigin::Local(0)));
block.add_phi(PhiNode::new(r2, VariableOrigin::Local(1)));
assert!(block.find_phi_defining(r1).is_some());
assert!(block.find_phi_defining(r2).is_some());
assert!(block.find_phi_defining(other).is_none());
}
#[test]
fn test_ssa_block_defined_variables() {
let mut block = SsaBlock::new(0);
let phi_result = SsaVarId::from_index(0);
let v0 = SsaVarId::from_index(1);
let v1 = SsaVarId::from_index(2);
let instr_result = SsaVarId::from_index(3);
let v2 = SsaVarId::from_index(4);
block.add_phi(PhiNode::new(phi_result, VariableOrigin::Local(0)));
let cil = make_test_cil_instruction("add", 2, 1);
let instr = SsaInstruction::new(
cil,
SsaOp::Add {
dest: instr_result,
left: v0,
right: v1,
},
);
block.add_instruction(instr);
let cil2 = make_test_cil_instruction("pop", 1, 0);
block.add_instruction(SsaInstruction::new(cil2, SsaOp::Pop { value: v2 }));
let defs: Vec<_> = block.defined_variables().collect();
assert_eq!(defs.len(), 2);
assert!(defs.contains(&phi_result));
assert!(defs.contains(&instr_result));
}
#[test]
fn test_ssa_block_used_variables() {
let mut block = SsaBlock::new(0);
let phi_result = SsaVarId::from_index(0);
let phi_op1 = SsaVarId::from_index(1);
let phi_op2 = SsaVarId::from_index(2);
let instr_op1 = SsaVarId::from_index(3);
let instr_op2 = SsaVarId::from_index(4);
let instr_result = SsaVarId::from_index(5);
let mut phi = PhiNode::new(phi_result, VariableOrigin::Local(0));
phi.add_operand(PhiOperand::new(phi_op1, 0));
phi.add_operand(PhiOperand::new(phi_op2, 1));
block.add_phi(phi);
let cil = make_test_cil_instruction("add", 2, 1);
let instr = SsaInstruction::new(
cil,
SsaOp::Add {
dest: instr_result,
left: instr_op1,
right: instr_op2,
},
);
block.add_instruction(instr);
let uses: Vec<_> = block.used_variables().collect();
assert_eq!(uses.len(), 4);
assert!(uses.contains(&phi_op1));
assert!(uses.contains(&phi_op2));
assert!(uses.contains(&instr_op1));
assert!(uses.contains(&instr_op2));
}
#[test]
fn test_ssa_block_display_empty() {
let block = SsaBlock::new(3);
let display = format!("{block}");
assert_eq!(display, "B3:\n");
}
#[test]
fn test_ssa_block_display_with_content() {
let mut block = SsaBlock::new(1);
let mut phi = PhiNode::new(SsaVarId::from_index(3), VariableOrigin::Local(0));
phi.add_operand(PhiOperand::new(SsaVarId::from_index(1), 0));
phi.add_operand(PhiOperand::new(SsaVarId::from_index(2), 2));
block.add_phi(phi);
let cil = make_test_cil_instruction("add", 2, 1);
let instr = SsaInstruction::new(
cil,
SsaOp::Add {
dest: SsaVarId::from_index(5),
left: SsaVarId::from_index(3),
right: SsaVarId::from_index(4),
},
);
block.add_instruction(instr);
let display = format!("{block}");
assert!(display.contains("B1:"));
assert!(display.contains("v3 = phi(v1 from B0, v2 from B2)"));
assert!(display.contains("v5 = add v3, v4"));
}
#[test]
fn test_ssa_block_mutable_access() {
let mut block = SsaBlock::new(0);
let result = SsaVarId::from_index(0);
let operand = SsaVarId::from_index(1);
block.add_phi(PhiNode::new(result, VariableOrigin::Local(0)));
if let Some(phi) = block.phi_mut(0) {
phi.add_operand(PhiOperand::new(operand, 1));
}
assert_eq!(block.phi(0).unwrap().operand_count(), 1);
}
#[test]
fn test_is_trampoline_unconditional_jump() {
let ssa = SsaFunctionBuilder::new(2, 0)
.build_with(|f| {
f.block(0, |b| b.jump(1));
f.block(1, |b| b.ret());
})
.unwrap();
assert_eq!(ssa.block(0).unwrap().is_trampoline(), Some(1));
assert_eq!(ssa.block(1).unwrap().is_trampoline(), None);
}
#[test]
fn test_is_trampoline_leave_instruction() {
let ssa = SsaFunctionBuilder::new(2, 0)
.build_with(|f| {
f.block(0, |b| b.leave(1));
f.block(1, |b| b.ret());
})
.unwrap();
assert_eq!(ssa.block(0).unwrap().is_trampoline(), Some(1));
}
#[test]
fn test_is_trampoline_blocked_by_phi_nodes() {
let mut ssa = SsaFunctionBuilder::new(2, 0)
.build_with(|f| {
f.block(0, |b| b.jump(1));
f.block(1, |b| b.ret());
})
.unwrap();
if let Some(block) = ssa.block_mut(0) {
block.add_phi(PhiNode::new(
SsaVarId::from_index(0),
VariableOrigin::Local(0),
));
}
assert_eq!(ssa.block(0).unwrap().is_trampoline(), None);
}
#[test]
fn test_is_trampoline_blocked_by_multiple_instructions() {
let ssa = SsaFunctionBuilder::new(2, 0)
.build_with(|f| {
f.block(0, |b| {
let _ = b.const_i32(42); b.jump(1);
});
f.block(1, |b| b.ret());
})
.unwrap();
assert_eq!(ssa.block(0).unwrap().is_trampoline(), None);
}
#[test]
fn test_is_trampoline_conditional_branch_not_trampoline() {
let ssa = SsaFunctionBuilder::new(2, 0)
.build_with(|f| {
f.block(0, |b| {
let cond = b.const_true();
b.branch(cond, 1, 1);
});
f.block(1, |b| b.ret());
})
.unwrap();
assert_eq!(ssa.block(0).unwrap().is_trampoline(), None);
}
}