celox-sir-opt 0.4.3

use crate::HashMap;
use crate::HashSet;
use crate::ir::*;

fn next_register_id(
    register_map: &mut HashMap<RegisterId, RegisterType>,
    counter: &mut usize,
) -> RegisterId {
    *counter += 1;
    while register_map.contains_key(&RegisterId(*counter)) {
        *counter += 1;
    }
    RegisterId(*counter)
}

/// Append the canonical four-state-to-control conversion required by SIR
/// `Branch`, returning an unsigned one-bit condition register.
pub(in crate::optimizer) fn normalize_branch_condition<A>(
    register_map: &mut HashMap<RegisterId, RegisterType>,
    instructions: &mut Vec<SIRInstruction<A>>,
    cond: RegisterId,
    reg_counter: &mut usize,
) -> RegisterId {
    let cond_type = register_map[&cond].clone();
    if matches!(
        cond_type,
        RegisterType::Bit {
            width: 1,
            signed: false
        }
    ) {
        return cond;
    }

    let truth = if cond_type.width() == 1 {
        cond
    } else {
        let truth = next_register_id(register_map, reg_counter);
        register_map.insert(truth, RegisterType::Logic { width: 1 });
        instructions.push(SIRInstruction::Unary(truth, UnaryOp::Or, cond));
        truth
    };
    let normalized = next_register_id(register_map, reg_counter);
    register_map.insert(
        normalized,
        RegisterType::Bit {
            width: 1,
            signed: false,
        },
    );
    instructions.push(SIRInstruction::Unary(
        normalized,
        UnaryOp::ToTwoState,
        truth,
    ));
    normalized
}

/// Returns `Some(RegisterId)` for instructions that define a register.
pub(in crate::optimizer) fn def_reg<A>(inst: &SIRInstruction<A>) -> Option<RegisterId> {
    match inst {
        SIRInstruction::Imm(dst, _)
        | SIRInstruction::Binary(dst, _, _, _)
        | SIRInstruction::Unary(dst, _, _)
        | SIRInstruction::Load(dst, _, _, _)
        | SIRInstruction::Concat(dst, _)
        | SIRInstruction::Slice(dst, _, _, _)
        | SIRInstruction::Mux(dst, _, _, _) => Some(*dst),
        SIRInstruction::Store(_, _, _, _, _, _)
        | SIRInstruction::Commit(_, _, _, _, _)
        | SIRInstruction::RuntimeEvent { .. }
        | SIRInstruction::CombCaptureEvent { .. }
        | SIRInstruction::CombCaptureEnableIfChanged { .. } => None,
    }
}

/// Try to extract a u64 value from a SIRValue that represents a 2-state constant.
pub(in crate::optimizer) fn sir_value_to_u64(val: &SIRValue) -> Option<u64> {
    if !val.mask.to_u64_digits().is_empty() {
        return None; // 4-state value
    }
    let digits = val.payload.to_u64_digits();
    match digits.len() {
        0 => Some(0),
        1 => Some(digits[0]),
        _ => None,
    }
}

/// Resolve transitive aliases: if A→B and B→C, produce A→C.
pub(in crate::optimizer) fn resolve_transitive_aliases(
    aliases: &HashMap<RegisterId, RegisterId>,
) -> HashMap<RegisterId, RegisterId> {
    let mut resolved = HashMap::default();
    for (&from, &to) in aliases {
        let mut target = to;
        while let Some(&next) = aliases.get(&target) {
            if next == target {
                break;
            }
            target = next;
        }
        resolved.insert(from, target);
    }
    resolved
}

/// Collect all registers that are used (read) anywhere in an execution unit,
/// including instruction operands and terminators.
pub(in crate::optimizer) fn collect_all_used_registers(
    eu: &ExecutionUnit<RegionedAbsoluteAddr>,
) -> HashSet<RegisterId> {
    let mut used = HashSet::default();
    for block in eu.blocks.values() {
        for inst in &block.instructions {
            collect_used_regs_into(inst, &mut used);
        }
        collect_terminator_used_regs(&block.terminator, &mut used);
    }
    used
}

