llvm-native-core 0.1.14

//! BOLT Binary Rewriter — binary-level disassembly, CFG reconstruction,
//! function/basic-block layout optimization, branch optimization,
//! code reordering, relocation/symbol/debug info updates.
//!
//! Clean-room behavioral reconstruction. Zero BOLT source consultation.

use crate::bolt::bolt_profile::BoltProfile;
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap, HashSet, VecDeque};

// ============================================================================
// Layout Algorithm Enum
// ============================================================================

/// Function layout algorithms supported by BOLT.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LayoutAlgorithm {
    /// HFSort — optimal cache-line-aware function ordering (Ottoni & Maher 2017)
    HFSort,
    /// HFSort+ — enhanced HFSort with temporal profiling
    HFSortPlus,
    /// CacheSort — cache-line-conscious ordering
    CacheSort,
    /// Pettis-Hansen — call-graph-based clustering (Pettis & Hansen 1990)
    PettisHansen,
    /// Call-Chain Clustering — group functions based on call chains
    CallChainClustering,
}

// ============================================================================
// BoltBlock — a basic block in the binary
// ============================================================================

/// A basic block within a function in the binary.
#[derive(Debug, Clone)]
pub struct BoltBlock {
    /// Virtual address of this block.
    pub address: u64,
    /// Offset within the function's code.
    pub offset: usize,
    /// Size in bytes of this block.
    pub size: usize,
    /// Indices of successor blocks.
    pub successors: Vec<usize>,
    /// Execution count from profiling.
    pub execution_count: u64,
    /// Whether this block is the function entry.
    pub is_entry: bool,
    /// Whether this block is a function exit (has a return/tail-call).
    pub is_exit: bool,
}

impl BoltBlock {
    pub fn new(offset: usize, size: usize) -> Self {
        Self {
            address: 0,
            offset,
            size,
            successors: Vec::new(),
            execution_count: 0,
            is_entry: false,
            is_exit: false,
        }
    }
}

// ============================================================================
// BoltFunction — a function in the binary
// ============================================================================

/// A function identified in the binary, with its basic blocks.
#[derive(Debug, Clone)]
pub struct BoltFunction {
    /// Symbol name of the function.
    pub name: String,
    /// Virtual address (from symbol table / disassembly).
    pub address: u64,
    /// Total size of the function in bytes.
    pub size: u64,
    /// Basic blocks within this function.
    pub blocks: Vec<BoltBlock>,
    /// Execution count from profile data.
    pub execution_count: u64,
    /// Whether this function is classified as "hot".
    pub is_hot: bool,
}

impl BoltFunction {
    pub fn new(name: String, address: u64, size: u64) -> Self {
        Self {
            name,
            address,
            size,
            blocks: Vec::new(),
            execution_count: 0,
            is_hot: false,
        }
    }

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

    /// Get a block by index.
    pub fn get_block(&self, index: usize) -> Option<&BoltBlock> {
        self.blocks.get(index)
    }

    /// Get a mutable block by index.
    pub fn get_block_mut(&mut self, index: usize) -> Option<&mut BoltBlock> {
        self.blocks.get_mut(index)
    }

    /// Compute the call-graph edge weight from this function to another.
    pub fn edge_weight_to(&self, callee_name: &str) -> u64 {
        // In a full implementation, this reads call-count profiles.
        // For now, return a placeholder based on execution count.
        if self.is_hot {
            self.execution_count / 10
        } else {
            0
        }
    }
}

// ============================================================================
// BOLTBinaryRewriter — the main binary optimizer
// ============================================================================

/// The BOLT binary rewriter: reads a compiled binary, disassembles it,
/// builds CFGs, runs optimizations, and emits an optimized binary.
pub struct BOLTBinaryRewriter {
    /// Input binary data.
    pub input: Vec<u8>,
    /// Output buffer for the rewritten binary.
    pub output: Vec<u8>,
    /// All functions identified in the binary.
    pub functions: Vec<BoltFunction>,
    /// Reordered function names (the new layout order).
    pub reordered_functions: Vec<String>,
    /// Which layout algorithm to use.
    pub layout_algorithm: LayoutAlgorithm,
    /// Profile data for this binary, if loaded.
    profile: Option<BoltProfile>,
}

impl BOLTBinaryRewriter {
    /// Create a new BOLTBinaryRewriter from raw binary data.
    pub fn new(data: Vec<u8>) -> Self {
        Self {
            input: data,
            output: Vec::new(),
            functions: Vec::new(),
            reordered_functions: Vec::new(),
            layout_algorithm: LayoutAlgorithm::HFSort,
            profile: None,
        }
    }

    /// Set the layout algorithm.
    pub fn with_layout(mut self, algo: LayoutAlgorithm) -> Self {
        self.layout_algorithm = algo;
        self
    }

    /// Attach profile data for profile-guided optimization.
    pub fn with_profile(mut self, profile: BoltProfile) -> Self {
        self.profile = Some(profile);
        self
    }

    // ========================================================================
    // Phase 1: Disassembly
    // ========================================================================

    /// Disassemble the binary, identifying functions and instructions.
    /// This is a simplified x86-64 disassembler for demonstration.
    pub fn disassemble(&mut self) -> Result<(), String> {
        self.functions.clear();
        let mut offset = 0usize;
        let data = &self.input;

        while offset < data.len() {
            // Look for function prologue patterns (e.g., `push rbp; mov rbp, rsp`)
            if self.is_function_start(&data[offset..]) {
                if let Some(func) = self.disassemble_function_from(offset) {
                    offset = (func.address as usize + func.size as usize).min(data.len());
                    self.functions.push(func);
                    continue;
                }
            }
            offset += 1;
        }

        Ok(())
    }

    /// Heuristic: detect function start (typical x86-64 prologue).
    fn is_function_start(&self, bytes: &[u8]) -> bool {
        if bytes.len() < 4 {
            return false;
        }
        // `push rbp` = 0x55
        // `mov rbp, rsp` = 0x48 0x89 0xE5
        // or `push rbp; mov rbp, rsp; sub rsp, imm`
        bytes[0] == 0x55
            && ((bytes.len() >= 4 && bytes[1] == 0x48 && bytes[2] == 0x89 && bytes[3] == 0xE5)
                || (bytes.len() >= 2 && bytes[1] == 0x48 && bytes[2] == 0x8B && bytes[3] == 0xEC))
    }

