llvm-native-core 0.1.4

//! X86 Rematerialization — identifies values that can be cheaply recomputed
//! instead of spilled to memory during register allocation. Implements remat
//! cost analysis, remat candidate identification, and remat insertion to
//! replace reload instructions with recomputation.
//!
//! Clean-room behavioral reconstruction from:
//! - Intel® 64 and IA-32 Architectures Optimization Reference Manual
//!   (Chapter 3: Instruction Latency and Throughput)
//! - Agner Fog's Instruction Tables (latency/throughput reference)
//! - Register Allocation: Spilling and Rematerialization literature
//! - X86-64 ABI specifications for frame layout
//!
//! Zero LLVM source code consultation. All behavior reconstructed from
//! published specifications and black-box oracle interrogation.

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

// ---------------------------------------------------------------------------
// X86 machine instruction representation (remat-focused)
// ---------------------------------------------------------------------------

/// Opcode identifiers for rematerializable instructions.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum RematOpcode {
    // Mov immediate
    MOV32ri,
    MOV64ri32,
    MOV64ri,
    MOV16ri,
    MOV8ri,
    // Load from RIP-relative
    MOV32rmRIP,
    MOV64rmRIP,
    MOV8rmNoRex,
    // LEA
    LEA64r,
    LEA32r,
    LEA16r,
    // Constant materialization patterns
    XOR32rr, // xor reg, reg -> 0
    XOR64rr,
    XOR8rr,
    XOR16rr,
    SUB32rr, // sub reg, reg -> 0
    SUB64rr,
    // Sign/zero extension
    MOVSX32rm8,
    MOVSX64rm8,
    MOVSX32rm16,
    MOVSX64rm16,
    MOVZX32rm8,
    MOVZX64rm8,
    MOVZX32rm16,
    MOVZX64rm16,
    // Frame/global address LEA patterns
    LEA64rFrame,
    LEA64rGlobal,
    // Constant pool load
    MOV32rmConstPool,
    MOV64rmConstPool,
    // Other
    Unknown,
}

impl RematOpcode {
    pub fn is_rematerializable(&self) -> bool {
        matches!(
            self,
            RematOpcode::MOV32ri
                | RematOpcode::MOV64ri32
                | RematOpcode::MOV64ri
                | RematOpcode::MOV16ri
                | RematOpcode::MOV8ri
                | RematOpcode::MOV32rmRIP
                | RematOpcode::MOV64rmRIP
                | RematOpcode::LEA64r
                | RematOpcode::LEA32r
                | RematOpcode::LEA16r
                | RematOpcode::XOR32rr
                | RematOpcode::XOR64rr
                | RematOpcode::SUB32rr
                | RematOpcode::SUB64rr
                | RematOpcode::LEA64rFrame
                | RematOpcode::LEA64rGlobal
                | RematOpcode::MOV32rmConstPool
                | RematOpcode::MOV64rmConstPool
                | RematOpcode::MOVZX32rm8
                | RematOpcode::MOVZX64rm8
                | RematOpcode::MOVZX32rm16
                | RematOpcode::MOVZX64rm16
        )
    }

    pub fn is_always_rematerializable(&self) -> bool {
        // Instructions that never depend on memory state
        matches!(
            self,
            RematOpcode::MOV32ri
                | RematOpcode::MOV64ri32
                | RematOpcode::MOV64ri
                | RematOpcode::MOV16ri
                | RematOpcode::MOV8ri
                | RematOpcode::XOR32rr
                | RematOpcode::XOR64rr
                | RematOpcode::SUB32rr
                | RematOpcode::SUB64rr
        )
    }

    pub fn is_cheap(&self) -> bool {
        // Low-latency, single-uop instructions
        matches!(
            self,
            RematOpcode::MOV32ri
                | RematOpcode::MOV64ri32
                | RematOpcode::XOR32rr
                | RematOpcode::XOR64rr
                | RematOpcode::SUB32rr
                | RematOpcode::SUB64rr
                | RematOpcode::LEA64r
                | RematOpcode::LEA32r
        )
    }
}

/// Operand types for remat analysis.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum RematOperand {
    /// Immediate constant.
    Immediate(i64),
    /// Register.
    Register(u32),
    /// Memory operand with addressing mode.
    Memory(RematMemoryOperand),
    /// Frame index reference.
    FrameIndex(i64),
    /// Global symbol reference.
    GlobalSymbol(String),
    /// Constant pool entry.
    ConstantPoolEntry(u32, u32), // (index, alignment)
    /// LEA operand: base + index*scale + displacement
    LeaOperand {
        base: Option<u32>,
        index: Option<u32>,
        scale: u8,
        displacement: i32,
    },
    /// No operand.
    None,
}

/// A simplified memory operand for rematerialization analysis.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RematMemoryOperand {
    pub base: Option<u32>,
    pub index: Option<u32>,
    pub scale: u8,
    pub displacement: i32,
    pub is_rip_relative: bool,
    pub segment: Option<u8>,
}

/// A single machine instruction as seen by the rematerializer.
#[derive(Debug, Clone)]
pub struct RematInstruction {
    pub opcode: RematOpcode,
    pub defs: Vec<u32>,          // Destination registers
    pub uses: Vec<RematOperand>, // Source operands
    pub flags_def: bool,         // Whether this instruction defines EFLAGS
    pub flags_use: bool,         // Whether this instruction uses EFLAGS
    pub has_side_effects: bool,  // Memory write, call, etc.
    pub machine_instr_id: u64,   // Unique ID in the function
    pub basic_block_id: u32,     // Which basic block
    pub slot_index: u32,         // Position within basic block
}

impl RematInstruction {
    pub fn new_imm(opcode: RematOpcode, dst: u32, imm: i64, id: u64, bb: u32, slot: u32) -> Self {
        RematInstruction {
            opcode,
            defs: vec![dst],
            uses: vec![RematOperand::Immediate(imm)],
            flags_def: false,
            flags_use: false,
            has_side_effects: false,
            machine_instr_id: id,
            basic_block_id: bb,
            slot_index: slot,
        }
    }

    pub fn new_xor_zero(opcode: RematOpcode, dst: u32, id: u64, bb: u32, slot: u32) -> Self {
        RematInstruction {
            opcode,
            defs: vec![dst],
            uses: vec![RematOperand::Register(dst), RematOperand::Register(dst)],
            flags_def: true,
            flags_use: false,
            has_side_effects: false,
            machine_instr_id: id,
            basic_block_id: bb,
            slot_index: slot,
        }
    }

    pub fn new_lea(
        dst: u32,
        base: Option<u32>,
        index: Option<u32>,
        scale: u8,
        disp: i32,
        id: u64,
        bb: u32,
        slot: u32,
        is_64bit: bool,
    ) -> Self {
        RematInstruction {
            opcode: if is_64bit {
                RematOpcode::LEA64r
            } else {
                RematOpcode::LEA32r
            },
            defs: vec![dst],
            uses: vec![RematOperand::LeaOperand {
                base,
                index,
                scale,
                displacement: disp,
            }],
            flags_def: false,
            flags_use: false,
            has_side_effects: false,
            machine_instr_id: id,
            basic_block_id: bb,
            slot_index: slot,
        }
    }

    pub fn new_rip_load(
        opcode: RematOpcode,
        dst: u32,
        disp: i32,
        id: u64,
        bb: u32,
        slot: u32,
    ) -> Self {
        let mem = RematMemoryOperand {
            base: None,
            index: None,
            scale: 1,
            displacement: disp,
            is_rip_relative: true,
            segment: None,
        };
        RematInstruction {
            opcode,
            defs: vec![dst],
            uses: vec![RematOperand::Memory(mem)],
            flags_def: false,
            flags_use: false,
            has_side_effects: false,
            machine_instr_id: id,
            basic_block_id: bb,
            slot_index: slot,
        }
    }

    pub fn new_frame_lea(
        dst: u32,
        frame_index: i64,
        offset: i32,
        id: u64,
        bb: u32,
        slot: u32,
    ) -> Self {
        RematInstruction {
            opcode: RematOpcode::LEA64rFrame,
            defs: vec![dst],
            uses: vec![
                RematOperand::FrameIndex(frame_index),
                RematOperand::Immediate(offset as i64),
            ],
            flags_def: false,
            flags_use: false,
            has_side_effects: false,
            machine_instr_id: id,
            basic_block_id: bb,
            slot_index: slot,
        }
    }
}

// ---------------------------------------------------------------------------
// Remat candidate and cost
// ---------------------------------------------------------------------------

/// A rematerializable value with its cost information.
#[derive(Debug, Clone)]
pub struct RematCandidate {
    /// The machine instruction ID of the defining instruction.
    pub instr_id: u64,
    /// The virtual register defined.
    pub vreg: u32,
    /// The opcode that can be rematerialized.
    pub opcode: RematOpcode,
    /// The operands needed to recompute.
    pub operands: Vec<RematOperand>,
    /// Cost to rematerialize (in cycles/throughput units).
    pub remat_cost: u32,
    /// Cost to spill+reload (in cycles).
    pub spill_reload_cost: u32,
    /// Whether the defining instruction has side effects.
    pub has_side_effects: bool,
    /// Whether the operands are always available.
    pub operands_always_available: bool,
    /// Whether the value is used in a loop (makes remat more attractive).
    pub in_loop: bool,
    /// Whether the defining instruction clobbers flags.
    pub clobbers_flags: bool,
    /// Estimated latency of rematerialization.
    pub latency: u32,
    /// Whether rematerialization would increase code size.
    pub increases_code_size: bool,
    /// Priority: higher = more favorable for remat.
    pub priority: i32,
}

