llvm-native-core 0.1.4

// basic_block_ext.rs — World-Class BasicBlock Extension
//
// Clean-room forensic-parity expansion (self-contained string-based CFG analysis):
//   - BlockFrequencyInfo
//   - BranchProbabilityInfo with heuristics
//   - CFG traversal orderings (preorder, postorder, RPO)
//   - Critical edge splitting
//   - Block merging utilities
//   - Extended loop information
//   - Block placement heuristics

use std::collections::{HashMap, HashSet, VecDeque};
use std::fmt;

// ============================================================================
// Section 1: BlockFrequencyInfo
// ============================================================================

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct BlockFrequency(u64);

impl BlockFrequency {
    pub const ZERO: BlockFrequency = BlockFrequency(0);
    pub const MAX: BlockFrequency = BlockFrequency(u64::MAX);

    pub fn new(val: u64) -> Self {
        BlockFrequency(val)
    }
    pub fn get_frequency(&self) -> u64 {
        self.0
    }
    pub fn is_zero(&self) -> bool {
        self.0 == 0
    }

    pub fn scale(&self, num: u64, den: u64) -> BlockFrequency {
        if den == 0 {
            return BlockFrequency::ZERO;
        }
        BlockFrequency(((self.0 as u128 * num as u128) / den as u128) as u64)
    }

    pub fn add(&self, other: BlockFrequency) -> BlockFrequency {
        BlockFrequency(self.0.saturating_add(other.0))
    }
}

impl fmt::Display for BlockFrequency {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

pub struct BlockFrequencyInfo {
    pub frequencies: HashMap<String, BlockFrequency>,
    pub entry_frequency: BlockFrequency,
    pub computed: bool,
}

impl BlockFrequencyInfo {
    pub fn new() -> Self {
        BlockFrequencyInfo {
            frequencies: HashMap::new(),
            entry_frequency: BlockFrequency::MAX,
            computed: false,
        }
    }

    pub fn get_block_frequency(&self, block_name: &str) -> BlockFrequency {
        self.frequencies
            .get(block_name)
            .copied()
            .unwrap_or(BlockFrequency::ZERO)
    }

    pub fn set_block_frequency(&mut self, block_name: &str, freq: BlockFrequency) {
        self.frequencies.insert(block_name.to_string(), freq);
    }

    pub fn max_frequency(&self) -> BlockFrequency {
        self.frequencies
            .values()
            .max()
            .copied()
            .unwrap_or(BlockFrequency::ZERO)
    }
}

impl Default for BlockFrequencyInfo {
    fn default() -> Self {
        BlockFrequencyInfo::new()
    }
}

// ============================================================================
// Section 2: BranchProbabilityInfo
// ============================================================================

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BranchProbability {
    num: u32,
    den: u32,
}

impl BranchProbability {
    pub const UNKNOWN: BranchProbability = BranchProbability { num: 0, den: 0 };
    pub const VERY_UNLIKELY: BranchProbability = BranchProbability {
        num: 1,
        den: 100000,
    };
    pub const UNLIKELY: BranchProbability = BranchProbability { num: 1, den: 100 };
    pub const EVEN: BranchProbability = BranchProbability { num: 1, den: 2 };
    pub const LIKELY: BranchProbability = BranchProbability { num: 99, den: 100 };
    pub const VERY_LIKELY: BranchProbability = BranchProbability {
        num: 99999,
        den: 100000,
    };
    pub const ALWAYS: BranchProbability = BranchProbability { num: 1, den: 1 };
    pub const NEVER: BranchProbability = BranchProbability { num: 0, den: 1 };

