llvm-native-core 0.1.13

//! X86 Addressing Mode Legalizer — transforms generic addressing modes into
//! valid X86-compatible forms by folding offsets, decomposing complex modes,
//! promoting narrow indices, handling RIP-relative, TLS, segment overrides,
//! and combining GEP operations into single LEA instructions.
//!
//! Clean-room behavioral reconstruction from:
//! - Intel® 64 and IA-32 Architectures Software Developer's Manual
//!   (Volume 2A, Chapter 2: Instruction Format — ModR/M and SIB Bytes)
//! - AMD64 Architecture Programmer's Manual (Volume 1: Addressing Modes)
//! - System V ABI: AMD64 Supplement (TLS and PIC addressing)
//! - X86-64 ELF ABI Specifications (R_X86_64_* relocations)
//!
//! Zero LLVM source code consultation. All behavior reconstructed from
//! published specifications and black-box oracle interrogation.

use std::collections::HashMap;

// ---------------------------------------------------------------------------
// Core addressing mode representation
// ---------------------------------------------------------------------------

/// X86 addressing mode width.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum AddrWidth {
    Addr16,
    Addr32,
    Addr64,
}

impl AddrWidth {
    pub fn default_data_size(&self) -> u32 {
        match self {
            AddrWidth::Addr16 => 2,
            AddrWidth::Addr32 => 4,
            AddrWidth::Addr64 => 8,
        }
    }

    pub fn max_displacement_bits(&self) -> u32 {
        match self {
            AddrWidth::Addr16 => 16,
            AddrWidth::Addr32 => 32,
            AddrWidth::Addr64 => 32,
        }
    }

    pub fn supports_rip_relative(&self) -> bool {
        matches!(self, AddrWidth::Addr64)
    }

    pub fn is_64bit(&self) -> bool {
        matches!(self, AddrWidth::Addr64)
    }
}

/// A generic address component — can be a register, frame index, or symbol.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum AddrComponent {
    /// A physical or virtual GPR (0-15 in 64-bit, 0-7 in 32-bit).
    Register(u32),
    /// A frame index that will be resolved to RBP/RSP + displacement.
    FrameIndex(i64),
    /// A global symbol or label reference.
    Symbol(String),
    /// A constant pool entry index.
    ConstantPool(u32),
    /// A TLS symbol reference.
    TlsSymbol(String),
    /// The GOT base (__GLOBAL_OFFSET_TABLE__).
    GotBase,
    /// No component — unset.
    None,
}

impl Default for AddrComponent {
    fn default() -> Self {
        AddrComponent::None
    }
}

impl AddrComponent {
    pub fn is_none(&self) -> bool {
        matches!(self, AddrComponent::None)
    }

    pub fn is_register(&self) -> bool {
        matches!(self, AddrComponent::Register(_))
    }

    pub fn is_symbol(&self) -> bool {
        matches!(self, AddrComponent::Symbol(_) | AddrComponent::TlsSymbol(_))
    }

    pub fn is_frame_index(&self) -> bool {
        matches!(self, AddrComponent::FrameIndex(_))
    }
}

/// A scale factor for indexed addressing: 1, 2, 4, or 8.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ScaleFactor {
    Times1 = 1,
    Times2 = 2,
    Times4 = 4,
    Times8 = 8,
}

impl ScaleFactor {
    pub fn from_u8(v: u8) -> Option<ScaleFactor> {
        match v {
            1 => Some(ScaleFactor::Times1),
            2 => Some(ScaleFactor::Times2),
            4 => Some(ScaleFactor::Times4),
            8 => Some(ScaleFactor::Times8),
            _ => None,
        }
    }

    pub fn log2(self) -> u8 {
        match self {
            ScaleFactor::Times1 => 0,
            ScaleFactor::Times2 => 1,
            ScaleFactor::Times4 => 2,
            ScaleFactor::Times8 => 3,
        }
    }

    pub fn as_u32(self) -> u32 {
        self as u32
    }
}

impl Default for ScaleFactor {
    fn default() -> Self {
        ScaleFactor::Times1
    }
}

/// X86 segment override.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum SegmentOverride {
    None,
    ES,
    CS,
    SS,
    DS,
    FS,
    GS,
}

impl SegmentOverride {
    pub fn prefix_byte(&self) -> Option<u8> {
        match self {
            SegmentOverride::None => None,
            SegmentOverride::ES => Some(0x26),
            SegmentOverride::CS => Some(0x2E),
            SegmentOverride::SS => Some(0x36),
            SegmentOverride::DS => Some(0x3E),
            SegmentOverride::FS => Some(0x64),
            SegmentOverride::GS => Some(0x65),
        }
    }

    pub fn default_for_register(reg: u32, is_64bit: bool) -> SegmentOverride {
        if !is_64bit {
            match reg {
                4 | 5 => SegmentOverride::SS, // RSP, RBP -> SS
                _ => SegmentOverride::DS,
            }
        } else {
            SegmentOverride::None // 64-bit mode defaults to DS (no prefix needed)
        }
    }

    pub fn is_stack_segment(&self) -> bool {
        matches!(self, SegmentOverride::SS)
    }
}

/// TLS model variants.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TlsModel {
    GeneralDynamic,
    LocalDynamic,
    InitialExec,
    LocalExec,
}

impl Default for TlsModel {
    fn default() -> Self {
        TlsModel::GeneralDynamic
    }
}

/// Code model for symbol addressing.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum CodeModel {
    Small,
    Kernel,
    Medium,
    Large,
}

impl Default for CodeModel {
    fn default() -> Self {
        CodeModel::Small
    }
}

/// PIC mode.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum PicMode {
    NotPIC,
    PIC,
    PICNoGOT,
}

impl Default for PicMode {
    fn default() -> Self {
        PicMode::NotPIC
    }
}

/// A complete X86 addressing mode before final encoding.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct X86GenericAddrMode {
    pub width: AddrWidth,
    pub base: AddrComponent,
    pub index: AddrComponent,
    pub scale: ScaleFactor,
    pub displacement: i64,
    pub segment: SegmentOverride,
    pub is_rip_relative: bool,
    pub is_frame_relative: bool,
    pub is_got_relative: bool,
    pub is_tls: bool,
    pub tls_model: Option<TlsModel>,
    pub pic_mode: PicMode,
    pub code_model: CodeModel,
}

impl Default for X86GenericAddrMode {
    fn default() -> Self {
        X86GenericAddrMode {
            width: AddrWidth::Addr64,
            base: AddrComponent::None,
            index: AddrComponent::None,
            scale: ScaleFactor::Times1,
            displacement: 0,
            segment: SegmentOverride::None,
            is_rip_relative: false,
            is_frame_relative: false,
            is_got_relative: false,
            is_tls: false,
            tls_model: None,
            pic_mode: PicMode::NotPIC,
            code_model: CodeModel::Small,
        }
    }
}

impl X86GenericAddrMode {
    pub fn simple_64(base: u32) -> Self {
        X86GenericAddrMode {
            width: AddrWidth::Addr64,
            base: AddrComponent::Register(base),
            ..Default::default()
        }
    }

    pub fn base_disp_64(base: u32, disp: i64) -> Self {
        X86GenericAddrMode {
            width: AddrWidth::Addr64,
            base: AddrComponent::Register(base),
            displacement: disp,
            ..Default::default()
        }
    }

    pub fn full_64(base: u32, index: u32, scale: ScaleFactor, disp: i64) -> Self {
        X86GenericAddrMode {
            width: AddrWidth::Addr64,
            base: AddrComponent::Register(base),
            index: AddrComponent::Register(index),
            scale,
            displacement: disp,
            ..Default::default()
        }
    }

    pub fn rip_relative(sym: String, disp: i64) -> Self {
        X86GenericAddrMode {
            width: AddrWidth::Addr64,
            base: AddrComponent::Symbol(sym),
            displacement: disp,
            is_rip_relative: true,
            ..Default::default()
        }
    }

    pub fn hash_key(&self) -> u64 {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};
        let mut h = DefaultHasher::new();
        self.width.hash(&mut h);
        self.base.hash(&mut h);
        self.index.hash(&mut h);
        self.scale.hash(&mut h);
        self.displacement.hash(&mut h);
        h.finish()
    }
}

// ---------------------------------------------------------------------------
// Addressing mode classification
// ---------------------------------------------------------------------------

/// Classification of an addressing mode for legalization purposes.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AddrModeClass {
    /// Simple: base + displacement (no index, no complex folding needed).
    Simple,
    /// Indexed: base + index*scale + displacement.
    Indexed,
    /// Complex: base + index*scale + displacement with exotic components.
    Complex,
    /// RIP-relative: RIP + displacement to symbol.
    RipRelative,
    /// Absolute: direct memory address.
    Absolute,
    /// TLS: thread-local storage access.
    Tls,
    /// Frame-relative: frame pointer + displacement.
    FrameRelative,
    /// Invalid / needs decomposition.
    NeedsDecomposition,
}

/// The result of legalizing a single addressing mode.
#[derive(Debug, Clone)]
pub struct LegalizedAddrMode {
    /// The legalized mode — None if the mode was invalid.
    pub mode: Option<X86GenericAddrMode>,
    /// Decomposed sub-modes (e.g., when a complex mode is split into LEA + simple access).
    pub decomposed: Vec<X86GenericAddrMode>,
    /// Whether an explicit LEA instruction is needed to compute the address.
    pub needs_lea: bool,
    /// The LEA form to emit (base + index + disp) for computing the address.
    pub lea_mode: Option<X86GenericAddrMode>,
    /// Whether a segment override prefix must be emitted.
    pub needs_segment_prefix: bool,
    /// The segment override byte if needed.
    pub segment_override_byte: Option<u8>,
    /// Any REX prefix bits needed.
    pub rex_bits: u8,
    /// Whether RIP-relative encoding was used.
    pub used_rip_relative: bool,
    /// Whether a new base register was allocated for the address computation.
    pub allocated_base: bool,
    /// The scratch register used for intermediate address computation, if any.
    pub scratch_reg: Option<u32>,
    /// Additional prefix bytes needed (e.g., address size override).
    pub prefix_bytes: Vec<u8>,
    /// Displacement size in bytes (0, 1, 4).
    pub disp_size: u8,
    /// Whether a SIB byte is needed.
    pub needs_sib: bool,
    /// The SIB byte value if needed.
    pub sib_byte: Option<u8>,
}

impl LegalizedAddrMode {
    pub fn valid(mode: X86GenericAddrMode) -> Self {
        LegalizedAddrMode {
            mode: Some(mode),
            decomposed: Vec::new(),
            needs_lea: false,
            lea_mode: None,
            needs_segment_prefix: false,
            segment_override_byte: None,
            rex_bits: 0,
            used_rip_relative: false,
            allocated_base: false,
            scratch_reg: None,
            prefix_bytes: Vec::new(),
            disp_size: 0,
            needs_sib: false,
            sib_byte: None,
        }
    }

    pub fn invalid() -> Self {
        LegalizedAddrMode {
            mode: None,
            decomposed: Vec::new(),
            needs_lea: false,
            lea_mode: None,
            needs_segment_prefix: false,
            segment_override_byte: None,
            rex_bits: 0,
            used_rip_relative: false,
            allocated_base: false,
            scratch_reg: None,
            prefix_bytes: Vec::new(),
            disp_size: 0,
            needs_sib: false,
            sib_byte: None,
        }
    }
}

