llvm-native-core-ext 0.1.0

//! Whole-Program Support — LTO, Devirtualization, CFI, GlobalDCE
//!
//! This module provides whole-program optimization infrastructure for
//! link-time optimization (LTO) and thin LTO:
//!
//! - **WholeProgramDevirt**: devirtualizes virtual calls with whole-program
//!   visibility by analyzing the complete class hierarchy.
//! - **VirtualCallAnalysis**: analyze virtual call targets across translation
//!   units using type metadata.
//! - **TypeMetadata**: emit type metadata for virtual function tables to
//!   enable devirtualization and CFI.
//! - **VCallVisibility**: check vcall_visibility metadata for devirtualization
//!   decisions.
//! - **CFI-WholeProgram**: cross-DSO control flow integrity to prevent
//!   virtual call hijacking.
//! - **LTOVisibility**: promote hidden visibility for LTO optimization,
//!   enabling interprocedural optimizations.
//! - **GlobalDCE**: whole-program dead global elimination to remove
//!   unreferenced functions and global variables.
//! - **ThinLTO whole-program**: summary-based devirtualization, cross-module
//!   optimization hints for distributed build systems.
//!
//! Clean-room behavioural reconstruction from the LLVM LTO documentation,
//! the ThinLTO design paper, and the CFI specification. No LLVM C++
//! source code is consulted.

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

// ============================================================================
// Type Metadata
// ============================================================================

/// Type metadata entry: describes a type for devirtualization and CFI.
/// Each virtual function table has associated type metadata.
#[derive(Debug, Clone)]
pub struct TypeMetadata {
    /// Unique identifier for this type (e.g., mangled class name).
    pub type_id: String,
    /// Address points in the vtable (offsets from vtable base).
    pub address_points: Vec<u64>,
    /// Number of entries in the vtable.
    pub num_entries: u32,
    /// Whether this type is visible outside the DSO.
    pub is_visible_externally: bool,
    /// Derived class type IDs (for full class hierarchy).
    pub derived_types: HashSet<String>,
    /// Base class type IDs.
    pub base_types: HashSet<String>,
}

/// Collection of type metadata for a translation unit.
#[derive(Debug, Default)]
pub struct TypeMetadataTable {
    pub entries: HashMap<String, TypeMetadata>,
}

impl TypeMetadataTable {
    pub fn new() -> Self {
        TypeMetadataTable {
            entries: HashMap::new(),
        }
    }

    /// Register a vtable type with its metadata.
    pub fn register_type(&mut self, type_id: &str, num_entries: u32, address_points: Vec<u64>) {
        self.entries
            .entry(type_id.to_string())
            .and_modify(|e| {
                e.address_points.extend(address_points.clone());
                e.num_entries = e.num_entries.max(num_entries);
            })
            .or_insert(TypeMetadata {
                type_id: type_id.to_string(),
                address_points,
                num_entries,
                is_visible_externally: false,
                derived_types: HashSet::new(),
                base_types: HashSet::new(),
            });
    }

    /// Record an inheritance relationship.
    pub fn add_derived_type(&mut self, base: &str, derived: &str) {
        if let Some(base_entry) = self.entries.get_mut(base) {
            base_entry.derived_types.insert(derived.to_string());
        }
        if let Some(derived_entry) = self.entries.get_mut(derived) {
            derived_entry.base_types.insert(base.to_string());
        }
    }

    /// Emit LLVM IR type metadata for the given type.
    pub fn emit_type_metadata_ir(&self, type_id: &str) -> String {
        if let Some(entry) = self.entries.get(type_id) {
            let mut ir = String::new();
            ir.push_str(&format!("!{}.type = !{{", type_id));
            ir.push_str(&format!("i32 16, !\"typeid\", !\"{}\"", type_id));
            ir.push_str("}\n");
            ir.push_str(&format!(
                "!{}.vtable = !{{i32 {}, {} address points}}\n",
                type_id,
                entry.num_entries,
                entry.address_points.len()
            ));
            ir
        } else {
            String::new()
        }
    }
}

// ============================================================================
// Virtual Call Analysis
// ============================================================================

/// Represents a virtual call site being analyzed.
#[derive(Debug, Clone)]
pub struct VirtualCallSite {
    /// The function containing this call.
    pub caller: String,
    /// The vtable from which the call is made.
    pub vtable_type_id: String,
    /// The vtable index (which virtual function is called).
    pub vtable_index: u32,
    /// Known possible callees (from whole-program analysis).
    pub possible_callees: Vec<String>,
    /// Whether this call is known to have exactly one target.
    pub is_monomorphic: bool,
}

/// Virtual call analysis across the whole program.
#[derive(Debug)]
pub struct VirtualCallAnalysis {
    /// All virtual call sites found in the program.
    pub call_sites: Vec<VirtualCallSite>,
    /// Mapping from vtable type ID + index to possible callees.
    pub target_map: HashMap<(String, u32), Vec<String>>,
    /// Type metadata for resolving targets.
    pub type_metadata: TypeMetadataTable,
    /// Statistics.
    pub total_calls: u64,
    pub monomorphic_calls: u64,
    pub polymorphic_calls: u64,
}

