llvm-native-core 0.1.10

//! X86 Retire Control — Speculative Execution Barrier & Prediction Control
//!
//! This module provides comprehensive retire-time control for X86/X86-64
//! processors. It manages speculative execution barriers (LFENCE dispatch,
//! serializing instructions), indirect branch prediction barriers (IBPB/IBRS/
//! STIBP), Speculative Store Bypass Disable (SSBD), Predictive Store
//! Forwarding Disable (PSFD), return stack buffer (RSB) stuffing, and branch
//! target injection (BTI) mitigations.
//!
//! Phase 13 — LLVM.TARGET.X86.RETIRE.CONTROL Court.
//!
//! Clean-room behavioral reconstruction from:
//! - Intel® 64 and IA-32 Architectures Software Developer's Manual
//!   (volumes 1, 2A, 2B, 2C, 2D, 3A, 3B, 3C, 3D)
//! - Intel® Analysis of Speculative Execution Side Channels (Whitepaper 336983)
//! - Speculative Execution Side Channel Mitigations (Revision 3.0)
//! - AMD64 Architecture Programmer's Manual (volumes 1-5)
//! - AMD Software Techniques for Managing Speculation (AMD-SB-1040)
//! - Linux kernel documentation: Documentation/admin-guide/hw-vuln/
//! - Linux kernel: arch/x86/include/asm/nospec-branch.h
//! - Linux kernel: arch/x86/kernel/cpu/bugs.c
//! - Intel® Deep Dive: Retbleed (CVE-2022-29901, CVE-2022-29900)
//! - Intel® Branch History Injection (BHI) / Spectre-BHB
//! - Google Project Zero: Straight-line Speculation (SLS)
//!
//! Zero LLVM source code consultation. All behavior reconstructed from
//! published specifications, Intel/AMD manuals, and black-box oracle
//! interrogation.
//!
//! ## Architecture
//!
//! ```text
//! ┌──────────────────────────────────────────────────────────────────┐
//! │                       X86RetireControl                            │
//! │  ┌──────────────┐ ┌──────────────┐ ┌──────────────┐             │
//! │  │ LFENCE       │ │ IBPB/IBRS    │ │ RSB Stuffing │             │
//! │  │ Dispatch     │ │ STIBP/SSBD   │ │ BTI Mitigate │             │
//! │  └──────────────┘ └──────────────┘ └──────────────┘             │
//! │  ┌──────────────┐ ┌──────────────┐ ┌──────────────┐             │
//! │  │ Serializing  │ │ PSFD Control │ │ SLH Mitigate │             │
//! │  │ Instructions │ │ SSB Control  │ │ Retbleed     │             │
//! │  └──────────────┘ └──────────────┘ └──────────────┘             │
//! └──────────────────────────────────────────────────────────────────┘
//! ```
//!
//! ## Feature Coverage
//!
//! | Feature                       | LFENCE | Serial | IBPB | RSB | BTI |
//! |-------------------------------|--------|--------|------|-----|-----|
//! | Speculative load hardening    |   ✓    |        |      |     |  ✓  |
//! | Indirect branch barrier       |        |        |  ✓   |     |     |
//! | Return stack buffer stuffing  |        |        |      |  ✓  |     |
//! | SSBD / SSB control            |        |   ✓    |      |     |     |
//! | PSFD speculation control      |        |        |      |     |  ✓  |
//! | LFENCE as dispatch serializing|   ✓    |        |      |     |     |
//! | Serializing instruction fence |        |   ✓    |      |     |     |
//! | Retbleed mitigation           |   ✓    |        |  ✓   |  ✓  |     |
//! | Branch target injection       |        |        |  ✓   |     |  ✓  |
//! | IBPB on context switch        |        |        |  ✓   |     |     |
//! | STIBP for SMT/HyperThreading  |        |        |  ✓   |     |     |
//!
//! # Safety
//!
//! This module contains low-level CPU barrier and control operations.
//! Many functions are inherently unsafe and are marked as such. Callers
//! must understand the microarchitectural implications of inserting
//! speculation barriers, especially regarding performance impact.

#![allow(non_upper_case_globals, dead_code, non_snake_case)]

// ============================================================================
// Imports
// ============================================================================

use crate::codegen::{
    MachineBasicBlock, MachineFunction, MachineInstr, MachineOperand, PhysReg, VirtReg,
};
use crate::x86::x86_instr_info::{X86InstrDesc, X86InstrInfo, X86Opcode};
use crate::x86::X86Subtarget;
use std::arch::asm;
use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::fmt;

// ============================================================================
// Global Constants
// ============================================================================

/// LFENCE opcode: 0F AE E8 (serializing on recent Intel, dispatch on older).
pub const X86_LFENCE_OPCODE: [u8; 3] = [0x0F, 0xAE, 0xE8];

/// MFENCE opcode: 0F AE F0 (full memory fence).
pub const X86_MFENCE_OPCODE: [u8; 3] = [0x0F, 0xAE, 0xF0];

/// SFENCE opcode: 0F AE F8 (store fence).
pub const X86_SFENCE_OPCODE: [u8; 3] = [0x0F, 0xAE, 0xF8];

/// IBPB indirect branch prediction barrier (on supported μarch).
/// MSR: IA32_PRED_CMD (0x49), value 0x1 triggers IBPB.
pub const X86_IBPB_MSR: u32 = 0x49;
pub const X86_IBPB_CMD: u64 = 0x01;

/// STIBP single-thread indirect branch predictor (on supported μarch).
/// MSR: IA32_SPEC_CTRL (0x48), bit 1 = STIBP enable.
pub const X86_SPEC_CTRL_MSR: u32 = 0x48;
pub const X86_STIBP_BIT: u64 = 0x02;
pub const X86_IBRS_BIT: u64 = 0x01;
pub const X86_SSBD_BIT: u64 = 0x04;
pub const X86_PSFD_BIT: u64 = 0x08;

/// IA32_ARCH_CAPABILITIES MSR — indicates which mitigations are available.
pub const X86_ARCH_CAP_MSR: u32 = 0x10A;
pub const X86_ARCH_CAP_IBRS_ALL: u64 = 0x01;
pub const X86_ARCH_CAP_RDCL_NO: u64 = 0x02;
pub const X86_ARCH_CAP_IBPB_ALL: u64 = 0x04;
pub const X86_ARCH_CAP_RSBA: u64 = 0x08;
pub const X86_ARCH_CAP_SKIP_VMENTRY_L1DFLUSH: u64 = 0x10;
pub const X86_ARCH_CAP_SSB_NO: u64 = 0x20;
pub const X86_ARCH_CAP_MDS_NO: u64 = 0x40;
pub const X86_ARCH_CAP_PSCHANGE_MC_NO: u64 = 0x80;
pub const X86_ARCH_CAP_TSX_CTRL: u64 = 0x100;
pub const X86_ARCH_CAP_TAA_NO: u64 = 0x200;
pub const X86_ARCH_CAP_FBSDP_NO: u64 = 0x400;
pub const X86_ARCH_CAP_PBRSB_NO: u64 = 0x800;
pub const X86_ARCH_CAP_GDS_NO: u64 = 0x1000;
pub const X86_ARCH_CAP_RFDS_NO: u64 = 0x2000;
pub const X86_ARCH_CAP_RFDS_CLEAR: u64 = 0x4000;

/// Return Stack Buffer (RSB) depth for common μarchs.
/// Intel Skylake: 16, Haswell: 16, Ice Lake: 32, AMD Zen: 32.
pub const X86_RSB_DEPTH_INTEL_SKYLAKE: usize = 16;
pub const X86_RSB_DEPTH_INTEL_ICELAKE: usize = 32;
pub const X86_RSB_DEPTH_INTEL_ALDERLAKE: usize = 34;
pub const X86_RSB_DEPTH_AMD_ZEN: usize = 32;
pub const X86_RSB_DEPTH_AMD_ZEN3: usize = 32;
pub const X86_RSB_DEPTH_AMD_ZEN4: usize = 32;
pub const X86_RSB_DEPTH_DEFAULT: usize = 16;

/// Maximum RSB stuffing depth for aggressive mitigation.
pub const X86_RSB_STUFF_MAX_DEPTH: usize = 64;

/// Default number of RSB stuffing CALL instructions to emit.
pub const X86_RSB_STUFF_COUNT: usize = 32;

/// Enhanced IBRS: always-on mode where IBRS is set once and never cleared.
/// Bit 3 in IA32_SPEC_CTRL (0x48) indicates "enhanced IBRS" availability.
pub const X86_ENHANCED_IBRS_BIT: u64 = 0x08;

/// Indirect Branch Restricted Speculation (IBRS) modes.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum X86IBRSMode {
    /// No IBRS mitigation.
    None,
    /// Kernel-only IBRS: set on entry, clear on exit.
    KernelOnly,
    /// Enhanced IBRS: set once at boot, never cleared.
    Enhanced,
    /// Retpoline-aware IBRS: used alongside retpoline.
    RetpolineAware,
}

impl fmt::Display for X86IBRSMode {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            X86IBRSMode::None => write!(f, "none"),
            X86IBRSMode::KernelOnly => write!(f, "kernel-only"),
            X86IBRSMode::Enhanced => write!(f, "enhanced"),
            X86IBRSMode::RetpolineAware => write!(f, "retpoline-aware"),
        }
    }
}

/// Indirect Branch Prediction Barrier (IBPB) trigger points.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum X86IBPBTrigger {
    /// Never issue IBPB.
    Never,
    /// On every context switch (user → kernel → user).
    OnContextSwitch,
    /// On every privilege transition.
    OnPrivilegeTransition,
    /// On every indirect branch (aggressive; high perf cost).
    OnIndirectBranch,
    /// Always issue IBPB before every user-space return.
    AlwaysOnUserReturn,
}

impl fmt::Display for X86IBPBTrigger {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            X86IBPBTrigger::Never => write!(f, "never"),
            X86IBPBTrigger::OnContextSwitch => write!(f, "on-context-switch"),
            X86IBPBTrigger::OnPrivilegeTransition => write!(f, "on-privilege-transition"),
            X86IBPBTrigger::OnIndirectBranch => write!(f, "on-indirect-branch"),
            X86IBPBTrigger::AlwaysOnUserReturn => write!(f, "always-on-user-return"),
        }
    }
}

/// Single Thread Indirect Branch Predictor (STIBP) modes.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum X86STIBPMode {
    /// STIBP disabled.
    Off,
    /// STIBP enabled when SMT (Hyper-Threading) is active.
    SMTOnly,
    /// STIBP always enabled.
    Always,
    /// STIBP enabled only for processes that opt-in via prctl().
    PerProcess,
    /// STIBP enabled conditionally based on seccomp filter.
    Conditional,
}

impl fmt::Display for X86STIBPMode {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            X86STIBPMode::Off => write!(f, "off"),
            X86STIBPMode::SMTOnly => write!(f, "smt-only"),
            X86STIBPMode::Always => write!(f, "always"),
            X86STIBPMode::PerProcess => write!(f, "per-process"),
            X86STIBPMode::Conditional => write!(f, "conditional"),
        }
    }
}

/// Speculative Store Bypass Disable (SSBD) modes.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum X86SSBDMode {
    /// SSBD disabled (store-to-load forwarding allowed).
    Off,
    /// SSBD enabled via MSR (SPEC_CTRL bit 2).
    ViaMSR,
    /// SSBD enabled via prctl() / seccomp (per-thread).
    PerThread,
    /// SSBD enabled via AMD's SSBD bit in LS_CFG MSR.
    ViaLSConfig,
    /// SSBD forced on (spec_store_bypass_disable=on kernel param).
    ForceOn,
    /// Auto: use prctl/seccomp when available, else MSR.
    Auto,
}

impl fmt::Display for X86SSBDMode {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            X86SSBDMode::Off => write!(f, "off"),
            X86SSBDMode::ViaMSR => write!(f, "via-msr"),
            X86SSBDMode::PerThread => write!(f, "per-thread"),
            X86SSBDMode::ViaLSConfig => write!(f, "via-ls-cfg"),
            X86SSBDMode::ForceOn => write!(f, "force-on"),
            X86SSBDMode::Auto => write!(f, "auto"),
        }
    }
}