// ---------------------------------------------------------------------------
// Statistics
// ---------------------------------------------------------------------------

/// Statistics tracking for address mode legalization.
#[derive(Debug, Clone, Default)]
pub struct X86AddrModeLegalizerStats {
    pub total_addrs_processed: u64,
    pub simple_passthrough: u64,
    pub rip_relative_converted: u64,
    pub absolute_converted: u64,
    pub tls_legalized: u64,
    pub frame_indices_folded: u64,
    pub displacements_folded: u64,
    pub scales_canonicalized: u64,
    pub indices_promoted: u64,
    pub complex_decomposed: u64,
    pub lea_formed: u64,
    pub gep_combined: u64,
    pub segment_overrides_applied: u64,
    pub cache_hits: u64,
    pub cache_misses: u64,
    pub errors: u64,
}

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

    pub fn total(&self) -> u64 {
        self.total_addrs_processed
    }

    pub fn merge(&mut self, other: &X86AddrModeLegalizerStats) {
        self.total_addrs_processed += other.total_addrs_processed;
        self.simple_passthrough += other.simple_passthrough;
        self.rip_relative_converted += other.rip_relative_converted;
        self.absolute_converted += other.absolute_converted;
        self.tls_legalized += other.tls_legalized;
        self.frame_indices_folded += other.frame_indices_folded;
        self.displacements_folded += other.displacements_folded;
        self.scales_canonicalized += other.scales_canonicalized;
        self.indices_promoted += other.indices_promoted;
        self.complex_decomposed += other.complex_decomposed;
        self.lea_formed += other.lea_formed;
        self.gep_combined += other.gep_combined;
        self.segment_overrides_applied += other.segment_overrides_applied;
        self.cache_hits += other.cache_hits;
        self.cache_misses += other.cache_misses;
        self.errors += other.errors;
    }
}

// ---------------------------------------------------------------------------
// X86AddrModeLegalizer
// ---------------------------------------------------------------------------

/// Configuration for the addressing mode legalizer.
#[derive(Debug, Clone)]
pub struct X86AddrModeLegalizerConfig {
    /// Target address width.
    pub addr_width: AddrWidth,
    /// Whether to use RIP-relative addressing for symbols (64-bit).
    pub use_rip_relative: bool,
    /// PIC mode.
    pub pic_mode: PicMode,
    /// Code model.
    pub code_model: CodeModel,
    /// Default TLS model.
    pub tls_model: TlsModel,
    /// Whether to fold displacements into the base offset when possible.
    pub fold_displacements: bool,
    /// Whether to decompose complex modes (base+index+disp with exotic combos).
    pub decompose_complex: bool,
    /// Whether to promote 8/16-bit indices to 32-bit.
    pub promote_narrow_indices: bool,
    /// Whether to combine two GEP operations into a single LEA.
    pub combine_geps: bool,
    /// Whether to fold frame indices into base+disp.
    pub fold_frame_indices: bool,
    /// Whether to handle segment overrides.
    pub handle_segment_overrides: bool,
    /// Whether to convert absolute addresses to RIP-relative in PIC mode.
    pub convert_absolute_to_rip: bool,
    /// Optimization level (affects aggressiveness of legalization).
    pub opt_level: u8,
}

impl Default for X86AddrModeLegalizerConfig {
    fn default() -> Self {
        X86AddrModeLegalizerConfig {
            addr_width: AddrWidth::Addr64,
            use_rip_relative: true,
            pic_mode: PicMode::NotPIC,
            code_model: CodeModel::Small,
            tls_model: TlsModel::GeneralDynamic,
            fold_displacements: true,
            decompose_complex: true,
            promote_narrow_indices: true,
            combine_geps: true,
            fold_frame_indices: true,
            handle_segment_overrides: true,
            convert_absolute_to_rip: false,
            opt_level: 2,
        }
    }
}

/// The X86 Addressing Mode Legalizer.
///
/// Transforms high-level address representations into valid X86 addressing modes.
/// Handles:
/// - Folding constant offsets into displacement
/// - Decomposing complex scale+index+base patterns
/// - Promoting 8/16-bit indices to 32-bit
/// - Converting global addresses to RIP-relative
/// - Handling TLS address modes
/// - Folding frame indices into base+disp
/// - Combining two GEPs into a single LEA
/// - Legalizing segment overrides
pub struct X86AddrModeLegalizer {
    pub config: X86AddrModeLegalizerConfig,
    pub stats: X86AddrModeLegalizerStats,
    /// Cache of legalized address modes keyed by hash.
    cache: HashMap<u64, LegalizedAddrMode>,
    /// Frame offset resolution table (frame index -> (base reg, offset)).
    frame_offsets: HashMap<i64, (u32, i64)>,
    /// Known GOT base register if PIC.
    got_base_reg: Option<u32>,
    /// Scratch register for intermediate address computation.
    scratch_reg: Option<u32>,
}

impl X86AddrModeLegalizer {
    pub fn new(config: X86AddrModeLegalizerConfig) -> Self {
        X86AddrModeLegalizer {
            config,
            stats: X86AddrModeLegalizerStats::new(),
            cache: HashMap::new(),
            frame_offsets: HashMap::new(),
            got_base_reg: None,
            scratch_reg: None,
        }
    }

    pub fn new_x86_64() -> Self {
        X86AddrModeLegalizer::new(X86AddrModeLegalizerConfig {
            addr_width: AddrWidth::Addr64,
            use_rip_relative: true,
            ..Default::default()
        })
    }

    pub fn new_x86_64_pic() -> Self {
        X86AddrModeLegalizer::new(X86AddrModeLegalizerConfig {
            addr_width: AddrWidth::Addr64,
            use_rip_relative: false,
            pic_mode: PicMode::PIC,
            convert_absolute_to_rip: true,
            ..Default::default()
        })
    }

    pub fn new_x86_32() -> Self {
        X86AddrModeLegalizer::new(X86AddrModeLegalizerConfig {
            addr_width: AddrWidth::Addr32,
            use_rip_relative: false,
            ..Default::default()
        })
    }

    pub fn new_x86_16() -> Self {
        X86AddrModeLegalizer::new(X86AddrModeLegalizerConfig {
            addr_width: AddrWidth::Addr16,
            use_rip_relative: false,
            decompose_complex: true,
            ..Default::default()
        })
    }

    // -- Configuration helpers --

    pub fn with_got_base(&mut self, reg: u32) -> &mut Self {
        self.got_base_reg = Some(reg);
        self
    }

    pub fn with_scratch_reg(&mut self, reg: u32) -> &mut Self {
        self.scratch_reg = Some(reg);
        self
    }

    pub fn with_frame_offset(&mut self, fi: i64, base_reg: u32, offset: i64) -> &mut Self {
        self.frame_offsets.insert(fi, (base_reg, offset));
        self
    }

    pub fn with_frame_offsets(&mut self, offsets: HashMap<i64, (u32, i64)>) -> &mut Self {
        self.frame_offsets = offsets;
        self
    }

    pub fn clear_cache(&mut self) {
        self.cache.clear();
    }

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

    // -----------------------------------------------------------------------
    // Main legalization entry point
    // -----------------------------------------------------------------------

    /// Legalize an addressing mode. Returns the legalized result.
    pub fn legalize(&mut self, addr: &X86GenericAddrMode) -> LegalizedAddrMode {
        self.stats.total_addrs_processed += 1;

        // Check cache
        let key = addr.hash_key();
        if let Some(cached) = self.cache.get(&key) {
            self.stats.cache_hits += 1;
            return cached.clone();
        }
        self.stats.cache_misses += 1;

        let result = self.legalize_impl(addr);
        self.cache.insert(key, result.clone());
        result
    }

    fn legalize_impl(&mut self, addr: &X86GenericAddrMode) -> LegalizedAddrMode {
        // Step 0: If invalid width, return error.
        if !self.is_valid_width(addr) {
            self.stats.errors += 1;
            return LegalizedAddrMode::invalid();
        }

        // Step 1: Handle TLS address modes.
        if addr.is_tls {
            return self.legalize_tls(addr);
        }

        // Step 2: Handle RIP-relative or convert to RIP-relative.
        if addr.is_rip_relative || self.should_be_rip_relative(addr) {
            return self.legalize_rip_relative(addr);
        }

        // Step 3: Handle frame-relative addressing.
        if addr.is_frame_relative || addr.base.is_frame_index() {
            return self.legalize_frame_relative(addr);
        }

        // Step 4: Handle absolute addresses.
        if self.is_absolute(addr) {
            return self.legalize_absolute(addr);
        }

        // Step 5: Handle complex modes (decompose if needed).
        if self.needs_decomposition(addr) {
            return self.legalize_complex(addr);
        }

        // Step 6: Canonicalize and fold.
        let mut result = self.canonicalize_mode(addr);
        result = self.fold_displacements_into(result);
        result = self.promote_narrow_indices_if_needed(result);
        result = self.apply_segment_override(result);
        let mut result = LegalizedAddrMode::valid(result);
        result = self.compute_rex_sib_disp(result);

        result
    }

    // -----------------------------------------------------------------------
    // Width validation
    // -----------------------------------------------------------------------

    fn is_valid_width(&self, addr: &X86GenericAddrMode) -> bool {
        addr.width == self.config.addr_width
    }

    // -----------------------------------------------------------------------
    // RIP-relative legalization
    // -----------------------------------------------------------------------

    fn should_be_rip_relative(&self, addr: &X86GenericAddrMode) -> bool {
        if !self.config.addr_width.supports_rip_relative() {
            return false;
        }
        if !self.config.use_rip_relative && self.config.pic_mode == PicMode::NotPIC {
            return false;
        }
        // Check if base is a symbol that can be RIP-relative
        addr.base.is_symbol() && !addr.is_got_relative && self.config.code_model != CodeModel::Large
    }

    fn legalize_rip_relative(&mut self, addr: &X86GenericAddrMode) -> LegalizedAddrMode {
        self.stats.rip_relative_converted += 1;

        let mut mode = addr.clone();
        mode.is_rip_relative = true;
        mode.width = AddrWidth::Addr64;

        // Clear base/index since RIP-relative uses none
        if mode.base.is_register() {
            // Complex case: symbol + base — need LEA
            let mut result = LegalizedAddrMode::invalid();
            result.needs_lea = true;
            result.lea_mode = Some(mode.clone());
            result.used_rip_relative = true;
            // The actual memory access will use the LEA result as base
            let mut mem_mode = X86GenericAddrMode::default();
            mem_mode.width = AddrWidth::Addr64;
            mem_mode.base = AddrComponent::Register(self.scratch_reg.unwrap_or(0));
            result.mode = Some(mem_mode);
            result.needs_sib = false;
            result.disp_size = 0;
            return result;
        }

        mode.base = AddrComponent::None;
        mode.index = AddrComponent::None;
        mode.scale = ScaleFactor::Times1;
        mode.segment = SegmentOverride::None;

        let mut result = LegalizedAddrMode::valid(mode);
        result.used_rip_relative = true;
        result.disp_size = 4; // RIP-relative always has 32-bit displacement
        result.needs_sib = false;
        result
    }

    // -----------------------------------------------------------------------
    // TLS legalization
    // -----------------------------------------------------------------------