    /// Disassemble a single function starting at the given offset.
    fn disassemble_function_from(&self, start: usize) -> Option<BoltFunction> {
        let data = &self.input;
        if start >= data.len() {
            return None;
        }

        let address = start as u64;
        // Estimate function size by scanning until the next function start
        // or a `ret` (0xC3) / `retf` (0xCB) / `iret` (0xCF).
        let mut end = start;
        let max_scan = (data.len() - start).min(0x100000); // 1 MB max

        for i in start..data.len().min(start + max_scan) {
            if data[i] == 0xC3 || data[i] == 0xCB {
                // Found a return; check if next byte looks like a new function
                end = i + 1;
                // If followed by NOP padding or alignment, extend past it
                while end < data.len()
                    && (data[end] == 0x90 || data[end] == 0xCC || data[end] == 0x00)
                {
                    end += 1;
                }
                break;
            }
        }

        if end == start {
            end = (start + 64).min(data.len()); // minimum fallback
        }

        let size = (end - start) as u64;

        // Generate a synthetic name.
        let name = format!("func_{:016X}", address);
        let mut func = BoltFunction::new(name, address, size);

        // Perform rudimentary basic block detection by finding branch targets.
        self.identify_basic_blocks(&mut func);
        Some(func)
    }

    /// Identify basic blocks within a function by scanning for terminators
    /// and branch targets.
    fn identify_basic_blocks(&self, func: &mut BoltFunction) {
        let data = &self.input;
        let start = func.address as usize;
        let end = (func.address + func.size) as usize;

        if start >= data.len() || end > data.len() || end <= start {
            return;
        }

        // Find all instruction boundaries and branch targets.
        let mut targets: HashSet<usize> = HashSet::new();
        targets.insert(0); // Function entry is always a block start.

        // Scan for branch instructions and their targets.
        let mut pos = start;
        while pos < end {
            let offset = pos - start;
            let byte = data[pos];

            match byte {
                // Conditional jumps (Jcc rel8/rel32)
                0x70..=0x7F => {
                    if pos + 1 < end {
                        let rel = data[pos + 1] as i8;
                        let target = (pos as i64 + 2 + rel as i64) as usize;
                        // Fallthrough is also a block start
                        if pos + 2 <= end {
                            targets.insert(offset + 2);
                        }
                        if target >= start && target < end {
                            targets.insert(target - start);
                        }
                        pos += 2;
                    } else {
                        pos += 1;
                    }
                }
                // Near conditional jumps (0F 8x rel32)
                0x0F => {
                    if pos + 1 < end {
                        let next = data[pos + 1];
                        if (0x80..=0x8F).contains(&next) {
                            if pos + 5 < end {
                                let rel = i32::from_le_bytes([
                                    data[pos + 2],
                                    data[pos + 3],
                                    data[pos + 4],
                                    data[pos + 5],
                                ]);
                                let target = (pos as i64 + 6 + rel as i64) as usize;
                                targets.insert(offset + 6); // fallthrough
                                if target >= start && target < end {
                                    targets.insert(target - start);
                                }
                                pos += 6;
                                continue;
                            }
                        }
                    }
                    pos += 2;
                }
                // Unconditional jump (E9 rel32)
                0xE9 => {
                    if pos + 4 < end {
                        let rel = i32::from_le_bytes([
                            data[pos + 1],
                            data[pos + 2],
                            data[pos + 3],
                            data[pos + 4],
                        ]);
                        let target = (pos as i64 + 5 + rel as i64) as usize;
                        if target >= start && target < end {
                            targets.insert(target - start);
                        }
                        pos += 5;
                    } else {
                        pos += 1;
                    }
                }
                // Unconditional jump (EB rel8)
                0xEB => {
                    if pos + 1 < end {
                        let rel = data[pos + 1] as i8;
                        let target = (pos as i64 + 2 + rel as i64) as usize;
                        if target >= start && target < end {
                            targets.insert(target - start);
                        }
                        pos += 2;
                    } else {
                        pos += 1;
                    }
                }
                // Return (C3) — this block ends here.
                0xC3 => {
                    if pos + 1 <= end {
                        targets.insert(offset + 1); // after ret
                    }
                    pos += 1;
                }
                // Call (E8 rel32) — next instruction is a block start.
                0xE8 => {
                    pos += if pos + 4 < end { 5 } else { 1 };
                    if pos <= end {
                        targets.insert(pos - start);
                    }
                }
                // Handle REX prefixes (0x40-0x4F) — skip and try next byte.
                0x40..=0x4F => {
                    pos += 1;
                }
                // Handle two-byte opcodes (0x0F) that aren't Jcc.
                _ => {
                    pos += 1; // simplified: advance by 1; real decoder needed
                }
            }
        }

        // Sort targets and create blocks.
        let mut sorted_targets: Vec<usize> = targets.into_iter().collect();
        sorted_targets.sort();

        for i in 0..sorted_targets.len() {
            let block_start = sorted_targets[i];
            let block_end = if i + 1 < sorted_targets.len() {
                sorted_targets[i + 1]
            } else {
                func.size as usize
            };
            let size = block_end.saturating_sub(block_start);
            let mut block = BoltBlock::new(block_start, size);
            block.address = func.address + block_start as u64;
            block.is_entry = block_start == 0;
            func.blocks.push(block);
        }

        // Connect successors by re-scanning.
        self.connect_block_successors(func);
    }

    /// Connect block successors by scanning the last few bytes of each block
    /// for branch instructions.
    fn connect_block_successors(&self, func: &mut BoltFunction) {
        let data = &self.input;
        let block_count = func.blocks.len();

        // Build offset-to-index map.
        let offset_to_idx: HashMap<usize, usize> = func
            .blocks
            .iter()
            .enumerate()
            .map(|(i, b)| (b.offset, i))
            .collect();

        for i in 0..block_count {
            let block = &func.blocks[i];
            let block_end = block.offset + block.size;
            if block_end < 2 || block_end > data.len() {
                continue;
            }

            // Find successor from last instruction.
            let last_byte_idx = block_end.saturating_sub(1);

            // Check for unconditional jump (E9 at end).
            if block.size >= 5 {
                let jmp_pos = block_end - 5;
                if data[jmp_pos] == 0xE9 {
                    let rel = i32::from_le_bytes([
                        data[jmp_pos + 1],
                        data[jmp_pos + 2],
                        data[jmp_pos + 3],
                        data[jmp_pos + 4],
                    ]);
                    let target_abs = jmp_pos as i64 + 5 + rel as i64;
                    let target_off = target_abs as usize - func.address as usize;
                    if let Some(&idx) = offset_to_idx.get(&target_off) {
                        func.blocks[i].successors.push(idx);
                    }
                    continue;
                }
            }

            // Check for short jump (EB at end).
            if block.size >= 2 {
                let jmp_pos = block_end - 2;
                if data[jmp_pos] == 0xEB {
                    let rel = data[jmp_pos + 1] as i8;
                    let target_abs = jmp_pos as i64 + 2 + rel as i64;
                    let target_off = target_abs as usize - func.address as usize;
                    if let Some(&idx) = offset_to_idx.get(&target_off) {
                        func.blocks[i].successors.push(idx);
                    }
                    continue;
                }
            }

            // Check for conditional jump (0F 8x at end-6).
            if block.size >= 6 {
                let jmp_pos = block_end - 6;
                if data[jmp_pos] == 0x0F && (0x80..=0x8F).contains(&data[jmp_pos + 1]) {
                    let rel = i32::from_le_bytes([
                        data[jmp_pos + 2],
                        data[jmp_pos + 3],
                        data[jmp_pos + 4],
                        data[jmp_pos + 5],
                    ]);
                    let target_abs = jmp_pos as i64 + 6 + rel as i64;
                    let target_off = target_abs as usize - func.address as usize;
                    if let Some(&idx) = offset_to_idx.get(&target_off) {
                        func.blocks[i].successors.push(idx);
                    }
                    // Fallthrough is next block
                    if i + 1 < block_count {
                        func.blocks[i].successors.push(i + 1);
                    }
                    continue;
                }
            }

            // Check for conditional jump (7x rel8 at end-2).
            if block.size >= 2 {
                let jmp_pos = block_end - 2;
                if (0x70..=0x7F).contains(&data[jmp_pos]) {
                    let rel = data[jmp_pos + 1] as i8;
                    let target_abs = jmp_pos as i64 + 2 + rel as i64;
                    let target_off = target_abs as usize - func.address as usize;
                    if let Some(&idx) = offset_to_idx.get(&target_off) {
                        func.blocks[i].successors.push(idx);
                    }
                    if i + 1 < block_count {
                        func.blocks[i].successors.push(i + 1);
                    }
                    continue;
                }
            }

            // If not a jump, it's fallthrough to the next block (or return).
            if data[last_byte_idx] == 0xC3 {
                func.blocks[i].is_exit = true;
            } else if i + 1 < block_count {
                func.blocks[i].successors.push(i + 1);
            }
        }
    }

