llvm-native-core-ext 0.1.0

//! ValueMapper — maps values from one module/function to another during
//! cloning, inlining, and module linking.
//! Clean-room behavioral reconstruction.
//!
//! @llvm_behavior: When cloning IR (for inlining, loop unrolling, function
//! specialization, or module linking), every Value in the source must be
//! remapped to a corresponding Value in the destination. The ValueMapper
//! maintains the mapping and provides utilities for cloning individual
//! values, instructions, basic blocks, and entire functions.
//!
//! Key operations:
//! - Map a value from source to destination context
//! - Map types (which may be different in the destination module's type system)
//! - Remap operands of cloned instructions to point to cloned values
//! - Clone basic blocks while preserving CFG structure
//! - Clone entire functions with optional name changes
//! - Clone entire modules (deep copy)
//!
//! The value_map uses the source Value's vid (unique ID) as the key,
//! mapping to the cloned ValueRef in the destination.

use llvm_native_core::types::{Type, TypeId, TypeKind};
use llvm_native_core::value::{valref, SubclassKind, Value, ValueRef};
use std::collections::HashMap;

// ============================================================================
// ValueMapper
// ============================================================================

/// Maps values from a source context to a destination context during cloning.
///
/// Maintains three maps:
/// - `value_map`: source Value vid → cloned ValueRef
/// - `type_map`: source TypeId → destination TypeId
/// - `local_map`: source local value ID → destination local ID
pub struct ValueMapper {
    /// Map from source value vid to cloned value.
    pub value_map: HashMap<usize, ValueRef>,
    /// Map from source type ID to destination type ID.
    pub type_map: HashMap<TypeId, TypeId>,
    /// Map from source local ID to destination local ID.
    pub local_map: HashMap<usize, usize>,
}

impl ValueMapper {
    /// Create a new empty ValueMapper.
    pub fn new() -> Self {
        Self {
            value_map: HashMap::new(),
            type_map: HashMap::new(),
            local_map: HashMap::new(),
        }
    }

    // ========================================================================
    // Value mapping
    // ========================================================================

    /// Map a value from source to destination context.
    ///
    /// If the value has already been mapped (exists in value_map), return
    /// the existing mapping. Otherwise, return the original value
    /// (identity mapping for values not yet cloned).
    pub fn map_value(&self, val: &ValueRef) -> ValueRef {
        let vid = val.borrow().vid as usize;
        self.value_map
            .get(&vid)
            .cloned()
            .unwrap_or_else(|| val.clone())
    }

    /// Map an instruction from source to destination.
    pub fn map_instruction(&self, inst: &ValueRef) -> ValueRef {
        self.map_value(inst)
    }

    /// Remap the operands of an instruction to point to cloned values.
    ///
    /// This updates the instruction's operand list so that each operand
    /// refers to its mapped counterpart in the destination.
    pub fn remap_operands(&self, inst: &mut ValueRef) {
        let mut i = inst.borrow_mut();
        let mut new_operands = Vec::with_capacity(i.operands.len());

        for operand in &i.operands {
            let mapped = self.map_value(operand);
            new_operands.push(mapped);
        }

        i.operands = new_operands;
        i.num_operands = i.operands.len();
    }

    /// Register a mapping from a source value to a cloned value.
    pub fn register_value(&mut self, src: &ValueRef, dest: ValueRef) {
        let src_vid = src.borrow().vid as usize;
        self.value_map.insert(src_vid, dest);
    }

    // ========================================================================
    // Type mapping
    // ========================================================================

    /// Map a type from source to destination context.
    ///
    /// If the type has been mapped, returns the mapped type. Otherwise
    /// returns the original type (identity mapping).
    pub fn map_type(&self, ty: &Type) -> Type {
        if let Some(&mapped_id) = self.type_map.get(&ty.id) {
            // Return a type with the mapped ID and same kind
            Type {
                id: mapped_id,
                kind: ty.kind.clone(),
            }
        } else {
            ty.clone()
        }
    }