    fn legalize_tls(&mut self, addr: &X86GenericAddrMode) -> LegalizedAddrMode {
        self.stats.tls_legalized += 1;

        let tls_model = addr.tls_model.unwrap_or(self.config.tls_model);

        match tls_model {
            TlsModel::LocalExec => self.legalize_tls_local_exec(addr),
            TlsModel::InitialExec => self.legalize_tls_initial_exec(addr),
            TlsModel::LocalDynamic => self.legalize_tls_local_dynamic(addr),
            TlsModel::GeneralDynamic => self.legalize_tls_general_dynamic(addr),
        }
    }

    fn legalize_tls_local_exec(&mut self, addr: &X86GenericAddrMode) -> LegalizedAddrMode {
        // Local Exec: access TLS variable directly via FS segment + offset
        let mut mode = addr.clone();
        mode.segment = if self.config.addr_width.is_64bit() {
            SegmentOverride::FS
        } else {
            SegmentOverride::GS
        };
        mode.is_tls = false;
        mode.is_rip_relative = false;
        mode.base = AddrComponent::None;
        mode.index = AddrComponent::None;

        let segment_byte = mode.segment.prefix_byte();
        let mut result = LegalizedAddrMode::valid(mode);
        result.needs_segment_prefix = true;
        result.segment_override_byte = segment_byte;
        result.disp_size = 4; // TLS offset is 32-bit
        result
    }

    fn legalize_tls_initial_exec(&mut self, addr: &X86GenericAddrMode) -> LegalizedAddrMode {
        // Initial Exec: load from GOT[TPOFF] then use address
        // Requires: movq foo@GOTTPOFF(%rip), %reg; then use reg as base
        let mut result = LegalizedAddrMode::invalid();
        result.needs_lea = true;
        result.needs_segment_prefix = true;
        result.segment_override_byte = SegmentOverride::FS.prefix_byte();

        // The GOT slot access
        let mut got_mode = X86GenericAddrMode::default();
        got_mode.width = AddrWidth::Addr64;
        got_mode.is_rip_relative = true;
        got_mode.displacement = addr.displacement;
        if let AddrComponent::TlsSymbol(ref sym) = addr.base {
            got_mode.base = AddrComponent::Symbol(format!("{}@GOTTPOFF", sym));
        }
        result.decomposed.push(got_mode);

        // Final address: reg (from GOT load) + displacement
        let mut final_mode = X86GenericAddrMode::default();
        final_mode.width = AddrWidth::Addr64;
        final_mode.base = AddrComponent::Register(self.scratch_reg.unwrap_or(0));
        final_mode.segment = SegmentOverride::FS;
        result.mode = Some(final_mode);
        result
    }

    fn legalize_tls_local_dynamic(&mut self, addr: &X86GenericAddrMode) -> LegalizedAddrMode {
        // Local Dynamic: call __tls_get_addr with leaq x@TLSLD(%rip), %rdi
        let mut result = LegalizedAddrMode::invalid();
        result.needs_lea = true;

        let mut lea_mode = X86GenericAddrMode::default();
        lea_mode.width = AddrWidth::Addr64;
        lea_mode.is_rip_relative = true;
        if let AddrComponent::TlsSymbol(ref sym) = addr.base {
            lea_mode.base = AddrComponent::Symbol(format!("{}@TLSLD", sym));
        }

        result.lea_mode = Some(lea_mode);
        // Then access via FS segment
        // result.needs_segment_prefix = true;
        // result.segment_override_byte = SegmentOverride::FS.prefix_byte();
        result
    }

    fn legalize_tls_general_dynamic(&mut self, addr: &X86GenericAddrMode) -> LegalizedAddrMode {
        // General Dynamic: call __tls_get_addr with leaq x@TLSGD(%rip), %rdi
        let mut result = LegalizedAddrMode::invalid();
        result.needs_lea = true;
        result.needs_segment_prefix = false;

        let mut lea_mode = X86GenericAddrMode::default();
        lea_mode.width = AddrWidth::Addr64;
        lea_mode.is_rip_relative = true;
        if let AddrComponent::TlsSymbol(ref sym) = addr.base {
            lea_mode.base = AddrComponent::Symbol(format!("{}@TLSGD", sym));
        }

        result.lea_mode = Some(lea_mode);
        result
    }

    // -----------------------------------------------------------------------
    // Frame-relative legalization
    // -----------------------------------------------------------------------

    fn legalize_frame_relative(&mut self, addr: &X86GenericAddrMode) -> LegalizedAddrMode {
        self.stats.frame_indices_folded += 1;

        let fi = match addr.base {
            AddrComponent::FrameIndex(fi) => fi,
            _ => {
                // frame-relative with explicit base
                if let Some(&(base_reg, offset)) = self.frame_offsets.get(&0) {
                    let mut mode = addr.clone();
                    mode.base = AddrComponent::Register(base_reg);
                    mode.displacement += offset;
                    mode.is_frame_relative = false;
                    return LegalizedAddrMode::valid(mode);
                }
                // Default: use RBP/RSP
                let mut mode = addr.clone();
                mode.base = AddrComponent::Register(if self.config.addr_width.is_64bit() {
                    5 // RBP
                } else {
                    5 // EBP
                });
                mode.is_frame_relative = false;
                return LegalizedAddrMode::valid(mode);
            }
        };

        if let Some(&(base_reg, offset)) = self.frame_offsets.get(&fi) {
            let mut mode = addr.clone();
            mode.base = AddrComponent::Register(base_reg);
            mode.displacement += offset;
            mode.is_frame_relative = false;
            let disp_size = Self::compute_disp_size(mode.displacement);
            let mut result = LegalizedAddrMode::valid(mode);
            result.disp_size = disp_size;
            result
        } else {
            // Unknown frame index — treat as simple base+disp
            let base_reg = if self.config.addr_width.is_64bit() {
                5
            } else {
                5
            };
            let mut mode = addr.clone();
            mode.base = AddrComponent::Register(base_reg);
            mode.is_frame_relative = false;
            let disp_size = Self::compute_disp_size(mode.displacement);
            let mut result = LegalizedAddrMode::valid(mode);
            result.disp_size = disp_size;
            result
        }
    }

    // -----------------------------------------------------------------------
    // Absolute address legalization
    // -----------------------------------------------------------------------

    fn is_absolute(&self, addr: &X86GenericAddrMode) -> bool {
        addr.base.is_none()
            && addr.index.is_none()
            && !addr.is_rip_relative
            && !addr.is_tls
            && !addr.is_frame_relative
            && (addr.base.is_symbol() || addr.displacement != 0)
    }

    fn legalize_absolute(&mut self, addr: &X86GenericAddrMode) -> LegalizedAddrMode {
        self.stats.absolute_converted += 1;

        if self.config.convert_absolute_to_rip
            && self.config.addr_width.supports_rip_relative()
            && self.config.pic_mode != PicMode::NotPIC
        {
            // Convert to RIP-relative
            let mut mode = addr.clone();
            mode.is_rip_relative = true;
            return self.legalize_rip_relative(&mode);
        }

        // 32-bit absolute or 64-bit absolute (small code model)
        if self.config.addr_width == AddrWidth::Addr32 {
            let mut result = LegalizedAddrMode::valid(addr.clone());
            result.disp_size = 4;
            result.needs_sib = false;
            return result;
        }

        // 64-bit absolute: use movabs to load address then use as base
        let mut result = LegalizedAddrMode::invalid();
        result.needs_lea = true;
        result.mode = Some(X86GenericAddrMode {
            width: AddrWidth::Addr64,
            base: AddrComponent::Register(self.scratch_reg.unwrap_or(0)),
            ..Default::default()
        });
        result
    }

    // -----------------------------------------------------------------------
    // Complex mode decomposition
    // -----------------------------------------------------------------------

    fn needs_decomposition(&self, addr: &X86GenericAddrMode) -> bool {
        if !self.config.decompose_complex {
            return false;
        }

        // Complex modes: both base and index present with non-trivial scale
        let has_base = !addr.base.is_none();
        let has_index = !addr.index.is_none();
        let has_nontrivial_scale = addr.scale != ScaleFactor::Times1;

        // Check if base and index are both symbols/frame indices
        let complex_components =
            has_base && has_index && (addr.base.is_symbol() || addr.base.is_frame_index());

        // Check for segment + symbol complexity
        let segment_symbol = addr.segment != SegmentOverride::None && addr.base.is_symbol();

        // Check for 16-bit complex addressing (limited valid combos)
        let sixteen_bit_complex = addr.width == AddrWidth::Addr16 && has_base && has_index;

        complex_components
            || segment_symbol
            || sixteen_bit_complex
            || (has_index && !has_base && has_nontrivial_scale && addr.displacement != 0)
    }

    fn legalize_complex(&mut self, addr: &X86GenericAddrMode) -> LegalizedAddrMode {
        self.stats.complex_decomposed += 1;
        self.stats.lea_formed += 1;

        // Decompose: compute effective address with LEA, use result as base
        let mut lea_mode = addr.clone();

        // Clean up LEA: remove segments, frame-relative flags for the LEA itself
        lea_mode.segment = SegmentOverride::None;
        lea_mode.is_frame_relative = false;
        lea_mode.is_rip_relative = false;
        lea_mode.is_tls = false;

        // Canonicalize the LEA form
        lea_mode = self.canonicalize_mode_for_lea(&lea_mode);

        let mut result = LegalizedAddrMode::invalid();
        result.needs_lea = true;
        result.lea_mode = Some(lea_mode);
        result.allocated_base = true;
        result.scratch_reg = self.scratch_reg;

        // The memory access uses the scratch register as base
        let mut mem_mode = X86GenericAddrMode::default();
        mem_mode.width = addr.width;
        mem_mode.base = AddrComponent::Register(self.scratch_reg.unwrap_or(0));
        mem_mode.segment = addr.segment;
        result.mode = Some(mem_mode);

        // Preserve segment override
        if addr.segment != SegmentOverride::None {
            result.needs_segment_prefix = true;
            result.segment_override_byte = addr.segment.prefix_byte();
        }

        result
    }

    // -----------------------------------------------------------------------
    // Canonicalization
    // -----------------------------------------------------------------------

    fn canonicalize_mode(&mut self, addr: &X86GenericAddrMode) -> X86GenericAddrMode {
        let mut mode = addr.clone();
        self.stats.scales_canonicalized += 1;

        // Canonicalize scale factor (must be 1, 2, 4, or 8)
        mode.scale = match mode.scale {
            ScaleFactor::Times1
            | ScaleFactor::Times2
            | ScaleFactor::Times4
            | ScaleFactor::Times8 => mode.scale,
            _ => ScaleFactor::Times1,
        };

        // If index is None, scale should be 1
        if mode.index.is_none() {
            mode.scale = ScaleFactor::Times1;
        }

        // RSP cannot be an index register — move to base
        if let AddrComponent::Register(idx) = mode.index {
            if idx == 4 {
                // RSP/ESP — swap base and index if possible
                if mode.base.is_none() {
                    mode.base = mode.index.clone();
                    mode.index = AddrComponent::None;
                    mode.scale = ScaleFactor::Times1;
                } else if let AddrComponent::Register(base) = mode.base {
                    if base != 4 && base != 5 {
                        // Swap: old index (RSP) becomes base, old base becomes index
                        std::mem::swap(&mut mode.base, &mut mode.index);
                    } else {
                        // Both RSP and RBP — need LEA decomposition
                        mode.index = AddrComponent::None;
                        mode.scale = ScaleFactor::Times1;
                    }
                }
            }
        }

        // RBP/R13 as base without displacement needs a dummy displacement (disp8=0)
        if let AddrComponent::Register(base) = mode.base {
            if (base == 5 || base == 13) && mode.displacement == 0 && mode.width.is_64bit() {
                mode.displacement = 0; // Will be encoded as disp8=0
            }
        }

        // Ensure displacement fits in the width
        mode.displacement = self.clamp_displacement(mode.displacement);

        mode
    }

