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;
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(),
}
}
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());
}
let _type_map = Self::merge_types(modules);
let global_map = Self::merge_globals(modules);
Self::merge_comdats(modules);
let _resolved = Self::resolve_linkage(&global_map);
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
}
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
}
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
}
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);
}
}
}
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()
}
}
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;
}
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
}
}
}
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)
}
#[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());
}
}