    /// Register a type mapping.
    pub fn register_type(&mut self, src_id: TypeId, dest_id: TypeId) {
        self.type_map.insert(src_id, dest_id);
    }

    // ========================================================================
    // Cloning
    // ========================================================================

    /// Clone a value into a destination context.
    ///
    /// Creates a new Value that is a copy of the source, registers the
    /// mapping, and returns the clone.
    pub fn clone_into(&self, src: &ValueRef, dest: &ValueRef) {
        // This is a simplified clone — in practice, we'd deep-clone the
        // Value and all its operands recursively.
        //
        // The destination is pre-created, and we copy over the relevant
        // properties:
        let src_val = src.borrow();
        let mut dest_val = dest.borrow_mut();

        dest_val.ty = self.map_type(&src_val.ty);
        dest_val.subclass = src_val.subclass;
        dest_val.opcode = src_val.opcode;
        dest_val.num_operands = src_val.num_operands;

        // Clone operands (shallow references that will be remapped)
        dest_val.operands = src_val.operands.clone();

        // Copy parent reference (may be remapped later)
        dest_val.parent = src_val.parent.clone();

        // Copy metadata
        dest_val.metadata = src_val.metadata.clone();
        dest_val.is_used_by_md = src_val.is_used_by_md;

        dest_val.subclass_data = src_val.subclass_data;
    }

    /// Clone a basic block from source to destination function.
    ///
    /// Creates a new basic block in the destination function that is a
    /// copy of the source basic block. All instructions are cloned and
    /// their operands are remapped.
    ///
    /// Returns the cloned basic block.
    pub fn clone_basic_block(&self, src_bb: &ValueRef, dest_func: &ValueRef) -> ValueRef {
        let src = src_bb.borrow();

        // Create the destination basic block
        let mut dest_bb = Value::new(Type::label())
            .with_subclass(SubclassKind::BasicBlock)
            .named(&src.name);

        dest_bb.parent = Some(dest_func.clone());

        // Clone each instruction
        for inst in &src.operands {
            let src_inst = inst.borrow();

            // Skip non-instructions
            if !src_inst.is_instruction() {
                continue;
            }

            let mut dest_inst = Value::new(self.map_type(&src_inst.ty))
                .with_subclass(src_inst.subclass)
                .named(&src_inst.name);

            dest_inst.opcode = src_inst.opcode;
            dest_inst.operands = src_inst.operands.clone();
            dest_inst.num_operands = src_inst.operands.len();
            dest_inst.parent = Some(dest_func.clone());
            dest_inst.metadata = src_inst.metadata.clone();
            dest_inst.is_used_by_md = src_inst.is_used_by_md;
            dest_inst.subclass_data = src_inst.subclass_data;

            let dest_inst_ref = valref(dest_inst);
            dest_bb.operands.push(dest_inst_ref);
        }

        valref(dest_bb)
    }

    /// Clone an entire function.
    ///
    /// Creates a new function that is a deep copy of the source function.
    /// All basic blocks and instructions are cloned, and operands are
    /// remapped to point to the cloned values.
    ///
    /// Returns the cloned function.
    pub fn clone_function(&self, src: &ValueRef, name: &str) -> ValueRef {
        let src_func = src.borrow();

        // Create the destination function value
        let mut dest_func = Value::new(src_func.ty.clone())
            .with_subclass(SubclassKind::Function)
            .named(name);

        dest_func.return_type = src_func.return_type.clone();

        let dest_func_ref = valref(dest_func);

        // Clone basic blocks
        for operand in &src_func.operands {
            if operand.borrow().is_basic_block() {
                let cloned_bb = self.clone_basic_block(operand, &dest_func_ref);
                dest_func_ref.borrow_mut().operands.push(cloned_bb);
            } else if operand.borrow().subclass == SubclassKind::Argument {
                // Clone argument
                let src_arg = operand.borrow();
                let mut arg = Value::new(self.map_type(&src_arg.ty))
                    .with_subclass(SubclassKind::Argument)
                    .named(&src_arg.name);
                arg.parent = Some(dest_func_ref.clone());
                dest_func_ref.borrow_mut().operands.push(valref(arg));
            }
        }

        let num_ops = dest_func_ref.borrow().operands.len();
        dest_func_ref.borrow_mut().num_operands = num_ops;

        dest_func_ref
    }
}

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