impl RematCandidate {
    pub fn is_profitable(&self) -> bool {
        self.remat_cost < self.spill_reload_cost && !self.has_side_effects
    }

    pub fn compute_priority(&self) -> i32 {
        let mut p: i32 = 0;
        // Higher spill cost -> higher priority for remat
        p += (self.spill_reload_cost as i32) - (self.remat_cost as i32);
        // Loop usage doubles the benefit
        if self.in_loop {
            p *= 2;
        }
        // Always-available operands are preferred
        if self.operands_always_available {
            p += 10;
        }
        // Side effects disqualify
        if self.has_side_effects {
            p = -1000;
        }
        // Always rematerializable (zero cost to compute)
        if self.opcode.is_always_rematerializable() {
            p += 50;
        }
        p
    }
}

// ---------------------------------------------------------------------------
// Remat cost model
// ---------------------------------------------------------------------------

/// Latency and throughput cost model for X86 instructions.
#[derive(Debug, Clone)]
pub struct RematCostModel {
    /// Base latency for move immediate (MOV ri, imm).
    pub mov_imm_latency: u32,
    /// Base latency for LEA.
    pub lea_latency: u32,
    /// Base latency for RIP-relative load.
    pub rip_load_latency: u32,
    /// Base latency for XOR zeroing.
    pub xor_latency: u32,
    /// Base latency for frame-relative LEA.
    pub frame_lea_latency: u32,
    /// Spill store latency.
    pub spill_store_latency: u32,
    /// Reload latency.
    pub reload_latency: u32,
    /// Whether to prefer size optimization over speed.
    pub opt_for_size: bool,
    /// Microarchitecture: Skylake, Zen3, etc.
    pub uarch: X86UArch,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum X86UArch {
    Generic,
    Skylake,
    IceLake,
    Zen2,
    Zen3,
    Zen4,
    AlderLake,
}

impl Default for RematCostModel {
    fn default() -> Self {
        RematCostModel::skylake()
    }
}

impl RematCostModel {
    pub fn skylake() -> Self {
        RematCostModel {
            mov_imm_latency: 1,
            lea_latency: 1,
            rip_load_latency: 5, // L1 hit ~4-5 cycles
            xor_latency: 1,
            frame_lea_latency: 1,
            spill_store_latency: 2, // Store + eventual L1 writeback
            reload_latency: 5,      // L1 load
            opt_for_size: false,
            uarch: X86UArch::Skylake,
        }
    }

    pub fn zen3() -> Self {
        RematCostModel {
            mov_imm_latency: 1,
            lea_latency: 1,
            rip_load_latency: 4,
            xor_latency: 1,
            frame_lea_latency: 1,
            spill_store_latency: 2,
            reload_latency: 4,
            opt_for_size: false,
            uarch: X86UArch::Zen3,
        }
    }

    pub fn size_optimized() -> Self {
        let mut m = Self::skylake();
        m.opt_for_size = true;
        m
    }

    /// Get the rematerialization cost for a given opcode.
    pub fn remat_cost(&self, opcode: RematOpcode) -> u32 {
        match opcode {
            RematOpcode::MOV32ri | RematOpcode::MOV16ri | RematOpcode::MOV8ri => {
                self.mov_imm_latency
            }
            RematOpcode::MOV64ri32 => self.mov_imm_latency,
            RematOpcode::MOV64ri => {
                // MOVABS is 2 uops, higher latency on most uarchs
                if self.opt_for_size {
                    10 // 10-byte instruction
                } else {
                    2
                }
            }
            RematOpcode::LEA64r | RematOpcode::LEA32r | RematOpcode::LEA16r => self.lea_latency,
            RematOpcode::LEA64rFrame | RematOpcode::LEA64rGlobal => self.frame_lea_latency,
            RematOpcode::MOV32rmRIP | RematOpcode::MOV64rmRIP => self.rip_load_latency,
            RematOpcode::MOV8rmNoRex => self.rip_load_latency,
            RematOpcode::XOR32rr
            | RematOpcode::XOR64rr
            | RematOpcode::XOR8rr
            | RematOpcode::XOR16rr
            | RematOpcode::SUB32rr
            | RematOpcode::SUB64rr => self.xor_latency,
            RematOpcode::MOV32rmConstPool | RematOpcode::MOV64rmConstPool => {
                self.rip_load_latency + 1 // Extra indirection
            }
            RematOpcode::MOVSX32rm8
            | RematOpcode::MOVSX64rm8
            | RematOpcode::MOVSX32rm16
            | RematOpcode::MOVSX64rm16
            | RematOpcode::MOVZX32rm8
            | RematOpcode::MOVZX64rm8
            | RematOpcode::MOVZX32rm16
            | RematOpcode::MOVZX64rm16 => self.rip_load_latency,
            RematOpcode::Unknown => u32::MAX,
        }
    }

    /// Compare rematerialization cost vs spill+reload.
    pub fn is_cheaper_to_remat(&self, opcode: RematOpcode) -> bool {
        let remat = self.remat_cost(opcode);
        let spill_reload = self.spill_store_latency + self.reload_latency;
        remat < spill_reload
    }

    /// Check if rematerialization is smaller in code size.
    pub fn is_smaller_to_remat(&self, opcode: RematOpcode) -> bool {
        // Rematerialization typically emits 2-7 bytes
        // Spill+reload is 2 instructions (store + load), ~6-8 bytes each
        match opcode {
            RematOpcode::XOR32rr | RematOpcode::XOR64rr => true, // 2-3 bytes
            RematOpcode::MOV32ri => true,                        // 5 bytes
            RematOpcode::MOV64ri32 => true,                      // 7 bytes (REX + op + imm32)
            RematOpcode::LEA64r if !self.opt_for_size => true,   // 4-7 bytes
            _ => true, // Generally remat is smaller or equal
        }
    }
}

// ---------------------------------------------------------------------------
// Live range information
// ---------------------------------------------------------------------------

/// Live interval information for a virtual register.
#[derive(Debug, Clone)]
pub struct LiveInterval {
    pub vreg: u32,
    /// List of (start_slot, end_slot) live ranges.
    pub ranges: Vec<(u32, u32)>,
    /// Whether this register crosses a loop backedge.
    pub crosses_loop: bool,
    /// Number of uses.
    pub num_uses: u32,
    /// Number of defs.
    pub num_defs: u32,
    /// Whether this register is used in a cold basic block.
    pub in_cold_region: bool,
    /// The defining instruction ID, if single-def.
    pub defining_instr: Option<u64>,
    /// Weight: higher = more important to keep in register.
    pub weight: f64,
}

impl LiveInterval {
    pub fn is_live_at(&self, slot: u32) -> bool {
        self.ranges
            .iter()
            .any(|&(start, end)| slot >= start && slot <= end)
    }

    pub fn spill_cost(&self, store_cost: u32, load_cost: u32) -> u32 {
        // Approximate: each use requires a reload, each def a store
        let store_count = if self.defining_instr.is_some() {
            1
        } else {
            self.num_defs
        };
        store_count * store_cost + self.num_uses * load_cost
    }
}

// ---------------------------------------------------------------------------
// X86Rematerialization
// ---------------------------------------------------------------------------

/// Configuration for the rematerialization pass.
#[derive(Debug, Clone)]
pub struct RematerializationConfig {
    pub cost_model: RematCostModel,
    pub enable_frame_index_remat: bool,
    pub enable_constant_pool_remat: bool,
    pub enable_global_addr_remat: bool,
    pub enable_rip_relative_remat: bool,
    pub enable_lea_remat: bool,
    pub enable_xor_zero_remat: bool,
    pub enable_immediate_remat: bool,
    pub enable_loop_aware_remat: bool,
    pub enable_cold_path_remat: bool,
    pub max_remat_distance: u32,
    pub min_spill_cost_for_remat: u32,
    pub remat_aggressiveness: u8, // 0=conservative, 1=moderate, 2=aggressive
}

impl Default for RematerializationConfig {
    fn default() -> Self {
        RematerializationConfig {
            cost_model: RematCostModel::default(),
            enable_frame_index_remat: true,
            enable_constant_pool_remat: true,
            enable_global_addr_remat: true,
            enable_rip_relative_remat: true,
            enable_lea_remat: true,
            enable_xor_zero_remat: true,
            enable_immediate_remat: true,
            enable_loop_aware_remat: true,
            enable_cold_path_remat: true,
            max_remat_distance: 100,
            min_spill_cost_for_remat: 5,
            remat_aggressiveness: 1,
        }
    }
}

/// Statistics for the rematerialization pass.
#[derive(Debug, Clone, Default)]
pub struct RematerializationStats {
    pub candidates_identified: u64,
    pub remat_inserted: u64,
    pub spills_avoided: u64,
    pub reloads_avoided: u64,
    pub instructions_added: u64,
    pub instructions_removed: u64,
    pub bytes_saved: u64,
    pub cycles_saved_estimate: u64,
    pub rip_relative_remats: u64,
    pub frame_index_remats: u64,
    pub immediate_remats: u64,
    pub xor_zero_remats: u64,
    pub lea_remats: u64,
    pub const_pool_remats: u64,
    pub global_addr_remats: u64,
    pub loop_remats: u64,
    pub cold_path_remats: u64,
    pub not_profitable: u64,
    pub blocked_by_side_effects: u64,
    pub blocked_by_flags: u64,
    pub blocked_by_operand_unavailable: u64,
}

impl RematerializationStats {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn total_remats(&self) -> u64 {
        self.remat_inserted
    }