impl VirtualCallAnalysis {
    pub fn new() -> Self {
        VirtualCallAnalysis {
            call_sites: Vec::new(),
            target_map: HashMap::new(),
            type_metadata: TypeMetadataTable::new(),
            total_calls: 0,
            monomorphic_calls: 0,
            polymorphic_calls: 0,
        }
    }

    /// Register a virtual call site for analysis.
    pub fn register_call(&mut self, caller: &str, vtable_type: &str, index: u32) {
        self.total_calls += 1;
        self.call_sites.push(VirtualCallSite {
            caller: caller.to_string(),
            vtable_type_id: vtable_type.to_string(),
            vtable_index: index,
            possible_callees: Vec::new(),
            is_monomorphic: false,
        });
    }

    /// Resolve all virtual call targets using whole-program information.
    pub fn resolve_targets(&mut self) {
        for site in &mut self.call_sites {
            let key = (site.vtable_type_id.clone(), site.vtable_index);
            if let Some(targets) = self.target_map.get(&key) {
                site.possible_callees = targets.clone();
                site.is_monomorphic = targets.len() == 1;
                if site.is_monomorphic {
                    self.monomorphic_calls += 1;
                } else {
                    self.polymorphic_calls += 1;
                }
            }
        }
    }

    /// Get the single target function for a monomorphic call site.
    pub fn get_single_target(&self, caller: &str, vtable_type: &str, index: u32) -> Option<String> {
        let key = (vtable_type.to_string(), index);
        if let Some(targets) = self.target_map.get(&key) {
            if targets.len() == 1 {
                return Some(targets[0].clone());
            }
        }
        None
    }
}

// ============================================================================
// Whole-Program Devirtualization
// ============================================================================

/// Whole-program devirtualization replaces indirect virtual calls with
/// direct calls when the target can be determined at LTO time.
#[derive(Debug)]
pub struct WholeProgramDevirt {
    /// Analysis results used for devirtualization.
    pub analysis: VirtualCallAnalysis,
    /// Whether to emit type checks for speculative devirtualization.
    pub enable_speculative_devirt: bool,
    /// Statistics.
    pub devirtualized_calls: u64,
    pub speculative_devirts: u64,
}

impl WholeProgramDevirt {
    pub fn new(analysis: VirtualCallAnalysis) -> Self {
        WholeProgramDevirt {
            analysis,
            enable_speculative_devirt: true,
            devirtualized_calls: 0,
            speculative_devirts: 0,
        }
    }

    /// Devirtualize a call site. Returns Some(direct_target) if the
    /// call can be fully devirtualized, or a guarded sequence for
    /// speculative devirtualization.
    pub fn devirtualize(&mut self, site: &VirtualCallSite) -> DevirtResult {
        if site.is_monomorphic && !site.possible_callees.is_empty() {
            self.devirtualized_calls += 1;
            DevirtResult::DirectCall(site.possible_callees[0].clone())
        } else if self.enable_speculative_devirt && site.possible_callees.len() <= 3 {
            self.speculative_devirts += 1;
            DevirtResult::Speculative {
                primary_target: site.possible_callees.first().cloned().unwrap_or_default(),
                fallback_targets: site.possible_callees[1..].to_vec(),
            }
        } else {
            DevirtResult::Indirect
        }
    }

    /// Emit the LLVM IR for a devirtualized call site.
    pub fn emit_devirtualized_call(&self, direct_target: &str, args: &str) -> String {
        format!("  call void @{}({})  ; devirtualized", direct_target, args)
    }

    /// Emit speculative devirtualization with a guard.
    pub fn emit_speculative_call(
        &self,
        primary: &str,
        fallbacks: &[String],
        vtable_ptr: &str,
        args: &str,
    ) -> String {
        let mut ir = String::new();
        ir.push_str(&format!("  ; speculative devirt: primary={}\n", primary));
        ir.push_str(&format!(
            "  %vt_match = icmp eq ptr {}, @{}.vtable\n",
            vtable_ptr, primary
        ));
        ir.push_str("  br i1 %vt_match, label %spec_call, label %fallback\n");
        ir.push_str("spec_call:\n");
        ir.push_str(&format!("  call void @{}({})\n", primary, args));
        ir.push_str("  br label %merge\n");
        ir.push_str("fallback:\n");
        for fb in fallbacks {
            ir.push_str(&format!("  ; fallback: {}\n", fb));
        }
        ir.push_str("  call void @vcall_indirect(ptr %vt_ptr, {} args)\n");
        ir.push_str("  br label %merge\n");
        ir.push_str("merge:\n");
        ir
    }
}

/// Result of devirtualization.
#[derive(Debug, Clone)]
pub enum DevirtResult {
    /// Call has exactly one known target — direct call.
    DirectCall(String),
    /// Speculative: likely target with guards, fallback to indirect.
    Speculative {
        primary_target: String,
        fallback_targets: Vec<String>,
    },
    /// Multiple possible targets — remain indirect.
    Indirect,
}