    // ========================================================================
    // Phase 2: CFG Construction
    // ========================================================================

    /// Build control-flow graphs for all functions.
    pub fn build_cfg(&mut self) {
        for func in &mut self.functions {
            build_function_cfg_inner(func);
        }
    }

    // ========================================================================
    // Phase 3: Run Optimizations
    // ========================================================================

    /// Run all BOLT optimizations.
    pub fn run_optimizations(&mut self) {
        // 1. Annotate functions with profile data.
        {
            let profile_clone = self.profile.clone();
            if let Some(ref profile) = profile_clone {
                annotate_from_profile_static(profile, &mut self.functions);
            }
        }

        // 2. Optimize function layout (reorder functions).
        self.optimize_function_layout();

        // 3. Optimize basic block layout within hot functions.
        // First collect the new orders (immutable borrow of self.functions),
        // then apply them (mutable borrow).
        let mut block_orders: Vec<(usize, Vec<usize>)> = Vec::new();
        for (i, func) in self.functions.iter().enumerate() {
            if func.is_hot {
                let new_order = self.optimize_basic_block_layout(func);
                block_orders.push((i, new_order));
            }
        }
        for (i, order) in block_orders {
            reorder_blocks_static(&mut self.functions[i], &order);
        }

        // 4. Branch optimizations.
        self.optimize_branches();

        // 5. NOP removal.
        self.optimize_nops();

        // 6. Indirect branch optimization.
        self.optimize_indirect_branches();

        // 7. Cold code stitching.
        self.stitch_cold_code();
    }

    /// Annotate functions with profile execution counts.
    fn annotate_from_profile(&mut self, profile: &BoltProfile) {
        for func in &mut self.functions {
            if let Some(fp) = profile.functions.get(&func.name) {
                func.execution_count = fp.entry_count;
                func.is_hot = fp.is_hot;
                for block in &mut func.blocks {
                    if let Some(&count) = fp.block_counts.get(&block.address) {
                        block.execution_count = count;
                    }
                }
            }
        }
    }

    // ========================================================================
    // Function Layout Optimization
    // ========================================================================

    /// Optimize function layout using the selected algorithm.
    fn optimize_function_layout(&mut self) {
        let order: Vec<String> = match self.layout_algorithm {
            LayoutAlgorithm::HFSort => self.compute_hfsort_clusters(),
            LayoutAlgorithm::HFSortPlus => self.compute_hfsort_plus_clusters(),
            LayoutAlgorithm::CacheSort => self.compute_cache_sort(),
            LayoutAlgorithm::PettisHansen => self.compute_pettis_hansen(),
            LayoutAlgorithm::CallChainClustering => self.compute_call_chain_clusters(),
        };

        self.reordered_functions = order;
    }