    pub fn merge(&mut self, other: &RematerializationStats) {
        self.candidates_identified += other.candidates_identified;
        self.remat_inserted += other.remat_inserted;
        self.spills_avoided += other.spills_avoided;
        self.reloads_avoided += other.reloads_avoided;
        self.instructions_added += other.instructions_added;
        self.instructions_removed += other.instructions_removed;
        self.bytes_saved += other.bytes_saved;
        self.cycles_saved_estimate += other.cycles_saved_estimate;
        self.rip_relative_remats += other.rip_relative_remats;
        self.frame_index_remats += other.frame_index_remats;
        self.immediate_remats += other.immediate_remats;
        self.xor_zero_remats += other.xor_zero_remats;
        self.lea_remats += other.lea_remats;
        self.const_pool_remats += other.const_pool_remats;
        self.global_addr_remats += other.global_addr_remats;
        self.loop_remats += other.loop_remats;
        self.cold_path_remats += other.cold_path_remats;
        self.not_profitable += other.not_profitable;
        self.blocked_by_side_effects += other.blocked_by_side_effects;
        self.blocked_by_flags += other.blocked_by_flags;
        self.blocked_by_operand_unavailable += other.blocked_by_operand_unavailable;
    }
}

/// The X86 Rematerialization pass.
///
/// Identifies values that can be inexpensively recomputed at each use site
/// rather than being spilled to and reloaded from memory. This is especially
/// valuable for constants, global addresses, frame indices, and simple LEA
/// computations.
///
/// Frequently rematerializable X86 instructions:
/// - MOV32ri / MOV64ri32: Move immediate (always remat)
/// - XOR32rr / XOR64rr: Zeroing idiom (always remat)
/// - LEA64r: Load Effective Address (remat if operands available)
/// - MOV32rm with RIP-relative: Load from constant pool (remat if no aliasing)
/// - MOV64rm with RIP-relative: Load 64-bit from constant pool
/// - MOV8rm_NOREX: 8-bit RIP-relative load
pub struct X86Rematerialization {
    pub config: RematerializationConfig,
    pub stats: RematerializationStats,
    /// Map from vreg to its remat candidate info.
    candidates: HashMap<u32, RematCandidate>,
    /// Set of vregs that have been selected for rematerialization.
    selected_remats: HashSet<u32>,
    /// Live intervals indexed by vreg.
    live_intervals: HashMap<u32, LiveInterval>,
    /// Instruction table indexed by machine_instr_id.
    instr_table: HashMap<u64, RematInstruction>,
    /// Loop header block IDs.
    loop_headers: HashSet<u32>,
    /// Cold block IDs.
    cold_blocks: HashSet<u32>,
}

impl X86Rematerialization {
    pub fn new(config: RematerializationConfig) -> Self {
        X86Rematerialization {
            config,
            stats: RematerializationStats::new(),
            candidates: HashMap::new(),
            selected_remats: HashSet::new(),
            live_intervals: HashMap::new(),
            instr_table: HashMap::new(),
            loop_headers: HashSet::new(),
            cold_blocks: HashSet::new(),
        }
    }

    pub fn new_default() -> Self {
        X86Rematerialization::new(RematerializationConfig::default())
    }

    pub fn new_aggressive() -> Self {
        X86Rematerialization::new(RematerializationConfig {
            remat_aggressiveness: 2,
            ..Default::default()
        })
    }

    pub fn new_conservative() -> Self {
        X86Rematerialization::new(RematerializationConfig {
            remat_aggressiveness: 0,
            ..Default::default()
        })
    }

    pub fn new_size_optimized() -> Self {
        X86Rematerialization::new(RematerializationConfig {
            cost_model: RematCostModel::size_optimized(),
            ..Default::default()
        })
    }

    pub fn with_loop_headers(&mut self, headers: HashSet<u32>) -> &mut Self {
        self.loop_headers = headers;
        self
    }

    pub fn with_cold_blocks(&mut self, cold: HashSet<u32>) -> &mut Self {
        self.cold_blocks = cold;
        self
    }

    pub fn take_stats(&mut self) -> RematerializationStats {
        std::mem::take(&mut self.stats)
    }

    // -----------------------------------------------------------------------
    // Candidate identification
    // -----------------------------------------------------------------------

    /// Scan all instructions and identify rematerializable definitions.
    pub fn identify_candidates(&mut self, instrs: &[RematInstruction]) {
        // Build instruction table
        self.instr_table.clear();
        for instr in instrs {
            self.instr_table
                .insert(instr.machine_instr_id, instr.clone());
        }

        // Identify remat candidates
        for instr in instrs {
            if let Some(candidate) = self.analyze_instruction(instr) {
                self.stats.candidates_identified += 1;
                self.candidates.insert(candidate.vreg, candidate);
            }
        }
    }

    /// Analyze a single instruction for rematerialization potential.
    fn analyze_instruction(&self, instr: &RematInstruction) -> Option<RematCandidate> {
        // Must define at least one register
        if instr.defs.is_empty() {
            return None;
        }

        // Must have no side effects
        if instr.has_side_effects {
            return None;
        }

        // Check if opcode is rematerializable
        if !instr.opcode.is_rematerializable() {
            return None;
        }

        // Check feature-specific remat enables
        let enabled = match instr.opcode {
            RematOpcode::LEA64rFrame => self.config.enable_frame_index_remat,
            RematOpcode::LEA64rGlobal => self.config.enable_global_addr_remat,
            RematOpcode::MOV32rmRIP | RematOpcode::MOV64rmRIP | RematOpcode::MOV8rmNoRex => {
                self.config.enable_rip_relative_remat
            }
            RematOpcode::MOV32rmConstPool | RematOpcode::MOV64rmConstPool => {
                self.config.enable_constant_pool_remat
            }
            RematOpcode::LEA64r | RematOpcode::LEA32r | RematOpcode::LEA16r => {
                self.config.enable_lea_remat
            }
            RematOpcode::XOR32rr
            | RematOpcode::XOR64rr
            | RematOpcode::XOR8rr
            | RematOpcode::XOR16rr
            | RematOpcode::SUB32rr
            | RematOpcode::SUB64rr => self.config.enable_xor_zero_remat,
            RematOpcode::MOV32ri
            | RematOpcode::MOV64ri32
            | RematOpcode::MOV64ri
            | RematOpcode::MOV16ri
            | RematOpcode::MOV8ri => self.config.enable_immediate_remat,
            _ => true,
        };

        if !enabled {
            return None;
        }

        let vreg = instr.defs[0];
        let remat_cost = self.config.cost_model.remat_cost(instr.opcode);
        let spill_reload_cost =
            self.config.cost_model.spill_store_latency + self.config.cost_model.reload_latency;

        // Check if operands are always available
        let always_available = instr.opcode.is_always_rematerializable()
            || instr.uses.iter().all(|op| {
                matches!(
                    op,
                    RematOperand::Immediate(_)
                        | RematOperand::FrameIndex(_)
                        | RematOperand::GlobalSymbol(_)
                        | RematOperand::ConstantPoolEntry(_, _)
                        | RematOperand::None
                )
            });

        // Check if we have live interval info
        let in_loop = self.is_in_loop(instr);
        let in_cold = self.is_in_cold_path(instr);

        let candidate = RematCandidate {
            instr_id: instr.machine_instr_id,
            vreg,
            opcode: instr.opcode,
            operands: instr.uses.clone(),
            remat_cost,
            spill_reload_cost,
            has_side_effects: false,
            operands_always_available: always_available,
            in_loop,
            clobbers_flags: instr.flags_def,
            latency: remat_cost,
            increases_code_size: false,
            priority: 0, // Computed below
        };

        let mut c = candidate;
        c.priority = c.compute_priority();

        Some(c)
    }

    fn is_in_loop(&self, instr: &RematInstruction) -> bool {
        self.loop_headers.contains(&instr.basic_block_id)
    }

    fn is_in_cold_path(&self, instr: &RematInstruction) -> bool {
        self.cold_blocks.contains(&instr.basic_block_id)
    }

    // -----------------------------------------------------------------------
    // Remat selection / profitability analysis
    // -----------------------------------------------------------------------

