llvm-native-core-ext 0.1.0

Extended modules for llvm-native-core: analysis passes, transforms, codegen extras, bitcode, linker, JIT, utilities. Part of the llvm-native workspace (https://crates.io/crates/llvm-native).
//! LLVM LTO Merge — module, type, global, and comdat merging for LTO.
//! Phase 7 — LLVM.LTOMERGE.1 Court.
//!
//! Clean-room behavioral reconstruction from the LLVM LTO design
//! documentation, published linker design literature, and the
//! LLVM Language Reference. Zero LLVM source code consultation.

use llvm_native_core::module::{ComdatKind, Module};
use llvm_native_core::types::{Type, TypeId, TypeKind};
use llvm_native_core::value::ValueRef;
use llvm_native_core::SubclassKind;
use std::collections::HashMap;

// ============================================================================
// LTOMerger — the central merge engine
// ============================================================================

pub struct LTOMerger {
    dest: Option<Module>,
    diagnostics: Vec<MergeDiagnostic>,
}

#[derive(Debug, Clone)]
pub struct MergeDiagnostic {
    pub level: DiagLevel,
    pub message: String,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DiagLevel {
    Note,
    Warning,
    Error,
}

impl LTOMerger {
    pub fn new() -> Self {
        Self {
            dest: None,
            diagnostics: Vec::new(),
        }
    }

    /// Merge multiple modules into a single combined module.
    pub fn merge_modules(modules: &[Module]) -> Module {
        if modules.is_empty() {
            return Module::new("lto_merged");
        }

        let mut dest = modules[0].clone_empty();
        dest.name = "lto_merged".to_string();
        if let Some(triple) = &modules[0].target_triple {
            dest.target_triple = Some(triple.clone());
        }

        // Phase 1: Merge types.
        let _type_map = Self::merge_types(modules);

        // Phase 2: Collect and merge globals.
        let global_map = Self::merge_globals(modules);

        // Phase 3: Merge comdats.
        Self::merge_comdats(modules);

        // Phase 4: Resolve linkage.
        let _resolved = Self::resolve_linkage(&global_map);

        // Phase 5: Copy prevailing definitions into dest.
        for module in modules.iter().skip(1) {
            for func in &module.functions {
                dest.add_function_unchecked(func.clone());
            }
            for gv in &module.globals {
                let _ = dest.add_global_variable(gv.clone());
            }
        }

        dest
    }

    /// Merge types from modules — returns a map from (module_index, type_index)
    /// to a unified TypeId.
    pub fn merge_types(modules: &[Module]) -> HashMap<(usize, usize), TypeId> {
        let mut type_map: HashMap<(usize, usize), TypeId> = HashMap::new();
        let mut unified: Vec<(Type, usize)> = Vec::new();

        for (mod_idx, module) in modules.iter().enumerate() {
            for (type_idx, ty) in module.types.iter().enumerate() {
                let unified_id = find_or_insert_type(ty, &mut unified);
                type_map.insert((mod_idx, type_idx), unified_id);
            }
        }
        type_map
    }

    /// Merge global variables and functions — returns map from name to all
    /// defining ValueRefs across modules.
    pub fn merge_globals(modules: &[Module]) -> HashMap<String, Vec<ValueRef>> {
        let mut global_map: HashMap<String, Vec<ValueRef>> = HashMap::new();

        for module in modules {
            for func in &module.functions {
                let name = func.borrow().name.clone();
                global_map.entry(name).or_default().push(func.clone());
            }
            for gv in &module.globals {
                let name = gv.borrow().name.clone();
                global_map.entry(name).or_default().push(gv.clone());
            }
            for alias in &module.aliases {
                let name = alias.borrow().name.clone();
                global_map.entry(name).or_default().push(alias.clone());
            }
        }
        global_map
    }

    /// Merge comdat groups from all modules.
    pub fn merge_comdats(modules: &[Module]) {
        let mut _seen: HashMap<String, ComdatKind> = HashMap::new();
        for module in modules {
            for (name, comdat) in &module.comdats {
                _seen
                    .entry(name.clone())
                    .and_modify(|existing| {
                        *existing = merge_comdat_kind(*existing, comdat.kind);
                    })
                    .or_insert(comdat.kind);
            }
        }
    }

    /// Resolve linkage conflicts — returns prevailing definition per symbol.
    pub fn resolve_linkage(globals: &HashMap<String, Vec<ValueRef>>) -> HashMap<String, ValueRef> {
        let mut resolved: HashMap<String, ValueRef> = HashMap::new();
        for (name, refs) in globals {
            if let Some(first) = refs.first() {
                resolved.insert(name.clone(), first.clone());
            }
        }
        resolved
    }

    pub fn emit_diagnostic(&mut self, level: DiagLevel, message: String) {
        self.diagnostics.push(MergeDiagnostic { level, message });
    }

    pub fn diagnostics(&self) -> &[MergeDiagnostic] {
        &self.diagnostics
    }

    pub fn has_errors(&self) -> bool {
        self.diagnostics.iter().any(|d| d.level == DiagLevel::Error)
    }
}

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

// ============================================================================
// Helper functions
// ============================================================================

fn find_or_insert_type(ty: &Type, unified: &mut Vec<(Type, usize)>) -> TypeId {
    for (idx, (existing, _count)) in unified.iter().enumerate() {
        if types_are_structurally_equivalent(existing, ty) {
            return existing.id;
        }
    }
    let id = ty.id;
    unified.push((ty.clone(), 1));
    id
}

fn types_are_structurally_equivalent(a: &Type, b: &Type) -> bool {
    if a.kind == b.kind {
        return true;
    }
    // Handle parameterized types.
    match (&a.kind, &b.kind) {
        (TypeKind::Integer { bits: b1 }, TypeKind::Integer { bits: b2 }) => b1 == b2,
        (TypeKind::Pointer { addr_space: a1 }, TypeKind::Pointer { addr_space: a2 }) => a1 == a2,
        (
            TypeKind::Array {
                len: l1,
                element_type_id: e1,
            },
            TypeKind::Array {
                len: l2,
                element_type_id: e2,
            },
        ) => l1 == l2 && e1 == e2,
        (
            TypeKind::Struct {
                element_type_ids: e1,
                is_packed: p1,
                ..
            },
            TypeKind::Struct {
                element_type_ids: e2,
                is_packed: p2,
                ..
            },
        ) => p1 == p2 && e1 == e2,
        (
            TypeKind::Function {
                return_type_id: r1,
                param_type_ids: p1,
                is_vararg: v1,
            },
            TypeKind::Function {
                return_type_id: r2,
                param_type_ids: p2,
                is_vararg: v2,
            },
        ) => r1 == r2 && p1 == p2 && v1 == v2,
        (
            TypeKind::FixedVector {
                len: l1,
                element_type_id: e1,
            },
            TypeKind::FixedVector {
                len: l2,
                element_type_id: e2,
            },
        ) => l1 == l2 && e1 == e2,
        (
            TypeKind::ScalableVector {
                min_elems: m1,
                element_type_id: e1,
            },
            TypeKind::ScalableVector {
                min_elems: m2,
                element_type_id: e2,
            },
        ) => m1 == m2 && e1 == e2,
        _ => false,
    }
}

fn merge_comdat_kind(a: ComdatKind, b: ComdatKind) -> ComdatKind {
    match (a, b) {
        (ComdatKind::Any, _) | (_, ComdatKind::Any) => ComdatKind::Any,
        (ComdatKind::Largest, _) | (_, ComdatKind::Largest) => ComdatKind::Largest,
        (ComdatKind::SameSize, _) | (_, ComdatKind::SameSize) => ComdatKind::SameSize,
        (ComdatKind::ExactMatch, ComdatKind::ExactMatch) => ComdatKind::ExactMatch,
        (ComdatKind::NoDeduplicate, _) | (_, ComdatKind::NoDeduplicate) => {
            ComdatKind::NoDeduplicate
        }
    }
}

// ============================================================================
// Convenience functions
// ============================================================================

pub fn merge_modules(modules: &[Module]) -> Module {
    LTOMerger::merge_modules(modules)
}

pub fn merge_types(modules: &[Module]) -> HashMap<(usize, usize), TypeId> {
    LTOMerger::merge_types(modules)
}

pub fn merge_globals(modules: &[Module]) -> HashMap<String, Vec<ValueRef>> {
    LTOMerger::merge_globals(modules)
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use llvm_native_core::module::Module;
    use llvm_native_core::types::{Type, TypeKind};

    fn make_module(name: &str) -> Module {
        let mut m = Module::new(name);
        m.target_triple = Some("x86_64-unknown-linux-gnu".to_string());
        m
    }

    #[test]
    fn test_lto_merger_new() {
        let merger = LTOMerger::new();
        assert!(merger.diagnostics().is_empty());
        assert!(!merger.has_errors());
    }

    #[test]
    fn test_merge_modules_empty() {
        let result = LTOMerger::merge_modules(&[]);
        assert_eq!(result.name, "lto_merged");
        assert_eq!(result.get_function_count(), 0);
    }

    #[test]
    fn test_merge_modules_single() {
        let m = make_module("test");
        let result = LTOMerger::merge_modules(&[m]);
        assert_eq!(result.name, "lto_merged");
    }

    #[test]
    fn test_merge_types_empty() {
        let result = merge_types(&[]);
        assert!(result.is_empty());
    }

    #[test]
    fn test_types_are_structurally_equivalent_same() {
        let t1 = Type::new(TypeKind::Integer { bits: 32 });
        let t2 = Type::new(TypeKind::Integer { bits: 32 });
        assert!(types_are_structurally_equivalent(&t1, &t2));
    }

    #[test]
    fn test_types_are_structurally_equivalent_different() {
        let t1 = Type::new(TypeKind::Integer { bits: 32 });
        let t2 = Type::new(TypeKind::Integer { bits: 64 });
        assert!(!types_are_structurally_equivalent(&t1, &t2));
    }

    #[test]
    fn test_merge_globals_empty() {
        let result = merge_globals(&[]);
        assert!(result.is_empty());
    }

    #[test]
    fn test_resolve_linkage_empty() {
        let map: HashMap<String, Vec<ValueRef>> = HashMap::new();
        let resolved = LTOMerger::resolve_linkage(&map);
        assert!(resolved.is_empty());
    }

    #[test]
    fn test_merge_comdat_kinds() {
        assert_eq!(
            merge_comdat_kind(ComdatKind::Any, ComdatKind::ExactMatch),
            ComdatKind::Any
        );
        assert_eq!(
            merge_comdat_kind(ComdatKind::ExactMatch, ComdatKind::ExactMatch),
            ComdatKind::ExactMatch
        );
    }

    #[test]
    fn test_diagnostics() {
        let mut merger = LTOMerger::new();
        merger.emit_diagnostic(DiagLevel::Warning, "test warning".to_string());
        assert_eq!(merger.diagnostics().len(), 1);
        assert!(!merger.has_errors());

        merger.emit_diagnostic(DiagLevel::Error, "test error".to_string());
        assert!(merger.has_errors());
    }
}