llvm-native-core 0.1.13

//! X86 Branch Relaxation — relax branch instructions when the branch target
//! exceeds the immediate range of the original encoding. Implements JCC rel8→rel32,
//! JMP rel8→rel32, LOOP→DEC+JNZ, CALL rel32 fallback, JMP indirection for very far
//! targets, branch range analysis per basic block, iterative relaxation with fixpoint,
//! fragment layout adjustment, alignment padding, and branch target recomputation.
//!
//! Clean-room behavioral reconstruction from:
//! - Intel® 64 and IA-32 Architectures Software Developer's Manual, Vol. 2 (Instruction Set)
//! - System V ABI: AMD64 Architecture Processor Supplement
//! - ELF x86-64 psABI (branch handling and relaxation)
//! - LLVM MC branch relaxation behavior (oracle interrogation)
//!
//! Coverage:
//! - Branch instruction classification: JCC (8-bit/32-bit), JMP (8-bit/32-bit),
//!   CALL (32-bit), LOOP/LOOPE/LOOPNE (8-bit only), JCXZ/JECXZ/JRCXZ
//! - Relaxation transforms: JCC rel8 → JCC rel32, JMP rel8 → JMP rel32,
//!   LOOP → DEC ECX + JNZ rel8, CALL → indirect call via thunk
//! - Fragment-based layout: each basic block is a fragment with size and alignment
//! - Fixpoint loop: recompute addresses, scan branches, relax any that overflow,
//!   repeat until no more changes
//! - Alignment: NOP padding between fragments, alignment constraints per fragment
//! - Branch target recomputation: after relaxation, update all branch offsets
//! - Iterative convergence detection
//! - Full .text section layout with relaxed branches
//!
//! Zero LLVM source code consultation. All behavior reconstructed from
//! published specifications and black-box oracle interrogation.

#![allow(non_upper_case_globals, dead_code)]

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

// ============================================================================
// Branch instruction kinds — x86 branch classification
// ============================================================================

/// Kinds of unconditional and conditional branch instructions.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum BranchKind {
    /// Unconditional jump: JMP rel8 / JMP rel32.
    Jmp,
    /// Conditional jump with 8-bit offset: JCC rel8 (e.g., JE, JNE, JG, JL, etc.).
    Jcc8,
    /// Conditional jump with 32-bit offset: JCC rel32.
    Jcc32,
    /// Near call: CALL rel32.
    Call32,
    /// LOOP: decrements ECX/RCX and jumps if not zero (only 8-bit offset).
    Loop,
    /// LOOPE/LOOPZ: decrements and jumps if ECX != 0 and ZF=1.
    Loope,
    /// LOOPNE/LOOPNZ: decrements and jumps if ECX != 0 and ZF=0.
    Loopne,
    /// JCXZ / JECXZ / JRCXZ: jump if ECX/RCX is zero (8-bit only).
    Jcxz,
    /// Indirect call: CALL [mem] or CALL reg.
    CallIndirect,
    /// Indirect jump: JMP [mem] or JMP reg.
    JmpIndirect,
    /// Return: RET (near or far).
    Ret,
    /// Unconditional jump via indirection thunk (for very far targets).
    JmpThunk,
}

impl BranchKind {
    /// Whether this branch instruction can be relaxed to a wider encoding.
    pub fn can_relax(self) -> bool {
        matches!(
            self,
            Self::Jmp | Self::Jcc8 | Self::Loop | Self::Loope | Self::Loopne | Self::Jcxz
        )
    }

    /// Maximum displacement range in bytes for this branch kind.
    /// Returns (min_offset, max_offset) relative to the end of the branch instruction.
    pub fn displacement_range(self) -> (i64, i64) {
        match self {
            Self::Jcc8 | Self::Loop | Self::Loope | Self::Loopne | Self::Jcxz => (-128, 127),
            Self::Jcc32 | Self::Jmp | Self::Call32 => (-2_147_483_648, 2_147_483_647),
            _ => (0, 0),
        }
    }

    /// The size in bytes of the branch instruction.
    pub fn instruction_size(self) -> u32 {
        match self {
            Self::Jmp => 2,          // EB xx
            Self::Jcc8 => 2,         // 7x xx
            Self::Jcc32 => 6,        // 0F 8x xx xx xx xx
            Self::Call32 => 5,       // E8 xx xx xx xx
            Self::Loop => 2,         // E2 xx
            Self::Loope => 2,        // E1 xx
            Self::Loopne => 2,       // E0 xx
            Self::Jcxz => 3,         // 67 E3 xx (addr override)
            Self::CallIndirect => 3, // FF /2 modrm
            Self::JmpIndirect => 3,  // FF /4 modrm
            Self::Ret => 1,          // C3
            Self::JmpThunk => 5,     // E9 xx xx xx xx (jmp rel32 to thunk)
        }
    }

    /// Size after relaxation.
    pub fn relaxed_size(self) -> u32 {
        match self {
            Self::Jmp => Self::Jmp.instruction_size() + 3, // JMP rel8 → JMP rel32: +3
            Self::Jcc8 => Self::Jcc32.instruction_size(),  // JCC rel8 → JCC rel32: 2 → 6
            Self::Loop | Self::Loope | Self::Loopne => {
                // LOOP → DEC ECX (1) + JNZ rel8 (2): 2 → 3
                3
            }
            Self::Jcxz => {
                // JCXZ → TEST ECX,ECX (2) + JZ rel8 (2): 3 → 4
                4
            }
            _ => self.instruction_size(),
        }
    }

    /// Whether this is a conditional branch.
    pub fn is_conditional(self) -> bool {
        matches!(
            self,
            Self::Jcc8 | Self::Jcc32 | Self::Loop | Self::Loope | Self::Loopne | Self::Jcxz
        )
    }

    /// Whether this is an unconditional branch.
    pub fn is_unconditional(self) -> bool {
        matches!(
            self,
            Self::Jmp | Self::JmpIndirect | Self::JmpThunk | Self::Ret
        )
    }

    /// Human-readable name.
    pub fn name(self) -> &'static str {
        match self {
            Self::Jmp => "JMP",
            Self::Jcc8 => "JCC8",
            Self::Jcc32 => "JCC32",
            Self::Call32 => "CALL32",
            Self::Loop => "LOOP",
            Self::Loope => "LOOPE",
            Self::Loopne => "LOOPNE",
            Self::Jcxz => "JCXZ",
            Self::CallIndirect => "CALL_INDIRECT",
            Self::JmpIndirect => "JMP_INDIRECT",
            Self::Ret => "RET",
            Self::JmpThunk => "JMP_THUNK",
        }
    }
}

// ============================================================================
// Condition codes for JCC instructions
// ============================================================================

/// x86 condition codes (the "CC" in JCC).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ConditionCode {
    O,  // overflow
    NO, // no overflow
    B,  // below (CF=1)
    AE, // above or equal (CF=0)
    E,  // equal (ZF=1)
    NE, // not equal (ZF=0)
    BE, // below or equal
    A,  // above
    S,  // sign
    NS, // not sign
    P,  // parity
    NP, // not parity
    L,  // less
    GE, // greater or equal
    LE, // less or equal
    G,  // greater
}

impl ConditionCode {
    /// The opcode byte for 8-bit JCC (0x70..0x7F).
    pub fn opcode8(self) -> u8 {
        match self {
            Self::O => 0x70,
            Self::NO => 0x71,
            Self::B => 0x72,
            Self::AE => 0x73,
            Self::E => 0x74,
            Self::NE => 0x75,
            Self::BE => 0x76,
            Self::A => 0x77,
            Self::S => 0x78,
            Self::NS => 0x79,
            Self::P => 0x7A,
            Self::NP => 0x7B,
            Self::L => 0x7C,
            Self::GE => 0x7D,
            Self::LE => 0x7E,
            Self::G => 0x7F,
        }
    }

    /// The second opcode byte for 32-bit JCC (0x80..0x8F).
    /// Full encoding: 0F 8x xx xx xx xx.
    pub fn opcode32(self) -> u8 {
        match self {
            Self::O => 0x80,
            Self::NO => 0x81,
            Self::B => 0x82,
            Self::AE => 0x83,
            Self::E => 0x84,
            Self::NE => 0x85,
            Self::BE => 0x86,
            Self::A => 0x87,
            Self::S => 0x88,
            Self::NS => 0x89,
            Self::P => 0x8A,
            Self::NP => 0x8B,
            Self::L => 0x8C,
            Self::GE => 0x8D,
            Self::LE => 0x8E,
            Self::G => 0x8F,
        }
    }
}

// ============================================================================
// Branch instruction — a single branch to be relaxed
// ============================================================================

/// A branch instruction within the code.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BranchInst {
    /// Offset of this instruction from the start of its fragment.
    pub fragment_offset: u32,
    /// The kind of branch.
    pub kind: BranchKind,
    /// The condition code (for JCC branches).
    pub cc: Option<ConditionCode>,
    /// Target fragment ID.
    pub target_fragment: usize,
    /// Whether this branch has been relaxed.
    pub relaxed: bool,
    /// Whether this branch has been inverted (for relaxation of LOOP → DEC+JNZ).
    pub inverted: bool,
}

impl BranchInst {
    /// Create a new branch instruction.
    pub fn new(
        fragment_offset: u32,
        kind: BranchKind,
        target_fragment: usize,
        cc: Option<ConditionCode>,
    ) -> Self {
        Self {
            fragment_offset,
            kind,
            cc,
            target_fragment,
            relaxed: false,
            inverted: false,
        }
    }

    /// Create a JCC rel8.
    pub fn jcc8(offset: u32, cc: ConditionCode, target: usize) -> Self {
        Self::new(offset, BranchKind::Jcc8, target, Some(cc))
    }

    /// Create a JMP rel8.
    pub fn jmp8(offset: u32, target: usize) -> Self {
        Self::new(offset, BranchKind::Jmp, target, None)
    }

    /// Create a LOOP.
    pub fn loop_(offset: u32, target: usize) -> Self {
        Self::new(offset, BranchKind::Loop, target, None)
    }

    /// Create a LOOPE.
    pub fn loope(offset: u32, target: usize) -> Self {
        Self::new(offset, BranchKind::Loope, target, None)
    }

    /// Create a LOOPNE.
    pub fn loopne(offset: u32, target: usize) -> Self {
        Self::new(offset, BranchKind::Loopne, target, None)
    }

    /// Create a CALL rel32.
    pub fn call32(offset: u32, target: usize) -> Self {
        Self::new(offset, BranchKind::Call32, target, None)
    }

    /// Current instruction size (before or after relaxation).
    pub fn size(&self) -> u32 {
        if self.relaxed {
            self.kind.relaxed_size()
        } else {
            self.kind.instruction_size()
        }
    }

    /// Compute displacement from this branch's end to the target fragment start.
    /// `branch_addr`: absolute address of this branch.
    /// `target_addr`: absolute address of the target fragment.
    pub fn compute_displacement(&self, branch_addr: u64, target_addr: u64) -> i64 {
        let branch_end = branch_addr + self.size() as u64;
        target_addr as i64 - branch_end as i64
    }

    /// Check if this branch needs relaxation given its displacement.
    pub fn needs_relaxation(&self, displacement: i64) -> bool {
        if self.relaxed {
            return false; // Already relaxed to max width.
        }
        if !self.kind.can_relax() {
            return false;
        }
        let (min, max) = self.kind.displacement_range();
        displacement < min || displacement > max
    }

    /// Relax this branch instruction.
    /// Returns the size increase in bytes.
    pub fn relax(&mut self) -> u32 {
        if self.relaxed || !self.kind.can_relax() {
            return 0;
        }
        let old_size = self.size();
        self.relaxed = true;
        self.size() - old_size
    }
}

// ============================================================================
// Code Fragment — a basic-block-level chunk of code
// ============================================================================

/// Alignment boundary types.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Alignment {
    /// No alignment requirement.
    None,
    /// Align to 2^N bytes (e.g., 16 means align to 16 bytes).
    Pow2(u32),
}

impl Alignment {
    /// Get the alignment value (1 = no alignment).
    pub fn value(self) -> u32 {
        match self {
            Self::None => 1,
            Self::Pow2(n) => n,
        }
    }

    /// Compute padding needed to align an address.
    pub fn pad_to(self, addr: u64) -> u32 {
        let align = self.value() as u64;
        if align <= 1 {
            return 0;
        }
        let remainder = addr % align;
        if remainder == 0 {
            0
        } else {
            (align - remainder) as u32
        }
    }
}

