neo-decompiler 0.10.1

use std::collections::{BTreeMap, BTreeSet};

use crate::decompiler::cfg::BlockId;

use super::super::form::{SsaBlock, SsaStmt, UseSite};
use super::super::variable::SsaVariable;
use super::state::SlotState;
use super::uses::collect_expr_uses;
use super::{SsaBuildResult, SsaBuilder};

impl SsaBuilder<'_> {
    /// Run the fixpoint that produces per-block phis, exit stacks, and the
    /// assembled [`super::super::form::SsaForm`] pieces.
    pub(super) fn build_ssa_blocks(&self) -> SsaBuildResult {
        let block_ids: Vec<BlockId> = self.cfg.blocks().map(|b| b.id).collect();

        let mut entry_stacks: BTreeMap<BlockId, Vec<SsaVariable>> = BTreeMap::new();
        let mut exit_stacks: BTreeMap<BlockId, Vec<SsaVariable>> = BTreeMap::new();
        let mut entry_slots: BTreeMap<BlockId, SlotState> = BTreeMap::new();
        let mut exit_slots: BTreeMap<BlockId, SlotState> = BTreeMap::new();
        let mut block_uses: BTreeMap<BlockId, Vec<(SsaVariable, usize)>> = BTreeMap::new();
        let mut versions: BTreeMap<String, usize> = BTreeMap::new();

        let max_iterations = block_ids.len() + 4;
        let mut changed = true;
        let mut iterations = 0usize;
        while changed && iterations <= max_iterations {
            changed = false;
            iterations += 1;
            versions.clear();
            for &bid in &block_ids {
                let (new_entry, _new_phis) = self.compute_join_entry(bid, &exit_stacks);
                let (new_slot_entry, _new_slot_phis) =
                    self.compute_join_slots(bid, &exit_slots, &mut versions);
                let exec = self.execute_block(bid, &new_entry, &new_slot_entry, &mut versions);

                let exit_changed = exit_stacks.get(&bid) != Some(&exec.exit_stack);
                let entry_changed = entry_stacks.get(&bid) != Some(&new_entry);
                let slot_exit_changed = exit_slots.get(&bid) != Some(&exec.exit_slots);
                let slot_entry_changed = entry_slots.get(&bid) != Some(&new_slot_entry);
                if exit_changed || entry_changed || slot_exit_changed || slot_entry_changed {
                    changed = true;
                }
                entry_stacks.insert(bid, new_entry);
                exit_stacks.insert(bid, exec.exit_stack);
                entry_slots.insert(bid, new_slot_entry);
                exit_slots.insert(bid, exec.exit_slots);
                block_uses.insert(bid, exec.uses);
            }
        }

        self.assemble_ssa_blocks(
            &block_ids,
            &entry_stacks,
            &exit_stacks,
            &entry_slots,
            &exit_slots,
            &block_uses,
        )
    }

    fn assemble_ssa_blocks(
        &self,
        block_ids: &[BlockId],
        entry_stacks: &BTreeMap<BlockId, Vec<SsaVariable>>,
        exit_stacks: &BTreeMap<BlockId, Vec<SsaVariable>>,
        entry_slots: &BTreeMap<BlockId, SlotState>,
        exit_slots: &BTreeMap<BlockId, SlotState>,
        block_uses: &BTreeMap<BlockId, Vec<(SsaVariable, usize)>>,
    ) -> SsaBuildResult {
        let mut ssa_blocks = BTreeMap::new();
        let mut definitions = BTreeMap::new();
        let mut uses: BTreeMap<SsaVariable, BTreeSet<UseSite>> = BTreeMap::new();
        let mut versions: BTreeMap<String, usize> = BTreeMap::new();

        for &bid in block_ids {
            let entry = entry_stacks.get(&bid).cloned().unwrap_or_default();
            let slot_entry = entry_slots.get(&bid).cloned().unwrap_or_default();
            let (_, stack_phis) = self.compute_join_entry(bid, exit_stacks);
            let (_, slot_phis) = self.compute_join_slots(bid, exit_slots, &mut versions);
            let exec = self.execute_block(bid, &entry, &slot_entry, &mut versions);

            let mut sb = SsaBlock::new();
            for phi in stack_phis.iter().chain(slot_phis.iter()) {
                definitions.insert(phi.target.clone(), bid);
                for var in phi.operands.values() {
                    uses.entry(var.clone())
                        .or_default()
                        .insert(UseSite::new(bid, 0));
                }
                sb.add_phi(phi.clone());
            }
            for (i, stmt) in exec.stmts.iter().enumerate() {
                if let SsaStmt::Assign { target, value } = stmt {
                    definitions.insert(target.clone(), bid);
                    for used in collect_expr_uses(value) {
                        uses.entry(used).or_default().insert(UseSite::new(bid, i));
                    }
                }
                sb.add_stmt(stmt.clone());
            }
            for (var, idx) in block_uses.get(&bid).cloned().unwrap_or_default() {
                uses.entry(var).or_default().insert(UseSite::new(bid, idx));
            }
            ssa_blocks.insert(bid, sb);
        }

        (ssa_blocks, definitions, uses)
    }
}