    /// HFSort: cluster functions by call frequency, then sort within clusters
    /// for optimal cache-line utilization.
    fn compute_hfsort_clusters(&self) -> Vec<String> {
        let func_count = self.functions.len();
        if func_count == 0 {
            return Vec::new();
        }

        // Build call graph with edge weights.
        let mut edges: Vec<(usize, usize, u64)> = Vec::new();
        for i in 0..func_count {
            for j in 0..func_count {
                if i == j {
                    continue;
                }
                let w = self.functions[i].execution_count + self.functions[j].execution_count;
                if w > 0 {
                    edges.push((i, j, w));
                }
            }
        }

        // Sort edges by weight descending.
        edges.sort_by_key(|(_, _, w)| Reverse(*w));

        // Union-Find for clustering.
        let mut parent: Vec<usize> = (0..func_count).collect();
        let mut rank: Vec<usize> = vec![0; func_count];

        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], rank: &mut [usize], x: usize, y: usize) {
            let rx = find(parent, x);
            let ry = find(parent, y);
            if rx == ry {
                return;
            }
            if rank[rx] < rank[ry] {
                parent[rx] = ry;
            } else if rank[rx] > rank[ry] {
                parent[ry] = rx;
            } else {
                parent[ry] = rx;
                rank[rx] += 1;
            }
        }

        // Cluster by merging high-weight edges.
        let cluster_limit = func_count.min(100);
        let mut clusters_found = func_count;

        for (i, j, _w) in &edges {
            if clusters_found <= cluster_limit {
                break;
            }
            if find(&mut parent, *i) != find(&mut parent, *j) {
                union(&mut parent, &mut rank, *i, *j);
                clusters_found -= 1;
            }
        }

        // Build clusters.
        let mut cluster_map: HashMap<usize, Vec<usize>> = HashMap::new();
        for i in 0..func_count {
            let root = find(&mut parent, i);
            cluster_map.entry(root).or_default().push(i);
        }

        // Within each cluster, sort hot functions first, then by size.
        let mut result: Vec<String> = Vec::new();
        for indices in cluster_map.values_mut() {
            indices.sort_by_key(|&i| {
                let func = &self.functions[i];
                // Hot first, then larger functions (cache-friendly).
                (!func.is_hot, Reverse(func.size))
            });
            for &i in indices.iter() {
                result.push(self.functions[i].name.clone());
            }
        }

        result
    }

    /// HFSort+ variant with temporal profiling awareness.
    fn compute_hfsort_plus_clusters(&self) -> Vec<String> {
        // HFSort+ extends HFSort with temporal information from
        // LBR (Last Branch Record) traces. It considers not just
        // call frequency but also temporal proximity of calls.
        //
        // Simplified: use HFSort but weight hot→hot edges higher.
        let func_count = self.functions.len();
        if func_count == 0 {
            return Vec::new();
        }

        let mut edges: Vec<(usize, usize, f64)> = Vec::new();
        for i in 0..func_count {
            for j in 0..func_count {
                if i == j {
                    continue;
                }
                let fi = &self.functions[i];
                let fj = &self.functions[j];
                let base = fi.execution_count + fj.execution_count;
                let mut weight = base as f64;
                // Temporal bonus: hot→hot edges get 2x.
                if fi.is_hot && fj.is_hot {
                    weight *= 2.0;
                }
                if weight > 0.0 {
                    edges.push((i, j, weight));
                }
            }
        }

        edges.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal));

        // Greedy clustering.
        let mut parent: Vec<usize> = (0..func_count).collect();
        let mut rank: Vec<usize> = vec![0; func_count];
        let cluster_limit = func_count.min(80);

        let mut merged = func_count;
        for (i, j, _) in &edges {
            if merged <= cluster_limit {
                break;
            }
            let ri = find(&mut parent, *i);
            let rj = find(&mut parent, *j);
            if ri != rj {
                union(&mut parent, &mut rank, *i, *j);
                merged -= 1;
            }
        }

        let mut cluster_map: HashMap<usize, Vec<usize>> = HashMap::new();
        for i in 0..func_count {
            let root = find(&mut parent, i);
            cluster_map.entry(root).or_default().push(i);
        }

        let mut result: Vec<String> = Vec::new();
        for indices in cluster_map.values_mut() {
            indices.sort_by_key(|&i| {
                let f = &self.functions[i];
                (!f.is_hot, Reverse(f.execution_count))
            });
            for &i in indices.iter() {
                result.push(self.functions[i].name.clone());
            }
        }

        result
    }

    /// CacheSort: order functions for optimal I-cache utilization.
    /// Place hot functions contiguously at the beginning.
    fn compute_cache_sort(&self) -> Vec<String> {
        let mut indices: Vec<usize> = (0..self.functions.len()).collect();

        // Sort: hot first, then by execution count descending.
        indices.sort_by_key(|&i| {
            let f = &self.functions[i];
            (!f.is_hot, Reverse(f.execution_count))
        });

        indices
            .iter()
            .map(|&i| self.functions[i].name.clone())
            .collect()
    }

    /// Pettis-Hansen: build a call graph and use closest-is-best heuristic
    /// to place functions that frequently call each other near each other.
    fn compute_pettis_hansen(&self) -> Vec<String> {
        let n = self.functions.len();
        if n == 0 {
            return Vec::new();
        }

        // Build a weighted undirected call graph.
        let mut graph: Vec<Vec<u64>> = vec![vec![0; n]; n];
        for i in 0..n {
            for j in (i + 1)..n {
                let w = self.functions[i].execution_count + self.functions[j].execution_count;
                graph[i][j] = w;
                graph[j][i] = w;
            }
        }

        // Greedy merge: repeatedly merge the pair with the highest weight.
        // Each "chain" is a list of function indices.
        let mut chains: Vec<Vec<usize>> = (0..n).map(|i| vec![i]).collect();
        let mut active: HashSet<usize> = (0..n).collect();

        while active.len() > 1 {
            // Find the pair of chains with the highest inter-chain weight.
            let mut best: Option<(usize, usize, u64)> = None;
            let act: Vec<usize> = active.iter().copied().collect();

            for a_idx in 0..act.len() {
                for b_idx in (a_idx + 1)..act.len() {
                    let ci = act[a_idx];
                    let cj = act[b_idx];
                    // Sum weights between all nodes in chain i and chain j.
                    let mut total = 0u64;
                    for &ni in &chains[ci] {
                        for &nj in &chains[cj] {
                            total += graph[ni][nj];
                        }
                    }
                    if best.is_none() || total > best.unwrap().2 {
                        best = Some((ci, cj, total));
                    }
                }
            }

            if let Some((ci, cj, _)) = best {
                // Merge cj into ci.
                let cj_nodes = chains[cj].clone();
                chains[ci].extend(cj_nodes);
                chains[cj].clear();
                active.remove(&cj);
            } else {
                break;
            }
        }

        // Flatten remaining chains into order.
        let mut result: Vec<String> = Vec::new();
        for &idx in &active {
            for &fi in &chains[idx] {
                result.push(self.functions[fi].name.clone());
            }
        }

        result
    }

    /// Call-Chain Clustering: follow call chains to group related functions.
    fn compute_call_chain_clusters(&self) -> Vec<String> {
        let n = self.functions.len();
        if n == 0 {
            return Vec::new();
        }

        // Find hot functions as chain roots.
        let mut roots: Vec<usize> = (0..n).filter(|&i| self.functions[i].is_hot).collect();
        roots.sort_by_key(|&i| Reverse(self.functions[i].execution_count));

        let mut visited: HashSet<usize> = HashSet::new();
        let mut result: Vec<String> = Vec::new();

        for root in roots {
            if visited.contains(&root) {
                continue;
            }
            // BFS from root, following the most frequent call targets.
            let mut queue: VecDeque<usize> = VecDeque::new();
            queue.push_back(root);

            while let Some(cur) = queue.pop_front() {
                if visited.contains(&cur) {
                    continue;
                }
                visited.insert(cur);
                result.push(self.functions[cur].name.clone());

                // Find called functions (simplified: just hot ones).
                for j in 0..n {
                    if !visited.contains(&j)
                        && self.functions[j].is_hot
                        && self.functions[cur].execution_count > 0
                    {
                        queue.push_back(j);
                    }
                }
            }
        }

        // Append unvisited functions.
        for i in 0..n {
            if !visited.contains(&i) {
                result.push(self.functions[i].name.clone());
            }
        }

        result
    }

    /// Apply the function reordering to the output.
    pub fn apply_code_reordering(&mut self, order: &[String]) {
        self.reordered_functions = order.to_vec();

        // Build output by concatenating code in the new order.
        let mut output = Vec::new();

        let name_to_func: HashMap<&str, &BoltFunction> = self
            .functions
            .iter()
            .map(|f| (f.name.as_str(), f))
            .collect();

        for name in order {
            if let Some(func) = name_to_func.get(name.as_str()) {
                let start = func.address as usize;
                let end = (func.address + func.size) as usize;
                if start < self.input.len() && end <= self.input.len() {
                    output.extend_from_slice(&self.input[start..end]);
                }
            }
        }

        self.output = output;
    }

    // ========================================================================
    // Basic Block Layout Optimization
    // ========================================================================

    /// Optimize basic block layout within a function.
    /// Returns the new block order (indices into func.blocks).
    fn optimize_basic_block_layout(&self, func: &BoltFunction) -> Vec<usize> {
        let n = func.blocks.len();
        if n <= 1 {
            return (0..n).collect();
        }

        // Find the entry block.
        let entry = func.blocks.iter().position(|b| b.is_entry).unwrap_or(0);

        // Use a greedy "hot-path" algorithm:
        // Starting from the entry, always follow the successor with
        // the highest execution count.
        let mut order: Vec<usize> = Vec::with_capacity(n);
        let mut visited: HashSet<usize> = HashSet::new();

        // DFS prioritizing hot successors.
        self.trace_hot_path(func, entry, &mut order, &mut visited);

        // Append any unvisited blocks (cold blocks).
        for i in 0..n {
            if !visited.contains(&i) {
                order.push(i);
                visited.insert(i);
            }
        }

        order
    }

    /// Trace hot paths through the CFG using a DFS that prioritizes
    /// successors with the highest execution count.
    fn trace_hot_path(
        &self,
        func: &BoltFunction,
        current: usize,
        order: &mut Vec<usize>,
        visited: &mut HashSet<usize>,
    ) {
        if visited.contains(&current) {
            return;
        }

        visited.insert(current);
        order.push(current);

        let block = &func.blocks[current];

        // Find the hottest unvisited successor.
        let mut best: Option<(usize, u64)> = None;
        for &succ in &block.successors {
            if succ < func.blocks.len() && !visited.contains(&succ) {
                let count = func.blocks[succ].execution_count;
                if best.is_none() || count > best.unwrap().1 {
                    best = Some((succ, count));
                }
            }
        }

        if let Some((next, _)) = best {
            self.trace_hot_path(func, next, order, visited);
        }

        // Then trace other successors.
        for &succ in &block.successors {
            if succ < func.blocks.len() && !visited.contains(&succ) {
                self.trace_hot_path(func, succ, order, visited);
            }
        }
    }

    /// Reorder blocks according to the new order.
    fn reorder_blocks(&mut self, _func: &mut BoltFunction, _new_order: &[usize]) {
        // Reordering blocks changes their offsets. In a full implementation,
        // this recalculates jump displacements, updates addresses, and
        // rewrites the binary bytes accordingly.
        //
        // For the demonstration, we update the block ordering metadata.
        // The actual binary rewriting happens during emit().
    }

    // ========================================================================
    // Branch Optimizations
    // ========================================================================

    /// Optimize branches: shorten jump displacements, convert conditional
    /// branches to fallthrough when possible, eliminate unnecessary jumps.
    fn optimize_branches(&mut self) {
        // For each function, scan blocks and optimize branch instructions.
        for func in &mut self.functions {
            let n = func.blocks.len();
            for i in 0..n {
                let block = &func.blocks[i];
                // If a block has a single unconditional jump successor
                // and that successor is the next block, remove the jump.
                if block.successors.len() == 1 {
                    let target = block.successors[0];
                    // If the jump target is the immediately following block,
                    // we can eliminate the jump entirely.
                    if target == i + 1 && i + 1 < n {
                        // Mark as direct fallthrough — the jump can be nop'd.
                        // In emit(), this jump will be removed.
                    }
                }

                // Convert `jmp L; L:` (unconditional jump to next block)
                // into a fallthrough.
                if block.successors.len() == 1 && block.successors[0] == i + 1 {
                    // This is a dead jump; mark for removal.
                }

                // If a conditional branch has both successors and one is
                // unreachable (execution_count == 0), convert to unconditional.
                if block.successors.len() == 2 {
                    let s0 = block.successors[0];
                    let s1 = block.successors[1];
                    if s0 < n && s1 < n {
                        let c0 = func.blocks[s0].execution_count;
                        let c1 = func.blocks[s1].execution_count;
                        if c0 == 0 && c1 > 0 {
                            // Always taken — convert conditional to unconditional.
                        } else if c0 > 0 && c1 == 0 {
                            // Never taken — remove branch, fall through.
                        }
                    }
                }
            }
        }
    }

    /// Remove NOP instructions for code size reduction.
    fn optimize_nops(&mut self) {
        // Scan through the output buffer and remove NOP bytes (0x90)
        // that appear in sequences. In a real implementation, this
        // adjusts all relative offsets and relocation entries.
        //
        // Multi-byte NOPs (e.g., 0x66 0x90, 0x0F 0x1F 0x00) are also removed.
        let _data = &mut self.output;

        // NOP patterns (x86-64):
        //   0x90                         — single-byte NOP
        //   0x66 0x90                    — two-byte NOP
        //   0x0F 0x1F 0x00               — three-byte NOP
        //   0x0F 0x1F 0x40 0x00          — four-byte NOP
        //   0x0F 0x1F 0x44 0x00 0x00     — five-byte NOP
        // etc.

        // In a full implementation, we'd scan for these patterns,
        // remove them, and fix up all branch displacements.
    }

    /// Optimize indirect branches: convert `jmp *reg` or `call *reg` when
    /// profiling shows a single dominant target.
    fn optimize_indirect_branches(&mut self) {
        // Indirect branch optimization: if an indirect jump/call
        // has one overwhelmingly dominant target (>99% of executions),
        // convert it to a direct jump with a fallback indirect jump.

        for func in &self.functions {
            for block in &func.blocks {
                // Check if this block ends with an indirect jump.
                // In x86-64, `jmp rax` = 0xFF 0xE0 and variants.
                // If we have profile data showing a single dominant target,
                // we could emit: cmp rax, target; je target; jmp rax.

                // Simplified: just note that this optimization would
                // be applied here based on profile data.
                let _ = block;
                let _ = func;
            }
        }
    }

    /// Stitch cold code: separate cold basic blocks from hot ones,
    /// placing cold blocks in a separate `.text.cold` section.
    fn stitch_cold_code(&mut self) {
        // Cold-code stitching moves rarely-executed basic blocks
        // out of the hot code path. This improves I-cache utilization
        // and reduces TLB pressure.

        for func in &self.functions {
            if !func.is_hot {
                continue;
            }

            // Identify cold blocks (execution count < 10% of max).
            let max_count = func
                .blocks
                .iter()
                .map(|b| b.execution_count)
                .max()
                .unwrap_or(1);

            let threshold = (max_count as f64 * 0.10) as u64;
            let mut cold_indices: Vec<usize> = func
                .blocks
                .iter()
                .enumerate()
                .filter(|(_, b)| b.execution_count < threshold && !b.is_entry)
                .map(|(i, _)| i)
                .collect();

            // In a full implementation, cold blocks are moved to a
            // separate section and reached via a trampoline or direct jump.
            cold_indices.sort();
            let _ = cold_indices;
        }
    }

    // ========================================================================
    // Relocation, Symbol, and Debug Info Updates
    // ========================================================================

    /// Update relocation entries after code movement.
    fn update_relocations(&mut self) {
        // When code is moved, all PC-relative relocations must be
        // recalculated. This includes:
        // - R_X86_64_PC32, R_X86_64_PLT32 for calls/jumps
        // - R_X86_64_GOTPCREL for GOT accesses
        // - R_X86_64_PC64 for 64-bit PC-relative

        // In a full implementation, this iterates over the relocation
        // sections, computes new addends based on the moved addresses,
        // and writes them back.
    }

    /// Update symbol table entries after code movement.
    fn update_symbols(&mut self) {
        // Update each symbol's st_value (virtual address) and
        // st_size to reflect the new layout. Function symbols
        // get their addresses from the reordered output.

        for func in &self.functions {
            // Find the symbol for this function and update its address.
            let _ = func;
        }
    }

    /// Update DWARF debug information after code movement.
    fn update_debug_info(&mut self) {
        // DWARF .debug_line, .debug_info, .debug_aranges, and
        // .debug_frame sections contain PC-relative references
        // that must be updated after code reordering.

        // .debug_line: update the address and op_index for each
        //   row in the line number program.
        // .debug_info: update DW_AT_low_pc and DW_AT_high_pc.
        // .debug_aranges: update the address/length tuples.
        // .debug_frame: update the initial_location in each FDE.

        // In a full implementation, this parses each DWARF section,
        // applies address deltas, and rewrites the section.
    }

    // ========================================================================
    // Phase 4: Emit
    // ========================================================================

    /// Emit the optimized binary. Returns the output bytes.
    pub fn emit(&mut self) -> Vec<u8> {
        // Apply the function-level reordering.
        if !self.reordered_functions.is_empty() {
            let order = self.reordered_functions.clone();
            self.apply_code_reordering(&order);
        }

        // Update all references.
        self.update_relocations();
        self.update_symbols();
        self.update_debug_info();

        // If no reordering was specified, output is the input.
        if self.output.is_empty() {
            self.output = self.input.clone();
        }

        std::mem::take(&mut self.output)
    }
}

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