fn collect_used_regs_into(
    inst: &SIRInstruction<RegionedAbsoluteAddr>,
    out: &mut HashSet<RegisterId>,
) {
    match inst {
        SIRInstruction::Imm(_, _) => {}
        SIRInstruction::Binary(_, lhs, _, rhs) => {
            out.insert(*lhs);
            out.insert(*rhs);
        }
        SIRInstruction::Unary(_, _, src) => {
            out.insert(*src);
        }
        SIRInstruction::Load(_, _, offset, _) => {
            out.extend(offset.dynamic_registers().into_iter().flatten());
        }
        SIRInstruction::Store(_, offset, _, src, _, _) => {
            out.extend(offset.dynamic_registers().into_iter().flatten());
            out.insert(*src);
        }
        SIRInstruction::Commit(_, _, offset, _, _) => {
            out.extend(offset.dynamic_registers().into_iter().flatten());
        }
        SIRInstruction::Concat(_, args) => {
            out.extend(args.iter().copied());
        }
        SIRInstruction::Slice(_, src, _, _) => {
            out.insert(*src);
        }
        SIRInstruction::Mux(_, cond, then_val, else_val) => {
            out.insert(*cond);
            out.insert(*then_val);
            out.insert(*else_val);
        }
        SIRInstruction::RuntimeEvent { args, .. }
        | SIRInstruction::CombCaptureEvent { args, .. } => {
            out.extend(args.iter().copied());
        }
        SIRInstruction::CombCaptureEnableIfChanged { old, new, .. } => {
            out.insert(*old);
            out.insert(*new);
        }
    }
}

fn collect_terminator_used_regs(term: &SIRTerminator, out: &mut HashSet<RegisterId>) {
    match term {
        SIRTerminator::Jump(_, args) => {
            out.extend(args.iter().copied());
        }
        SIRTerminator::Branch {
            cond,
            true_block,
            false_block,
        } => {
            out.insert(*cond);
            out.extend(true_block.1.iter().copied());
            out.extend(false_block.1.iter().copied());
        }
        SIRTerminator::Switch { selector, .. } => {
            out.insert(*selector);
        }
        SIRTerminator::Return | SIRTerminator::Error(_) => {}
    }
}

pub(in crate::optimizer) fn replace_reg_in_terminator(
    term: &mut SIRTerminator,
    from: RegisterId,
    to: RegisterId,
) {
    match term {
        SIRTerminator::Jump(_, args) => {
            for arg in args {
                if *arg == from {
                    *arg = to;
                }
            }
        }
        SIRTerminator::Branch {
            cond,
            true_block,
            false_block,
        } => {
            if *cond == from {
                *cond = to;
            }
            for arg in &mut true_block.1 {
                if *arg == from {
                    *arg = to;
                }
            }
            for arg in &mut false_block.1 {
                if *arg == from {
                    *arg = to;
                }
            }
        }
        SIRTerminator::Switch { selector, .. } => {
            if *selector == from {
                *selector = to;
            }
        }
        SIRTerminator::Return | SIRTerminator::Error(_) => {}
    }
}