    /// Select which candidates should be rematerialized based on profitability.
    pub fn select_remats(&mut self, intervals: &HashMap<u32, LiveInterval>) {
        self.live_intervals = intervals.clone();

        let vregs: Vec<u32> = intervals.keys().copied().collect();
        for vreg in vregs {
            let li = &intervals[&vreg];
            let candidate = self.candidates.get(&vreg).cloned();
            if let Some(candidate) = candidate {
                if self.should_remat(&candidate, li) {
                    self.selected_remats.insert(vreg);
                } else {
                    self.stats.not_profitable += 1;
                }
            }
        }
    }

    /// Decide whether a candidate should be rematerialized.
    fn should_remat(&mut self, candidate: &RematCandidate, li: &LiveInterval) -> bool {
        // Must be profitable: remat cost < spill+reload cost
        if !candidate.is_profitable() {
            return false;
        }

        // Must not have side effects
        if candidate.has_side_effects {
            self.stats.blocked_by_side_effects += 1;
            return false;
        }

        // In conservative mode, require significant benefit
        if self.config.remat_aggressiveness == 0 {
            let benefit = candidate
                .spill_reload_cost
                .saturating_sub(candidate.remat_cost);
            if benefit < self.config.min_spill_cost_for_remat {
                return false;
            }
        }

        // Check if flags clobbering would be a problem
        if candidate.clobbers_flags {
            // XOR zeroing clobbers flags but is very cheap — still allow
            if !candidate.opcode.is_always_rematerializable() {
                self.stats.blocked_by_flags += 1;
                if self.config.remat_aggressiveness < 2 {
                    return false;
                }
            }
        }

        // For RIP-relative and const pool loads, be more conservative
        match candidate.opcode {
            RematOpcode::MOV32rmRIP
            | RematOpcode::MOV64rmRIP
            | RematOpcode::MOV8rmNoRex
            | RematOpcode::MOV32rmConstPool
            | RematOpcode::MOV64rmConstPool => {
                // These go to memory — only remat if the live range is short
                // or if the value is used in a loop (avoiding spill in hot code)
                if !li.in_cold_region && !self.config.enable_loop_aware_remat {
                    if self.config.remat_aggressiveness < 1 {
                        return false;
                    }
                }
            }
            _ => {}
        }

        // Always rematerializable values (xor zero, mov imm) should always be remat'd
        if candidate.opcode.is_always_rematerializable() {
            return true;
        }

        // LEA remat: check if operands are available
        if matches!(
            candidate.opcode,
            RematOpcode::LEA64r | RematOpcode::LEA32r | RematOpcode::LEA16r
        ) {
            if !candidate.operands_always_available {
                if self.config.remat_aggressiveness < 2 {
                    return false;
                }
            }
        }

        true
    }

    // -----------------------------------------------------------------------
    // Remat insertion (replacing reloads with recomputation)
    // -----------------------------------------------------------------------

    /// Replace a reload instruction with rematerialized computation.
    /// Returns the replacement instruction sequence.
    pub fn replace_reload_with_remat(
        &mut self,
        vreg: u32,
        reload_instr: &RematInstruction,
    ) -> Option<Vec<RematInstruction>> {
        if !self.selected_remats.contains(&vreg) {
            return None;
        }

        let candidate = self.candidates.get(&vreg)?;
        let defining_instr = self.instr_table.get(&candidate.instr_id)?;

        self.stats.remat_inserted += 1;
        self.stats.reloads_avoided += 1;

        // Rematerialize: clone the defining instruction, redirect destination
        // to the reload instruction's destination
        let mut remat_instrs = Vec::new();
        if let Some(&dst) = reload_instr.defs.first() {
            let mut remat = defining_instr.clone();
            remat.defs = vec![dst];
            remat.machine_instr_id = reload_instr.machine_instr_id;
            remat.basic_block_id = reload_instr.basic_block_id;
            remat.slot_index = reload_instr.slot_index;
            remat_instrs.push(remat);

            // Track stats
            match candidate.opcode {
                RematOpcode::MOV32rmRIP | RematOpcode::MOV64rmRIP | RematOpcode::MOV8rmNoRex => {
                    self.stats.rip_relative_remats += 1;
                }
                RematOpcode::LEA64rFrame => {
                    self.stats.frame_index_remats += 1;
                }
                RematOpcode::MOV32ri
                | RematOpcode::MOV64ri32
                | RematOpcode::MOV64ri
                | RematOpcode::MOV16ri
                | RematOpcode::MOV8ri => {
                    self.stats.immediate_remats += 1;
                }
                RematOpcode::XOR32rr
                | RematOpcode::XOR64rr
                | RematOpcode::XOR8rr
                | RematOpcode::XOR16rr => {
                    self.stats.xor_zero_remats += 1;
                }
                RematOpcode::LEA64r | RematOpcode::LEA32r | RematOpcode::LEA16r => {
                    self.stats.lea_remats += 1;
                }
                RematOpcode::MOV32rmConstPool | RematOpcode::MOV64rmConstPool => {
                    self.stats.const_pool_remats += 1;
                }
                RematOpcode::LEA64rGlobal => {
                    self.stats.global_addr_remats += 1;
                }
                _ => {}
            }

            if candidate.in_loop {
                self.stats.loop_remats += 1;
            }
            if self.is_in_cold_path(reload_instr) {
                self.stats.cold_path_remats += 1;
            }
        }

        Some(remat_instrs)
    }

    /// Check if a specific vreg is marked for rematerialization.
    pub fn is_remat_candidate(&self, vreg: u32) -> bool {
        self.selected_remats.contains(&vreg)
    }

    /// Get the remat candidate info for a vreg.
    pub fn get_candidate(&self, vreg: u32) -> Option<&RematCandidate> {
        self.candidates.get(&vreg)
    }

    /// Get all selected remat vregs.
    pub fn get_selected_remats(&self) -> &HashSet<u32> {
        &self.selected_remats
    }

    /// Get estimated cycles saved by rematerialization.
    pub fn estimate_cycles_saved(&self) -> u64 {
        let mut total: u64 = 0;
        for vreg in &self.selected_remats {
            if let (Some(cand), Some(li)) =
                (self.candidates.get(vreg), self.live_intervals.get(vreg))
            {
                let benefit = cand.spill_reload_cost.saturating_sub(cand.remat_cost) as u64;
                let use_count = li.num_uses as u64;
                total += benefit * use_count;
            }
        }
        total
    }
}

// ---------------------------------------------------------------------------
// Remat analysis for specific X86 patterns
// ---------------------------------------------------------------------------

impl X86Rematerialization {
    /// Check if an instruction is a zeroing idiom (xor reg, reg or sub reg, reg).
    pub fn is_zeroing_idiom(opcode: RematOpcode) -> bool {
        matches!(
            opcode,
            RematOpcode::XOR32rr
                | RematOpcode::XOR64rr
                | RematOpcode::XOR8rr
                | RematOpcode::XOR16rr
                | RematOpcode::SUB32rr
                | RematOpcode::SUB64rr
        )
    }

    /// Check if an instruction is an immediate move (always remat, zero latency cost).
    pub fn is_immediate_move(opcode: RematOpcode) -> bool {
        matches!(
            opcode,
            RematOpcode::MOV32ri
                | RematOpcode::MOV64ri32
                | RematOpcode::MOV64ri
                | RematOpcode::MOV16ri
                | RematOpcode::MOV8ri
        )
    }

    /// Check if an instruction is an LEA that computes a frame address.
    pub fn is_frame_lea(opcode: RematOpcode) -> bool {
        matches!(opcode, RematOpcode::LEA64rFrame)
    }

    /// Check if an instruction loads from a RIP-relative address (constant pool, global).
    pub fn is_rip_relative_load(opcode: RematOpcode) -> bool {
        matches!(
            opcode,
            RematOpcode::MOV32rmRIP | RematOpcode::MOV64rmRIP | RematOpcode::MOV8rmNoRex
        )
    }

    /// Classify the rematerialization type for an opcode.
    pub fn classify_remat_type(opcode: RematOpcode) -> RematType {
        if Self::is_zeroing_idiom(opcode) {
            RematType::ZeroIdiom
        } else if Self::is_immediate_move(opcode) {
            RematType::Immediate
        } else if Self::is_frame_lea(opcode) {
            RematType::FrameAddress
        } else if Self::is_rip_relative_load(opcode) {
            RematType::RipRelative
        } else if matches!(opcode, RematOpcode::LEA64rGlobal) {
            RematType::GlobalAddress
        } else if matches!(
            opcode,
            RematOpcode::MOV32rmConstPool | RematOpcode::MOV64rmConstPool
        ) {
            RematType::ConstantPool
        } else if matches!(
            opcode,
            RematOpcode::LEA64r | RematOpcode::LEA32r | RematOpcode::LEA16r
        ) {
            RematType::Lea
        } else {
            RematType::Other
        }
    }
}

/// Rematerialization type classification.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RematType {
    ZeroIdiom,
    Immediate,
    FrameAddress,
    GlobalAddress,
    ConstantPool,
    RipRelative,
    Lea,
    Other,
}

