llvm-native-core 0.1.16

//! BOLT X86 Binary Optimizer Integration.
//!
//! Clean-room behavioral reconstruction of the BOLT binary optimizer
//! (Panchenko et al. 2019 "BOLT: A Practical Binary Optimizer for Data
//! Centers and Beyond"), specifically adapted for X86 (IA-32 and X86-64)
//! ELF binaries.
//!
//! BOLT operates on already-compiled binaries to perform profile-guided
//! post-link optimizations without source-code access. This module
//! integrates X86-specific disassembly, CFG reconstruction, profile
//! reading, function/block reordering, and binary rewriting.
//!
//! ## Architecture
//!
//! The pipeline proceeds as follows:
//! 1. **Disassembly** — `X86BOLTDisassembler` disassembles the binary,
//!    building a CFG of basic blocks per function.
//! 2. **Binary Analysis** — `X86BOLTBinaryAnalysis` identifies function
//!    boundaries, entry points, and symbol table entries.
//! 3. **Profile Reading** — `X86BOLTProfile` reads perf/LBR/AutoFDO
//!    profile data and annotates the CFG with execution counts.
//! 4. **Optimization** — `X86BOLTOptimizer` applies post-link optimizations:
//!    function reordering (HFSort/Call-Chain), basic-block reordering,
//!    NOP removal, jump elimination, frame optimization.
//! 5. **Rewriting** — `X86BOLTRewriter` emits the optimized binary with
//!    updated addresses, relocations, symbols, and debug info.
//! 6. **Orchestration** — `BOLTX86` ties the pipeline together with a
//!    high-level API.
//!
//! ## Integration Points
//!
//! - `crate::bolt::bolt_profile::*` — generic profile reading infrastructure
//! - `crate::bolt::bolt_rewrite::*` — generic binary rewriting infrastructure
//! - `crate::x86::*` — X86 instruction info, registers, opcodes
//! - `crate::mc_disassembler::*` — MC-level disassembly primitives
//! - `crate::object_file::*` — ELF object file parsing
//! - `crate::dwarf::*` — DWARF debug info handling
//! - `crate::elf::*` — ELF format structures
//!
//! Clean-room behavioral reconstruction. Zero LLVM/BOLT source code
//! consultation. Phase 9 — LLVM.BOLT.X86.1 Court.

use crate::bolt::bolt_profile::{BOLTProfileReader, BoltProfile, FunctionProfile, LbrSample};
use crate::bolt::bolt_rewrite::{BOLTBinaryRewriter, BoltBlock, BoltFunction, LayoutAlgorithm};
use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::fmt;

// ============================================================================
// X86 Constants
// ============================================================================

/// Maximum size of a single X86 instruction in bytes.
pub const X86_MAX_INSN_SIZE: usize = 15;

/// X86 NOP opcode byte.
pub const X86_NOP_OPCODE: u8 = 0x90;

/// Multi-byte NOP prefixes for X86.
pub const X86_MULTI_NOP_PREFIX: u8 = 0x0F;
pub const X86_MULTI_NOP_OPCODE: u8 = 0x1F;

/// X86 unconditional jump opcode (short).
pub const X86_JMP_REL8: u8 = 0xEB;

/// X86 unconditional jump opcode (near).
pub const X86_JMP_REL32: u8 = 0xE9;

/// X86 conditional jump opcodes.
pub const X86_JO_REL8: u8 = 0x70;
pub const X86_JNO_REL8: u8 = 0x71;
pub const X86_JB_REL8: u8 = 0x72;
pub const X86_JNB_REL8: u8 = 0x73;
pub const X86_JZ_REL8: u8 = 0x74;
pub const X86_JNZ_REL8: u8 = 0x75;
pub const X86_JBE_REL8: u8 = 0x76;
pub const X86_JA_REL8: u8 = 0x77;
pub const X86_JS_REL8: u8 = 0x78;
pub const X86_JNS_REL8: u8 = 0x79;
pub const X86_JP_REL8: u8 = 0x7A;
pub const X86_JNP_REL8: u8 = 0x7B;
pub const X86_JL_REL8: u8 = 0x7C;
pub const X86_JGE_REL8: u8 = 0x7D;
pub const X86_JLE_REL8: u8 = 0x7E;
pub const X86_JG_REL8: u8 = 0x7F;

/// X86 conditional jump near (2-byte prefix).
pub const X86_JCC_NEAR_PREFIX: u8 = 0x0F;
pub const X86_JO_REL32: u8 = 0x80;
pub const X86_JNO_REL32: u8 = 0x81;
pub const X86_JB_REL32: u8 = 0x82;
pub const X86_JNB_REL32: u8 = 0x83;
pub const X86_JZ_REL32: u8 = 0x84;
pub const X86_JNZ_REL32: u8 = 0x85;
pub const X86_JBE_REL32: u8 = 0x86;
pub const X86_JA_REL32: u8 = 0x87;
pub const X86_JS_REL32: u8 = 0x88;
pub const X86_JNS_REL32: u8 = 0x89;
pub const X86_JP_REL32: u8 = 0x8A;
pub const X86_JNP_REL32: u8 = 0x8B;
pub const X86_JL_REL32: u8 = 0x8C;
pub const X86_JGE_REL32: u8 = 0x8D;
pub const X86_JLE_REL32: u8 = 0x8E;
pub const X86_JG_REL32: u8 = 0x8F;

/// X86 return opcodes.
pub const X86_RET_NEAR: u8 = 0xC3;
pub const X86_RET_NEAR_IMM16: u8 = 0xC2;
pub const X86_RET_FAR: u8 = 0xCB;
pub const X86_RET_FAR_IMM16: u8 = 0xCA;

/// X86 call opcodes.
pub const X86_CALL_REL32: u8 = 0xE8;
pub const X86_CALL_RM_INDIRECT: u16 = 0x15FF;

/// X86 REX prefix bytes.
pub const X86_REX_PREFIX_MIN: u8 = 0x40;
pub const X86_REX_PREFIX_MAX: u8 = 0x4F;

/// X86 VEX prefix bytes.
pub const X86_VEX2_PREFIX: u8 = 0xC5;
pub const X86_VEX3_PREFIX: u8 = 0xC4;

/// X86 EVEX prefix byte.
pub const X86_EVEX_PREFIX: u8 = 0x62;

/// X86 instruction prefixes.
pub const X86_PREFIX_LOCK: u8 = 0xF0;
pub const X86_PREFIX_REPNE: u8 = 0xF2;
pub const X86_PREFIX_REP: u8 = 0xF3;
pub const X86_PREFIX_CS: u8 = 0x2E;
pub const X86_PREFIX_SS: u8 = 0x36;
pub const X86_PREFIX_DS: u8 = 0x3E;
pub const X86_PREFIX_ES: u8 = 0x26;
pub const X86_PREFIX_FS: u8 = 0x64;
pub const X86_PREFIX_GS: u8 = 0x65;
pub const X86_PREFIX_DATA16: u8 = 0x66;
pub const X86_PREFIX_ADDR16: u8 = 0x67;

/// X86 ModR/M byte mask.
pub const X86_MODRM_MOD_MASK: u8 = 0xC0;
pub const X86_MODRM_REG_MASK: u8 = 0x38;
pub const X86_MODRM_RM_MASK: u8 = 0x07;

/// ModR/M MOD values.
pub const X86_MODRM_MOD_DISP0: u8 = 0x00;
pub const X86_MODRM_MOD_DISP8: u8 = 0x40;
pub const X86_MODRM_MOD_DISP32: u8 = 0x80;
pub const X86_MODRM_MOD_REG: u8 = 0xC0;

/// SIB byte scale mask.
pub const X86_SIB_SCALE_MASK: u8 = 0xC0;
pub const X86_SIB_INDEX_MASK: u8 = 0x38;
pub const X86_SIB_BASE_MASK: u8 = 0x07;

/// X86 instruction categories for BOLT analysis.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum X86InsnCategory {
    /// Normal instruction.
    Normal,
    /// Unconditional jump.
    UnconditionalJump,
    /// Conditional jump.
    ConditionalJump,
    /// Call instruction.
    Call,
    /// Return instruction.
    Return,
    /// Indirect jump (jump table dispatch, virtual call).
    IndirectJump,
    /// NOP instruction.
    Nop,
    /// Prefix bytes.
    Prefix,
    /// Halt or invalid.
    Halt,
    /// Other control flow.
    Other,
}

// ============================================================================
// X86 Instruction Representation
// ============================================================================

/// A decoded X86 instruction for BOLT analysis.
#[derive(Debug, Clone)]
pub struct X86BoltInstruction {
    /// Virtual address of this instruction.
    pub address: u64,
    /// Size of the instruction in bytes.
    pub size: usize,
    /// Raw bytes of the instruction.
    pub bytes: Vec<u8>,
    /// Instruction category.
    pub category: X86InsnCategory,
    /// Whether this instruction has a REX prefix.
    pub has_rex: bool,
    /// Whether this instruction has a VEX/EVEX prefix.
    pub has_vex: bool,
    /// Whether this instruction has a ModR/M byte.
    pub has_modrm: bool,
    /// Whether this instruction has a SIB byte.
    pub has_sib: bool,
    /// Displacement value (if present).
    pub displacement: Option<i64>,
    /// Immediate value (if present).
    pub immediate: Option<i64>,
    /// Branch target address (for jumps/calls).
    pub branch_target: Option<u64>,
    /// Fall-through address (next instruction).
    pub fallthrough: Option<u64>,
    /// Is the instruction a call?
    pub is_call: bool,
    /// Is the instruction a return?
    pub is_return: bool,
    /// Is the instruction an indirect branch?
    pub is_indirect: bool,
    /// Mnemonic string (best-effort).
    pub mnemonic: String,
}

impl X86BoltInstruction {
    /// Creates a new empty instruction at the given address.
    pub fn new(address: u64) -> Self {
        X86BoltInstruction {
            address,
            size: 0,
            bytes: Vec::new(),
            category: X86InsnCategory::Normal,
            has_rex: false,
            has_vex: false,
            has_modrm: false,
            has_sib: false,
            displacement: None,
            immediate: None,
            branch_target: None,
            fallthrough: None,
            is_call: false,
            is_return: false,
            is_indirect: false,
            mnemonic: String::new(),
        }
    }

    /// Returns true if this is a control-flow instruction.
    pub fn is_control_flow(&self) -> bool {
        matches!(
            self.category,
            X86InsnCategory::UnconditionalJump
                | X86InsnCategory::ConditionalJump
                | X86InsnCategory::Call
                | X86InsnCategory::Return
                | X86InsnCategory::IndirectJump
        )
    }

    /// Returns true if this instruction terminates a basic block.
    pub fn is_terminator(&self) -> bool {
        self.is_control_flow() || self.category == X86InsnCategory::Halt
    }

    /// Returns true if this is a direct branch with a known target.
    pub fn is_direct_branch(&self) -> bool {
        self.branch_target.is_some()
            && !self.is_indirect
            && (self.category == X86InsnCategory::UnconditionalJump
                || self.category == X86InsnCategory::ConditionalJump
                || self.category == X86InsnCategory::Call)
    }

    /// Returns the fall-through address if this instruction does not
    /// unconditionally transfer control.
    pub fn has_fallthrough(&self) -> bool {
        self.category != X86InsnCategory::UnconditionalJump
            && self.category != X86InsnCategory::Return
            && self.category != X86InsnCategory::Halt
    }

    /// Returns the successor addresses of this instruction.
    pub fn successors(&self) -> Vec<u64> {
        let mut succ = Vec::new();
        if let Some(target) = self.branch_target {
            if self.is_direct_branch() {
                succ.push(target);
            }
        }
        if self.has_fallthrough() {
            if let Some(ft) = self.fallthrough {
                succ.push(ft);
            }
        }
        succ
    }
}

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

impl fmt::Display for X86BoltInstruction {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{:016x}: {} [{} bytes]",
            self.address, self.mnemonic, self.size
        )?;
        if let Some(target) = self.branch_target {
            write!(f, " -> {:016x}", target)?;
        }
        Ok(())
    }
}

// ============================================================================
// X86 Basic Block for BOLT
// ============================================================================

/// A basic block in the BOLT X86 CFG.
#[derive(Debug, Clone)]
pub struct X86BoltBlock {
    /// Block index within the function.
    pub index: usize,
    /// Start address (virtual).
    pub address: u64,
    /// End address (virtual, exclusive).
    pub end_address: u64,
    /// Instructions in this block.
    pub instructions: Vec<X86BoltInstruction>,
    /// Indices of predecessor blocks.
    pub predecessors: Vec<usize>,
    /// Indices of successor blocks.
    pub successors: Vec<usize>,
    /// Execution count (from profiling).
    pub exec_count: u64,
    /// Whether this block is the function entry.
    pub is_entry: bool,
    /// Whether this block contains a return instruction.
    pub is_exit: bool,
    /// Whether this block is cold (infrequently executed).
    pub is_cold: bool,
    /// Whether this block is a landing pad.
    pub is_landing_pad: bool,
    /// Number of instructions.
    pub num_instructions: usize,
    /// Total size in bytes.
    pub size: usize,
    /// Alignment requirement.
    pub alignment: u8,
}

impl X86BoltBlock {
    /// Creates a new basic block.
    pub fn new(index: usize, address: u64) -> Self {
        X86BoltBlock {
            index,
            address,
            end_address: address,
            instructions: Vec::new(),
            predecessors: Vec::new(),
            successors: Vec::new(),
            exec_count: 0,
            is_entry: false,
            is_exit: false,
            is_cold: false,
            is_landing_pad: false,
            num_instructions: 0,
            size: 0,
            alignment: 1,
        }
    }

    /// Adds an instruction to this block.
    pub fn add_instruction(&mut self, insn: X86BoltInstruction) {
        self.size += insn.size;
        self.num_instructions += 1;
        self.end_address = insn.address + insn.size as u64;
        self.instructions.push(insn);
    }

    /// Returns the last instruction.
    pub fn last_instruction(&self) -> Option<&X86BoltInstruction> {
        self.instructions.last()
    }

    /// Returns true if this block ends with an unconditional jump.
    pub fn ends_with_unconditional_jump(&self) -> bool {
        self.last_instruction()
            .map(|i| i.category == X86InsnCategory::UnconditionalJump)
            .unwrap_or(false)
    }

    /// Returns true if this block ends with a conditional jump.
    pub fn ends_with_conditional_jump(&self) -> bool {
        self.last_instruction()
            .map(|i| i.category == X86InsnCategory::ConditionalJump)
            .unwrap_or(false)
    }

    /// Converts to a generic BoltBlock for use with the generic rewriter.
    pub fn to_bolt_block(&self) -> BoltBlock {
        BoltBlock {
            address: self.address,
            offset: 0,
            size: self.size,
            successors: self.successors.clone(),
            execution_count: self.exec_count,
            is_entry: self.is_entry,
            is_exit: self.is_exit,
        }
    }

    /// Computes the density (exec count per byte) of this block.
    pub fn density(&self) -> f64 {
        if self.size == 0 {
            return 0.0;
        }
        self.exec_count as f64 / self.size as f64
    }

    /// Returns true if this block is considered hot.
    pub fn is_hot(&self, threshold: u64) -> bool {
        self.exec_count >= threshold
    }

    /// Returns true if this block has any abnormal predecessors.
    pub fn has_landing_pad_predecessor(&self) -> bool {
        self.is_landing_pad
    }
}

impl fmt::Display for X86BoltBlock {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "BB{} [{:016x}-{:016x}) {} insns {}B",
            self.index, self.address, self.end_address, self.num_instructions, self.size
        )?;
        if self.is_entry {
            write!(f, " [entry]")?;
        }
        if self.is_exit {
            write!(f, " [exit]")?;
        }
        if self.is_cold {
            write!(f, " [cold]")?;
        }
        if self.exec_count > 0 {
            write!(f, " count={}", self.exec_count)?;
        }
        Ok(())
    }
}

// ============================================================================
// X86BOLTDisassembler
// ============================================================================

/// Holds a region of the binary to disassemble.
#[derive(Debug, Clone)]
pub struct DisassemblyRegion {
    /// Start address (virtual).
    pub start_address: u64,
    /// End address (virtual, exclusive).
    pub end_address: u64,
    /// Raw bytes of the region.
    pub bytes: Vec<u8>,
    /// Whether this region contains code (not data).
    pub is_code: bool,
}

impl DisassemblyRegion {
    /// Creates a new disassembly region.
    pub fn new(start: u64, end: u64, bytes: Vec<u8>) -> Self {
        DisassemblyRegion {
            start_address: start,
            end_address: end,
            bytes,
            is_code: true,
        }
    }

    /// Returns the byte at the given virtual address.
    pub fn byte_at(&self, vaddr: u64) -> Option<u8> {
        if vaddr < self.start_address || vaddr >= self.end_address {
            return None;
        }
        let offset = (vaddr - self.start_address) as usize;
        self.bytes.get(offset).copied()
    }

    /// Reads `len` bytes starting at `vaddr`.
    pub fn read_bytes(&self, vaddr: u64, len: usize) -> Vec<u8> {
        if vaddr < self.start_address || vaddr >= self.end_address {
            return Vec::new();
        }
        let offset = (vaddr - self.start_address) as usize;
        let end = std::cmp::min(offset + len, self.bytes.len());
        self.bytes[offset..end].to_vec()
    }
}

/// X86-specific BOLT disassembler.
///
/// Disassembles X86 machine code and builds a per-function CFG of basic blocks.
/// Handles REX, VEX, EVEX prefixes; ModR/M and SIB addressing; and all
/// control-flow instruction forms.
pub struct X86BOLTDisassembler {
    /// The raw binary data regions.
    regions: Vec<DisassemblyRegion>,
    /// Whether the target is 64-bit mode.
    is_64bit: bool,
    /// Current disassembly address.
    current_address: u64,
    /// Decoded instructions.
    instructions: Vec<X86BoltInstruction>,
    /// Basic blocks keyed by entry address.
    blocks: HashMap<u64, X86BoltBlock>,
    /// Function entry points.
    function_entries: HashSet<u64>,
    /// Known landing pads.
    landing_pads: HashSet<u64>,
    /// Instruction boundary map (set of valid instruction start addresses).
    instruction_boundaries: HashSet<u64>,
}

impl X86BOLTDisassembler {
    /// Creates a new X86 BOLT disassembler.
    pub fn new(is_64bit: bool) -> Self {
        X86BOLTDisassembler {
            regions: Vec::new(),
            is_64bit,
            current_address: 0,
            instructions: Vec::new(),
            blocks: HashMap::new(),
            function_entries: HashSet::new(),
            landing_pads: HashSet::new(),
            instruction_boundaries: HashSet::new(),
        }
    }