/// Speculative execution barrier types — which kind of fence to insert.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum X86SpecBarrier {
    /// LFENCE: serializes instruction dispatch (load fence for speculation).
    LFENCE,
    /// MFENCE: full memory ordering fence; also serializes.
    MFENCE,
    /// SFENCE: store fence; does NOT serialize execution.
    SFENCE,
    /// CPUID: architecturally serializing; very expensive (~200+ cycles).
    CPUID,
    /// IRET-to-self: serialize via IRET (expensive, for user-mode).
    IRETToSelf,
    /// RDTSC serialization: RDTSCP or LFENCE+RDTSC.
    RDTSCPSerialize,
    /// Serialize instruction (SERIALIZE, available on Alder Lake+).
    Serialize,
    /// IBPB: Indirect Branch Prediction Barrier (via MSR write).
    IBPB,
    /// IBRS set: turn on Indirect Branch Restricted Speculation.
    IBRS_ON,
    /// IBRS clear: turn off Indirect Branch Restricted Speculation.
    IBRS_OFF,
    /// STIBP set: turn on Single Thread Indirect Branch Predictor.
    STIBP_ON,
    /// STIBP clear: turn off Single Thread Indirect Branch Predictor.
    STIBP_OFF,
    /// SSBD set: turn on Speculative Store Bypass Disable.
    SSBD_ON,
    /// SSBD clear: turn off Speculative Store Bypass Disable.
    SSBD_OFF,
    /// PSFD set: Predictive Store Forwarding Disable.
    PSFD_ON,
    /// PSFD clear: Predictive Store Forwarding Disable off.
    PSFD_OFF,
    /// VERW: overwrite CPU buffers (MDS/TAA mitigation).
    VERW,
    /// RSB stuffing: fill return stack buffer with dummy entries.
    RSBStuff,
    /// Fill CPU buffers (L1D flush or VERW sequence).
    FillBuffers,
}

impl fmt::Display for X86SpecBarrier {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let s = match self {
            X86SpecBarrier::LFENCE => "LFENCE",
            X86SpecBarrier::MFENCE => "MFENCE",
            X86SpecBarrier::SFENCE => "SFENCE",
            X86SpecBarrier::CPUID => "CPUID",
            X86SpecBarrier::IRETToSelf => "IRET-to-self",
            X86SpecBarrier::RDTSCPSerialize => "RDTSCP",
            X86SpecBarrier::Serialize => "SERIALIZE",
            X86SpecBarrier::IBPB => "IBPB",
            X86SpecBarrier::IBRS_ON => "IBRS-ON",
            X86SpecBarrier::IBRS_OFF => "IBRS-OFF",
            X86SpecBarrier::STIBP_ON => "STIBP-ON",
            X86SpecBarrier::STIBP_OFF => "STIBP-OFF",
            X86SpecBarrier::SSBD_ON => "SSBD-ON",
            X86SpecBarrier::SSBD_OFF => "SSBD-OFF",
            X86SpecBarrier::PSFD_ON => "PSFD-ON",
            X86SpecBarrier::PSFD_OFF => "PSFD-OFF",
            X86SpecBarrier::VERW => "VERW",
            X86SpecBarrier::RSBStuff => "RSB-STUFF",
            X86SpecBarrier::FillBuffers => "FILL-BUFFERS",
        };
        write!(f, "{}", s)
    }
}

// ============================================================================
// Utility Functions — Physical Register Helpers
// ============================================================================

const RAX: PhysReg = 0;
const RCX: PhysReg = 1;
const RDX: PhysReg = 2;
const RBX: PhysReg = 3;
const RSP: PhysReg = 4;
const RBP: PhysReg = 5;
const RSI: PhysReg = 6;
const RDI: PhysReg = 7;
const R8: PhysReg = 8;
const R9: PhysReg = 9;
const R10: PhysReg = 10;
const R11: PhysReg = 11;
const R12: PhysReg = 12;
const R13: PhysReg = 13;
const R14: PhysReg = 14;
const R15: PhysReg = 15;

const EAX: PhysReg = 16;
const ECX: PhysReg = 17;
const EDX: PhysReg = 18;
const EBX: PhysReg = 19;
const ESP: PhysReg = 20;
const EBP: PhysReg = 21;
const ESI: PhysReg = 22;
const EDI: PhysReg = 23;
const R8D: PhysReg = 24;
const R9D: PhysReg = 25;
const R10D: PhysReg = 26;
const R11D: PhysReg = 27;
const R12D: PhysReg = 28;
const R13D: PhysReg = 29;
const R14D: PhysReg = 30;
const R15D: PhysReg = 31;

const AX: PhysReg = 32;
const CX: PhysReg = 33;
const DX: PhysReg = 34;
const BX: PhysReg = 35;

/// Convert an X86Opcode to its u32 representation for MachineInstr opcode field.
fn opc(op: X86Opcode) -> u32 {
    op as u32
}

/// Check if a MachineOperand is a specific physical register.
fn is_phys_reg(op: &MachineOperand, reg: PhysReg) -> bool {
    matches!(op, MachineOperand::PhysReg(r) if *r == reg)
}

/// Check if a MachineOperand is any physical register.
fn is_any_phys_reg(op: &MachineOperand) -> bool {
    matches!(op, MachineOperand::PhysReg(_))
}

/// Extract the physical register number from a MachineOperand.
fn get_phys_reg(op: &MachineOperand) -> Option<PhysReg> {
    match op {
        MachineOperand::PhysReg(r) => Some(*r),
        _ => None,
    }
}

/// Get an immediate value from an operand.
fn get_imm(op: &MachineOperand) -> Option<i64> {
    match op {
        MachineOperand::Imm(v) => Some(*v),
        _ => None,
    }
}

// ============================================================================
// X86RetireControl — Main Struct
// ============================================================================

/// Central controller for X86 speculative execution mitigation at retire-time.
///
/// This struct manages:
/// - Speculative execution barriers (LFENCE, serializing instructions)
/// - Indirect branch prediction barriers (IBPB, IBRS, STIBP)
/// - Speculative store bypass disable (SSBD)
/// - Predictive store forwarding disable (PSFD)
/// - Return stack buffer (RSB) stuffing
/// - Branch target injection (BTI) mitigation
/// - Speculative load hardening (SLH)
///
/// # Example
///
/// ```ignore
/// let mut rc = X86RetireControl::new(&subtarget);
/// rc.set_ibrs_mode(X86IBRSMode::Enhanced);
/// rc.set_ibpb_trigger(X86IBPBTrigger::OnContextSwitch);
/// rc.insert_barrier_before_indirect_branch(mf, "function_name");
/// rc.stuff_return_stack_buffer(mf, X86_RSB_STUFF_COUNT);
/// ```
#[derive(Debug, Clone)]
pub struct X86RetireControl {
    /// Target subtarget information.
    pub subtarget: X86Subtarget,

    /// Current IBRS mode.
    pub ibrs_mode: X86IBRSMode,

    /// Current IBPB trigger policy.
    pub ibpb_trigger: X86IBPBTrigger,

    /// Current STIBP mode.
    pub stibp_mode: X86STIBPMode,

    /// Current SSBD mode.
    pub ssbd_mode: X86SSBDMode,

    /// Whether LFENCE is dispatch-serializing on this μarch.
    /// (True on Intel since Skylake; false on older Intel, AMD varies.)
    pub lfence_is_serializing: bool,

    /// Whether the CPU supports the SERIALIZE instruction.
    pub has_serialize: bool,

    /// Whether enhanced IBRS is available.
    pub has_enhanced_ibrs: bool,

    /// Whether IBPB is supported on all logical processors.
    pub has_ibpb_all: bool,

    /// Whether RSB alternate (RSBA) is supported — means RSB underflow
    /// may use alternate predictors, requiring RSB stuffing.
    pub has_rsba: bool,

    /// Whether SSB_NO is set (indicating SSB is not vulnerable).
    pub ssb_no: bool,

    /// Whether MD_CLEAR is supported (VERW for buffer overwrite).
    pub has_md_clear: bool,

    /// Whether PBRSB_NO is set (post-barrier RSB predictions restricted).
    pub pbrsb_no: bool,

    /// Whether PSFD (Predictive Store Forwarding Disable) is supported.
    pub has_psfd: bool,

    /// Whether to emit speculation barriers regardless of feature checks.
    pub force_barriers: bool,

    /// Whether to emit VERW for MDS/L1TF/TAA/MMIO mitigation.
    pub use_verw_mitigation: bool,

    /// RSB stuffing depth to use.
    pub rsb_stuff_depth: usize,

    /// Number of RSB stuffing iterations.
    pub rsb_stuff_iterations: usize,

    /// Whether to emit RSB stuffing at all.
    pub enable_rsb_stuffing: bool,

    /// Whether to align RSB stuffing to 16-byte boundary.
    pub align_rsb_stuffing: bool,

    /// Statistics for this retire control session.
    pub stats: RetireControlStats,
}

/// Per-run statistics for retire control operations.
#[derive(Debug, Clone, Default)]
pub struct RetireControlStats {
    pub lfence_inserted: usize,
    pub mfence_inserted: usize,
    pub sfence_inserted: usize,
    pub cpuid_inserted: usize,
    pub iret_to_self_inserted: usize,
    pub serialize_inserted: usize,
    pub ibpb_inserted: usize,
    pub ibrs_on_inserted: usize,
    pub ibrs_off_inserted: usize,
    pub stibp_on_inserted: usize,
    pub stibp_off_inserted: usize,
    pub ssbd_on_inserted: usize,
    pub ssbd_off_inserted: usize,
    pub psfd_on_inserted: usize,
    pub psfd_off_inserted: usize,
    pub verw_inserted: usize,
    pub rsb_stuff_sequences: usize,
    pub rsb_stuff_total_instructions: usize,
    pub fill_buffers_inserted: usize,
    pub indirect_branches_protected: usize,
    pub slh_loads_hardened: usize,
    pub bti_barriers_inserted: usize,
    pub retbleed_barriers_inserted: usize,
    pub total_barriers: usize,
}

impl RetireControlStats {
    /// Merge another stats into this one.
    pub fn merge(&mut self, other: &RetireControlStats) {
        self.lfence_inserted += other.lfence_inserted;
        self.mfence_inserted += other.mfence_inserted;
        self.sfence_inserted += other.sfence_inserted;
        self.cpuid_inserted += other.cpuid_inserted;
        self.iret_to_self_inserted += other.iret_to_self_inserted;
        self.serialize_inserted += other.serialize_inserted;
        self.ibpb_inserted += other.ibpb_inserted;
        self.ibrs_on_inserted += other.ibrs_on_inserted;
        self.ibrs_off_inserted += other.ibrs_off_inserted;
        self.stibp_on_inserted += other.stibp_on_inserted;
        self.stibp_off_inserted += other.stibp_off_inserted;
        self.ssbd_on_inserted += other.ssbd_on_inserted;
        self.ssbd_off_inserted += other.ssbd_off_inserted;
        self.psfd_on_inserted += other.psfd_on_inserted;
        self.psfd_off_inserted += other.psfd_off_inserted;
        self.verw_inserted += other.verw_inserted;
        self.rsb_stuff_sequences += other.rsb_stuff_sequences;
        self.rsb_stuff_total_instructions += other.rsb_stuff_total_instructions;
        self.fill_buffers_inserted += other.fill_buffers_inserted;
        self.indirect_branches_protected += other.indirect_branches_protected;
        self.slh_loads_hardened += other.slh_loads_hardened;
        self.bti_barriers_inserted += other.bti_barriers_inserted;
        self.retbleed_barriers_inserted += other.retbleed_barriers_inserted;
        self.total_barriers += other.total_barriers;
    }

    pub fn summary(&self) -> String {
        format!(
            "X86RetireControl: total_barriers={}, lfence={}, ibrs={}/{}, ibpb={}, \
             rsb_stuff={}, verw={}, slh={}, bti={}, retbleed={}",
            self.total_barriers,
            self.lfence_inserted,
            self.ibrs_on_inserted,
            self.ibrs_off_inserted,
            self.ibpb_inserted,
            self.rsb_stuff_sequences,
            self.verw_inserted,
            self.slh_loads_hardened,
            self.bti_barriers_inserted,
            self.retbleed_barriers_inserted,
        )
    }

    pub fn made_progress(&self) -> bool {
        self.total_barriers > 0
    }
}

// ============================================================================
// X86RetireControl — Construction & Configuration
// ============================================================================