impl RematType {
    pub fn typical_cost(&self) -> u32 {
        match self {
            RematType::ZeroIdiom => 0,     // xor reg,reg: recognized as zeroing idiom
            RematType::Immediate => 1,     // mov reg, imm: 1 cycle
            RematType::FrameAddress => 1,  // lea reg, [rbp+disp]
            RematType::GlobalAddress => 5, // rip-relative load
            RematType::ConstantPool => 5,  // rip-relative load
            RematType::RipRelative => 5,   // load from memory
            RematType::Lea => 1,           // lea: 1 cycle
            RematType::Other => 5,         // conservative
        }
    }

    pub fn always_available(&self) -> bool {
        matches!(
            self,
            RematType::ZeroIdiom | RematType::Immediate | RematType::FrameAddress
        )
    }
}

// ---------------------------------------------------------------------------
// Rematerialization during register allocation
// ---------------------------------------------------------------------------

/// Context for remat decision during live range splitting / allocation.
#[derive(Debug, Clone)]
pub struct RematAllocationContext {
    /// The vreg being allocated.
    pub vreg: u32,
    /// Available physical registers at this point.
    pub available_regs: Vec<u32>,
    /// Current spill slot assignment (if any).
    pub spill_slot: Option<u32>,
    /// Remat candidate info.
    pub candidate: Option<RematCandidate>,
    /// Distance to nearest use.
    pub distance_to_use: u32,
    /// Distance to nearest def.
    pub distance_to_def: u32,
    /// Whether we're in a hot loop.
    pub in_hot_loop: bool,
    /// Whether a reload would be needed.
    pub needs_reload: bool,
    /// Whether a spill would be needed.
    pub needs_spill: bool,
}

impl X86Rematerialization {
    /// Make a rematerialization decision during register allocation.
    ///
    /// Returns true if rematerialization is the better choice over
    /// spilling + reloading.
    pub fn decide_remat_vs_spill(&self, ctx: &RematAllocationContext) -> RematDecision {
        let candidate = match &ctx.candidate {
            Some(c) => c,
            None => return RematDecision::MustSpill,
        };

        // Can't remat if we need to spill (value still live)
        if ctx.needs_spill && !candidate.operands_always_available {
            return RematDecision::MustSpill;
        }

        // Always remat zeroing idioms and immediate moves
        if candidate.opcode.is_always_rematerializable() {
            return RematDecision::Rematerialize;
        }

        // In hot loops, prefer remat to avoid memory traffic
        if ctx.in_hot_loop && self.config.enable_loop_aware_remat {
            let remat_cost = candidate.remat_cost;
            let spill_reload = candidate.spill_reload_cost;
            if remat_cost * 10 <= spill_reload * 10 {
                // Actually, in loop body, the formula is:
                // spill outside loop + reload inside loop vs remat inside loop
                return RematDecision::Rematerialize;
            }
        }

        // If distance to use is very short, remat may be cheaper
        if ctx.distance_to_use <= 2 && candidate.remat_cost <= 1 {
            return RematDecision::Rematerialize;
        }

        // Default: use cost model
        if candidate.is_profitable() {
            RematDecision::Rematerialize
        } else {
            RematDecision::MustSpill
        }
    }
}

/// Decision result for remat vs spill.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RematDecision {
    Rematerialize,
    MustSpill,
    CanSpill,
    NoAction,
}

// ---------------------------------------------------------------------------
// Remat pass: main entry point
// ---------------------------------------------------------------------------

/// Result of running the rematerialization pass on a function.
#[derive(Debug, Clone)]
pub struct RematerializationResult {
    /// Modified instruction sequence with remat replacements.
    pub instructions: Vec<RematInstruction>,
    /// Spill code that was removed.
    pub spills_removed: u64,
    /// Reload code that was replaced by remat.
    pub reloads_rematd: u64,
    /// Stats from the pass.
    pub stats: RematerializationStats,
}

impl X86Rematerialization {
    /// Run the full rematerialization pass on a function's instructions.
    ///
    /// 1. Identify rematerializable definitions
    /// 2. Select profitable candidates
    /// 3. Replace reload instructions with rematerialized computations
    pub fn run(
        &mut self,
        instrs: &[RematInstruction],
        intervals: &HashMap<u32, LiveInterval>,
    ) -> RematerializationResult {
        self.candidates.clear();
        self.selected_remats.clear();
        self.instr_table.clear();

        // Step 1: Identify candidates
        self.identify_candidates(instrs);

        // Step 2: Select remats
        self.select_remats(intervals);

        // Step 3: Replace reloads with rematerialized instructions
        let mut new_instrs: Vec<RematInstruction> = Vec::new();
        let mut reloads_rematd: u64 = 0;
        let mut spills_removed: u64 = 0;

        for instr in instrs {
            // Check if this is a reload (load from spill slot)
            if self.is_reload_instruction(instr) {
                // If any def vreg is a selected remat, replace the reload
                let mut replaced = false;
                for &vreg in &instr.defs {
                    if self.is_remat_candidate(vreg) {
                        if let Some(remat_instrs) = self.replace_reload_with_remat(vreg, instr) {
                            new_instrs.extend(remat_instrs);
                            reloads_rematd += 1;
                            replaced = true;
                            break;
                        }
                    }
                }
                if !replaced {
                    new_instrs.push(instr.clone());
                }
            } else if self.is_spill_instruction(instr) {
                // Check if the spilled vreg is a remat candidate
                let mut skip_spill = false;
                for &vreg in &instr
                    .uses
                    .iter()
                    .filter_map(|op| {
                        if let RematOperand::Register(r) = op {
                            Some(*r)
                        } else {
                            None
                        }
                    })
                    .collect::<Vec<_>>()
                {
                    if self.is_remat_candidate(vreg) && instr.defs.is_empty() {
                        // Pure spill: we can skip it if value will be remat'd at uses
                        skip_spill = true;
                        spills_removed += 1;
                        break;
                    }
                }
                if !skip_spill {
                    new_instrs.push(instr.clone());
                }
            } else {
                new_instrs.push(instr.clone());
            }
        }

        self.stats.spills_avoided += spills_removed;
        self.stats.reloads_avoided += reloads_rematd;

        RematerializationResult {
            instructions: new_instrs,
            spills_removed,
            reloads_rematd,
            stats: self.stats.clone(),
        }
    }

    /// Heuristic: detect if an instruction is a reload from a spill slot.
    fn is_reload_instruction(&self, instr: &RematInstruction) -> bool {
        // A reload is typically: mov reg, [rbp/rsp + offset] where offset >= 0
        // and the instruction has a memory operand as source
        if instr.defs.is_empty() {
            return false;
        }
        instr
            .uses
            .iter()
            .any(|op| matches!(op, RematOperand::Memory(_)))
            && !instr.has_side_effects
            && instr.flags_def == false
    }

    /// Heuristic: detect if an instruction is a spill to a spill slot.
    fn is_spill_instruction(&self, instr: &RematInstruction) -> bool {
        // A spill is typically: mov [rbp/rsp + offset], reg
        // The instruction stores a register to a memory operand
        instr.defs.is_empty()
            && instr
                .uses
                .iter()
                .any(|op| matches!(op, RematOperand::Memory(_)))
            && !instr.has_side_effects
    }
}

// ---------------------------------------------------------------------------
// Additional analysis: Remat viability for complex expressions
// ---------------------------------------------------------------------------

/// Description of a complex value that might be rematerializable.
#[derive(Debug, Clone)]
pub struct ComplexRematExpr {
    /// The root vreg of the expression.
    pub root_vreg: u32,
    /// The chain of instructions that compute this value.
    pub instr_chain: Vec<RematInstruction>,
    /// Total cost to rematerialize the entire chain.
    pub total_remat_cost: u32,
    /// Whether all operands in the chain are leaves (always available).
    pub all_leaves_available: bool,
    /// Whether the chain is short enough to be profitable.
    pub chain_too_long: bool,
    /// Maximum chain length for profitability.
    pub max_chain_length: usize,
}

impl X86Rematerialization {
    /// Check if a value defined by a short chain of cheap instructions
    /// (e.g., LEA followed by ADD) is worth rematerializing as a unit.
    pub fn analyze_complex_remat(
        &self,
        root_vreg: u32,
        max_depth: usize,
    ) -> Option<ComplexRematExpr> {
        let candidate = self.candidates.get(&root_vreg)?;
        let mut chain = Vec::new();
        let mut total_cost: u32 = 0;

        // Walk backwards through the def-use chain
        let mut current_id = candidate.instr_id;
        let mut depth = 0;

        while depth < max_depth {
            let instr = self.instr_table.get(&current_id)?;
            chain.push(instr.clone());
            total_cost += self.config.cost_model.remat_cost(instr.opcode);

            // Check if operands are all leaves
            let mut all_leaves = true;
            for op in &instr.uses {
                match op {
                    RematOperand::Register(r) => {
                        // Check if this register is itself rematerializable
                        if !self.candidates.contains_key(r) {
                            all_leaves = false;
                            break;
                        }
                        // Follow into the defining instruction
                        if let Some(def_cand) = self.candidates.get(r) {
                            current_id = def_cand.instr_id;
                            depth += 1;
                        } else {
                            all_leaves = false;
                            break;
                        }
                    }
                    _ => {} // Immediates, frame indices, etc. are leaves
                }
            }

            if !all_leaves {
                break;
            }

            if depth > max_depth {
                break;
            }
        }

        let chain_too_long = chain.len() > max_depth;
        let all_leaves = chain.iter().all(|i| {
            i.uses
                .iter()
                .all(|op| !matches!(op, RematOperand::Register(_)) || op == &RematOperand::None)
        });

        Some(ComplexRematExpr {
            root_vreg,
            instr_chain: chain,
            total_remat_cost: total_cost,
            all_leaves_available: all_leaves,
            chain_too_long,
            max_chain_length: max_depth,
        })
    }
}