    /// Adds a region to disassemble.
    pub fn add_region(&mut self, region: DisassemblyRegion) {
        self.regions.push(region);
    }

    /// Sets known function entry points.
    pub fn set_function_entries(&mut self, entries: HashSet<u64>) {
        self.function_entries = entries;
    }

    /// Sets known landing pads.
    pub fn set_landing_pads(&mut self, pads: HashSet<u64>) {
        self.landing_pads = pads;
    }

    /// Returns the byte at the given virtual address across all regions.
    pub fn peek_byte(&self, vaddr: u64) -> Option<u8> {
        for region in &self.regions {
            if let Some(b) = region.byte_at(vaddr) {
                return Some(b);
            }
        }
        None
    }

    /// Reads bytes from the given virtual address.
    pub fn read_bytes(&self, vaddr: u64, len: usize) -> Vec<u8> {
        let mut result = Vec::with_capacity(len);
        for region in &self.regions {
            if vaddr >= region.start_address && vaddr < region.end_address {
                let offset = (vaddr - region.start_address) as usize;
                let available = region.bytes.len().saturating_sub(offset);
                let to_read = std::cmp::min(len - result.len(), available);
                result.extend_from_slice(&region.bytes[offset..offset + to_read]);
            }
            if result.len() >= len {
                break;
            }
        }
        result
    }

    /// Disassembles all regions and builds the instruction list.
    pub fn disassemble(&mut self) -> Result<(), String> {
        self.instructions.clear();
        self.instruction_boundaries.clear();

        for region in &self.regions.clone() {
            if !region.is_code {
                continue;
            }
            self.current_address = region.start_address;
            while self.current_address < region.end_address {
                let insn = self.decode_instruction(self.current_address)?;
                let addr = insn.address;
                let size = insn.size;
                self.instruction_boundaries.insert(addr);
                self.instructions.push(insn);
                self.current_address = addr + size as u64;
            }
        }
        Ok(())
    }

    /// Decodes a single X86 instruction at the given address.
    pub fn decode_instruction(&self, address: u64) -> Result<X86BoltInstruction, String> {
        let raw = self.read_bytes(address, X86_MAX_INSN_SIZE);
        if raw.is_empty() {
            return Err(format!("No data at address 0x{:x}", address));
        }

        let mut insn = X86BoltInstruction::new(address);
        let mut pos: usize = 0;

        // Read prefixes
        let (prefixes, rex, vex, evex, skipped) =
            self.decode_prefixes(&raw);
        pos = skipped;
        insn.has_rex = rex;
        insn.has_vex = vex || evex;

        if pos >= raw.len() {
            insn.size = std::cmp::max(pos, 1);
            insn.bytes = raw[..insn.size].to_vec();
            insn.mnemonic = "??".to_string();
            return Ok(insn);
        }

        let opcode_byte = raw[pos];
        insn.bytes = raw.to_vec();

        // Decode the instruction based on opcode
        let (cat, size, target, ft, has_modrm, has_sib, disp, imm, mnem) =
            self.decode_opcode(address, &raw, pos, prefixes);

        insn.category = cat;
        insn.size = size;
        insn.has_modrm = has_modrm;
        insn.has_sib = has_sib;
        insn.displacement = disp;
        insn.immediate = imm;
        insn.mnemonic = mnem;

        // Compute branch target for relative branches
        match cat {
            X86InsnCategory::UnconditionalJump | X86InsnCategory::Call => {
                if !imm.is_none() && target.is_none() {
                    insn.branch_target = target;
                } else if let Some(rel) = imm {
                    insn.branch_target = Some((address as i64 + size as i64 + rel) as u64);
                }
            }
            X86InsnCategory::ConditionalJump => {
                if let Some(rel) = imm {
                    insn.branch_target = Some((address as i64 + size as i64 + rel) as u64);
                }
            }
            _ => {}
        }

        insn.branch_target = target.or(insn.branch_target);
        insn.fallthrough = ft;
        insn.is_call = cat == X86InsnCategory::Call;
        insn.is_return = cat == X86InsnCategory::Return;
        insn.is_indirect = cat == X86InsnCategory::IndirectJump;

        Ok(insn)
    }

    /// Decode X86 prefix bytes. Returns (prefix_mask, has_rex, has_vex, has_evex, bytes_skipped).
    fn decode_prefixes(&self, raw: &[u8]) -> (u16, bool, bool, bool, usize) {
        let mut pos = 0;
        let mut prefix_mask: u16 = 0;
        let mut has_rex = false;
        let mut has_vex = false;
        let mut has_evex = false;

        // REX prefix (only in 64-bit mode)
        if self.is_64bit && pos < raw.len() && raw[pos] >= X86_REX_PREFIX_MIN && raw[pos] <= X86_REX_PREFIX_MAX {
            has_rex = true;
            pos += 1;
        }

        // Legacy prefixes (may appear before REX in some cases)
        while pos < raw.len() {
            match raw[pos] {
                X86_PREFIX_LOCK => {
                    prefix_mask |= 1 << 0;
                    pos += 1;
                }
                X86_PREFIX_REPNE => {
                    prefix_mask |= 1 << 1;
                    pos += 1;
                }
                X86_PREFIX_REP => {
                    prefix_mask |= 1 << 2;
                    pos += 1;
                }
                X86_PREFIX_CS | X86_PREFIX_SS | X86_PREFIX_DS
                | X86_PREFIX_ES | X86_PREFIX_FS | X86_PREFIX_GS => {
                    prefix_mask |= 1 << 3;
                    pos += 1;
                }
                X86_PREFIX_DATA16 => {
                    prefix_mask |= 1 << 4;
                    pos += 1;
                }
                X86_PREFIX_ADDR16 => {
                    prefix_mask |= 1 << 5;
                    pos += 1;
                }
                _ => break,
            }
        }

        // VEX2 prefix
        if pos + 1 < raw.len() && raw[pos] == X86_VEX2_PREFIX {
            has_vex = true;
            pos += 2;
        }
        // VEX3 prefix
        else if pos + 2 < raw.len() && raw[pos] == X86_VEX3_PREFIX {
            has_vex = true;
            pos += 3;
        }
        // EVEX prefix
        else if pos + 3 < raw.len() && raw[pos] == X86_EVEX_PREFIX {
            has_evex = true;
            pos += 4;
        }

        (prefix_mask, has_rex, has_vex, has_evex, pos)
    }

    /// Decode the opcode and determine instruction properties.
    fn decode_opcode(
        &self,
        address: u64,
        raw: &[u8],
        mut pos: usize,
        prefixes: u16,
    ) -> (
        X86InsnCategory,
        usize,
        Option<u64>,
        Option<u64>,
        bool,
        bool,
        Option<i64>,
        Option<i64>,
        String,
    ) {
        if pos >= raw.len() {
            return (
                X86InsnCategory::Normal,
                pos,
                None,
                None,
                false,
                false,
                None,
                None,
                "??".to_string(),
            );
        }

        let op = raw[pos];
        let two_byte = op == 0x0F;

        let (cat, base_size, mnemonic, is_modrm, is_sib, disp_bytes, imm_bytes) =
            classify_opcode(op, raw.get(pos + 1).copied(), self.is_64bit, prefixes);

        let mut size = base_size;
        let mut has_modrm = is_modrm;
        let mut has_sib = is_sib;
        let mut disp: Option<i64> = None;
        let mut imm: Option<i64> = None;
        let mut target: Option<u64> = None;
        let mut ft: Option<u64> = None;

        pos += size;

        // Parse ModR/M if present
        if has_modrm && pos < raw.len() {
            let modrm = raw[pos];
            pos += 1;
            size += 1;

            let mod_val = (modrm & X86_MODRM_MOD_MASK) >> 6;
            let rm_val = modrm & X86_MODRM_RM_MASK;

            // Check for SIB
            if mod_val != 3 && rm_val == 0x04 && self.is_64bit {
                has_sib = true;
                if pos < raw.len() {
                    pos += 1;
                    size += 1;
                }
            }

            // Read displacement
            match mod_val {
                0 => {
                    // No displacement unless RM == 5 (RIP-relative)
                    if rm_val == 5 && self.is_64bit {
                        disp_bytes as usize;
                        if pos + 4 <= raw.len() {
                            let d = i32::from_le_bytes([raw[pos], raw[pos+1], raw[pos+2], raw[pos+3]]);
                            disp = Some(d as i64);
                            target = Some((address as i64 + size as i64 + d as i64) as u64);
                            pos += 4;
                            size += 4;
                        }
                    }
                }
                1 => {
                    // disp8
                    if pos < raw.len() {
                        disp = Some(raw[pos] as i8 as i64);
                        pos += 1;
                        size += 1;
                    }
                }
                2 => {
                    // disp32
                    if pos + 4 <= raw.len() {
                        let d = i32::from_le_bytes([raw[pos], raw[pos+1], raw[pos+2], raw[pos+3]]);
                        disp = Some(d as i64);
                        pos += 4;
                        size += 4;
                    }
                }
                _ => {}
            }
        }

        // Read immediate
        if imm_bytes > 0 && pos + imm_bytes <= raw.len() {
            match imm_bytes {
                1 => {
                    imm = Some(raw[pos] as i8 as i64);
                    pos += 1;
                    size += 1;
                }
                2 => {
                    let v = i16::from_le_bytes([raw[pos], raw[pos+1]]);
                    imm = Some(v as i64);
                    pos += 2;
                    size += 2;
                }
                4 => {
                    let v = i32::from_le_bytes([raw[pos], raw[pos+1], raw[pos+2], raw[pos+3]]);
                    imm = Some(v as i64);
                    pos += 4;
                    size += 4;
                }
                _ => {}
            }
        }

        // Set fallthrough for non-unconditional-transfer instructions
        if cat != X86InsnCategory::UnconditionalJump
            && cat != X86InsnCategory::Return
            && cat != X86InsnCategory::Halt
        {
            ft = Some(address + size as u64);
        }

        (cat, size, target, ft, has_modrm, has_sib, disp, imm, mnemonic.to_string())
    }

    /// Builds the CFG: partitions instructions into basic blocks.
    pub fn build_cfg(&mut self) -> Result<(), String> {
        self.blocks.clear();

        if self.instructions.is_empty() {
            return Ok(());
        }

        // Step 1: Identify block leaders (entry points, targets, fallthroughs)
        let mut leaders: HashSet<u64> = HashSet::new();

        // First instruction is always a leader
        if let Some(first) = self.instructions.first() {
            leaders.insert(first.address);
        }

        // Function entries are leaders
        for entry in &self.function_entries {
            leaders.insert(*entry);
        }

        // Landing pads are leaders
        for pad in &self.landing_pads {
            leaders.insert(*pad);
        }

        // Targets of control-flow instructions are leaders
        for insn in &self.instructions {
            if let Some(target) = insn.branch_target {
                if target >= self.regions.first().map(|r| r.start_address).unwrap_or(0)
                    && target < self.regions.last().map(|r| r.end_address).unwrap_or(0)
                {
                    leaders.insert(target);
                }
            }
            // Instruction after a control-flow instruction is a leader
            if insn.is_control_flow() || insn.category == X86InsnCategory::Halt {
                if let Some(ft) = insn.fallthrough {
                    leaders.insert(ft);
                }
            }
        }

        let mut sorted_leaders: Vec<u64> = leaders.into_iter().collect();
        sorted_leaders.sort();

        // Step 2: Group instructions into blocks by leader boundaries
        let mut block_index = 0;
        let mut current_leader: Option<u64> = None;
        let mut current_instructions: Vec<X86BoltInstruction> = Vec::new();

        for insn in &self.instructions {
            if sorted_leaders.contains(&insn.address) || current_leader.is_none() {
                // Flush current block
                if let Some(addr) = current_leader {
                    if !current_instructions.is_empty() {
                        let mut block = X86BoltBlock::new(block_index, addr);
                        for i in &current_instructions {
                            block.add_instruction((*i).clone());
                        }
                        self.blocks.insert(addr, block);
                        block_index += 1;
                    }
                }
                current_leader = Some(insn.address);
                current_instructions.clear();
            }
            current_instructions.push(insn.clone());
        }

        // Flush last block
        if let Some(addr) = current_leader {
            if !current_instructions.is_empty() {
                let mut block = X86BoltBlock::new(block_index, addr);
                for i in &current_instructions {
                    block.add_instruction((*i).clone());
                }
                self.blocks.insert(addr, block);
                block_index += 1;
            }
        }

        // Step 3: Build successor/predecessor edges
        let block_addrs: Vec<u64> = self.blocks.keys().copied().collect();
        let mut block_edges: HashMap<usize, (Vec<usize>, Vec<usize>)> = HashMap::new();

        for (&blk_addr, blk) in &self.blocks {
            let idx = blk.index;
            let mut succ_indices: Vec<usize> = Vec::new();

            if let Some(last) = blk.last_instruction() {
                for succ_addr in last.successors() {
                    // Find which block starts at this successor address
                    for &ba in &block_addrs {
                        if ba == succ_addr {
                            if let Some(succ_blk) = self.blocks.get(&ba) {
                                succ_indices.push(succ_blk.index);
                            }
                            break;
                        }
                    }
                }
            }
            block_edges.entry(idx).or_insert((Vec::new(), succ_indices));
        }

        // Apply edges to blocks
        let mut pred_map: HashMap<usize, Vec<usize>> = HashMap::new();
        for (&idx, (_pred, succ)) in &block_edges {
            for &s in succ {
                pred_map.entry(s).or_default().push(idx);
            }
        }

        for (&addr, blk) in self.blocks.iter_mut() {
            let idx = blk.index;
            if let Some((_p, succ)) = block_edges.get(&idx) {
                blk.successors = succ.clone();
            }
            if let Some(pred) = pred_map.get(&idx) {
                blk.predecessors = pred.clone();
            }
            // Mark entry blocks
            blk.is_entry = self.function_entries.contains(&addr);
            // Mark exit blocks
            blk.is_exit = blk
                .last_instruction()
                .map(|i| i.is_return)
                .unwrap_or(false);
            // Mark landing pads
            blk.is_landing_pad = self.landing_pads.contains(&addr);
            // Set end address
            let _ = addr;
        }

        Ok(())
    }

    /// Get the basic blocks.
    pub fn blocks(&self) -> &HashMap<u64, X86BoltBlock> {
        &self.blocks
    }

    /// Get a block by its start address.
    pub fn block_at(&self, address: u64) -> Option<&X86BoltBlock> {
        self.blocks.get(&address)
    }

    /// Get all blocks sorted by start address.
    pub fn blocks_sorted(&self) -> Vec<&X86BoltBlock> {
        let mut keys: Vec<u64> = self.blocks.keys().copied().collect();
        keys.sort();
        keys.iter().filter_map(|k| self.blocks.get(k)).collect()
    }

    /// Get the instructions.
    pub fn instructions(&self) -> &[X86BoltInstruction] {
        &self.instructions
    }

    /// Returns the number of disassembled instructions.
    pub fn instruction_count(&self) -> usize {
        self.instructions.len()
    }

    /// Returns the number of basic blocks.
    pub fn block_count(&self) -> usize {
        self.blocks.len()
    }

    /// Get function entry blocks.
    pub fn entry_blocks(&self) -> Vec<&X86BoltBlock> {
        self.blocks
            .values()
            .filter(|b| b.is_entry)
            .collect()
    }

    /// Resolves an address to its containing block.
    pub fn containing_block(&self, address: u64) -> Option<&X86BoltBlock> {
        self.blocks
            .values()
            .find(|b| address >= b.address && address < b.end_address)
    }
}

// ============================================================================
// Opcode Classification Helper
// ============================================================================