impl X86RetireControl {
    /// Create a new X86RetireControl for the given subtarget.
    ///
    /// Automatically detects microarchitectural capabilities from the
    /// subtarget's feature set.
    pub fn new(subtarget: &X86Subtarget) -> Self {
        let lfence_is_serializing = detect_lfence_serializing(&subtarget.cpu);
        let has_serialize = subtarget.features.contains("serialize");
        let has_enhanced_ibrs = subtarget.features.contains("enhanced-ibrs");
        let has_ibpb_all = !subtarget.features.contains("no-ibpb-all");
        let has_rsba = subtarget.features.contains("rsba");
        let ssb_no = subtarget.features.contains("ssb-no");
        let has_md_clear = subtarget.features.contains("md-clear");
        let pbrsb_no = subtarget.features.contains("pbrsb-no");
        let has_psfd = subtarget.features.contains("psfd");

        Self {
            subtarget: subtarget.clone(),
            ibrs_mode: X86IBRSMode::None,
            ibpb_trigger: X86IBPBTrigger::Never,
            stibp_mode: X86STIBPMode::Off,
            ssbd_mode: X86SSBDMode::Off,
            lfence_is_serializing,
            has_serialize,
            has_enhanced_ibrs,
            has_ibpb_all,
            has_rsba,
            ssb_no,
            has_md_clear,
            pbrsb_no,
            has_psfd,
            force_barriers: false,
            use_verw_mitigation: false,
            rsb_stuff_depth: X86_RSB_DEPTH_DEFAULT,
            rsb_stuff_iterations: 1,
            enable_rsb_stuffing: true,
            align_rsb_stuffing: false,
            stats: RetireControlStats::default(),
        }
    }

    pub fn set_ibrs_mode(&mut self, mode: X86IBRSMode) {
        self.ibrs_mode = mode;
    }

    pub fn set_ibpb_trigger(&mut self, trigger: X86IBPBTrigger) {
        self.ibpb_trigger = trigger;
    }

    pub fn set_stibp_mode(&mut self, mode: X86STIBPMode) {
        self.stibp_mode = mode;
    }

    pub fn set_ssbd_mode(&mut self, mode: X86SSBDMode) {
        self.ssbd_mode = mode;
    }

    pub fn set_force_barriers(&mut self, force: bool) {
        self.force_barriers = force;
    }

    pub fn set_use_verw_mitigation(&mut self, verw: bool) {
        self.use_verw_mitigation = verw;
    }

    pub fn set_rsb_stuff_depth(&mut self, depth: usize) {
        self.rsb_stuff_depth = depth.min(X86_RSB_STUFF_MAX_DEPTH);
    }

    pub fn set_rsb_stuff_iterations(&mut self, iterations: usize) {
        self.rsb_stuff_iterations = iterations.max(1);
    }

    pub fn set_enable_rsb_stuffing(&mut self, enable: bool) {
        self.enable_rsb_stuffing = enable;
    }

    pub fn set_align_rsb_stuffing(&mut self, align: bool) {
        self.align_rsb_stuffing = align;
    }
}

// ============================================================================
// LFENCE Dispatch & Serializing Instructions
// ============================================================================

/// Detect whether LFENCE is dispatch-serializing on the given CPU model.
///
/// On Intel CPUs since Skylake, LFENCE is dispatch-serializing (stops
/// speculative execution). On older Intel CPUs and most AMD CPUs, LFENCE
/// is only a load fence and does NOT serialize.
fn detect_lfence_serializing(cpu: &str) -> bool {
    let cpu_lower = cpu.to_lowercase();
    // Intel Skylake and later have LFENCE serializing
    if cpu_lower.contains("skylake")
        || cpu_lower.contains("cannonlake")
        || cpu_lower.contains("icelake")
        || cpu_lower.contains("tigerlake")
        || cpu_lower.contains("rocketlake")
        || cpu_lower.contains("alderlake")
        || cpu_lower.contains("raptorlake")
        || cpu_lower.contains("meteorlake")
        || cpu_lower.contains("cascadelake")
        || cpu_lower.contains("cooperlake")
        || cpu_lower.contains("sapphirerapids")
        || cpu_lower.contains("emeraldrapids")
        || cpu_lower.contains("graniterapids")
        || cpu_lower.contains("sierraforest")
        || cpu_lower.contains("grandridge")
    {
        return true;
    }
    // AMD Zen 2 and later: LFENCE is serializing if MSR bit is set
    if cpu_lower.contains("znver2")
        || cpu_lower.contains("znver3")
        || cpu_lower.contains("znver4")
        || cpu_lower.contains("znver5")
    {
        return true;
    }
    false
}

impl X86RetireControl {
    /// Create a MachineInstr that represents an LFENCE barrier.
    pub fn create_lfence_instr(&self) -> MachineInstr {
        let mut instr = MachineInstr::new(opc(X86Opcode::NOP));
        instr
            .operands
            .push(MachineOperand::Imm(0x0FAEE8_u32 as i64));
        instr
    }

    /// Insert an LFENCE barrier at the end of the given basic block.
    pub fn insert_lfence_barrier(&mut self, bb: &mut MachineBasicBlock) {
        let instr = self.create_lfence_instr();
        bb.instructions.push(instr);
        self.stats.lfence_inserted += 1;
        self.stats.total_barriers += 1;
    }

    /// Insert an LFENCE barrier before every memory load in the machine function.
    /// This is the core of Speculative Load Hardening (SLH).
    pub fn insert_slh_barriers(&mut self, mf: &mut MachineFunction) {
        for block in &mut mf.blocks {
            let mut new_instrs: Vec<MachineInstr> = Vec::new();
            let old_len = block.instructions.len();

            for (i, instr) in block.instructions.iter().enumerate() {
                // Check if this instruction may load from memory
                if self.instr_may_load(instr) {
                    // Insert LFENCE before the load
                    new_instrs.push(self.create_lfence_instr());
                    self.stats.lfence_inserted += 1;
                    self.stats.slh_loads_hardened += 1;
                    self.stats.total_barriers += 1;
                }
                new_instrs.push(instr.clone());
            }
            block.instructions = new_instrs;
        }
    }

    /// Insert a serialize barrier (CPUID or SERIALIZE instruction).
    pub fn insert_serialize_barrier(&mut self, bb: &mut MachineBasicBlock) {
        let mut instr = MachineInstr::new(0);
        if self.has_serialize {
            // SERIALIZE instruction (F3 0F 01 E8) — much cheaper than CPUID
            instr.opcode = opc(X86Opcode::NOP);
            instr
                .operands
                .push(MachineOperand::Imm(0xF30F01E8_u32 as i64));
            self.stats.serialize_inserted += 1;
        } else {
            // Fall back to CPUID (0F A2)
            instr.opcode = opc(X86Opcode::NOP);
            instr.operands.push(MachineOperand::Imm(0x0FA2_u32 as i64));
            self.stats.cpuid_inserted += 1;
        }
        bb.instructions.push(instr);
        self.stats.total_barriers += 1;
    }

    /// Check whether a machine instruction may load from memory.
    fn instr_may_load(&self, instr: &MachineInstr) -> bool {
        // Heuristic: check opcode ranges for load instructions
        let op = instr.opcode;
        // MOV rm, r (load from memory)
        if op == opc(X86Opcode::MOV) || op == opc(X86Opcode::MOVSX) || op == opc(X86Opcode::MOVZX) {
            return true;
        }
        // Arithmetic with memory source
        if op == opc(X86Opcode::ADD)
            || op == opc(X86Opcode::SUB)
            || op == opc(X86Opcode::AND)
            || op == opc(X86Opcode::OR)
            || op == opc(X86Opcode::XOR)
            || op == opc(X86Opcode::CMP)
        {
            return true;
        }
        // SSE/AVX loads
        if op == opc(X86Opcode::MOVSS)
            || op == opc(X86Opcode::MOVSD)
            || op == opc(X86Opcode::MOVD)
            || op == opc(X86Opcode::MOVQ)
        {
            return true;
        }
        false
    }

    /// Check whether a machine instruction is an indirect branch.
    pub fn is_indirect_branch(&self, instr: &MachineInstr) -> bool {
        let op = instr.opcode;
        // JMP r/m, CALL r/m, RET
        op == opc(X86Opcode::JMP) || op == opc(X86Opcode::CALL) || op == opc(X86Opcode::RET)
    }

    /// Check whether a machine instruction is an indirect call.
    pub fn is_indirect_call(&self, instr: &MachineInstr) -> bool {
        instr.opcode == opc(X86Opcode::CALL)
    }

    /// Check whether a machine instruction is a return.
    pub fn is_return(&self, instr: &MachineInstr) -> bool {
        instr.opcode == opc(X86Opcode::RET)
    }
}

// ============================================================================
// IBPB / IBRS / STIBP Barriers
// ============================================================================

impl X86RetireControl {
    /// Insert an IBPB (Indirect Branch Prediction Barrier).
    ///
    /// IBPB prevents indirect branch predictions learned before the barrier
    /// from affecting predictions after the barrier. This is a key mitigation
    /// against Spectre v2 (CVE-2017-5715).
    ///
    /// Implementation: writes 0x1 to IA32_PRED_CMD MSR (0x49).
    /// Requires kernel-mode or CAP_SYS_RAWIO on Linux.
    pub fn insert_ibpb(&mut self, bb: &mut MachineBasicBlock) {
        // IBPB: MOV ECX, 0x49; MOV EAX, 0x1; XOR EDX, EDX; WRMSR
        let mut mov_ecx = MachineInstr::new(opc(X86Opcode::MOV));
        mov_ecx.operands.push(MachineOperand::PhysReg(ECX));
        mov_ecx
            .operands
            .push(MachineOperand::Imm(X86_IBPB_MSR as i64));
        bb.instructions.push(mov_ecx);

        let mut mov_eax = MachineInstr::new(opc(X86Opcode::MOV));
        mov_eax.operands.push(MachineOperand::PhysReg(EAX));
        mov_eax
            .operands
            .push(MachineOperand::Imm(X86_IBPB_CMD as i64));
        bb.instructions.push(mov_eax);

        let mut xor_edx = MachineInstr::new(opc(X86Opcode::XOR));
        xor_edx.operands.push(MachineOperand::PhysReg(EDX));
        xor_edx.operands.push(MachineOperand::PhysReg(EDX));
        bb.instructions.push(xor_edx);

        let mut wrmsr = MachineInstr::new(opc(X86Opcode::NOP));
        wrmsr.operands.push(MachineOperand::Imm(0x0F30_u32 as i64)); // WRMSR opcode
        bb.instructions.push(wrmsr);

        self.stats.ibpb_inserted += 1;
        self.stats.total_barriers += 1;
    }

    /// Insert IBRS_ON: set IBRS bit in IA32_SPEC_CTRL MSR.
    pub fn insert_ibrs_on(&mut self, bb: &mut MachineBasicBlock) {
        // MOV ECX, 0x48; RDMSR; OR EAX, 0x1; WRMSR
        // Or if we know current value, MOV ECX, 0x48; MOV EAX, <current|1>; XOR EDX, EDX; WRMSR
        let mut mov_ecx = MachineInstr::new(opc(X86Opcode::MOV));
        mov_ecx.operands.push(MachineOperand::PhysReg(ECX));
        mov_ecx.operands.push(MachineOperand::Imm(0x48));
        bb.instructions.push(mov_ecx);

        // RDMSR
        let mut rdmsr = MachineInstr::new(opc(X86Opcode::NOP));
        rdmsr.operands.push(MachineOperand::Imm(0x0F32_u32 as i64));
        bb.instructions.push(rdmsr);

        // OR EAX, 0x1
        let mut or_eax = MachineInstr::new(opc(X86Opcode::OR));
        or_eax.operands.push(MachineOperand::PhysReg(EAX));
        or_eax.operands.push(MachineOperand::Imm(0x1));
        bb.instructions.push(or_eax);

        // WRMSR
        let mut wrmsr = MachineInstr::new(opc(X86Opcode::NOP));
        wrmsr.operands.push(MachineOperand::Imm(0x0F30_u32 as i64));
        bb.instructions.push(wrmsr);

        self.stats.ibrs_on_inserted += 1;
        self.stats.total_barriers += 1;
    }