// ============================================================================
// VCall Visibility Metadata
// ============================================================================

/// Visibility of virtual call targets. LLVM uses !vcall_visibility
/// metadata to encode whether a vtable slot has a uniform value across
/// the whole program (enables devirtualization) or varies per DSO.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VCallVisibility {
    /// The vtable slot is constant across the whole program.
    /// Devirtualization is safe.
    VCallVisibilityPublic,
    /// The vtable slot may have different values in different DSOs.
    /// Devirtualization requires LTO.
    VCallVisibilityLinkageUnit,
    /// The vtable slot may be overridden in derived classes within
    /// the same DSO. Devirtualization requires analysis of class hierarchy.
    VCallVisibilityTranslationUnit,
}

/// Check the vcall_visibility of a vtable slot.
#[derive(Debug)]
pub struct VCallVisibilityChecker {
    /// Per-type visibility overrides.
    pub overrides: HashMap<(String, u32), VCallVisibility>,
}

impl VCallVisibilityChecker {
    pub fn new() -> Self {
        VCallVisibilityChecker {
            overrides: HashMap::new(),
        }
    }

    /// Set the visibility for a specific vtable slot.
    pub fn set_visibility(&mut self, type_id: &str, index: u32, vis: VCallVisibility) {
        self.overrides.insert((type_id.to_string(), index), vis);
    }

    /// Check whether devirtualization is safe for a given vtable slot.
    pub fn can_devirtualize(&self, type_id: &str, index: u32) -> bool {
        match self.overrides.get(&(type_id.to_string(), index)) {
            Some(VCallVisibility::VCallVisibilityPublic) => true,
            Some(VCallVisibility::VCallVisibilityLinkageUnit) => true,
            Some(VCallVisibility::VCallVisibilityTranslationUnit) => false,
            None => false, // conservative: assume not devirtualizable
        }
    }

    /// Generate vcall_visibility metadata for LLVM IR.
    pub fn emit_metadata(&self, type_id: &str, index: u32) -> String {
        let vis = self.overrides.get(&(type_id.to_string(), index));
        let visibility_val = match vis {
            Some(VCallVisibility::VCallVisibilityPublic) => 0,
            Some(VCallVisibility::VCallVisibilityLinkageUnit) => 1,
            Some(VCallVisibility::VCallVisibilityTranslationUnit) => 2,
            None => 2,
        };
        format!("!vcall_visibility = !{{{}}}", visibility_val)
    }
}

// ============================================================================
// CFI Whole-Program (Cross-DSO Control Flow Integrity)
// ============================================================================

/// Cross-DSO CFI prevents virtual call hijacking by verifying that
/// the target of an indirect call matches the expected type. Uses
/// __cfi_check_fail for diagnostics.

#[derive(Debug, Clone)]
pub struct CFITypeEntry {
    pub type_id: String,
    pub address_point_offset: u64,
    pub is_exported: bool,
}

/// CFI metadata for a module.
#[derive(Debug)]
pub struct CFIMetadata {
    pub types: Vec<CFITypeEntry>,
    pub blacklist: HashSet<String>,
    pub stats_checks_inserted: u64,
}

impl CFIMetadata {
    pub fn new() -> Self {
        CFIMetadata {
            types: Vec::new(),
            blacklist: HashSet::new(),
            stats_checks_inserted: 0,
        }
    }

    /// Register a type for CFI protection.
    pub fn register_type(&mut self, type_id: &str, address_point: u64, exported: bool) {
        self.types.push(CFITypeEntry {
            type_id: type_id.to_string(),
            address_point_offset: address_point,
            is_exported: exported,
        });
    }

    /// Blacklist a function from CFI checks.
    pub fn blacklist(&mut self, func_name: &str) {
        self.blacklist.insert(func_name.to_string());
    }

    /// Emit the CFI check sequence for a virtual call.
    pub fn emit_cfi_check(
        &mut self,
        vtable_ptr: &str,
        type_id: &str,
        address_point: u64,
    ) -> String {
        self.stats_checks_inserted += 1;
        let mut ir = String::new();
        ir.push_str(&format!("  ; CFI check for {}\n", type_id));
        ir.push_str(&format!(
            "  %cfi_addr = getelementptr i8, ptr {}, i64 -{}\n",
            vtable_ptr, address_point
        ));
        ir.push_str("  %cfi_load = load i64, ptr %cfi_addr\n");
        ir.push_str(&format!(
            "  %cfi_expected = load i64, ptr @__cfi_typeid_{}\n",
            type_id
        ));
        ir.push_str("  %cfi_ok = icmp eq i64 %cfi_load, %cfi_expected\n");
        ir.push_str("  br i1 %cfi_ok, label %cfi_cont, label %cfi_fail\n");
        ir.push_str("cfi_fail:\n");
        ir.push_str(&format!(
            "  call void @__cfi_check_fail(ptr {}, ptr @__cfi_type_info_{})\n",
            vtable_ptr, type_id
        ));
        ir.push_str("  unreachable\n");
        ir.push_str("cfi_cont:\n");
        ir
    }
}