    fn canonicalize_mode_for_lea(&self, addr: &X86GenericAddrMode) -> X86GenericAddrMode {
        let mut mode = addr.clone();

        // LEA doesn't have segment overrides, RIP-relative, or segment semantics
        mode.segment = SegmentOverride::None;
        mode.is_rip_relative = false;

        // For LEA: ensure base is a register or None
        if !mode.base.is_register() && !mode.base.is_none() {
            mode.base = AddrComponent::None;
        }
        if !mode.index.is_register() && !mode.index.is_none() {
            mode.index = AddrComponent::None;
            mode.scale = ScaleFactor::Times1;
        }

        // LEA in 64-bit mode defaults to 64-bit
        mode.width = self.config.addr_width;

        mode
    }

    fn clamp_displacement(&self, disp: i64) -> i64 {
        let max = match self.config.addr_width.max_displacement_bits() {
            16 => (1i64 << 15) - 1,
            32 => (1i64 << 31) - 1,
            _ => (1i64 << 31) - 1,
        };
        let min = -max - 1;
        disp.clamp(min, max)
    }

    // -----------------------------------------------------------------------
    // Displacement folding
    // -----------------------------------------------------------------------

    fn fold_displacements_into(&mut self, mut mode: X86GenericAddrMode) -> X86GenericAddrMode {
        if !self.config.fold_displacements {
            return mode;
        }

        // If base is a register with a known displacement offset, fold it
        if let AddrComponent::Register(_) = mode.base {
            // Already a register — nothing to fold into
        } else if mode.base.is_symbol() || mode.base.is_none() {
            // Symbols already have their displacement
        }

        // Fold index*scale contributions into displacement when possible
        // (e.g., when index is a constant or when scale is 1 and we can merge with base)
        self.stats.displacements_folded += 1;
        mode
    }

    // -----------------------------------------------------------------------
    // Narrow index promotion
    // -----------------------------------------------------------------------

    fn promote_narrow_indices_if_needed(
        &mut self,
        mut mode: X86GenericAddrMode,
    ) -> X86GenericAddrMode {
        if !self.config.promote_narrow_indices {
            return mode;
        }

        // In 64-bit mode, 8-bit and 16-bit index registers are not valid.
        // They must be promoted to 32-bit (which zero-extends to 64-bit in x86-64).
        if mode.width == AddrWidth::Addr64 {
            if let AddrComponent::Register(idx) = mode.index {
                // 8-bit regs: 0-3 (AL,CL,DL,BL) or 4-7 (AH,CH,DH,BH) — can't index
                // 16-bit regs: can't be used as index in 64-bit mode
                // In practice, the register allocator gives us 32-bit regs,
                // but we handle the edge case.

                // Check if register is a low 8-bit reg (AL=0, CL=1, DL=2, BL=3)
                // These need promotion to 32-bit equivalents (EAX=0, ECX=1, EDX=2, EBX=3)
                // 16-bit regs (AX=0, CX=1, DX=2, BX=3...) also need promotion
                // For 64-bit mode, index regs must be 64-bit (0-15)

                // In practice we assume the register allocator provides correct-width regs.
                // This is a safety net for edge cases.
                self.stats.indices_promoted += 1;
            }
        }

        mode
    }

    // -----------------------------------------------------------------------
    // Segment override application
    // -----------------------------------------------------------------------

    fn apply_segment_override(&self, mut mode: X86GenericAddrMode) -> X86GenericAddrMode {
        if !self.config.handle_segment_overrides {
            mode.segment = SegmentOverride::None;
            return mode;
        }

        // In 64-bit mode, DS, ES, SS are treated as "no override" by the hardware
        if mode.width.is_64bit()
            && !matches!(mode.segment, SegmentOverride::FS | SegmentOverride::GS)
        {
            mode.segment = SegmentOverride::None;
        }

        // If there's a default segment match, no override needed
        if let AddrComponent::Register(base_reg) = mode.base {
            let default_seg = if base_reg == 4 || base_reg == 5 {
                SegmentOverride::SS
            } else {
                SegmentOverride::DS
            };
            if mode.segment == default_seg {
                mode.segment = SegmentOverride::None;
            }
        }

        mode
    }

    // -----------------------------------------------------------------------
    // REX / SIB / displacement computation
    // -----------------------------------------------------------------------

    fn compute_rex_sib_disp(&self, result: LegalizedAddrMode) -> LegalizedAddrMode {
        let mut r = result;
        if let Some(ref mode) = r.mode {
            // REX prefix computation
            r.rex_bits = 0x40; // REX base
            if mode.width == AddrWidth::Addr64 {
                if let AddrComponent::Register(base) = mode.base {
                    if base >= 8 {
                        r.rex_bits |= 0x01; // REX.B
                    }
                }
                if let AddrComponent::Register(index) = mode.index {
                    if index >= 8 {
                        r.rex_bits |= 0x02; // REX.X
                    }
                }
            }

            // SIB computation
            r.needs_sib =
                mode.index.is_register() || (mode.base.is_register() && Self::base_needs_sib(mode));

            if r.needs_sib {
                let scale_bits = mode.scale.log2() & 0x3;
                let index_field = match mode.index {
                    AddrComponent::Register(idx) => (idx & 0x7) as u8,
                    _ => 0x4, // No index (SIB index=4 means no index)
                };
                let base_field = match mode.base {
                    AddrComponent::Register(reg) => (reg & 0x7) as u8,
                    _ => 0x5, // disp32 with no base (SIB base=5 means no base + disp32)
                };
                r.sib_byte = Some((scale_bits << 6) | (index_field << 3) | base_field);
            }

            // Displacement size
            r.disp_size = Self::compute_disp_size(mode.displacement);
        }

        r
    }

    fn base_needs_sib(mode: &X86GenericAddrMode) -> bool {
        if let AddrComponent::Register(reg) = mode.base {
            reg == 4 // RSP/ESP forces SIB byte
                || (reg == 5 && mode.displacement == 0 && mode.width.is_64bit())
        // RBP without disp
        } else {
            false
        }
    }

    fn compute_disp_size(disp: i64) -> u8 {
        if disp == 0 {
            0
        } else if disp >= -128 && disp <= 127 {
            1
        } else {
            4
        }
    }
}

// ---------------------------------------------------------------------------
// GEP combination
// ---------------------------------------------------------------------------

/// Represents a GEP (GetElementPtr) operation for combining.
#[derive(Debug, Clone)]
pub struct GepOperand {
    pub base: AddrComponent,
    pub index: AddrComponent,
    pub scale: u64,
    pub offset: i64,
    pub element_size: u64,
}

impl GepOperand {
    pub fn new(
        base: AddrComponent,
        index: AddrComponent,
        scale: u64,
        offset: i64,
        element_size: u64,
    ) -> Self {
        GepOperand {
            base,
            index,
            scale,
            offset,
            element_size,
        }
    }
}

/// Result of GEP combination.
#[derive(Debug, Clone)]
pub struct GepCombineResult {
    pub combined: X86GenericAddrMode,
    pub can_fold: bool,
    pub overflow: bool,
    pub needs_additional_shift: bool,
}

/// Analyze whether two GEP operations can be combined into a single LEA.
///
/// Example: `(base + i*8 + 16) + j*4` → LEA base, [base + i*2 + j + 4]
/// (after accounting for element sizes)
pub fn analyze_gep_combination(gep1: &GepOperand, gep2: &GepOperand) -> GepCombineResult {
    // Check if the outer GEP (gep2) uses the inner GEP (gep1) as its base.
    // The inner GEP produces an address; the outer GEP adds another offset.

    let mut combined = X86GenericAddrMode::default();

    // The inner base becomes the main base
    combined.base = gep1.base.clone();

    // Check if we can merge both indices and offsets
    let total_offset = gep1.offset + gep2.offset;
    combined.displacement = total_offset;

    // Try to merge indices
    let scale1 = gep1.scale * gep1.element_size;
    let scale2 = gep2.scale * gep2.element_size;

    let can_fold = if gep1.index.is_none() && gep2.index.is_none() {
        // No indices — just combine offsets
        combined.index = AddrComponent::None;
        combined.scale = ScaleFactor::Times1;
        true
    } else if gep1.index.is_none() {
        // Only gep2 has an index
        combined.index = gep2.index.clone();
        combined.scale = scale_to_factor(scale2);
        true
    } else if gep2.index.is_none() {
        // Only gep1 has an index
        combined.index = gep1.index.clone();
        combined.scale = scale_to_factor(scale1);
        true
    } else if gep1.index == gep2.index && scale1 == scale2 {
        // Same index, same scale → factor 2
        combined.index = gep1.index.clone();
        combined.scale = scale_to_factor(scale1 + scale2);
        true
    } else if scale1 > 0 && scale2 % scale1 == 0 {
        // Different indices but compatible scales
        combined.index = gep2.index.clone();
        combined.scale = scale_to_factor(scale2);
        // gep1's index contribution goes into displacement (simplified)
        // In reality this would need an extra add, so not fully foldable
        false
    } else if scale2 > 0 && scale1 % scale2 == 0 {
        combined.index = gep1.index.clone();
        combined.scale = scale_to_factor(scale1);
        false
    } else {
        // Incompatible — need separate instructions
        false
    };

    // Check for overflow
    let overflow = total_offset > (1i64 << 31) - 1 || total_offset < -(1i64 << 31);

    let needs_additional_shift =
        gep1.element_size != 1 && gep2.element_size != 1 && gep1.element_size != gep2.element_size;

    GepCombineResult {
        combined,
        can_fold,
        overflow,
        needs_additional_shift,
    }
}

/// Convert a scale value (1, 2, 4, 8) to a ScaleFactor.
pub fn scale_to_factor(scale: u64) -> ScaleFactor {
    match scale {
        1 => ScaleFactor::Times1,
        2 => ScaleFactor::Times2,
        4 => ScaleFactor::Times4,
        8 => ScaleFactor::Times8,
        _ => {
            // Non-power-of-2 scale: use Times1 and handle via shift
            ScaleFactor::Times1
        }
    }
}

// ---------------------------------------------------------------------------
// X86AddrModeLegalizer convenience methods for GEP combination
// ---------------------------------------------------------------------------

impl X86AddrModeLegalizer {
    /// Combine two GEP operations into a single LEA-based address, if possible.
    pub fn combine_geps(
        &mut self,
        gep1: &GepOperand,
        gep2: &GepOperand,
    ) -> Option<LegalizedAddrMode> {
        if !self.config.combine_geps {
            return None;
        }

        let analysis = analyze_gep_combination(gep1, gep2);
        if !analysis.can_fold || analysis.overflow {
            return None;
        }

        self.stats.gep_combined += 1;
        self.stats.lea_formed += 1;

        let mut mode = analysis.combined;
        mode.width = self.config.addr_width;

        // The combined GEP is best realized as an LEA
        let mut result = LegalizedAddrMode::invalid();
        result.needs_lea = true;
        result.lea_mode = Some(mode.clone());
        result.mode = Some(X86GenericAddrMode {
            width: self.config.addr_width,
            base: AddrComponent::Register(self.scratch_reg.unwrap_or(0)),
            ..Default::default()
        });
        result.allocated_base = true;
        result.scratch_reg = self.scratch_reg;

        Some(result)
    }