// ---------------------------------------------------------------------------
// Remat heuristics for specific scenarios
// ---------------------------------------------------------------------------

impl X86Rematerialization {
    /// Determine if a constant should be rematerialized instead of spilled.
    /// Constants are among the best candidates since MOV reg, imm is very cheap.
    pub fn should_remat_constant(&self, imm: i64, li: &LiveInterval) -> bool {
        if !self.config.enable_immediate_remat {
            return false;
        }

        // Small constants (fits in 32-bit signed): mov r32, imm32 = 5 bytes, 1 cycle
        let fits_in_32 = imm >= -(1i64 << 31) && imm <= (1i64 << 31) - 1;

        let spill_cost = li.spill_cost(
            self.config.cost_model.spill_store_latency,
            self.config.cost_model.reload_latency,
        );

        // Even for constants that need movabs (10 bytes, 2 cycles), it's often
        // cheaper than spill+reload (store + load ≈ 7+ bytes, 7+ cycles)
        if fits_in_32 {
            // MOV ri32: 1 cycle, 5 bytes
            // Spill+reload per use: 7+ cycles, 12+ bytes
            spill_cost > 1
        } else {
            // MOVABS: 2 cycles, 10 bytes
            spill_cost > 2
        }
    }

    /// Determine if a zero value (xor reg, reg) should be rematerialized.
    /// Zeroing idioms are recognized by the CPU as dependency-breaking and
    /// are essentially free in modern out-of-order cores.
    pub fn should_remat_zero(&self, li: &LiveInterval) -> bool {
        if !self.config.enable_xor_zero_remat {
            return false;
        }
        // Zeroing is always cheaper than spilling
        true
    }

    /// Determine if a frame index computation should be rematerialized.
    /// Frame addresses (LEA reg, [RBP+disp]) are cheap LEA operations.
    pub fn should_remat_frame_index(&self, li: &LiveInterval) -> bool {
        if !self.config.enable_frame_index_remat {
            return false;
        }
        let spill_cost = li.spill_cost(
            self.config.cost_model.spill_store_latency,
            self.config.cost_model.reload_latency,
        );
        // LEA is 1 cycle — remat is cheaper unless live range has 1 use
        // (in which case spilling might be fine too).
        // But in practice, always remat frame indices.
        li.num_uses > 1 || self.config.remat_aggressiveness >= 1
    }

    /// Determine if a global address load should be rematerialized.
    /// Global address loads go through memory (RIP-relative), so they're
    /// more expensive than spills in some cases.
    pub fn should_remat_global_addr(&self, li: &LiveInterval) -> bool {
        if !self.config.enable_global_addr_remat {
            return false;
        }

        let remat_cost = self.config.cost_model.rip_load_latency;
        let spill_reload =
            self.config.cost_model.spill_store_latency + self.config.cost_model.reload_latency;

        // Only profitable if remat is cheaper than spill+reload AND
        // the value is used multiple times (otherwise 1 spill+reload might be OK)
        if li.num_uses <= 1 {
            return false;
        }

        remat_cost < spill_reload
    }
}

// ---------------------------------------------------------------------------
// Factory functions
// ---------------------------------------------------------------------------

pub fn make_x86_rematerialization() -> X86Rematerialization {
    X86Rematerialization::new_default()
}

pub fn make_x86_rematerialization_aggressive() -> X86Rematerialization {
    X86Rematerialization::new_aggressive()
}

pub fn make_x86_rematerialization_conservative() -> X86Rematerialization {
    X86Rematerialization::new_conservative()
}

pub fn make_x86_rematerialization_size_opt() -> X86Rematerialization {
    X86Rematerialization::new_size_optimized()
}