pub(in crate::optimizer) fn hoist_common_branch_loads(
    eu: &mut ExecutionUnit<RegionedAbsoluteAddr>,
) {
    let mut reg_counter: usize = eu.register_map.keys().map(|r| r.0).max().unwrap_or(0);

    #[derive(Clone, Copy)]
    struct Candidate {
        pred: BlockId,
        true_block: BlockId,
        false_block: BlockId,
        offset: usize,
        bits: usize,
        dst_true: RegisterId,
        dst_false: RegisterId,
        addr: RegionedAbsoluteAddr,
    }

    loop {
        let mut candidates = Vec::new();

        let block_ids: Vec<_> = eu.blocks.keys().copied().collect();
        for bid in block_ids {
            let Some(block) = eu.blocks.get(&bid) else {
                continue;
            };

            let SIRTerminator::Branch {
                true_block,
                false_block,
                ..
            } = &block.terminator
            else {
                continue;
            };

            let Some(t_block) = eu.blocks.get(&true_block.0) else {
                continue;
            };
            let Some(f_block) = eu.blocks.get(&false_block.0) else {
                continue;
            };

            let Some(SIRInstruction::Load(dst_t, addr_t, SIROffset::Static(off_t), bits_t)) =
                t_block.instructions.first()
            else {
                continue;
            };
            let Some(SIRInstruction::Load(dst_f, addr_f, SIROffset::Static(off_f), bits_f)) =
                f_block.instructions.first()
            else {
                continue;
            };

            if addr_t == addr_f && off_t == off_f && bits_t == bits_f {
                candidates.push(Candidate {
                    pred: bid,
                    true_block: true_block.0,
                    false_block: false_block.0,
                    offset: *off_t,
                    bits: *bits_t,
                    dst_true: *dst_t,
                    dst_false: *dst_f,
                    addr: *addr_t,
                });
            }
        }

        if candidates.is_empty() {
            break;
        }

        // Batch: collect all replacements, then apply once
        let mut replacement_map: HashMap<RegisterId, RegisterId> = HashMap::default();
        let mut changed = false;

        for c in candidates {
            let can_apply = if let (Some(t_block), Some(f_block)) =
                (eu.blocks.get(&c.true_block), eu.blocks.get(&c.false_block))
            {
                let t_ok = matches!(
                    t_block.instructions.first(),
                    Some(SIRInstruction::Load(dst, addr, SIROffset::Static(off), bits))
                        if *dst == c.dst_true && *addr == c.addr && *off == c.offset && *bits == c.bits
                );
                let f_ok = matches!(
                    f_block.instructions.first(),
                    Some(SIRInstruction::Load(dst, addr, SIROffset::Static(off), bits))
                        if *dst == c.dst_false && *addr == c.addr && *off == c.offset && *bits == c.bits
                );
                t_ok && f_ok
            } else {
                false
            };

            if !can_apply {
                continue;
            }

            let existing_reg = eu.blocks.get(&c.pred).and_then(|pred_block| {
                pred_block.instructions.iter().find_map(|inst| match inst {
                    SIRInstruction::Load(dst, addr, SIROffset::Static(off), bits)
                        if *addr == c.addr && *off == c.offset && *bits == c.bits =>
                    {
                        Some(*dst)
                    }
                    _ => None,
                })
            });

            let hoisted_reg = existing_reg.unwrap_or_else(|| {
                let new_reg = next_register_id(&mut eu.register_map, &mut reg_counter);
                eu.register_map
                    .insert(new_reg, RegisterType::Logic { width: c.bits });

                if let Some(pred_block) = eu.blocks.get_mut(&c.pred) {
                    pred_block.instructions.push(SIRInstruction::Load(
                        new_reg,
                        c.addr,
                        SIROffset::Static(c.offset),
                        c.bits,
                    ));
                }
                new_reg
            });

            if let Some(t_block) = eu.blocks.get_mut(&c.true_block) {
                t_block.instructions.remove(0);
            }
            if let Some(f_block) = eu.blocks.get_mut(&c.false_block) {
                f_block.instructions.remove(0);
            }

            replacement_map.insert(c.dst_true, hoisted_reg);
            replacement_map.insert(c.dst_false, hoisted_reg);
            changed = true;
        }

        if !changed {
            break;
        }

        // Resolve transitive replacements
        let mut final_map = HashMap::default();
        for &from in replacement_map.keys() {
            let mut to = replacement_map[&from];
            let mut depth = 0;
            while let Some(&next_to) = replacement_map.get(&to) {
                if next_to == to || depth > replacement_map.len() {
                    break;
                }
                to = next_to;
                depth += 1;
            }
            final_map.insert(from, to);
        }

        // Batch apply all replacements in a single pass over all blocks
        for block in eu.blocks.values_mut() {
            for p in &mut block.params {
                if let Some(&to) = final_map.get(p) {
                    *p = to;
                }
            }
            for inst in &mut block.instructions {
                batch_replace_in_inst(inst, &final_map);
            }
            batch_replace_in_terminator(&mut block.terminator, &final_map);
        }
    }
}