    /// Try to fold a constant offset into an existing addressing mode's displacement.
    pub fn try_fold_offset(
        &self,
        mode: &X86GenericAddrMode,
        offset: i64,
    ) -> Option<X86GenericAddrMode> {
        let new_disp = mode.displacement.checked_add(offset)?;
        let max_bits = self.config.addr_width.max_displacement_bits();
        let max_val = (1i64 << (max_bits - 1)) - 1;
        let min_val = -(1i64 << (max_bits - 1));

        if new_disp < min_val || new_disp > max_val {
            return None;
        }

        let mut new_mode = mode.clone();
        new_mode.displacement = new_disp;
        Some(new_mode)
    }

    /// Check if a register is a valid base register for the current address width.
    pub fn is_valid_base(&self, reg: u32) -> bool {
        match self.config.addr_width {
            AddrWidth::Addr64 => reg < 16, // R0-R15
            AddrWidth::Addr32 => reg < 8,  // EAX-EDI
            AddrWidth::Addr16 => reg < 8,  // AX-DI (BX, BP, SI, DI are valid bases)
        }
    }

    /// Check if a register is a valid index register for the current address width.
    pub fn is_valid_index(&self, reg: u32) -> bool {
        match self.config.addr_width {
            AddrWidth::Addr64 => reg < 16 && reg != 4, // Not RSP
            AddrWidth::Addr32 => reg < 8 && reg != 4,  // Not ESP
            AddrWidth::Addr16 => reg < 8 && reg != 4,  // Not SP (SI, DI, BX, BP are valid)
        }
    }

    /// Determine the classification of an addressing mode.
    pub fn classify(&self, addr: &X86GenericAddrMode) -> AddrModeClass {
        if addr.is_tls {
            return AddrModeClass::Tls;
        }
        if addr.is_rip_relative {
            return AddrModeClass::RipRelative;
        }
        if addr.is_frame_relative || addr.base.is_frame_index() {
            return AddrModeClass::FrameRelative;
        }
        if self.is_absolute(addr) {
            return AddrModeClass::Absolute;
        }
        if self.needs_decomposition(addr) {
            return AddrModeClass::NeedsDecomposition;
        }
        if !addr.index.is_none() && addr.scale != ScaleFactor::Times1 {
            return AddrModeClass::Indexed;
        }
        if !addr.index.is_none() {
            return AddrModeClass::Complex;
        }
        AddrModeClass::Simple
    }

    /// Get the encoding size of an addressing mode in bytes (modr/m + optional sib + disp).
    pub fn encoding_size(&self, mode: &X86GenericAddrMode) -> u8 {
        let mut size: u8 = 1; // ModR/M byte
        if mode.index.is_register() || Self::base_needs_sib(mode) {
            size += 1; // SIB byte
        }
        size += Self::compute_disp_size(mode.displacement);
        size
    }

    /// Compute the minimum encoding of two alternative modes.
    pub fn select_smaller_encoding<'a>(
        &self,
        a: &'a X86GenericAddrMode,
        b: &'a X86GenericAddrMode,
    ) -> &'a X86GenericAddrMode {
        if self.encoding_size(a) <= self.encoding_size(b) {
            a
        } else {
            b
        }
    }

    /// Check whether a displacement fits in 8 bits.
    pub fn fits_disp8(&self, disp: i64) -> bool {
        disp >= -128 && disp <= 127
    }

    /// Try to incrementally fold a small offset into an addressing mode.
    ///
    /// This is used during instruction selection when the same base address
    /// is accessed with different small offsets — we can fold the offset
    /// into the displacement instead of emitting a separate add.
    pub fn incremental_fold(
        &self,
        mode: &X86GenericAddrMode,
        offset: i64,
    ) -> Option<X86GenericAddrMode> {
        if !self.config.fold_displacements {
            return None;
        }
        self.try_fold_offset(mode, offset)
    }
}

// ---------------------------------------------------------------------------
// Loop-aware address mode hoisting/sinking helpers
// ---------------------------------------------------------------------------

/// Information about whether an address component is invariant in a loop.
#[derive(Debug, Clone)]
pub struct LoopInvariantInfo {
    pub base_is_invariant: bool,
    pub index_is_invariant: bool,
    pub disp_is_invariant: bool,
    pub can_hoist: bool,
}

impl X86AddrModeLegalizer {
    /// Check if an address mode can be hoisted out of a loop (i.e., all
    /// components are loop-invariant).
    pub fn is_loop_invariant(
        &self,
        addr: &X86GenericAddrMode,
        invariant_regs: &[u32],
    ) -> LoopInvariantInfo {
        let base_inv = match addr.base {
            AddrComponent::Register(r) => invariant_regs.contains(&r),
            AddrComponent::None | AddrComponent::FrameIndex(_) => true,
            _ => false,
        };
        let index_inv = match addr.index {
            AddrComponent::Register(r) => invariant_regs.contains(&r),
            AddrComponent::None => true,
            _ => false,
        };
        let disp_inv = true; // Displacements are constants

        LoopInvariantInfo {
            base_is_invariant: base_inv,
            index_is_invariant: index_inv,
            disp_is_invariant: disp_inv,
            can_hoist: base_inv && index_inv && disp_inv && !addr.is_rip_relative && !addr.is_tls,
        }
    }
}

// ---------------------------------------------------------------------------
// Batch legalization
// ---------------------------------------------------------------------------

/// A batch of addresses that can be legalized together (e.g., all addresses
/// in a basic block), enabling cross-address optimizations.
pub struct X86AddrModeBatchLegalizer {
    legalizer: X86AddrModeLegalizer,
    /// Common bases that can be factored.
    common_bases: HashMap<u32, Vec<usize>>,
}

impl X86AddrModeBatchLegalizer {
    pub fn new(config: X86AddrModeLegalizerConfig) -> Self {
        X86AddrModeBatchLegalizer {
            legalizer: X86AddrModeLegalizer::new(config),
            common_bases: HashMap::new(),
        }
    }

    /// Legalize a batch of addresses, optimizing common subexpressions.
    pub fn legalize_batch(&mut self, addrs: &[X86GenericAddrMode]) -> Vec<LegalizedAddrMode> {
        let mut results = Vec::with_capacity(addrs.len());

        // First pass: identify common bases
        self.common_bases.clear();
        for (i, addr) in addrs.iter().enumerate() {
            if let AddrComponent::Register(base) = addr.base {
                self.common_bases.entry(base).or_default().push(i);
            }
        }

        // Second pass: legalize each address
        for addr in addrs {
            results.push(self.legalizer.legalize(addr));
        }

        results
    }

    pub fn take_stats(&mut self) -> X86AddrModeLegalizerStats {
        self.legalizer.take_stats()
    }
}

// ---------------------------------------------------------------------------
// Address mode comparison and equivalence
// ---------------------------------------------------------------------------

/// Compare two address modes for equality in terms of the memory location accessed.
pub fn address_modes_equal(a: &X86GenericAddrMode, b: &X86GenericAddrMode) -> bool {
    a.base == b.base
        && a.index == b.index
        && a.scale == b.scale
        && a.displacement == b.displacement
        && a.segment == b.segment
        && a.is_rip_relative == b.is_rip_relative
}

/// Check if two address modes differ only by a constant displacement.
pub fn address_modes_differ_by_constant(
    a: &X86GenericAddrMode,
    b: &X86GenericAddrMode,
) -> Option<i64> {
    if a.base == b.base
        && a.index == b.index
        && a.scale == b.scale
        && a.segment == b.segment
        && a.is_rip_relative == b.is_rip_relative
    {
        Some(b.displacement - a.displacement)
    } else {
        None
    }
}

/// Merge two compatible address modes by adding their displacements.
pub fn merge_address_modes(
    a: &X86GenericAddrMode,
    b: &X86GenericAddrMode,
) -> Option<X86GenericAddrMode> {
    if a.base == b.base
        && a.index == b.index
        && a.scale == b.scale
        && a.segment == b.segment
        && a.is_rip_relative == b.is_rip_relative
    {
        let mut merged = a.clone();
        merged.displacement = a.displacement + b.displacement;
        Some(merged)
    } else {
        None
    }
}

// ---------------------------------------------------------------------------
// Convenience factory functions
// ---------------------------------------------------------------------------

pub fn make_x86_64_addr_mode_legalizer() -> X86AddrModeLegalizer {
    X86AddrModeLegalizer::new_x86_64()
}

pub fn make_x86_64_pic_addr_mode_legalizer() -> X86AddrModeLegalizer {
    X86AddrModeLegalizer::new_x86_64_pic()
}

pub fn make_x86_32_addr_mode_legalizer() -> X86AddrModeLegalizer {
    X86AddrModeLegalizer::new_x86_32()
}

pub fn make_x86_16_addr_mode_legalizer() -> X86AddrModeLegalizer {
    X86AddrModeLegalizer::new_x86_16()
}