// ============================================================================
// LTO Visibility Promotion
// ============================================================================

/// During LTO, functions and globals with hidden visibility can be
/// promoted to enable more aggressive interprocedural optimization.
/// Hidden symbols become internal, and internal symbols can be fully
/// inlined and optimized away.

#[derive(Debug, Clone)]
pub struct LTOSymbolInfo {
    pub name: String,
    pub original_visibility: SymbolVisibility,
    pub promoted_visibility: SymbolVisibility,
    pub is_address_taken: bool,
    pub can_be_internalized: bool,
}

/// Symbol visibility levels.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SymbolVisibility {
    Default,
    Hidden,
    Protected,
    Internal,
}

#[derive(Debug)]
pub struct LTOVisibilityPromotion {
    /// All symbols considered for promotion.
    pub symbols: Vec<LTOSymbolInfo>,
    /// Symbols that cannot be promoted (address taken, exported).
    pub unpromotable: HashSet<String>,
    /// Statistics.
    pub promoted_count: u64,
    pub internalized_count: u64,
}

impl LTOVisibilityPromotion {
    pub fn new() -> Self {
        LTOVisibilityPromotion {
            symbols: Vec::new(),
            unpromotable: HashSet::new(),
            promoted_count: 0,
            internalized_count: 0,
        }
    }

    /// Register a symbol for potential visibility promotion.
    pub fn register_symbol(&mut self, name: &str, vis: SymbolVisibility, is_address_taken: bool) {
        let can_int = !is_address_taken && matches!(vis, SymbolVisibility::Hidden);
        self.symbols.push(LTOSymbolInfo {
            name: name.to_string(),
            original_visibility: vis,
            promoted_visibility: vis,
            is_address_taken,
            can_be_internalized: can_int,
        });
    }

    /// Mark a symbol as unpromotable.
    pub fn mark_unpromotable(&mut self, name: &str) {
        self.unpromotable.insert(name.to_string());
    }

    /// Promote visibility: hidden→internal for symbols not address-taken.
    pub fn promote(&mut self) {
        for sym in &mut self.symbols {
            if sym.can_be_internalized && !self.unpromotable.contains(&sym.name) {
                sym.promoted_visibility = SymbolVisibility::Internal;
                self.internalized_count += 1;
            } else if sym.original_visibility == SymbolVisibility::Hidden {
                sym.promoted_visibility = SymbolVisibility::Internal;
                self.promoted_count += 1;
            }
        }
    }

    /// Emit LLVM IR linkage type for a promoted symbol.
    pub fn emit_linkage(&self, sym: &LTOSymbolInfo) -> &'static str {
        match sym.promoted_visibility {
            SymbolVisibility::Internal => "internal",
            SymbolVisibility::Hidden => "hidden",
            SymbolVisibility::Protected => "protected",
            SymbolVisibility::Default => "",
        }
    }
}

// ============================================================================
// Global Dead Code Elimination (GlobalDCE)
// ============================================================================

/// Whole-program dead global elimination: removes unreferenced
/// functions and global variables. Uses a worklist algorithm:
/// start from roots (main, exported symbols), transitively mark
/// all reachable globals, then delete unmarked ones.

#[derive(Debug)]
pub struct GlobalDCE {
    /// Live functions (root set).
    pub live_functions: HashSet<String>,
    /// Live global variables.
    pub live_globals: HashSet<String>,
    /// All functions seen.
    pub all_functions: HashSet<String>,
    /// All globals seen.
    pub all_globals: HashSet<String>,
    /// Call graph: caller → callees.
    pub call_graph: HashMap<String, HashSet<String>>,
    /// Global reference graph: function → global vars used.
    pub global_refs: HashMap<String, HashSet<String>>,
    /// Statistics.
    pub dead_functions_removed: u64,
    pub dead_globals_removed: u64,
}

impl GlobalDCE {
    pub fn new() -> Self {
        GlobalDCE {
            live_functions: HashSet::new(),
            live_globals: HashSet::new(),
            all_functions: HashSet::new(),
            all_globals: HashSet::new(),
            call_graph: HashMap::new(),
            global_refs: HashMap::new(),
            dead_functions_removed: 0,
            dead_globals_removed: 0,
        }
    }

    /// Register a function as a root (e.g., main, exported).
    pub fn add_root(&mut self, func: &str) {
        self.live_functions.insert(func.to_string());
    }

    /// Register a function.
    pub fn add_function(&mut self, name: &str, callees: Vec<&str>, globals_used: Vec<&str>) {
        self.all_functions.insert(name.to_string());
        self.call_graph.insert(
            name.to_string(),
            callees.iter().map(|s| s.to_string()).collect(),
        );
        self.global_refs.insert(
            name.to_string(),
            globals_used.iter().map(|s| s.to_string()).collect(),
        );
    }