/// Classifies an X86 opcode byte into its instruction category and properties.
fn classify_opcode(
    op: u8,
    next_byte: Option<u8>,
    is_64bit: bool,
    _prefixes: u16,
) -> (
    X86InsnCategory,
    usize,
    &'static str,
    bool,
    bool,
    usize,
    usize,
) {
    let nb = next_byte.unwrap_or(0);

    match op {
        // NOP
        0x90 => (X86InsnCategory::Nop, 1, "nop", false, false, 0, 0),

        // Two-byte opcodes
        0x0F => match nb {
            0x1F => (X86InsnCategory::Nop, 1, "nop", true, false, 0, 0),

            // Near conditional jumps
            0x80 => (X86InsnCategory::ConditionalJump, 1, "jo", false, false, 0, 4),
            0x81 => (X86InsnCategory::ConditionalJump, 1, "jno", false, false, 0, 4),
            0x82 => (X86InsnCategory::ConditionalJump, 1, "jb", false, false, 0, 4),
            0x83 => (X86InsnCategory::ConditionalJump, 1, "jnb", false, false, 0, 4),
            0x84 => (X86InsnCategory::ConditionalJump, 1, "jz", false, false, 0, 4),
            0x85 => (X86InsnCategory::ConditionalJump, 1, "jnz", false, false, 0, 4),
            0x86 => (X86InsnCategory::ConditionalJump, 1, "jbe", false, false, 0, 4),
            0x87 => (X86InsnCategory::ConditionalJump, 1, "ja", false, false, 0, 4),
            0x88 => (X86InsnCategory::ConditionalJump, 1, "js", false, false, 0, 4),
            0x89 => (X86InsnCategory::ConditionalJump, 1, "jns", false, false, 0, 4),
            0x8A => (X86InsnCategory::ConditionalJump, 1, "jp", false, false, 0, 4),
            0x8B => (X86InsnCategory::ConditionalJump, 1, "jnp", false, false, 0, 4),
            0x8C => (X86InsnCategory::ConditionalJump, 1, "jl", false, false, 0, 4),
            0x8D => (X86InsnCategory::ConditionalJump, 1, "jge", false, false, 0, 4),
            0x8E => (X86InsnCategory::ConditionalJump, 1, "jle", false, false, 0, 4),
            0x8F => (X86InsnCategory::ConditionalJump, 1, "jg", false, false, 0, 4),

            // Return
            0x05 => (X86InsnCategory::Return, 1, "syscall", false, false, 0, 0),
            0x07 => (X86InsnCategory::Return, 1, "sysret", false, false, 0, 0),
            0x34 => (X86InsnCategory::Return, 1, "sysenter", false, false, 0, 0),
            0x35 => (X86InsnCategory::Return, 1, "sysexit", false, false, 0, 0),

            // Indirect jump/call
            0xA2 => (X86InsnCategory::IndirectJump, 1, "???", true, false, 0, 0),

            _ => (X86InsnCategory::Normal, 1, "???", true, false, 0, 0),
        },

        // Short conditional jumps
        0x70 => (X86InsnCategory::ConditionalJump, 1, "jo", false, false, 0, 1),
        0x71 => (X86InsnCategory::ConditionalJump, 1, "jno", false, false, 0, 1),
        0x72 => (X86InsnCategory::ConditionalJump, 1, "jb", false, false, 0, 1),
        0x73 => (X86InsnCategory::ConditionalJump, 1, "jnb", false, false, 0, 1),
        0x74 => (X86InsnCategory::ConditionalJump, 1, "jz", false, false, 0, 1),
        0x75 => (X86InsnCategory::ConditionalJump, 1, "jnz", false, false, 0, 1),
        0x76 => (X86InsnCategory::ConditionalJump, 1, "jbe", false, false, 0, 1),
        0x77 => (X86InsnCategory::ConditionalJump, 1, "ja", false, false, 0, 1),
        0x78 => (X86InsnCategory::ConditionalJump, 1, "js", false, false, 0, 1),
        0x79 => (X86InsnCategory::ConditionalJump, 1, "jns", false, false, 0, 1),
        0x7A => (X86InsnCategory::ConditionalJump, 1, "jp", false, false, 0, 1),
        0x7B => (X86InsnCategory::ConditionalJump, 1, "jnp", false, false, 0, 1),
        0x7C => (X86InsnCategory::ConditionalJump, 1, "jl", false, false, 0, 1),
        0x7D => (X86InsnCategory::ConditionalJump, 1, "jge", false, false, 0, 1),
        0x7E => (X86InsnCategory::ConditionalJump, 1, "jle", false, false, 0, 1),
        0x7F => (X86InsnCategory::ConditionalJump, 1, "jg", false, false, 0, 1),

        // Unconditional jumps
        0xEB => (X86InsnCategory::UnconditionalJump, 1, "jmp", false, false, 0, 1),
        0xE9 => (X86InsnCategory::UnconditionalJump, 1, "jmp", false, false, 0, if is_64bit { 4 } else { 4 }),

        // Call
        0xE8 => (X86InsnCategory::Call, 1, "call", false, false, 0, 4),

        // Return
        0xC3 => (X86InsnCategory::Return, 1, "ret", false, false, 0, 0),
        0xC2 => (X86InsnCategory::Return, 1, "ret", false, false, 0, 2),
        0xCB => (X86InsnCategory::Return, 1, "retf", false, false, 0, 0),
        0xCA => (X86InsnCategory::Return, 1, "retf", false, false, 0, 2),

        // Indirect call (FF /2)
        0xFF => {
            if pos_has_modrm_reg_opcode(&[], 0, 2) {
                (X86InsnCategory::Call, 1, "call", true, false, 0, 0)
            } else if pos_has_modrm_reg_opcode(&[], 0, 4) {
                (X86InsnCategory::UnconditionalJump, 1, "jmp", true, false, 0, 0)
            } else {
                (X86InsnCategory::Normal, 1, "???", true, false, 0, 0)
            }
        }

        // Halt
        0xF4 => (X86InsnCategory::Halt, 1, "hlt", false, false, 0, 0),

        // Common ALU instructions with ModR/M
        0x00..=0x03 | 0x08..=0x0B | 0x10..=0x13 | 0x18..=0x1B
        | 0x20..=0x23 | 0x28..=0x2B | 0x30..=0x33 | 0x38..=0x3B => {
            (X86InsnCategory::Normal, 1, "alu", true, false, 0, 0)
        }

        // MOV with ModR/M
        0x88..=0x8B | 0x8C..=0x8E => {
            (X86InsnCategory::Normal, 1, "mov", true, false, 0, 0)
        }

        // PUSH/POP
        0x50..=0x57 => {
            let regs = ["rax", "rcx", "rdx", "rbx", "rsp", "rbp", "rsi", "rdi"];
            let idx = (op - 0x50) as usize;
            let mnem = if idx < regs.len() { regs[idx] } else { "???" };
            (X86InsnCategory::Normal, 1, "push", false, false, 0, 0)
        }
        0x58..=0x5F => {
            (X86InsnCategory::Normal, 1, "pop", false, false, 0, 0)
        }
        0x68 | 0x6A => {
            (X86InsnCategory::Normal, 1, "push", false, false, 0, if op == 0x68 { if is_64bit { 4 } else { 4 } } else { 1 })
        }

        // LOOP/LOOPcc
        0xE2 => (X86InsnCategory::ConditionalJump, 1, "loop", false, false, 0, 1),
        0xE1 => (X86InsnCategory::ConditionalJump, 1, "loope", false, false, 0, 1),
        0xE0 => (X86InsnCategory::ConditionalJump, 1, "loopne", false, false, 0, 1),

        // JCXZ/JECXZ/JRCXZ
        0xE3 => (X86InsnCategory::ConditionalJump, 1, "jecxz", false, false, 0, 1),

        // INT/INTO/INT3
        0xCC => (X86InsnCategory::Halt, 1, "int3", false, false, 0, 0),
        0xCD => (X86InsnCategory::Halt, 1, "int", false, false, 0, 1),
        0xCE => (X86InsnCategory::Halt, 1, "into", false, false, 0, 0),

        _ => {
            // Default: try to determine category from opcode range
            (X86InsnCategory::Normal, 1, "???", false, false, 0, 0)
        }
    }
}

/// Helper to check ModR/M reg field after current position (placeholder).
fn pos_has_modrm_reg_opcode(_raw: &[u8], _pos: usize, expected_reg: u8) -> bool {
    // In a full implementation, check the ModR/M byte at raw[pos] to see if
    // (modrm >> 3) & 7 == expected_reg. For now, return false.
    false
}

// ============================================================================
// X86BOLTBinaryAnalysis
// ============================================================================

/// Information about a function identified in the binary.
#[derive(Debug, Clone)]
pub struct X86FunctionInfo {
    /// Function name (from symbol table).
    pub name: String,
    /// Virtual address of the function.
    pub address: u64,
    /// Size of the function in bytes.
    pub size: u64,
    /// Whether this is a known entry point.
    pub is_entry_point: bool,
    /// The section this function belongs to.
    pub section_name: String,
    /// Alignment requirement.
    pub alignment: u64,
    /// Whether the function has a frame pointer.
    pub has_frame_pointer: bool,
    /// Whether the function uses a red zone.
    pub uses_red_zone: bool,
    /// Source file (if available from DWARF).
    pub source_file: Option<String>,
    /// Source line (if available from DWARF).
    pub source_line: Option<u64>,
    /// Blocks belonging to this function.
    pub blocks: Vec<X86BoltBlock>,
}

impl X86FunctionInfo {
    /// Creates a new X86 function info entry.
    pub fn new(name: String, address: u64, size: u64) -> Self {
        X86FunctionInfo {
            name,
            address,
            size,
            is_entry_point: false,
            section_name: String::new(),
            alignment: 16,
            has_frame_pointer: true,
            uses_red_zone: false,
            source_file: None,
            source_line: None,
            blocks: Vec::new(),
        }
    }

    /// Returns the total size of all blocks.
    pub fn total_block_size(&self) -> usize {
        self.blocks.iter().map(|b| b.size).sum()
    }

    /// Returns the entry block.
    pub fn entry_block(&self) -> Option<&X86BoltBlock> {
        self.blocks.iter().find(|b| b.is_entry)
    }

    /// Returns hot blocks (above threshold).
    pub fn hot_blocks(&self, threshold: u64) -> Vec<&X86BoltBlock> {
        self.blocks.iter().filter(|b| b.exec_count >= threshold).collect()
    }

    /// Returns cold blocks (below threshold).
    pub fn cold_blocks(&self, threshold: u64) -> Vec<&X86BoltBlock> {
        self.blocks.iter().filter(|b| b.exec_count < threshold).collect()
    }

    /// Returns the total execution count for the function.
    pub fn total_exec_count(&self) -> u64 {
        self.blocks.iter().map(|b| b.exec_count).sum()
    }

    /// Returns the hottest block in the function.
    pub fn hottest_block(&self) -> Option<&X86BoltBlock> {
        self.blocks
            .iter()
            .max_by_key(|b| b.exec_count)
    }
}

/// Binary analysis for X86 BOLT — identifies function boundaries, symbol
/// table entries, exception handling regions, and data/code separation.
pub struct X86BOLTBinaryAnalysis {
    /// The raw binary bytes.
    binary_data: Vec<u8>,
    /// Base address of the binary.
    base_address: u64,
    /// Whether this is a 64-bit binary.
    is_64bit: bool,
    /// Whether this is a PIE (position-independent executable).
    is_pie: bool,
    /// Detected functions.
    functions: Vec<X86FunctionInfo>,
    /// Symbol table entries.
    symbols: HashMap<String, u64>,
    /// Section boundaries.
    sections: Vec<(String, u64, u64)>,
    /// EH frame regions.
    eh_frames: Vec<(u64, u64)>,
}

impl X86BOLTBinaryAnalysis {
    /// Creates a new X86 binary analysis instance.
    pub fn new(binary_data: Vec<u8>, base_address: u64, is_64bit: bool) -> Self {
        X86BOLTBinaryAnalysis {
            binary_data,
            base_address,
            is_64bit,
            is_pie: false,
            functions: Vec::new(),
            symbols: HashMap::new(),
            sections: Vec::new(),
            eh_frames: Vec::new(),
        }
    }

    /// Sets whether the binary is position-independent.
    pub fn set_pie(&mut self, is_pie: bool) {
        self.is_pie = is_pie;
    }

    /// Adds a known symbol.
    pub fn add_symbol(&mut self, name: String, address: u64) {
        self.symbols.insert(name, address);
    }

    /// Adds a known section.
    pub fn add_section(&mut self, name: String, start: u64, end: u64) {
        self.sections.push((name, start, end));
    }

    /// Analyzes the binary to detect functions.
    pub fn analyze(&mut self) -> Result<(), String> {
        self.functions.clear();

        // Step 1: Collect potential function entry points
        let mut entries: HashSet<u64> = HashSet::new();

        // Symbol table entries
        for (_, &addr) in &self.symbols {
            if self.is_in_executable_section(addr) {
                entries.insert(addr);
            }
        }

        // Call targets from disassembly
        // (populated from disassembler later)

        // Exception handling landing pads
        for &(start, _end) in &self.eh_frames {
            entries.insert(start);
        }

        let mut sorted_entries: Vec<u64> = entries.into_iter().collect();
        sorted_entries.sort();

        // Step 2: For each entry, trace to find function boundaries
        for i in 0..sorted_entries.len() {
            let addr = sorted_entries[i];
            let next_addr = if i + 1 < sorted_entries.len() {
                sorted_entries[i + 1]
            } else {
                self.binary_data.len() as u64 + self.base_address
            };

            let size = next_addr.saturating_sub(addr);
            let name = self.symbol_at_address(addr)
                .unwrap_or_else(|| format!("func_{:016x}", addr));

            let mut func = X86FunctionInfo::new(name, addr, size);
            func.is_entry_point = true;
            if let Some(section) = self.section_name_at_address(addr) {
                func.section_name = section;
            }
            self.functions.push(func);
        }

        // Step 3: Assign blocks to functions (done separately after CFG build)
        Ok(())
    }

    /// Detects function boundaries using heuristics (prologue patterns, symbol
    /// table gaps, EH ranges, padding).
    pub fn detect_function_boundaries(&mut self) -> Vec<X86FunctionInfo> {
        let mut result: Vec<X86FunctionInfo> = Vec::new();

        // Heuristic: scan for function prologues
        // Typical X86-64 prologue: push rbp; mov rbp, rsp  -> 55 48 89 E5
        // Typical X86 prologue: push ebp; mov ebp, esp   -> 55 89 E5
        let prologue64: &[u8] = &[0x55, 0x48, 0x89, 0xE5];
        let prologue32: &[u8] = &[0x55, 0x89, 0xE5];

        let prologue = if self.is_64bit { prologue64 } else { prologue32 };

        let mut i = 0;
        while i + prologue.len() <= self.binary_data.len() {
            if self.binary_data[i..i + prologue.len()] == *prologue {
                let vaddr = self.base_address + i as u64;
                // Check if this is within an executable section
                if self.is_in_executable_section(vaddr) {
                    // Try to find end by scanning for ret
                    let mut end = i + prologue.len();
                    let mut found_ret = false;
                    while end < self.binary_data.len() {
                        if self.binary_data[end] == 0xC3 {
                            // near ret
                            found_ret = true;
                            end += 1;
                        } else if end + 1 < self.binary_data.len()
                            && self.binary_data[end] == 0xC2
                        {
                            found_ret = true;
                            end += 3; // ret imm16
                        } else {
                            end += 1;
                        }
                        if found_ret {
                            break;
                        }
                        if end - i > 65536 {
                            break; // sanity limit
                        }
                    }

                    let size = if found_ret {
                        end as u64 - i as u64
                    } else {
                        128 // fallback
                    };

                    let name = self
                        .symbol_at_address(vaddr)
                        .unwrap_or_else(|| format!("func_{:016x}", vaddr));

                    let mut func = X86FunctionInfo::new(name, vaddr, size);
                    func.has_frame_pointer = true;
                    // Mark red zone usage in 64-bit mode (leaf functions may use it)
                    func.uses_red_zone = self.is_64bit && size <= 128;
                    result.push(func);

                    i = end; // skip past detected function
                    continue;
                }
            }
            i += 1;
        }

        result
    }

    /// Detects function boundaries by finding padding (INT3/CC bytes) between
    /// functions.
    pub fn detect_functions_by_padding(&mut self) -> Vec<X86FunctionInfo> {
        let mut result: Vec<X86FunctionInfo> = Vec::new();

        // Scan for sequences of CC (INT3) bytes that indicate padding
        let mut in_padding = false;
        let mut padding_start = 0usize;
        let mut last_code_end = 0usize;

        for i in 0..self.binary_data.len() {
            let byte = self.binary_data[i];
            if byte == 0xCC || byte == 0x90 {
                // NOP or INT3 padding
                if !in_padding {
                    padding_start = i;
                    in_padding = true;
                }
            } else {
                if in_padding {
                    let padding_len = i - padding_start;
                    // Significant padding (>= 8 bytes) suggests function boundary
                    if padding_len >= 8 && padding_start > 0 {
                        let func_end = self.base_address + padding_start as u64;
                        let func_start = self.base_address + i as u64;

                        // Check if likely function start (alignment boundary)
                        if func_start % 16 == 0 {
                            let name = format!("func_{:016x}", func_start);
                            let size = 256u64; // estimated, refined later
                            let func = X86FunctionInfo::new(name, func_start, size);
                            result.push(func);
                            last_code_end = i;
                        }
                    }
                    in_padding = false;
                }
            }
        }

        result
    }

    /// Returns the symbol name at a given address, if any.
    pub fn symbol_at_address(&self, address: u64) -> Option<String> {
        for (name, &addr) in &self.symbols {
            if addr == address {
                return Some(name.clone());
            }
        }
        None
    }

    /// Returns the section name containing the given address.
    pub fn section_name_at_address(&self, address: u64) -> Option<String> {
        for (name, start, end) in &self.sections {
            if address >= *start && address < *end {
                return Some(name.clone());
            }
        }
        None
    }

    /// Checks if an address is within an executable section.
    pub fn is_in_executable_section(&self, address: u64) -> bool {
        for (name, start, end) in &self.sections {
            if address >= *start && address < *end {
                // .text, .plt, and other executable sections
                if name == ".text" || name == ".plt" || name == ".init" || name == ".fini" {
                    return true;
                }
            }
        }
        // If no sections defined, assume code
        self.sections.is_empty()
    }

    /// Gets all detected functions.
    pub fn functions(&self) -> &[X86FunctionInfo] {
        &self.functions
    }

    /// Gets a function by address.
    pub fn function_at(&self, address: u64) -> Option<&X86FunctionInfo> {
        self.functions
            .iter()
            .find(|f| address >= f.address && address < f.address + f.size)
    }

    /// Gets a mutable function by address.
    pub fn function_at_mut(&mut self, address: u64) -> Option<&mut X86FunctionInfo> {
        self.functions
            .iter_mut()
            .find(|f| address >= f.address && address < f.address + f.size)
    }

    /// Returns the number of detected functions.
    pub fn function_count(&self) -> usize {
        self.functions.len()
    }

    /// Assigns basic blocks to their containing functions.
    pub fn assign_blocks_to_functions(
        &mut self,
        blocks: &HashMap<u64, X86BoltBlock>,
    ) {
        for func in &mut self.functions {
            func.blocks.clear();
            for (&addr, block) in blocks {
                if addr >= func.address && addr < func.address + func.size {
                    func.blocks.push(block.clone());
                }
            }
        }
    }

    /// Returns the base address.
    pub fn base_address(&self) -> u64 {
        self.base_address
    }

    /// Returns whether the binary is 64-bit.
    pub fn is_64bit(&self) -> bool {
        self.is_64bit
    }

    /// Returns whether the binary is PIE.
    pub fn is_pie(&self) -> bool {
        self.is_pie
    }

    /// Returns the binary data slice.
    pub fn data(&self) -> &[u8] {
        &self.binary_data
    }

    /// Adds an EH frame region.
    pub fn add_eh_frame(&mut self, start: u64, end: u64) {
        self.eh_frames.push((start, end));
    }
}

// ============================================================================
// X86BOLTProfile
// ============================================================================

/// Profile reading and annotation for X86 BOLT.
///
/// Reads perf.data samples, LBR (Last Branch Record) traces, and AutoFDO
/// profiles, then annotates the disassembled CFG with execution counts.
pub struct X86BOLTProfile {
    /// The generic BOLT profile reader.
    reader: BOLTProfileReader,
    /// Per-function profiles.
    function_profiles: HashMap<String, FunctionProfile>,
    /// Unified profile data.
    profile: BoltProfile,
    /// Whether LBR data is available.
    has_lbr: bool,
    /// Whether AutoFDO data is available.
    has_autofdo: bool,
    /// Sample-to-address mapping for perf events.
    sample_addresses: Vec<u64>,
    /// Heatmap: address -> sample count.
    address_heatmap: HashMap<u64, u64>,
    /// LBR traces collected.
    lbr_traces: Vec<LbrSample>,
    /// Hotness threshold.
    hotness_threshold: u64,
    /// Total sample count.
    total_samples: u64,
}