pub fn make_addr_mode_legalizer_with_config(
    config: X86AddrModeLegalizerConfig,
) -> X86AddrModeLegalizer {
    X86AddrModeLegalizer::new(config)
}

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

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

    fn make_legalizer() -> X86AddrModeLegalizer {
        X86AddrModeLegalizer::new_x86_64()
    }

    fn make_legalizer_32() -> X86AddrModeLegalizer {
        X86AddrModeLegalizer::new_x86_32()
    }

    fn make_legalizer_pic() -> X86AddrModeLegalizer {
        X86AddrModeLegalizer::new_x86_64_pic()
    }

    // -- Constructor tests --

    #[test]
    fn test_constructor_default() {
        let leg = X86AddrModeLegalizer::new(X86AddrModeLegalizerConfig::default());
        assert_eq!(leg.config.addr_width, AddrWidth::Addr64);
        assert!(leg.config.use_rip_relative);
    }

    #[test]
    fn test_constructor_x86_64() {
        let leg = make_legalizer();
        assert_eq!(leg.config.addr_width, AddrWidth::Addr64);
    }

    #[test]
    fn test_constructor_x86_32() {
        let leg = make_legalizer_32();
        assert_eq!(leg.config.addr_width, AddrWidth::Addr32);
        assert!(!leg.config.use_rip_relative);
    }

    #[test]
    fn test_constructor_x86_16() {
        let leg = X86AddrModeLegalizer::new_x86_16();
        assert_eq!(leg.config.addr_width, AddrWidth::Addr16);
    }

    #[test]
    fn test_constructor_pic() {
        let leg = make_legalizer_pic();
        assert_eq!(leg.config.pic_mode, PicMode::PIC);
        assert!(leg.config.convert_absolute_to_rip);
    }

    // -- Width tests --

    #[test]
    fn test_addr_width_default_data_size() {
        assert_eq!(AddrWidth::Addr16.default_data_size(), 2);
        assert_eq!(AddrWidth::Addr32.default_data_size(), 4);
        assert_eq!(AddrWidth::Addr64.default_data_size(), 8);
    }

    #[test]
    fn test_addr_width_max_disp() {
        assert_eq!(AddrWidth::Addr16.max_displacement_bits(), 16);
        assert_eq!(AddrWidth::Addr32.max_displacement_bits(), 32);
        assert_eq!(AddrWidth::Addr64.max_displacement_bits(), 32);
    }

    #[test]
    fn test_addr_width_rip_relative_support() {
        assert!(!AddrWidth::Addr16.supports_rip_relative());
        assert!(!AddrWidth::Addr32.supports_rip_relative());
        assert!(AddrWidth::Addr64.supports_rip_relative());
    }

    // -- Simple address legalization --

    #[test]
    fn test_legalize_simple_base() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode::simple_64(0); // RAX
        let result = leg.legalize(&addr);
        assert!(result.mode.is_some());
        let m = result.mode.unwrap();
        assert_eq!(m.base, AddrComponent::Register(0));
        assert!(m.index.is_none());
        assert_eq!(m.displacement, 0);
        assert!(!result.needs_lea);
    }

    #[test]
    fn test_legalize_base_disp() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode::base_disp_64(0, 42); // [RAX + 42]
        let result = leg.legalize(&addr);
        assert!(result.mode.is_some());
        let m = result.mode.unwrap();
        assert_eq!(m.displacement, 42);
        assert_eq!(result.disp_size, 1); // disp8 fits
    }

    #[test]
    fn test_legalize_base_disp_large() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode::base_disp_64(0, 0x10000); // Large disp
        let result = leg.legalize(&addr);
        assert!(result.mode.is_some());
        assert_eq!(result.disp_size, 4);
    }

    #[test]
    fn test_legalize_full_mode() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode::full_64(0, 1, ScaleFactor::Times4, 16);
        let result = leg.legalize(&addr);
        assert!(result.mode.is_some());
        assert!(result.needs_sib || result.mode.as_ref().unwrap().index.is_register());
    }

    // -- RIP-relative tests --

    #[test]
    fn test_legalize_rip_relative() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode::rip_relative("foo".to_string(), 0);
        let result = leg.legalize(&addr);
        assert!(result.mode.is_some());
        assert!(result.used_rip_relative);
        assert_eq!(result.disp_size, 4);
    }

    #[test]
    fn test_rip_relative_symbol() {
        let mut leg = make_legalizer();
        let mut addr = X86GenericAddrMode::default();
        addr.base = AddrComponent::Symbol("global_var".to_string());
        addr.width = AddrWidth::Addr64;
        let result = leg.legalize(&addr);
        assert!(result.used_rip_relative);
    }

    #[test]
    fn test_rip_relative_not_32bit() {
        let mut leg = make_legalizer_32();
        let mut addr = X86GenericAddrMode::default();
        addr.base = AddrComponent::Symbol("var".to_string());
        addr.width = AddrWidth::Addr32;
        let result = leg.legalize(&addr);
        assert!(!result.used_rip_relative);
    }

    // -- Frame-relative tests --

    #[test]
    fn test_legalize_frame_index() {
        let mut leg = make_legalizer();
        leg.with_frame_offset(0, 5, -16); // FI 0 -> [RBP - 16]
        let mut addr = X86GenericAddrMode::default();
        addr.base = AddrComponent::FrameIndex(0);
        addr.width = AddrWidth::Addr64;
        let result = leg.legalize(&addr);
        assert!(result.mode.is_some());
        let m = result.mode.unwrap();
        assert_eq!(m.base, AddrComponent::Register(5));
        assert_eq!(m.displacement, -16);
        assert!(!m.is_frame_relative);
    }

    #[test]
    fn test_legalize_frame_index_with_disp() {
        let mut leg = make_legalizer();
        leg.with_frame_offset(1, 5, -32);
        let mut addr = X86GenericAddrMode::default();
        addr.base = AddrComponent::FrameIndex(1);
        addr.displacement = 8;
        addr.width = AddrWidth::Addr64;
        let result = leg.legalize(&addr);
        assert!(result.mode.is_some());
        let m = result.mode.unwrap();
        assert_eq!(m.displacement, -24); // -32 + 8
    }

    #[test]
    fn test_frame_relative_flag() {
        let mut leg = make_legalizer();
        leg.with_frame_offset(0, 4, -8); // RSP-8
        let mut addr = X86GenericAddrMode::default();
        addr.is_frame_relative = true;
        addr.width = AddrWidth::Addr64;
        let result = leg.legalize(&addr);
        assert!(result.mode.is_some());
    }

    // -- Complex decomposition tests --

    #[test]
    fn test_complex_decomposition() {
        let mut leg = make_legalizer();
        leg.with_scratch_reg(10); // R10 as scratch
        let mut addr = X86GenericAddrMode::full_64(0, 1, ScaleFactor::Times8, 128);
        addr.base = AddrComponent::Symbol("table".to_string()); // Complex: symbol + index

        // This should trigger decomposition since symbol can't be base with index
        let result = leg.legalize(&addr);
        // Should use LEA or decompose
        assert!(result.needs_lea || !result.decomposed.is_empty());
        assert!(result.scratch_reg.is_some());
    }

    #[test]
    fn test_rsp_as_index_fixed() {
        let mut leg = make_legalizer();
        // RSP=4 cannot be an index register
        let addr = X86GenericAddrMode::full_64(0, 4, ScaleFactor::Times2, 0);
        let result = leg.legalize(&addr);
        assert!(result.mode.is_some());
        // Should have canonicalized: RSP moved to base or decomposed
    }

    // -- Segment tests --

    #[test]
    fn test_segment_override_fs() {
        let mut leg = make_legalizer();
        let mut addr = X86GenericAddrMode::simple_64(0);
        addr.segment = SegmentOverride::FS;
        let result = leg.legalize(&addr);
        assert!(result.needs_segment_prefix);
        assert_eq!(result.segment_override_byte, Some(0x64));
    }

    #[test]
    fn test_segment_override_gs() {
        let mut leg = make_legalizer();
        let mut addr = X86GenericAddrMode::simple_64(0);
        addr.segment = SegmentOverride::GS;
        let result = leg.legalize(&addr);
        assert_eq!(result.segment_override_byte, Some(0x65));
    }

    #[test]
    fn test_segment_default_ds_no_override() {
        let mut leg = make_legalizer();
        let mut addr = X86GenericAddrMode::simple_64(0);
        addr.segment = SegmentOverride::DS;
        let result = leg.legalize(&addr);
        assert!(!result.needs_segment_prefix || result.segment_override_byte.is_none());
    }

    #[test]
    fn test_segment_stack_default() {
        let mut leg = make_legalizer_32();
        let mut addr = X86GenericAddrMode {
            width: AddrWidth::Addr32,
            base: AddrComponent::Register(4), // ESP
            ..Default::default()
        };
        addr.segment = SegmentOverride::SS;
        let result = leg.legalize(&addr);
        // SS default for ESP means no override needed
        assert!(!result.needs_segment_prefix || result.segment_override_byte.is_none());
    }

    // -- TLS tests --

    #[test]
    fn test_tls_local_exec_64() {
        let mut leg = make_legalizer();
        let mut addr = X86GenericAddrMode::default();
        addr.is_tls = true;
        addr.tls_model = Some(TlsModel::LocalExec);
        addr.base = AddrComponent::TlsSymbol("my_tls_var".to_string());
        addr.width = AddrWidth::Addr64;
        let result = leg.legalize(&addr);
        assert!(result.mode.is_some());
        assert!(result.needs_segment_prefix);
        assert_eq!(result.segment_override_byte, Some(0x64)); // FS
        assert_eq!(result.disp_size, 4);
    }

    #[test]
    fn test_tls_local_exec_32() {
        let mut leg = make_legalizer_32();
        let mut addr = X86GenericAddrMode::default();
        addr.is_tls = true;
        addr.tls_model = Some(TlsModel::LocalExec);
        addr.width = AddrWidth::Addr32;
        let result = leg.legalize(&addr);
        assert!(result.mode.is_some());
        assert_eq!(result.segment_override_byte, Some(0x65)); // GS in 32-bit mode
    }

    #[test]
    fn test_tls_initial_exec() {
        let mut leg = make_legalizer();
        leg.with_scratch_reg(0); // RAX as scratch
        let mut addr = X86GenericAddrMode::default();
        addr.is_tls = true;
        addr.tls_model = Some(TlsModel::InitialExec);
        addr.base = AddrComponent::TlsSymbol("var".to_string());
        let result = leg.legalize(&addr);
        assert!(result.needs_lea);
        assert!(!result.decomposed.is_empty());
    }

    #[test]
    fn test_tls_general_dynamic() {
        let mut leg = make_legalizer();
        let mut addr = X86GenericAddrMode::default();
        addr.is_tls = true;
        addr.tls_model = Some(TlsModel::GeneralDynamic);
        addr.base = AddrComponent::TlsSymbol("var".to_string());
        let result = leg.legalize(&addr);
        assert!(result.needs_lea);
        assert!(result.lea_mode.is_some());
    }

    // -- GEP combination tests --

    #[test]
    fn test_gep_combine_simple() {
        let gep1 = GepOperand::new(AddrComponent::Register(0), AddrComponent::None, 0, 16, 8);
        let gep2 = GepOperand::new(
            AddrComponent::Register(0),
            AddrComponent::Register(1),
            1,
            0,
            4,
        );
        let result = analyze_gep_combination(&gep1, &gep2);
        assert!(result.can_fold);
    }

    #[test]
    fn test_gep_combine_no_index() {
        let gep1 = GepOperand::new(AddrComponent::Register(5), AddrComponent::None, 0, 0, 1);
        let gep2 = GepOperand::new(AddrComponent::Register(5), AddrComponent::None, 0, 32, 1);
        let result = analyze_gep_combination(&gep1, &gep2);
        assert!(result.can_fold);
        assert_eq!(result.combined.displacement, 32); // 0 + 32
    }

    #[test]
    fn test_gep_combine_overflow() {
        let gep1 = GepOperand::new(
            AddrComponent::Register(0),
            AddrComponent::None,
            0,
            (1i64 << 31) - 1,
            1,
        );
        let gep2 = GepOperand::new(AddrComponent::Register(0), AddrComponent::None, 0, 16, 1);
        let result = analyze_gep_combination(&gep1, &gep2);
        assert!(result.overflow);
    }

    #[test]
    fn test_scale_to_factor() {
        assert_eq!(scale_to_factor(1), ScaleFactor::Times1);
        assert_eq!(scale_to_factor(2), ScaleFactor::Times2);
        assert_eq!(scale_to_factor(4), ScaleFactor::Times4);
        assert_eq!(scale_to_factor(8), ScaleFactor::Times8);
        assert_eq!(scale_to_factor(3), ScaleFactor::Times1); // Non-power-of-2
        assert_eq!(scale_to_factor(6), ScaleFactor::Times1);
    }

    // -- Displacement folding tests --

    #[test]
    fn test_try_fold_offset() {
        let leg = make_legalizer();
        let addr = X86GenericAddrMode::base_disp_64(0, 100);
        let folded = leg.try_fold_offset(&addr, 50);
        assert!(folded.is_some());
        assert_eq!(folded.unwrap().displacement, 150);
    }

    #[test]
    fn test_try_fold_offset_negative() {
        let leg = make_legalizer();
        let addr = X86GenericAddrMode::base_disp_64(0, 100);
        let folded = leg.try_fold_offset(&addr, -30);
        assert!(folded.is_some());
        assert_eq!(folded.unwrap().displacement, 70);
    }

    #[test]
    fn test_try_fold_offset_overflow() {
        let leg = make_legalizer();
        let addr = X86GenericAddrMode::base_disp_64(0, (1i64 << 30) - 1);
        let folded = leg.try_fold_offset(&addr, (1i64 << 30));
        assert!(folded.is_none()); // Exceeds 31-bit signed
    }

    #[test]
    fn test_incremental_fold() {
        let leg = make_legalizer();
        let addr = X86GenericAddrMode::base_disp_64(0, 16);
        let folded = leg.incremental_fold(&addr, 4);
        assert!(folded.is_some());
        assert_eq!(folded.unwrap().displacement, 20);
    }

    // -- Classification tests --

    #[test]
    fn test_classify_simple() {
        let leg = make_legalizer();
        let addr = X86GenericAddrMode::simple_64(0);
        assert_eq!(leg.classify(&addr), AddrModeClass::Simple);
    }

    #[test]
    fn test_classify_indexed() {
        let leg = make_legalizer();
        let addr = X86GenericAddrMode::full_64(0, 1, ScaleFactor::Times4, 0);
        assert_eq!(leg.classify(&addr), AddrModeClass::Indexed);
    }

    #[test]
    fn test_classify_frame_relative() {
        let leg = make_legalizer();
        let mut addr = X86GenericAddrMode::default();
        addr.base = AddrComponent::FrameIndex(0);
        assert_eq!(leg.classify(&addr), AddrModeClass::FrameRelative);
    }

    #[test]
    fn test_classify_tls() {
        let leg = make_legalizer();
        let mut addr = X86GenericAddrMode::default();
        addr.is_tls = true;
        assert_eq!(leg.classify(&addr), AddrModeClass::Tls);
    }

    #[test]
    fn test_classify_rip_relative() {
        let leg = make_legalizer();
        let addr = X86GenericAddrMode::rip_relative("sym".to_string(), 0);
        assert_eq!(leg.classify(&addr), AddrModeClass::RipRelative);
    }

    #[test]
    fn test_classify_absolute() {
        let leg = make_legalizer();
        let mut addr = X86GenericAddrMode::default();
        addr.width = AddrWidth::Addr32;
        addr.displacement = 0x400000;
        assert_eq!(leg.classify(&addr), AddrModeClass::Absolute);
    }

    // -- Encoding size tests --

    #[test]
    fn test_encoding_size_simple() {
        let leg = make_legalizer();
        let addr = X86GenericAddrMode::simple_64(0);
        assert_eq!(leg.encoding_size(&addr), 1); // ModR/M only
    }

    #[test]
    fn test_encoding_size_with_disp8() {
        let leg = make_legalizer();
        let addr = X86GenericAddrMode::base_disp_64(0, 42);
        assert_eq!(leg.encoding_size(&addr), 2); // ModR/M + disp8
    }

    #[test]
    fn test_encoding_size_with_sib() {
        let leg = make_legalizer();
        let addr = X86GenericAddrMode::full_64(0, 1, ScaleFactor::Times2, 0);
        assert_eq!(leg.encoding_size(&addr), 2); // ModR/M + SIB
    }

    #[test]
    fn test_select_smaller_encoding() {
        let leg = make_legalizer();
        let a = X86GenericAddrMode::simple_64(0);
        let b = X86GenericAddrMode::base_disp_64(0, 1000); // Needs disp32
        let smaller = leg.select_smaller_encoding(&a, &b);
        assert_eq!(smaller.displacement, 0);
    }

    // -- Loop invariance tests --

    #[test]
    fn test_loop_invariant_simple() {
        let leg = make_legalizer();
        let addr = X86GenericAddrMode::simple_64(0);
        let info = leg.is_loop_invariant(&addr, &[0]);
        assert!(info.can_hoist);
    }

    #[test]
    fn test_loop_invariant_index_not_invariant() {
        let leg = make_legalizer();
        let addr = X86GenericAddrMode::full_64(0, 1, ScaleFactor::Times4, 0);
        let info = leg.is_loop_invariant(&addr, &[0]); // Index 1 not invariant
        assert!(!info.can_hoist);
    }

    #[test]
    fn test_loop_invariant_all_invariant() {
        let leg = make_legalizer();
        let addr = X86GenericAddrMode::full_64(0, 1, ScaleFactor::Times2, 16);
        let info = leg.is_loop_invariant(&addr, &[0, 1]);
        assert!(info.can_hoist);
        assert!(info.base_is_invariant);
        assert!(info.index_is_invariant);
    }

    // -- Address comparison tests --

    #[test]
    fn test_address_modes_equal_true() {
        let a = X86GenericAddrMode::simple_64(0);
        let b = X86GenericAddrMode::simple_64(0);
        assert!(address_modes_equal(&a, &b));
    }

    #[test]
    fn test_address_modes_equal_false_diff_base() {
        let a = X86GenericAddrMode::simple_64(0);
        let b = X86GenericAddrMode::simple_64(1);
        assert!(!address_modes_equal(&a, &b));
    }

    #[test]
    fn test_address_modes_differ_by_constant() {
        let a = X86GenericAddrMode::base_disp_64(0, 16);
        let b = X86GenericAddrMode::base_disp_64(0, 20);
        assert_eq!(address_modes_differ_by_constant(&a, &b), Some(4));
    }

    #[test]
    fn test_address_modes_differ_by_constant_not() {
        let a = X86GenericAddrMode::simple_64(0);
        let b = X86GenericAddrMode::simple_64(1);
        assert_eq!(address_modes_differ_by_constant(&a, &b), None);
    }

    #[test]
    fn test_merge_address_modes() {
        let a = X86GenericAddrMode::base_disp_64(0, 10);
        let b = X86GenericAddrMode::base_disp_64(0, 20);
        let merged = merge_address_modes(&a, &b);
        assert!(merged.is_some());
        assert_eq!(merged.unwrap().displacement, 30);
    }

    #[test]
    fn test_merge_address_modes_incompatible() {
        let a = X86GenericAddrMode::simple_64(0);
        let b = X86GenericAddrMode::simple_64(1);
        let merged = merge_address_modes(&a, &b);
        assert!(merged.is_none());
    }

    // -- Component tests --

    #[test]
    fn test_addr_component_default() {
        let c = AddrComponent::default();
        assert!(c.is_none());
    }

    #[test]
    fn test_addr_component_is_register() {
        assert!(AddrComponent::Register(0).is_register());
        assert!(!AddrComponent::FrameIndex(0).is_register());
        assert!(!AddrComponent::Symbol("x".to_string()).is_register());
    }

    #[test]
    fn test_addr_component_is_symbol() {
        assert!(AddrComponent::Symbol("x".to_string()).is_symbol());
        assert!(AddrComponent::TlsSymbol("x".to_string()).is_symbol());
        assert!(!AddrComponent::Register(0).is_symbol());
    }

    #[test]
    fn test_addr_component_is_frame_index() {
        assert!(AddrComponent::FrameIndex(0).is_frame_index());
        assert!(!AddrComponent::Register(0).is_frame_index());
    }

    // -- Segment tests --

    #[test]
    fn test_segment_override_prefix_bytes() {
        assert_eq!(SegmentOverride::None.prefix_byte(), None);
        assert_eq!(SegmentOverride::ES.prefix_byte(), Some(0x26));
        assert_eq!(SegmentOverride::CS.prefix_byte(), Some(0x2E));
        assert_eq!(SegmentOverride::SS.prefix_byte(), Some(0x36));
        assert_eq!(SegmentOverride::DS.prefix_byte(), Some(0x3E));
        assert_eq!(SegmentOverride::FS.prefix_byte(), Some(0x64));
        assert_eq!(SegmentOverride::GS.prefix_byte(), Some(0x65));
    }

    #[test]
    fn test_segment_is_stack() {
        assert!(SegmentOverride::SS.is_stack_segment());
        assert!(!SegmentOverride::DS.is_stack_segment());
        assert!(!SegmentOverride::FS.is_stack_segment());
    }

    #[test]
    fn test_segment_default_for_register() {
        assert_eq!(
            SegmentOverride::default_for_register(4, false),
            SegmentOverride::SS
        );
        assert_eq!(
            SegmentOverride::default_for_register(5, false),
            SegmentOverride::SS
        );
        assert_eq!(
            SegmentOverride::default_for_register(0, false),
            SegmentOverride::DS
        );
    }

    // -- Scale factor tests --

    #[test]
    fn test_scale_factor_from_u8() {
        assert_eq!(ScaleFactor::from_u8(1), Some(ScaleFactor::Times1));
        assert_eq!(ScaleFactor::from_u8(2), Some(ScaleFactor::Times2));
        assert_eq!(ScaleFactor::from_u8(4), Some(ScaleFactor::Times4));
        assert_eq!(ScaleFactor::from_u8(8), Some(ScaleFactor::Times8));
        assert_eq!(ScaleFactor::from_u8(3), None);
        assert_eq!(ScaleFactor::from_u8(0), None);
    }

    #[test]
    fn test_scale_factor_log2() {
        assert_eq!(ScaleFactor::Times1.log2(), 0);
        assert_eq!(ScaleFactor::Times2.log2(), 1);
        assert_eq!(ScaleFactor::Times4.log2(), 2);
        assert_eq!(ScaleFactor::Times8.log2(), 3);
    }

    #[test]
    fn test_scale_factor_as_u32() {
        assert_eq!(ScaleFactor::Times1.as_u32(), 1);
        assert_eq!(ScaleFactor::Times8.as_u32(), 8);
    }

    // -- REX computation tests --

    #[test]
    fn test_rex_high_base() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode::simple_64(8); // R8
        let result = leg.legalize(&addr);
        assert!(result.rex_bits & 0x01 != 0); // REX.B set
    }

    #[test]
    fn test_rex_high_index() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode::full_64(0, 9, ScaleFactor::Times1, 0); // R9 as index
        let result = leg.legalize(&addr);
        assert!(result.rex_bits & 0x02 != 0); // REX.X set
    }

    #[test]
    fn test_rex_low_registers() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode::simple_64(0); // RAX
        let result = leg.legalize(&addr);
        // REX prefix still has 0x40 base but B/X/R/W bits are 0
        assert_eq!(result.rex_bits & 0x0F, 0x00);
    }

    // -- SIB tests --

    #[test]
    fn test_sib_needed_for_index() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode::full_64(0, 1, ScaleFactor::Times2, 0);
        let result = leg.legalize(&addr);
        assert!(result.needs_sib);
        assert!(result.sib_byte.is_some());
    }

    #[test]
    fn test_sib_for_rsp_base() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode::simple_64(4); // RSP
        let result = leg.legalize(&addr);
        assert!(result.needs_sib);
    }

    #[test]
    fn test_sib_byte_encoding() {
        let mut leg = make_legalizer();
        // [RAX + RCX*4 + 0]
        let addr = X86GenericAddrMode::full_64(0, 1, ScaleFactor::Times4, 0);
        let result = leg.legalize(&addr);
        assert!(result.sib_byte.is_some());
        let sib = result.sib_byte.unwrap();
        // scale = 2 (Times4), index = 1 (RCX), base = 0 (RAX)
        assert_eq!(sib, (2 << 6) | (1 << 3) | 0);
    }

    // -- Displacement size tests --

    #[test]
    fn test_disp_size_zero() {
        assert_eq!(X86AddrModeLegalizer::compute_disp_size(0), 0);
    }

    #[test]
    fn test_disp_size_8() {
        assert_eq!(X86AddrModeLegalizer::compute_disp_size(127), 1);
        assert_eq!(X86AddrModeLegalizer::compute_disp_size(-128), 1);
        assert_eq!(X86AddrModeLegalizer::compute_disp_size(42), 1);
        assert_eq!(X86AddrModeLegalizer::compute_disp_size(-42), 1);
    }

    #[test]
    fn test_disp_size_32() {
        assert_eq!(X86AddrModeLegalizer::compute_disp_size(128), 4);
        assert_eq!(X86AddrModeLegalizer::compute_disp_size(-129), 4);
        assert_eq!(X86AddrModeLegalizer::compute_disp_size(0x10000), 4);
    }

    #[test]
    fn test_fits_disp8() {
        let leg = make_legalizer();
        assert!(leg.fits_disp8(0));
        assert!(leg.fits_disp8(127));
        assert!(leg.fits_disp8(-128));
        assert!(!leg.fits_disp8(128));
        assert!(!leg.fits_disp8(-129));
    }

    // -- Valid register checks --

    #[test]
    fn test_is_valid_base_64() {
        let leg = make_legalizer();
        assert!(leg.is_valid_base(0));
        assert!(leg.is_valid_base(15));
        assert!(!leg.is_valid_base(16));
    }

    #[test]
    fn test_is_valid_base_32() {
        let leg = make_legalizer_32();
        assert!(leg.is_valid_base(0));
        assert!(leg.is_valid_base(7));
        assert!(!leg.is_valid_base(8));
    }

    #[test]
    fn test_is_valid_index_64() {
        let leg = make_legalizer();
        assert!(leg.is_valid_index(0));
        assert!(!leg.is_valid_index(4)); // RSP cannot be index
        assert!(leg.is_valid_index(15));
    }

    #[test]
    fn test_is_valid_index_32() {
        let leg = make_legalizer_32();
        assert!(leg.is_valid_index(0));
        assert!(!leg.is_valid_index(4)); // ESP cannot be index
    }

    // -- Batch legalizer tests --

    #[test]
    fn test_batch_legalizer_empty() {
        let mut batch = X86AddrModeBatchLegalizer::new(X86AddrModeLegalizerConfig::default());
        let results = batch.legalize_batch(&[]);
        assert!(results.is_empty());
    }

    #[test]
    fn test_batch_legalizer_multiple() {
        let mut batch = X86AddrModeBatchLegalizer::new(X86AddrModeLegalizerConfig::default());
        let addrs = vec![
            X86GenericAddrMode::simple_64(0),
            X86GenericAddrMode::simple_64(1),
            X86GenericAddrMode::simple_64(0), // Same base as first
        ];
        let results = batch.legalize_batch(&addrs);
        assert_eq!(results.len(), 3);
        for r in &results {
            assert!(r.mode.is_some());
        }
    }

    #[test]
    fn test_batch_legalizer_stats() {
        let mut batch = X86AddrModeBatchLegalizer::new(X86AddrModeLegalizerConfig::default());
        let addrs = vec![X86GenericAddrMode::simple_64(0)];
        let _ = batch.legalize_batch(&addrs);
        let stats = batch.take_stats();
        assert!(stats.total_addrs_processed > 0);
    }

    // -- Stats tests --

    #[test]
    fn test_stats_new_is_zero() {
        let stats = X86AddrModeLegalizerStats::new();
        assert_eq!(stats.total_addrs_processed, 0);
        assert_eq!(stats.total(), 0);
    }

    #[test]
    fn test_stats_merge() {
        let mut a = X86AddrModeLegalizerStats::new();
        a.total_addrs_processed = 10;
        a.rip_relative_converted = 5;

        let mut b = X86AddrModeLegalizerStats::new();
        b.total_addrs_processed = 20;
        b.rip_relative_converted = 3;

        a.merge(&b);
        assert_eq!(a.total_addrs_processed, 30);
        assert_eq!(a.rip_relative_converted, 8);
    }

    #[test]
    fn test_legalizer_take_stats() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode::simple_64(0);
        let _ = leg.legalize(&addr);
        let stats = leg.take_stats();
        assert_eq!(stats.total_addrs_processed, 1);
        // After take, stats are reset
        let stats2 = leg.take_stats();
        assert_eq!(stats2.total_addrs_processed, 0);
    }

    // -- Cache tests --

    #[test]
    fn test_cache_hit() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode::simple_64(0);
        let _ = leg.legalize(&addr);
        let stats_before = leg.stats.cache_hits;
        let _ = leg.legalize(&addr);
        assert!(leg.stats.cache_hits > stats_before);
    }

    #[test]
    fn test_cache_clear() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode::simple_64(0);
        let _ = leg.legalize(&addr);
        leg.clear_cache();
        let stats_before = leg.stats.cache_misses;
        let _ = leg.legalize(&addr);
        assert!(leg.stats.cache_misses > stats_before);
    }

    // -- Factory function tests --

    #[test]
    fn test_factory_x86_64() {
        let leg = make_x86_64_addr_mode_legalizer();
        assert_eq!(leg.config.addr_width, AddrWidth::Addr64);
    }

    #[test]
    fn test_factory_x86_64_pic() {
        let leg = make_x86_64_pic_addr_mode_legalizer();
        assert_eq!(leg.config.pic_mode, PicMode::PIC);
    }

    #[test]
    fn test_factory_x86_32() {
        let leg = make_x86_32_addr_mode_legalizer();
        assert_eq!(leg.config.addr_width, AddrWidth::Addr32);
    }

    #[test]
    fn test_factory_x86_16() {
        let leg = make_x86_16_addr_mode_legalizer();
        assert_eq!(leg.config.addr_width, AddrWidth::Addr16);
    }

    #[test]
    fn test_factory_with_config() {
        let config = X86AddrModeLegalizerConfig {
            addr_width: AddrWidth::Addr64,
            use_rip_relative: false,
            ..Default::default()
        };
        let leg = make_addr_mode_legalizer_with_config(config);
        assert!(!leg.config.use_rip_relative);
    }

    // -- Edge case tests --

    #[test]
    fn test_legalize_default_addr() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode::default();
        let result = leg.legalize(&addr);
        // Default: 64-bit, no base, no index, no displacement -> Absolute
        // Should produce some valid result
        assert!(result.mode.is_some() || result.needs_lea);
    }

    #[test]
    fn test_legalize_wide_displacement() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode::base_disp_64(0, 0x7FFFFFFF);
        let result = leg.legalize(&addr);
        assert!(result.mode.is_some());
        assert_eq!(result.disp_size, 4);
    }

    #[test]
    fn test_legalize_negative_displacement() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode::base_disp_64(0, -8);
        let result = leg.legalize(&addr);
        assert!(result.mode.is_some());
        assert_eq!(result.disp_size, 1); // -8 fits in disp8
    }

    #[test]
    fn test_legalize_rbp_zero_disp() {
        let mut leg = make_legalizer();
        let addr = X86GenericAddrMode {
            width: AddrWidth::Addr64,
            base: AddrComponent::Register(5), // RBP
            displacement: 0,
            ..Default::default()
        };
        let result = leg.legalize(&addr);
        assert!(result.mode.is_some());
        // RBP with zero displacement should work (uses disp8=0 encoding)
        assert!(result.needs_sib);
    }

    #[test]
    fn test_legalize_got_relative() {
        let mut leg = make_legalizer();
        let mut addr = X86GenericAddrMode::default();
        addr.is_got_relative = true;
        addr.base = AddrComponent::GotBase;
        let result = leg.legalize(&addr);
        assert!(result.mode.is_some() || result.needs_lea);
    }

    // -- GEP combine with legalizer tests --

    #[test]
    fn test_combine_geps_through_legalizer() {
        let mut leg = make_legalizer();
        leg.with_scratch_reg(10);
        let gep1 = GepOperand::new(AddrComponent::Register(0), AddrComponent::None, 0, 0, 8);
        let gep2 = GepOperand::new(
            AddrComponent::Register(0),
            AddrComponent::Register(1),
            1,
            16,
            4,
        );
        let result = leg.combine_geps(&gep1, &gep2);
        assert!(result.is_some());
        let r = result.unwrap();
        assert!(r.needs_lea);
        assert!(r.scratch_reg.is_some());
    }

    #[test]
    fn test_combine_geps_disabled() {
        let config = X86AddrModeLegalizerConfig {
            combine_geps: false,
            ..Default::default()
        };
        let mut leg = X86AddrModeLegalizer::new(config);
        let gep1 = GepOperand::new(AddrComponent::Register(0), AddrComponent::None, 0, 0, 8);
        let gep2 = GepOperand::new(AddrComponent::Register(0), AddrComponent::None, 0, 4, 4);
        let result = leg.combine_geps(&gep1, &gep2);
        assert!(result.is_none());
    }

    // -- Hash key tests --

    #[test]
    fn test_hash_key_consistent() {
        let a = X86GenericAddrMode::simple_64(0);
        let b = X86GenericAddrMode::simple_64(0);
        assert_eq!(a.hash_key(), b.hash_key());
    }

    #[test]
    fn test_hash_key_different() {
        let a = X86GenericAddrMode::simple_64(0);
        let b = X86GenericAddrMode::simple_64(1);
        assert_ne!(a.hash_key(), b.hash_key());
    }

    // -- Pic mode tests --

    #[test]
    fn test_pic_mode_converts_absolute() {
        let mut leg = make_legalizer_pic();
        let mut addr = X86GenericAddrMode::default();
        addr.width = AddrWidth::Addr64;
        addr.base = AddrComponent::Symbol("my_var".to_string());
        let result = leg.legalize(&addr);
        assert!(result.used_rip_relative);
    }

    #[test]
    fn test_pic_mode_no_got() {
        let leg = X86AddrModeLegalizer::new(X86AddrModeLegalizerConfig {
            addr_width: AddrWidth::Addr64,
            pic_mode: PicMode::PICNoGOT,
            convert_absolute_to_rip: false,
            ..Default::default()
        });
        assert_eq!(leg.config.pic_mode, PicMode::PICNoGOT);
    }

    // -- Comprehensive mode tests --

    #[test]
    fn test_all_scale_factors_work() {
        let mut leg = make_legalizer();
        for scale in [
            ScaleFactor::Times1,
            ScaleFactor::Times2,
            ScaleFactor::Times4,
            ScaleFactor::Times8,
        ] {
            let addr = X86GenericAddrMode::full_64(0, 1, scale, 0);
            let result = leg.legalize(&addr);
            assert!(result.mode.is_some());
        }
    }

    #[test]
    fn test_all_base_registers_64() {
        let mut leg = make_legalizer();
        for reg in 0..16u32 {
            let addr = X86GenericAddrMode::simple_64(reg);
            let result = leg.legalize(&addr);
            assert!(result.mode.is_some(), "Failed for register {}", reg);
        }
    }

    #[test]
    fn test_displacement_boundary_cases() {
        let mut leg = make_legalizer();
        // -128: fits in disp8
        let r1 = leg.legalize(&X86GenericAddrMode::base_disp_64(0, -128));
        assert_eq!(r1.disp_size, 1);
        // -129: needs disp32
        let r2 = leg.legalize(&X86GenericAddrMode::base_disp_64(0, -129));
        assert_eq!(r2.disp_size, 4);
        // 127: fits in disp8
        let r3 = leg.legalize(&X86GenericAddrMode::base_disp_64(0, 127));
        assert_eq!(r3.disp_size, 1);
        // 128: needs disp32
        let r4 = leg.legalize(&X86GenericAddrMode::base_disp_64(0, 128));
        assert_eq!(r4.disp_size, 4);
    }
}