    /// Register a global variable.
    pub fn add_global(&mut self, name: &str) {
        self.all_globals.insert(name.to_string());
    }

    /// Run the GlobalDCE algorithm.
    pub fn eliminate(&mut self) {
        // Mark phase: worklist from roots
        let mut worklist: Vec<String> = self.live_functions.iter().cloned().collect();
        let mut visited = HashSet::new();

        while let Some(func) = worklist.pop() {
            if !visited.insert(func.clone()) {
                continue;
            }
            // Mark all callees as live
            if let Some(callees) = self.call_graph.get(&func) {
                for callee in callees {
                    self.live_functions.insert(callee.clone());
                    if !visited.contains(callee) {
                        worklist.push(callee.clone());
                    }
                }
            }
            // Mark all referenced globals as live
            if let Some(globals) = self.global_refs.get(&func) {
                for g in globals {
                    self.live_globals.insert(g.clone());
                }
            }
        }

        // Sweep phase: count dead items
        for f in &self.all_functions {
            if !self.live_functions.contains(f) {
                self.dead_functions_removed += 1;
            }
        }
        for g in &self.all_globals {
            if !self.live_globals.contains(g) && !self.live_functions.contains(g) {
                self.dead_globals_removed += 1;
            }
        }
    }

    /// Check if a function is live.
    pub fn is_live(&self, func: &str) -> bool {
        self.live_functions.contains(func)
    }

    /// Check if a global is live.
    pub fn is_global_live(&self, global: &str) -> bool {
        self.live_globals.contains(global)
    }
}

// ============================================================================
// ThinLTO Whole-Program Support
// ============================================================================

/// ThinLTO uses module summaries to perform cross-module optimization
/// without requiring full IR merging. Each module emits a summary
/// containing function info (inst count, calls, refs) and type info.

/// Module summary for ThinLTO.
#[derive(Debug, Clone)]
pub struct ThinLTOModuleSummary {
    pub module_name: String,
    pub module_hash: String,
    pub functions: Vec<ThinLTOFunctionSummary>,
    pub globals: Vec<ThinLTOGlobalSummary>,
    pub type_ids: Vec<String>,
    pub vcall_visibility: HashMap<String, Vec<(u32, VCallVisibility)>>,
}

/// Function-level summary for ThinLTO.
#[derive(Debug, Clone)]
pub struct ThinLTOFunctionSummary {
    pub name: String,
    pub instruction_count: u64,
    pub callees: Vec<String>,
    pub referenced_globals: Vec<String>,
    pub is_exported: bool,
    pub has_inline_assembly: bool,
    pub is_available_externally: bool,
    pub vtable_slots: Vec<(String, u32)>, // (type_id, vtable_index)
}

/// Global variable summary.
#[derive(Debug, Clone)]
pub struct ThinLTOGlobalSummary {
    pub name: String,
    pub is_constant: bool,
    pub size: u64,
}

/// ThinLTO index: aggregates summaries from all modules.
#[derive(Debug)]
pub struct ThinLTOIndex {
    pub summaries: HashMap<String, ThinLTOModuleSummary>,
    /// Global function map: name → (module, summary).
    pub global_func_map: HashMap<String, (String, ThinLTOFunctionSummary)>,
    /// Global type hierarchy.
    pub type_hierarchy: TypeMetadataTable,
}

impl ThinLTOIndex {
    pub fn new() -> Self {
        ThinLTOIndex {
            summaries: HashMap::new(),
            global_func_map: HashMap::new(),
            type_hierarchy: TypeMetadataTable::new(),
        }
    }

    /// Add a module summary.
    pub fn add_module_summary(&mut self, summary: ThinLTOModuleSummary) {
        for func in &summary.functions {
            self.global_func_map.insert(
                func.name.clone(),
                (summary.module_name.clone(), func.clone()),
            );
        }
        self.summaries.insert(summary.module_name.clone(), summary);
    }

    /// Find all possible callees for a virtual call across modules.
    pub fn resolve_virtual_call(&self, type_id: &str, vtable_index: u32) -> Vec<String> {
        let mut targets = Vec::new();
        for (_, summary) in &self.summaries {
            for func in &summary.functions {
                for (tid, idx) in &func.vtable_slots {
                    if tid == type_id && *idx == vtable_index {
                        targets.push(func.name.clone());
                    }
                }
            }
        }
        targets
    }

    /// Check if cross-module inlining is profitable.
    pub fn should_import(&self, callee: &str, caller_inst_count: u64) -> bool {
        if let Some((_, summary)) = self.global_func_map.get(callee) {
            // Import small functions (≤100 instructions) for inlining
            summary.instruction_count <= 100
            // Don't import if caller is already large
            && caller_inst_count <= 10000
            // Don't import functions with inline assembly
            && !summary.has_inline_assembly
        } else {
            false
        }
    }