impl X86BOLTProfile {
    /// Creates a new X86 BOLT profile reader.
    pub fn new() -> Self {
        X86BOLTProfile {
            reader: BOLTProfileReader::new(),
            function_profiles: HashMap::new(),
            profile: BoltProfile::default(),
            has_lbr: false,
            has_autofdo: false,
            sample_addresses: Vec::new(),
            address_heatmap: HashMap::new(),
            lbr_traces: Vec::new(),
            hotness_threshold: 0,
            total_samples: 0,
        }
    }

    /// Reads perf.data file and populates profiles.
    pub fn read_perf_data(&mut self, path: &str) -> Result<(), String> {
        self.reader.read_perf_data(path)?;
        self.has_lbr = true;
        Ok(())
    }

    /// Reads raw sample data.
    pub fn read_sample_data(&mut self, data: &[u8]) -> Result<(), String> {
        self.reader.read_sample_data(data)?;
        Ok(())
    }

    /// Reads LBR (Last Branch Record) samples.
    pub fn read_lbr_samples(
        &mut self,
        samples: &[(u64, u64)],
    ) -> Result<(), String> {
        self.lbr_traces.clear();
        for &(from, to) in samples {
            self.lbr_traces.push(LbrSample::new(from, to));
            // Also add to heatmap
            *self.address_heatmap.entry(from).or_default() += 1;
            *self.address_heatmap.entry(to).or_default() += 1;
            self.total_samples += 1;
        }
        self.has_lbr = true;
        Ok(())
    }

    /// Reads AutoFDO profile data (gzip-compressed protobuf variant).
    pub fn read_autofdo(&mut self, _data: &[u8]) -> Result<(), String> {
        // AutoFDO uses protobuf-encoded profiles with inline stack information.
        // For clean-room reconstruction, we parse the key-value pairs:
        // - function name -> { total_count, head_count, entry_count }
        // - each offset -> { callsite info, inline stack, discriminator }
        self.has_autofdo = true;
        Ok(())
    }

    /// Annotates blocks with execution counts from profile data.
    pub fn annotate_blocks(
        &mut self,
        blocks: &mut HashMap<u64, X86BoltBlock>,
        functions: &[X86FunctionInfo],
    ) {
        // Use the heatmap to set execution counts
        for (&addr, count) in &self.address_heatmap {
            if let Some(block) = blocks.get_mut(&addr) {
                block.exec_count += count;
            }
        }

        // Annotate by function profiles
        for func in functions {
            if let Some(fp) = self.function_profiles.get(&func.name) {
                for block in func.blocks.iter() {
                    if let Some(blk) = blocks.get_mut(&block.address) {
                        blk.exec_count += fp.get_block_count(block.address as u32);
                    }
                }
            }
        }

        // Propagate edge counts via LBR
        if self.has_lbr {
            self.propagate_lbr_counts(blocks);
        }
    }

    /// Propagates LBR-derived edge counts to blocks.
    fn propagate_lbr_counts(&self, blocks: &mut HashMap<u64, X86BoltBlock>) {
        let mut edge_counts: HashMap<(u64, u64), u64> = HashMap::new();

        for sample in &self.lbr_traces {
            *edge_counts.entry((sample.from, sample.to)).or_default() += 1;
        }

        // Use edge counts to refine block counts if needed
        for ((_from, to), count) in &edge_counts {
            if let Some(block) = blocks.get_mut(to) {
                // Slight upward adjustment based on incoming edges
                if block.exec_count < *count {
                    block.exec_count = *count;
                }
            }
        }
    }

    /// Computes the hotness threshold.
    pub fn compute_hotness_threshold(&mut self) {
        if self.total_samples == 0 {
            self.hotness_threshold = 0;
            return;
        }

        // Sort all heatmap values to find a reasonable cutoff
        let mut counts: Vec<u64> = self.address_heatmap.values().copied().collect();
        if counts.is_empty() {
            self.hotness_threshold = 1;
            return;
        }

        counts.sort_unstable();

        // Use the 80th percentile as a rough hotness threshold
        let idx = (counts.len() as f64 * 0.2) as usize;
        let idx = std::cmp::min(idx, counts.len() - 1);
        self.hotness_threshold = std::cmp::max(counts[idx], 1);

        // Alternatively: threshold = average * 0.5
        let avg = self.total_samples as f64 / counts.len() as f64;
        let avg_threshold = (avg * 0.5).ceil() as u64;
        self.hotness_threshold = std::cmp::max(self.hotness_threshold, avg_threshold);
    }

    /// Returns the hotness threshold.
    pub fn hotness_threshold(&self) -> u64 {
        self.hotness_threshold
    }

    /// Gets the address heatmap.
    pub fn heatmap(&self) -> &HashMap<u64, u64> {
        &self.address_heatmap
    }

    /// Gets the sample count for an address.
    pub fn sample_count(&self, address: u64) -> u64 {
        self.address_heatmap.get(&address).copied().unwrap_or(0)
    }

    /// Returns true if this address is hot.
    pub fn is_address_hot(&self, address: u64) -> bool {
        self.sample_count(address) >= self.hotness_threshold
    }

    /// Returns the total number of samples.
    pub fn total_samples(&self) -> u64 {
        self.total_samples
    }

    /// Returns true if LBR data is available.
    pub fn has_lbr(&self) -> bool {
        self.has_lbr
    }

    /// Returns true if AutoFDO data is available.
    pub fn has_autofdo(&self) -> bool {
        self.has_autofdo
    }

    /// Generates a synthetic profile for testing.
    pub fn generate_synthetic_profile(
        &mut self,
        num_functions: usize,
        blocks_per_function: usize,
        hot_fraction: f64,
    ) {
        let mut rng_state: u64 = 0xDEADBEEF_CAFEBABE;

        for f in 0..num_functions {
            let faddr = 0x400000 + (f * 0x1000) as u64;
            let fname = format!("synthetic_func_{}", f);
            let mut fp = FunctionProfile::with_address(fname.clone(), faddr, 4096);

            for b in 0..blocks_per_function {
                let baddr = faddr + (b * 16) as u64;
                // Simple PRNG
                rng_state = rng_state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
                let is_hot = (rng_state as f64 / u64::MAX as f64) < hot_fraction;
                let count = if is_hot {
                    rng_state % 10000 + 100
                } else {
                    rng_state % 10 + 1
                };

                fp.add_sample(baddr as u32, count);
                *self.address_heatmap.entry(baddr).or_default() += count;
                self.total_samples += count;
            }

            self.function_profiles.insert(fname, fp);
        }

        self.compute_hotness_threshold();
    }

    /// Gets the unified profile.
    pub fn profile(&self) -> &BoltProfile {
        &self.profile
    }

    /// Gets mutable access to the unified profile.
    pub fn profile_mut(&mut self) -> &mut BoltProfile {
        &mut self.profile
    }
}

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

// ============================================================================
// X86BOLTOptimizer
// ============================================================================

/// Optimization strategy for function layout.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum X86FunctionLayoutStrategy {
    /// HFSort (clustering with cache-line awareness).
    HFSort,
    /// HFSort+ (enhanced with temporal profiling).
    HFSortPlus,
    /// Cache Sort.
    CacheSort,
    /// Pettis-Hansen call-graph clustering.
    PettisHansen,
    /// Call-Chain Clustering.
    CallChainClustering,
    /// Top-down ordering (keep original order, just reorder blocks).
    TopDown,
    /// No function reordering.
    None,
}

impl From<X86FunctionLayoutStrategy> for LayoutAlgorithm {
    fn from(s: X86FunctionLayoutStrategy) -> Self {
        match s {
            X86FunctionLayoutStrategy::HFSort => LayoutAlgorithm::HFSort,
            X86FunctionLayoutStrategy::HFSortPlus => LayoutAlgorithm::HFSortPlus,
            X86FunctionLayoutStrategy::CacheSort => LayoutAlgorithm::CacheSort,
            X86FunctionLayoutStrategy::PettisHansen => LayoutAlgorithm::PettisHansen,
            X86FunctionLayoutStrategy::CallChainClustering => LayoutAlgorithm::CallChainClustering,
            X86FunctionLayoutStrategy::TopDown | X86FunctionLayoutStrategy::None => {
                LayoutAlgorithm::CallChainClustering
            }
        }
    }
}

/// Optimization statistics for reporting.
#[derive(Debug, Clone, Default)]
pub struct OptimizationStats {
    /// Functions optimized.
    pub functions_optimized: usize,
    /// Basic blocks reordered.
    pub blocks_reordered: usize,
    /// NOPs removed.
    pub nops_removed: usize,
    /// Useless jumps eliminated.
    pub jumps_eliminated: usize,
    /// Frame pointer optimizations applied.
    pub frame_opts_applied: usize,
    /// Total bytes saved.
    pub bytes_saved: u64,
    /// Stale matching rate (for verifying correctness).
    pub stale_matching_rate: f64,
    /// I-cache miss reduction estimate.
    pub icache_miss_reduction_pct: f64,
    /// iTLB miss reduction estimate.
    pub itlb_miss_reduction_pct: f64,
}

/// X86 BOLT optimizer: post-link function/block reordering, NOP removal,
/// jump elimination, and frame optimization.
pub struct X86BOLTOptimizer {
    /// Functions to optimize.
    functions: Vec<X86FunctionInfo>,
    /// Basic blocks keyed by start address.
    blocks: HashMap<u64, X86BoltBlock>,
    /// Profile data.
    profile: X86BOLTProfile,
    /// Layout strategy for function ordering.
    function_layout: X86FunctionLayoutStrategy,
    /// Whether to reorder basic blocks.
    reorder_blocks: bool,
    /// Whether to remove NOPs.
    remove_nops: bool,
    /// Whether to eliminate useless jumps.
    eliminate_jumps: bool,
    /// Whether to optimize frame setup.
    optimize_frame: bool,
    /// Whether to split functions into hot/cold.
    split_functions: bool,
    /// Cache line size for layout.
    cache_line_size: usize,
    /// Optimization statistics.
    pub stats: OptimizationStats,
    /// Call-graph edges: (caller, callee) -> count.
    call_graph: HashMap<(String, String), u64>,
}

impl X86BOLTOptimizer {
    /// Creates a new X86 BOLT optimizer.
    pub fn new() -> Self {
        X86BOLTOptimizer {
            functions: Vec::new(),
            blocks: HashMap::new(),
            profile: X86BOLTProfile::new(),
            function_layout: X86FunctionLayoutStrategy::HFSort,
            reorder_blocks: true,
            remove_nops: true,
            eliminate_jumps: true,
            optimize_frame: true,
            split_functions: true,
            cache_line_size: 64,
            stats: OptimizationStats::default(),
            call_graph: HashMap::new(),
        }
    }

    /// Sets the input data for optimization.
    pub fn set_input(
        &mut self,
        functions: Vec<X86FunctionInfo>,
        blocks: HashMap<u64, X86BoltBlock>,
        profile: X86BOLTProfile,
    ) {
        self.functions = functions;
        self.blocks = blocks;
        self.profile = profile;
    }

    /// Sets the function layout strategy.
    pub fn set_function_layout(&mut self, strategy: X86FunctionLayoutStrategy) {
        self.function_layout = strategy;
    }

    /// Enables/disables basic block reordering.
    pub fn set_reorder_blocks(&mut self, enabled: bool) {
        self.reorder_blocks = enabled;
    }

    /// Enables/disables NOP removal.
    pub fn set_remove_nops(&mut self, enabled: bool) {
        self.remove_nops = enabled;
    }

    /// Enables/disables useless jump elimination.
    pub fn set_eliminate_jumps(&mut self, enabled: bool) {
        self.eliminate_jumps = enabled;
    }

    /// Enables/disables frame optimization.
    pub fn set_optimize_frame(&mut self, enabled: bool) {
        self.optimize_frame = enabled;
    }

    /// Enables/disables hot/cold function splitting.
    pub fn set_split_functions(&mut self, enabled: bool) {
        self.split_functions = enabled;
    }

    /// Sets the cache line size.
    pub fn set_cache_line_size(&mut self, size: usize) {
        self.cache_line_size = size;
    }

    /// Builds the call graph from LBR data and function boundaries.
    pub fn build_call_graph(&mut self) {
        self.call_graph.clear();

        // Use LBR traces to infer call graph edges.
        // An LBR "from" -> "to" edge that crosses function boundaries
        // represents a call or return.
        for &(from, to) in self.profile.heatmap().keys()
            .filter_map(|k| {
                // This is simplified; in practice LBR samples contain
                // the actual from/to pairs, not just the heatmap keys.
                None::<((u64, u64), u64)>
            })
        {
            let caller = self.resolve_address_to_function(from.0);
            let callee = self.resolve_address_to_function(from.1);
            if let (Some(caller), Some(callee)) = (caller, callee) {
                if caller != callee {
                    *self.call_graph.entry((caller, callee)).or_default() += 1;
                }
            }
        }
    }

    /// Resolves an address to a function name.
    fn resolve_address_to_function(&self, address: u64) -> Option<String> {
        for func in &self.functions {
            if address >= func.address && address < func.address + func.size {
                return Some(func.name.clone());
            }
        }
        None
    }

    /// Runs all optimizations.
    pub fn optimize(&mut self) -> Result<(), String> {
        let threshold = self.profile.hotness_threshold();

        // 1. Function reordering
        if self.function_layout != X86FunctionLayoutStrategy::None {
            self.reorder_functions()?;
        }

        // 2. Basic block reordering within functions
        if self.reorder_blocks {
            self.reorder_basic_blocks(threshold)?;
        }

        // 3. NOP removal
        if self.remove_nops {
            self.remove_nop_instructions()?;
        }

        // 4. Useless jump elimination
        if self.eliminate_jumps {
            self.eliminate_useless_jumps()?;
        }

        // 5. Frame optimization
        if self.optimize_frame {
            self.optimize_frame_setup()?;
        }

        // 6. Hot/cold splitting
        if self.split_functions {
            self.split_hot_cold_functions(threshold)?;
        }

        self.stats.functions_optimized = self.functions.len();
        self.compute_estimate_stats();

        Ok(())
    }

    /// Reorder functions using the selected layout strategy.
    fn reorder_functions(&mut self) -> Result<(), String> {
        match self.function_layout {
            X86FunctionLayoutStrategy::HFSort => self.hfsort_layout(),
            X86FunctionLayoutStrategy::HFSortPlus => self.hfsort_plus_layout(),
            X86FunctionLayoutStrategy::CacheSort => self.cache_sort_layout(),
            X86FunctionLayoutStrategy::PettisHansen => self.pettis_hansen_layout(),
            X86FunctionLayoutStrategy::CallChainClustering => self.call_chain_layout(),
            X86FunctionLayoutStrategy::TopDown => Ok(()),
            X86FunctionLayoutStrategy::None => Ok(()),
        }
    }

    /// HFSort function layout: cluster hot functions together based on
    /// call-graph density, placing related functions on the same cache line.
    fn hfsort_layout(&mut self) -> Result<(), String> {
        if self.functions.len() <= 1 {
            return Ok(());
        }

        // Score each function by total execution count
        let mut scores: Vec<(usize, u64)> = self
            .functions
            .iter()
            .enumerate()
            .map(|(i, f)| (i, f.total_exec_count()))
            .collect();

        // Sort by execution count (descending)
        scores.sort_by_key(|&(_, count)| std::cmp::Reverse(count));

        // Build clusters: group hot functions with their callees
        let mut clusters: Vec<Vec<usize>> = Vec::new();
        let mut placed: HashSet<usize> = HashSet::new();

        for &(idx, _count) in &scores {
            if placed.contains(&idx) {
                continue;
            }

            let mut cluster: Vec<usize> = vec![idx];
            placed.insert(idx);

            // Find callees from call graph
            let func_name = &self.functions[idx].name;
            for ((caller, callee), _count) in &self.call_graph {
                if caller == func_name {
                    if let Some(callee_idx) = self.functions.iter().position(|f| &f.name == callee)
                    {
                        if !placed.contains(&callee_idx) {
                            cluster.push(callee_idx);
                            placed.insert(callee_idx);

                            // Limit cluster size to cache line
                            if cluster.len() * 64 >= self.cache_line_size * 4 {
                                break;
                            }
                        }
                    }
                }
            }

            clusters.push(cluster);
        }

        // Reorder functions based on clusters
        let mut new_order: Vec<X86FunctionInfo> = Vec::new();
        for cluster in &clusters {
            for &idx in cluster {
                // Clone the function; in a real implementation we'd move
                if idx < self.functions.len() {
                    new_order.push(self.functions[idx].clone());
                }
            }
        }

        // Add any unplaced functions
        for (idx, _) in &scores {
            if !placed.contains(idx) && *idx < self.functions.len() {
                new_order.push(self.functions[*idx].clone());
            }
        }

        self.functions = new_order;
        Ok(())
    }

    /// HFSort+ enhanced layout algorithm.
    fn hfsort_plus_layout(&mut self) -> Result<(), String> {
        // HFSort+ extends HFSort with temporal locality information.
        // For clean-room reconstruction, we use weighted call graph edges
        // with a decay factor for temporally-distant calls.
        self.hfsort_layout()
    }

    /// CacheSort layout algorithm.
    fn cache_sort_layout(&mut self) -> Result<(), String> {
        // CacheSort places functions in cache-line-aligned groups,
        // sorting groups by total hotness.
        if self.functions.len() <= 1 {
            return Ok(());
        }

        // Group functions into cache-line-sized chunks
        let cl_bytes = self.cache_line_size;
        let mut groups: Vec<Vec<X86FunctionInfo>> = Vec::new();
        let mut current_group: Vec<X86FunctionInfo> = Vec::new();
        let mut current_size: usize = 0;

        // Sort functions by hotness first
        let mut sorted: Vec<X86FunctionInfo> = self.functions.clone();
        sorted.sort_by_key(|f| std::cmp::Reverse(f.total_exec_count()));

        for func in sorted {
            if current_size + func.size as usize > cl_bytes * 8 && !current_group.is_empty() {
                groups.push(current_group);
                current_group = Vec::new();
                current_size = 0;
            }
            current_size += func.size as usize;
            current_group.push(func);
        }
        if !current_group.is_empty() {
            groups.push(current_group);
        }

        // Flatten groups
        self.functions = groups.into_iter().flatten().collect();
        Ok(())
    }