// ============================================================================
// Free Functions
// ============================================================================

/// Clone an entire module (deep copy).
///
/// Creates a new module that is a deep copy of the source module with
/// all functions, globals, and metadata cloned.
pub fn clone_module(src: &llvm_native_core::module::Module, name: &str) -> llvm_native_core::module::Module {
    let mut dest = llvm_native_core::module::Module::new(name);
    let mut mapper = ValueMapper::new();

    // Copy module metadata
    dest.source_filename = src.source_filename.clone();
    dest.target_triple = src.target_triple.clone();
    dest.data_layout = src.data_layout.clone();

    // Clone types
    for ty in &src.types {
        let cloned_ty = ty.clone(); // Types are value types, safe to clone
        dest.types.push(cloned_ty);
    }

    // Clone named types
    dest.named_types = src.named_types.clone();

    // Clone globals
    for global in &src.globals {
        let g = global.borrow();
        let mut cloned = Value::new(g.ty.clone())
            .with_subclass(SubclassKind::GlobalVariable)
            .named(&g.name);
        cloned.subclass_data = g.subclass_data;

        let cloned_ref = valref(cloned);
        mapper.register_value(global, cloned_ref.clone());
        dest.globals.push(cloned_ref);
    }

    // Clone functions
    for func in &src.functions {
        let f = func.borrow();
        let cloned = mapper.clone_function(func, &f.name);
        dest.functions.push(cloned);
    }

    // Clone aliases
    for alias in &src.aliases {
        let a = alias.borrow();
        let mut cloned = Value::new(a.ty.clone())
            .with_subclass(SubclassKind::GlobalAlias)
            .named(&a.name);
        let cloned_ref = valref(cloned);
        mapper.register_value(alias, cloned_ref.clone());
        dest.aliases.push(cloned_ref);
    }

    // Clone ifuncs
    for ifunc in &src.ifuncs {
        let i = ifunc.borrow();
        let mut cloned = Value::new(i.ty.clone())
            .with_subclass(SubclassKind::GlobalIFunc)
            .named(&i.name);
        let cloned_ref = valref(cloned);
        mapper.register_value(ifunc, cloned_ref.clone());
        dest.ifuncs.push(cloned_ref);
    }

    // Clone comdats
    dest.comdats = src.comdats.clone();

    // Clone named metadata
    dest.named_metadata = src.named_metadata.clone();

    // Clone attribute groups
    dest.attr_groups.clone_from(&src.attr_groups);

    // Clone module flags
    dest.flags.clone_from(&src.flags);

    // Clone imported function names
    dest.imported_functions = src.imported_functions.clone();

    dest
}

