neo-decompiler 0.10.1

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

use super::{SsaBlock, UseSite};
use crate::decompiler::cfg::ssa::dominance::DominanceInfo;
use crate::decompiler::cfg::ssa::variable::SsaVariable;
use crate::decompiler::cfg::{BlockId, Cfg};

/// A control flow graph in Static Single Assignment form.
///
/// SSA form guarantees that each variable is assigned exactly once, making
/// data flow analysis and optimizations significantly simpler.
#[derive(Debug, Clone)]
pub struct SsaForm {
    /// The original control flow graph.
    pub cfg: Cfg,

    /// Dominance information (immediate dominators, dominator tree, dominance frontiers).
    pub dominance: DominanceInfo,

    /// SSA blocks indexed by block ID.
    pub blocks: BTreeMap<BlockId, SsaBlock>,

    /// Mapping from SSA variables to the block where they are defined.
    pub definitions: BTreeMap<SsaVariable, BlockId>,

    /// Mapping from SSA variables to all their use sites.
    pub uses: BTreeMap<SsaVariable, BTreeSet<UseSite>>,
}

impl SsaForm {
    /// Create a new empty SSA form.
    #[must_use]
    pub fn new(cfg: Cfg, dominance: DominanceInfo) -> Self {
        Self {
            cfg,
            dominance,
            blocks: BTreeMap::new(),
            definitions: BTreeMap::new(),
            uses: BTreeMap::new(),
        }
    }

    /// Add a block to the SSA form.
    pub fn add_block(&mut self, id: BlockId, block: SsaBlock) {
        self.blocks.insert(id, block);
    }

    /// Get a block by ID.
    #[must_use]
    pub fn block(&self, id: BlockId) -> Option<&SsaBlock> {
        self.blocks.get(&id)
    }

    /// Get all blocks in SSA form.
    pub fn blocks_iter(&self) -> impl Iterator<Item = (&BlockId, &SsaBlock)> {
        self.blocks.iter()
    }

    /// Get the number of blocks.
    #[must_use]
    pub fn block_count(&self) -> usize {
        self.blocks.len()
    }

    /// Record a variable definition.
    pub fn add_definition(&mut self, var: SsaVariable, block: BlockId) {
        self.definitions.insert(var, block);
    }

    /// Record a variable use.
    pub fn add_use(&mut self, var: SsaVariable, site: UseSite) {
        self.uses.entry(var).or_default().insert(site);
    }

    /// Get all use sites for a variable.
    #[must_use]
    pub fn uses_of(&self, var: &SsaVariable) -> Option<&BTreeSet<UseSite>> {
        self.uses.get(var)
    }

    /// Render the SSA form as a string for debugging/display.
    #[must_use]
    pub fn render(&self) -> String {
        let mut output = String::new();
        use std::fmt::Write;

        writeln!(output, "// SSA Form - {} blocks", self.block_count()).unwrap();
        writeln!(output).unwrap();

        for (block_id, block) in self.blocks_iter() {
            writeln!(output, "block {:?}:", block_id).unwrap();
            write!(output, "{}", block).unwrap();
        }

        output
    }

    /// Get statistics about the SSA form.
    #[must_use]
    pub fn stats(&self) -> SsaStats {
        let total_phi_nodes: usize = self.blocks.values().map(SsaBlock::phi_count).sum();
        let total_statements: usize = self.blocks.values().map(SsaBlock::stmt_count).sum();
        let total_variables = self.definitions.len();

        SsaStats {
            block_count: self.block_count(),
            total_phi_nodes,
            total_statements,
            total_variables,
        }
    }
}

/// Statistics about an SSA form.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SsaStats {
    /// Number of blocks in the SSA form.
    pub block_count: usize,
    /// Total number of φ nodes across all blocks.
    pub total_phi_nodes: usize,
    /// Total number of SSA statements (excluding φ nodes).
    pub total_statements: usize,
    /// Total number of unique SSA variables.
    pub total_variables: usize,
}

impl fmt::Display for SsaStats {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "SSA Stats: {} blocks, {} φ nodes, {} statements, {} variables",
            self.block_count, self.total_phi_nodes, self.total_statements, self.total_variables
        )
    }
}