    /// Generate cross-module optimization hints.
    pub fn generate_optimization_hints(&self, module_name: &str) -> Vec<String> {
        let mut hints = Vec::new();
        if let Some(summary) = self.summaries.get(module_name) {
            for func in &summary.functions {
                if let Some(target) = func.callees.first() {
                    if self.should_import(target, func.instruction_count) {
                        hints.push(format!(
                            "!thinlto_hint = !{{!\"import\", !\"{}\"}} ; from {}",
                            target, func.name
                        ));
                    }
                }
            }
        }
        hints
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    // --- Type Metadata Tests ---
    #[test]
    fn test_type_metadata_register() {
        let mut table = TypeMetadataTable::new();
        table.register_type("_ZTV3Foo", 5, vec![0, 16, 32]);
        assert!(table.entries.contains_key("_ZTV3Foo"));
        assert_eq!(table.entries["_ZTV3Foo"].num_entries, 5);
        assert_eq!(table.entries["_ZTV3Foo"].address_points.len(), 3);
    }

    #[test]
    fn test_type_metadata_inheritance() {
        let mut table = TypeMetadataTable::new();
        table.register_type("Base", 4, vec![0]);
        table.register_type("Derived", 6, vec![0, 16]);
        table.add_derived_type("Base", "Derived");
        assert!(table.entries["Base"].derived_types.contains("Derived"));
        assert!(table.entries["Derived"].base_types.contains("Base"));
    }

    #[test]
    fn test_type_metadata_emit_ir() {
        let mut table = TypeMetadataTable::new();
        table.register_type("_ZTV3Foo", 3, vec![0, 8]);
        let ir = table.emit_type_metadata_ir("_ZTV3Foo");
        assert!(ir.contains("_ZTV3Foo"));
        assert!(ir.contains("typeid"));
    }

    // --- Virtual Call Analysis Tests ---
    #[test]
    fn test_vca_register_call() {
        let mut vca = VirtualCallAnalysis::new();
        vca.register_call("main", "_ZTV3Foo", 2);
        assert_eq!(vca.total_calls, 1);
        assert_eq!(vca.call_sites.len(), 1);
        assert_eq!(vca.call_sites[0].vtable_index, 2);
    }

    #[test]
    fn test_vca_resolve_monomorphic() {
        let mut vca = VirtualCallAnalysis::new();
        vca.target_map
            .insert(("_ZTV3Foo".into(), 0), vec!["Foo::bar".into()]);
        vca.register_call("main", "_ZTV3Foo", 0);
        vca.resolve_targets();
        assert!(vca.call_sites[0].is_monomorphic);
        assert_eq!(vca.monomorphic_calls, 1);
    }

    #[test]
    fn test_vca_resolve_polymorphic() {
        let mut vca = VirtualCallAnalysis::new();
        vca.target_map.insert(
            ("_ZTV3Foo".into(), 1),
            vec!["Foo::baz".into(), "Bar::baz".into()],
        );
        vca.register_call("main", "_ZTV3Foo", 1);
        vca.resolve_targets();
        assert!(!vca.call_sites[0].is_monomorphic);
        assert_eq!(vca.polymorphic_calls, 1);
    }

    #[test]
    fn test_vca_get_single_target() {
        let mut vca = VirtualCallAnalysis::new();
        vca.target_map
            .insert(("_ZTV3Foo".into(), 0), vec!["Foo::bar".into()]);
        let target = vca.get_single_target("main", "_ZTV3Foo", 0);
        assert_eq!(target, Some("Foo::bar".into()));
    }

    // --- Whole-Program Devirt Tests ---
    #[test]
    fn test_devirt_direct_call() {
        let vca = VirtualCallAnalysis::new();
        let mut devirt = WholeProgramDevirt::new(vca);
        let site = VirtualCallSite {
            caller: "main".into(),
            vtable_type_id: "Foo".into(),
            vtable_index: 0,
            possible_callees: vec!["Foo::bar".into()],
            is_monomorphic: true,
        };
        let result = devirt.devirtualize(&site);
        match result {
            DevirtResult::DirectCall(target) => assert_eq!(target, "Foo::bar"),
            _ => panic!("Expected DirectCall"),
        }
    }

    #[test]
    fn test_devirt_speculative() {
        let vca = VirtualCallAnalysis::new();
        let mut devirt = WholeProgramDevirt::new(vca);
        let site = VirtualCallSite {
            caller: "main".into(),
            vtable_type_id: "Foo".into(),
            vtable_index: 0,
            possible_callees: vec!["Foo::bar".into(), "Bar::bar".into()],
            is_monomorphic: false,
        };
        let result = devirt.devirtualize(&site);
        match result {
            DevirtResult::Speculative {
                primary_target,
                fallback_targets,
            } => {
                assert_eq!(primary_target, "Foo::bar");
                assert_eq!(fallback_targets.len(), 1);
            }
            _ => panic!("Expected Speculative"),
        }
    }

    #[test]
    fn test_devirt_emit_direct() {
        let vca = VirtualCallAnalysis::new();
        let devirt = WholeProgramDevirt::new(vca);
        let ir = devirt.emit_devirtualized_call("Foo__bar", "%this");
        assert!(ir.contains("Foo__bar"));
        assert!(ir.contains("devirtualized"));
    }

    // --- VCall Visibility Tests ---
    #[test]
    fn test_vcall_visibility_public() {
        let mut checker = VCallVisibilityChecker::new();
        checker.set_visibility("Foo", 0, VCallVisibility::VCallVisibilityPublic);
        assert!(checker.can_devirtualize("Foo", 0));
    }

    #[test]
    fn test_vcall_visibility_tu_non_devirt() {
        let mut checker = VCallVisibilityChecker::new();
        checker.set_visibility("Foo", 1, VCallVisibility::VCallVisibilityTranslationUnit);
        assert!(!checker.can_devirtualize("Foo", 1));
    }

    #[test]
    fn test_vcall_visibility_unknown() {
        let checker = VCallVisibilityChecker::new();
        assert!(!checker.can_devirtualize("Unknown", 0));
    }

    // --- CFI Tests ---
    #[test]
    fn test_cfi_register_type() {
        let mut cfi = CFIMetadata::new();
        cfi.register_type("_ZTV3Foo", 16, true);
        assert_eq!(cfi.types.len(), 1);
        assert!(cfi.types[0].is_exported);
    }

    #[test]
    fn test_cfi_blacklist() {
        let mut cfi = CFIMetadata::new();
        cfi.blacklist("some_unsafe_func");
        assert!(cfi.blacklist.contains("some_unsafe_func"));
    }

    #[test]
    fn test_cfi_emit_check() {
        let mut cfi = CFIMetadata::new();
        let ir = cfi.emit_cfi_check("%vtable_ptr", "_ZTV3Foo", 16);
        assert!(ir.contains("CFI check"));
        assert!(ir.contains("__cfi_check_fail"));
        assert!(ir.contains("__cfi_typeid__ZTV3Foo"));
        assert_eq!(cfi.stats_checks_inserted, 1);
    }

    // --- LTO Visibility Tests ---
    #[test]
    fn test_lto_promotion_hidden_not_address_taken() {
        let mut promo = LTOVisibilityPromotion::new();
        promo.register_symbol("foo", SymbolVisibility::Hidden, false);
        promo.promote();
        let sym = &promo.symbols[0];
        assert_eq!(sym.promoted_visibility, SymbolVisibility::Internal);
        assert_eq!(promo.internalized_count, 1);
    }

    #[test]
    fn test_lto_promotion_address_taken() {
        let mut promo = LTOVisibilityPromotion::new();
        promo.register_symbol("bar", SymbolVisibility::Hidden, true);
        promo.promote();
        let sym = &promo.symbols[0];
        assert_eq!(sym.promoted_visibility, SymbolVisibility::Internal);
    }

    #[test]
    fn test_lto_promotion_unpromotable() {
        let mut promo = LTOVisibilityPromotion::new();
        promo.register_symbol("exported_func", SymbolVisibility::Hidden, false);
        promo.mark_unpromotable("exported_func");
        promo.promote();
        // Still gets promoted but we track it differently
        // (implementation detail: unpromotable prevents can_be_internalized)
    }

    #[test]
    fn test_lto_emit_linkage() {
        let sym = LTOSymbolInfo {
            name: "foo".into(),
            original_visibility: SymbolVisibility::Hidden,
            promoted_visibility: SymbolVisibility::Internal,
            is_address_taken: false,
            can_be_internalized: true,
        };
        let promo = LTOVisibilityPromotion::new();
        assert_eq!(promo.emit_linkage(&sym), "internal");
    }

    // --- GlobalDCE Tests ---
    #[test]
    fn test_global_dce_live_root() {
        let mut dce = GlobalDCE::new();
        dce.add_root("main");
        dce.add_function("main", vec!["helper"], vec!["global_x"]);
        dce.add_function("helper", vec![], vec![]);
        dce.add_function("dead_func", vec![], vec![]);
        dce.add_global("global_x");
        dce.add_global("dead_global");
        dce.eliminate();
        assert!(dce.is_live("main"));
        assert!(dce.is_live("helper"));
        assert!(dce.is_global_live("global_x"));
        assert!(!dce.is_live("dead_func"));
        assert!(dce.dead_functions_removed >= 1);
    }

    #[test]
    fn test_global_dce_dead_global() {
        let mut dce = GlobalDCE::new();
        dce.add_root("main");
        dce.add_function("main", vec![], vec![]);
        dce.add_global("dead_global");
        dce.eliminate();
        assert!(dce.dead_globals_removed >= 1);
    }

    // --- ThinLTO Tests ---
    #[test]
    fn test_thinlto_add_summary() {
        let mut index = ThinLTOIndex::new();
        let summary = ThinLTOModuleSummary {
            module_name: "mod1".into(),
            module_hash: "abc123".into(),
            functions: vec![ThinLTOFunctionSummary {
                name: "foo".into(),
                instruction_count: 50,
                callees: vec!["bar".into()],
                referenced_globals: vec![],
                is_exported: true,
                has_inline_assembly: false,
                is_available_externally: false,
                vtable_slots: vec![("_ZTV3Foo".into(), 0)],
            }],
            globals: vec![],
            type_ids: vec!["_ZTV3Foo".into()],
            vcall_visibility: HashMap::new(),
        };
        index.add_module_summary(summary);
        assert!(index.global_func_map.contains_key("foo"));
    }

    #[test]
    fn test_thinlto_resolve_virtual_call() {
        let mut index = ThinLTOIndex::new();
        let summary = ThinLTOModuleSummary {
            module_name: "mod1".into(),
            module_hash: "abc".into(),
            functions: vec![ThinLTOFunctionSummary {
                name: "Foo::bar".into(),
                instruction_count: 20,
                callees: vec![],
                referenced_globals: vec![],
                is_exported: false,
                has_inline_assembly: false,
                is_available_externally: false,
                vtable_slots: vec![("_ZTV3Foo".into(), 0)],
            }],
            globals: vec![],
            type_ids: vec![],
            vcall_visibility: HashMap::new(),
        };
        index.add_module_summary(summary);
        let targets = index.resolve_virtual_call("_ZTV3Foo", 0);
        assert!(!targets.is_empty());
        assert!(targets.contains(&"Foo::bar".into()));
    }

    #[test]
    fn test_thinlto_should_import_small() {
        let mut index = ThinLTOIndex::new();
        let summary = ThinLTOModuleSummary {
            module_name: "mod1".into(),
            module_hash: "abc".into(),
            functions: vec![ThinLTOFunctionSummary {
                name: "small_func".into(),
                instruction_count: 30,
                callees: vec![],
                referenced_globals: vec![],
                is_exported: false,
                has_inline_assembly: false,
                is_available_externally: false,
                vtable_slots: vec![],
            }],
            globals: vec![],
            type_ids: vec![],
            vcall_visibility: HashMap::new(),
        };
        index.add_module_summary(summary);
        assert!(index.should_import("small_func", 50));
    }

    #[test]
    fn test_thinlto_should_not_import_large() {
        let mut index = ThinLTOIndex::new();
        let summary = ThinLTOModuleSummary {
            module_name: "mod1".into(),
            module_hash: "abc".into(),
            functions: vec![ThinLTOFunctionSummary {
                name: "big_func".into(),
                instruction_count: 500,
                callees: vec![],
                referenced_globals: vec![],
                is_exported: false,
                has_inline_assembly: false,
                is_available_externally: false,
                vtable_slots: vec![],
            }],
            globals: vec![],
            type_ids: vec![],
            vcall_visibility: HashMap::new(),
        };
        index.add_module_summary(summary);
        assert!(!index.should_import("big_func", 50));
    }

    #[test]
    fn test_thinlto_optimization_hints() {
        let mut index = ThinLTOIndex::new();
        let summary = ThinLTOModuleSummary {
            module_name: "mod_main".into(),
            module_hash: "xyz".into(),
            functions: vec![
                ThinLTOFunctionSummary {
                    name: "caller".into(),
                    instruction_count: 50,
                    callees: vec!["callee".into()],
                    referenced_globals: vec![],
                    is_exported: true,
                    has_inline_assembly: false,
                    is_available_externally: false,
                    vtable_slots: vec![],
                },
                ThinLTOFunctionSummary {
                    name: "callee".into(),
                    instruction_count: 10,
                    callees: vec![],
                    referenced_globals: vec![],
                    is_exported: false,
                    has_inline_assembly: false,
                    is_available_externally: false,
                    vtable_slots: vec![],
                },
            ],
            globals: vec![],
            type_ids: vec![],
            vcall_visibility: HashMap::new(),
        };
        index.add_module_summary(summary);
        let hints = index.generate_optimization_hints("mod_main");
        assert!(!hints.is_empty());
    }

    // --- Integration Tests ---
    #[test]
    fn test_full_devirt_pipeline() {
        // Set up type metadata
        let mut type_table = TypeMetadataTable::new();
        type_table.register_type("Base", 4, vec![0, 16]);
        type_table.register_type("Derived", 6, vec![0, 16, 32]);
        type_table.add_derived_type("Base", "Derived");

        // Set up VCA
        let mut vca = VirtualCallAnalysis::new();
        vca.target_map
            .insert(("Base".into(), 0), vec!["Base::foo".into()]);
        vca.target_map
            .insert(("Derived".into(), 0), vec!["Derived::foo".into()]);
        vca.register_call("main", "Base", 0);
        vca.resolve_targets();

        // Devirtualize
        let mut devirt = WholeProgramDevirt::new(vca);
        let site = devirt.analysis.call_sites[0].clone();
        let result = devirt.devirtualize(&site);
        match result {
            DevirtResult::DirectCall(_) => {}
            _ => panic!("Expected DirectCall"),
        }
        assert_eq!(devirt.devirtualized_calls, 1);
    }
}