    pub fn new(num: u32, den: u32) -> Self {
        if den == 0 {
            BranchProbability::UNKNOWN
        } else {
            BranchProbability { num, den }
        }
    }
    pub fn numerator(&self) -> u64 {
        self.num as u64
    }
    pub fn denominator(&self) -> u64 {
        self.den as u64
    }
    pub fn is_unknown(&self) -> bool {
        self.den == 0
    }
    pub fn as_f64(&self) -> f64 {
        if self.den == 0 {
            0.5
        } else {
            self.num as f64 / self.den as f64
        }
    }
    pub fn complement(&self) -> BranchProbability {
        if self.den == 0 {
            return BranchProbability::UNKNOWN;
        }
        BranchProbability {
            num: self.den - self.num,
            den: self.den,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BranchHeuristic {
    LoopBranch,
    LoopExit,
    Return,
    Unwind,
    ColdCall,
    NonNull,
    NonZero,
    Likely,
    Default,
}

pub struct BranchProbabilityInfo {
    pub edge_weights: HashMap<(String, String), BranchProbability>,
}

impl BranchProbabilityInfo {
    pub fn new() -> Self {
        BranchProbabilityInfo {
            edge_weights: HashMap::new(),
        }
    }

    pub fn get_edge_probability(&self, from: &str, to: &str) -> BranchProbability {
        self.edge_weights
            .get(&(from.to_string(), to.to_string()))
            .copied()
            .unwrap_or(BranchProbability::UNKNOWN)
    }

    pub fn set_edge_probability(&mut self, from: &str, to: &str, prob: BranchProbability) {
        self.edge_weights
            .insert((from.to_string(), to.to_string()), prob);
    }

    pub fn analyze(&mut self, successors: &HashMap<String, Vec<String>>) {
        for (block, succ_list) in successors {
            let num_succ = succ_list.len();
            if num_succ == 0 {
                continue;
            }
            if num_succ == 1 {
                self.set_edge_probability(block, &succ_list[0], BranchProbability::ALWAYS);
                continue;
            }
            for succ in succ_list {
                let prob = if succ.contains("unwind") || succ.contains("lpad") {
                    BranchProbability::VERY_UNLIKELY
                } else if succ.contains("return") {
                    BranchProbability::LIKELY
                } else {
                    BranchProbability::EVEN
                };
                self.set_edge_probability(block, succ, prob);
            }
        }
    }
}

impl Default for BranchProbabilityInfo {
    fn default() -> Self {
        BranchProbabilityInfo::new()
    }
}

// ============================================================================
// Section 3: CFG Traversal Orderings
// ============================================================================

pub struct CFGTraversal;

impl CFGTraversal {
    pub fn preorder(entry: &str, successors: &HashMap<String, Vec<String>>) -> Vec<String> {
        let mut result = Vec::new();
        let mut visited = HashSet::new();
        let mut stack = vec![entry.to_string()];
        while let Some(block) = stack.pop() {
            if !visited.insert(block.clone()) {
                continue;
            }
            result.push(block.clone());
            if let Some(succs) = successors.get(&block) {
                for succ in succs.iter().rev() {
                    if !visited.contains(succ) {
                        stack.push(succ.clone());
                    }
                }
            }
        }
        result
    }

    pub fn postorder(entry: &str, successors: &HashMap<String, Vec<String>>) -> Vec<String> {
        let mut result = Vec::new();
        let mut visited = HashSet::new();
        Self::postorder_dfs(entry, successors, &mut visited, &mut result);
        result
    }

    fn postorder_dfs(
        block: &str,
        succs: &HashMap<String, Vec<String>>,
        visited: &mut HashSet<String>,
        result: &mut Vec<String>,
    ) {
        if !visited.insert(block.to_string()) {
            return;
        }
        if let Some(succ_list) = succs.get(block) {
            for succ in succ_list {
                Self::postorder_dfs(succ, succs, visited, result);
            }
        }
        result.push(block.to_string());
    }

    pub fn reverse_postorder(
        entry: &str,
        successors: &HashMap<String, Vec<String>>,
    ) -> Vec<String> {
        let mut po = Self::postorder(entry, successors);
        po.reverse();
        po
    }
}

// ============================================================================
// Section 4: Critical Edge Splitting
// ============================================================================

#[derive(Debug, Clone)]
pub struct CriticalEdgeSplitter {
    pub successors: HashMap<String, Vec<String>>,
    pub predecessors: HashMap<String, Vec<String>>,
    pub new_blocks: Vec<String>,
}

impl CriticalEdgeSplitter {
    pub fn new(
        successors: HashMap<String, Vec<String>>,
        predecessors: HashMap<String, Vec<String>>,
    ) -> Self {
        CriticalEdgeSplitter {
            successors,
            predecessors,
            new_blocks: Vec::new(),
        }
    }

    pub fn is_critical_edge(&self, from: &str, to: &str) -> bool {
        let from_succs = self.successors.get(from).map(|v| v.len()).unwrap_or(0) > 1;
        let to_preds = self.predecessors.get(to).map(|v| v.len()).unwrap_or(0) > 1;
        from_succs && to_preds
    }

    pub fn find_critical_edges(&self) -> Vec<(String, String)> {
        let mut edges = Vec::new();
        for (from, succs) in &self.successors {
            for to in succs {
                if self.is_critical_edge(from, to) {
                    edges.push((from.clone(), to.clone()));
                }
            }
        }
        edges
    }

    pub fn split_edge(&mut self, from: &str, to: &str) -> String {
        let new_name = format!("{}.{}.critedge", from, to);
        self.new_blocks.push(new_name.clone());
        if let Some(succs) = self.successors.get_mut(from) {
            if let Some(pos) = succs.iter().position(|s| s == to) {
                succs[pos] = new_name.clone();
            }
        }
        self.successors
            .insert(new_name.clone(), vec![to.to_string()]);
        if let Some(preds) = self.predecessors.get_mut(to) {
            if let Some(pos) = preds.iter().position(|p| p == from) {
                preds[pos] = new_name.clone();
            }
        }
        self.predecessors
            .insert(new_name.clone(), vec![from.to_string()]);
        new_name
    }
}

// ============================================================================
// Section 5: Block Merging
// ============================================================================

pub struct BlockMerger;

impl BlockMerger {
    pub fn can_merge_blocks(
        from: &str,
        to: &str,
        successors: &HashMap<String, Vec<String>>,
        predecessors: &HashMap<String, Vec<String>>,
    ) -> bool {
        let succs = successors.get(from).map(|v| v.as_slice()).unwrap_or(&[]);
        if succs.len() != 1 || succs[0] != to {
            return false;
        }
        let preds = predecessors.get(to).map(|v| v.as_slice()).unwrap_or(&[]);
        if preds.len() != 1 || preds[0] != from {
            return false;
        }
        if from == "entry" || from == "0" {
            return false;
        }
        true
    }
}

// ============================================================================
// Section 6: Extended Loop Info
// ============================================================================

#[derive(Debug, Clone)]
pub struct ExtendedLoopInfo {
    pub header: String,
    pub blocks: Vec<String>,
    pub latches: HashSet<String>,
    pub exiting_blocks: HashSet<String>,
    pub exit_blocks: HashSet<String>,
    pub depth: u32,
    pub trip_count: Option<u64>,
}

impl ExtendedLoopInfo {
    pub fn new(header: String, blocks: Vec<String>, depth: u32) -> Self {
        ExtendedLoopInfo {
            header,
            blocks,
            latches: HashSet::new(),
            exiting_blocks: HashSet::new(),
            exit_blocks: HashSet::new(),
            depth,
            trip_count: None,
        }
    }

    pub fn analyze(
        &mut self,
        successors: &HashMap<String, Vec<String>>,
        predecessors: &HashMap<String, Vec<String>>,
    ) {
        for block in &self.blocks {
            if let Some(succs) = successors.get(block) {
                for succ in succs {
                    if !self.blocks.contains(succ) {
                        self.exiting_blocks.insert(block.clone());
                        self.exit_blocks.insert(succ.clone());
                    }
                }
            }
        }
        if let Some(preds) = predecessors.get(&self.header) {
            for pred in preds {
                if self.blocks.contains(pred) {
                    self.latches.insert(pred.clone());
                }
            }
        }
    }
}

// ============================================================================
// Section 7: Block Placement
// ============================================================================

pub struct BlockPlacement {
    pub chains: HashMap<String, String>,
    pub edge_weights: HashMap<(String, String), f64>,
}

impl BlockPlacement {
    pub fn new() -> Self {
        BlockPlacement {
            chains: HashMap::new(),
            edge_weights: HashMap::new(),
        }
    }

    pub fn compute_layout(
        &mut self,
        entry: &str,
        successors: &HashMap<String, Vec<String>>,
        bfi: &BlockFrequencyInfo,
        bpi: &BranchProbabilityInfo,
    ) -> Vec<String> {
        let all_blocks: Vec<String> = successors.keys().cloned().collect();
        for block in &all_blocks {
            self.chains.insert(block.clone(), block.clone());
        }

        for (from, succs) in successors {
            let from_freq = bfi.get_block_frequency(from).get_frequency() as f64;
            for to in succs {
                let prob = bpi.get_edge_probability(from, to).as_f64();
                self.edge_weights
                    .insert((from.clone(), to.clone()), from_freq * prob);
            }
        }

        let mut edges: Vec<((String, String), f64)> = self
            .edge_weights
            .iter()
            .map(|(k, v)| (k.clone(), *v))
            .collect();
        edges.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        for ((from, to), _) in &edges {
            let from_chain = self.get_chain_leader(from);
            let to_chain = self.get_chain_leader(to);
            if from_chain != to_chain && !self.would_create_cycle(from, to) {
                self.chains.insert(from_chain, to_chain);
            }
        }

        let mut visited = HashSet::new();
        let mut layout = Vec::new();
        let mut current = entry.to_string();
        loop {
            if visited.contains(&current) {
                break;
            }
            visited.insert(current.clone());
            layout.push(current.clone());
            let mut found = false;
            if let Some(succs) = successors.get(&current) {
                for succ in succs {
                    if !visited.contains(succ) {
                        current = succ.clone();
                        found = true;
                        break;
                    }
                }
            }
            if !found {
                break;
            }
        }
        layout
    }

    fn get_chain_leader(&self, block: &str) -> String {
        let mut current = block.to_string();
        while let Some(leader) = self.chains.get(&current) {
            if leader == &current {
                break;
            }
            current = leader.clone();
        }
        current
    }

    fn would_create_cycle(&self, from: &str, to: &str) -> bool {
        let mut current = to.to_string();
        let mut visited = HashSet::new();
        while let Some(leader) = self.chains.get(&current) {
            if leader == &current {
                break;
            }
            if leader == from {
                return true;
            }
            if !visited.insert(leader.clone()) {
                break;
            }
            current = leader.clone();
        }
        false
    }
}

impl Default for BlockPlacement {
    fn default() -> Self {
        BlockPlacement::new()
    }
}

// ============================================================================
// Section 8: Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_block_frequency() {
        let f = BlockFrequency::new(1000);
        assert!(!f.is_zero());
        assert_eq!(f.get_frequency(), 1000);
        let s = f.scale(1, 2);
        assert_eq!(s.get_frequency(), 500);
    }

    #[test]
    fn test_branch_probability() {
        let p = BranchProbability::new(3, 4);
        assert_eq!(p.numerator(), 3);
        assert_eq!(p.denominator(), 4);
        assert!((p.as_f64() - 0.75).abs() < 0.001);
    }

    #[test]
    fn test_cfg_postorder() {
        let mut succs: HashMap<String, Vec<String>> = HashMap::new();
        succs.insert("entry".to_string(), vec!["body".to_string()]);
        succs.insert("body".to_string(), vec!["exit".to_string()]);
        succs.insert("exit".to_string(), vec![]);
        let po = CFGTraversal::postorder("entry", &succs);
        assert_eq!(po[0], "exit");
        assert_eq!(po[2], "entry");
    }

    #[test]
    fn test_reverse_postorder() {
        let mut succs: HashMap<String, Vec<String>> = HashMap::new();
        succs.insert("entry".to_string(), vec!["body".to_string()]);
        succs.insert("body".to_string(), vec!["exit".to_string()]);
        succs.insert("exit".to_string(), vec![]);
        let rpo = CFGTraversal::reverse_postorder("entry", &succs);
        assert_eq!(rpo[0], "entry");
        assert_eq!(rpo[2], "exit");
    }

    #[test]
    fn test_critical_edge() {
        let mut succs: HashMap<String, Vec<String>> = HashMap::new();
        let mut preds: HashMap<String, Vec<String>> = HashMap::new();
        succs.insert("A".to_string(), vec!["B".to_string(), "C".to_string()]);
        succs.insert("B".to_string(), vec!["D".to_string()]);
        succs.insert("C".to_string(), vec!["D".to_string()]);
        preds.insert("D".to_string(), vec!["B".to_string(), "C".to_string()]);
        let splitter = CriticalEdgeSplitter::new(succs, preds);
        let critical = splitter.find_critical_edges();
        assert!(!critical.is_empty());
    }

    #[test]
    fn test_block_placement() {
        let mut succs: HashMap<String, Vec<String>> = HashMap::new();
        succs.insert("entry".to_string(), vec!["A".to_string()]);
        succs.insert("A".to_string(), vec!["B".to_string()]);
        succs.insert("B".to_string(), vec![]);
        let mut bfi = BlockFrequencyInfo::new();
        bfi.set_block_frequency("entry", BlockFrequency::new(1000));
        bfi.set_block_frequency("A", BlockFrequency::new(1000));
        bfi.set_block_frequency("B", BlockFrequency::new(1000));
        let bpi = BranchProbabilityInfo::new();
        let mut placement = BlockPlacement::new();
        let layout = placement.compute_layout("entry", &succs, &bfi, &bpi);
        assert_eq!(layout.len(), 3);
    }

    #[test]
    fn test_extended_loop_info() {
        let mut eli = ExtendedLoopInfo::new(
            "header".to_string(),
            vec!["header".to_string(), "body".to_string()],
            1,
        );
        let mut succs = HashMap::new();
        succs.insert("header".to_string(), vec!["body".to_string()]);
        succs.insert("body".to_string(), vec!["exit".to_string()]);
        let mut preds = HashMap::new();
        preds.insert("body".to_string(), vec!["header".to_string()]);
        eli.analyze(&succs, &preds);
        assert_eq!(eli.blocks.len(), 2);
    }

    #[test]
    fn test_block_merger() {
        let mut succs: HashMap<String, Vec<String>> = HashMap::new();
        let mut preds: HashMap<String, Vec<String>> = HashMap::new();
        succs.insert("A".to_string(), vec!["B".to_string()]);
        preds.insert("B".to_string(), vec!["A".to_string()]);
        assert!(BlockMerger::can_merge_blocks("A", "B", &succs, &preds));
    }
}