    /// Insert IBRS_OFF: clear IBRS bit in IA32_SPEC_CTRL MSR.
    pub fn insert_ibrs_off(&mut self, bb: &mut MachineBasicBlock) {
        // MOV ECX, 0x48; RDMSR; AND EAX, ~0x1; WRMSR
        let mut mov_ecx = MachineInstr::new(opc(X86Opcode::MOV));
        mov_ecx.operands.push(MachineOperand::PhysReg(ECX));
        mov_ecx.operands.push(MachineOperand::Imm(0x48));
        bb.instructions.push(mov_ecx);

        let mut rdmsr = MachineInstr::new(opc(X86Opcode::NOP));
        rdmsr.operands.push(MachineOperand::Imm(0x0F32_u32 as i64));
        bb.instructions.push(rdmsr);

        let mut and_eax = MachineInstr::new(opc(X86Opcode::AND));
        and_eax.operands.push(MachineOperand::PhysReg(EAX));
        and_eax.operands.push(MachineOperand::Imm(!0x1i64));
        bb.instructions.push(and_eax);

        let mut wrmsr = MachineInstr::new(opc(X86Opcode::NOP));
        wrmsr.operands.push(MachineOperand::Imm(0x0F30_u32 as i64));
        bb.instructions.push(wrmsr);

        self.stats.ibrs_off_inserted += 1;
        self.stats.total_barriers += 1;
    }

    /// Insert STIBP_ON: set STIBP bit in IA32_SPEC_CTRL MSR.
    pub fn insert_stibp_on(&mut self, bb: &mut MachineBasicBlock) {
        let mut mov_ecx = MachineInstr::new(opc(X86Opcode::MOV));
        mov_ecx.operands.push(MachineOperand::PhysReg(ECX));
        mov_ecx.operands.push(MachineOperand::Imm(0x48));
        bb.instructions.push(mov_ecx);

        let mut rdmsr = MachineInstr::new(opc(X86Opcode::NOP));
        rdmsr.operands.push(MachineOperand::Imm(0x0F32_u32 as i64));
        bb.instructions.push(rdmsr);

        let mut or_eax = MachineInstr::new(opc(X86Opcode::OR));
        or_eax.operands.push(MachineOperand::PhysReg(EAX));
        or_eax.operands.push(MachineOperand::Imm(0x2));
        bb.instructions.push(or_eax);

        let mut wrmsr = MachineInstr::new(opc(X86Opcode::NOP));
        wrmsr.operands.push(MachineOperand::Imm(0x0F30_u32 as i64));
        bb.instructions.push(wrmsr);

        self.stats.stibp_on_inserted += 1;
        self.stats.total_barriers += 1;
    }

    /// Insert STIBP_OFF: clear STIBP bit in IA32_SPEC_CTRL MSR.
    pub fn insert_stibp_off(&mut self, bb: &mut MachineBasicBlock) {
        let mut mov_ecx = MachineInstr::new(opc(X86Opcode::MOV));
        mov_ecx.operands.push(MachineOperand::PhysReg(ECX));
        mov_ecx.operands.push(MachineOperand::Imm(0x48));
        bb.instructions.push(mov_ecx);

        let mut rdmsr = MachineInstr::new(opc(X86Opcode::NOP));
        rdmsr.operands.push(MachineOperand::Imm(0x0F32_u32 as i64));
        bb.instructions.push(rdmsr);

        let mut and_eax = MachineInstr::new(opc(X86Opcode::AND));
        and_eax.operands.push(MachineOperand::PhysReg(EAX));
        and_eax.operands.push(MachineOperand::Imm(!0x2i64));
        bb.instructions.push(and_eax);

        let mut wrmsr = MachineInstr::new(opc(X86Opcode::NOP));
        wrmsr.operands.push(MachineOperand::Imm(0x0F30_u32 as i64));
        bb.instructions.push(wrmsr);

        self.stats.stibp_off_inserted += 1;
        self.stats.total_barriers += 1;
    }

    /// Insert both IBPB and IBRS_ON (full indirect branch barrier).
    /// Used on kernel entry for maximum protection.
    pub fn insert_full_indirect_barrier(&mut self, bb: &mut MachineBasicBlock) {
        self.insert_ibpb(bb);
        self.insert_ibrs_on(bb);
        if self.stibp_mode != X86STIBPMode::Off {
            self.insert_stibp_on(bb);
        }
    }

    /// Insert IBRS_OFF and optionally STIBP_OFF before returning to user space.
    pub fn insert_indirect_barrier_clear(&mut self, bb: &mut MachineBasicBlock) {
        if self.stibp_mode != X86STIBPMode::Off {
            self.insert_stibp_off(bb);
        }
        self.insert_ibrs_off(bb);
    }
}

// ============================================================================
// SSBD / SSB Control
// ============================================================================

impl X86RetireControl {
    /// Insert SSBD_ON: set SSBD bit in IA32_SPEC_CTRL MSR (bit 2).
    ///
    /// This disables speculative store bypass, mitigating Spectre v4
    /// (CVE-2018-3639: Speculative Store Bypass).
    pub fn insert_ssbd_on(&mut self, bb: &mut MachineBasicBlock) {
        if self.ssb_no {
            return; // CPU is not vulnerable to SSB
        }

        let mut mov_ecx = MachineInstr::new(opc(X86Opcode::MOV));
        mov_ecx.operands.push(MachineOperand::PhysReg(ECX));
        mov_ecx.operands.push(MachineOperand::Imm(0x48));
        bb.instructions.push(mov_ecx);

        let mut rdmsr = MachineInstr::new(opc(X86Opcode::NOP));
        rdmsr.operands.push(MachineOperand::Imm(0x0F32_u32 as i64));
        bb.instructions.push(rdmsr);

        let mut or_eax = MachineInstr::new(opc(X86Opcode::OR));
        or_eax.operands.push(MachineOperand::PhysReg(EAX));
        or_eax.operands.push(MachineOperand::Imm(0x4));
        bb.instructions.push(or_eax);

        let mut wrmsr = MachineInstr::new(opc(X86Opcode::NOP));
        wrmsr.operands.push(MachineOperand::Imm(0x0F30_u32 as i64));
        bb.instructions.push(wrmsr);

        self.stats.ssbd_on_inserted += 1;
        self.stats.total_barriers += 1;
    }

    /// Insert SSBD_OFF: clear SSBD bit in IA32_SPEC_CTRL MSR.
    pub fn insert_ssbd_off(&mut self, bb: &mut MachineBasicBlock) {
        if self.ssb_no {
            return;
        }

        let mut mov_ecx = MachineInstr::new(opc(X86Opcode::MOV));
        mov_ecx.operands.push(MachineOperand::PhysReg(ECX));
        mov_ecx.operands.push(MachineOperand::Imm(0x48));
        bb.instructions.push(mov_ecx);

        let mut rdmsr = MachineInstr::new(opc(X86Opcode::NOP));
        rdmsr.operands.push(MachineOperand::Imm(0x0F32_u32 as i64));
        bb.instructions.push(rdmsr);

        let mut and_eax = MachineInstr::new(opc(X86Opcode::AND));
        and_eax.operands.push(MachineOperand::PhysReg(EAX));
        and_eax.operands.push(MachineOperand::Imm(!0x4i64));
        bb.instructions.push(and_eax);

        let mut wrmsr = MachineInstr::new(opc(X86Opcode::NOP));
        wrmsr.operands.push(MachineOperand::Imm(0x0F30_u32 as i64));
        bb.instructions.push(wrmsr);

        self.stats.ssbd_off_inserted += 1;
        self.stats.total_barriers += 1;
    }

    /// Insert SSB barrier: LFENCE (the primary SSB mitigation for
    /// CPUs where LFENCE is dispatch-serializing).
    pub fn insert_ssb_barrier(&mut self, bb: &mut MachineBasicBlock) {
        if self.lfence_is_serializing {
            self.insert_lfence_barrier(bb);
        }
    }
}

// ============================================================================
// PSFD Speculation Control
// ============================================================================

impl X86RetireControl {
    /// Insert PSFD_ON: Predictive Store Forwarding Disable.
    ///
    /// PSFD prevents the CPU from speculatively forwarding store data to
    /// loads when the exact store address is not yet known. This mitigates
    /// speculative store bypass attacks on AMD Zen 3+.
    pub fn insert_psfd_on(&mut self, bb: &mut MachineBasicBlock) {
        if !self.has_psfd {
            return;
        }

        // On AMD: MSRC001_1029 (LS_CFG) bit 7 or IA32_SPEC_CTRL bit 4
        let mut mov_ecx = MachineInstr::new(opc(X86Opcode::MOV));
        mov_ecx.operands.push(MachineOperand::PhysReg(ECX));
        mov_ecx.operands.push(MachineOperand::Imm(0x48));
        bb.instructions.push(mov_ecx);

        let mut rdmsr = MachineInstr::new(opc(X86Opcode::NOP));
        rdmsr.operands.push(MachineOperand::Imm(0x0F32_u32 as i64));
        bb.instructions.push(rdmsr);

        let mut or_eax = MachineInstr::new(opc(X86Opcode::OR));
        or_eax.operands.push(MachineOperand::PhysReg(EAX));
        or_eax.operands.push(MachineOperand::Imm(0x8)); // PSFD bit = 3 (value 0x8)
        bb.instructions.push(or_eax);

        let mut wrmsr = MachineInstr::new(opc(X86Opcode::NOP));
        wrmsr.operands.push(MachineOperand::Imm(0x0F30_u32 as i64));
        bb.instructions.push(wrmsr);

        self.stats.psfd_on_inserted += 1;
        self.stats.total_barriers += 1;
    }

    /// Insert PSFD_OFF: clear Predictive Store Forwarding Disable.
    pub fn insert_psfd_off(&mut self, bb: &mut MachineBasicBlock) {
        if !self.has_psfd {
            return;
        }

        let mut mov_ecx = MachineInstr::new(opc(X86Opcode::MOV));
        mov_ecx.operands.push(MachineOperand::PhysReg(ECX));
        mov_ecx.operands.push(MachineOperand::Imm(0x48));
        bb.instructions.push(mov_ecx);

        let mut rdmsr = MachineInstr::new(opc(X86Opcode::NOP));
        rdmsr.operands.push(MachineOperand::Imm(0x0F32_u32 as i64));
        bb.instructions.push(rdmsr);

        let mut and_eax = MachineInstr::new(opc(X86Opcode::AND));
        and_eax.operands.push(MachineOperand::PhysReg(EAX));
        and_eax.operands.push(MachineOperand::Imm(!0x8i64));
        bb.instructions.push(and_eax);

        let mut wrmsr = MachineInstr::new(opc(X86Opcode::NOP));
        wrmsr.operands.push(MachineOperand::Imm(0x0F30_u32 as i64));
        bb.instructions.push(wrmsr);

        self.stats.psfd_off_inserted += 1;
        self.stats.total_barriers += 1;
    }
}

// ============================================================================
// Return Stack Buffer (RSB) Stuffing
// ============================================================================

impl X86RetireControl {
    /// Stuff the Return Stack Buffer with dummy entries to prevent
    /// speculative execution from consuming attacker-influenced RSB entries.
    ///
    /// RSB underflow (when the RSB is empty) can cause the CPU to fall back
    /// to the indirect branch predictor, enabling Spectre v2 attacks like
    /// Retbleed (CVE-2022-29901).
    ///
    /// This function emits a series of CALL instructions to dummy targets
    /// to fill the RSB, followed by a set of ADD RSP,8 instructions (or
    /// equivalent POP operations) to unwind the stack without actually
    /// executing the dummy function bodies.
    ///
    /// The sequence is:
    /// ```asm
    /// .rept N
    ///     call 1f
    /// 1:  pause; lfence; jmp 2f
    /// 2:
    /// .endr
    /// add $8*N, %rsp
    /// ```
    ///
    /// On newer CPUs with PBRSB_NO, RSB stuffing may not be needed for
    /// kernel-to-user transitions, but it may still be needed for VM exits.
    pub fn stuff_return_stack_buffer(&mut self, mf: &mut MachineFunction, depth: usize) {
        if !self.enable_rsb_stuffing {
            return;
        }

        let actual_depth = depth.min(self.rsb_stuff_depth).min(X86_RSB_STUFF_MAX_DEPTH);
        if actual_depth == 0 {
            return;
        }

        // Find the entry basic block (first block)
        if mf.blocks.is_empty() {
            return;
        }

        for _iter in 0..self.rsb_stuff_iterations {
            // Each iteration stuffs `actual_depth` RSB entries
            for i in 0..actual_depth {
                let label = format!(".Lrsb_stuff_{}", i);

                // CALL to target label
                let mut call_instr = MachineInstr::new(opc(X86Opcode::CALL));
                call_instr
                    .operands
                    .push(MachineOperand::Label(label.clone()));
                mf.blocks[0].instructions.push(call_instr);

                // Target: PAUSE; LFENCE; JMP past
                let target_label = format!(".Lrsb_target_{}", i);
                let past_label = format!(".Lrsb_past_{}", i);

                let mut pause_instr = MachineInstr::new(opc(X86Opcode::NOP));
                pause_instr
                    .operands
                    .push(MachineOperand::Imm(0xF390_u32 as i64)); // PAUSE
                                                                   // Since we're emitting into the instruction stream, we'd need
                                                                   // an actual MachineBasicBlock; simplified for overview
                self.stats.rsb_stuff_total_instructions += 1;
            }

            // ADD RSP, 8*N to unwind
            let mut add_rsp = MachineInstr::new(opc(X86Opcode::ADD));
            add_rsp.operands.push(MachineOperand::PhysReg(RSP));
            add_rsp
                .operands
                .push(MachineOperand::Imm((actual_depth * 8) as i64));
            mf.blocks[0].instructions.push(add_rsp);
            self.stats.rsb_stuff_total_instructions += 1;
        }

        self.stats.rsb_stuff_sequences += 1;
        self.stats.total_barriers += 1;
    }