/// A code fragment represents a contiguous block of code — typically a basic block
/// or a function — that participates in branch relaxation.
#[derive(Debug, Clone)]
pub struct CodeFragment {
    /// Unique fragment ID.
    pub id: usize,
    /// The instructions in this fragment (only branches matter for relaxation,
    /// but we track all instructions for size computation).
    pub instructions: Vec<BranchInst>,
    /// Total byte size of non-branch instructions in this fragment.
    pub fixed_size: u32,
    /// Alignment requirement for this fragment.
    pub alignment: Alignment,
    /// Assigned address (computed during layout).
    pub address: Option<u64>,
    /// Whether this fragment has been laid out.
    pub laid_out: bool,
    /// Labels within this fragment (offset → label name).
    pub labels: BTreeMap<u32, String>,
    /// Is this fragment a thunk (indirection for far jumps)?
    pub is_thunk: bool,
    /// Is this fragment a function entry point?
    pub is_function_entry: bool,
}

impl CodeFragment {
    /// Create a new code fragment.
    pub fn new(id: usize) -> Self {
        Self {
            id,
            instructions: Vec::new(),
            fixed_size: 0,
            alignment: Alignment::None,
            address: None,
            laid_out: false,
            labels: BTreeMap::new(),
            is_thunk: false,
            is_function_entry: false,
        }
    }

    /// Total size of this fragment (fixed + all instruction sizes).
    pub fn total_size(&self) -> u32 {
        self.fixed_size
            + self
                .instructions
                .iter()
                .map(|inst| inst.size())
                .sum::<u32>()
    }

    /// End address of this fragment.
    pub fn end_address(&self) -> Option<u64> {
        self.address.map(|addr| addr + self.total_size() as u64)
    }

    /// Add a branch instruction.
    pub fn add_branch(&mut self, branch: BranchInst) {
        self.instructions.push(branch);
    }

    /// Add a label at a given offset.
    pub fn add_label(&mut self, offset: u32, name: String) {
        self.labels.insert(offset, name);
    }

    /// Set fixed instruction bytes (non-branch instructions).
    pub fn set_fixed_size(&mut self, size: u32) {
        self.fixed_size = size;
    }

    /// Set alignment requirement.
    pub fn set_alignment(&mut self, align: Alignment) {
        self.alignment = align;
    }

    /// Mark as a function entry (gets higher alignment).
    pub fn mark_function_entry(&mut self) {
        self.is_function_entry = true;
        if matches!(self.alignment, Alignment::None) {
            self.alignment = Alignment::Pow2(16);
        }
    }

    /// Mark as a thunk.
    pub fn mark_thunk(&mut self) {
        self.is_thunk = true;
    }
}

// ============================================================================
// Fragment Layout — the section-level layout
// ============================================================================

/// The layout of all fragments in a section.
#[derive(Debug, Clone)]
pub struct FragmentLayout {
    /// All fragments (in layout order).
    pub fragments: Vec<CodeFragment>,
    /// Map from fragment ID to index in the layout order.
    id_to_index: HashMap<usize, usize>,
    /// Base address of the section.
    pub base_address: u64,
    /// Total section size after layout.
    pub section_size: u64,
    /// Alignment padding bytes inserted during layout.
    pub padding_bytes: u64,
}

impl FragmentLayout {
    /// Create a new layout with a given base address.
    pub fn new(base_address: u64) -> Self {
        Self {
            fragments: Vec::new(),
            id_to_index: HashMap::new(),
            base_address,
            section_size: 0,
            padding_bytes: 0,
        }
    }

    /// Add a fragment to the layout.
    pub fn add_fragment(&mut self, mut fragment: CodeFragment) {
        let idx = self.fragments.len();
        self.id_to_index.insert(fragment.id, idx);
        self.fragments.push(fragment);
    }

    /// Get fragment by ID.
    pub fn get_fragment(&self, id: usize) -> Option<&CodeFragment> {
        self.id_to_index.get(&id).map(|&idx| &self.fragments[idx])
    }

    /// Get mutable fragment by ID.
    pub fn get_fragment_mut(&mut self, id: usize) -> Option<&mut CodeFragment> {
        self.id_to_index
            .get(&id)
            .map(|&idx| &mut self.fragments[idx])
    }

    /// Get the index of a fragment in layout order.
    pub fn index_of(&self, id: usize) -> Option<usize> {
        self.id_to_index.get(&id).copied()
    }

    /// Number of fragments.
    pub fn len(&self) -> usize {
        self.fragments.len()
    }

    /// Whether the layout is empty.
    pub fn is_empty(&self) -> bool {
        self.fragments.is_empty()
    }
}

// ============================================================================
// Branch Relaxation Engine
// ============================================================================

/// Statistics for branch relaxation.
#[derive(Debug, Clone, Default)]
pub struct RelaxationStats {
    /// Number of iteration passes.
    pub passes: u32,
    /// Total branches inspected.
    pub branches_inspected: u64,
    /// Total branches relaxed.
    pub branches_relaxed: u64,
    /// Total size increase from relaxation (bytes).
    pub size_increase: u64,
    /// Number of thunks created for far jumps.
    pub thunks_created: u64,
}

/// The main branch relaxation engine.
#[derive(Debug, Clone)]
pub struct X86BranchRelaxation {
    /// Fragment layout.
    pub layout: FragmentLayout,
    /// Thunk fragments (created for far jumps).
    pub thunks: Vec<CodeFragment>,
    /// Next fragment ID.
    next_fragment_id: usize,
    /// Statistics.
    pub stats: RelaxationStats,
    /// Maximum number of relaxation iterations (safety limit).
    pub max_iterations: u32,
    /// Threshold for creating thunks instead of relaxing (bytes).
    pub thunk_threshold: i64,
    /// Enable/disable relaxation of LOOP instructions.
    pub relax_loops: bool,
    /// Whether NOP padding can be removed to make room.
    pub aggressive_padding_removal: bool,
}

impl Default for X86BranchRelaxation {
    fn default() -> Self {
        Self {
            layout: FragmentLayout::new(0),
            thunks: Vec::new(),
            next_fragment_id: 1,
            stats: RelaxationStats::default(),
            max_iterations: 100,
            thunk_threshold: 2_147_483_648,
            relax_loops: true,
            aggressive_padding_removal: false,
        }
    }
}

impl X86BranchRelaxation {
    /// Create a new branch relaxation engine.
    pub fn new(base_address: u64) -> Self {
        Self {
            layout: FragmentLayout::new(base_address),
            ..Default::default()
        }
    }

    /// Allocate a new fragment ID.
    pub fn next_fragment_id(&mut self) -> usize {
        let id = self.next_fragment_id;
        self.next_fragment_id += 1;
        id
    }

    /// Add a code fragment to the layout.
    pub fn add_fragment(&mut self, fragment: CodeFragment) {
        self.layout.add_fragment(fragment);
    }

    /// Create and add a new fragment.
    pub fn create_fragment(&mut self) -> usize {
        let id = self.next_fragment_id();
        let fragment = CodeFragment::new(id);
        self.layout.add_fragment(fragment);
        id
    }

    /// Create a thunk fragment for indirect jump targets (very far jumps).
    /// A thunk is a small piece of code: `JMP [RIP+offset]` or `JMP rel32` to the real target.
    pub fn create_thunk(&mut self, target_id: usize) -> usize {
        let thunk_id = self.next_fragment_id();
        let mut thunk = CodeFragment::new(thunk_id);
        thunk.mark_thunk();
        // A thunk is just a single JMP rel32 to the target.
        thunk.add_branch(BranchInst::new(0, BranchKind::Jmp, target_id, None));
        thunk.fixed_size = 0;

        self.thunks.push(thunk.clone());
        self.layout.add_fragment(thunk);
        self.stats.thunks_created += 1;

        thunk_id
    }

    // ------------------------------------------------------------------
    // Layout computation
    // ------------------------------------------------------------------

    /// Compute the layout: assign addresses to all fragments in order,
    /// inserting alignment padding as needed.
    pub fn compute_layout(&mut self) {
        let mut current_addr = self.layout.base_address;
        let mut total_padding: u64 = 0;

        for fragment in self.layout.fragments.iter_mut() {
            // Apply alignment.
            let pad = fragment.alignment.pad_to(current_addr);
            if pad > 0 {
                total_padding += pad as u64;
                current_addr += pad as u64;
            }
            fragment.address = Some(current_addr);
            fragment.laid_out = true;
            current_addr += fragment.total_size() as u64;
        }

        self.layout.section_size = current_addr - self.layout.base_address;
        self.layout.padding_bytes = total_padding;
    }

    // ------------------------------------------------------------------
    // Branch scanning — find branches that need relaxation
    // ------------------------------------------------------------------

    /// Scan all branches and return those that need relaxation.
    pub fn find_overreaching_branches(&mut self) -> Vec<(usize, usize, i64)> {
        // (fragment_index, branch_index, displacement)
        let mut overreaching = Vec::new();

        for (fi, fragment) in self.layout.fragments.iter().enumerate() {
            let frag_addr = match fragment.address {
                Some(a) => a,
                None => continue,
            };

            for (bi, branch) in fragment.instructions.iter().enumerate() {
                self.stats.branches_inspected += 1;

                // Compute the absolute address of this branch inside its fragment.
                let inst_offset = branch.fragment_offset;
                // Accumulate sizes of instructions before this one to find the
                // offset of this branch within the fragment's instruction stream.
                let mut acc_offset: u32 = 0;
                for prev in &fragment.instructions[..bi] {
                    acc_offset += prev.size();
                }
                acc_offset += inst_offset;

                let branch_abs_addr = frag_addr + acc_offset as u64;

                let target_fragment = match self.layout.get_fragment(branch.target_fragment) {
                    Some(f) => f,
                    None => continue,
                };
                let target_addr = match target_fragment.address {
                    Some(a) => a,
                    None => continue,
                };

                let displacement = branch.compute_displacement(branch_abs_addr, target_addr);

                if branch.needs_relaxation(displacement) {
                    overreaching.push((fi, bi, displacement));
                }
            }
        }

        overreaching
    }

    // ------------------------------------------------------------------
    // Relaxation pass — relax one set of branches
    // ------------------------------------------------------------------

    /// Run a single relaxation pass.
    /// Returns the number of branches relaxed in this pass.
    pub fn relax_pass(&mut self) -> u32 {
        // First compute layout to get accurate addresses.
        self.compute_layout();

        let overreaching = self.find_overreaching_branches();
        let mut relaxed_count: u32 = 0;

        // Relax branches that need it.
        for (fi, bi, displacement) in overreaching {
            let fragment = match self.layout.fragments.get_mut(fi) {
                Some(f) => f,
                None => continue,
            };

            let branch = match fragment.instructions.get_mut(bi) {
                Some(b) => b,
                None => continue,
            };

            // Check if we need a thunk instead.
            if branch.kind == BranchKind::Jmp && displacement.abs() >= self.thunk_threshold {
                // Create a thunk and redirect.
                let old_target = branch.target_fragment;
                let bi_copy = bi;
                let fi_copy = fi;
                // Release the borrows on fragment/branch before create_thunk
                drop(branch);
                drop(fragment);
                let thunk_id = self.create_thunk(old_target);
                // Re-borrow to update the branch target
                let fragment = self.layout.fragments.get_mut(fi_copy).unwrap();
                let branch = fragment.instructions.get_mut(bi_copy).unwrap();
                branch.target_fragment = thunk_id;
                // The thunk itself contains a JMP (which is already relaxed from the start).
                // This JMP to thunk rewrite works for the relaxation too.
                // Relax the JMP to rel32.
                let size_increase = branch.relax() as u64;
                self.stats.size_increase += size_increase;
                relaxed_count += 1;
                self.stats.branches_relaxed += 1;
            } else if branch.kind.can_relax() {
                let size_increase = branch.relax() as u64;
                self.stats.size_increase += size_increase;
                relaxed_count += 1;
                self.stats.branches_relaxed += 1;
            }
        }

        relaxed_count
    }

    // ------------------------------------------------------------------
    // Fixpoint loop — iterate until no more branches need relaxation
    // ------------------------------------------------------------------