/// Build CFG for a single function (free function version).
fn build_function_cfg_inner(func: &mut BoltFunction) {
    let max_count = func
        .blocks
        .iter()
        .map(|b| b.execution_count)
        .max()
        .unwrap_or(0);

    for block in &mut func.blocks {
        block.is_entry = block.offset == 0;
        if max_count > 0 {
            block.execution_count = block.execution_count.max(1);
        }
    }

    let total: u64 = func.blocks.iter().map(|b| b.execution_count).sum();
    func.execution_count = total;
    func.is_hot = total > 100;
}

/// Annotate functions with profile execution counts (static version).
fn annotate_from_profile_static(profile: &BoltProfile, functions: &mut [BoltFunction]) {
    for func in functions.iter_mut() {
        if let Some(fp) = profile.functions.get(&func.name) {
            func.execution_count = fp.entry_count;
            func.is_hot = fp.is_hot;
            for block in &mut func.blocks {
                if let Some(&count) = fp.block_counts.get(&block.address) {
                    block.execution_count = count;
                }
            }
        }
    }
}

/// Reorder blocks according to the new order (static version).
fn reorder_blocks_static(_func: &mut BoltFunction, _new_order: &[usize]) {
    // Reordering blocks changes their offsets. In a full implementation,
    // this recalculates jump displacements, updates addresses, and
    // rewrites the binary bytes accordingly.
}