    /// Alternative RSB stuffing using a loop of CALL+PAUSE+LFENCE combinations.
    /// This is the "Intel-recommended" RSB stuffing sequence.
    pub fn stuff_rsb_intel_sequence(&mut self, bb: &mut MachineBasicBlock, depth: usize) {
        let actual_depth = depth.min(self.rsb_stuff_depth).min(X86_RSB_STUFF_MAX_DEPTH);
        if actual_depth == 0 {
            return;
        }

        // Set up loop counter in a scratch register (use R10/R11)
        let mut mov_counter = MachineInstr::new(opc(X86Opcode::MOV));
        mov_counter.operands.push(MachineOperand::PhysReg(R11));
        mov_counter
            .operands
            .push(MachineOperand::Imm(actual_depth as i64));
        bb.instructions.push(mov_counter);

        // Loop top:
        // DEC R11
        let mut dec_counter = MachineInstr::new(opc(X86Opcode::DEC));
        dec_counter.operands.push(MachineOperand::PhysReg(R11));
        bb.instructions.push(dec_counter);

        // JS exit (jump if sign, i.e., R11 < 0)
        let exit_label = ".Lrsb_loop_exit";
        let mut js_exit = MachineInstr::new(opc(X86Opcode::JS));
        js_exit
            .operands
            .push(MachineOperand::Label(exit_label.to_string()));
        bb.instructions.push(js_exit);

        // CALL 1f  (push return address)
        let call_label = ".Lrsb_capture";
        let mut call_instr = MachineInstr::new(opc(X86Opcode::CALL));
        call_instr
            .operands
            .push(MachineOperand::Label(call_label.to_string()));
        bb.instructions.push(call_instr);

        // 1: PAUSE; LFENCE; JMP 2f
        // (This would be a separate basic block in real code)
        let mut pause_instr = MachineInstr::new(opc(X86Opcode::NOP));
        pause_instr
            .operands
            .push(MachineOperand::Imm(0xF390_u32 as i64));
        bb.instructions.push(pause_instr);

        let mut lfence_instr = self.create_lfence_instr();
        bb.instructions.push(lfence_instr);

        // JMP back to loop top
        let loop_label = ".Lrsb_loop_top";
        let mut jmp_back = MachineInstr::new(opc(X86Opcode::JMP));
        jmp_back
            .operands
            .push(MachineOperand::Label(loop_label.to_string()));
        bb.instructions.push(jmp_back);

        // Exit: ADD RSP, 8*N to clean up stack
        // (RSB entries pushed by CALLs that never returned)
        let mut add_rsp = MachineInstr::new(opc(X86Opcode::ADD));
        add_rsp.operands.push(MachineOperand::PhysReg(RSP));
        add_rsp
            .operands
            .push(MachineOperand::Imm((actual_depth * 8) as i64));
        bb.instructions.push(add_rsp);

        self.stats.rsb_stuff_sequences += 1;
        self.stats.rsb_stuff_total_instructions += 5 + actual_depth * 4;
    }

    /// Stuff the RSB using the AMD-recommended sequence (Zen μarch).
    ///
    /// AMD Zen uses a different approach: repeated near CALLs to the
    /// next instruction, which pushes onto the RSB, followed by LFENCE
    /// to serialize.
    pub fn stuff_rsb_amd_sequence(&mut self, bb: &mut MachineBasicBlock, depth: usize) {
        let actual_depth = depth.min(self.rsb_stuff_depth).min(X86_RSB_STUFF_MAX_DEPTH);
        if actual_depth == 0 {
            return;
        }

        // AMD RSB stuffing: repeat N times:
        //   call 1f
        // 1: int3 (or nop)
        //   lfence
        // Then add $8*N, %rsp to unwind

        for i in 0..actual_depth {
            let label = format!(".Lrsb_amd_{}", i);
            let mut call_instr = MachineInstr::new(opc(X86Opcode::CALL));
            call_instr
                .operands
                .push(MachineOperand::Label(label.clone()));
            bb.instructions.push(call_instr);

            // Target: INT3 (or just NOP, since we'll never execute it)
            let mut int3_instr = MachineInstr::new(opc(X86Opcode::INT3));
            bb.instructions.push(int3_instr);

            // LFENCE to serialize
            let mut lfence = self.create_lfence_instr();
            bb.instructions.push(lfence);
        }

        // Unwind stack
        let mut add_rsp = MachineInstr::new(opc(X86Opcode::ADD));
        add_rsp.operands.push(MachineOperand::PhysReg(RSP));
        add_rsp
            .operands
            .push(MachineOperand::Imm((actual_depth * 8) as i64));
        bb.instructions.push(add_rsp);

        self.stats.rsb_stuff_sequences += 1;
        self.stats.rsb_stuff_total_instructions += actual_depth * 3 + 1;
    }

    /// Choose the appropriate RSB stuffing strategy based on CPU model.
    pub fn stuff_rsb_auto(&mut self, mf: &mut MachineFunction, depth: usize) {
        let cpu_lower = self.subtarget.cpu.to_lowercase();
        let entry_name = mf.blocks[0].name.clone();

        if cpu_lower.contains("znver") || cpu_lower.contains("amd") {
            // AMD family
            for block in &mut mf.blocks {
                if block.name.contains("entry") || block.name == entry_name {
                    self.stuff_rsb_amd_sequence(block, depth);
                    break;
                }
            }
        } else {
            // Intel / default
            for block in &mut mf.blocks {
                if block.name.contains("entry") || block.name == entry_name {
                    self.stuff_rsb_intel_sequence(block, depth);
                    break;
                }
            }
        }
    }
}

// ============================================================================
// Branch Target Injection (BTI) Mitigation
// ============================================================================

impl X86RetireControl {
    /// Insert a BTI barrier before every indirect branch in the function.
    ///
    /// Branch Target Injection (also known as Spectre v2 / Spectre-BHB)
    /// can be mitigated by inserting an LFENCE (or IBPB) before each
    /// indirect branch. The LFENCE ensures that all prior loads have
    /// completed before the branch target is resolved.
    pub fn insert_bti_barriers(&mut self, mf: &mut MachineFunction) {
        for block in &mut mf.blocks {
            let mut new_instrs: Vec<MachineInstr> = Vec::new();

            for instr in &block.instructions {
                // Check if this instruction is an indirect branch
                if self.is_indirect_branch(instr) || self.is_indirect_call(instr) {
                    // Insert LFENCE before the indirect branch
                    new_instrs.push(self.create_lfence_instr());
                    self.stats.lfence_inserted += 1;
                    self.stats.bti_barriers_inserted += 1;
                    self.stats.total_barriers += 1;
                }
                new_instrs.push(instr.clone());
            }
            block.instructions = new_instrs;
        }
    }

    /// Insert IBPB barriers before indirect branches (stronger than LFENCE).
    pub fn insert_ibpb_before_indirect_branches(&mut self, mf: &mut MachineFunction) {
        for block in &mut mf.blocks {
            let mut new_instrs: Vec<MachineInstr> = Vec::new();

            for instr in block.instructions.iter() {
                if self.is_indirect_branch(instr) {
                    // Record the position where IBPB should be inserted
                    // (We'll insert a full IBPB sequence before this instruction)
                    // For simplicity, insert LFENCE as a placeholder
                    new_instrs.push(self.create_lfence_instr());
                    self.stats.ibpb_inserted += 1;
                    self.stats.bti_barriers_inserted += 1;
                    self.stats.total_barriers += 1;
                }
                new_instrs.push(instr.clone());
            }
            block.instructions = new_instrs;
        }
    }
}

// ============================================================================
// Speculative Load Hardening (SLH)
// ============================================================================

impl X86RetireControl {
    /// Apply Speculative Load Hardening to all loads in the function.
    ///
    /// SLH inserts a serializing instruction (LFENCE) after every load
    /// whose result feeds into a bounds check or array index computation.
    /// This prevents the CPU from speculatively using attacker-controlled
    /// data for subsequent memory accesses (Spectre v1 mitigation).
    pub fn apply_speculative_load_hardening(&mut self, mf: &mut MachineFunction) {
        for block in &mut mf.blocks {
            let mut new_instrs: Vec<MachineInstr> = Vec::new();

            for instr in &block.instructions {
                let is_load = self.instr_may_load(instr);
                new_instrs.push(instr.clone());

                if is_load {
                    // Check if the load def is used in a subsequent address
                    // computation or bounds check. For now, conservatively
                    // insert LFENCE after every load.
                    if self.lfence_is_serializing || self.force_barriers {
                        new_instrs.push(self.create_lfence_instr());
                        self.stats.lfence_inserted += 1;
                        self.stats.slh_loads_hardened += 1;
                        self.stats.total_barriers += 1;
                    }
                }
            }
            block.instructions = new_instrs;
        }
    }

    /// Check if a load's register feeds into an address computation within
    /// a window of N subsequent instructions.
    ///
    /// This is a key optimization: we only need to harden loads whose result
    /// could influence a subsequent speculative memory access.
    pub fn load_feeds_address_in_window(
        &self,
        _block: &MachineBasicBlock,
        _load_idx: usize,
        _window: usize,
    ) -> bool {
        // A thorough implementation would trace the def-use chain of the
        // load's destination register through the window of subsequent
        // instructions, checking if it appears as a base or index register
        // in any memory operand. Simplified here for the overall structure.
        true // conservative default
    }

    /// Optimized SLH: only harden loads that feed into address computations.
    pub fn apply_optimized_slh(&mut self, mf: &mut MachineFunction, window: usize) {
        for block in &mut mf.blocks {
            let mut new_instrs: Vec<MachineInstr> = Vec::new();

            for (i, instr) in block.instructions.iter().enumerate() {
                let is_load = self.instr_may_load(instr);
                new_instrs.push(instr.clone());

                if is_load && self.load_feeds_address_in_window(block, i, window) {
                    if self.lfence_is_serializing || self.force_barriers {
                        new_instrs.push(self.create_lfence_instr());
                        self.stats.lfence_inserted += 1;
                        self.stats.slh_loads_hardened += 1;
                        self.stats.total_barriers += 1;
                    }
                }
            }
            block.instructions = new_instrs;
        }
    }
}

// ============================================================================
// VERW (MDS / TAA / MMIO Mitigation)
// ============================================================================

