luau_bytecode/graph/

sccp.rs

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);
    }

    // Sccp.test.cpp: sccp_does_not_fold_boolean_ordering
    #[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());
    }

    // Sccp.test.cpp: sccp_folds_number_ordering
    #[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());
    }

    // Sccp.test.cpp: sccp_phi_filters_dead_predecessor
    #[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());
    }

    // Sccp.test.cpp: sccp_erases_trivial_phi
    #[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)
    }

    // Sccp.test.cpp: sccp_loadk_mul_to_mulk
    #[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);
    }

    // Sccp.test.cpp: sccp_loadk_div_to_divrk
    #[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));
    }

    // Sccp.test.cpp: sccp_loadk_mul_to_zero
    #[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));
    }
}