    /// Run the full iterative relaxation algorithm to fixpoint.
    /// Returns the total number of passes executed.
    pub fn run_relaxation(&mut self) -> u32 {
        let mut passes: u32 = 0;

        loop {
            let relaxed = self.relax_pass();
            passes += 1;

            if relaxed == 0 {
                break;
            }
            if passes >= self.max_iterations {
                break;
            }
        }

        // Final layout pass.
        self.compute_layout();

        self.stats.passes = passes;
        passes
    }

    // ------------------------------------------------------------------
    // Post-relaxation: recompute branch targets
    // ------------------------------------------------------------------

    /// After relaxation, recompute all branch displacements to produce
    /// the final encoded branch instructions.
    pub fn recompute_displacements(&self) -> Vec<(usize, usize, i64)> {
        let mut displacements = Vec::new();

        for (fi, fragment) in self.layout.fragments.iter().enumerate() {
            let frag_addr = match fragment.address {
                Some(a) => a,
                None => continue,
            };

            for (bi, branch) in fragment.instructions.iter().enumerate() {
                // Find the absolute address of this branch.
                let mut acc_offset: u32 = 0;
                for prev in &fragment.instructions[..bi] {
                    acc_offset += prev.size();
                }
                acc_offset += branch.fragment_offset;

                let branch_abs = frag_addr + acc_offset as u64;

                let target_addr = self
                    .layout
                    .get_fragment(branch.target_fragment)
                    .and_then(|f| f.address)
                    .unwrap_or(0);

                let disp = branch.compute_displacement(branch_abs, target_addr);

                displacements.push((fi, bi, disp));
            }
        }

        displacements
    }

    // ------------------------------------------------------------------
    // Relaxation of specific branch types
    // ------------------------------------------------------------------

    /// Relax a JCC rel8 to JCC rel32, given the condition code.
    /// Returns the old and new instruction byte sequences.
    pub fn relax_jcc(cc: ConditionCode, old_offset: i8) -> (Vec<u8>, Vec<u8>) {
        let old = vec![cc.opcode8(), old_offset as u8];
        let new_offset = old_offset as i32;
        let new = vec![
            0x0F,
            cc.opcode32(),
            new_offset as u8,
            (new_offset >> 8) as u8,
            (new_offset >> 16) as u8,
            (new_offset >> 24) as u8,
        ];
        (old, new)
    }

    /// Relax a JMP rel8 to JMP rel32.
    pub fn relax_jmp(old_offset: i8) -> (Vec<u8>, Vec<u8>) {
        let old = vec![0xEB, old_offset as u8];
        let new_offset = old_offset as i32;
        let new = vec![
            0xE9,
            new_offset as u8,
            (new_offset >> 8) as u8,
            (new_offset >> 16) as u8,
            (new_offset >> 24) as u8,
        ];
        (old, new)
    }

    /// Relax a LOOP to DEC ECX + JNZ rel8.
    pub fn relax_loop(old_offset: i8) -> (Vec<u8>, Vec<u8>) {
        let old = vec![0xE2, old_offset as u8];
        // DEC ECX: 0x49 (REX.WB) 0xFF 0xC9 in 64-bit, simplified to 66 FF C9 for 32-bit.
        // For relaxed displacement we use JNZ rel8.
        let new = vec![0x66, 0xFF, 0xC9, 0x75, old_offset as u8];
        (old, new)
    }

    /// Relax a JCXZ to TEST ECX,ECX + JZ rel8.
    pub fn relax_jcxz(old_offset: i8) -> (Vec<u8>, Vec<u8>) {
        let old = vec![0x67, 0xE3, old_offset as u8];
        let new = vec![0x85, 0xC9, 0x74, old_offset as u8];
        (old, new)
    }

    /// Emit NOP padding bytes for alignment.
    pub fn emit_alignment_padding(count: u32) -> Vec<u8> {
        match count {
            0 => vec![],
            1 => vec![0x90],
            2 => vec![0x66, 0x90],
            3 => vec![0x0F, 0x1F, 0x00],
            4 => vec![0x0F, 0x1F, 0x40, 0x00],
            5 => vec![0x0F, 0x1F, 0x44, 0x00, 0x00],
            6 => vec![0x66, 0x0F, 0x1F, 0x44, 0x00, 0x00],
            7 => vec![0x0F, 0x1F, 0x80, 0x00, 0x00, 0x00, 0x00],
            8 => vec![0x0F, 0x1F, 0x84, 0x00, 0x00, 0x00, 0x00, 0x00],
            9 => {
                let mut v = vec![0x66, 0x0F, 0x1F, 0x84, 0x00, 0x00, 0x00, 0x00, 0x00];
                v.push(0x90);
                v
            }
            _ => {
                // For larger padding, emit a sequence of 9-byte NOPs then a smaller one.
                let mut result = Vec::with_capacity(count as usize);
                let mut remaining = count;
                while remaining >= 9 {
                    result
                        .extend_from_slice(&[0x66, 0x0F, 0x1F, 0x84, 0x00, 0x00, 0x00, 0x00, 0x00]);
                    remaining -= 9;
                }
                result.extend_from_slice(&Self::emit_alignment_padding(remaining));
                result
            }
        }
    }
}

// ============================================================================
// Full section encoder — produces relaxed code bytes
// ============================================================================

/// Encode an entire section with relaxed branches and alignment.
#[derive(Debug, Clone, Default)]
pub struct SectionEncoder {
    /// The relaxed layout.
    pub relaxation: Option<X86BranchRelaxation>,
    /// Encoded bytes.
    pub encoded_bytes: Vec<u8>,
}

impl SectionEncoder {
    /// Create a new section encoder.
    pub fn new() -> Self {
        Self::default()
    }

    /// Encode the section after relaxation.
    pub fn encode(
        &mut self,
        relaxation: &X86BranchRelaxation,
        displacements: &[(usize, usize, i64)],
    ) {
        let mut buf = Vec::new();
        let mut current_addr = relaxation.layout.base_address;

        // Build a quick lookup: (fragment_index, branch_index) → displacement.
        let disp_map: HashMap<(usize, usize), i64> = displacements
            .iter()
            .map(|&(fi, bi, disp)| ((fi, bi), disp))
            .collect();

        for fragment in &relaxation.layout.fragments {
            // Alignment padding.
            let pad = fragment.alignment.pad_to(current_addr);
            if pad > 0 {
                buf.extend_from_slice(&X86BranchRelaxation::emit_alignment_padding(pad));
                current_addr += pad as u64;
            }

            let frag_start = buf.len();

            // For simplicity, we encode branches as their relaxed forms with computed
            // displacements. Non-branch bytes are placeholders (NOPs for now).
            for (bi, branch) in fragment.instructions.iter().enumerate() {
                let disp = disp_map.get(&(fragment.id, bi)).copied().unwrap_or(0);

                match branch.kind {
                    BranchKind::Jmp => {
                        if branch.relaxed {
                            // JMP rel32.
                            buf.push(0xE9);
                            buf.extend_from_slice(&(disp as i32).to_le_bytes());
                        } else {
                            // JMP rel8.
                            buf.push(0xEB);
                            buf.push(disp as u8);
                        }
                    }
                    BranchKind::Jcc8 => {
                        if branch.relaxed {
                            // JCC rel32: 0F 8x xx xx xx xx.
                            if let Some(cc) = branch.cc {
                                buf.push(0x0F);
                                buf.push(cc.opcode32());
                                buf.extend_from_slice(&(disp as i32).to_le_bytes());
                            }
                        } else {
                            // JCC rel8: 7x xx.
                            if let Some(cc) = branch.cc {
                                buf.push(cc.opcode8());
                                buf.push(disp as u8);
                            }
                        }
                    }
                    BranchKind::Jcc32 => {
                        // Already wide; just encode.
                        if let Some(cc) = branch.cc {
                            buf.push(0x0F);
                            buf.push(cc.opcode32());
                            buf.extend_from_slice(&(disp as i32).to_le_bytes());
                        }
                    }
                    BranchKind::Call32 => {
                        buf.push(0xE8);
                        buf.extend_from_slice(&(disp as i32).to_le_bytes());
                    }
                    BranchKind::Loop => {
                        if branch.relaxed && relaxation.relax_loops {
                            // DEC ECX + JNZ rel8
                            buf.extend_from_slice(&[0x66, 0xFF, 0xC9, 0x75]);
                            buf.push(disp as u8);
                        } else {
                            buf.push(0xE2);
                            buf.push(disp as u8);
                        }
                    }
                    BranchKind::Loope => {
                        buf.push(0xE1);
                        buf.push(disp as u8);
                    }
                    BranchKind::Loopne => {
                        buf.push(0xE0);
                        buf.push(disp as u8);
                    }
                    BranchKind::Jcxz => {
                        buf.push(0x67);
                        buf.push(0xE3);
                        buf.push(disp as u8);
                    }
                    _ => {
                        // Placeholder for other instruction types.
                        for _ in 0..branch.size() {
                            buf.push(0x90); // NOP placeholder
                        }
                    }
                }
            }

            // Add fixed-size bytes as NOP placeholders.
            for _ in 0..fragment.fixed_size {
                buf.push(0x90);
            }

            let frag_end = buf.len();
            current_addr += (frag_end - frag_start) as u64;
        }

        self.encoded_bytes = buf;
        self.relaxation = Some(relaxation.clone());
    }
}

// ============================================================================
// High-level relaxation pipeline
// ============================================================================

/// A convenience pipeline that bundles fragment creation, layout, relaxation,
/// and encoding into a single flow.
#[derive(Debug, Clone)]
pub struct RelaxationPipeline {
    /// The relaxation engine.
    pub engine: X86BranchRelaxation,
    /// The section encoder.
    pub encoder: SectionEncoder,
}

impl RelaxationPipeline {
    /// Create a new pipeline.
    pub fn new(base_address: u64) -> Self {
        Self {
            engine: X86BranchRelaxation::new(base_address),
            encoder: SectionEncoder::new(),
        }
    }

    /// Add a fragment with branch instructions.
    pub fn add_fragment_with_branches(
        &mut self,
        fixed_size: u32,
        branches: Vec<BranchInst>,
        alignment: Alignment,
        is_entry: bool,
    ) -> usize {
        let id = self.engine.next_fragment_id();
        let mut fragment = CodeFragment::new(id);
        fragment.fixed_size = fixed_size;
        fragment.alignment = alignment;
        if is_entry {
            fragment.mark_function_entry();
        }
        for branch in branches {
            fragment.add_branch(branch);
        }
        self.engine.add_fragment(fragment);
        id
    }

    /// Run the full relaxation and encoding pipeline.
    pub fn run(&mut self) -> Result<Vec<u8>, String> {
        let passes = self.engine.run_relaxation();
        if passes >= self.engine.max_iterations {
            return Err(format!(
                "Relaxation did not converge after {} passes",
                passes
            ));
        }

        let displacements = self.engine.recompute_displacements();
        self.encoder.encode(&self.engine, &displacements);

        Ok(self.encoder.encoded_bytes.clone())
    }
}

// ============================================================================
// Branch Prefix Handling
// ============================================================================

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BranchPrefix {
    None,
    Rex(u8),
    SegmentOverride(u8),
    AddrSizeOverride,
    OperandSizeOverride,
    Vex2(u8, u8),
    Vex3(u8, u8, u8),
    Evex(u8, u8, u8, u8),
}

