use super::*;
use crate::function::{BytecodeFunction, BytecodeFunctionConstant};
use std::collections::{HashMap, HashSet, VecDeque};
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum ConstantLattice {
#[default]
Undetermined,
NotConstant,
VmConstant(ConstantIndex),
Immediate(BytecodeImmediate),
}
impl ConstantLattice {
fn merge(self, other: Self) -> Self {
match (self, other) {
(Self::Undetermined, value) | (value, Self::Undetermined) => value,
(left, right) if left == right => left,
_ => Self::NotConstant,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ConditionState {
AlwaysFalse,
AlwaysTrue,
Unknown,
}
#[derive(Debug, Clone, Copy)]
struct JumpTarget {
dead: bool,
block: BytecodeBlockId,
}
pub struct Sccp<'function, 'table> {
function: &'function mut BytecodeFunction<'table>,
lattice: HashMap<BytecodeOperand, ConstantLattice>,
block_uses: HashMap<BytecodeBlockId, HashSet<BytecodeBlockId>>,
flow_worklist: VecDeque<BytecodeBlockId>,
visited_blocks: HashSet<BytecodeBlockId>,
ssa_worklist: VecDeque<BytecodeOperand>,
}
impl<'function, 'table> Sccp<'function, 'table> {
pub fn new(function: &'function mut BytecodeFunction<'table>) -> Self {
Self {
function,
lattice: HashMap::new(),
block_uses: HashMap::new(),
flow_worklist: VecDeque::new(),
visited_blocks: HashSet::new(),
ssa_worklist: VecDeque::new(),
}
}
pub fn lattice(&self, operand: BytecodeOperand) -> ConstantLattice {
self.operand_lattice(operand)
}
pub fn propagate(&mut self) {
let entry = self.function.entry_block;
let exit = self.function.exit_block;
self.block_uses.entry(entry).or_default().insert(entry);
self.block_uses.entry(exit).or_default().insert(exit);
self.flow_worklist.push_back(entry);
while !self.flow_worklist.is_empty() || !self.ssa_worklist.is_empty() {
while let Some(block) = self.flow_worklist.pop_front() {
if !self.visited_blocks.insert(block) {
continue;
}
let phis = self.function.block(block).phis.clone();
for phi in phis {
self.visit_phi(phi);
}
let instructions = self.function.block(block).instruction_ids.clone();
for instruction in instructions.iter().copied() {
self.visit_instruction(instruction);
}
let ends_with_branch = instructions
.last()
.is_some_and(|id| !self.jump_targets(*id).is_empty());
let successors = self.function.block(block).successors.clone();
for edge in successors {
if edge.kind != BytecodeEdgeKind::Fallthrough {
continue;
}
if ends_with_branch
&& !self
.block_uses
.get(&edge.target)
.is_some_and(|uses| uses.contains(&block))
{
continue;
}
self.mark_flow_edge(block, edge.target);
}
}
while let Some(operand) = self.ssa_worklist.pop_front() {
match operand {
BytecodeOperand::Instruction(id) => self.visit_instruction(id),
BytecodeOperand::Phi(id) => self.visit_phi(id),
_ => {}
}
}
}
}
pub fn rewrite(&mut self) {
self.rewrite_arithmetic_constants();
self.replace_folded_instructions();
self.simplify_phis();
self.update_block_reachability();
rebuild_uses(&mut self.function.instructions, &mut self.function.phis);
}
fn operand_lattice(&self, operand: BytecodeOperand) -> ConstantLattice {
if matches!(
operand,
BytecodeOperand::Projection(_)
| BytecodeOperand::VmRegister(_)
| BytecodeOperand::VmUpvalue(_)
) {
ConstantLattice::NotConstant
} else {
self.lattice.get(&operand).copied().unwrap_or_default()
}
}
fn unknown_condition(&self, operands: &[BytecodeOperand]) -> ConstantLattice {
if operands
.iter()
.any(|operand| self.operand_lattice(*operand) == ConstantLattice::NotConstant)
{
ConstantLattice::NotConstant
} else {
ConstantLattice::Undetermined
}
}
fn visit_phi(&mut self, id: BytecodePhiId) {
let mut value = ConstantLattice::Undetermined;
for operand in self.function.phi(id).operands.iter().copied() {
value = self.operand_lattice(operand).merge(value);
}
let operand = BytecodeOperand::Phi(id);
let previous = self.lattice.get(&operand).copied().unwrap_or_default();
if value != previous {
self.ssa_worklist
.extend(self.function.phi(id).users.iter().copied());
self.lattice.insert(operand, value);
}
}
fn visit_instruction(&mut self, id: BytecodeInstructionId) {
let instruction = self.function.graph_instruction(id).clone();
if instruction.opcode == Opcode::Capture
&& instruction.operands.len() >= 2
&& self.immediate(instruction.operands[0]) == Some(BytecodeImmediate::Int(1))
{
let source = instruction.operands[1];
if matches!(
source,
BytecodeOperand::Instruction(_) | BytecodeOperand::Phi(_)
) && self.operand_lattice(source) != ConstantLattice::NotConstant
{
self.lattice.insert(source, ConstantLattice::NotConstant);
self.ssa_worklist.extend(self.users(source));
}
}
let operand = BytecodeOperand::Instruction(id);
let value = self.evaluate_instruction(&instruction);
let previous = self.lattice.get(&operand).copied().unwrap_or_default();
let merged = value.merge(previous);
if merged != previous {
self.ssa_worklist.extend(instruction.users.iter().copied());
}
for target in self.jump_targets(id) {
if !target.dead {
self.mark_flow_edge(instruction.block, target.block);
}
}
self.lattice.insert(operand, merged);
}
fn mark_flow_edge(&mut self, source: BytecodeBlockId, target: BytecodeBlockId) {
self.block_uses.entry(target).or_default().insert(source);
if !self.visited_blocks.contains(&target) {
self.flow_worklist.push_back(target);
}
}
fn evaluate_instruction(&mut self, instruction: &BytecodeInstruction) -> ConstantLattice {
match instruction.opcode {
Opcode::LoadK | Opcode::LoadKx => match instruction.operands.first().copied() {
Some(BytecodeOperand::VmConstant(index)) => ConstantLattice::VmConstant(index),
_ => ConstantLattice::NotConstant,
},
Opcode::LoadB | Opcode::LoadN => instruction
.operands
.first()
.and_then(|operand| self.immediate(*operand))
.map(ConstantLattice::Immediate)
.unwrap_or(ConstantLattice::NotConstant),
Opcode::LoadNil => ConstantLattice::VmConstant(
self.find_or_add_constant(BytecodeFunctionConstant::Nil),
),
Opcode::Add
| Opcode::Sub
| Opcode::Mul
| Opcode::Div
| Opcode::Mod
| Opcode::Pow
| Opcode::IDiv => self.evaluate_arithmetic(instruction),
Opcode::Move => self.operand_lattice(instruction.operands[0]),
Opcode::JumpIf | Opcode::JumpIfNot => {
let condition = self.evaluate_condition(instruction.operands[0]);
if condition == ConditionState::Unknown {
return self.unknown_condition(&instruction.operands[..1]);
}
let jumps_on_true = instruction.opcode == Opcode::JumpIf;
ConstantLattice::Immediate(BytecodeImmediate::Boolean(
(condition == ConditionState::AlwaysTrue) == jumps_on_true,
))
}
Opcode::JumpIfEq
| Opcode::JumpIfLe
| Opcode::JumpIfLt
| Opcode::JumpIfNotEq
| Opcode::JumpIfNotLe
| Opcode::JumpIfNotLt => {
let condition = self.evaluate_comparison(
instruction.opcode,
instruction.operands[0],
instruction.operands[1],
);
if condition == ConditionState::Unknown {
return self.unknown_condition(&instruction.operands[..2]);
}
let negated = matches!(
instruction.opcode,
Opcode::JumpIfNotEq | Opcode::JumpIfNotLe | Opcode::JumpIfNotLt
);
ConstantLattice::Immediate(BytecodeImmediate::Boolean(
(condition == ConditionState::AlwaysTrue) != negated,
))
}
Opcode::JumpXEqKNil | Opcode::JumpXEqKB | Opcode::JumpXEqKN | Opcode::JumpXEqKS => {
let condition = self.evaluate_constant_comparison(instruction);
if condition == ConditionState::Unknown {
return self.unknown_condition(&instruction.operands[..1]);
}
let negated = self
.immediate(instruction.operands[1])
.is_some_and(|value| value == BytecodeImmediate::Boolean(true));
ConstantLattice::Immediate(BytecodeImmediate::Boolean(
(condition == ConditionState::AlwaysTrue) != negated,
))
}
_ => ConstantLattice::NotConstant,
}
}
fn evaluate_arithmetic(&mut self, instruction: &BytecodeInstruction) -> ConstantLattice {
let left = self.operand_lattice(instruction.operands[0]);
let right = self.operand_lattice(instruction.operands[1]);
match (left, right) {
(
ConstantLattice::Immediate(BytecodeImmediate::Int(left)),
ConstantLattice::Immediate(BytecodeImmediate::Int(right)),
) => {
if right == 0
&& matches!(instruction.opcode, Opcode::Div | Opcode::Mod | Opcode::IDiv)
{
return ConstantLattice::NotConstant;
}
if matches!(instruction.opcode, Opcode::Div | Opcode::Pow) {
return ConstantLattice::NotConstant;
}
let value = match instruction.opcode {
Opcode::Add => i64::from(left) + i64::from(right),
Opcode::Sub => i64::from(left) - i64::from(right),
Opcode::Mul => i64::from(left) * i64::from(right),
Opcode::Mod => {
let mut value = i64::from(left) % i64::from(right);
if value != 0 && (left < 0) != (right < 0) {
value += i64::from(right);
}
value
}
Opcode::IDiv => i64::from(left).div_euclid(i64::from(right)),
_ => return ConstantLattice::NotConstant,
};
i16::try_from(value)
.ok()
.map(|value| {
ConstantLattice::Immediate(BytecodeImmediate::Int(i32::from(value)))
})
.unwrap_or(ConstantLattice::NotConstant)
}
(ConstantLattice::VmConstant(left), ConstantLattice::VmConstant(right)) => {
let (
Some(BytecodeFunctionConstant::Number(left)),
Some(BytecodeFunctionConstant::Number(right)),
) = (self.constant(left), self.constant(right))
else {
return ConstantLattice::NotConstant;
};
let value = match instruction.opcode {
Opcode::Add => left + right,
Opcode::Sub => left - right,
Opcode::Mul => left * right,
Opcode::Div => left / right,
Opcode::Mod if *right != 0.0 => left - (left / right).floor() * right,
Opcode::Pow => left.powf(*right),
Opcode::IDiv => (left / right).floor(),
_ => return ConstantLattice::NotConstant,
};
ConstantLattice::VmConstant(
self.find_or_add_constant(BytecodeFunctionConstant::Number(value)),
)
}
(ConstantLattice::Undetermined, ConstantLattice::Undetermined) => {
ConstantLattice::Undetermined
}
_ => ConstantLattice::NotConstant,
}
}
fn evaluate_condition(&self, operand: BytecodeOperand) -> ConditionState {
match self.operand_lattice(operand) {
ConstantLattice::VmConstant(index) => {
if self.constant_is_falsey(index) {
ConditionState::AlwaysFalse
} else {
ConditionState::AlwaysTrue
}
}
ConstantLattice::Immediate(BytecodeImmediate::Boolean(value)) => {
if value {
ConditionState::AlwaysTrue
} else {
ConditionState::AlwaysFalse
}
}
_ => ConditionState::Unknown,
}
}
fn evaluate_comparison(
&self,
opcode: Opcode,
left: BytecodeOperand,
right: BytecodeOperand,
) -> ConditionState {
let left = self.operand_lattice(left);
let right = self.operand_lattice(right);
let ordering = matches!(
opcode,
Opcode::JumpIfLt | Opcode::JumpIfLe | Opcode::JumpIfNotLt | Opcode::JumpIfNotLe
);
if ordering && (!self.is_orderable(left) || !self.is_orderable(right)) {
return ConditionState::Unknown;
}
let comparison = match (left, right) {
(ConstantLattice::VmConstant(left), ConstantLattice::VmConstant(right)) => {
let (Some(left), Some(right)) = (self.constant(left), self.constant(right)) else {
return ConditionState::Unknown;
};
if std::mem::discriminant(left) != std::mem::discriminant(right) {
return ConditionState::Unknown;
}
compare_constants(left, right)
}
(ConstantLattice::Immediate(left), ConstantLattice::Immediate(right)) => {
compare_immediates(left, right)
}
(ConstantLattice::VmConstant(left), ConstantLattice::Immediate(right)) => self
.constant(left)
.and_then(|left| compare_constant_immediate(left, right)),
(ConstantLattice::Immediate(left), ConstantLattice::VmConstant(right)) => self
.constant(right)
.and_then(|right| compare_constant_immediate(right, left))
.map(|ordering| ordering.reverse()),
_ => None,
};
let Some(comparison) = comparison else {
return ConditionState::Unknown;
};
let value = match opcode {
Opcode::JumpIfEq | Opcode::JumpIfNotEq => comparison.is_eq(),
Opcode::JumpIfLt | Opcode::JumpIfNotLt => comparison.is_lt(),
Opcode::JumpIfLe | Opcode::JumpIfNotLe => comparison.is_le(),
_ => return ConditionState::Unknown,
};
if value {
ConditionState::AlwaysTrue
} else {
ConditionState::AlwaysFalse
}
}
fn evaluate_constant_comparison(
&mut self,
instruction: &BytecodeInstruction,
) -> ConditionState {
let value = self.operand_lattice(instruction.operands[0]);
let equal = match instruction.opcode {
Opcode::JumpXEqKNil => match value {
ConstantLattice::VmConstant(index) => self
.constant(index)
.map(|constant| matches!(constant, BytecodeFunctionConstant::Nil)),
ConstantLattice::Immediate(_) => Some(false),
_ => None,
},
Opcode::JumpXEqKB => {
let expected = self.immediate(instruction.operands[3]);
match (value, expected) {
(
ConstantLattice::Immediate(BytecodeImmediate::Boolean(value)),
Some(BytecodeImmediate::Boolean(expected)),
) => Some(value == expected),
(ConstantLattice::VmConstant(index), Some(expected)) => self
.constant(index)
.and_then(|value| constant_equals_immediate(value, expected)),
_ => None,
}
}
Opcode::JumpXEqKN | Opcode::JumpXEqKS => {
let BytecodeOperand::VmConstant(expected) = instruction.operands[3] else {
return ConditionState::Unknown;
};
match value {
ConstantLattice::VmConstant(value) => {
constants_equal(self.constant(value), self.constant(expected))
}
ConstantLattice::Immediate(immediate) => self
.constant(expected)
.and_then(|constant| constant_equals_immediate(constant, immediate)),
_ => None,
}
}
_ => None,
};
match equal {
Some(true) => ConditionState::AlwaysTrue,
Some(false) => ConditionState::AlwaysFalse,
None => ConditionState::Unknown,
}
}
fn jump_targets(&mut self, id: BytecodeInstructionId) -> Vec<JumpTarget> {
let instruction = self.function.graph_instruction(id).clone();
match instruction.opcode {
Opcode::Jump | Opcode::JumpBack => instruction
.operands
.first()
.and_then(|operand| block_operand(*operand))
.map(|block| vec![JumpTarget { dead: false, block }])
.unwrap_or_default(),
Opcode::JumpIf | Opcode::JumpIfNot => {
let condition = self.evaluate_condition(instruction.operands[0]);
self.conditional_targets(
&instruction,
instruction.operands[1],
condition,
instruction.opcode == Opcode::JumpIf,
)
}
Opcode::JumpIfEq
| Opcode::JumpIfLe
| Opcode::JumpIfLt
| Opcode::JumpIfNotEq
| Opcode::JumpIfNotLe
| Opcode::JumpIfNotLt => {
let condition = self.evaluate_comparison(
instruction.opcode,
instruction.operands[0],
instruction.operands[1],
);
let negated = matches!(
instruction.opcode,
Opcode::JumpIfNotEq | Opcode::JumpIfNotLe | Opcode::JumpIfNotLt
);
self.conditional_targets(&instruction, instruction.operands[2], condition, !negated)
}
Opcode::JumpXEqKNil | Opcode::JumpXEqKB | Opcode::JumpXEqKN | Opcode::JumpXEqKS => {
let condition = self.evaluate_constant_comparison(&instruction);
let negated = self
.immediate(instruction.operands[1])
.is_some_and(|value| value == BytecodeImmediate::Boolean(true));
self.conditional_targets(&instruction, instruction.operands[2], condition, !negated)
}
Opcode::ForNPrep
| Opcode::ForNLoop
| Opcode::ForGPrep
| Opcode::ForGPrepNext
| Opcode::ForGPrepInext => self.conditional_targets(
&instruction,
instruction.operands[3],
ConditionState::Unknown,
true,
),
Opcode::ForGLoop => self.conditional_targets(
&instruction,
instruction.operands[5],
ConditionState::Unknown,
true,
),
Opcode::CmpProto => self.conditional_targets(
&instruction,
instruction.operands[2],
ConditionState::Unknown,
true,
),
_ => Vec::new(),
}
}
fn conditional_targets(
&self,
instruction: &BytecodeInstruction,
target: BytecodeOperand,
condition: ConditionState,
target_taken_on_true: bool,
) -> Vec<JumpTarget> {
let Some(target) = block_operand(target) else {
return Vec::new();
};
let Some(fallthrough) = self
.function
.block(instruction.block)
.successors
.iter()
.find(|edge| edge.kind == BytecodeEdgeKind::Fallthrough)
.map(|edge| edge.target)
else {
return Vec::new();
};
let (target_dead, fallthrough_dead) = match condition {
ConditionState::AlwaysTrue => (!target_taken_on_true, target_taken_on_true),
ConditionState::AlwaysFalse => (target_taken_on_true, !target_taken_on_true),
ConditionState::Unknown => (false, false),
};
vec![
JumpTarget {
dead: target_dead,
block: target,
},
JumpTarget {
dead: fallthrough_dead,
block: fallthrough,
},
]
}
fn rewrite_arithmetic_constants(&mut self) {
for block_index in 0..self.function.blocks.len() {
let block = BytecodeBlockId::new(block_index);
if self.block_uses.get(&block).is_none_or(HashSet::is_empty) {
continue;
}
let instructions = self.function.block(block).instruction_ids.clone();
for id in instructions {
let instruction = self.function.graph_instruction(id).clone();
let Some(mut opcode) = arithmetic_constant_opcode(instruction.opcode) else {
continue;
};
if instruction.operands.len() != 2 {
continue;
}
let left = instruction.operands[0];
let right = instruction.operands[1];
let left_lattice = self.operand_lattice(left);
let right_lattice = self.operand_lattice(right);
let is_number = |lattice| match lattice {
ConstantLattice::VmConstant(index) => {
self.constant(index).is_some_and(|constant| {
matches!(constant, BytecodeFunctionConstant::Number(_))
})
}
_ => false,
};
let (nonconstant, constant, old_constant, reverse) =
if is_number(right_lattice) && left_lattice == ConstantLattice::NotConstant {
(left, right_lattice, right, false)
} else if is_number(left_lattice)
&& right_lattice == ConstantLattice::NotConstant
&& matches!(
instruction.opcode,
Opcode::Add | Opcode::Mul | Opcode::Sub | Opcode::Div
)
{
if instruction.opcode == Opcode::Sub {
opcode = Opcode::SubRK;
} else if instruction.opcode == Opcode::Div {
opcode = Opcode::DivRK;
}
(
right,
left_lattice,
left,
matches!(instruction.opcode, Opcode::Sub | Opcode::Div),
)
} else {
continue;
};
let ConstantLattice::VmConstant(constant_index) = constant else {
continue;
};
let Some(BytecodeFunctionConstant::Number(number)) = self.constant(constant_index)
else {
continue;
};
let number = *number;
let replacement = if number == 0.0 && !reverse {
match instruction.opcode {
Opcode::Add | Opcode::Sub => Some((Opcode::Move, vec![nonconstant])),
Opcode::Mul => Some((
Opcode::LoadN,
vec![self.add_immediate(BytecodeImmediate::Int(0))],
)),
Opcode::Pow => Some((
Opcode::LoadN,
vec![self.add_immediate(BytecodeImmediate::Int(1))],
)),
_ => None,
}
} else if number == 1.0 && !reverse {
match instruction.opcode {
Opcode::Mul | Opcode::Pow | Opcode::Div => {
Some((Opcode::Move, vec![nonconstant]))
}
_ => None,
}
} else {
None
};
let (opcode, operands) = replacement.unwrap_or_else(|| {
if reverse {
(
opcode,
vec![BytecodeOperand::VmConstant(constant_index), nonconstant],
)
} else {
(
opcode,
vec![nonconstant, BytecodeOperand::VmConstant(constant_index)],
)
}
});
self.function.instructions[id.index()].opcode = opcode;
self.set_instruction_operands(id, operands);
self.erase_dead_producer(old_constant);
}
}
}
fn replace_folded_instructions(&mut self) {
let folded = self
.lattice
.iter()
.filter_map(|(operand, lattice)| match (operand, lattice) {
(BytecodeOperand::Instruction(id), ConstantLattice::VmConstant(_))
| (BytecodeOperand::Instruction(id), ConstantLattice::Immediate(_)) => {
Some((*id, *lattice))
}
_ => None,
})
.collect::<Vec<_>>();
for (id, lattice) in folded {
let opcode = self.function.graph_instruction(id).opcode;
if matches!(
opcode,
Opcode::LoadK | Opcode::LoadKx | Opcode::LoadN | Opcode::LoadB | Opcode::LoadNil
) {
continue;
}
if opcode.is_jump_d() {
self.remove_dead_edges(id);
self.erase_instruction(id);
} else {
let (opcode, operand) = match lattice {
ConstantLattice::VmConstant(index) => {
(Opcode::LoadK, BytecodeOperand::VmConstant(index))
}
ConstantLattice::Immediate(immediate) => {
let opcode = if matches!(immediate, BytecodeImmediate::Boolean(_)) {
Opcode::LoadB
} else {
Opcode::LoadN
};
(opcode, self.add_immediate(immediate))
}
_ => continue,
};
self.function.instructions[id.index()].opcode = opcode;
self.set_instruction_operands(id, vec![operand]);
}
}
}
fn remove_dead_edges(&mut self, id: BytecodeInstructionId) {
let targets = self.jump_targets(id);
let block = self.function.graph_instruction(id).block;
let mut live = None;
for target in targets {
if target.dead {
self.function.blocks[block.index()]
.successors
.retain(|edge| edge.target != target.block);
} else {
live = Some(target.block);
}
}
if let Some(live) = live {
if let Some(edge) = self.function.blocks[block.index()]
.successors
.iter_mut()
.find(|edge| edge.target == live)
{
edge.kind = BytecodeEdgeKind::Fallthrough;
} else {
self.function.blocks[block.index()]
.successors
.push(BytecodeEdge::new(BytecodeEdgeKind::Fallthrough, live));
}
}
}
fn simplify_phis(&mut self) {
let blocks = self.visited_blocks.iter().copied().collect::<Vec<_>>();
for block in blocks {
let phis = self.function.block(block).phis.clone();
for id in phis {
let operands = self.function.phi(id).operands.clone();
let Some(unique) = operands.first().copied() else {
continue;
};
if operands.iter().all(|operand| *operand == unique) {
let phi = BytecodeOperand::Phi(id);
for user in self.users(phi) {
self.replace_user_operand(user, phi, unique);
}
self.function.phis[id.index()].users.clear();
self.function.blocks[block.index()]
.phis
.retain(|phi| *phi != id);
}
}
}
}
fn update_block_reachability(&mut self) {
let mut reachable = HashSet::new();
let mut worklist = vec![self.function.entry_block];
reachable.insert(self.function.entry_block);
reachable.insert(self.function.exit_block);
while let Some(block) = worklist.pop() {
for edge in self.function.block(block).successors.iter() {
if reachable.insert(edge.target) {
worklist.push(edge.target);
}
}
}
for (index, block) in self.function.blocks.iter_mut().enumerate() {
let id = BytecodeBlockId::new(index);
block.use_count = self.block_uses.get(&id).map_or(0, |uses| uses.len() as u32);
block.dead = !reachable.contains(&id);
}
}
fn set_instruction_operands(
&mut self,
id: BytecodeInstructionId,
operands: Vec<BytecodeOperand>,
) {
let user = BytecodeOperand::Instruction(id);
let old = std::mem::replace(
&mut self.function.instructions[id.index()].operands,
operands,
);
for operand in old {
self.erase_use(operand, user);
}
let new = self.function.instructions[id.index()].operands.clone();
for operand in new {
record_use(
&mut self.function.instructions,
&mut self.function.phis,
operand,
user,
);
}
}
fn replace_user_operand(
&mut self,
user: BytecodeOperand,
old: BytecodeOperand,
new: BytecodeOperand,
) {
match user {
BytecodeOperand::Instruction(id) => {
for operand in &mut self.function.instructions[id.index()].operands {
if *operand == old {
*operand = new;
}
}
}
BytecodeOperand::Phi(id) => {
for operand in &mut self.function.phis[id.index()].operands {
if *operand == old {
*operand = new;
}
}
}
_ => return,
}
self.erase_use(old, user);
record_use(
&mut self.function.instructions,
&mut self.function.phis,
new,
user,
);
}
fn erase_use(&mut self, used: BytecodeOperand, user: BytecodeOperand) {
match used {
BytecodeOperand::Instruction(id) => self.function.instructions[id.index()]
.users
.retain(|candidate| *candidate != user),
BytecodeOperand::Phi(id) => self.function.phis[id.index()]
.users
.retain(|candidate| *candidate != user),
_ => {}
}
}
fn erase_instruction(&mut self, id: BytecodeInstructionId) {
let block = self.function.graph_instruction(id).block;
self.function.blocks[block.index()]
.instruction_ids
.retain(|candidate| *candidate != id);
}
fn erase_dead_producer(&mut self, operand: BytecodeOperand) {
let BytecodeOperand::Instruction(id) = operand else {
return;
};
let instruction = self.function.graph_instruction(id);
if !matches!(
instruction.opcode,
Opcode::LoadK
| Opcode::LoadKx
| Opcode::LoadN
| Opcode::LoadB
| Opcode::LoadNil
| Opcode::GetUpval
) || !instruction.users.is_empty()
{
return;
}
self.erase_instruction(id);
}
fn users(&self, operand: BytecodeOperand) -> Vec<BytecodeOperand> {
match operand {
BytecodeOperand::Instruction(id) => {
self.function.instructions[id.index()].users.clone()
}
BytecodeOperand::Phi(id) => self.function.phis[id.index()].users.clone(),
_ => Vec::new(),
}
}
fn immediate(&self, operand: BytecodeOperand) -> Option<BytecodeImmediate> {
let BytecodeOperand::Immediate(id) = operand else {
return None;
};
self.function.immediates.get(id.index()).copied()
}
fn add_immediate(&mut self, immediate: BytecodeImmediate) -> BytecodeOperand {
if let Some(index) = self
.function
.immediates
.iter()
.position(|existing| *existing == immediate)
{
BytecodeOperand::Immediate(BytecodeImmediateId::new(index))
} else {
self.function.immediates.push(immediate);
BytecodeOperand::Immediate(BytecodeImmediateId::new(self.function.immediates.len() - 1))
}
}
fn constant(&self, index: ConstantIndex) -> Option<&BytecodeFunctionConstant<'table>> {
usize::try_from(index)
.ok()
.and_then(|index| self.function.constants.get(index))
}
fn find_or_add_constant(
&mut self,
constant: BytecodeFunctionConstant<'table>,
) -> ConstantIndex {
if let Some(index) = self
.function
.constants
.iter()
.position(|existing| *existing == constant)
{
index as ConstantIndex
} else {
self.function.constants.push(constant);
(self.function.constants.len() - 1) as ConstantIndex
}
}
fn constant_is_falsey(&self, index: ConstantIndex) -> bool {
self.constant(index).is_some_and(|constant| {
matches!(
constant,
BytecodeFunctionConstant::Nil | BytecodeFunctionConstant::Boolean(false)
)
})
}
fn is_orderable(&self, lattice: ConstantLattice) -> bool {
match lattice {
ConstantLattice::VmConstant(index) => self.constant(index).is_some_and(|constant| {
matches!(
constant,
BytecodeFunctionConstant::Number(_)
| BytecodeFunctionConstant::Integer(_)
| BytecodeFunctionConstant::String(_)
)
}),
ConstantLattice::Immediate(BytecodeImmediate::Int(_)) => true,
_ => false,
}
}
}
pub fn fold_constants(function: &mut BytecodeFunction<'_>) {
let mut sccp = Sccp::new(function);
sccp.propagate();
sccp.rewrite();
}
fn block_operand(operand: BytecodeOperand) -> Option<BytecodeBlockId> {
let BytecodeOperand::Block(block) = operand else {
return None;
};
Some(block)
}
fn arithmetic_constant_opcode(opcode: Opcode) -> Option<Opcode> {
match opcode {
Opcode::Add => Some(Opcode::AddK),
Opcode::Sub => Some(Opcode::SubK),
Opcode::Mul => Some(Opcode::MulK),
Opcode::Div => Some(Opcode::DivK),
Opcode::Mod => Some(Opcode::ModK),
Opcode::Pow => Some(Opcode::PowK),
_ => None,
}
}
fn compare_constants(
left: &BytecodeFunctionConstant<'_>,
right: &BytecodeFunctionConstant<'_>,
) -> Option<std::cmp::Ordering> {
match (left, right) {
(BytecodeFunctionConstant::Number(left), BytecodeFunctionConstant::Number(right)) => {
left.partial_cmp(right)
}
(BytecodeFunctionConstant::Integer(left), BytecodeFunctionConstant::Integer(right)) => {
Some(left.cmp(right))
}
(BytecodeFunctionConstant::Boolean(left), BytecodeFunctionConstant::Boolean(right)) => {
Some(if left == right {
std::cmp::Ordering::Equal
} else {
std::cmp::Ordering::Greater
})
}
(BytecodeFunctionConstant::String(left), BytecodeFunctionConstant::String(right)) => {
Some(left.cmp(right))
}
_ => None,
}
}
fn compare_immediates(
left: BytecodeImmediate,
right: BytecodeImmediate,
) -> Option<std::cmp::Ordering> {
match (left, right) {
(BytecodeImmediate::Int(left), BytecodeImmediate::Int(right)) => Some(left.cmp(&right)),
(BytecodeImmediate::Boolean(left), BytecodeImmediate::Boolean(right)) => {
Some(if left == right {
std::cmp::Ordering::Equal
} else {
std::cmp::Ordering::Greater
})
}
_ => None,
}
}
fn compare_constant_immediate(
constant: &BytecodeFunctionConstant<'_>,
immediate: BytecodeImmediate,
) -> Option<std::cmp::Ordering> {
match (constant, immediate) {
(BytecodeFunctionConstant::Number(left), BytecodeImmediate::Int(right)) => {
left.partial_cmp(&f64::from(right))
}
(BytecodeFunctionConstant::Integer(left), BytecodeImmediate::Int(right)) => {
Some(left.cmp(&i64::from(right)))
}
(BytecodeFunctionConstant::Boolean(left), BytecodeImmediate::Boolean(right)) => {
Some(if *left == right {
std::cmp::Ordering::Equal
} else {
std::cmp::Ordering::Greater
})
}
_ => None,
}
}
fn constants_equal(
left: Option<&BytecodeFunctionConstant<'_>>,
right: Option<&BytecodeFunctionConstant<'_>>,
) -> Option<bool> {
match (left?, right?) {
(BytecodeFunctionConstant::Number(left), BytecodeFunctionConstant::Integer(right)) => {
Some(*left == *right as f64)
}
(BytecodeFunctionConstant::Integer(left), BytecodeFunctionConstant::Number(right)) => {
Some(*left as f64 == *right)
}
(left, right) => compare_constants(left, right).map(std::cmp::Ordering::is_eq),
}
}
fn constant_equals_immediate(
constant: &BytecodeFunctionConstant<'_>,
immediate: BytecodeImmediate,
) -> Option<bool> {
compare_constant_immediate(constant, immediate).map(std::cmp::Ordering::is_eq)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::model::{BytecodeTypedLocal, Register};
struct TestGraph {
function: BytecodeFunction<'static>,
}
impl TestGraph {
fn new() -> Self {
Self {
function: BytecodeFunction {
max_stack_size: 0,
num_params: 0,
upvalue_count: 0,
is_vararg: false,
flags: 0,
type_info: Vec::new(),
upvalue_types: Vec::new(),
local_types: Vec::<BytecodeTypedLocal>::new(),
blocks: Vec::new(),
instructions: Vec::new(),
constants: Vec::new(),
immediates: Vec::new(),
phis: Vec::new(),
projections: Vec::new(),
registers: HashMap::new(),
table_shapes: Vec::new(),
class_shapes: Vec::new(),
entry_block: BytecodeBlockId::new(0),
exit_block: BytecodeBlockId::new(0),
pc_to_block: Vec::new(),
pc_to_instruction: Vec::new(),
protos: Vec::new(),
line_defined: 0,
debug_name: &[],
lines: Vec::new(),
locals: Vec::new(),
upvalue_names: Vec::new(),
},
}
}
fn block(&mut self) -> BytecodeBlockId {
let id = BytecodeBlockId::new(self.function.blocks.len());
self.function
.blocks
.push(BytecodeBlock::new(InstructionPc::new(usize::MAX)));
id
}
fn connect(
&mut self,
source: BytecodeBlockId,
target: BytecodeBlockId,
kind: BytecodeEdgeKind,
) {
connect_blocks(&mut self.function.blocks, source, target, kind);
}
fn immediate(&mut self, value: BytecodeImmediate) -> BytecodeOperand {
self.function.immediates.push(value);
BytecodeOperand::Immediate(BytecodeImmediateId::new(self.function.immediates.len() - 1))
}
fn constant(&mut self, value: BytecodeFunctionConstant<'static>) -> BytecodeOperand {
self.function.constants.push(value);
BytecodeOperand::VmConstant((self.function.constants.len() - 1) as i32)
}
fn instruction(
&mut self,
block: BytecodeBlockId,
opcode: Opcode,
operands: Vec<BytecodeOperand>,
register: Option<Register>,
) -> BytecodeOperand {
let id = BytecodeInstructionId::new(self.function.instructions.len());
self.function.instructions.push(BytecodeInstruction {
pc: InstructionPc::new(usize::MAX),
opcode,
line: 0,
block,
operands,
users: Vec::new(),
});
self.function.blocks[block.index()].instruction_ids.push(id);
let operand = BytecodeOperand::Instruction(id);
if let Some(register) = register {
self.function.registers.insert(operand, register);
}
operand
}
fn phi(
&mut self,
block: BytecodeBlockId,
operands: Vec<BytecodeOperand>,
register: Register,
) -> BytecodeOperand {
let id = BytecodePhiId::new(self.function.phis.len());
self.function.phis.push(BytecodePhi {
operands,
users: Vec::new(),
});
self.function.blocks[block.index()].phis.push(id);
let operand = BytecodeOperand::Phi(id);
self.function.registers.insert(operand, register);
operand
}
fn finish(
mut self,
entry: BytecodeBlockId,
exit: BytecodeBlockId,
) -> BytecodeFunction<'static> {
self.function.entry_block = entry;
self.function.exit_block = exit;
rebuild_uses(&mut self.function.instructions, &mut self.function.phis);
self.function
}
}
fn add_return(graph: &mut TestGraph, block: BytecodeBlockId, exit: BytecodeBlockId) {
graph.instruction(block, Opcode::Return, Vec::new(), None);
graph.connect(block, exit, BytecodeEdgeKind::Fallthrough);
}
#[test]
fn does_not_fold_boolean_ordering() {
let mut graph = TestGraph::new();
let entry = graph.block();
let on_true = graph.block();
let on_false = graph.block();
let exit = graph.block();
let value = graph.immediate(BytecodeImmediate::Boolean(true));
let load = graph.instruction(entry, Opcode::LoadB, vec![value], Some(0));
graph.instruction(
entry,
Opcode::JumpIfLt,
vec![load, load, BytecodeOperand::Block(on_true)],
None,
);
graph.connect(entry, on_true, BytecodeEdgeKind::Branch);
graph.connect(entry, on_false, BytecodeEdgeKind::Fallthrough);
add_return(&mut graph, on_true, exit);
add_return(&mut graph, on_false, exit);
let mut function = graph.finish(entry, exit);
fold_constants(&mut function);
assert!(!function.block(on_true).is_dead());
assert!(!function.block(on_false).is_dead());
}
#[test]
fn folds_number_ordering() {
let mut graph = TestGraph::new();
let entry = graph.block();
let on_true = graph.block();
let on_false = graph.block();
let exit = graph.block();
let one = graph.immediate(BytecodeImmediate::Int(1));
let two = graph.immediate(BytecodeImmediate::Int(2));
let load_one = graph.instruction(entry, Opcode::LoadN, vec![one], Some(0));
let load_two = graph.instruction(entry, Opcode::LoadN, vec![two], Some(1));
graph.instruction(
entry,
Opcode::JumpIfLt,
vec![load_one, load_two, BytecodeOperand::Block(on_true)],
None,
);
graph.connect(entry, on_true, BytecodeEdgeKind::Branch);
graph.connect(entry, on_false, BytecodeEdgeKind::Fallthrough);
add_return(&mut graph, on_true, exit);
add_return(&mut graph, on_false, exit);
let mut function = graph.finish(entry, exit);
fold_constants(&mut function);
assert!(!function.block(on_true).is_dead());
assert!(function.block(on_false).is_dead());
}
#[test]
fn phi_filters_dead_predecessor() {
let mut graph = TestGraph::new();
let entry = graph.block();
let on_true = graph.block();
let on_false = graph.block();
let merge = graph.block();
let exit = graph.block();
let condition = graph.immediate(BytecodeImmediate::Boolean(true));
let condition = graph.instruction(entry, Opcode::LoadB, vec![condition], Some(0));
graph.instruction(
entry,
Opcode::JumpIf,
vec![condition, BytecodeOperand::Block(on_true)],
None,
);
graph.connect(entry, on_true, BytecodeEdgeKind::Branch);
graph.connect(entry, on_false, BytecodeEdgeKind::Fallthrough);
let forty_two = graph.constant(BytecodeFunctionConstant::Number(42.0));
let ninety_nine = graph.constant(BytecodeFunctionConstant::Number(99.0));
let load_true = graph.instruction(on_true, Opcode::LoadK, vec![forty_two], Some(1));
let load_false = graph.instruction(on_false, Opcode::LoadK, vec![ninety_nine], Some(1));
graph.connect(on_true, merge, BytecodeEdgeKind::Fallthrough);
graph.connect(on_false, merge, BytecodeEdgeKind::Fallthrough);
let phi = graph.phi(merge, vec![load_true, load_false], 1);
graph.instruction(merge, Opcode::Return, vec![phi], None);
graph.connect(merge, exit, BytecodeEdgeKind::Fallthrough);
let mut function = graph.finish(entry, exit);
let mut sccp = Sccp::new(&mut function);
sccp.propagate();
assert_eq!(sccp.lattice(phi), ConstantLattice::VmConstant(0));
sccp.rewrite();
assert!(function.block(on_false).is_dead());
}
#[test]
fn erases_trivial_phi() {
let mut graph = TestGraph::new();
let entry = graph.block();
let exit = graph.block();
let value = graph.constant(BytecodeFunctionConstant::Number(42.0));
let load = graph.instruction(entry, Opcode::LoadK, vec![value], Some(0));
graph.connect(entry, exit, BytecodeEdgeKind::Fallthrough);
let phi = graph.phi(exit, vec![load, load], 0);
graph.instruction(exit, Opcode::Return, vec![phi], None);
let mut function = graph.finish(entry, exit);
fold_constants(&mut function);
assert!(function.block(exit).phis().is_empty());
}
fn arithmetic_graph(
opcode: Opcode,
value: f64,
constant_on_left: bool,
) -> (BytecodeFunction<'static>, BytecodeBlockId, BytecodeOperand) {
let mut graph = TestGraph::new();
let entry = graph.block();
let exit = graph.block();
let constant = graph.constant(BytecodeFunctionConstant::Number(value));
let load = graph.instruction(entry, Opcode::LoadK, vec![constant], Some(0));
let upvalue = graph.instruction(
entry,
Opcode::GetUpval,
vec![BytecodeOperand::VmUpvalue(0)],
Some(1),
);
let operands = if constant_on_left {
vec![load, upvalue]
} else {
vec![upvalue, load]
};
let arithmetic = graph.instruction(entry, opcode, operands, Some(2));
graph.instruction(entry, Opcode::Return, vec![arithmetic], None);
graph.connect(entry, exit, BytecodeEdgeKind::Fallthrough);
(graph.finish(entry, exit), entry, arithmetic)
}
#[test]
fn loadk_mul_to_mulk() {
let (mut function, entry, arithmetic) = arithmetic_graph(Opcode::Mul, 42.0, true);
fold_constants(&mut function);
assert_eq!(function.block(entry).graph_instructions().len(), 3);
let BytecodeOperand::Instruction(id) = arithmetic else {
unreachable!();
};
assert_eq!(function.graph_instruction(id).opcode(), Opcode::MulK);
}
#[test]
fn loadk_div_to_divrk() {
let (mut function, entry, arithmetic) = arithmetic_graph(Opcode::Div, 42.0, true);
fold_constants(&mut function);
assert_eq!(function.block(entry).graph_instructions().len(), 3);
let BytecodeOperand::Instruction(id) = arithmetic else {
unreachable!();
};
let instruction = function.graph_instruction(id);
assert_eq!(instruction.opcode(), Opcode::DivRK);
assert_eq!(instruction.operands()[0], BytecodeOperand::VmConstant(0));
}
#[test]
fn loadk_mul_to_zero() {
let (mut function, entry, arithmetic) = arithmetic_graph(Opcode::Mul, 0.0, true);
fold_constants(&mut function);
assert_eq!(function.block(entry).graph_instructions().len(), 3);
let BytecodeOperand::Instruction(id) = arithmetic else {
unreachable!();
};
let instruction = function.graph_instruction(id);
assert_eq!(instruction.opcode(), Opcode::LoadN);
let BytecodeOperand::Immediate(value) = instruction.operands()[0] else {
unreachable!();
};
assert_eq!(*function.immediate(value), BytecodeImmediate::Int(0));
}
}