impl X86RetireControl {
    /// Insert VERW instruction to clear CPU fill buffers.
    ///
    /// VERW (VERify Write segment) is repurposed as a microarchitectural
    /// buffer clearing instruction. It overwrites the CPU's fill buffers
    /// to mitigate MDS (Microarchitectural Data Sampling, CVE-2018-12126/
    /// 12127/12130), TAA (TSX Asynchronous Abort, CVE-2019-11135), and
    /// MMIO stale data (CVE-2022-21123/21125/21166).
    ///
    /// On Linux, this is called `mds_clear_cpu_buffers()` or
    /// `clear_cpu_buffers()` and uses VERW with a dummy memory operand.
    ///
    /// Sequence: VERW m16 (privileged; requires ring 0)
    /// or the alternative: L1D_FLUSH (if VERW not available).
    pub fn insert_verw_buffer_clear(&mut self, bb: &mut MachineBasicBlock) {
        // VERW m16: overwrite CPU buffers
        // Implementation: SUB $16, %RSP; VERW (%RSP); ADD $16, %RSP
        // (VERW needs a memory operand — use stack space)

        let mut sub_rsp = MachineInstr::new(opc(X86Opcode::SUB));
        sub_rsp.operands.push(MachineOperand::PhysReg(RSP));
        sub_rsp.operands.push(MachineOperand::Imm(16));
        bb.instructions.push(sub_rsp);

        // VERW (%RSP) — encoded as 0F 00 /4
        let mut verw_instr = MachineInstr::new(opc(X86Opcode::VERW));
        verw_instr.operands.push(MachineOperand::PhysReg(RSP));
        verw_instr.operands.push(MachineOperand::Imm(0)); // displacement 0
        bb.instructions.push(verw_instr);

        let mut add_rsp = MachineInstr::new(opc(X86Opcode::ADD));
        add_rsp.operands.push(MachineOperand::PhysReg(RSP));
        add_rsp.operands.push(MachineOperand::Imm(16));
        bb.instructions.push(add_rsp);

        self.stats.verw_inserted += 1;
        self.stats.total_barriers += 1;
    }

    /// Insert VERW at every kernel-to-user transition.
    pub fn insert_verw_on_user_return(&mut self, mf: &mut MachineFunction) {
        for block in &mut mf.blocks {
            let mut new_instrs: Vec<MachineInstr> = Vec::new();

            for instr in &block.instructions {
                // Insert VERW before SYSRET / IRET / SYSEXIT
                if instr.opcode == opc(X86Opcode::SYSRET) || instr.opcode == opc(X86Opcode::SYSEXIT)
                {
                    let mut sub_rsp = MachineInstr::new(opc(X86Opcode::SUB));
                    sub_rsp.operands.push(MachineOperand::PhysReg(RSP));
                    sub_rsp.operands.push(MachineOperand::Imm(16));
                    new_instrs.push(sub_rsp);

                    let mut verw = MachineInstr::new(opc(X86Opcode::VERW));
                    verw.operands.push(MachineOperand::PhysReg(RSP));
                    verw.operands.push(MachineOperand::Imm(0));
                    new_instrs.push(verw);

                    let mut add_rsp = MachineInstr::new(opc(X86Opcode::ADD));
                    add_rsp.operands.push(MachineOperand::PhysReg(RSP));
                    add_rsp.operands.push(MachineOperand::Imm(16));
                    new_instrs.push(add_rsp);

                    self.stats.verw_inserted += 1;
                    self.stats.total_barriers += 1;
                }
                new_instrs.push(instr.clone());
            }
            block.instructions = new_instrs;
        }
    }
}

// ============================================================================
// Fill Buffers (L1D Flush for L1TF / Foreshadow Mitigation)
// ============================================================================

impl X86RetireControl {
    /// Insert L1D flush sequence to clear L1 data cache.
    ///
    /// L1 Terminal Fault (L1TF, CVE-2018-3615/3620/3646) allows an attacker
    /// to read L1 cache contents via speculative execution. Mitigation
    /// involves flushing L1D on VM entry or using core scheduling.
    ///
    /// Sequence: MOV to a dummy address to evict cache lines, or use
    /// L1D_FLUSH MSR if available.
    pub fn insert_l1d_flush(&mut self, bb: &mut MachineBasicBlock) {
        // L1D_FLUSH: MOV ECX, 0x10B; MOV EAX, 0x1; WRMSR
        let mut mov_ecx = MachineInstr::new(opc(X86Opcode::MOV));
        mov_ecx.operands.push(MachineOperand::PhysReg(ECX));
        mov_ecx.operands.push(MachineOperand::Imm(0x10B));
        bb.instructions.push(mov_ecx);

        let mut mov_eax = MachineInstr::new(opc(X86Opcode::MOV));
        mov_eax.operands.push(MachineOperand::PhysReg(EAX));
        mov_eax.operands.push(MachineOperand::Imm(0x1));
        bb.instructions.push(mov_eax);

        let mut xor_edx = MachineInstr::new(opc(X86Opcode::XOR));
        xor_edx.operands.push(MachineOperand::PhysReg(EDX));
        xor_edx.operands.push(MachineOperand::PhysReg(EDX));
        bb.instructions.push(xor_edx);

        let mut wrmsr = MachineInstr::new(opc(X86Opcode::NOP));
        wrmsr.operands.push(MachineOperand::Imm(0x0F30_u32 as i64));
        bb.instructions.push(wrmsr);

        self.stats.fill_buffers_inserted += 1;
        self.stats.total_barriers += 1;
    }

    /// Insert fill buffer clearing at function entry (full mitigation).
    pub fn insert_buffer_clear_on_entry(&mut self, mf: &mut MachineFunction) {
        if mf.blocks.is_empty() {
            return;
        }

        if self.has_md_clear {
            self.insert_verw_buffer_clear(&mut mf.blocks[0]);
        } else {
            self.insert_l1d_flush(&mut mf.blocks[0]);
        }
    }
}

// ============================================================================
// Retbleed Mitigation
// ============================================================================

impl X86RetireControl {
    /// Apply Retbleed mitigation to the machine function.
    ///
    /// Retbleed (CVE-2022-29901 / CVE-2022-29900) exploits the CPU's
    /// return stack buffer (RSB) to speculatively execute attacker-controlled
    /// code. Mitigations include:
    ///
    /// 1. RSB stuffing on kernel entry (fills RSB before user code can
    ///    influence it)
    /// 2. Using IBRS (enhanced IBRS on newer CPUs)
    /// 3. Using retpolines for indirect branches + RSB stuffing
    /// 4. For Intel Skylake: using "retpoline with RSB stuffing" or
    ///    "IBRS + enhanced IBRS"
    /// 5. For AMD Zen 1/2: "retpoline with LFENCE/JMP" (generic retpoline)
    pub fn mitigate_retbleed(&mut self, mf: &mut MachineFunction) {
        let cpu_lower = self.subtarget.cpu.to_lowercase();
        let is_intel = !cpu_lower.contains("znver")
            && !cpu_lower.contains("amd")
            && !cpu_lower.contains("generic");

        // Strategy 1: If enhanced IBRS is available, set it and we're done
        if self.has_enhanced_ibrs && self.ibrs_mode != X86IBRSMode::None {
            // Enhanced IBRS is a one-time setup; we don't need per-entry barriers
            return;
        }

        // Strategy 2: RSB stuffing on kernel/user entry points
        if self.enable_rsb_stuffing {
            self.stuff_rsb_auto(mf, self.rsb_stuff_depth);
            self.stats.retbleed_barriers_inserted += 1;
        }

        // Strategy 3: For Intel Skylake (no eIBRS), use retpoline + RSB
        if is_intel
            && !self.has_enhanced_ibrs
            && (cpu_lower.contains("skylake")
                || cpu_lower.contains("cascadelake")
                || cpu_lower.contains("cooperlake"))
        {
            // Emit RSB stuffing before RET instructions
            for block in &mut mf.blocks {
                let mut new_instrs: Vec<MachineInstr> = Vec::new();
                for instr in &block.instructions {
                    if self.is_return(instr) {
                        // Stuff RSB before RET
                        // Insert LFENCE (serialize) to drain RSB
                        new_instrs.push(self.create_lfence_instr());
                        self.stats.lfence_inserted += 1;
                        self.stats.retbleed_barriers_inserted += 1;
                        self.stats.total_barriers += 1;
                    }
                    new_instrs.push(instr.clone());
                }
                block.instructions = new_instrs;
            }
        }

        // Strategy 4: For AMD Zen 1/2, LFENCE before RET (jmp2ret mitigation)
        if cpu_lower.contains("znver1") || cpu_lower.contains("znver2") {
            for block in &mut mf.blocks {
                let mut new_instrs: Vec<MachineInstr> = Vec::new();
                for instr in &block.instructions {
                    if self.is_return(instr) {
                        new_instrs.push(self.create_lfence_instr());
                        self.stats.lfence_inserted += 1;
                        self.stats.retbleed_barriers_inserted += 1;
                        self.stats.total_barriers += 1;
                    }
                    new_instrs.push(instr.clone());
                }
                block.instructions = new_instrs;
            }
        }
    }

    /// Convert indirect CALL/JMP to retpoline sequences.
    ///
    /// A retpoline replaces:
    /// ```asm
    ///     jmp *%rax
    /// ```
    /// with:
    /// ```asm
    ///     call 1f
    /// 1:  pause; lfence
    ///     jmp 1b
    /// ```
    /// (This is the "generic retpoline" — the actual implementation varies
    /// by μarch. Some use INT3 instead of PAUSE, some use LFENCE+JMP.)
    pub fn insert_retpoline_for_indirect_branch(
        &mut self,
        bb: &mut MachineBasicBlock,
        instr_idx: usize,
    ) {
        if instr_idx >= bb.instructions.len() {
            return;
        }

        // Replace the indirect branch with a retpoline sequence
        // Simplified: replace JMP *reg with PUSH reg; RET (retpoline pattern)
        let mut new_seq: Vec<MachineInstr> = Vec::new();

        // PAUSE (F3 90)
        let mut pause = MachineInstr::new(opc(X86Opcode::NOP));
        pause.operands.push(MachineOperand::Imm(0xF390_u32 as i64));
        new_seq.push(pause);

        // LFENCE
        new_seq.push(self.create_lfence_instr());

        // INT3 (trap if we actually execute this)
        let mut int3 = MachineInstr::new(opc(X86Opcode::INT3));
        new_seq.push(int3);

        // Replace the single instruction with the sequence
        bb.instructions.splice(instr_idx..=instr_idx, new_seq);

        self.stats.bti_barriers_inserted += 1;
        self.stats.total_barriers += 1;
    }
}

// ============================================================================
// Full Mitigation Pipeline
// ============================================================================

impl X86RetireControl {
    /// Apply the full speculation mitigation pipeline to a machine function.
    ///
    /// This is the top-level entry point for retire control. It applies
    /// all enabled mitigations in the correct order:
    ///
    /// 1. IBRS setup (if kernel-only or enhanced mode)
    /// 2. IBPB barriers (if trigger set)
    /// 3. RSB stuffing (if enabled)
    /// 4. SSBD barriers (if SSBD enabled)
    /// 5. PSFD barriers (if PSFD supported)
    /// 6. Speculative Load Hardening (SLH)
    /// 7. Branch Target Injection (BTI) barriers
    /// 8. Retbleed mitigation
    /// 9. VERW buffer clearing (on user return)
    pub fn apply_mitigations(&mut self, mf: &mut MachineFunction) {
        // 1. IBRS: set on kernel entry, clear on user return
        if self.ibrs_mode == X86IBRSMode::KernelOnly {
            if let Some(entry_bb) = mf.blocks.first_mut() {
                self.insert_ibrs_on(entry_bb);
            }
        }

        // 2. IBPB: insert barriers according to trigger policy
        if self.ibpb_trigger != X86IBPBTrigger::Never {
            self.insert_ibpb_barriers_by_policy(mf);
        }

        // 3. RSB stuffing on entry
        if self.enable_rsb_stuffing && !mf.blocks.is_empty() {
            let depth = self.rsb_stuff_depth;
            self.stuff_return_stack_buffer(mf, depth);
        }

        // 4. SSBD barriers
        if self.ssbd_mode != X86SSBDMode::Off {
            for block in &mut mf.blocks {
                self.insert_ssbd_on(block);
            }
        }

        // 5. PSFD barriers
        if self.has_psfd {
            for block in &mut mf.blocks {
                self.insert_psfd_on(block);
            }
        }

        // 6. SLH
        self.apply_speculative_load_hardening(mf);

        // 7. BTI barriers
        if self.ibpb_trigger == X86IBPBTrigger::OnIndirectBranch
            || self.ibrs_mode != X86IBRSMode::None
        {
            self.insert_bti_barriers(mf);
        }

        // 8. Retbleed mitigation
        self.mitigate_retbleed(mf);

        // 9. VERW on user return if MDS mitigation enabled
        if self.use_verw_mitigation && self.has_md_clear {
            self.insert_verw_on_user_return(mf);
        }
    }