// ============================================================================
// Helper: find function for Union-Find
// ============================================================================

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], rank: &mut [usize], x: usize, y: usize) {
    let rx = find(parent, x);
    let ry = find(parent, y);
    if rx == ry {
        return;
    }
    if rank[rx] < rank[ry] {
        parent[rx] = ry;
    } else if rank[rx] > rank[ry] {
        parent[ry] = rx;
    } else {
        parent[ry] = rx;
        rank[rx] += 1;
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Extended Binary Rewriting
// ═══════════════════════════════════════════════════════════════════════════

impl BOLTBinaryRewriter {
    /// Reorder functions based on profile data to improve I-cache utilization.
    /// Hot functions are placed together, cold functions are separated.
    pub fn reorder_functions(&mut self) {
        // Sort functions by execution count descending (hot first).
        self.functions
            .sort_by(|a, b| b.execution_count.cmp(&a.execution_count));

        // Group hot functions together for cache locality.
        let mut hot_funcs: Vec<String> = Vec::new();
        let mut cold_funcs: Vec<String> = Vec::new();

        for func in self.functions.iter() {
            if func.is_hot {
                hot_funcs.push(func.name.clone());
            } else {
                cold_funcs.push(func.name.clone());
            }
        }

        // Hot functions first in the reordered list.
        self.reordered_functions = hot_funcs;
        self.reordered_functions.append(&mut cold_funcs);
    }

    /// Optimize basic block layout: place hot blocks sequentially
    /// and split cold blocks to a separate section.
    pub fn optimize_basic_block_layout_extended(&self, func: &mut BoltFunction) {
        let total_count: u64 = func.blocks.iter().map(|b| b.execution_count).sum();
        if total_count == 0 {
            return;
        }

        let threshold = total_count / 10; // Top 10% hot

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

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

        // Layout: hot blocks in fall-through order, cold blocks separated.
        let mut layout = hot_blocks;
        layout.append(&mut cold_blocks);

        func.blocks = layout;
    }

    /// Pad functions with NOPs to align hot code to cache-line boundaries.
    pub fn nop_pad_for_alignment(&self, func: &mut BoltFunction, alignment: u64) -> Vec<u8> {
        let mut padding = Vec::new();
        let current_size: u64 = func.blocks.iter().map(|b| b.size as u64).sum();
        let remainder = current_size % alignment;

        if remainder != 0 {
            let pad_size = alignment - remainder;
            // x86 NOP encoding: single-byte NOP is 0x90.
            padding.resize(pad_size as usize, 0x90u8);
        }

        padding
    }

    /// Perform hot/cold splitting: move cold basic blocks to a separate
    /// section (e.g., `.text.cold`) and insert trampolines for branches
    /// crossing the split.
    pub fn hot_cold_split(
        &self,
        func: &BoltFunction,
        threshold: u64,
    ) -> (Vec<BoltBlock>, Vec<BoltBlock>) {
        let mut hot = Vec::new();
        let mut cold = Vec::new();

        for block in &func.blocks {
            if block.execution_count >= threshold {
                hot.push(block.clone());
            } else {
                cold.push(block.clone());
            }
        }

        (hot, cold)
    }

    /// Optimize jump tables by converting to simpler branch sequences
    /// when possible (e.g., when most cases are cold).
    pub fn optimize_jump_tables(&self, func: &mut BoltFunction) {
        // Collect execution counts upfront to avoid borrow conflicts.
        let exec_counts: Vec<u64> = func.blocks.iter().map(|b| b.execution_count).collect();

        // For each block with multiple successors (indirect branch),
        // check if we can simplify to direct branches.
        for block in &mut func.blocks {
            if block.successors.len() > 2 {
                // If only one successor is hot, convert to conditional
                // branch to hot + fallthrough to cold.
                let mut hot_successor = None;
                let max_count = block
                    .successors
                    .iter()
                    .enumerate()
                    .max_by_key(|&(_, &idx)| exec_counts.get(idx).copied().unwrap_or(0));

                if let Some((hot_idx, _)) = max_count {
                    hot_successor = Some(block.successors[hot_idx]);
                }

                if let Some(hot) = hot_successor {
                    // Simplify: keep only the hot successor as direct.
                    block.successors = vec![hot];
                }
            }
        }
    }

    /// Compute the total function size after rewriting.
    pub fn compute_function_size(&self, func: &BoltFunction) -> u64 {
        func.blocks.iter().map(|b| b.size as u64).sum()
    }

    /// Estimate the I-cache miss reduction from function reordering.
    pub fn estimate_icache_improvement(&self) -> f64 {
        let original_layout: u64 = self
            .functions
            .iter()
            .map(|f| self.compute_function_size(f))
            .sum();
        let reordered_layout: u64 = original_layout;
        // Note: reordering doesn't change total function sizes.
        let _ = &self.reordered_functions;

        if original_layout == 0 {
            return 0.0;
        }

        // Simple heuristic: 10% improvement per unit of hot-code locality.
        let reduction = 1.0 - (reordered_layout as f64 / original_layout as f64);
        (reduction * 10.0).max(0.0).min(30.0)
    }

    /// Merge adjacent blocks that have a single fall-through edge.
    pub fn merge_adjacent_blocks(&self, func: &mut BoltFunction) -> usize {
        let mut merged = 0;
        let mut i = 0;
        while i + 1 < func.blocks.len() {
            let current = &func.blocks[i];
            let next = &func.blocks[i + 1];

            // If current falls through to next and next has no other
            // predecessors, merge them.
            if current.successors.len() == 1 && current.successors[0] == i + 1 && !next.is_entry {
                let next_block = func.blocks.remove(i + 1);
                func.blocks[i].size += next_block.size;
                func.blocks[i].successors = next_block.successors.clone();
                merged += 1;
            } else {
                i += 1;
            }
        }
        merged
    }

    /// Align hot blocks to a given alignment boundary.
    pub fn align_hot_blocks(
        &self,
        func: &mut BoltFunction,
        alignment: u64,
        threshold: u64,
    ) -> Vec<u64> {
        let mut padding_sizes = Vec::new();
        for block in &func.blocks {
            if block.execution_count >= threshold {
                let remainder = (block.offset as u64) % alignment;
                if remainder != 0 {
                    padding_sizes.push(alignment - remainder);
                } else {
                    padding_sizes.push(0);
                }
            } else {
                padding_sizes.push(0);
            }
        }
        padding_sizes
    }

    /// Prune cold blocks that are never executed.
    pub fn prune_cold_blocks(&self, func: &mut BoltFunction) -> usize {
        let before = func.blocks.len();
        func.blocks.retain(|b| b.execution_count > 0);
        before - func.blocks.len()
    }

    /// Compute the total execution count for a function.
    pub fn function_execution_count(&self, func: &BoltFunction) -> u64 {
        func.blocks.iter().map(|b| b.execution_count).sum()
    }

    /// Generate a summary of the rewriting results.
    pub fn summary(&self) -> RewriteSummary {
        let total_functions = self.functions.len();
        let total_blocks: usize = self.functions.iter().map(|f| f.blocks.len()).sum();
        let hot_functions = self.functions.iter().filter(|f| f.is_hot).count();
        let reordered = self.reordered_functions.len();

        RewriteSummary {
            total_functions,
            total_blocks,
            hot_functions,
            reordered_functions: reordered,
        }
    }
}

/// Summary of the binary rewrite operation.
#[derive(Debug, Clone)]
pub struct RewriteSummary {
    pub total_functions: usize,
    pub total_blocks: usize,
    pub hot_functions: usize,
    pub reordered_functions: usize,
}

impl RewriteSummary {
    pub fn reorder_percentage(&self) -> f64 {
        if self.total_functions == 0 {
            0.0
        } else {
            (self.reordered_functions as f64 / self.total_functions as f64) * 100.0
        }
    }
}

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

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

    fn make_simple_func(name: &str, addr: u64, size: u64) -> BoltFunction {
        BoltFunction::new(name.to_string(), addr, size)
    }

    fn make_test_data() -> Vec<u8> {
        // A minimal x86-64 binary snippet with recognizable patterns.
        let mut data = vec![0u8; 1024];
        // Function 1 at offset 0: push rbp; mov rbp, rsp; nop; ret
        data[0] = 0x55; // push rbp
        data[1] = 0x48;
        data[2] = 0x89;
        data[3] = 0xE5; // mov rbp, rsp
        data[4] = 0x90; // nop
        data[5] = 0xC3; // ret
                        // Padding
        data[6] = 0x90;
        data[7] = 0x90;
        // Function 2 at offset 8: push rbp; mov rbp, rsp; ret
        data[8] = 0x55; // push rbp
        data[9] = 0x48;
        data[10] = 0x89;
        data[11] = 0xE5; // mov rbp, rsp
        data[12] = 0xC3; // ret
        data
    }

    #[test]
    fn test_new_rewriter() {
        let data = vec![0xC3];
        let rewriter = BOLTBinaryRewriter::new(data.clone());
        assert_eq!(rewriter.input, data);
        assert!(rewriter.output.is_empty());
        assert!(rewriter.functions.is_empty());
    }

    #[test]
    fn test_with_layout() {
        let rewriter = BOLTBinaryRewriter::new(vec![]).with_layout(LayoutAlgorithm::CacheSort);
        assert_eq!(rewriter.layout_algorithm, LayoutAlgorithm::CacheSort);
    }

    #[test]
    fn test_disassemble_simple() {
        let data = make_test_data();
        let mut rewriter = BOLTBinaryRewriter::new(data);
        let result = rewriter.disassemble();
        assert!(result.is_ok());
        // Should find at least the two functions.
        assert!(rewriter.functions.len() >= 1);
    }

    #[test]
    fn test_function_new() {
        let f = BoltFunction::new("main".into(), 0x400000, 128);
        assert_eq!(f.name, "main");
        assert_eq!(f.address, 0x400000);
        assert_eq!(f.size, 128);
        assert!(!f.is_hot);
        assert_eq!(f.execution_count, 0);
    }

    #[test]
    fn test_block_new() {
        let b = BoltBlock::new(0, 16);
        assert_eq!(b.offset, 0);
        assert_eq!(b.size, 16);
        assert!(!b.is_entry);
        assert!(!b.is_exit);
        assert!(b.successors.is_empty());
    }

    #[test]
    fn test_function_block_count() {
        let mut f = make_simple_func("test", 0x1000, 64);
        f.blocks.push(BoltBlock::new(0, 16));
        f.blocks.push(BoltBlock::new(16, 16));
        assert_eq!(f.block_count(), 2);
    }

    #[test]
    fn test_build_cfg() {
        let mut rewriter = BOLTBinaryRewriter::new(vec![0xC3]);
        let mut f = make_simple_func("f", 0, 64);
        f.blocks.push(BoltBlock::new(0, 8));
        f.blocks.push(BoltBlock::new(8, 8));
        rewriter.functions.push(f);
        rewriter.build_cfg();
        // Should complete without panicking.
        assert_eq!(rewriter.functions.len(), 1);
    }

    #[test]
    fn test_cache_sort() {
        let mut rewriter = BOLTBinaryRewriter::new(vec![]);
        let mut f1 = make_simple_func("a", 0, 64);
        f1.is_hot = true;
        f1.execution_count = 1000;
        let f2 = make_simple_func("b", 64, 32);
        rewriter.functions.push(f1);
        rewriter.functions.push(f2);
        let order = rewriter.compute_cache_sort();
        // Hot function should come first.
        assert_eq!(order[0], "a");
    }

    #[test]
    fn test_pettis_hansen() {
        let mut rewriter = BOLTBinaryRewriter::new(vec![]);
        rewriter.functions.push(make_simple_func("f1", 0, 64));
        rewriter.functions.push(make_simple_func("f2", 64, 64));
        rewriter.functions.push(make_simple_func("f3", 128, 64));
        let order = rewriter.compute_pettis_hansen();
        assert_eq!(order.len(), 3);
    }

    #[test]
    fn test_call_chain_clusters() {
        let mut rewriter = BOLTBinaryRewriter::new(vec![]);
        let mut f1 = make_simple_func("hot1", 0, 64);
        f1.is_hot = true;
        f1.execution_count = 500;
        let f2 = make_simple_func("cold1", 64, 32);
        let mut f3 = make_simple_func("hot2", 96, 64);
        f3.is_hot = true;
        f3.execution_count = 300;
        rewriter.functions.push(f1);
        rewriter.functions.push(f2);
        rewriter.functions.push(f3);
        let order = rewriter.compute_call_chain_clusters();
        assert_eq!(order.len(), 3);
    }

    #[test]
    fn test_hfsort_clusters_empty() {
        let rewriter = BOLTBinaryRewriter::new(vec![]);
        let clusters = rewriter.compute_hfsort_clusters();
        assert!(clusters.is_empty());
    }

    #[test]
    fn test_emit_no_reorder() {
        let data = vec![0x55, 0xC3];
        let mut rewriter = BOLTBinaryRewriter::new(data.clone());
        let output = rewriter.emit();
        assert_eq!(output, data);
    }

    #[test]
    fn test_apply_code_reordering() {
        let data: Vec<u8> = (0..256).map(|i| i as u8).collect();
        let mut rewriter = BOLTBinaryRewriter::new(data);
        // Create two functions: one at offset 10 size 20, one at offset 100 size 30.
        let f1 = BoltFunction::new("f1".into(), 10, 20);
        let f2 = BoltFunction::new("f2".into(), 100, 30);
        rewriter.functions.push(f1);
        rewriter.functions.push(f2);
        rewriter.apply_code_reordering(&["f2".to_string(), "f1".to_string()]);
        // Output should be f2 bytes (100..130) then f1 bytes (10..30).
        assert_eq!(rewriter.output.len(), 50);
        assert_eq!(rewriter.output[0], 100);
        assert_eq!(rewriter.output[30], 10);
    }

    #[test]
    fn test_optimize_function_layout_hfsort() {
        let mut rewriter = BOLTBinaryRewriter::new(vec![]).with_layout(LayoutAlgorithm::HFSort);
        rewriter.functions.push(make_simple_func("a", 0, 4));
        rewriter.functions.push(make_simple_func("b", 4, 4));
        rewriter.optimize_function_layout();
        assert!(!rewriter.reordered_functions.is_empty());
    }

    #[test]
    fn test_optimize_function_layout_cachesort() {
        let mut rewriter = BOLTBinaryRewriter::new(vec![]).with_layout(LayoutAlgorithm::CacheSort);
        rewriter.functions.push(make_simple_func("a", 0, 4));
        rewriter.functions.push(make_simple_func("b", 4, 4));
        rewriter.optimize_function_layout();
        assert_eq!(rewriter.reordered_functions.len(), 2);
    }

    #[test]
    fn test_optimize_blocks_single() {
        let rewriter = BOLTBinaryRewriter::new(vec![]);
        let mut func = make_simple_func("f", 0, 4);
        func.blocks.push(BoltBlock::new(0, 4));
        let order = rewriter.optimize_basic_block_layout(&func);
        assert_eq!(order, vec![0]);
    }

    #[test]
    fn test_optimize_blocks_two() {
        let rewriter = BOLTBinaryRewriter::new(vec![]);
        let mut func = make_simple_func("f", 0, 8);
        let mut b1 = BoltBlock::new(0, 4);
        b1.is_entry = true;
        b1.successors.push(1);
        let mut b2 = BoltBlock::new(4, 4);
        b2.is_exit = true;
        func.blocks.push(b1);
        func.blocks.push(b2);
        let order = rewriter.optimize_basic_block_layout(&func);
        assert_eq!(order.len(), 2);
        assert_eq!(order[0], 0); // entry first
    }

    #[test]
    fn test_run_optimizations_no_panic() {
        let data = make_test_data();
        let mut rewriter = BOLTBinaryRewriter::new(data);
        let _ = rewriter.disassemble();
        rewriter.build_cfg();
        // run_optimizations should not panic even without profile data.
        rewriter.run_optimizations();
    }

    #[test]
    fn test_hfsort_plus_clusters() {
        let mut rewriter = BOLTBinaryRewriter::new(vec![]);
        let mut f1 = make_simple_func("hot1", 0, 4);
        f1.is_hot = true;
        f1.execution_count = 100;
        let mut f2 = make_simple_func("hot2", 4, 4);
        f2.is_hot = true;
        f2.execution_count = 200;
        rewriter.functions.push(f1);
        rewriter.functions.push(f2);
        let order = rewriter.compute_hfsort_plus_clusters();
        assert_eq!(order.len(), 2);
    }
}