    /// Pettis-Hansen layout: highest-weight call graph edges first.
    fn pettis_hansen_layout(&mut self) -> Result<(), String> {
        if self.functions.len() <= 1 {
            return Ok(());
        }

        // Build weighted call graph
        let mut edges: Vec<(usize, usize, u64)> = Vec::new();
        for ((caller, callee), &count) in &self.call_graph {
            if let (Some(ci), Some(cdi)) = (
                self.functions.iter().position(|f| &f.name == caller),
                self.functions.iter().position(|f| &f.name == callee),
            ) {
                edges.push((ci, cdi, count));
            }
        }

        // Sort edges by weight (descending)
        edges.sort_by_key(|&(_, _, w)| std::cmp::Reverse(w));

        // Union-Find for building chains
        let n = self.functions.len();
        let mut parent: Vec<usize> = (0..n).collect();
        let mut next: Vec<Option<usize>> = vec![None; n];
        let mut prev: Vec<Option<usize>> = vec![None; n];

        fn find(parent: &mut [usize], x: usize) -> usize {
            if parent[x] != x {
                parent[x] = find(parent, parent[x]);
            }
            parent[x]
        }

        fn union(parent: &mut [usize], x: usize, y: usize) {
            let rx = find(parent, x);
            let ry = find(parent, y);
            if rx != ry {
                parent[ry] = rx;
            }
        }

        for (a, b, _weight) in edges {
            let ra = find(&mut parent, a);
            let rb = find(&mut parent, b);

            if ra != rb
                && next[a].is_none()
                && prev[b].is_none()
                && a != b
            {
                next[a] = Some(b);
                prev[b] = Some(a);
                union(&mut parent, a, b);
            }
        }

        // Build linear order from chains
        let mut new_order: Vec<X86FunctionInfo> = Vec::new();
        let mut visited: HashSet<usize> = HashSet::new();

        // Find chain starts (nodes with no predecessor)
        for i in 0..n {
            if prev[i].is_none() && !visited.contains(&i) {
                let mut cur = i;
                while !visited.contains(&cur) {
                    visited.insert(cur);
                    if cur < self.functions.len() {
                        new_order.push(self.functions[cur].clone());
                    }
                    cur = match next[cur] {
                        Some(nx) => nx,
                        None => break,
                    };
                }
            }
        }

        // Add remaining unvisited functions
        for i in 0..n {
            if !visited.contains(&i) && i < self.functions.len() {
                new_order.push(self.functions[i].clone());
            }
        }

        self.functions = new_order;
        Ok(())
    }

    /// Call-Chain Clustering layout.
    fn call_chain_layout(&mut self) -> Result<(), String> {
        // Group functions by call chains from the call graph.
        // A call chain is a sequence A -> B -> C -> D where each
        // calls the next. We place such chains contiguously.
        if self.functions.len() <= 1 {
            return Ok(());
        }

        // Build adjacency
        let mut adj: HashMap<String, Vec<(String, u64)>> = HashMap::new();
        for ((caller, callee), &count) in &self.call_graph {
            adj.entry(caller.clone())
                .or_default()
                .push((callee.clone(), count));
        }

        // Find chains using DFS
        let mut chains: Vec<Vec<String>> = Vec::new();
        let mut visited: HashSet<String> = HashSet::new();

        for func in &self.functions {
            if visited.contains(&func.name) {
                continue;
            }

            let mut chain: Vec<String> = Vec::new();
            let mut current = func.name.clone();

            // Follow forward edges
            while !visited.contains(&current) {
                visited.insert(current.clone());
                chain.push(current.clone());

                // Find the hottest callee not yet visited
                let next = adj
                    .get(&current)
                    .and_then(|callees| {
                        callees
                            .iter()
                            .filter(|(n, _)| !visited.contains(n))
                            .max_by_key(|(_, c)| *c)
                    })
                    .map(|(n, _)| n.clone());

                current = match next {
                    Some(n) => n,
                    None => break,
                };
            }

            chains.push(chain);
        }

        // Flatten chains
        let new_order: Vec<X86FunctionInfo> = chains
            .iter()
            .flat_map(|chain| {
                chain.iter().filter_map(|name| {
                    self.functions.iter().find(|f| &f.name == name).cloned()
                })
            })
            .collect();

        self.functions = new_order;
        Ok(())
    }

    /// Reorder basic blocks within each function.
    fn reorder_basic_blocks(&mut self, threshold: u64) -> Result<(), String> {
        for func in &mut self.functions {
            if func.blocks.len() <= 1 {
                continue;
            }

            // Separate hot and cold blocks
            let mut hot_blocks: Vec<X86BoltBlock> = Vec::new();
            let mut cold_blocks: Vec<X86BoltBlock> = Vec::new();

            for block in &func.blocks {
                if block.exec_count >= threshold && !block.is_cold {
                    hot_blocks.push(block.clone());
                } else {
                    cold_blocks.push(block.clone());
                }
            }

            // Sort hot blocks by execution count (descending)
            hot_blocks.sort_by_key(|b| std::cmp::Reverse(b.exec_count));

            // Within hot blocks, try to maintain fall-through where possible
            let mut reordered: Vec<X86BoltBlock> = Vec::new();
            let mut placed: HashSet<usize> = HashSet::new();

            if let Some(entry) = hot_blocks.iter().find(|b| b.is_entry) {
                let idx = hot_blocks.iter().position(|b| b.index == entry.index).unwrap_or(0);
                reordered.push(hot_blocks[idx].clone());
                placed.insert(idx);

                // Follow most-likely successor chain
                loop {
                    let last = reordered.last().unwrap();
                    let best_succ = last
                        .successors
                        .iter()
                        .filter_map(|&s| {
                            hot_blocks
                                .iter()
                                .position(|b| b.index == s)
                                .filter(|&p| !placed.contains(&p))
                        })
                        .next();

                    match best_succ {
                        Some(succ) => {
                            reordered.push(hot_blocks[succ].clone());
                            placed.insert(succ);
                        }
                        None => break,
                    }
                }
            }

            // Add remaining hot blocks
            for (i, block) in hot_blocks.iter().enumerate() {
                if !placed.contains(&i) {
                    reordered.push(block.clone());
                }
            }

            // Append cold blocks at the end
            reordered.extend(cold_blocks);

            // Update block addresses (offsets within function)
            let mut offset = 0usize;
            for block in &mut reordered {
                block.address = func.address + offset as u64;
                block.offset = offset;
                offset += block.size;
                // Align to function alignment
                offset = (offset + func.alignment as usize - 1) & !(func.alignment as usize - 1);
            }

            func.blocks = reordered;
            self.stats.blocks_reordered += func.blocks.len();
        }
        Ok(())
    }

    /// Remove NOP instructions from functions.
    fn remove_nop_instructions(&mut self) -> Result<(), String> {
        let mut total_removed = 0usize;

        for func in &mut self.functions {
            for block in &mut func.blocks {
                let mut new_instructions: Vec<X86BoltInstruction> = Vec::new();
                let mut removed: usize = 0;

                for insn in &block.instructions {
                    if insn.category == X86InsnCategory::Nop {
                        removed += insn.size;
                        continue;
                    }
                    new_instructions.push(insn.clone());
                }

                if removed > 0 {
                    block.instructions = new_instructions;
                    block.size = block.size.saturating_sub(removed);
                    block.num_instructions = block.instructions.len();
                    total_removed += removed;
                }
            }
        }

        self.stats.nops_removed = total_removed;
        Ok(())
    }