impl BranchPrefix {
    pub fn size(self) -> u32 {
        match self {
            Self::None => 0,
            Self::Rex(_)
            | Self::SegmentOverride(_)
            | Self::AddrSizeOverride
            | Self::OperandSizeOverride => 1,
            Self::Vex2(_, _) => 2,
            Self::Vex3(_, _, _) => 3,
            Self::Evex(_, _, _, _) => 4,
        }
    }
    pub fn encode(self) -> Vec<u8> {
        match self {
            Self::None => vec![],
            Self::Rex(b) => vec![b],
            Self::SegmentOverride(b) => vec![b],
            Self::AddrSizeOverride => vec![0x67],
            Self::OperandSizeOverride => vec![0x66],
            Self::Vex2(b1, b2) => vec![0xC5, b1, b2],
            Self::Vex3(b1, b2, b3) => vec![0xC4, b1, b2, b3],
            Self::Evex(b1, b2, b3, b4) => vec![0x62, b1, b2, b3, b4],
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PrefixedBranch {
    pub prefix: BranchPrefix,
    pub branch: BranchInst,
    pub mandatory_prefix: bool,
}

impl PrefixedBranch {
    pub fn new(prefix: BranchPrefix, branch: BranchInst, mandatory: bool) -> Self {
        Self {
            prefix,
            branch,
            mandatory_prefix: mandatory,
        }
    }
    pub fn total_size(&self) -> u32 {
        self.prefix.size() + self.branch.size()
    }
    pub fn compute_displacement(&self, branch_addr: u64, target_addr: u64) -> i64 {
        let branch_end = branch_addr + self.total_size() as u64;
        target_addr as i64 - branch_end as i64
    }
}

// ============================================================================
// Branch inversion
// ============================================================================

pub fn invert_condition(cc: ConditionCode) -> ConditionCode {
    match cc {
        ConditionCode::O => ConditionCode::NO,
        ConditionCode::NO => ConditionCode::O,
        ConditionCode::B => ConditionCode::AE,
        ConditionCode::AE => ConditionCode::B,
        ConditionCode::E => ConditionCode::NE,
        ConditionCode::NE => ConditionCode::E,
        ConditionCode::BE => ConditionCode::A,
        ConditionCode::A => ConditionCode::BE,
        ConditionCode::S => ConditionCode::NS,
        ConditionCode::NS => ConditionCode::S,
        ConditionCode::P => ConditionCode::NP,
        ConditionCode::NP => ConditionCode::P,
        ConditionCode::L => ConditionCode::GE,
        ConditionCode::GE => ConditionCode::L,
        ConditionCode::LE => ConditionCode::G,
        ConditionCode::G => ConditionCode::LE,
    }
}

pub fn invert_branch_with_jmp(
    cc: ConditionCode,
    target_disp: i64,
    fallthrough_disp: i64,
) -> Option<(Vec<u8>, u32)> {
    if fallthrough_disp < -128 || fallthrough_disp > 127 {
        return None;
    }
    let inverted_cc = invert_condition(cc);
    let mut bytes = vec![inverted_cc.opcode8(), fallthrough_disp as u8];
    bytes.push(0xE9);
    bytes.extend_from_slice(&(target_disp as i32).to_le_bytes());
    Some((bytes, 7))
}

// ============================================================================
// Branch-to-branch forwarding
// ============================================================================

#[derive(Debug, Clone, Default)]
pub struct BranchForwarder {
    forwarding: HashMap<usize, usize>,
    pub branches_forwarded: u64,
}

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

    pub fn forward_branches(&mut self, layout: &mut FragmentLayout) {
        self.forwarding.clear();
        for fragment in &layout.fragments {
            if fragment.instructions.len() == 1 && fragment.fixed_size == 0 {
                let inst = &fragment.instructions[0];
                if inst.kind == BranchKind::Jmp || inst.kind == BranchKind::JmpThunk {
                    self.forwarding.insert(fragment.id, inst.target_fragment);
                }
            }
        }
        let mut changed = true;
        while changed {
            changed = false;
            let keys: Vec<usize> = self.forwarding.keys().copied().collect();
            for key in keys {
                let target = self.forwarding[&key];
                if let Some(&next) = self.forwarding.get(&target) {
                    if next != key {
                        self.forwarding.insert(key, next);
                        changed = true;
                    }
                }
            }
        }
        for fragment in layout.fragments.iter_mut() {
            for branch in fragment.instructions.iter_mut() {
                if let Some(&ultimate) = self.forwarding.get(&branch.target_fragment) {
                    branch.target_fragment = ultimate;
                    self.branches_forwarded += 1;
                }
            }
        }
    }
}

// ============================================================================
// Alignment optimization
// ============================================================================

#[derive(Debug, Clone)]
pub struct AlignmentConfig {
    pub function_alignment: u32,
    pub loop_alignment: u32,
    pub branch_target_alignment: u32,
    pub max_padding: u32,
    pub align_branch_targets: bool,
}

impl Default for AlignmentConfig {
    fn default() -> Self {
        Self {
            function_alignment: 16,
            loop_alignment: 16,
            branch_target_alignment: 8,
            max_padding: 15,
            align_branch_targets: false,
        }
    }
}

#[derive(Debug, Clone)]
pub struct AlignmentOptimizer {
    pub config: AlignmentConfig,
    pub hot_fragments: HashMap<usize, u64>,
    pub loop_headers: HashSet<usize>,
}

impl AlignmentOptimizer {
    pub fn new(config: AlignmentConfig) -> Self {
        Self {
            config,
            hot_fragments: HashMap::new(),
            loop_headers: HashSet::new(),
        }
    }
    pub fn mark_loop_header(&mut self, fragment_id: usize) {
        self.loop_headers.insert(fragment_id);
    }
    pub fn mark_frequency(&mut self, fragment_id: usize, freq: u64) {
        self.hot_fragments.insert(fragment_id, freq);
    }
    pub fn apply(&self, layout: &mut FragmentLayout) {
        for fragment in layout.fragments.iter_mut() {
            if fragment.is_function_entry {
                fragment.alignment = Alignment::Pow2(self.config.function_alignment);
            } else if self.loop_headers.contains(&fragment.id) {
                fragment.alignment = Alignment::Pow2(self.config.loop_alignment);
            } else if self.config.align_branch_targets
                && self.hot_fragments.contains_key(&fragment.id)
            {
                fragment.alignment = Alignment::Pow2(self.config.branch_target_alignment);
            }
        }
    }
}

// ============================================================================
// NOP strategies
// ============================================================================

pub mod nop_strategies {
    pub fn intel_optimal(count: u32) -> Vec<u8> {
        super::X86BranchRelaxation::emit_alignment_padding(count)
    }
    pub fn amd_optimal(count: u32) -> Vec<u8> {
        vec![0x90; count as usize]
    }
    pub fn minimal(count: u32) -> Vec<u8> {
        vec![0x90; count as usize]
    }
    pub fn compatible(count: u32) -> Vec<u8> {
        super::X86BranchRelaxation::emit_alignment_padding(count)
    }
}

// ============================================================================
// Fragment reordering
// ============================================================================

#[derive(Debug, Clone, Default)]
pub struct FragmentReorderer {
    edge_frequencies: HashMap<(usize, usize), u64>,
}

impl FragmentReorderer {
    pub fn new() -> Self {
        Self::default()
    }
    pub fn add_edge(&mut self, source: usize, target: usize, freq: u64) {
        *self.edge_frequencies.entry((source, target)).or_insert(0) += freq;
    }
    pub fn reorder(&self, layout: &mut FragmentLayout) {
        if layout.fragments.len() <= 1 {
            return;
        }
        let mut best_successor: HashMap<usize, usize> = HashMap::new();
        for &(src, tgt) in self.edge_frequencies.keys() {
            let freq = self.edge_frequencies.get(&(src, tgt)).copied().unwrap_or(0);
            let current_best = best_successor
                .get(&src)
                .and_then(|&bst| self.edge_frequencies.get(&(src, bst)).copied())
                .unwrap_or(0);
            if freq > current_best {
                best_successor.insert(src, tgt);
            }
        }
        let mut placed: HashSet<usize> = HashSet::new();
        let mut new_order: Vec<CodeFragment> = Vec::new();
        while placed.len() < layout.fragments.len() {
            let start = layout
                .fragments
                .iter()
                .find(|f| !placed.contains(&f.id))
                .map(|f| f.id);
            if let Some(mut current) = start {
                while !placed.contains(&current) {
                    placed.insert(current);
                    if let Some(idx) = layout.index_of(current) {
                        new_order.push(layout.fragments[idx].clone());
                    }
                    if let Some(&next) = best_successor.get(&current) {
                        if !placed.contains(&next) {
                            current = next;
                            continue;
                        }
                    }
                    break;
                }
            }
        }
        layout.fragments = new_order;
        layout.id_to_index.clear();
        for (i, frag) in layout.fragments.iter().enumerate() {
            layout.id_to_index.insert(frag.id, i);
        }
    }
}

// ============================================================================
// Relaxation report
// ============================================================================

#[derive(Debug, Clone, Default)]
pub struct RelaxationReport {
    pub stats: RelaxationStats,
    pub fragment_sizes: Vec<(usize, u32, u32)>,
    pub relaxed_branches: Vec<(usize, usize, u32, u32)>,
    pub aligned_fragments: Vec<(usize, u32)>,
    pub total_before: u64,
    pub total_after: u64,
    pub section_size: u64,
    pub padding_bytes: u64,
}

impl X86BranchRelaxation {
    pub fn generate_report(&self) -> RelaxationReport {
        let mut report = RelaxationReport {
            stats: self.stats.clone(),
            ..Default::default()
        };
        for fragment in &self.layout.fragments {
            let size = fragment.total_size();
            report.total_after += size as u64;
            report.fragment_sizes.push((fragment.id, size, size));
            for (bi, branch) in fragment.instructions.iter().enumerate() {
                if branch.relaxed {
                    let new_size = branch.size();
                    let old_size = branch.kind.instruction_size();
                    report
                        .relaxed_branches
                        .push((fragment.id, bi, old_size, new_size));
                    report.stats.size_increase += (new_size - old_size) as u64;
                }
            }
        }
        report.section_size = self.layout.section_size;
        report.padding_bytes = self.layout.padding_bytes;
        report
    }
}

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

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

    // === BranchKind tests ===

    #[test]
    fn test_branch_kind_can_relax() {
        assert!(BranchKind::Jcc8.can_relax());
        assert!(BranchKind::Jmp.can_relax());
        assert!(BranchKind::Loop.can_relax());
        assert!(!BranchKind::Call32.can_relax());
        assert!(!BranchKind::Jcc32.can_relax());
        assert!(!BranchKind::Ret.can_relax());
    }

    #[test]
    fn test_branch_kind_displacement_range() {
        assert_eq!(BranchKind::Jcc8.displacement_range(), (-128, 127));
        assert_eq!(
            BranchKind::Jcc32.displacement_range(),
            (-2_147_483_648, 2_147_483_647)
        );
        assert_eq!(BranchKind::Ret.displacement_range(), (0, 0));
    }

    #[test]
    fn test_branch_kind_instruction_size() {
        assert_eq!(BranchKind::Jmp.instruction_size(), 2);
        assert_eq!(BranchKind::Jcc8.instruction_size(), 2);
        assert_eq!(BranchKind::Jcc32.instruction_size(), 6);
        assert_eq!(BranchKind::Call32.instruction_size(), 5);
        assert_eq!(BranchKind::Loop.instruction_size(), 2);
        assert_eq!(BranchKind::Ret.instruction_size(), 1);
    }

    #[test]
    fn test_branch_kind_relaxed_size() {
        assert_eq!(BranchKind::Jcc8.relaxed_size(), 6); // 2 → 6
        assert_eq!(BranchKind::Jmp.relaxed_size(), 5); // 2 → 5
        assert_eq!(BranchKind::Loop.relaxed_size(), 3);
    }

    #[test]
    fn test_branch_kind_is_conditional() {
        assert!(BranchKind::Jcc8.is_conditional());
        assert!(BranchKind::Jcc32.is_conditional());
        assert!(BranchKind::Loop.is_conditional());
        assert!(!BranchKind::Jmp.is_conditional());
        assert!(!BranchKind::Call32.is_conditional());
        assert!(!BranchKind::Ret.is_conditional());
    }

    #[test]
    fn test_branch_kind_name() {
        assert_eq!(BranchKind::Jmp.name(), "JMP");
        assert_eq!(BranchKind::Jcc8.name(), "JCC8");
        assert_eq!(BranchKind::Call32.name(), "CALL32");
    }

    // === ConditionCode tests ===

    #[test]
    fn test_condition_code_opcode8() {
        assert_eq!(ConditionCode::E.opcode8(), 0x74);
        assert_eq!(ConditionCode::NE.opcode8(), 0x75);
        assert_eq!(ConditionCode::G.opcode8(), 0x7F);
        assert_eq!(ConditionCode::L.opcode8(), 0x7C);
    }

    #[test]
    fn test_condition_code_opcode32() {
        assert_eq!(ConditionCode::E.opcode32(), 0x84);
        assert_eq!(ConditionCode::NE.opcode32(), 0x85);
        assert_eq!(ConditionCode::G.opcode32(), 0x8F);
        assert_eq!(ConditionCode::L.opcode32(), 0x8C);
    }

    // === BranchInst tests ===

    #[test]
    fn test_branch_inst_compute_displacement_forward() {
        let branch = BranchInst::new(0, BranchKind::Jcc8, 1, Some(ConditionCode::E));
        // Branch at address 0x1000, size 2 → end at 0x1002.
        // Target at 0x1050.
        let disp = branch.compute_displacement(0x1000, 0x1050);
        assert_eq!(disp, 0x1050 - 0x1002); // 0x4E = 78
    }

    #[test]
    fn test_branch_inst_compute_displacement_backward() {
        let branch = BranchInst::new(0, BranchKind::Jmp, 1, None);
        // Branch at 0x1050, size 2 → end at 0x1052.
        // Target at 0x1000.
        let disp = branch.compute_displacement(0x1050, 0x1000);
        assert_eq!(disp, 0x1000i64 - 0x1052i64); // -82
    }

    #[test]
    fn test_branch_inst_needs_relaxation_far_forward() {
        let branch = BranchInst::new(0, BranchKind::Jcc8, 1, Some(ConditionCode::NE));
        assert!(branch.needs_relaxation(200)); // > 127
    }

    #[test]
    fn test_branch_inst_needs_relaxation_far_backward() {
        let branch = BranchInst::new(0, BranchKind::Jcc8, 1, Some(ConditionCode::NE));
        assert!(branch.needs_relaxation(-200)); // < -128
    }

    #[test]
    fn test_branch_inst_no_relaxation_in_range() {
        let branch = BranchInst::new(0, BranchKind::Jcc8, 1, Some(ConditionCode::NE));
        assert!(!branch.needs_relaxation(50));
        assert!(!branch.needs_relaxation(-50));
        assert!(!branch.needs_relaxation(0));
    }

    #[test]
    fn test_branch_inst_relax() {
        let mut branch = BranchInst::new(0, BranchKind::Jcc8, 1, Some(ConditionCode::E));
        assert_eq!(branch.size(), 2);
        let increase = branch.relax();
        assert_eq!(increase, 4); // 6 - 2
        assert!(branch.relaxed);
        assert_eq!(branch.size(), 6);
        // Second relax has no effect.
        let increase2 = branch.relax();
        assert_eq!(increase2, 0);
    }

    #[test]
    fn test_branch_inst_factory_methods() {
        let jcc = BranchInst::jcc8(0, ConditionCode::G, 0);
        assert_eq!(jcc.kind, BranchKind::Jcc8);
        assert_eq!(jcc.cc, Some(ConditionCode::G));

        let jmp = BranchInst::jmp8(0, 0);
        assert_eq!(jmp.kind, BranchKind::Jmp);

        let loop_ = BranchInst::loop_(0, 0);
        assert_eq!(loop_.kind, BranchKind::Loop);

        let call = BranchInst::call32(0, 0);
        assert_eq!(call.kind, BranchKind::Call32);
    }

    // === CodeFragment tests ===

    #[test]
    fn test_fragment_new() {
        let f = CodeFragment::new(5);
        assert_eq!(f.id, 5);
        assert_eq!(f.total_size(), 0);
        assert!(f.address.is_none());
    }

    #[test]
    fn test_fragment_total_size() {
        let mut f = CodeFragment::new(0);
        f.fixed_size = 10;
        f.add_branch(BranchInst::jmp8(0, 1));
        f.add_branch(BranchInst::jcc8(2, ConditionCode::E, 2));
        // 10 + 2 + 2 = 14
        assert_eq!(f.total_size(), 14);
    }

    #[test]
    fn test_fragment_end_address() {
        let mut f = CodeFragment::new(0);
        f.fixed_size = 10;
        f.address = Some(0x1000);
        assert_eq!(f.end_address(), Some(0x100A));
    }

    #[test]
    fn test_fragment_labels() {
        let mut f = CodeFragment::new(0);
        f.add_label(5, "loop_top".to_string());
        f.add_label(20, "loop_end".to_string());
        assert_eq!(f.labels.get(&5).unwrap(), "loop_top");
        assert_eq!(f.labels.get(&20).unwrap(), "loop_end");
    }

    #[test]
    fn test_fragment_alignment() {
        let mut f = CodeFragment::new(0);
        f.set_alignment(Alignment::Pow2(16));
        assert_eq!(f.alignment.value(), 16);
    }

    #[test]
    fn test_fragment_function_entry() {
        let mut f = CodeFragment::new(0);
        f.mark_function_entry();
        assert!(f.is_function_entry);
        assert_eq!(f.alignment.value(), 16);
    }

    #[test]
    fn test_fragment_thunk() {
        let mut f = CodeFragment::new(0);
        f.mark_thunk();
        assert!(f.is_thunk);
    }

    // === Alignment tests ===

    #[test]
    fn test_alignment_pad_to() {
        assert_eq!(Alignment::None.pad_to(5), 0);
        assert_eq!(Alignment::Pow2(16).pad_to(16), 0);
        assert_eq!(Alignment::Pow2(16).pad_to(17), 15);
        assert_eq!(Alignment::Pow2(16).pad_to(0), 0);
        assert_eq!(Alignment::Pow2(8).pad_to(5), 3);
    }

    // === FragmentLayout tests ===

    #[test]
    fn test_layout_add_and_get() {
        let mut layout = FragmentLayout::new(0x1000);
        let f = CodeFragment::new(1);
        layout.add_fragment(f);

        assert_eq!(layout.len(), 1);
        assert!(layout.get_fragment(1).is_some());
        assert!(layout.get_fragment(99).is_none());
    }

    #[test]
    fn test_layout_index_of() {
        let mut layout = FragmentLayout::new(0);
        layout.add_fragment(CodeFragment::new(10));
        layout.add_fragment(CodeFragment::new(20));

        assert_eq!(layout.index_of(10), Some(0));
        assert_eq!(layout.index_of(20), Some(1));
        assert_eq!(layout.index_of(30), None);
    }

    // === X86BranchRelaxation tests ===

    #[test]
    fn test_relaxation_create_fragment() {
        let mut engine = X86BranchRelaxation::new(0);
        let id = engine.create_fragment();
        assert_eq!(id, 1);

        let frag = engine.layout.get_fragment(id).unwrap();
        assert_eq!(frag.id, 1);
    }

    #[test]
    fn test_relaxation_add_fragment() {
        let mut engine = X86BranchRelaxation::new(0);
        let f = CodeFragment::new(42);
        engine.add_fragment(f);
        assert_eq!(engine.layout.len(), 1);
        assert!(engine.layout.get_fragment(42).is_some());
    }

    #[test]
    fn test_relaxation_compute_layout_simple() {
        let mut engine = X86BranchRelaxation::new(0x1000);

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 10;
        engine.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 20;
        f2.set_alignment(Alignment::Pow2(16));
        engine.add_fragment(f2);

        engine.compute_layout();

        let f1_addr = engine.layout.get_fragment(1).unwrap().address.unwrap();
        assert_eq!(f1_addr, 0x1000);

        let f2_addr = engine.layout.get_fragment(2).unwrap().address.unwrap();
        // f1 ends at 0x1000 + 10 = 0x100A. Next 16-byte boundary is 0x1010.
        assert_eq!(f2_addr, 0x1010);
    }

    #[test]
    fn test_relaxation_run_no_relaxation_needed() {
        let mut engine = X86BranchRelaxation::new(0x1000);

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 5;
        let target = 2;
        f1.add_branch(BranchInst::jcc8(5, ConditionCode::E, target));
        engine.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 5;
        engine.add_fragment(f2);

        let passes = engine.run_relaxation();
        assert_eq!(passes, 1); // One pass, no relaxation.

        // Check that the 8-bit branch fits.
        let displacements = engine.recompute_displacements();
        let (_fi, _bi, disp) = displacements[0];
        assert!(disp >= -128 && disp <= 127);
    }

    #[test]
    fn test_relaxation_jcc_overreach() {
        let mut engine = X86BranchRelaxation::new(0x1000);

        // Fragment 1: a JCC that needs to jump past a large gap.
        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 10;
        let target_id = 2;
        f1.add_branch(BranchInst::jcc8(10, ConditionCode::NE, target_id));
        engine.add_fragment(f1);

        // Fragment 2: a large "gap" fragment.
        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 500; // large enough that JCC rel8 won't reach.
        engine.add_fragment(f2);

        let mut f3 = CodeFragment::new(3);
        f3.fixed_size = 10;
        engine.add_fragment(f3);

        // Make fragment 2 the target.
        let branch = &mut engine.layout.get_fragment_mut(1).unwrap().instructions[0];
        branch.target_fragment = 3; // Jump from f1 to f3, over f2.

        let passes = engine.run_relaxation();
        assert!(passes >= 1);
        // The JCC should have been relaxed.
        let branch_final = &engine.layout.get_fragment(1).unwrap().instructions[0];
        assert!(branch_final.relaxed);
    }

    #[test]
    fn test_relaxation_jmp_overreach() {
        let mut engine = X86BranchRelaxation::new(0x1000);

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 3;
        f1.add_branch(BranchInst::jmp8(3, 2));
        engine.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 300;
        engine.add_fragment(f2);

        let passes = engine.run_relaxation();
        let branch = &engine.layout.get_fragment(1).unwrap().instructions[0];
        assert!(branch.relaxed);
        assert!(passes >= 1);
    }

    #[test]
    fn test_relaxation_loop() {
        let mut engine = X86BranchRelaxation::new(0x1000);
        engine.relax_loops = true;

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 5;
        f1.add_branch(BranchInst::loop_(5, 2));
        engine.add_fragment(f1);

        // Fragment 2 is far away.
        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 500;
        engine.add_fragment(f2);

        let mut f3 = CodeFragment::new(3);
        f3.fixed_size = 10;
        engine.add_fragment(f3);

        // Make LOOP target fragment 3.
        let branch = &mut engine.layout.get_fragment_mut(1).unwrap().instructions[0];
        branch.target_fragment = 3;

        let passes = engine.run_relaxation();
        assert!(passes >= 1);
    }

    #[test]
    fn test_relaxation_convergence_limit() {
        let mut engine = X86BranchRelaxation::new(0x1000);
        engine.max_iterations = 3;

        // Create many fragments with cross-jumps to force many relaxation passes.
        for i in 1..=10 {
            let mut f = CodeFragment::new(i);
            f.fixed_size = 100;
            if i < 10 {
                f.add_branch(BranchInst::jcc8(100, ConditionCode::E, i + 1));
            }
            engine.add_fragment(f);
        }

        let passes = engine.run_relaxation();
        // Should converge within max_iterations.
        assert!(passes <= 3);
    }

    // === Relaxation thunk tests ===

    #[test]
    fn test_create_thunk() {
        let mut engine = X86BranchRelaxation::new(0x1000);
        engine.thunk_threshold = 100; // Low threshold to trigger thunk creation.

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 5;
        let target = 2;
        f1.add_branch(BranchInst::jmp8(5, target));
        engine.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 200; // Enough that jmp8 can't reach but also triggers thunk.
        engine.add_fragment(f2);

        // Before relaxation, create a thunk manually for testing.
        let thunk_id = engine.create_thunk(2);
        assert!(engine.thunks.len() == 1);
        let thunk = engine.layout.get_fragment(thunk_id).unwrap();
        assert!(thunk.is_thunk);

        // Redirect branch to thunk.
        let branch = &mut engine.layout.get_fragment_mut(1).unwrap().instructions[0];
        branch.target_fragment = thunk_id;

        let passes = engine.run_relaxation();
        assert!(passes >= 1);
        assert_eq!(engine.stats.thunks_created, 1);
    }

    // === RelaxJCC / RelaxJMP tests ===

    #[test]
    fn test_relax_jcc_encoding() {
        let (old, new) = X86BranchRelaxation::relax_jcc(ConditionCode::E, 5);
        assert_eq!(old, vec![0x74, 5]);
        assert_eq!(new.len(), 6);
        assert_eq!(new[0], 0x0F);
        assert_eq!(new[1], 0x84); // JZ rel32
    }

    #[test]
    fn test_relax_jmp_encoding() {
        let (old, new) = X86BranchRelaxation::relax_jmp(-10);
        assert_eq!(old, vec![0xEB, 0xF6]); // -10 = 0xF6
        assert_eq!(new.len(), 5);
        assert_eq!(new[0], 0xE9);
    }

    #[test]
    fn test_relax_loop_encoding() {
        let (old, new) = X86BranchRelaxation::relax_loop(3);
        assert_eq!(old, vec![0xE2, 3]);
        assert_eq!(new[0], 0x66);
        assert_eq!(new[1], 0xFF);
        assert_eq!(new[2], 0xC9); // DEC ECX
        assert_eq!(new[3], 0x75); // JNZ
        assert_eq!(new[4], 3);
    }

    #[test]
    fn test_relax_jcxz_encoding() {
        let (old, new) = X86BranchRelaxation::relax_jcxz(7);
        assert_eq!(old, vec![0x67, 0xE3, 7]);
        assert_eq!(new[0], 0x85);
        assert_eq!(new[1], 0xC9); // TEST ECX,ECX
        assert_eq!(new[2], 0x74); // JZ
    }

    // === Alignment padding tests ===

    #[test]
    fn test_emit_alignment_padding_small() {
        assert_eq!(X86BranchRelaxation::emit_alignment_padding(0), vec![]);
        assert_eq!(X86BranchRelaxation::emit_alignment_padding(1), vec![0x90]);
        assert_eq!(
            X86BranchRelaxation::emit_alignment_padding(2),
            vec![0x66, 0x90]
        );
        assert_eq!(
            X86BranchRelaxation::emit_alignment_padding(3),
            vec![0x0F, 0x1F, 0x00]
        );
    }

    #[test]
    fn test_emit_alignment_padding_large() {
        let pad = X86BranchRelaxation::emit_alignment_padding(15);
        // 15 = 9 + 6: 9-byte NOP + 6-byte NOP.
        assert_eq!(pad.len(), 15);

        let pad2 = X86BranchRelaxation::emit_alignment_padding(20);
        // 20 = 9 + 9 + 2.
        assert_eq!(pad2.len(), 20);
    }

    // === SectionEncoder tests ===

    #[test]
    fn test_section_encoder_empty() {
        let mut encoder = SectionEncoder::new();
        let relaxation = X86BranchRelaxation::new(0);
        encoder.encode(&relaxation, &[]);
        assert!(encoder.encoded_bytes.is_empty());
    }

    #[test]
    fn test_section_encoder_basic() {
        let mut engine = X86BranchRelaxation::new(0x1000);

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 0;
        f1.add_branch(BranchInst::jmp8(0, 2));
        engine.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 10;
        engine.add_fragment(f2);

        engine.run_relaxation();
        let displacements = engine.recompute_displacements();

        let mut encoder = SectionEncoder::new();
        encoder.encode(&engine, &displacements);
        assert!(!encoder.encoded_bytes.is_empty());
        // Should have JMP rel8 (EB xx).
        assert_eq!(encoder.encoded_bytes[0], 0xEB);
    }

    // === RelaxationPipeline tests ===

    #[test]
    fn test_pipeline_simple() {
        let mut pipeline = RelaxationPipeline::new(0x1000);
        let id1 = pipeline.add_fragment_with_branches(
            0,
            vec![BranchInst::jmp8(0, 0)], // placeholder target
            Alignment::None,
            true,
        );
        let id2 = pipeline.add_fragment_with_branches(10, vec![], Alignment::None, false);

        // Fix the branch target now that we have both IDs.
        pipeline
            .engine
            .layout
            .get_fragment_mut(id1)
            .unwrap()
            .instructions[0]
            .target_fragment = id2;

        let result = pipeline.run();
        assert!(result.is_ok());
        let bytes = result.unwrap();
        assert!(!bytes.is_empty());
    }

    #[test]
    fn test_pipeline_with_relaxation() {
        let mut pipeline = RelaxationPipeline::new(0x1000);

        let id_target = pipeline.add_fragment_with_branches(10, vec![], Alignment::None, false);

        // Fragment with JCC8 that needs to be relaxed.
        let branches = vec![BranchInst::jcc8(0, ConditionCode::NE, id_target)];
        let id_source = pipeline.add_fragment_with_branches(0, branches, Alignment::None, true);

        // Insert a large fragment between source and target to force relaxation.
        let _id_gap = pipeline.add_fragment_with_branches(500, vec![], Alignment::None, false);

        let result = pipeline.run();
        assert!(result.is_ok());
    }

    // === Relaxation statistics tests ===

    #[test]
    fn test_stats_initial() {
        let engine = X86BranchRelaxation::new(0);
        assert_eq!(engine.stats.passes, 0);
        assert_eq!(engine.stats.branches_relaxed, 0);
        assert_eq!(engine.stats.thunks_created, 0);
    }

    #[test]
    fn test_stats_after_relaxation() {
        let mut engine = X86BranchRelaxation::new(0x1000);

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 0;
        f1.add_branch(BranchInst::jcc8(0, ConditionCode::E, 3));
        engine.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 300;
        engine.add_fragment(f2);

        let mut f3 = CodeFragment::new(3);
        f3.fixed_size = 10;
        engine.add_fragment(f3);

        engine.run_relaxation();
        assert!(engine.stats.passes >= 1);
        assert!(engine.stats.branches_relaxed >= 1);
        assert!(engine.stats.branches_inspected >= 1);
    }

    // === Edge case: zero displacement ===

    #[test]
    fn test_zero_displacement_jcc() {
        let branch = BranchInst::jcc8(0, ConditionCode::E, 0);
        // Set branch and target at same address, so displacement is negative of size.
        let disp = branch.compute_displacement(0x1000, 0x1000);
        assert_eq!(disp, -(branch.size() as i64));
        assert!(!branch.needs_relaxation(disp)); // -2 is within range
    }

    // === Edge case: exact boundaries of JCC8 range ===

    #[test]
    fn test_jcc8_boundary_forward() {
        let branch = BranchInst::jcc8(0, ConditionCode::E, 0);
        // 127 is the max forward offset.
        assert!(!branch.needs_relaxation(127));
        assert!(branch.needs_relaxation(128));
    }

    #[test]
    fn test_jcc8_boundary_backward() {
        let branch = BranchInst::jcc8(0, ConditionCode::E, 0);
        assert!(!branch.needs_relaxation(-128));
        assert!(branch.needs_relaxation(-129));
    }

    // === Edge case: JMP target within range ===

    #[test]
    fn test_jmp_in_range() {
        let branch = BranchInst::jmp8(0, 0);
        assert!(!branch.needs_relaxation(50));
        assert!(!branch.needs_relaxation(-50));
    }

    // === Edge case: already relaxed should not re-relax ===

    #[test]
    fn test_already_relaxed() {
        let mut branch = BranchInst::jcc8(0, ConditionCode::E, 0);
        branch.relax();
        // Should not need relaxation even if displacement is out of range.
        assert!(!branch.needs_relaxation(500));
    }

    // === Fragment edge cases ===

    #[test]
    fn test_fragment_with_many_branches() {
        let mut f = CodeFragment::new(0);
        for i in 0..50 {
            f.add_branch(BranchInst::jcc8(i * 2, ConditionCode::NE, i as usize));
        }
        assert_eq!(f.instructions.len(), 50);
        assert_eq!(f.total_size(), 100); // 50 * 2 bytes
    }

    // === Sizes after relaxation of different branch types ===

    #[test]
    fn test_sizes_before_and_after() {
        // JCC8
        let mut jcc = BranchInst::jcc8(0, ConditionCode::E, 0);
        assert_eq!(jcc.size(), 2);
        jcc.relax();
        assert_eq!(jcc.size(), 6);

        // JMP
        let mut jmp = BranchInst::jmp8(0, 0);
        assert_eq!(jmp.size(), 2);
        jmp.relax();
        assert_eq!(jmp.size(), 5);

        // LOOP
        let mut lp = BranchInst::loop_(0, 0);
        assert_eq!(lp.size(), 2);
        lp.relax();
        assert_eq!(lp.size(), 3);
    }

    // === Fragment layout with nested alignment ===

    #[test]
    fn test_layout_with_various_alignments() {
        let mut engine = X86BranchRelaxation::new(0);

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 5;
        f1.set_alignment(Alignment::None);
        engine.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 10;
        f2.set_alignment(Alignment::Pow2(32));
        engine.add_fragment(f2);

        engine.compute_layout();

        // f1 at 0.
        let f1_addr = engine.layout.get_fragment(1).unwrap().address.unwrap();
        assert_eq!(f1_addr, 0);

        // f1 ends at 5, f2 needs 32-byte alignment → 32.
        let f2_addr = engine.layout.get_fragment(2).unwrap().address.unwrap();
        assert_eq!(f2_addr, 32);
    }

    // === recompute_displacements test ===

    #[test]
    fn test_recompute_displacements_basic() {
        let mut engine = X86BranchRelaxation::new(0x1000);

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 0;
        f1.add_branch(BranchInst::jmp8(0, 2));
        engine.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 20;
        engine.add_fragment(f2);

        engine.run_relaxation();
        let disps = engine.recompute_displacements();

        assert_eq!(disps.len(), 1);
        let (_fi, _bi, disp) = disps[0];
        // Branch at 0x1000 + 0, size 2 → end at 0x1002, target at 0x1002 (after f1).
        // Actually: f1 at 0x1000, f2 at 0x1002.
        // displacement = 0x1002 - (0x1000 + 2) = 0.
        assert_eq!(disp, 0);
    }

    // === BranchPrefix tests ===

    #[test]
    fn test_prefix_sizes() {
        assert_eq!(BranchPrefix::None.size(), 0);
        assert_eq!(BranchPrefix::Rex(0x48).size(), 1);
        assert_eq!(BranchPrefix::Vex2(0xF8, 0x00).size(), 2);
        assert_eq!(BranchPrefix::Vex3(0xF8, 0x00, 0x00).size(), 3);
        assert_eq!(BranchPrefix::Evex(0x00, 0x00, 0x00, 0x00).size(), 4);
    }

    #[test]
    fn test_prefix_encode() {
        assert_eq!(BranchPrefix::None.encode(), vec![]);
        assert_eq!(BranchPrefix::Rex(0x48).encode(), vec![0x48]);
        assert_eq!(BranchPrefix::AddrSizeOverride.encode(), vec![0x67]);
        assert_eq!(BranchPrefix::OperandSizeOverride.encode(), vec![0x66]);
        assert_eq!(
            BranchPrefix::Vex2(0xF8, 0x00).encode(),
            vec![0xC5, 0xF8, 0x00]
        );
        assert_eq!(
            BranchPrefix::Vex3(0xF8, 0x00, 0x01).encode(),
            vec![0xC4, 0xF8, 0x00, 0x01]
        );
        assert_eq!(
            BranchPrefix::Evex(0x00, 0x01, 0x02, 0x03).encode(),
            vec![0x62, 0x00, 0x01, 0x02, 0x03]
        );
    }

    #[test]
    fn test_prefixed_branch_size() {
        let branch = BranchInst::jcc8(0, ConditionCode::E, 0);
        let pb = PrefixedBranch::new(BranchPrefix::Rex(0x48), branch, true);
        assert_eq!(pb.total_size(), 3); // 1 prefix + 2 JCC8
    }

    // === Branch inversion tests ===

    #[test]
    fn test_invert_condition_all() {
        assert_eq!(invert_condition(ConditionCode::E), ConditionCode::NE);
        assert_eq!(invert_condition(ConditionCode::NE), ConditionCode::E);
        assert_eq!(invert_condition(ConditionCode::G), ConditionCode::LE);
        assert_eq!(invert_condition(ConditionCode::LE), ConditionCode::G);
        assert_eq!(invert_condition(ConditionCode::L), ConditionCode::GE);
        assert_eq!(invert_condition(ConditionCode::GE), ConditionCode::L);
        assert_eq!(invert_condition(ConditionCode::A), ConditionCode::BE);
        assert_eq!(invert_condition(ConditionCode::B), ConditionCode::AE);
    }

    #[test]
    fn test_invert_branch_fallthrough_in_range() {
        let result = invert_branch_with_jmp(ConditionCode::E, 1000, 5);
        assert!(result.is_some());
        let (bytes, size) = result.unwrap();
        assert_eq!(size, 7);
        assert_eq!(bytes[0], ConditionCode::NE.opcode8());
        assert_eq!(bytes[2], 0xE9); // JMP rel32
    }

    #[test]
    fn test_invert_branch_fallthrough_out_of_range() {
        let result = invert_branch_with_jmp(ConditionCode::E, 1000, 200);
        assert!(result.is_none());
    }

    // === BranchForwarder tests ===

    #[test]
    fn test_forwarder_single_jmp_chain() {
        let mut layout = FragmentLayout::new(0x1000);

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 5;
        f1.add_branch(BranchInst::jmp8(5, 2));
        layout.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 0;
        f2.add_branch(BranchInst::jmp8(0, 3));
        layout.add_fragment(f2);

        let mut f3 = CodeFragment::new(3);
        f3.fixed_size = 10;
        layout.add_fragment(f3);

        let mut forwarder = BranchForwarder::new();
        forwarder.forward_branches(&mut layout);

        // f1's branch should now target f3 directly.
        let branch = &layout.fragments[0].instructions[0];
        assert_eq!(branch.target_fragment, 3);
        assert!(forwarder.branches_forwarded >= 1);
    }

    #[test]
    fn test_forwarder_self_loop_no_change() {
        let mut layout = FragmentLayout::new(0);
        let mut f1 = CodeFragment::new(1);
        f1.add_branch(BranchInst::jmp8(0, 1));
        layout.add_fragment(f1);

        let mut forwarder = BranchForwarder::new();
        forwarder.forward_branches(&mut layout);
        // Should remain a self-loop (forwarding resolves to 1 → 1 which we guard against).
        let branch = &layout.fragments[0].instructions[0];
        assert_eq!(branch.target_fragment, 1);
    }

    // === Alignment optimizer tests ===

    #[test]
    fn test_alignment_optimizer_function_entry() {
        let config = AlignmentConfig::default();
        let opt = AlignmentOptimizer::new(config);

        let mut layout = FragmentLayout::new(0);
        let mut f = CodeFragment::new(1);
        f.mark_function_entry();
        layout.add_fragment(f);

        opt.apply(&mut layout);
        assert_eq!(layout.fragments[0].alignment.value(), 16);
    }

    #[test]
    fn test_alignment_optimizer_loop_header() {
        let mut opt = AlignmentOptimizer::new(AlignmentConfig::default());
        opt.mark_loop_header(2);

        let mut layout = FragmentLayout::new(0);
        layout.add_fragment(CodeFragment::new(1));
        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 10;
        layout.add_fragment(f2);

        opt.apply(&mut layout);
        assert_eq!(layout.fragments[1].alignment.value(), 16);
    }

    // === NOP strategy tests ===

    #[test]
    fn test_nop_intel_optimal() {
        let pad = nop_strategies::intel_optimal(5);
        assert_eq!(pad.len(), 5);
    }

    #[test]
    fn test_nop_amd_optimal() {
        let pad = nop_strategies::amd_optimal(10);
        assert_eq!(pad.len(), 10);
        assert_eq!(pad, vec![0x90; 10]);
    }

    #[test]
    fn test_nop_minimal() {
        let pad = nop_strategies::minimal(3);
        assert_eq!(pad, vec![0x90, 0x90, 0x90]);
    }

    // === FragmentReorderer tests ===

    #[test]
    fn test_reorder_empty() {
        let reorderer = FragmentReorderer::new();
        let mut layout = FragmentLayout::new(0);
        reorderer.reorder(&mut layout);
        assert!(layout.is_empty());
    }

    #[test]
    fn test_reorder_single_fragment() {
        let reorderer = FragmentReorderer::new();
        let mut layout = FragmentLayout::new(0);
        layout.add_fragment(CodeFragment::new(1));
        reorderer.reorder(&mut layout);
        assert_eq!(layout.len(), 1);
    }

    #[test]
    fn test_reorder_with_edges() {
        let mut reorderer = FragmentReorderer::new();
        reorderer.add_edge(1, 2, 100);
        reorderer.add_edge(2, 3, 50);

        let mut layout = FragmentLayout::new(0);
        layout.add_fragment(CodeFragment::new(1));
        layout.add_fragment(CodeFragment::new(3));
        layout.add_fragment(CodeFragment::new(2));

        reorderer.reorder(&mut layout);
        // Should have placed 1, then 2 (hot successor), then 3.
        assert_eq!(layout.fragments[0].id, 1);
        assert_eq!(layout.fragments[1].id, 2);
        assert_eq!(layout.fragments[2].id, 3);
    }

    // === RelaxationReport tests ===

    #[test]
    fn test_report_no_relaxation() {
        let mut engine = X86BranchRelaxation::new(0x1000);
        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 10;
        engine.add_fragment(f1);
        engine.run_relaxation();

        let report = engine.generate_report();
        assert_eq!(report.stats.branches_relaxed, 0);
        assert!(report.fragment_sizes.len() >= 1);
    }

    #[test]
    fn test_report_with_relaxation() {
        let mut engine = X86BranchRelaxation::new(0x1000);
        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 0;
        f1.add_branch(BranchInst::jcc8(0, ConditionCode::NE, 3));
        engine.add_fragment(f1);
        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 300;
        engine.add_fragment(f2);
        let mut f3 = CodeFragment::new(3);
        f3.fixed_size = 10;
        engine.add_fragment(f3);
        engine.run_relaxation();

        let report = engine.generate_report();
        assert!(report.stats.branches_relaxed >= 1);
        assert!(!report.relaxed_branches.is_empty());
    }

    // === Conditional branch stress tests ===

    #[test]
    fn test_all_condition_codes_jcc8() {
        let all_ccs = vec![
            ConditionCode::O,
            ConditionCode::NO,
            ConditionCode::B,
            ConditionCode::AE,
            ConditionCode::E,
            ConditionCode::NE,
            ConditionCode::BE,
            ConditionCode::A,
            ConditionCode::S,
            ConditionCode::NS,
            ConditionCode::P,
            ConditionCode::NP,
            ConditionCode::L,
            ConditionCode::GE,
            ConditionCode::LE,
            ConditionCode::G,
        ];
        for cc in all_ccs {
            let branch = BranchInst::jcc8(0, cc, 0);
            assert_eq!(branch.size(), 2);
            let mut relaxed = branch.clone();
            relaxed.relax();
            assert_eq!(relaxed.size(), 6);
        }
    }

    #[test]
    fn test_all_condition_codes_opcode_consistency() {
        // 8-bit and 32-bit opcodes should encode the same condition.
        let pairs = vec![
            (ConditionCode::E, 0x74, 0x84),
            (ConditionCode::NE, 0x75, 0x85),
            (ConditionCode::L, 0x7C, 0x8C),
            (ConditionCode::G, 0x7F, 0x8F),
        ];
        for (cc, op8, op32) in pairs {
            assert_eq!(cc.opcode8(), op8);
            assert_eq!(cc.opcode32(), op32);
        }
    }

    // === Large section stress test ===

    #[test]
    fn test_large_section_relaxation() {
        let mut engine = X86BranchRelaxation::new(0x1000);
        engine.max_iterations = 50;

        // Create many fragments with cross-branches.
        for i in 1..=50 {
            let mut f = CodeFragment::new(i);
            f.fixed_size = 10;
            if i < 50 {
                // Every other fragment has a backward branch.
                if i % 2 == 0 {
                    f.add_branch(BranchInst::jcc8(10, ConditionCode::E, i - 1));
                }
            }
            engine.add_fragment(f);
        }

        let passes = engine.run_relaxation();
        // Should converge.
        assert!(passes < engine.max_iterations);
    }

    // === Fragment size correctness after multiple relaxations ===

    #[test]
    fn test_size_increases_monotonically() {
        let mut engine = X86BranchRelaxation::new(0x1000);

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 0;
        // Three JCC8's, all targeting far fragments.
        f1.add_branch(BranchInst::jcc8(0, ConditionCode::E, 4));
        f1.add_branch(BranchInst::jcc8(2, ConditionCode::NE, 5));
        f1.add_branch(BranchInst::jcc8(4, ConditionCode::G, 6));
        engine.add_fragment(f1);

        // Gaps to force relaxation.
        for i in 2..=6 {
            let mut f = CodeFragment::new(i);
            f.fixed_size = 300;
            engine.add_fragment(f);
        }

        let initial_size = engine.layout.fragments[0].total_size();
        engine.run_relaxation();
        let final_size = engine.layout.fragments[0].total_size();

        // Each JCC8 grows from 2 to 6 → +4 each, total +12.
        assert_eq!(final_size, initial_size + 12);
    }

    // === JMP relaxation boundary tests ===

    #[test]
    fn test_jmp_relax_at_boundary() {
        let mut branch = BranchInst::jmp8(0, 0);
        assert!(!branch.needs_relaxation(127));
        assert!(branch.needs_relaxation(128));
        assert!(!branch.needs_relaxation(-128));
        assert!(branch.needs_relaxation(-129));
    }

    #[test]
    fn test_jmp_relax_encoding_roundtrip() {
        let (old, new) = X86BranchRelaxation::relax_jmp(42);
        assert_eq!(old.len(), 2);
        assert_eq!(new.len(), 5);
        assert_eq!(old[0], 0xEB);
        assert_eq!(new[0], 0xE9);
        let disp = i32::from_le_bytes(new[1..5].try_into().unwrap());
        assert_eq!(disp, 42);
    }

    // === Section encoder post-relaxation byte verification ===

    #[test]
    fn test_encoder_jcc_relaxed_uses_rel32() {
        let mut engine = X86BranchRelaxation::new(0x1000);

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 0;
        f1.add_branch(BranchInst::jcc8(0, ConditionCode::E, 3));
        engine.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 300;
        engine.add_fragment(f2);

        let mut f3 = CodeFragment::new(3);
        f3.fixed_size = 10;
        engine.add_fragment(f3);

        engine.run_relaxation();
        let disps = engine.recompute_displacements();

        let mut encoder = SectionEncoder::new();
        encoder.encode(&engine, &disps);
        let bytes = &encoder.encoded_bytes;

        // After relaxation, JCC should be rel32: 0F 8x xx xx xx xx (6 bytes).
        assert!(!bytes.is_empty());
        assert_eq!(bytes[0], 0x0F);
        // The second byte depends on the condition code.
        assert_eq!(bytes[1], ConditionCode::E.opcode32());
    }

    #[test]
    fn test_encoder_jcc_not_relaxed_uses_rel8() {
        let mut engine = X86BranchRelaxation::new(0x1000);

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 0;
        f1.add_branch(BranchInst::jcc8(0, ConditionCode::NE, 2));
        engine.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 10;
        engine.add_fragment(f2);

        engine.run_relaxation();
        let disps = engine.recompute_displacements();

        let mut encoder = SectionEncoder::new();
        encoder.encode(&engine, &disps);
        let bytes = &encoder.encoded_bytes;

        // Short JCC: 7x xx (2 bytes).
        assert_eq!(bytes.len(), 12); // f1 branch(2) + f2 fixed(10)
        assert_eq!(bytes[0], ConditionCode::NE.opcode8());
    }

    #[test]
    fn test_encoder_loop_relaxed() {
        let mut engine = X86BranchRelaxation::new(0x1000);
        engine.relax_loops = true;

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 0;
        f1.add_branch(BranchInst::loop_(0, 3));
        engine.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 300;
        engine.add_fragment(f2);

        let mut f3 = CodeFragment::new(3);
        f3.fixed_size = 10;
        engine.add_fragment(f3);

        engine.run_relaxation();
        let disps = engine.recompute_displacements();

        let mut encoder = SectionEncoder::new();
        encoder.encode(&engine, &disps);
        let bytes = &encoder.encoded_bytes;

        // Relaxed LOOP → DEC ECX (66 FF C9) + JNZ rel8 (75 xx) = 5 bytes.
        assert!(bytes.len() >= 5);
        assert_eq!(bytes[0], 0x66);
        assert_eq!(bytes[1], 0xFF);
        assert_eq!(bytes[2], 0xC9);
        assert_eq!(bytes[3], 0x75);
    }

    // === Loope/Loopne encoding tests ===

    #[test]
    fn test_loope_encoding() {
        let branch = BranchInst::loope(0, 1);
        assert_eq!(branch.kind, BranchKind::Loope);
        assert_eq!(branch.size(), 2);
        assert!(branch.kind.can_relax());
    }

    #[test]
    fn test_loopne_encoding() {
        let branch = BranchInst::loopne(0, 1);
        assert_eq!(branch.kind, BranchKind::Loopne);
        assert_eq!(branch.size(), 2);
    }

    // === Call32 encoding tests ===

    #[test]
    fn test_call32_displacement() {
        let branch = BranchInst::call32(0, 1);
        assert_eq!(branch.kind, BranchKind::Call32);
        assert_eq!(branch.size(), 5);
        assert!(!branch.kind.can_relax()); // CALL32 can't be relaxed further.
    }

    // === Fragment with mixed branch types ===

    #[test]
    fn test_mixed_branch_types_in_fragment() {
        let mut f = CodeFragment::new(0);
        f.add_branch(BranchInst::jmp8(0, 1));
        f.add_branch(BranchInst::jcc8(2, ConditionCode::E, 2));
        f.add_branch(BranchInst::call32(4, 3));
        f.add_branch(BranchInst::loop_(9, 4));
        // sizes: 2 + 2 + 5 + 2 = 11
        assert_eq!(f.total_size(), 11);
    }

    // === Alignment padding large count test ===

    #[test]
    fn test_alignment_padding_exact_sizes() {
        for n in 0..=15 {
            let pad = X86BranchRelaxation::emit_alignment_padding(n);
            assert_eq!(
                pad.len(),
                n as usize,
                "Padding for {n} should produce exactly {n} bytes"
            );
        }
    }

    #[test]
    fn test_alignment_padding_very_large() {
        let pad = X86BranchRelaxation::emit_alignment_padding(100);
        assert_eq!(pad.len(), 100);
    }

    // === Relaxation with alignment interactions ===

    #[test]
    fn test_relaxation_with_alignment_expands() {
        let mut engine = X86BranchRelaxation::new(0x1000);

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 0;
        f1.add_branch(BranchInst::jcc8(0, ConditionCode::E, 3));
        f1.set_alignment(Alignment::Pow2(64)); // large alignment
        engine.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 300;
        engine.add_fragment(f2);

        let mut f3 = CodeFragment::new(3);
        f3.fixed_size = 10;
        engine.add_fragment(f3);

        engine.run_relaxation();
        // Should still converge.
        assert!(engine.stats.passes >= 1);
    }

    // === BranchInst clone / equality tests ===

    #[test]
    fn test_branch_inst_equality() {
        let b1 = BranchInst::jcc8(5, ConditionCode::E, 10);
        let b2 = BranchInst::jcc8(5, ConditionCode::E, 10);
        assert_eq!(b1, b2);

        let b3 = BranchInst::jcc8(5, ConditionCode::NE, 10);
        assert_ne!(b1, b3);
    }

    // === Fragment clone test ===

    #[test]
    fn test_fragment_clone() {
        let mut f = CodeFragment::new(42);
        f.fixed_size = 10;
        f.add_branch(BranchInst::jmp8(0, 99));
        let cloned = f.clone();
        assert_eq!(cloned.id, f.id);
        assert_eq!(cloned.fixed_size, f.fixed_size);
        assert_eq!(cloned.instructions.len(), f.instructions.len());
    }

    // === Relaxation pipeline non-convergence ===

    #[test]
    fn test_pipeline_non_convergence() {
        let mut pipeline = RelaxationPipeline::new(0x1000);
        pipeline.engine.max_iterations = 1; // Force non-convergence.

        // Create fragments that will need multiple passes.
        let id1 = pipeline.add_fragment_with_branches(
            0,
            vec![BranchInst::jcc8(0, ConditionCode::E, 0)],
            Alignment::None,
            true,
        );
        let id2 = pipeline.add_fragment_with_branches(500, vec![], Alignment::None, false);
        let id3 = pipeline.add_fragment_with_branches(10, vec![], Alignment::None, false);

        // Set the JCC target to id3.
        pipeline
            .engine
            .layout
            .get_fragment_mut(id1)
            .unwrap()
            .instructions[0]
            .target_fragment = id3;

        let result = pipeline.run();
        assert!(result.is_err());
    }

    // === Forwarder transitive chains ===

    #[test]
    fn test_forwarder_transitive_chain() {
        let mut layout = FragmentLayout::new(0);

        // f1: JMP → f2 (JMP → f3 (real code)).
        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 0;
        f1.add_branch(BranchInst::jmp8(0, 2));
        layout.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 0;
        f2.add_branch(BranchInst::jmp8(0, 3));
        layout.add_fragment(f2);

        let mut f3 = CodeFragment::new(3);
        f3.fixed_size = 10;
        layout.add_fragment(f3);

        let mut forwarder = BranchForwarder::new();
        forwarder.forward_branches(&mut layout);

        // f1 should now target f3.
        assert_eq!(layout.fragments[0].instructions[0].target_fragment, 3);
        assert!(forwarder.branches_forwarded >= 2);
    }

    // === AlignmentOptimizer hot fragments ===

    #[test]
    fn test_alignment_optimizer_hot_fragment() {
        let mut config = AlignmentConfig::default();
        config.align_branch_targets = true;
        let mut opt = AlignmentOptimizer::new(config);
        opt.mark_frequency(2, 1000);

        let mut layout = FragmentLayout::new(0);
        layout.add_fragment(CodeFragment::new(1));
        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 10;
        layout.add_fragment(f2);

        opt.apply(&mut layout);
        assert_eq!(layout.fragments[1].alignment.value(), 8);
    }

    // === FragmentReorderer edge cases ===

    #[test]
    fn test_reorder_preserves_all_fragments() {
        let reorderer = FragmentReorderer::new();
        let mut layout = FragmentLayout::new(0);
        let ids: Vec<usize> = vec![10, 20, 30, 40, 50];
        for &id in &ids {
            layout.add_fragment(CodeFragment::new(id));
        }
        reorderer.reorder(&mut layout);

        // All fragments should still be present.
        let mut present: HashSet<usize> = layout.fragments.iter().map(|f| f.id).collect();
        for &id in &ids {
            assert!(present.contains(&id));
            present.remove(&id);
        }
        assert!(present.is_empty());
    }

    // === PrefixedBranch displacement test ===

    #[test]
    fn test_prefixed_branch_displacement() {
        let branch = BranchInst::jmp8(0, 1);
        let pb = PrefixedBranch::new(BranchPrefix::Rex(0x48), branch, false);
        // Branch at 0x1000, prefix + JMP rel8 = 1 + 2 = 3 bytes.
        // End at 0x1003, target at 0x1010.
        let disp = pb.compute_displacement(0x1000, 0x1010);
        assert_eq!(disp, 0x1010 - 0x1003);
    }

    // === Report correctness after no relaxation ===

    #[test]
    fn test_report_empty() {
        let mut engine = X86BranchRelaxation::new(0);
        engine.run_relaxation();
        let report = engine.generate_report();
        assert_eq!(report.stats.passes, 1);
        assert_eq!(report.stats.branches_relaxed, 0);
    }

    // === Edge: fragment with zero instructions and zero fixed size ===

    #[test]
    fn test_empty_fragment() {
        let mut f = CodeFragment::new(99);
        assert_eq!(f.total_size(), 0);
        assert_eq!(f.end_address(), None);
        f.address = Some(0x1000);
        assert_eq!(f.end_address(), Some(0x1000));
    }

    // === Prefix VEX2/VEX3/EVEV sizes correctly accounted ===

    #[test]
    fn test_vex_evex_prefix_sizes_in_prefixed_branch() {
        let branch = BranchInst::jmp8(0, 1);

        let pb_vex2 = PrefixedBranch::new(BranchPrefix::Vex2(0xF8, 0x00), branch.clone(), false);
        assert_eq!(pb_vex2.total_size(), 4); // 2 prefix + 2 JMP

        let pb_vex3 =
            PrefixedBranch::new(BranchPrefix::Vex3(0xF8, 0x00, 0x01), branch.clone(), false);
        assert_eq!(pb_vex3.total_size(), 5); // 3 prefix + 2 JMP

        let pb_evex =
            PrefixedBranch::new(BranchPrefix::Evex(0x00, 0x01, 0x02, 0x03), branch, false);
        assert_eq!(pb_evex.total_size(), 6); // 4 prefix + 2 JMP
    }

    // === Segment override prefix test ===

    #[test]
    fn test_segment_override_prefix() {
        let prefix = BranchPrefix::SegmentOverride(0x64); // FS segment
        assert_eq!(prefix.size(), 1);
        assert_eq!(prefix.encode(), vec![0x64]);
    }

    // === Relaxation with conditional branch to immediate next fragment ===

    #[test]
    fn test_jcc_to_next_fragment_no_relaxation() {
        let mut engine = X86BranchRelaxation::new(0x1000);
        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 0;
        f1.add_branch(BranchInst::jcc8(0, ConditionCode::E, 2));
        engine.add_fragment(f1);
        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 10;
        engine.add_fragment(f2);

        engine.run_relaxation();
        let branch = &engine.layout.fragments[0].instructions[0];
        assert!(
            !branch.relaxed,
            "JCC to next fragment should not need relaxation"
        );
    }

    // === Backward branch within range ===

    #[test]
    fn test_backward_branch_in_range() {
        let mut engine = X86BranchRelaxation::new(0x1000);

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 10;
        engine.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 0;
        f2.add_branch(BranchInst::jcc8(0, ConditionCode::NE, 1)); // back to f1.
        engine.add_fragment(f2);

        engine.run_relaxation();
        let branch = &engine.layout.fragments[1].instructions[0];
        assert!(!branch.relaxed);
    }

    // === Multiple branches in single fragment all needing relaxation ===

    #[test]
    fn test_multiple_branches_all_relax() {
        let mut engine = X86BranchRelaxation::new(0x1000);

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 0;
        f1.add_branch(BranchInst::jcc8(0, ConditionCode::E, 4));
        f1.add_branch(BranchInst::jcc8(2, ConditionCode::NE, 5));
        f1.add_branch(BranchInst::jcc8(4, ConditionCode::G, 6));
        f1.add_branch(BranchInst::jmp8(6, 7));
        engine.add_fragment(f1);

        for i in 2..=7 {
            let mut f = CodeFragment::new(i);
            f.fixed_size = 500;
            engine.add_fragment(f);
        }

        engine.run_relaxation();
        let final_f1 = &engine.layout.fragments[0];
        for branch in &final_f1.instructions {
            assert!(branch.relaxed, "Branch {:?} should be relaxed", branch.kind);
        }
    }

    // === Verify JCC32 is never relaxed ===

    #[test]
    fn test_jcc32_never_relaxed() {
        let branch = BranchInst::new(0, BranchKind::Jcc32, 0, Some(ConditionCode::E));
        assert!(!branch.kind.can_relax());
        assert!(!branch.needs_relaxation(2_147_483_648)); // beyond 32-bit range but can't relax
    }

    // === Verify section_size after layout with multiple fragments ===

    #[test]
    fn test_section_size_multi_fragment() {
        let mut engine = X86BranchRelaxation::new(0x1000);
        for i in 1..=5 {
            let mut f = CodeFragment::new(i);
            f.fixed_size = 100;
            engine.add_fragment(f);
        }
        engine.compute_layout();
        // 5 fragments * 100 = 500.
        assert_eq!(engine.layout.section_size, 500);
    }

    // === Alignment padding in layout ===

    #[test]
    fn test_alignment_padding_bytes() {
        let mut engine = X86BranchRelaxation::new(0x1000);

        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 5;
        f1.set_alignment(Alignment::None);
        engine.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 10;
        f2.set_alignment(Alignment::Pow2(16));
        engine.add_fragment(f2);

        engine.compute_layout();
        // f1 at 0x1000..0x1005, f2 needs 16-alignment → 0x1010.
        // Padding from 0x1005 to 0x1010 = 11 bytes.
        assert_eq!(engine.layout.padding_bytes, 11);
    }

    // === Forwarder with cycle detection ===

    #[test]
    fn test_forwarder_cycle() {
        let mut layout = FragmentLayout::new(0);

        // f1 → f2, f2 → f1 (cycle).
        let mut f1 = CodeFragment::new(1);
        f1.fixed_size = 0;
        f1.add_branch(BranchInst::jmp8(0, 2));
        layout.add_fragment(f1);

        let mut f2 = CodeFragment::new(2);
        f2.fixed_size = 0;
        f2.add_branch(BranchInst::jmp8(0, 1));
        layout.add_fragment(f2);

        let mut forwarder = BranchForwarder::new();
        forwarder.forward_branches(&mut layout);

        // Cycle should not cause infinite loop; the guard prevents self-redirect.
        // After forwarding: 1→2→1 becomes 1→2 (since 2→1 is guarded).
        // The exact behavior depends on implementation; either way it should not panic.
        assert!(forwarder.branches_forwarded <= 2);
    }

    // === Fragment layout preserves order on no reorder ===

    #[test]
    fn test_fragment_layout_preserves_add_order() {
        let mut layout = FragmentLayout::new(0);
        layout.add_fragment(CodeFragment::new(5));
        layout.add_fragment(CodeFragment::new(3));
        layout.add_fragment(CodeFragment::new(1));

        assert_eq!(layout.index_of(5), Some(0));
        assert_eq!(layout.index_of(3), Some(1));
        assert_eq!(layout.index_of(1), Some(2));
    }
}