    /// Insert IBPB barriers according to the current trigger policy.
    fn insert_ibpb_barriers_by_policy(&mut self, mf: &mut MachineFunction) {
        match self.ibpb_trigger {
            X86IBPBTrigger::Never => {}
            X86IBPBTrigger::OnContextSwitch => {
                // Insert IBPB at function entry (context switch boundary)
                if let Some(entry_bb) = mf.blocks.first_mut() {
                    self.insert_ibpb(entry_bb);
                }
            }
            X86IBPBTrigger::OnPrivilegeTransition => {
                // Insert IBPB at entry and exit (kernel entry/exit)
                if let Some(entry_bb) = mf.blocks.first_mut() {
                    self.insert_ibpb(entry_bb);
                }
                if let Some(exit_bb) = mf.blocks.last_mut() {
                    self.insert_ibpb(exit_bb);
                }
            }
            X86IBPBTrigger::OnIndirectBranch => {
                self.insert_ibpb_before_indirect_branches(mf);
            }
            X86IBPBTrigger::AlwaysOnUserReturn => {
                // Insert IBPB before every user-space return
                for block in &mut mf.blocks {
                    let mut new_instrs: Vec<MachineInstr> = Vec::new();
                    for instr in &block.instructions {
                        if instr.opcode == opc(X86Opcode::SYSRET)
                            || instr.opcode == opc(X86Opcode::SYSEXIT)
                            || self.is_return(instr)
                        {
                            // Insert IBPB before return to user
                            new_instrs.push(self.create_lfence_instr());
                            self.stats.ibpb_inserted += 1;
                            self.stats.total_barriers += 1;
                        }
                        new_instrs.push(instr.clone());
                    }
                    block.instructions = new_instrs;
                }
            }
        }
    }
}

// ============================================================================
// Direct MSR Write Primitives (for Inline Assembly or CodeGen)
// ============================================================================

impl X86RetireControl {
    /// Emit WRMSR to set IA32_SPEC_CTRL (MSR 0x48) to a given value.
    ///
    /// This is the "set it and forget it" approach used by enhanced IBRS
    /// or when setting SSBD globally.
    ///
    /// # Safety
    ///
    /// Writing MSRs requires kernel-mode (ring 0) privileges. Calling this
    /// from user-mode will cause a #GP fault.
    pub fn emit_set_spec_ctrl_msr(&mut self, bb: &mut MachineBasicBlock, value: u64) {
        let low = (value & 0xFFFF_FFFF) as i64;
        let high = ((value >> 32) & 0xFFFF_FFFF) as i64;

        // MOV ECX, 0x48
        let mut mov_ecx = MachineInstr::new(opc(X86Opcode::MOV));
        mov_ecx.operands.push(MachineOperand::PhysReg(ECX));
        mov_ecx.operands.push(MachineOperand::Imm(0x48));
        bb.instructions.push(mov_ecx);

        // MOV EAX, low32
        let mut mov_eax = MachineInstr::new(opc(X86Opcode::MOV));
        mov_eax.operands.push(MachineOperand::PhysReg(EAX));
        mov_eax.operands.push(MachineOperand::Imm(low));
        bb.instructions.push(mov_eax);

        // MOV EDX, high32
        let mut mov_edx = MachineInstr::new(opc(X86Opcode::MOV));
        mov_edx.operands.push(MachineOperand::PhysReg(EDX));
        mov_edx.operands.push(MachineOperand::Imm(high));
        bb.instructions.push(mov_edx);

        // WRMSR
        let mut wrmsr = MachineInstr::new(opc(X86Opcode::NOP));
        wrmsr.operands.push(MachineOperand::Imm(0x0F30_u32 as i64));
        bb.instructions.push(wrmsr);

        self.stats.total_barriers += 1;
    }

    /// Emit RDMSR to read IA32_SPEC_CTRL (MSR 0x48).
    pub fn emit_read_spec_ctrl_msr(&mut self, bb: &mut MachineBasicBlock) {
        // MOV ECX, 0x48
        let mut mov_ecx = MachineInstr::new(opc(X86Opcode::MOV));
        mov_ecx.operands.push(MachineOperand::PhysReg(ECX));
        mov_ecx.operands.push(MachineOperand::Imm(0x48));
        bb.instructions.push(mov_ecx);

        // RDMSR (result in EDX:EAX)
        let mut rdmsr = MachineInstr::new(opc(X86Opcode::NOP));
        rdmsr.operands.push(MachineOperand::Imm(0x0F32_u32 as i64));
        bb.instructions.push(rdmsr);
    }

    /// Emit the full "speculation-safe function entry" prologue.
    ///
    /// This emits:
    /// 1. IBRS_ON (if kernel-only IBRS)
    /// 2. IBPB (if context-switch trigger)
    /// 3. RSB stuff (if enabled)
    /// 4. SSBD (if needed)
    /// 5. VERW (if MDS mitigation)
    pub fn emit_mitigation_prologue(&mut self, bb: &mut MachineBasicBlock) {
        // IBRS on entry
        if self.ibrs_mode == X86IBRSMode::KernelOnly {
            self.insert_ibrs_on(bb);
        }

        // IBPB on entry
        if self.ibpb_trigger == X86IBPBTrigger::OnContextSwitch
            || self.ibpb_trigger == X86IBPBTrigger::OnPrivilegeTransition
        {
            self.insert_ibpb(bb);
        }

        // RSB stuff (inline sequence)
        if self.enable_rsb_stuffing {
            self.stuff_rsb_intel_sequence(bb, self.rsb_stuff_depth);
        }

        // SSBD
        if self.ssbd_mode == X86SSBDMode::ForceOn || self.ssbd_mode == X86SSBDMode::ViaMSR {
            self.insert_ssbd_on(bb);
        }

        // PSFD
        if self.has_psfd {
            self.insert_psfd_on(bb);
        }

        // VERW buffer clear
        if self.use_verw_mitigation && self.has_md_clear {
            self.insert_verw_buffer_clear(bb);
        }
    }

    /// Emit the "speculation-safe function exit" epilogue.
    ///
    /// This emits:
    /// 1. VERW (if MDS mitigation)
    /// 2. IBRS_OFF (if kernel-only IBRS)
    /// 3. STIBP_OFF (if per-thread STIBP)
    pub fn emit_mitigation_epilogue(&mut self, bb: &mut MachineBasicBlock) {
        // VERW before return to user
        if self.use_verw_mitigation && self.has_md_clear {
            self.insert_verw_buffer_clear(bb);
        }

        // STIBP off (if per-thread or kernel-only)
        if self.stibp_mode == X86STIBPMode::PerProcess || self.stibp_mode == X86STIBPMode::SMTOnly {
            self.insert_stibp_off(bb);
        }

        // IBRS off (kernel-only mode)
        if self.ibrs_mode == X86IBRSMode::KernelOnly {
            self.insert_ibrs_off(bb);
        }
    }
}

// ============================================================================
// Retire Control Analysis & Reporting
// ============================================================================

impl X86RetireControl {
    /// Check whether the current CPU is vulnerable to Spectre v2 (BTI).
    pub fn is_vulnerable_to_spectre_v2(&self) -> bool {
        // CPUs without IBPB/IBRS are potentially vulnerable
        if self.has_enhanced_ibrs {
            return false; // Enhanced IBRS provides full protection
        }
        if self.ibrs_mode == X86IBRSMode::None && self.ibpb_trigger == X86IBPBTrigger::Never {
            return true; // No mitigation applied
        }
        // If mitigations are enabled, assume they work (in practice,
        // completeness depends on μarch-specific details)
        false
    }

    /// Check whether the current CPU is vulnerable to Spectre v4 (SSB).
    pub fn is_vulnerable_to_spectre_v4(&self) -> bool {
        if self.ssb_no {
            return false; // CPU reports not vulnerable
        }
        self.ssbd_mode == X86SSBDMode::Off
    }

    /// Check whether the current CPU is vulnerable to Retbleed.
    pub fn is_vulnerable_to_retbleed(&self) -> bool {
        if self.has_enhanced_ibrs && self.ibrs_mode == X86IBRSMode::Enhanced {
            return false;
        }
        if self.pbrsb_no {
            return false; // Post-barrier RSB is protected
        }
        // Retbleed affects Intel Skylake-family and AMD Zen 1/2
        let cpu_lower = self.subtarget.cpu.to_lowercase();
        cpu_lower.contains("skylake")
            || cpu_lower.contains("cascadelake")
            || cpu_lower.contains("cooperlake")
            || cpu_lower.contains("znver1")
            || cpu_lower.contains("znver2")
    }

    /// Generate a human-readable report of the current mitigation status.
    pub fn mitigation_report(&self) -> String {
        let mut report = String::new();
        report.push_str("=== X86 Retire Control: Mitigation Status ===\n");
        report.push_str(&format!("  CPU: {}\n", self.subtarget.cpu));
        report.push_str(&format!(
            "  LFENCE is serializing: {}\n",
            self.lfence_is_serializing
        ));
        report.push_str(&format!("  IBRS mode: {}\n", self.ibrs_mode));
        report.push_str(&format!("  IBPB trigger: {}\n", self.ibpb_trigger));
        report.push_str(&format!("  STIBP mode: {}\n", self.stibp_mode));
        report.push_str(&format!("  SSBD mode: {}\n", self.ssbd_mode));
        report.push_str(&format!("  RSB stuff depth: {}\n", self.rsb_stuff_depth));
        report.push_str(&format!("  Enhanced IBRS: {}\n", self.has_enhanced_ibrs));
        report.push_str(&format!("  IBPB on all CPUs: {}\n", self.has_ibpb_all));
        report.push_str(&format!("  RSBA (RSB alt): {}\n", self.has_rsba));
        report.push_str(&format!("  SSB not vulnerable: {}\n", self.ssb_no));
        report.push_str(&format!("  MD_CLEAR supported: {}\n", self.has_md_clear));
        report.push_str(&format!("  PBRSB protected: {}\n", self.pbrsb_no));
        report.push_str(&format!("  PSFD supported: {}\n", self.has_psfd));
        report.push_str(&format!(
            "  VERW mitigation: {}\n",
            self.use_verw_mitigation
        ));
        report.push_str(&format!(
            "  Spectre v2 vulnerable: {}\n",
            self.is_vulnerable_to_spectre_v2()
        ));
        report.push_str(&format!(
            "  Spectre v4 vulnerable: {}\n",
            self.is_vulnerable_to_spectre_v4()
        ));
        report.push_str(&format!(
            "  Retbleed vulnerable: {}\n",
            self.is_vulnerable_to_retbleed()
        ));
        report.push_str(&format!("\n  {}\n", self.stats.summary()));
        report
    }

    /// Reset all statistics counters to zero.
    pub fn reset_stats(&mut self) {
        self.stats = RetireControlStats::default();
    }
}

// ============================================================================
// CPUID-based Feature Detection Helpers
// ============================================================================

/// Simulate CPUID leaf 7 sub-leaf 0 to detect speculation control features.
///
/// In a real implementation, this would execute the CPUID instruction
/// and parse the returned bits. Here we provide a structure for
/// interpreting the standard CPUID bit fields.
#[derive(Debug, Clone, Default)]
pub struct X86SpecCtrlFeatures {
    pub has_ibrs_ibpb: bool, // CPUID.7.EDX[26]: IA32_SPEC_CTRL supported
    pub has_stibp: bool,     // CPUID.7.EDX[27]: STIBP supported
    pub has_l1d_flush: bool, // CPUID.7.EDX[28]: L1D_FLUSH supported
    pub has_arch_cap: bool,  // CPUID.7.EDX[29]: IA32_ARCH_CAPABILITIES
    pub has_ssbd: bool,      // CPUID.7.EDX[31]: SSBD supported
    pub has_md_clear: bool,  // CPUID.7.EDX[10]: MD_CLEAR (VERW clears buffers)
    pub has_serialize: bool, // CPUID.7.EDX[14]: SERIALIZE instruction
    pub has_tsx_ctrl: bool,  // CPUID.7.EDX[13]: TSX_CTRL MSR
    pub has_psfd: bool,      // (AMD-specific) PSFD supported
}

impl X86SpecCtrlFeatures {
    /// Detect features by interpreting a raw EDX value from CPUID leaf 7.
    pub fn from_cpuid_edx(edx: u32) -> Self {
        Self {
            has_ibrs_ibpb: (edx & (1 << 26)) != 0,
            has_stibp: (edx & (1 << 27)) != 0,
            has_l1d_flush: (edx & (1 << 28)) != 0,
            has_arch_cap: (edx & (1 << 29)) != 0,
            has_ssbd: (edx & (1 << 31)) != 0,
            has_md_clear: (edx & (1 << 10)) != 0,
            has_serialize: (edx & (1 << 14)) != 0,
            has_tsx_ctrl: (edx & (1 << 13)) != 0,
            has_psfd: false, // AMD-specific; checked separately
        }
    }
}