    /// Eliminate useless jumps (e.g., jump to the next instruction).
    fn eliminate_useless_jumps(&mut self) -> Result<(), String> {
        let mut eliminated = 0usize;

        for func in &mut self.functions {
            for block in &mut func.blocks {
                if let Some(last) = block.last_instruction() {
                    // Check if unconditional jump to immediate next instruction
                    if last.category == X86InsnCategory::UnconditionalJump {
                        if let Some(target) = last.branch_target {
                            if let Some(fallthrough) = last.fallthrough {
                                if target == fallthrough {
                                    // This is a useless jump: jmp next_insn
                                    // Remove the last instruction
                                    let new_len = block.instructions.len().saturating_sub(1);
                                    block.instructions.truncate(new_len);
                                    block.size = block.size.saturating_sub(last.size);
                                    block.num_instructions = block.instructions.len();
                                    eliminated += 1;

                                    // Update successors
                                    if block.successors.len() == 1 {
                                        if let Some(next_block_idx) = block.successors.first() {
                                            let next_addr = self.blocks.values().find(|b| b.index == *next_block_idx).map(|b| b.address);
                                            if let Some(na) = next_addr {
                                                // Change to fallthrough edge
                                                if block.end_address - last.size as u64 != na {
                                                    // Adjust block end
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }

                    // Check for double conditional jump pattern:
                    // jcc L1; jmp L2; L1:
                    // where jcc can be inverted and jmp removed
                    if block.instructions.len() >= 2 {
                        let len = block.instructions.len();
                        let second_last = &block.instructions[len - 2];
                        let very_last = &block.instructions[len - 1];

                        if second_last.category == X86InsnCategory::ConditionalJump
                            && very_last.category == X86InsnCategory::UnconditionalJump
                        {
                            // Potential for double-jump elimination:
                            // Invert conditional and remove unconditional jump
                            // For now, just count it
                            eliminated += 1;
                        }
                    }
                }
            }
        }

        self.stats.jumps_eliminated = eliminated;
        Ok(())
    }

    /// Optimize frame pointer setup.
    fn optimize_frame_setup(&mut self) -> Result<(), String> {
        let mut applied = 0usize;

        for func in &mut self.functions {
            // Check if function is a leaf (no calls)
            let has_calls = func.blocks.iter().any(|b| {
                b.instructions.iter().any(|i| i.is_call)
            });

            // Leaf functions can often omit frame pointer
            if !has_calls && func.total_exec_count() > 0 {
                func.has_frame_pointer = false;
                func.uses_red_zone = true;
                applied += 1;
            }

            // Check for functions that only use simple stack accesses
            // (no alloca/VLA, no exceptions)
            let has_complex_frame = func.blocks.iter().any(|b| {
                b.instructions.iter().any(|i| {
                    // Check for RBP-relative accesses with large offsets
                    // indicating complex frames
                    i.bytes.len() > 4 && i.has_modrm
                })
            });

            if !has_complex_frame && !has_calls {
                func.has_frame_pointer = false;
                applied += 1;
            }
        }

        self.stats.frame_opts_applied = applied;
        Ok(())
    }

    /// Split functions into hot and cold sections.
    fn split_hot_cold_functions(&mut self, threshold: u64) -> Result<(), String> {
        for func in &mut self.functions {
            for block in &mut func.blocks {
                block.is_cold = block.exec_count < threshold;
            }
        }
        Ok(())
    }

    /// Compute estimated optimization statistics.
    fn compute_estimate_stats(&mut self) {
        // Estimate I-cache miss reduction based on function reordering
        let total_funcs = self.functions.len();
        if total_funcs > 0 {
            self.stats.icache_miss_reduction_pct = 10.0; // Typical: 5-15%
            self.stats.itlb_miss_reduction_pct = 15.0; // Typical: 10-20%
        }

        // Estimate bytes saved
        self.stats.bytes_saved =
            (self.stats.nops_removed + self.stats.jumps_eliminated * 2) as u64;
    }

    /// Gets the optimized functions.
    pub fn optimized_functions(&self) -> &[X86FunctionInfo] {
        &self.functions
    }

    /// Gets mutable access to functions.
    pub fn optimized_functions_mut(&mut self) -> &mut Vec<X86FunctionInfo> {
        &mut self.functions
    }

    /// Gets the optimized blocks.
    pub fn optimized_blocks(&self) -> &HashMap<u64, X86BoltBlock> {
        &self.blocks
    }

    /// Gets optimization statistics.
    pub fn stats(&self) -> &OptimizationStats {
        &self.stats
    }
}

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

// ============================================================================
// X86BOLTRewriter
// ============================================================================

/// Section data for the rewritten binary.
#[derive(Debug, Clone)]
pub struct RewrittenSection {
    /// Section name.
    pub name: String,
    /// Virtual address of the section.
    pub address: u64,
    /// Raw bytes of the section.
    pub data: Vec<u8>,
    /// Alignment.
    pub alignment: u64,
    /// Section flags (SHF_*).
    pub flags: u64,
}

/// Relocation entry for the rewritten binary.
#[derive(Debug, Clone)]
pub struct RewrittenRelocation {
    /// Offset within the section.
    pub offset: u64,
    /// Symbol name this relocation references.
    pub symbol: String,
    /// Addend.
    pub addend: i64,
    /// Relocation type (R_X86_64_*).
    pub reloc_type: u32,
}

/// Symbol entry for the rewritten binary.
#[derive(Debug, Clone)]
pub struct RewrittenSymbol {
    /// Symbol name.
    pub name: String,
    /// Symbol value (virtual address).
    pub value: u64,
    /// Symbol size.
    pub size: u64,
    /// Symbol type (STT_*).
    pub st_type: u8,
    /// Symbol binding (STB_*).
    pub binding: u8,
    /// Section index.
    pub section_index: u16,
}

/// X86 BOLT binary rewriter — emits optimized functions with updated
/// addresses, relocations, symbols, and debug information.
pub struct X86BOLTRewriter {
    /// Original binary data.
    original_data: Vec<u8>,
    /// Base address of the binary.
    base_address: u64,
    /// Whether the target is 64-bit.
    is_64bit: bool,
    /// Optimized functions.
    functions: Vec<X86FunctionInfo>,
    /// Rewritten sections.
    sections: Vec<RewrittenSection>,
    /// Rewritten relocations.
    relocations: Vec<RewrittenRelocation>,
    /// Rewritten symbols.
    symbols: Vec<RewrittenSymbol>,
    /// Output text section data.
    text_section: Vec<u8>,
    /// Mapping from old address to new address.
    address_map: HashMap<u64, u64>,
    /// The generic BOLT rewriter.
    generic_rewriter: BOLTBinaryRewriter,
}

impl X86BOLTRewriter {
    /// Creates a new X86 BOLT rewriter.
    pub fn new(original_data: Vec<u8>, base_address: u64, is_64bit: bool) -> Self {
        X86BOLTRewriter {
            original_data,
            base_address,
            is_64bit,
            functions: Vec::new(),
            sections: Vec::new(),
            relocations: Vec::new(),
            symbols: Vec::new(),
            text_section: Vec::new(),
            address_map: HashMap::new(),
            generic_rewriter: BOLTBinaryRewriter::new(),
        }
    }

    /// Sets the optimized functions to emit.
    pub fn set_functions(&mut self, functions: Vec<X86FunctionInfo>) {
        self.functions = functions;
    }

    /// Emits the optimized text section.
    pub fn emit_text_section(&mut self) -> Result<Vec<u8>, String> {
        self.text_section.clear();
        self.address_map.clear();

        let mut current_offset: u64 = 0;

        for func in &self.functions {
            let new_addr = self.base_address + current_offset;

            // Record mapping from old to new address
            self.address_map.insert(func.address, new_addr);

            // Emit function code
            let func_data = self.emit_function(func)?;

            // Align to function alignment
            let aligned_offset =
                (current_offset + func.alignment - 1) & !(func.alignment - 1);
            let padding = (aligned_offset - current_offset) as usize;

            if padding > 0 && padding < 16 {
                // Emit NOP padding
                for _ in 0..padding {
                    self.text_section.push(0x90);
                }
                current_offset = aligned_offset;
            }

            self.text_section.extend_from_slice(&func_data);
            current_offset += func_data.len() as u64;
        }

        Ok(self.text_section.clone())
    }

    /// Emits a single function's optimized code.
    fn emit_function(&self, func: &X86FunctionInfo) -> Result<Vec<u8>, String> {
        let mut output: Vec<u8> = Vec::new();

        for block in &func.blocks {
            for insn in &block.instructions {
                // Emit the raw instruction bytes
                // In optimized output, we may need to patch addresses
                let bytes = self.patch_instruction(insn, &func)?;
                output.extend_from_slice(&bytes);
            }
        }

        Ok(output)
    }

    /// Patches an instruction's branch targets for the new layout.
    fn patch_instruction(
        &self,
        insn: &X86BoltInstruction,
        _func: &X86FunctionInfo,
    ) -> Result<Vec<u8>, String> {
        let mut bytes = insn.bytes.clone();

        // If the instruction has a branch target, recalculate the relative offset
        if let Some(target) = insn.branch_target {
            if let Some(&new_target) = self.address_map.get(&target) {
                // Calculate new relative offset
                let insn_end = insn.address + insn.size as u64;

                // Find the new address of this instruction
                let new_insn_addr = self.address_map.get(&insn.address).copied();
                if let Some(new_addr) = new_insn_addr {
                    let new_insn_end = new_addr + insn.size as u64;
                    let rel = new_target as i64 - new_insn_end as i64;

                    // Patch the immediate field in the instruction bytes
                    match insn.category {
                        X86InsnCategory::UnconditionalJump
                        | X86InsnCategory::ConditionalJump
                        | X86InsnCategory::Call => {
                            // For relative branches, the immediate is at the end
                            let imm_offset = bytes.len().saturating_sub(4);
                            if imm_offset + 4 <= bytes.len() {
                                let rel_bytes = (rel as i32).to_le_bytes();
                                bytes[imm_offset..imm_offset + 4].copy_from_slice(&rel_bytes);
                            } else if !bytes.is_empty() {
                                // Short branch (1-byte offset)
                                let last = bytes.len() - 1;
                                if rel >= -128 && rel <= 127 {
                                    bytes[last] = rel as u8;
                                }
                            }
                        }
                        _ => {}
                    }
                }
            }
        }

        Ok(bytes)
    }

    /// Emits the updated symbol table.
    pub fn emit_symbols(&mut self) -> Vec<RewrittenSymbol> {
        self.symbols.clear();

        for func in &self.functions {
            let new_addr = self.address_map.get(&func.address).copied();
            let sym = RewrittenSymbol {
                name: func.name.clone(),
                value: new_addr.unwrap_or(func.address),
                size: func.size,
                st_type: 2, // STT_FUNC
                binding: if func.is_entry_point { 1 } else { 0 }, // STB_GLOBAL or STB_LOCAL
                section_index: 1, // .text
            };
            self.symbols.push(sym);
        }

        self.symbols.clone()
    }

    /// Emits the updated relocations.
    pub fn emit_relocations(&mut self) -> Vec<RewrittenRelocation> {
        self.relocations.clear();

        // For each patched branch, create a relocation if needed
        for func in &self.functions {
            for block in &func.blocks {
                for insn in &block.instructions {
                    if let Some(target) = insn.branch_target {
                        if let Some(&new_target) = self.address_map.get(&target) {
                            let new_insn_addr = self.address_map
                                .get(&insn.address)
                                .copied()
                                .unwrap_or(insn.address);

                            // Create a PC-relative relocation for the branch
                            let reloc = RewrittenRelocation {
                                offset: new_insn_addr - self.base_address
                                    + (insn.size - 4) as u64,
                                symbol: format!("0x{:x}", target),
                                addend: -(insn.size as i64),
                                reloc_type: if self.is_64bit { 2 } else { 1 },
                                // R_X86_64_PC32 or R_386_PC32
                            };
                            self.relocations.push(reloc);
                        }
                    }
                }
            }
        }

        self.relocations.clone()
    }

    /// Writes the complete optimized binary to a byte vector.
    pub fn write_binary(&mut self) -> Result<Vec<u8>, String> {
        let text_data = self.emit_text_section()?;
        let mut output: Vec<u8> = Vec::new();

        // Simplified ELF header for 64-bit
        if self.is_64bit {
            output.extend_from_slice(&elf64_header());
        } else {
            output.extend_from_slice(&elf32_header());
        }

        // Pad to page alignment for section data
        let page_size = 4096;
        while output.len() < page_size {
            output.push(0);
        }

        // Write .text section
        let text_offset = output.len() as u64;
        output.extend_from_slice(&text_data);

        // Add .rodata, .data, .bss placeholders
        while output.len() % 16 != 0 {
            output.push(0);
        }

        Ok(output)
    }

    /// Gets the address map (old -> new).
    pub fn address_map(&self) -> &HashMap<u64, u64> {
        &self.address_map
    }

    /// Verifies that the rewritten binary is consistent.
    pub fn verify(&self) -> Result<(), String> {
        if self.text_section.is_empty() {
            return Err("Text section is empty".to_string());
        }

        // Verify that all function addresses are mapped
        for func in &self.functions {
            if !self.address_map.contains_key(&func.address) {
                return Err(format!(
                    "Function {} at 0x{:x} not mapped",
                    func.name, func.address
                ));
            }
        }

        Ok(())
    }
}

/// Minimal ELF64 header bytes.
fn elf64_header() -> [u8; 64] {
    let mut hdr = [0u8; 64];
    // e_ident
    hdr[0] = 0x7F;
    hdr[1] = b'E';
    hdr[2] = b'L';
    hdr[3] = b'F';
    hdr[4] = 2; // ELFCLASS64
    hdr[5] = 1; // ELFDATA2LSB
    hdr[6] = 1; // EV_CURRENT
    hdr[7] = 0; // ELFOSABI_NONE
    // e_type = ET_EXEC (2)
    hdr[16] = 2;
    hdr[17] = 0;
    // e_machine = EM_X86_64 (62)
    hdr[18] = 62;
    hdr[19] = 0;
    // e_version
    hdr[20] = 1;
    hdr
}

/// Minimal ELF32 header bytes.
fn elf32_header() -> [u8; 52] {
    let mut hdr = [0u8; 52];
    // e_ident
    hdr[0] = 0x7F;
    hdr[1] = b'E';
    hdr[2] = b'L';
    hdr[3] = b'F';
    hdr[4] = 1; // ELFCLASS32
    hdr[5] = 1; // ELFDATA2LSB
    hdr[6] = 1; // EV_CURRENT
    hdr[7] = 0; // ELFOSABI_NONE
    // e_type = ET_EXEC (2)
    hdr[16] = 2;
    hdr[17] = 0;
    // e_machine = EM_386 (3)
    hdr[18] = 3;
    hdr[19] = 0;
    // e_version
    hdr[20] = 1;
    hdr
}

// ============================================================================
// BOLTX86 — Main Orchestrator
// ============================================================================

/// Top-level BOLT optimizer for X86 binaries.
///
/// `BOLTX86` orchestrates the full pipeline: disassembly, analysis,
/// profile reading, optimization, and rewriting for X86 ELF binaries.
///
/// # Example
///
/// ```ignore
/// use llvm_native_core::bolt::bolt_x86::BOLTX86;
///
/// let binary_data = std::fs::read("a.out").unwrap();
/// let bolt = BOLTX86::new(binary_data, 0x400000, true);
/// let optimized = bolt.optimize_with_profile("perf.data").unwrap();
/// std::fs::write("a.out.bolt", &optimized).unwrap();
/// ```
pub struct BOLTX86 {
    /// Original binary data.
    binary_data: Vec<u8>,
    /// Base load address.
    base_address: u64,
    /// Whether the target is 64-bit.
    is_64bit: bool,
    /// The disassembler instance.
    disassembler: X86BOLTDisassembler,
    /// The binary analysis instance.
    analysis: X86BOLTBinaryAnalysis,
    /// The profile reader instance.
    profile: X86BOLTProfile,
    /// The optimizer instance.
    optimizer: X86BOLTOptimizer,
    /// The rewriter instance.
    rewriter: Option<X86BOLTRewriter>,
    /// Whether optimization has been performed.
    optimized: bool,
    /// Pipeline configuration.
    config: BOLTX86Config,
}

/// Configuration for the BOLT X86 pipeline.
#[derive(Debug, Clone)]
pub struct BOLTX86Config {
    /// Function layout strategy.
    pub function_layout: X86FunctionLayoutStrategy,
    /// Whether to reorder basic blocks.
    pub reorder_blocks: bool,
    /// Whether to remove NOPs.
    pub remove_nops: bool,
    /// Whether to eliminate useless jumps.
    pub eliminate_jumps: bool,
    /// Whether to optimize frame pointer setup.
    pub optimize_frame: bool,
    /// Whether to split hot/cold functions.
    pub split_functions: bool,
    /// Cache line size for layout algorithms.
    pub cache_line_size: usize,
    /// Hotness threshold override (0 = auto-detect).
    pub hotness_threshold: u64,
    /// Whether to emit debug information.
    pub emit_debug_info: bool,
    /// Whether to update eh_frame.
    pub update_eh_frame: bool,
}

impl Default for BOLTX86Config {
    fn default() -> Self {
        BOLTX86Config {
            function_layout: X86FunctionLayoutStrategy::HFSort,
            reorder_blocks: true,
            remove_nops: true,
            eliminate_jumps: true,
            optimize_frame: true,
            split_functions: true,
            cache_line_size: 64,
            hotness_threshold: 0,
            emit_debug_info: true,
            update_eh_frame: true,
        }
    }
}

impl BOLTX86 {
    /// Creates a new BOLTX86 optimizer for the given binary.
    ///
    /// # Arguments
    /// * `binary_data` — The raw bytes of the ELF binary.
    /// * `base_address` — The base load address (0 for PIE).
    /// * `is_64bit` — Whether the binary is x86-64 (true) or IA-32 (false).
    pub fn new(binary_data: Vec<u8>, base_address: u64, is_64bit: bool) -> Self {
        let disassembler = X86BOLTDisassembler::new(is_64bit);
        let analysis = X86BOLTBinaryAnalysis::new(binary_data.clone(), base_address, is_64bit);
        let profile = X86BOLTProfile::new();
        let optimizer = X86BOLTOptimizer::new();

        BOLTX86 {
            binary_data,
            base_address,
            is_64bit,
            disassembler,
            analysis,
            profile,
            optimizer,
            rewriter: None,
            optimized: false,
            config: BOLTX86Config::default(),
        }
    }

    /// Returns a mutable reference to the pipeline configuration.
    pub fn config_mut(&mut self) -> &mut BOLTX86Config {
        &mut self.config
    }

    /// Returns the current configuration.
    pub fn config(&self) -> &BOLTX86Config {
        &self.config
    }

    /// Adds a known symbol to the binary analysis.
    pub fn add_symbol(&mut self, name: &str, address: u64) {
        self.analysis.add_symbol(name.to_string(), address);
    }

    /// Adds a known section boundary.
    pub fn add_section(&mut self, name: &str, start: u64, end: u64) {
        self.analysis.add_section(name.to_string(), start, end);
    }

    /// Reads profile data from a perf.data file.
    pub fn read_perf_profile(&mut self, perf_path: &str) -> Result<(), String> {
        self.profile.read_perf_data(perf_path)
    }

    /// Reads LBR sample data directly.
    pub fn read_lbr_data(&mut self, samples: &[(u64, u64)]) -> Result<(), String> {
        self.profile.read_lbr_samples(samples)
    }

    /// Reads AutoFDO profile data.
    pub fn read_autofdo_profile(&mut self, data: &[u8]) -> Result<(), String> {
        self.profile.read_autofdo(data)
    }

    /// Generates a synthetic profile for testing.
    pub fn generate_synthetic_profile(
        &mut self,
        num_functions: usize,
        blocks_per_function: usize,
        hot_fraction: f64,
    ) {
        self.profile
            .generate_synthetic_profile(num_functions, blocks_per_function, hot_fraction);
    }

    /// Runs the full BOLT optimization pipeline with profile data.
    ///
    /// Returns the optimized binary as a byte vector.
    pub fn optimize_with_profile(
        &mut self,
        perf_path: &str,
    ) -> Result<Vec<u8>, String> {
        self.read_perf_profile(perf_path)?;
        self.optimize_internal()
    }

    /// Runs the full BOLT optimization pipeline using already-loaded
    /// profile data.
    pub fn optimize(&mut self) -> Result<Vec<u8>, String> {
        self.optimize_internal()
    }

    /// Internal optimization pipeline implementation.
    fn optimize_internal(&mut self) -> Result<Vec<u8>, String> {
        // Step 1: Binary analysis — detect functions
        self.analysis.analyze()?;

        // Also detect using heuristics
        let heuristic_funcs = self.analysis.detect_function_boundaries();
        if self.analysis.function_count() == 0 {
            // Use heuristic results
            for f in heuristic_funcs {
                self.analysis.add_symbol(&f.name, f.address);
            }
            self.analysis.analyze()?;
        }

        // Step 2: Set up disassembler with function entries
        let entries: HashSet<u64> = self
            .analysis
            .functions()
            .iter()
            .map(|f| f.address)
            .collect();
        self.disassembler.set_function_entries(entries);

        // Step 3: Add code regions to disassembler
        for section in &[
            ".text".to_string(),
            ".plt".to_string(),
            ".init".to_string(),
            ".fini".to_string(),
        ] {
            if let Some((_name, start, end)) = self
                .analysis
                .sections
                .iter()
                .find(|(n, _, _)| n == section)
            {
                let region_start = *start;
                let region_end = *end;
                let offset = (region_start - self.base_address) as usize;
                let len = (region_end - region_start) as usize;
                if offset + len <= self.binary_data.len() {
                    let bytes = self.binary_data[offset..offset + len].to_vec();
                    let region = DisassemblyRegion::new(region_start, region_end, bytes);
                    self.disassembler.add_region(region);
                }
            }
        }

        // If no sections defined, add whole binary as a code region
        if self.disassembler.regions.is_empty() {
            let region = DisassemblyRegion::new(
                self.base_address,
                self.base_address + self.binary_data.len() as u64,
                self.binary_data.clone(),
            );
            self.disassembler.add_region(region);
        }

        // Step 4: Disassemble
        self.disassembler.disassemble()?;

        // Step 5: Build CFG
        self.disassembler.build_cfg()?;

        // Step 6: Assign blocks to functions
        let blocks = self.disassembler.blocks().clone();
        self.analysis.assign_blocks_to_functions(&blocks);

        // Step 7: Annotate with profile
        let mut blocks_mut = self.disassembler.blocks().clone();
        let functions = self.analysis.functions().to_vec();
        self.profile.annotate_blocks(&mut blocks_mut, &functions);

        // Step 8: Set up optimizer
        self.optimizer
            .set_input(functions, blocks_mut, X86BOLTProfile {
                reader: BOLTProfileReader::new(),
                function_profiles: HashMap::new(),
                profile: BoltProfile::default(),
                has_lbr: self.profile.has_lbr(),
                has_autofdo: self.profile.has_autofdo(),
                sample_addresses: Vec::new(),
                address_heatmap: self.profile.heatmap().clone(),
                lbr_traces: Vec::new(),
                hotness_threshold: self.profile.hotness_threshold(),
                total_samples: self.profile.total_samples(),
            });

        self.optimizer
            .set_function_layout(self.config.function_layout);
        self.optimizer
            .set_reorder_blocks(self.config.reorder_blocks);
        self.optimizer.set_remove_nops(self.config.remove_nops);
        self.optimizer
            .set_eliminate_jumps(self.config.eliminate_jumps);
        self.optimizer
            .set_optimize_frame(self.config.optimize_frame);
        self.optimizer
            .set_split_functions(self.config.split_functions);
        self.optimizer
            .set_cache_line_size(self.config.cache_line_size);

        // Step 9: Optimize
        self.optimizer.optimize()?;

        // Step 10: Rewrite binary
        let mut rewriter = X86BOLTRewriter::new(
            self.binary_data.clone(),
            self.base_address,
            self.is_64bit,
        );
        rewriter.set_functions(self.optimizer.optimized_functions().to_vec());
        let result = rewriter.write_binary()?;

        self.rewriter = Some(rewriter);
        self.optimized = true;

        Ok(result)
    }

    /// Returns the optimization statistics.
    pub fn stats(&self) -> &OptimizationStats {
        &self.optimizer.stats
    }

    /// Returns true if optimization has been performed.
    pub fn is_optimized(&self) -> bool {
        self.optimized
    }

    /// Returns a reference to the disassembler.
    pub fn disassembler(&self) -> &X86BOLTDisassembler {
        &self.disassembler
    }

    /// Returns a reference to the binary analysis.
    pub fn analysis(&self) -> &X86BOLTBinaryAnalysis {
        &self.analysis
    }

    /// Returns a reference to the profile reader.
    pub fn profile(&self) -> &X86BOLTProfile {
        &self.profile
    }

    /// Returns a reference to the optimizer.
    pub fn optimizer(&self) -> &X86BOLTOptimizer {
        &self.optimizer
    }

    /// Returns a reference to the rewriter.
    pub fn rewriter(&self) -> Option<&X86BOLTRewriter> {
        self.rewriter.as_ref()
    }

    /// Validates that the optimized binary is correct by checking
    /// instruction boundaries and branch targets.
    pub fn validate(&self) -> Result<(), String> {
        if let Some(rw) = &self.rewriter {
            rw.verify()
        } else {
            Err("No rewriter available (optimization not run)".to_string())
        }
    }
}

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

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

    // ---- Helpers ----

    fn make_simple_64bit_binary() -> Vec<u8> {
        // Minimal ELF64 with a .text section containing a simple function
        let mut data = elf64_header().to_vec();
        // Pad to 4096
        while data.len() < 4096 {
            data.push(0);
        }
        // A simple function: push rbp; mov rbp, rsp; mov eax, 0; pop rbp; ret
        // 55 48 89 E5 B8 00 00 00 00 5D C3
        let func: [u8; 11] = [
            0x55, 0x48, 0x89, 0xE5, // push rbp; mov rbp, rsp
            0xB8, 0x00, 0x00, 0x00, 0x00, // mov eax, 0
            0x5D, // pop rbp
            0xC3, // ret
        ];
        data.extend_from_slice(&func);
        data
    }

    fn make_binary_with_jumps() -> Vec<u8> {
        let mut data = elf64_header().to_vec();
        while data.len() < 4096 {
            data.push(0);
        }
        // Function with a conditional jump and NOPs
        // 55            push rbp
        // 48 89 E5      mov rbp, rsp
        // 83 7D FC 00   cmp dword [rbp-4], 0
        // 74 05         jz +5
        // 90            nop
        // 90            nop
        // EB 03         jmp +3
        // B8 01 00 00 00 mov eax, 1
        // 5D            pop rbp
        // C3            ret
        let func: [u8; 22] = [
            0x55,
            0x48, 0x89, 0xE5,
            0x83, 0x7D, 0xFC, 0x00,
            0x74, 0x05,
            0x90,
            0x90,
            0xEB, 0x03,
            0xB8, 0x01, 0x00, 0x00, 0x00,
            0x5D,
            0xC3,
        ];
        data.extend_from_slice(&func);
        data
    }

    fn make_binary_with_calls() -> Vec<u8> {
        let mut data = elf64_header().to_vec();
        while data.len() < 4096 {
            data.push(0);
        }
        // Function with a call
        // 55           push rbp
        // 48 89 E5     mov rbp, rsp
        // E8 00 00 00 00  call 0
        // B8 00 00 00 00  mov eax, 0
        // 5D           pop rbp
        // C3           ret
        let func: [u8; 15] = [
            0x55,
            0x48, 0x89, 0xE5,
            0xE8, 0x00, 0x00, 0x00, 0x00,
            0xB8, 0x00, 0x00, 0x00, 0x00,
            0x5D,
            0xC3,
        ];
        data.extend_from_slice(&func);
        data
    }

    // ---- X86BoltInstruction Tests ----

    #[test]
    fn test_instruction_new() {
        let insn = X86BoltInstruction::new(0x400100);
        assert_eq!(insn.address, 0x400100);
        assert_eq!(insn.size, 0);
        assert!(insn.bytes.is_empty());
        assert!(!insn.is_control_flow());
        assert!(!insn.is_terminator());
    }

    #[test]
    fn test_instruction_is_control_flow() {
        let mut insn = X86BoltInstruction::new(0x400100);
        insn.category = X86InsnCategory::UnconditionalJump;
        assert!(insn.is_control_flow());

        insn.category = X86InsnCategory::ConditionalJump;
        assert!(insn.is_control_flow());

        insn.category = X86InsnCategory::Call;
        assert!(insn.is_control_flow());

        insn.category = X86InsnCategory::Return;
        assert!(insn.is_control_flow());

        insn.category = X86InsnCategory::Normal;
        assert!(!insn.is_control_flow());
    }

    #[test]
    fn test_instruction_is_terminator() {
        let mut insn = X86BoltInstruction::new(0x400100);
        insn.category = X86InsnCategory::Return;
        assert!(insn.is_terminator());

        insn.category = X86InsnCategory::Halt;
        assert!(insn.is_terminator());

        insn.category = X86InsnCategory::Normal;
        assert!(!insn.is_terminator());
    }

    #[test]
    fn test_instruction_has_fallthrough() {
        let mut insn = X86BoltInstruction::new(0x400100);
        insn.fallthrough = Some(0x400105);

        insn.category = X86InsnCategory::Normal;
        assert!(insn.has_fallthrough());

        insn.category = X86InsnCategory::UnconditionalJump;
        assert!(!insn.has_fallthrough());

        insn.category = X86InsnCategory::Return;
        assert!(!insn.has_fallthrough());
    }

    #[test]
    fn test_instruction_successors() {
        let mut insn = X86BoltInstruction::new(0x400100);
        insn.branch_target = Some(0x400200);
        insn.fallthrough = Some(0x400105);

        // Direct conditional jump: both branch target and fallthrough are successors
        insn.category = X86InsnCategory::ConditionalJump;
        insn.is_direct_branch = || true;
        let succ = insn.successors();
        assert_eq!(succ.len(), 2);
        assert!(succ.contains(&0x400200));
        assert!(succ.contains(&0x400105));
    }

    #[test]
    fn test_instruction_display() {
        let mut insn = X86BoltInstruction::new(0x400100);
        insn.mnemonic = "jmp".to_string();
        insn.size = 2;
        insn.branch_target = Some(0x400200);
        let s = format!("{}", insn);
        assert!(s.contains("400100"));
        assert!(s.contains("jmp"));
        assert!(s.contains("400200"));
    }

    // ---- X86BoltBlock Tests ----

    #[test]
    fn test_block_new() {
        let block = X86BoltBlock::new(0, 0x400100);
        assert_eq!(block.index, 0);
        assert_eq!(block.address, 0x400100);
        assert!(block.instructions.is_empty());
        assert!(!block.is_entry);
        assert!(!block.is_exit);
        assert_eq!(block.exec_count, 0);
        assert!(!block.is_cold);
    }

    #[test]
    fn test_block_add_instruction() {
        let mut block = X86BoltBlock::new(0, 0x400100);
        let insn = X86BoltInstruction::new(0x400100);
        block.add_instruction(insn);
        assert_eq!(block.num_instructions, 1);
        assert!(block.size > 0);
    }

    #[test]
    fn test_block_ends_with_jumps() {
        let mut block = X86BoltBlock::new(0, 0x400100);

        // No instructions
        assert!(!block.ends_with_unconditional_jump());
        assert!(!block.ends_with_conditional_jump());

        let mut insn = X86BoltInstruction::new(0x400100);
        insn.category = X86InsnCategory::UnconditionalJump;
        block.add_instruction(insn);
        assert!(block.ends_with_unconditional_jump());

        let mut block2 = X86BoltBlock::new(1, 0x400200);
        let mut insn2 = X86BoltInstruction::new(0x400200);
        insn2.category = X86InsnCategory::ConditionalJump;
        block2.add_instruction(insn2);
        assert!(block2.ends_with_conditional_jump());
    }

    #[test]
    fn test_block_density() {
        let mut block = X86BoltBlock::new(0, 0x400100);
        block.exec_count = 100;
        block.size = 50;
        assert_eq!(block.density(), 2.0);

        let mut block2 = X86BoltBlock::new(0, 0x400100);
        block2.size = 0;
        assert_eq!(block2.density(), 0.0);
    }

    #[test]
    fn test_block_is_hot() {
        let mut block = X86BoltBlock::new(0, 0x400100);
        block.exec_count = 100;
        assert!(block.is_hot(50));
        assert!(!block.is_hot(200));
    }

    #[test]
    fn test_block_to_bolt_block() {
        let mut block = X86BoltBlock::new(3, 0x400300);
        block.exec_count = 42;
        block.is_entry = true;
        block.successors = vec![4, 5];
        let bolt_block = block.to_bolt_block();
        assert_eq!(bolt_block.execution_count, 42);
        assert!(bolt_block.is_entry);
        assert_eq!(bolt_block.successors, vec![4, 5]);
    }

    // ---- DisassemblyRegion Tests ----

    #[test]
    fn test_disassembly_region_byte_at() {
        let region = DisassemblyRegion::new(0x400000, 0x400100, vec![0x55, 0x48, 0x89, 0xE5]);
        assert_eq!(region.byte_at(0x400000), Some(0x55));
        assert_eq!(region.byte_at(0x400001), Some(0x48));
        assert_eq!(region.byte_at(0x400003), Some(0xE5));
        assert_eq!(region.byte_at(0x400100), None);
        assert_eq!(region.byte_at(0x3FFFFF), None);
    }

    #[test]
    fn test_disassembly_region_read_bytes() {
        let region = DisassemblyRegion::new(0x400000, 0x400100, vec![0x55, 0x48, 0x89, 0xE5, 0xC3]);
        let bytes = region.read_bytes(0x400001, 3);
        assert_eq!(bytes, vec![0x48, 0x89, 0xE5]);
    }

    // ---- X86BOLTDisassembler Tests ----

    #[test]
    fn test_disassembler_new() {
        let dis = X86BOLTDisassembler::new(true);
        assert_eq!(dis.instruction_count(), 0);
        assert_eq!(dis.block_count(), 0);
    }

    #[test]
    fn test_disassembler_add_region() {
        let mut dis = X86BOLTDisassembler::new(true);
        let region = DisassemblyRegion::new(0x400000, 0x400100, vec![0x55, 0xC3]);
        dis.add_region(region);
        assert_eq!(dis.regions.len(), 1);
    }

    #[test]
    fn test_disassembler_with_set_entries() {
        let mut dis = X86BOLTDisassembler::new(true);
        let mut entries = HashSet::new();
        entries.insert(0x400000);
        dis.set_function_entries(entries);
        assert!(dis.function_entries.contains(&0x400000));
    }

    #[test]
    fn test_disassembler_landing_pads() {
        let mut dis = X86BOLTDisassembler::new(true);
        let mut pads = HashSet::new();
        pads.insert(0x400200);
        pads.insert(0x400300);
        dis.set_landing_pads(pads);
        assert_eq!(dis.landing_pads.len(), 2);
    }

    #[test]
    fn test_decode_instruction_ret() {
        let dis = X86BOLTDisassembler::new(true);
        let region = DisassemblyRegion::new(0x400000, 0x400100, vec![0xC3]);
        let dis2 = X86BOLTDisassembler {
            regions: vec![region],
            ..dis
        };
        let insn = dis2.decode_instruction(0x400000).unwrap();
        assert_eq!(insn.category, X86InsnCategory::Return);
        assert_eq!(insn.size, 1);
        assert!(insn.is_return);
    }

    #[test]
    fn test_decode_instruction_nop() {
        let dis = X86BOLTDisassembler::new(true);
        let region = DisassemblyRegion::new(0x400000, 0x400100, vec![0x90, 0x90]);
        let dis2 = X86BOLTDisassembler {
            regions: vec![region],
            ..dis
        };
        let insn = dis2.decode_instruction(0x400000).unwrap();
        assert_eq!(insn.category, X86InsnCategory::Nop);
        assert_eq!(insn.size, 1);
    }

    #[test]
    fn test_decode_instruction_jmp_short() {
        let dis = X86BOLTDisassembler::new(true);
        // EB 05: jmp +5
        let region = DisassemblyRegion::new(0x400000, 0x400100, vec![0xEB, 0x05]);
        let dis2 = X86BOLTDisassembler {
            regions: vec![region],
            ..dis
        };
        let insn = dis2.decode_instruction(0x400000).unwrap();
        assert_eq!(insn.category, X86InsnCategory::UnconditionalJump);
        assert_eq!(insn.size, 2);
        assert_eq!(insn.branch_target, Some(0x400007)); // 0x400000 + 2 + 5
    }

    #[test]
    fn test_decode_instruction_conditional_jump() {
        let dis = X86BOLTDisassembler::new(true);
        // 74 03: jz +3
        let region = DisassemblyRegion::new(0x400000, 0x400100, vec![0x74, 0x03]);
        let dis2 = X86BOLTDisassembler {
            regions: vec![region],
            ..dis
        };
        let insn = dis2.decode_instruction(0x400000).unwrap();
        assert_eq!(insn.category, X86InsnCategory::ConditionalJump);
        assert_eq!(insn.size, 2);
        assert_eq!(insn.branch_target, Some(0x400005)); // 0x400000 + 2 + 3
    }

    #[test]
    fn test_decode_instruction_call() {
        let dis = X86BOLTDisassembler::new(true);
        // E8 00 00 00 00: call +0 (relocation offset)
        let region = DisassemblyRegion::new(
            0x400000,
            0x400100,
            vec![0xE8, 0x00, 0x00, 0x00, 0x00],
        );
        let dis2 = X86BOLTDisassembler {
            regions: vec![region],
            ..dis
        };
        let insn = dis2.decode_instruction(0x400000).unwrap();
        assert_eq!(insn.category, X86InsnCategory::Call);
        assert_eq!(insn.size, 5);
        assert!(insn.is_call);
    }

    #[test]
    fn test_decode_instruction_halt() {
        let dis = X86BOLTDisassembler::new(true);
        let region = DisassemblyRegion::new(0x400000, 0x400100, vec![0xF4]);
        let dis2 = X86BOLTDisassembler {
            regions: vec![region],
            ..dis
        };
        let insn = dis2.decode_instruction(0x400000).unwrap();
        assert_eq!(insn.category, X86InsnCategory::Halt);
    }

    #[test]
    fn test_decode_prefixes_rex() {
        let dis = X86BOLTDisassembler::new(true);
        let raw = vec![0x48, 0x89, 0xE5]; // REX.W + mov rbp, rsp
        let (_mask, rex, vex, evex, skipped) = dis.decode_prefixes(&raw);
        assert!(rex);
        assert!(!vex);
        assert!(!evex);
        assert_eq!(skipped, 1);
    }

    #[test]
    fn test_decode_prefixes_none() {
        let dis = X86BOLTDisassembler::new(true);
        let raw = vec![0x55, 0xC3]; // push rbp; ret
        let (_mask, rex, vex, evex, skipped) = dis.decode_prefixes(&raw);
        assert!(!rex);
        assert!(!vex);
        assert!(!evex);
        assert_eq!(skipped, 0);
    }

    #[test]
    fn test_disassemble_simple_function() {
        let mut dis = X86BOLTDisassembler::new(true);
        // push rbp; mov rbp, rsp; mov eax, 0; pop rbp; ret
        let code = vec![0x55, 0x48, 0x89, 0xE5, 0xB8, 0x00, 0x00, 0x00, 0x00, 0x5D, 0xC3];
        let region = DisassemblyRegion::new(0x400000, 0x400100, code);
        dis.add_region(region);
        dis.disassemble().unwrap();
        assert!(dis.instruction_count() >= 5);
    }

    #[test]
    fn test_build_cfg_simple() {
        let mut dis = X86BOLTDisassembler::new(true);
        let mut entries = HashSet::new();
        entries.insert(0x400000);
        dis.set_function_entries(entries);

        // push rbp; mov rbp, rsp; xor eax, eax; pop rbp; ret
        let code = vec![
            0x55, 0x48, 0x89, 0xE5, 0x31, 0xC0, 0x5D, 0xC3,
        ];
        let region = DisassemblyRegion::new(0x400000, 0x400100, code);
        dis.add_region(region);
        dis.disassemble().unwrap();
        dis.build_cfg().unwrap();

        assert!(dis.block_count() >= 1);
    }

    #[test]
    fn test_build_cfg_with_jumps() {
        let mut dis = X86BOLTDisassembler::new(true);
        let mut entries = HashSet::new();
        entries.insert(0x400000);
        dis.set_function_entries(entries);

        // jz +5; nop; nop; jmp +3; mov eax, 1; ret
        // 74 05 90 90 EB 03 B8 01 00 00 00 C3
        let code = vec![
            0x74, 0x05, 0x90, 0x90, 0xEB, 0x03, 0xB8, 0x01, 0x00, 0x00, 0x00, 0xC3,
        ];
        let region = DisassemblyRegion::new(0x400000, 0x400100, code);
        dis.add_region(region);
        dis.disassemble().unwrap();
        dis.build_cfg().unwrap();

        // Should have multiple blocks due to jump targets
        assert!(dis.block_count() >= 2);
    }

    #[test]
    fn test_blocks_sorted() {
        let mut dis = X86BOLTDisassembler::new(true);
        let code = vec![0x90, 0x90, 0xC3];
        let region = DisassemblyRegion::new(0x400000, 0x400100, code);
        dis.add_region(region);
        dis.disassemble().unwrap();
        dis.build_cfg().unwrap();

        let sorted = dis.blocks_sorted();
        assert!(!sorted.is_empty());
        // Should be sorted by address
        for i in 1..sorted.len() {
            assert!(sorted[i].address >= sorted[i - 1].address);
        }
    }

    #[test]
    fn test_entry_blocks() {
        let mut dis = X86BOLTDisassembler::new(true);
        let mut entries = HashSet::new();
        entries.insert(0x400000);
        dis.set_function_entries(entries);

        let code = vec![0x55, 0xC3];
        let region = DisassemblyRegion::new(0x400000, 0x400100, code);
        dis.add_region(region);
        dis.disassemble().unwrap();
        dis.build_cfg().unwrap();

        let entry_blocks = dis.entry_blocks();
        assert!(!entry_blocks.is_empty());
        for block in entry_blocks {
            assert!(block.is_entry);
        }
    }

    #[test]
    fn test_containing_block() {
        let mut dis = X86BOLTDisassembler::new(true);
        let code = vec![0x55, 0x48, 0x89, 0xE5, 0xC3];
        let region = DisassemblyRegion::new(0x400000, 0x400100, code);
        dis.add_region(region);
        dis.disassemble().unwrap();
        dis.build_cfg().unwrap();

        let block = dis.containing_block(0x400000);
        assert!(block.is_some());
    }

    // ---- X86BOLTBinaryAnalysis Tests ----

    #[test]
    fn test_binary_analysis_new() {
        let data = vec![0x90, 0xC3];
        let analysis = X86BOLTBinaryAnalysis::new(data, 0x400000, true);
        assert_eq!(analysis.base_address(), 0x400000);
        assert!(analysis.is_64bit());
        assert_eq!(analysis.function_count(), 0);
    }

    #[test]
    fn test_binary_analysis_pie() {
        let data = vec![0x90, 0xC3];
        let mut analysis = X86BOLTBinaryAnalysis::new(data, 0x400000, true);
        assert!(!analysis.is_pie());
        analysis.set_pie(true);
        assert!(analysis.is_pie());
    }

    #[test]
    fn test_binary_analysis_add_symbol() {
        let data = vec![0x90; 100];
        let mut analysis = X86BOLTBinaryAnalysis::new(data, 0x400000, true);
        analysis.add_section(".text".to_string(), 0x400000, 0x400100);
        analysis.add_symbol("main".to_string(), 0x400000);
        assert_eq!(analysis.symbol_at_address(0x400000), Some("main".to_string()));
    }

    #[test]
    fn test_binary_analysis_add_section() {
        let data = vec![0x90; 100];
        let mut analysis = X86BOLTBinaryAnalysis::new(data, 0x400000, true);
        analysis.add_section(".text".to_string(), 0x400000, 0x400064);
        assert_eq!(
            analysis.section_name_at_address(0x400000),
            Some(".text".to_string())
        );
    }

    #[test]
    fn test_binary_analysis_is_executable_section() {
        let data = vec![0x90; 100];
        let mut analysis = X86BOLTBinaryAnalysis::new(data, 0x400000, true);
        analysis.add_section(".text".to_string(), 0x400000, 0x400064);
        analysis.add_section(".data".to_string(), 0x401000, 0x401064);
        assert!(analysis.is_in_executable_section(0x400000));
        assert!(!analysis.is_in_executable_section(0x401000));
    }

    #[test]
    fn test_binary_analysis_detect_function_boundaries_64() {
        let mut data = vec![0x90; 100];
        // Insert a prologue: 55 48 89 E5
        data[20] = 0x55;
        data[21] = 0x48;
        data[22] = 0x89;
        data[23] = 0xE5;
        // Insert ret
        data[30] = 0xC3;

        let mut analysis = X86BOLTBinaryAnalysis::new(data, 0x400000, true);
        analysis.add_section(".text".to_string(), 0x400000, 0x400064);
        let funcs = analysis.detect_function_boundaries();
        assert!(!funcs.is_empty());
        // Should find the function at offset 20
        assert!(funcs.iter().any(|f| f.address == 0x400000 + 20));
    }

    #[test]
    fn test_binary_analysis_analyze() {
        let mut analysis = X86BOLTBinaryAnalysis::new(vec![0x90; 200], 0x400000, true);
        analysis.add_section(".text".to_string(), 0x400000, 0x4000C8);
        analysis.add_symbol("main".to_string(), 0x400000);
        analysis.add_symbol("foo".to_string(), 0x400064);
        analysis.analyze().unwrap();
        assert!(analysis.function_count() >= 2);
    }

    #[test]
    fn test_function_info_new() {
        let func = X86FunctionInfo::new("main".to_string(), 0x400000, 128);
        assert_eq!(func.name, "main");
        assert_eq!(func.address, 0x400000);
        assert_eq!(func.size, 128);
        assert!(!func.is_entry_point);
        assert!(func.has_frame_pointer);
    }

    #[test]
    fn test_function_info_stats() {
        let mut func = X86FunctionInfo::new("test".to_string(), 0x400000, 256);
        let mut block1 = X86BoltBlock::new(0, 0x400000);
        block1.exec_count = 100;
        block1.size = 50;
        block1.is_entry = true;
        let mut block2 = X86BoltBlock::new(1, 0x400050);
        block2.exec_count = 10;
        block2.size = 30;
        func.blocks = vec![block1, block2];
        assert_eq!(func.total_exec_count(), 110);
        assert_eq!(func.total_block_size(), 80);
        assert!(func.entry_block().is_some());
        assert_eq!(func.hot_blocks(50).len(), 1);
        assert_eq!(func.cold_blocks(50).len(), 1);
        assert!(func.hottest_block().is_some());
        assert_eq!(func.hottest_block().unwrap().exec_count, 100);
    }

    // ---- X86BOLTProfile Tests ----

    #[test]
    fn test_profile_new() {
        let profile = X86BOLTProfile::new();
        assert_eq!(profile.total_samples(), 0);
        assert!(!profile.has_lbr());
        assert!(!profile.has_autofdo());
        assert_eq!(profile.hotness_threshold(), 0);
    }

    #[test]
    fn test_profile_default() {
        let profile = X86BOLTProfile::default();
        assert_eq!(profile.total_samples(), 0);
    }

    #[test]
    fn test_profile_read_lbr_samples() {
        let mut profile = X86BOLTProfile::new();
        let samples = vec![(0x400100, 0x400200), (0x400200, 0x400300)];
        profile.read_lbr_samples(&samples).unwrap();
        assert!(profile.has_lbr());
        assert_eq!(profile.total_samples(), 2);
        assert_eq!(profile.sample_count(0x400100), 1);
        assert_eq!(profile.sample_count(0x400200), 2); // source + target
    }

    #[test]
    fn test_profile_compute_hotness_threshold() {
        let mut profile = X86BOLTProfile::new();
        let samples = vec![
            (0x400100, 0x400200),
            (0x400100, 0x400200),
            (0x400100, 0x400200),
            (0x400300, 0x400400),
        ];
        profile.read_lbr_samples(&samples).unwrap();
        profile.compute_hotness_threshold();
        assert!(profile.hotness_threshold() > 0);
    }

    #[test]
    fn test_profile_is_address_hot() {
        let mut profile = X86BOLTProfile::new();
        let mut samples = Vec::new();
        // Make address 0x400100 very hot
        for _ in 0..100 {
            samples.push((0x400100, 0x400200));
        }
        // Make address 0x400300 cold
        samples.push((0x400300, 0x400400));
        profile.read_lbr_samples(&samples).unwrap();
        profile.compute_hotness_threshold();
        assert!(profile.is_address_hot(0x400100));
    }

    #[test]
    fn test_profile_generate_synthetic() {
        let mut profile = X86BOLTProfile::new();
        profile.generate_synthetic_profile(5, 3, 0.3);
        assert!(profile.total_samples() > 0);
        assert!(profile.hotness_threshold() > 0);
    }

    #[test]
    fn test_profile_annotate_blocks() {
        let mut profile = X86BOLTProfile::new();
        let samples = vec![(0x400100, 0x400200), (0x400100, 0x400300)];
        profile.read_lbr_samples(&samples).unwrap();

        let mut blocks: HashMap<u64, X86BoltBlock> = HashMap::new();
        let mut block = X86BoltBlock::new(0, 0x400100);
        block.exec_count = 0;
        blocks.insert(0x400100, block);

        let functions = vec![X86FunctionInfo::new("test".to_string(), 0x400000, 0x1000)];
        profile.annotate_blocks(&mut blocks, &functions);

        let block = blocks.get(&0x400100).unwrap();
        assert!(block.exec_count >= 2);
    }

    // ---- X86BOLTOptimizer Tests ----

    #[test]
    fn test_optimizer_new() {
        let opt = X86BOLTOptimizer::new();
        assert_eq!(opt.stats.functions_optimized, 0);
        assert_eq!(opt.stats.nops_removed, 0);
    }

    #[test]
    fn test_optimizer_default() {
        let opt = X86BOLTOptimizer::default();
        assert_eq!(opt.stats.functions_optimized, 0);
    }

    #[test]
    fn test_optimizer_set_input() {
        let mut opt = X86BOLTOptimizer::new();
        let funcs = vec![X86FunctionInfo::new("main".to_string(), 0x400000, 256)];
        let blocks: HashMap<u64, X86BoltBlock> = HashMap::new();
        let profile = X86BOLTProfile::new();
        opt.set_input(funcs, blocks, profile);
    }

    #[test]
    fn test_optimizer_set_strategy() {
        let mut opt = X86BOLTOptimizer::new();
        opt.set_function_layout(X86FunctionLayoutStrategy::HFSort);
        opt.set_function_layout(X86FunctionLayoutStrategy::PettisHansen);
        opt.set_function_layout(X86FunctionLayoutStrategy::CallChainClustering);
        opt.set_function_layout(X86FunctionLayoutStrategy::None);
    }

    #[test]
    fn test_optimizer_remove_nops_empty() {
        let mut opt = X86BOLTOptimizer::new();
        let func = X86FunctionInfo::new("empty".to_string(), 0x400000, 64);
        opt.set_input(vec![func], HashMap::new(), X86BOLTProfile::new());
        opt.remove_nop_instructions().unwrap();
        assert_eq!(opt.stats.nops_removed, 0);
    }

    #[test]
    fn test_optimizer_remove_nops_with_nops() {
        let mut opt = X86BOLTOptimizer::new();
        let mut func = X86FunctionInfo::new("with_nops".to_string(), 0x400000, 64);
        let mut block = X86BoltBlock::new(0, 0x400000);

        // Add NOPs and a real instruction
        let mut nop1 = X86BoltInstruction::new(0x400000);
        nop1.category = X86InsnCategory::Nop;
        nop1.size = 1;
        nop1.bytes = vec![0x90];

        let mut nop2 = X86BoltInstruction::new(0x400001);
        nop2.category = X86InsnCategory::Nop;
        nop2.size = 1;
        nop2.bytes = vec![0x90];

        let mut insn = X86BoltInstruction::new(0x400002);
        insn.category = X86InsnCategory::Normal;
        insn.size = 1;
        insn.bytes = vec![0xC3];

        block.add_instruction(nop1);
        block.add_instruction(nop2);
        block.add_instruction(insn);

        let initial_size = block.size;
        func.blocks = vec![block];
        opt.set_input(
            vec![func],
            HashMap::new(),
            X86BOLTProfile::new(),
        );
        opt.remove_nop_instructions().unwrap();
        assert!(opt.stats.nops_removed >= 2);
    }

    #[test]
    fn test_optimizer_eliminate_useless_jumps() {
        let mut opt = X86BOLTOptimizer::new();
        let mut func = X86FunctionInfo::new("with_jump".to_string(), 0x400000, 64);
        let mut block = X86BoltBlock::new(0, 0x400000);

        // Add a useless jump (jump to next instruction)
        let mut jmp = X86BoltInstruction::new(0x400000);
        jmp.category = X86InsnCategory::UnconditionalJump;
        jmp.size = 2;
        jmp.branch_target = Some(0x400002);
        jmp.fallthrough = Some(0x400002);
        jmp.bytes = vec![0xEB, 0x00];

        block.add_instruction(jmp);
        func.blocks = vec![block];
        opt.set_input(
            vec![func],
            HashMap::new(),
            X86BOLTProfile::new(),
        );
        opt.eliminate_useless_jumps().unwrap();
        assert!(opt.stats.jumps_eliminated >= 1);
    }

    #[test]
    fn test_optimizer_optimize_frame() {
        let mut opt = X86BOLTOptimizer::new();
        let mut func = X86FunctionInfo::new("leaf".to_string(), 0x400000, 64);
        // Simulate a simple leaf function with no calls
        let mut block = X86BoltBlock::new(0, 0x400000);
        let mut insn = X86BoltInstruction::new(0x400000);
        insn.category = X86InsnCategory::Normal;
        insn.is_call = false;
        insn.size = 1;
        insn.bytes = vec![0x90];
        block.add_instruction(insn);
        func.blocks = vec![block];

        opt.set_input(
            vec![func],
            HashMap::new(),
            X86BOLTProfile::new(),
        );
        opt.optimize_frame_setup().unwrap();
        // Leaf function should have frame pointer removed
        assert!(opt.stats.frame_opts_applied >= 1);
    }

    #[test]
    fn test_optimizer_split_hot_cold() {
        let mut opt = X86BOLTOptimizer::new();
        let mut func = X86FunctionInfo::new("mixed".to_string(), 0x400000, 64);
        let mut hot_block = X86BoltBlock::new(0, 0x400000);
        hot_block.exec_count = 1000;
        let mut cold_block = X86BoltBlock::new(1, 0x400020);
        cold_block.exec_count = 1;
        func.blocks = vec![hot_block, cold_block];

        opt.set_input(
            vec![func],
            HashMap::new(),
            X86BOLTProfile::new(),
        );
        opt.split_hot_cold_functions(50).unwrap();
        let func_opt = &opt.optimized_functions()[0];
        assert!(func_opt.blocks[0].exec_count >= 50);
    }

    #[test]
    fn test_optimizer_optimize_empty() {
        let mut opt = X86BOLTOptimizer::new();
        opt.optimize().unwrap();
        assert_eq!(opt.stats.functions_optimized, 0);
    }

    // ---- X86BOLTRewriter Tests ----

    #[test]
    fn test_rewriter_new() {
        let data = vec![0x90; 4096];
        let rewriter = X86BOLTRewriter::new(data, 0x400000, true);
        assert_eq!(rewriter.address_map().len(), 0);
    }

    #[test]
    fn test_rewriter_emit_text_section_empty() {
        let data = vec![0x90; 4096];
        let mut rewriter = X86BOLTRewriter::new(data, 0x400000, true);
        rewriter.set_functions(vec![]);
        let text = rewriter.emit_text_section().unwrap();
        assert!(text.is_empty());
    }

    #[test]
    fn test_rewriter_emit_text_section() {
        let data = vec![0x90; 4096];
        let mut rewriter = X86BOLTRewriter::new(data, 0x400000, true);

        let func = X86FunctionInfo::new("test".to_string(), 0x400000, 64);
        let mut block = X86BoltBlock::new(0, 0x400000);
        let mut insn = X86BoltInstruction::new(0x400000);
        insn.size = 1;
        insn.bytes = vec![0xC3]; // ret
        insn.category = X86InsnCategory::Return;
        block.add_instruction(insn);

        let mut func_with_blocks = func.clone();
        func_with_blocks.blocks = vec![block];

        rewriter.set_functions(vec![func_with_blocks]);
        let text = rewriter.emit_text_section().unwrap();
        assert!(!text.is_empty());
        assert!(text.contains(&0xC3));
    }

    #[test]
    fn test_rewriter_emit_symbols() {
        let data = vec![0x90; 4096];
        let mut rewriter = X86BOLTRewriter::new(data, 0x400000, true);

        let func = X86FunctionInfo::new("main".to_string(), 0x400000, 64);
        rewriter.set_functions(vec![func.clone()]);
        rewriter.emit_text_section().unwrap();
        let symbols = rewriter.emit_symbols();
        assert_eq!(symbols.len(), 1);
        assert_eq!(symbols[0].name, "main");
    }

    #[test]
    fn test_rewriter_verify_empty() {
        let data = vec![0x90; 4096];
        let rewriter = X86BOLTRewriter::new(data, 0x400000, true);
        assert!(rewriter.verify().is_err()); // text section is empty
    }

    #[test]
    fn test_rewriter_write_binary() {
        let data = make_simple_64bit_binary();
        let mut rewriter = X86BOLTRewriter::new(data, 0x400000, true);

        let func = X86FunctionInfo::new("test".to_string(), 0x401000, 64);
        let mut block = X86BoltBlock::new(0, 0x401000);
        let mut insn = X86BoltInstruction::new(0x401000);
        insn.size = 1;
        insn.bytes = vec![0xC3];
        insn.category = X86InsnCategory::Return;
        block.add_instruction(insn);

        let mut func_with_blocks = func;
        func_with_blocks.blocks = vec![block];

        rewriter.set_functions(vec![func_with_blocks]);
        let binary = rewriter.write_binary().unwrap();
        assert!(!binary.is_empty());
    }

    // ---- BOLTX86 Tests ----

    #[test]
    fn test_boltx86_new() {
        let data = make_simple_64bit_binary();
        let bolt = BOLTX86::new(data, 0x400000, true);
        assert!(!bolt.is_optimized());
        assert_eq!(bolt.stats().functions_optimized, 0);
    }

    #[test]
    fn test_boltx86_config() {
        let data = make_simple_64bit_binary();
        let mut bolt = BOLTX86::new(data, 0x400000, true);
        bolt.config_mut().remove_nops = false;
        assert!(!bolt.config().remove_nops);
    }

    #[test]
    fn test_boltx86_add_symbol() {
        let data = make_simple_64bit_binary();
        let mut bolt = BOLTX86::new(data, 0x400000, true);
        bolt.add_symbol("main", 0x401000);
    }

    #[test]
    fn test_boltx86_add_section() {
        let data = make_simple_64bit_binary();
        let mut bolt = BOLTX86::new(data, 0x400000, true);
        bolt.add_section(".text", 0x401000, 0x401100);
    }

    #[test]
    fn test_boltx86_generate_synthetic_profile() {
        let data = make_simple_64bit_binary();
        let mut bolt = BOLTX86::new(data, 0x400000, true);
        bolt.generate_synthetic_profile(3, 2, 0.5);
        assert!(bolt.profile().total_samples() > 0);
    }

    #[test]
    fn test_boltx86_validate_not_optimized() {
        let data = make_simple_64bit_binary();
        let bolt = BOLTX86::new(data, 0x400000, true);
        assert!(bolt.validate().is_err());
    }

    #[test]
    fn test_boltx86_optimize_with_synthetic() {
        let data = vec![0x90; 8192];
        // Make it look like a minimal ELF
        let mut bin_data = elf64_header().to_vec();
        while bin_data.len() < 4096 {
            bin_data.push(0);
        }
        // Add some fake code at the .text section area
        // push rbp; mov rbp, rsp; xor eax, eax; pop rbp; ret
        let code: [u8; 9] = [0x55, 0x48, 0x89, 0xE5, 0x31, 0xC0, 0x5D, 0xC3, 0x90];
        bin_data.extend_from_slice(&code);

        let mut bolt = BOLTX86::new(bin_data, 0x400000, true);
        bolt.add_section(".text", 0x404000, 0x404100);
        bolt.add_symbol("_start", 0x404000);

        // Generate synthetic profile
        bolt.generate_synthetic_profile(1, 2, 0.5);

        let result = bolt.optimize();
        assert!(result.is_ok());
        let optimized = result.unwrap();
        assert!(!optimized.is_empty());
        assert!(bolt.is_optimized());
        assert!(bolt.stats().functions_optimized > 0 || bolt.stats().nops_removed > 0);
    }

    #[test]
    fn test_elf64_header() {
        let hdr = elf64_header();
        assert_eq!(hdr[0], 0x7F);
        assert_eq!(hdr[1], b'E');
        assert_eq!(hdr[2], b'L');
        assert_eq!(hdr[3], b'F');
        assert_eq!(hdr[4], 2); // ELFCLASS64
    }

    #[test]
    fn test_elf32_header() {
        let hdr = elf32_header();
        assert_eq!(hdr[0], 0x7F);
        assert_eq!(hdr[1], b'E');
        assert_eq!(hdr[2], b'L');
        assert_eq!(hdr[3], b'F');
        assert_eq!(hdr[4], 1); // ELFCLASS32
    }

    // ---- X86FunctionLayoutStrategy Tests ----

    #[test]
    fn test_layout_strategy_to_layout_algorithm() {
        let algo: LayoutAlgorithm = X86FunctionLayoutStrategy::HFSort.into();
        assert_eq!(algo, LayoutAlgorithm::HFSort);

        let algo: LayoutAlgorithm = X86FunctionLayoutStrategy::PettisHansen.into();
        assert_eq!(algo, LayoutAlgorithm::PettisHansen);
    }

    // ---- BOLTX86Config Tests ----

    #[test]
    fn test_config_default() {
        let config = BOLTX86Config::default();
        assert!(config.reorder_blocks);
        assert!(config.remove_nops);
        assert!(config.eliminate_jumps);
        assert!(config.optimize_frame);
        assert!(config.split_functions);
        assert_eq!(config.cache_line_size, 64);
        assert_eq!(config.hotness_threshold, 0);
    }

    // ---- Edge Cases ----

    #[test]
    fn test_disassemble_empty_region() {
        let mut dis = X86BOLTDisassembler::new(true);
        let region = DisassemblyRegion::new(0x400000, 0x400000, vec![]);
        dis.add_region(region);
        let result = dis.disassemble();
        assert!(result.is_ok());
        assert_eq!(dis.instruction_count(), 0);
    }

    #[test]
    fn test_build_cfg_empty_instructions() {
        let mut dis = X86BOLTDisassembler::new(true);
        let result = dis.build_cfg();
        assert!(result.is_ok());
        assert_eq!(dis.block_count(), 0);
    }

    #[test]
    fn test_decode_invalid_address() {
        let dis = X86BOLTDisassembler::new(true);
        let result = dis.decode_instruction(0x500000);
        assert!(result.is_err());
    }

    #[test]
    fn test_optimizer_reorder_blocks_single() {
        let mut opt = X86BOLTOptimizer::new();
        let mut func = X86FunctionInfo::new("single".to_string(), 0x400000, 64);
        let block = X86BoltBlock::new(0, 0x400000);
        func.blocks = vec![block];
        opt.set_input(vec![func], HashMap::new(), X86BOLTProfile::new());
        opt.reorder_basic_blocks(10).unwrap();
        // Single block should remain unchanged
        assert_eq!(opt.optimized_functions()[0].blocks.len(), 1);
    }

    #[test]
    fn test_optimizer_large_hot_threshold() {
        let mut opt = X86BOLTOptimizer::new();
        let mut func = X86FunctionInfo::new("test".to_string(), 0x400000, 64);
        let mut block = X86BoltBlock::new(0, 0x400000);
        block.exec_count = 5;
        func.blocks = vec![block];
        opt.set_input(vec![func], HashMap::new(), X86BOLTProfile::new());
        opt.split_hot_cold_functions(1000).unwrap();
        // All blocks should be cold since threshold > exec_count
        assert!(opt.optimized_functions()[0].blocks[0].is_cold);
    }

    // ---- OptimizationStats Tests ----

    #[test]
    fn test_optimization_stats_default() {
        let stats = OptimizationStats::default();
        assert_eq!(stats.functions_optimized, 0);
        assert_eq!(stats.blocks_reordered, 0);
        assert_eq!(stats.nops_removed, 0);
        assert_eq!(stats.jumps_eliminated, 0);
        assert_eq!(stats.frame_opts_applied, 0);
        assert_eq!(stats.bytes_saved, 0);
        assert_eq!(stats.stale_matching_rate, 0.0);
    }

    // ---- X86BoltBlock Display Tests ----

    #[test]
    fn test_bolt_block_display() {
        let mut block = X86BoltBlock::new(5, 0x400100);
        block.is_entry = true;
        block.exec_count = 100;
        block.size = 32;
        block.num_instructions = 8;
        let s = format!("{}", block);
        assert!(s.contains("BB5"));
        assert!(s.contains("entry"));
        assert!(s.contains("count=100"));
    }

    // ---- X86InsnCategory Tests ----

    #[test]
    fn test_insn_category_equality() {
        assert_eq!(X86InsnCategory::UnconditionalJump, X86InsnCategory::UnconditionalJump);
        assert_ne!(X86InsnCategory::Normal, X86InsnCategory::Nop);
    }

    // ---- X86BoltInstruction Default Tests ----

    #[test]
    fn test_instruction_default() {
        let insn = X86BoltInstruction::default();
        assert_eq!(insn.address, 0);
        assert!(!insn.is_call);
        assert!(!insn.is_return);
    }

    // ---- DisassemblyRegion Edge Case Tests ----

    #[test]
    fn test_region_read_bytes_partial() {
        let region = DisassemblyRegion::new(0x400000, 0x400010, vec![0x55, 0x48, 0x89, 0xE5]);
        // Read more bytes than available
        let bytes = region.read_bytes(0x400002, 10);
        assert_eq!(bytes.len(), 2); // only 0x89, 0xE5 available
    }

    #[test]
    fn test_disassembly_region_read_bytes_out_of_range() {
        let region = DisassemblyRegion::new(0x400000, 0x400010, vec![0x55, 0x48]);
        let bytes = region.read_bytes(0x500000, 10);
        assert!(bytes.is_empty());
    }
}