pub(in crate::optimizer) fn replace_offset_registers(
    offset: &mut SIROffset,
    map: &HashMap<RegisterId, RegisterId>,
) {
    match offset {
        SIROffset::Static(_) | SIROffset::PackedElements { .. } => {}
        SIROffset::Dynamic(register) => {
            if let Some(&replacement) = map.get(register) {
                *register = replacement;
            }
        }
        SIROffset::Element {
            index,
            dynamic_bit_offset,
            ..
        } => {
            if let Some(&replacement) = map.get(index) {
                *index = replacement;
            }
            if let Some(dynamic) = dynamic_bit_offset
                && let Some(&replacement) = map.get(dynamic)
            {
                *dynamic = replacement;
            }
        }
    }
}

pub(in crate::optimizer) fn batch_replace_in_inst(
    inst: &mut SIRInstruction<RegionedAbsoluteAddr>,
    map: &HashMap<RegisterId, RegisterId>,
) {
    match inst {
        SIRInstruction::Imm(_, _) => {}
        SIRInstruction::Binary(_, lhs, _, rhs) => {
            if let Some(&to) = map.get(lhs) {
                *lhs = to;
            }
            if let Some(&to) = map.get(rhs) {
                *rhs = to;
            }
        }
        SIRInstruction::Unary(_, _, src) => {
            if let Some(&to) = map.get(src) {
                *src = to;
            }
        }
        SIRInstruction::Load(_, _, offset, _) => replace_offset_registers(offset, map),
        SIRInstruction::Store(_, offset, _, src, _, _) => {
            replace_offset_registers(offset, map);
            if let Some(&to) = map.get(src) {
                *src = to;
            }
        }
        SIRInstruction::Commit(_, _, offset, _, _) => replace_offset_registers(offset, map),
        SIRInstruction::Concat(_, args) => {
            for arg in args {
                if let Some(&to) = map.get(arg) {
                    *arg = to;
                }
            }
        }
        SIRInstruction::Slice(_, src, _, _) => {
            if let Some(&to) = map.get(src) {
                *src = to;
            }
        }
        SIRInstruction::Mux(_, cond, then_val, else_val) => {
            if let Some(&to) = map.get(cond) {
                *cond = to;
            }
            if let Some(&to) = map.get(then_val) {
                *then_val = to;
            }
            if let Some(&to) = map.get(else_val) {
                *else_val = to;
            }
        }
        SIRInstruction::RuntimeEvent { args, .. }
        | SIRInstruction::CombCaptureEvent { args, .. } => {
            for arg in args {
                if let Some(&to) = map.get(arg) {
                    *arg = to;
                }
            }
        }
        SIRInstruction::CombCaptureEnableIfChanged { old, new, .. } => {
            if let Some(&to) = map.get(old) {
                *old = to;
            }
            if let Some(&to) = map.get(new) {
                *new = to;
            }
        }
    }
}

pub(in crate::optimizer) fn batch_replace_in_terminator(
    term: &mut SIRTerminator,
    map: &HashMap<RegisterId, RegisterId>,
) {
    match term {
        SIRTerminator::Jump(_, args) => {
            for arg in args {
                if let Some(&to) = map.get(arg) {
                    *arg = to;
                }
            }
        }
        SIRTerminator::Branch {
            cond,
            true_block,
            false_block,
        } => {
            if let Some(&to) = map.get(cond) {
                *cond = to;
            }
            for arg in &mut true_block.1 {
                if let Some(&to) = map.get(arg) {
                    *arg = to;
                }
            }
            for arg in &mut false_block.1 {
                if let Some(&to) = map.get(arg) {
                    *arg = to;
                }
            }
        }
        SIRTerminator::Switch { selector, .. } => {
            if let Some(&to) = map.get(selector) {
                *selector = to;
            }
        }
        SIRTerminator::Return | SIRTerminator::Error(_) => {}
    }
}