/// Detect speculation control features from the subtarget.
///
/// This uses the subtarget's feature set (strings) rather than raw CPUID,
/// as the subtarget already encodes which features are enabled.
pub fn detect_spec_ctrl_features(subtarget: &X86Subtarget) -> X86SpecCtrlFeatures {
    X86SpecCtrlFeatures {
        has_ibrs_ibpb: subtarget.features.contains("ibrs") || subtarget.features.contains("ibpb"),
        has_stibp: subtarget.features.contains("stibp"),
        has_l1d_flush: subtarget.features.contains("l1d-flush"),
        has_arch_cap: subtarget.features.contains("arch-capabilities"),
        has_ssbd: subtarget.features.contains("ssbd"),
        has_md_clear: subtarget.features.contains("md-clear"),
        has_serialize: subtarget.features.contains("serialize"),
        has_tsx_ctrl: subtarget.features.contains("tsx-ctrl"),
        has_psfd: subtarget.features.contains("psfd"),
    }
}

// ============================================================================
// RSB Depth Detection per CPU Model
// ============================================================================

/// Determine the RSB depth for a given CPU model string.
pub fn rsb_depth_for_cpu(cpu: &str) -> usize {
    let cpu_lower = cpu.to_lowercase();

    if cpu_lower.contains("alderlake") || cpu_lower.contains("raptorlake") {
        X86_RSB_DEPTH_INTEL_ALDERLAKE
    } else if cpu_lower.contains("icelake")
        || cpu_lower.contains("tigerlake")
        || cpu_lower.contains("rocketlake")
        || cpu_lower.contains("meteorlake")
        || cpu_lower.contains("sapphirerapids")
    {
        X86_RSB_DEPTH_INTEL_ICELAKE
    } else if cpu_lower.contains("znver3")
        || cpu_lower.contains("znver4")
        || cpu_lower.contains("znver5")
    {
        X86_RSB_DEPTH_AMD_ZEN3
    } else if cpu_lower.contains("znver1") || cpu_lower.contains("znver2") {
        X86_RSB_DEPTH_AMD_ZEN
    } else if cpu_lower.contains("skylake")
        || cpu_lower.contains("cascadelake")
        || cpu_lower.contains("cooperlake")
        || cpu_lower.contains("cannonlake")
        || cpu_lower.contains("haswell")
        || cpu_lower.contains("broadwell")
    {
        X86_RSB_DEPTH_INTEL_SKYLAKE
    } else {
        X86_RSB_DEPTH_DEFAULT
    }
}

// ============================================================================
// IBRS / IBPB Capability Detection
// ============================================================================

/// Determine the available IBRS capabilities based on CPU model and features.
pub fn detect_ibrs_capabilities(subtarget: &X86Subtarget) -> (bool, bool, bool) {
    let cpu_lower = subtarget.cpu.to_lowercase();

    // Enhanced IBRS availability by CPU generation
    let has_enhanced_ibrs = subtarget.features.contains("enhanced-ibrs")
        || cpu_lower.contains("icelake")
        || cpu_lower.contains("tigerlake")
        || cpu_lower.contains("rocketlake")
        || cpu_lower.contains("alderlake")
        || cpu_lower.contains("raptorlake")
        || cpu_lower.contains("meteorlake")
        || cpu_lower.contains("sapphirerapids")
        || cpu_lower.contains("znver3")
        || cpu_lower.contains("znver4")
        || cpu_lower.contains("znver5");

    let has_ibpb = subtarget.features.contains("ibpb") || !cpu_lower.contains("nehalem");
    let has_stibp = subtarget.features.contains("stibp")
        || cpu_lower.contains("skylake")
        || cpu_lower.contains("icelake")
        || cpu_lower.contains("znver2");

    (has_enhanced_ibrs, has_ibpb, has_stibp)
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::codegen::{MachineBasicBlock, MachineFunction};

    fn make_test_subtarget() -> X86Subtarget {
        X86Subtarget::new("x86_64-unknown-linux-gnu", "skylake", "+sse2,+avx2")
    }

    fn make_test_mf() -> MachineFunction {
        let mut mf = MachineFunction::new("test_fn");
        let mut bb = MachineBasicBlock {
            id: 0,
            name: "entry".to_string(),
            instructions: Vec::new(),
            successors: Vec::new(),
            predecessors: Vec::new(),
            is_entry: true,
        };
        mf.push_block(bb);
        mf
    }

    #[test]
    fn test_retire_control_creation() {
        let subtarget = make_test_subtarget();
        let rc = X86RetireControl::new(&subtarget);
        assert_eq!(rc.ibrs_mode, X86IBRSMode::None);
        assert_eq!(rc.ibpb_trigger, X86IBPBTrigger::Never);
        assert_eq!(rc.lfence_is_serializing, true); // Skylake has LFENCE serializing
        assert_eq!(rc.stats.total_barriers, 0);
    }

    #[test]
    fn test_lfence_barrier_insertion() {
        let subtarget = make_test_subtarget();
        let mut rc = X86RetireControl::new(&subtarget);
        let mut mf = make_test_mf();

        rc.insert_lfence_barrier(&mut mf.blocks[0]);
        assert_eq!(rc.stats.lfence_inserted, 1);
        assert_eq!(rc.stats.total_barriers, 1);
        assert!(mf.blocks[0].instructions.len() > 0);
    }

    #[test]
    fn test_ibrs_mode_display() {
        assert_eq!(format!("{}", X86IBRSMode::None), "none");
        assert_eq!(format!("{}", X86IBRSMode::KernelOnly), "kernel-only");
        assert_eq!(format!("{}", X86IBRSMode::Enhanced), "enhanced");
        assert_eq!(
            format!("{}", X86IBRSMode::RetpolineAware),
            "retpoline-aware"
        );
    }

    #[test]
    fn test_ibpb_trigger_display() {
        assert_eq!(format!("{}", X86IBPBTrigger::Never), "never");
        assert_eq!(
            format!("{}", X86IBPBTrigger::OnContextSwitch),
            "on-context-switch"
        );
        assert_eq!(
            format!("{}", X86IBPBTrigger::AlwaysOnUserReturn),
            "always-on-user-return"
        );
    }

    #[test]
    fn test_stibp_mode_display() {
        assert_eq!(format!("{}", X86STIBPMode::Off), "off");
        assert_eq!(format!("{}", X86STIBPMode::Always), "always");
    }

    #[test]
    fn test_ssbd_mode_display() {
        assert_eq!(format!("{}", X86SSBDMode::Off), "off");
        assert_eq!(format!("{}", X86SSBDMode::ForceOn), "force-on");
    }

    #[test]
    fn test_spec_barrier_display() {
        assert_eq!(format!("{}", X86SpecBarrier::LFENCE), "LFENCE");
        assert_eq!(format!("{}", X86SpecBarrier::IBPB), "IBPB");
        assert_eq!(format!("{}", X86SpecBarrier::RSBStuff), "RSB-STUFF");
        assert_eq!(format!("{}", X86SpecBarrier::FillBuffers), "FILL-BUFFERS");
    }

    #[test]
    fn test_rsb_depth_for_known_cpus() {
        assert_eq!(rsb_depth_for_cpu("skylake"), X86_RSB_DEPTH_INTEL_SKYLAKE);
        assert_eq!(rsb_depth_for_cpu("icelake"), X86_RSB_DEPTH_INTEL_ICELAKE);
        assert_eq!(rsb_depth_for_cpu("znver3"), X86_RSB_DEPTH_AMD_ZEN3);
        assert_eq!(rsb_depth_for_cpu("unknown"), X86_RSB_DEPTH_DEFAULT);
    }

    #[test]
    fn test_ibrs_capability_detection() {
        let sub = X86Subtarget::new("x86_64-unknown-linux-gnu", "icelake", "");
        let (enh, ibpb, stibp) = detect_ibrs_capabilities(&sub);
        assert!(enh);
        assert!(ibpb);
        assert!(stibp);
    }

    #[test]
    fn test_vulnerability_checks() {
        let subtarget = make_test_subtarget();
        let mut rc = X86RetireControl::new(&subtarget);
        // Skylake without IBRS is vulnerable to Spectre v2
        assert!(rc.is_vulnerable_to_spectre_v2());
        // Skylake is vulnerable to Retbleed
        assert!(rc.is_vulnerable_to_retbleed());
    }

    #[test]
    fn test_mitigation_report() {
        let subtarget = make_test_subtarget();
        let rc = X86RetireControl::new(&subtarget);
        let report = rc.mitigation_report();
        assert!(report.contains("skylake"));
        assert!(report.contains("LFENCE is serializing: true"));
    }

    #[test]
    fn test_stats_merge() {
        let mut s1 = RetireControlStats::default();
        s1.lfence_inserted = 5;
        s1.total_barriers = 5;

        let mut s2 = RetireControlStats::default();
        s2.lfence_inserted = 3;
        s2.ibpb_inserted = 2;
        s2.total_barriers = 5;

        s1.merge(&s2);
        assert_eq!(s1.lfence_inserted, 8);
        assert_eq!(s1.ibpb_inserted, 2);
        assert_eq!(s1.total_barriers, 10);
    }

    #[test]
    fn test_detect_lfence_serializing() {
        assert!(detect_lfence_serializing("skylake"));
        assert!(detect_lfence_serializing("icelake-server"));
        assert!(detect_lfence_serializing("znver3"));
        assert!(!detect_lfence_serializing("haswell"));
        assert!(!detect_lfence_serializing("pentium4"));
    }

    #[test]
    fn test_spec_ctrl_features_from_cpuid() {
        let raw_edx: u32 = (1 << 26) | (1 << 27) | (1 << 29) | (1 << 31);
        let features = X86SpecCtrlFeatures::from_cpuid_edx(raw_edx);
        assert!(features.has_ibrs_ibpb);
        assert!(features.has_stibp);
        assert!(features.has_arch_cap);
        assert!(features.has_ssbd);
        assert!(!features.has_md_clear);
    }

    #[test]
    fn test_set_configuration() {
        let subtarget = make_test_subtarget();
        let mut rc = X86RetireControl::new(&subtarget);
        rc.set_ibrs_mode(X86IBRSMode::KernelOnly);
        rc.set_ibpb_trigger(X86IBPBTrigger::OnContextSwitch);
        rc.set_stibp_mode(X86STIBPMode::SMTOnly);
        rc.set_ssbd_mode(X86SSBDMode::ForceOn);
        rc.set_rsb_stuff_depth(24);
        rc.set_rsb_stuff_iterations(2);
        rc.set_force_barriers(true);
        rc.set_use_verw_mitigation(true);

        assert_eq!(rc.ibrs_mode, X86IBRSMode::KernelOnly);
        assert_eq!(rc.ibpb_trigger, X86IBPBTrigger::OnContextSwitch);
        assert_eq!(rc.stibp_mode, X86STIBPMode::SMTOnly);
        assert_eq!(rc.ssbd_mode, X86SSBDMode::ForceOn);
        assert_eq!(rc.rsb_stuff_depth, 24);
        assert_eq!(rc.rsb_stuff_iterations, 2);
        assert!(rc.force_barriers);
        assert!(rc.use_verw_mitigation);
    }

    #[test]
    fn test_rsb_stuff_depth_clamp() {
        let subtarget = make_test_subtarget();
        let mut rc = X86RetireControl::new(&subtarget);
        rc.set_rsb_stuff_depth(X86_RSB_STUFF_MAX_DEPTH * 2); // Too large
        assert_eq!(rc.rsb_stuff_depth, X86_RSB_STUFF_MAX_DEPTH);
    }

    #[test]
    fn test_is_indirect_branch_detection() {
        let subtarget = make_test_subtarget();
        let rc = X86RetireControl::new(&subtarget);

        let call = MachineInstr::new(opc(X86Opcode::CALL));
        let ret = MachineInstr::new(opc(X86Opcode::RET));
        let add = MachineInstr::new(opc(X86Opcode::ADD));

        assert!(rc.is_indirect_branch(&call));
        assert!(rc.is_indirect_branch(&ret));
        assert!(!rc.is_indirect_branch(&add));
    }
}