pub fn make_rematerialization_with_config(config: RematerializationConfig) -> X86Rematerialization {
    X86Rematerialization::new(config)
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    fn make_remat() -> X86Rematerialization {
        X86Rematerialization::new_default()
    }

    fn make_test_instr_imm(id: u64, dst: u32, imm: i64) -> RematInstruction {
        RematInstruction::new_imm(RematOpcode::MOV32ri, dst, imm, id, 0, 0)
    }

    fn make_test_instr_xor(id: u64, dst: u32) -> RematInstruction {
        RematInstruction::new_xor_zero(RematOpcode::XOR32rr, dst, id, 0, 0)
    }

    fn make_test_instr_lea(id: u64, dst: u32, base: u32, disp: i32) -> RematInstruction {
        RematInstruction::new_lea(dst, Some(base), None, 1, disp, id, 0, 0, true)
    }

    fn make_test_interval(vreg: u32, num_uses: u32, num_defs: u32) -> LiveInterval {
        LiveInterval {
            vreg,
            ranges: vec![(0, 100)],
            crosses_loop: false,
            num_uses,
            num_defs,
            in_cold_region: false,
            defining_instr: Some(vreg as u64),
            weight: 1.0,
        }
    }

    // -- Constructor tests --

    #[test]
    fn test_constructor_default() {
        let remat = make_remat();
        assert_eq!(remat.config.remat_aggressiveness, 1);
    }

    #[test]
    fn test_constructor_aggressive() {
        let remat = X86Rematerialization::new_aggressive();
        assert_eq!(remat.config.remat_aggressiveness, 2);
    }

    #[test]
    fn test_constructor_conservative() {
        let remat = X86Rematerialization::new_conservative();
        assert_eq!(remat.config.remat_aggressiveness, 0);
    }

    // -- Candidate identification tests --

    #[test]
    fn test_identify_immediate_candidate() {
        let mut remat = make_remat();
        let instrs = vec![make_test_instr_imm(1, 10, 42)];
        remat.identify_candidates(&instrs);
        assert!(remat.candidates.contains_key(&10));
        let c = remat.candidates.get(&10).unwrap();
        assert_eq!(c.opcode, RematOpcode::MOV32ri);
        assert!(c.operands_always_available);
    }

    #[test]
    fn test_identify_xor_zero_candidate() {
        let mut remat = make_remat();
        let instrs = vec![make_test_instr_xor(1, 20)];
        remat.identify_candidates(&instrs);
        assert!(remat.candidates.contains_key(&20));
        let c = remat.candidates.get(&20).unwrap();
        assert_eq!(c.opcode, RematOpcode::XOR32rr);
        assert!(c.is_profitable());
    }

    #[test]
    fn test_identify_lea_candidate() {
        let mut remat = make_remat();
        let instrs = vec![make_test_instr_lea(1, 30, 5, -16)];
        remat.identify_candidates(&instrs);
        assert!(remat.candidates.contains_key(&30));
    }

    #[test]
    fn test_identify_rip_relative_candidate() {
        let mut remat = make_remat();
        let instr = RematInstruction::new_rip_load(RematOpcode::MOV32rmRIP, 40, 0, 1, 0, 0);
        remat.identify_candidates(&[instr]);
        assert!(remat.candidates.contains_key(&40));
    }

    #[test]
    fn test_skip_side_effect_instruction() {
        let mut remat = make_remat();
        let mut instr = RematInstruction::new_imm(RematOpcode::MOV32ri, 50, 42, 1, 0, 0);
        instr.has_side_effects = true;
        remat.identify_candidates(&[instr]);
        assert!(!remat.candidates.contains_key(&50));
    }

    #[test]
    fn test_skip_non_remat_opcode() {
        let mut remat = make_remat();
        let mut instr = RematInstruction::new_imm(RematOpcode::Unknown, 60, 0, 1, 0, 0);
        instr.opcode = RematOpcode::Unknown;
        remat.identify_candidates(&[instr]);
        assert!(!remat.candidates.contains_key(&60));
    }

    // -- Selection tests --

    #[test]
    fn test_select_immediate_remat() {
        let mut remat = make_remat();
        let instrs = vec![make_test_instr_imm(1, 10, 42)];
        remat.identify_candidates(&instrs);
        let intervals: HashMap<u32, LiveInterval> =
            [(10, make_test_interval(10, 5, 1))].into_iter().collect();
        remat.select_remats(&intervals);
        assert!(remat.selected_remats.contains(&10));
    }

    #[test]
    fn test_select_xor_zero_remat() {
        let mut remat = make_remat();
        let instrs = vec![make_test_instr_xor(1, 20)];
        remat.identify_candidates(&instrs);
        let intervals: HashMap<u32, LiveInterval> =
            [(20, make_test_interval(20, 3, 1))].into_iter().collect();
        remat.select_remats(&intervals);
        assert!(remat.selected_remats.contains(&20));
    }

    #[test]
    fn test_select_lea_remat() {
        let mut remat = make_remat();
        let instrs = vec![make_test_instr_lea(1, 30, 5, 0)];
        remat.identify_candidates(&instrs);
        let intervals: HashMap<u32, LiveInterval> =
            [(30, make_test_interval(30, 4, 1))].into_iter().collect();
        remat.select_remats(&intervals);
        assert!(remat.selected_remats.contains(&30));
    }

    #[test]
    fn test_not_profitable_remat() {
        let mut remat = make_remat();
        // RIP-relative load with single use — not profitable
        let instr = RematInstruction::new_rip_load(RematOpcode::MOV32rmRIP, 40, 0, 1, 0, 0);
        remat.identify_candidates(&[instr]);
        let intervals: HashMap<u32, LiveInterval> =
            [(40, make_test_interval(40, 1, 1))].into_iter().collect();
        remat.select_remats(&intervals);
        // Single use RIP-relative may not be selected
        // (depends on cost model)
    }

    // -- Remat replacement tests --

    #[test]
    fn test_replace_reload_with_remat_immediate() {
        let mut remat = make_remat();
        let defining = make_test_instr_imm(1, 10, 42);
        remat.identify_candidates(&[defining.clone()]);

        let intervals: HashMap<u32, LiveInterval> =
            [(10, make_test_interval(10, 3, 1))].into_iter().collect();
        remat.select_remats(&intervals);

        // Simulate a reload instruction
        let reload = RematInstruction::new_imm(RematOpcode::MOV32rmRIP, 15, 0, 2, 0, 1);
        let replacement = remat.replace_reload_with_remat(10, &reload);
        assert!(replacement.is_some());
        let instrs = replacement.unwrap();
        assert_eq!(instrs.len(), 1);
        assert_eq!(instrs[0].defs, vec![15]); // Destination redirected
    }

    #[test]
    fn test_replace_reload_not_selected() {
        let mut remat = make_remat();
        let defining = make_test_instr_imm(1, 10, 42);
        remat.identify_candidates(&[defining]);

        // Don't select — leave selected_remats empty
        let reload = RematInstruction::new_imm(RematOpcode::MOV32rmRIP, 15, 0, 2, 0, 1);
        let replacement = remat.replace_reload_with_remat(10, &reload);
        assert!(replacement.is_none());
    }

    // -- Zeroing idiom tests --

    #[test]
    fn test_is_zeroing_idiom() {
        assert!(X86Rematerialization::is_zeroing_idiom(RematOpcode::XOR32rr));
        assert!(X86Rematerialization::is_zeroing_idiom(RematOpcode::XOR64rr));
        assert!(!X86Rematerialization::is_zeroing_idiom(
            RematOpcode::MOV32ri
        ));
    }

    #[test]
    fn test_is_immediate_move() {
        assert!(X86Rematerialization::is_immediate_move(
            RematOpcode::MOV32ri
        ));
        assert!(X86Rematerialization::is_immediate_move(
            RematOpcode::MOV64ri32
        ));
        assert!(!X86Rematerialization::is_immediate_move(
            RematOpcode::LEA64r
        ));
    }

    // -- Classification tests --

    #[test]
    fn test_classify_remat_type() {
        assert_eq!(
            X86Rematerialization::classify_remat_type(RematOpcode::XOR32rr),
            RematType::ZeroIdiom
        );
        assert_eq!(
            X86Rematerialization::classify_remat_type(RematOpcode::MOV32ri),
            RematType::Immediate
        );
        assert_eq!(
            X86Rematerialization::classify_remat_type(RematOpcode::LEA64rFrame),
            RematType::FrameAddress
        );
        assert_eq!(
            X86Rematerialization::classify_remat_type(RematOpcode::LEA64rGlobal),
            RematType::GlobalAddress
        );
        assert_eq!(
            X86Rematerialization::classify_remat_type(RematOpcode::MOV32rmRIP),
            RematType::RipRelative
        );
        assert_eq!(
            X86Rematerialization::classify_remat_type(RematOpcode::LEA64r),
            RematType::Lea
        );
    }

    // -- RematType tests --

    #[test]
    fn test_remat_type_typical_cost() {
        assert_eq!(RematType::ZeroIdiom.typical_cost(), 0);
        assert_eq!(RematType::Immediate.typical_cost(), 1);
        assert_eq!(RematType::Lea.typical_cost(), 1);
        assert_eq!(RematType::RipRelative.typical_cost(), 5);
    }

    #[test]
    fn test_remat_type_always_available() {
        assert!(RematType::ZeroIdiom.always_available());
        assert!(RematType::Immediate.always_available());
        assert!(RematType::FrameAddress.always_available());
        assert!(!RematType::RipRelative.always_available());
        assert!(!RematType::Lea.always_available());
    }

    // -- Cost model tests --

    #[test]
    fn test_cost_model_default() {
        let model = RematCostModel::default();
        assert!(model.remat_cost(RematOpcode::MOV32ri) < model.remat_cost(RematOpcode::MOV32rmRIP));
    }

    #[test]
    fn test_cost_model_zen3() {
        let model = RematCostModel::zen3();
        assert_eq!(model.lea_latency, 1);
    }

    #[test]
    fn test_is_cheaper_to_remat() {
        let model = RematCostModel::default();
        assert!(model.is_cheaper_to_remat(RematOpcode::MOV32ri));
        assert!(model.is_cheaper_to_remat(RematOpcode::XOR32rr));
        assert!(model.is_cheaper_to_remat(RematOpcode::LEA64r));
    }

    // -- Decision tests --

    #[test]
    fn test_decide_remat_immediate() {
        let remat = make_remat();
        let ctx = RematAllocationContext {
            vreg: 10,
            available_regs: vec![0, 1, 2],
            spill_slot: None,
            candidate: Some(RematCandidate {
                instr_id: 1,
                vreg: 10,
                opcode: RematOpcode::MOV32ri,
                operands: vec![RematOperand::Immediate(42)],
                remat_cost: 1,
                spill_reload_cost: 7,
                has_side_effects: false,
                operands_always_available: true,
                in_loop: false,
                clobbers_flags: false,
                latency: 1,
                increases_code_size: false,
                priority: 50,
            }),
            distance_to_use: 5,
            distance_to_def: 10,
            in_hot_loop: false,
            needs_reload: true,
            needs_spill: false,
        };
        assert_eq!(
            remat.decide_remat_vs_spill(&ctx),
            RematDecision::Rematerialize
        );
    }

    #[test]
    fn test_decide_must_spill_no_candidate() {
        let remat = make_remat();
        let ctx = RematAllocationContext {
            vreg: 10,
            available_regs: vec![],
            spill_slot: None,
            candidate: None,
            distance_to_use: 5,
            distance_to_def: 10,
            in_hot_loop: false,
            needs_reload: true,
            needs_spill: true,
        };
        assert_eq!(remat.decide_remat_vs_spill(&ctx), RematDecision::MustSpill);
    }

    // -- Run pass tests --

    #[test]
    fn test_run_full_pass() {
        let mut remat = make_remat();
        let instrs = vec![
            make_test_instr_imm(1, 10, 42),
            // Simulated reload of vreg 10
            RematInstruction {
                opcode: RematOpcode::MOV32rmRIP,
                defs: vec![10],
                uses: vec![RematOperand::Memory(RematMemoryOperand {
                    base: None,
                    index: None,
                    scale: 1,
                    displacement: 0,
                    is_rip_relative: false,
                    segment: None,
                })],
                flags_def: false,
                flags_use: false,
                has_side_effects: false,
                machine_instr_id: 2,
                basic_block_id: 0,
                slot_index: 1,
            },
        ];
        let intervals: HashMap<u32, LiveInterval> =
            [(10, make_test_interval(10, 2, 1))].into_iter().collect();

        let result = remat.run(&instrs, &intervals);
        assert!(result.reloads_rematd > 0 || result.instructions.len() <= instrs.len());
    }

    // -- LiveInterval tests --

    #[test]
    fn test_live_interval_is_live_at() {
        let li = LiveInterval {
            vreg: 1,
            ranges: vec![(0, 10), (20, 30)],
            crosses_loop: false,
            num_uses: 1,
            num_defs: 1,
            in_cold_region: false,
            defining_instr: Some(0),
            weight: 1.0,
        };
        assert!(li.is_live_at(5));
        assert!(li.is_live_at(25));
        assert!(!li.is_live_at(15));
    }

    #[test]
    fn test_live_interval_spill_cost() {
        let li = make_test_interval(1, 5, 1);
        let cost = li.spill_cost(2, 5);
        // store_cost (1 def * 2) + load_cost (5 uses * 5) = 2 + 25 = 27
        assert_eq!(cost, 27);
    }

    // -- RematCandidate tests --

    #[test]
    fn test_candidate_is_profitable() {
        let c = RematCandidate {
            instr_id: 1,
            vreg: 10,
            opcode: RematOpcode::MOV32ri,
            operands: vec![],
            remat_cost: 1,
            spill_reload_cost: 7,
            has_side_effects: false,
            operands_always_available: true,
            in_loop: false,
            clobbers_flags: false,
            latency: 1,
            increases_code_size: false,
            priority: 0,
        };
        assert!(c.is_profitable());
    }

    #[test]
    fn test_candidate_not_profitable() {
        let c = RematCandidate {
            instr_id: 1,
            vreg: 10,
            opcode: RematOpcode::MOV32rmRIP,
            operands: vec![],
            remat_cost: 10,
            spill_reload_cost: 7,
            has_side_effects: false,
            operands_always_available: false,
            in_loop: false,
            clobbers_flags: false,
            latency: 10,
            increases_code_size: false,
            priority: 0,
        };
        assert!(!c.is_profitable());
    }

    #[test]
    fn test_candidate_compute_priority() {
        let c = RematCandidate {
            instr_id: 1,
            vreg: 10,
            opcode: RematOpcode::MOV32ri,
            operands: vec![],
            remat_cost: 1,
            spill_reload_cost: 10,
            has_side_effects: false,
            operands_always_available: true,
            in_loop: true,
            clobbers_flags: false,
            latency: 1,
            increases_code_size: false,
            priority: 0,
        };
        let p = c.compute_priority();
        assert!(p > 0);
    }

    // -- Stats tests --

    #[test]
    fn test_stats_new_zero() {
        let stats = RematerializationStats::new();
        assert_eq!(stats.total_remats(), 0);
    }

    #[test]
    fn test_stats_merge() {
        let mut a = RematerializationStats::new();
        a.remat_inserted = 10;
        a.spills_avoided = 5;

        let mut b = RematerializationStats::new();
        b.remat_inserted = 3;
        b.spills_avoided = 2;

        a.merge(&b);
        assert_eq!(a.remat_inserted, 13);
        assert_eq!(a.spills_avoided, 7);
    }

    #[test]
    fn test_take_stats() {
        let mut remat = make_remat();
        let instrs = vec![make_test_instr_imm(1, 10, 42)];
        remat.identify_candidates(&instrs);
        let stats = remat.take_stats();
        assert!(stats.candidates_identified > 0);
        // After take, stats are reset
        let stats2 = remat.take_stats();
        assert_eq!(stats2.candidates_identified, 0);
    }

    // -- Heuristic tests --

    #[test]
    fn test_should_remat_constant() {
        let remat = make_remat();
        let li = make_test_interval(1, 5, 1);
        assert!(remat.should_remat_constant(42, &li));
        assert!(remat.should_remat_constant(-128, &li));
    }

    #[test]
    fn test_should_remat_zero() {
        let remat = make_remat();
        let li = make_test_interval(1, 1, 1);
        assert!(remat.should_remat_zero(&li));
    }

    #[test]
    fn test_should_remat_frame_index() {
        let remat = make_remat();
        let li = make_test_interval(1, 3, 1);
        assert!(remat.should_remat_frame_index(&li));
    }

    // -- Opcode tests --

    #[test]
    fn test_remat_opcode_is_rematerializable() {
        assert!(RematOpcode::MOV32ri.is_rematerializable());
        assert!(RematOpcode::MOV64ri32.is_rematerializable());
        assert!(RematOpcode::LEA64r.is_rematerializable());
        assert!(RematOpcode::XOR32rr.is_rematerializable());
        assert!(!RematOpcode::Unknown.is_rematerializable());
    }

    #[test]
    fn test_remat_opcode_is_always_rematerializable() {
        assert!(RematOpcode::MOV32ri.is_always_rematerializable());
        assert!(RematOpcode::XOR32rr.is_always_rematerializable());
        assert!(!RematOpcode::LEA64r.is_always_rematerializable());
        assert!(!RematOpcode::MOV32rmRIP.is_always_rematerializable());
    }

    #[test]
    fn test_remat_opcode_is_cheap() {
        assert!(RematOpcode::MOV32ri.is_cheap());
        assert!(RematOpcode::XOR32rr.is_cheap());
        assert!(RematOpcode::LEA64r.is_cheap());
        assert!(!RematOpcode::MOV32rmRIP.is_cheap());
    }

    // -- Frame LEA tests --

    #[test]
    fn test_frame_lea_candidate() {
        let mut remat = make_remat();
        let instr = RematInstruction::new_frame_lea(1, 0, -8, 1, 0, 0);
        remat.identify_candidates(&[instr]);
        assert!(remat.candidates.contains_key(&1));
        let c = remat.candidates.get(&1).unwrap();
        assert_eq!(c.opcode, RematOpcode::LEA64rFrame);
    }

    // -- Feature enable/disable tests --

    #[test]
    fn test_disable_immediate_remat() {
        let config = RematerializationConfig {
            enable_immediate_remat: false,
            ..Default::default()
        };
        let mut remat = X86Rematerialization::new(config);
        let instrs = vec![make_test_instr_imm(1, 10, 42)];
        remat.identify_candidates(&instrs);
        assert!(!remat.candidates.contains_key(&10));
    }

    #[test]
    fn test_disable_xor_zero_remat() {
        let config = RematerializationConfig {
            enable_xor_zero_remat: false,
            ..Default::default()
        };
        let mut remat = X86Rematerialization::new(config);
        let instrs = vec![make_test_instr_xor(1, 20)];
        remat.identify_candidates(&instrs);
        assert!(!remat.candidates.contains_key(&20));
    }

    #[test]
    fn test_disable_lea_remat() {
        let config = RematerializationConfig {
            enable_lea_remat: false,
            ..Default::default()
        };
        let mut remat = X86Rematerialization::new(config);
        let instrs = vec![make_test_instr_lea(1, 30, 5, 0)];
        remat.identify_candidates(&instrs);
        assert!(!remat.candidates.contains_key(&30));
    }

    // -- Estimate cycles saved --

    #[test]
    fn test_estimate_cycles_saved() {
        let mut remat = make_remat();
        let instrs = vec![make_test_instr_imm(1, 10, 42)];
        remat.identify_candidates(&instrs);
        let intervals: HashMap<u32, LiveInterval> =
            [(10, make_test_interval(10, 5, 1))].into_iter().collect();
        remat.select_remats(&intervals);
        let saved = remat.estimate_cycles_saved();
        assert!(saved > 0);
    }

    // -- Edge cases --

    #[test]
    fn test_empty_instructions() {
        let mut remat = make_remat();
        remat.identify_candidates(&[]);
        assert!(remat.candidates.is_empty());
    }

    #[test]
    fn test_multiple_defs_no_remat() {
        let mut remat = make_remat();
        let mut instr = RematInstruction::new_imm(RematOpcode::MOV32ri, 10, 42, 1, 0, 0);
        instr.defs = vec![]; // No defs
        remat.identify_candidates(&[instr]);
        assert!(remat.candidates.is_empty());
    }

    #[test]
    fn test_instr_with_flags_def() {
        let mut remat = make_remat();
        let mut instr = make_test_instr_xor(1, 20);
        instr.flags_def = true; // XOR always defines flags
        remat.identify_candidates(&[instr]);
        assert!(remat.candidates.contains_key(&20));
    }

    // -- Factory tests --

    #[test]
    fn test_factory_default() {
        let remat = make_x86_rematerialization();
        assert_eq!(remat.config.remat_aggressiveness, 1);
    }

    #[test]
    fn test_factory_aggressive() {
        let remat = make_x86_rematerialization_aggressive();
        assert_eq!(remat.config.remat_aggressiveness, 2);
    }

    #[test]
    fn test_factory_conservative() {
        let remat = make_x86_rematerialization_conservative();
        assert_eq!(remat.config.remat_aggressiveness, 0);
    }

    #[test]
    fn test_factory_size_opt() {
        let remat = make_x86_rematerialization_size_opt();
        assert!(remat.config.cost_model.opt_for_size);
    }

    #[test]
    fn test_factory_with_config() {
        let config = RematerializationConfig {
            remat_aggressiveness: 0,
            ..Default::default()
        };
        let remat = make_rematerialization_with_config(config);
        assert_eq!(remat.config.remat_aggressiveness, 0);
    }

    // -- Complex remat analysis tests --

    #[test]
    fn test_analyze_complex_remat_empty() {
        let remat = make_remat();
        let result = remat.analyze_complex_remat(999, 3);
        assert!(result.is_none());
    }

    // -- Comprehensive tests --

    #[test]
    fn test_all_remat_opcodes_classified() {
        let opcodes = [
            RematOpcode::MOV32ri,
            RematOpcode::MOV64ri32,
            RematOpcode::MOV64ri,
            RematOpcode::LEA64r,
            RematOpcode::LEA32r,
            RematOpcode::XOR32rr,
            RematOpcode::XOR64rr,
            RematOpcode::MOV32rmRIP,
            RematOpcode::MOV64rmRIP,
            RematOpcode::LEA64rFrame,
            RematOpcode::LEA64rGlobal,
            RematOpcode::MOV32rmConstPool,
            RematOpcode::MOV64rmConstPool,
        ];
        for op in &opcodes {
            let remat_type = X86Rematerialization::classify_remat_type(*op);
            assert_ne!(remat_type, RematType::Other);
        }
    }
}