/// Clone a function into a destination module.
///
/// Creates a copy of the source function in the destination module
/// with all operands remapped. This is useful for inlining, function
/// specialization, and inter-module code motion.
pub fn clone_function_into(
    src: &ValueRef,
    dest_module: &mut llvm_native_core::module::Module,
    name: &str,
) -> ValueRef {
    let mapper = ValueMapper::new();
    let cloned = mapper.clone_function(src, name);

    // Update the cloned function's operands to use the destination module's
    // types and values (remap)
    for operand in &cloned.borrow().operands.clone() {
        let mut op = operand.borrow_mut();
        // Remap types
        op.ty = mapper.map_type(&op.ty);
        // Remap parent
        if op.parent.is_some() {
            // Parent stays as the cloned function
        }
    }

    dest_module.functions.push(cloned.clone());
    cloned
}

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

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

    fn make_value(ty: Type) -> ValueRef {
        valref(Value::new(ty).named("test"))
    }

    #[test]
    fn test_value_mapper_new() {
        let mapper = ValueMapper::new();
        assert!(mapper.value_map.is_empty());
        assert!(mapper.type_map.is_empty());
        assert!(mapper.local_map.is_empty());
    }

    #[test]
    fn test_default() {
        let mapper = ValueMapper::default();
        assert!(mapper.value_map.is_empty());
    }

    #[test]
    fn test_register_and_map_value() {
        let mut mapper = ValueMapper::new();
        let src = make_value(Type::i32());
        let dest = make_value(Type::i32());

        mapper.register_value(&src, dest.clone());
        let mapped = mapper.map_value(&src);
        assert!(Rc::ptr_eq(&mapped, &dest));
    }

    #[test]
    fn test_map_unregistered_value() {
        let mapper = ValueMapper::new();
        let val = make_value(Type::i32());
        // Unregistered values map to themselves
        let mapped = mapper.map_value(&val);
        assert!(Rc::ptr_eq(&mapped, &val));
    }

    #[test]
    fn test_map_instruction() {
        let mut mapper = ValueMapper::new();
        let src = make_value(Type::new(TypeKind::Integer { bits: 32 }));
        src.borrow_mut().subclass = SubclassKind::Instruction;
        let dest = make_value(Type::new(TypeKind::Integer { bits: 32 }));
        dest.borrow_mut().subclass = SubclassKind::Instruction;

        mapper.register_value(&src, dest.clone());
        assert!(Rc::ptr_eq(&mapper.map_instruction(&src), &dest));
    }

    #[test]
    fn test_map_type_registered() {
        let mut mapper = ValueMapper::new();
        let ty1 = Type::i32();
        let ty2 = Type::i64();
        mapper.register_type(ty1.id, ty2.id);

        let mapped = mapper.map_type(&ty1);
        assert_eq!(mapped.id, ty2.id);
    }

    #[test]
    fn test_map_type_unregistered() {
        let mapper = ValueMapper::new();
        let ty = Type::i32();
        let mapped = mapper.map_type(&ty);
        assert_eq!(mapped.id, ty.id);
        assert_eq!(mapped.kind, ty.kind);
    }

    #[test]
    fn test_remap_operands() {
        let mut mapper = ValueMapper::new();

        let op1 = make_value(Type::i32());
        let op1_mapped = make_value(Type::i32());
        mapper.register_value(&op1, op1_mapped.clone());

        let inst = make_value(Type::i32());
        inst.borrow_mut().operands.push(op1.clone());
        inst.borrow_mut().num_operands = 1;

        let mut inst_ref = inst.clone();
        mapper.remap_operands(&mut inst_ref);

        let mapped_inst = inst_ref.borrow();
        assert_eq!(mapped_inst.operands.len(), 1);
        assert!(Rc::ptr_eq(&mapped_inst.operands[0], &op1_mapped));
    }

    #[test]
    fn test_clone_basic_block() {
        let mapper = ValueMapper::new();

        // Create source BB
        let src_bb = valref(
            Value::new(Type::label())
                .with_subclass(SubclassKind::BasicBlock)
                .named("src_bb"),
        );

        // Add an instruction
        let inst = valref(
            Value::new(Type::i32())
                .with_subclass(SubclassKind::Instruction)
                .named("add"),
        );
        inst.borrow_mut().opcode = Some(llvm_native_core::opcode::Opcode::Add);
        src_bb.borrow_mut().operands.push(inst);

        // Create destination function
        let dest_func = valref(
            Value::new(Type::void())
                .with_subclass(SubclassKind::Function)
                .named("dest_func"),
        );

        let cloned = mapper.clone_basic_block(&src_bb, &dest_func);
        let c = cloned.borrow();
        assert_eq!(c.name, "src_bb");
        assert!(c.is_basic_block());
        assert_eq!(c.operands.len(), 1); // One instruction cloned
    }

    #[test]
    fn test_clone_function() {
        let mapper = ValueMapper::new();

        // Create source function
        let src_func = valref(
            Value::new(Type::void())
                .with_subclass(SubclassKind::Function)
                .named("src_func"),
        );

        // Add a basic block
        let bb = valref(
            Value::new(Type::label())
                .with_subclass(SubclassKind::BasicBlock)
                .named("entry"),
        );
        src_func.borrow_mut().operands.push(bb);

        let cloned = mapper.clone_function(&src_func, "cloned_func");
        let c = cloned.borrow();
        assert_eq!(c.name, "cloned_func");
        assert!(c.subclass == SubclassKind::Function);
        assert_eq!(c.operands.len(), 1); // One BB cloned
    }

    #[test]
    fn test_clone_module_empty() {
        let src = llvm_native_core::module::Module::new("src_module");
        let dest = clone_module(&src, "dest_module");
        assert_eq!(dest.name, "dest_module");
        assert!(dest.functions.is_empty());
        assert!(dest.globals.is_empty());
    }

    #[test]
    fn test_clone_module_with_function() {
        let mut src = llvm_native_core::module::Module::new("src_module");
        let func = valref(
            Value::new(Type::void())
                .with_subclass(SubclassKind::Function)
                .named("my_func"),
        );
        src.functions.push(func);

        let dest = clone_module(&src, "dest_module");
        assert_eq!(dest.name, "dest_module");
        assert_eq!(dest.functions.len(), 1);
        assert_eq!(dest.functions[0].borrow().name, "my_func");
    }

    #[test]
    fn test_clone_module_with_globals() {
        let mut src = llvm_native_core::module::Module::new("src_module");
        let global = valref(
            Value::new(Type::i32())
                .with_subclass(SubclassKind::GlobalVariable)
                .named("@g"),
        );
        src.globals.push(global);

        let dest = clone_module(&src, "dest_module");
        assert_eq!(dest.globals.len(), 1);
        assert_eq!(dest.globals[0].borrow().name, "@g");
    }

    #[test]
    fn test_clone_function_into_module() {
        let mut src_module = llvm_native_core::module::Module::new("src");
        let src_func = valref(
            Value::new(Type::void())
                .with_subclass(SubclassKind::Function)
                .named("original"),
        );
        let bb = valref(
            Value::new(Type::label())
                .with_subclass(SubclassKind::BasicBlock)
                .named("entry"),
        );
        src_func.borrow_mut().operands.push(bb);
        src_module.functions.push(src_func);

        let mut dest_module = llvm_native_core::module::Module::new("dest");
        let cloned = clone_function_into(&src_module.functions[0], &mut dest_module, "renamed");

        assert_eq!(cloned.borrow().name, "renamed");
        assert_eq!(dest_module.functions.len(), 1);
    }

    #[test]
    fn test_clone_module_preserves_metadata() {
        let mut src = llvm_native_core::module::Module::new("src");
        src.set_source_filename("test.c");
        src.set_target_triple("x86_64-unknown-linux-gnu");
        src.set_data_layout(
            "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128",
        );

        let dest = clone_module(&src, "dest");
        assert_eq!(dest.source_filename, "test.c");
        assert_eq!(dest.target_triple.unwrap(), "x86_64-unknown-linux-gnu");
        assert!(dest.data_layout.is_some());
    }

    #[test]
    fn test_register_type_multiple() {
        let mut mapper = ValueMapper::new();
        let t1 = Type::i32();
        let t2 = Type::i32();
        let t3 = Type::i64();

        mapper.register_type(t1.id, t3.id);
        mapper.register_type(t2.id, t3.id);

        assert_eq!(mapper.map_type(&t1).id, t3.id);
        assert_eq!(mapper.map_type(&t2).id, t3.id);
    }
}

use std::rc::Rc;