strat9_kernel/arch/x86_64/

smp.rs

//! SMP (Symmetric Multi-Processing) boot for x86_64.
//!
//! Boots Application Processors (APs) using the legacy INIT+SIPI sequence.
//! Inspired by Redox-OS's approach: a minimal trampoline does the 16=>64 bit
//! mode switch, then jumps directly to `smp_main` in Rust.
//!
//! Data layout after the trampoline code (written by BSP, read by AP):
//!   offset +0: PML4 physical address (CR3)
//!   offset +8: kernel stack top virtual address (RSP)
//!
//! Synchronization: AP increments `BOOTED_CORES` after finishing per-CPU init.
//! BSP spins until all expected APs are online.

use core::{
    arch::global_asm,
    sync::atomic::{AtomicBool, AtomicUsize, Ordering},
};

use alloc::{vec, vec::Vec};
use x86_64::{
    structures::paging::{Page, PageTableFlags, PhysFrame, Size4KiB},
    PhysAddr, VirtAddr,
};

use crate::{
    acpi::madt,
    arch::x86_64::{apic, idt, io::io_wait, percpu, timer},
    memory,
};

/// Physical address where the SMP trampoline is copied.
pub const TRAMPOLINE_PHYS_ADDR: u64 = 0x8000;

/// Number of booted cores (starts at 1 for BSP).
static BOOTED_CORES: AtomicUsize = AtomicUsize::new(1);
/// Counter for synchronization barriers.
static SYNC_BARRIER: AtomicUsize = AtomicUsize::new(0);
/// Target count for the rendezvous barrier (set by BSP before barrier).
static BARRIER_TARGET: AtomicUsize = AtomicUsize::new(0);
/// Gate used by BSP to release APs into scheduler/timer start.
static AP_SCHED_GATE_OPEN: AtomicBool = AtomicBool::new(false);
/// AP ACK: set to apic_id by each AP at start of smp_main(), BSP waits for it.
/// Sentinel usize::MAX means "no ACK yet" (avoids collision with APIC ID 0).
static AP_REACHED_RUST: AtomicUsize = AtomicUsize::new(usize::MAX);

// ---------------------------------------------------------------------------
// Trampoline: 16-bit => 32-bit => 64-bit mode switch.
//
// The AP starts in real mode at the SIPI vector (0x8000).  This stub:
//   1. Loads a GDT embedded at known physical offsets.
//   2. Enables protected mode, then PAE + long mode + paging.
//   3. Loads the kernel PML4 (CR3) and kernel stack (RSP) from the data area.
//   4. Jumps to smp_main (64-bit Rust code).
//
// LAYOUT (physical addresses, copied to 0x8000):
//   0x8000:  cli ; cld ; ljmp 0, 0x8040         (16-bit)
//   0x8010:  _gdt_table  (32 bytes: null, code64, data, code32)
//   0x8030:  _gdt        (GDTR: limit=31, base=0x8010)
//   0x8040:  real-mode setup (xor ax,ax; lgdt; enter PM)
//   SMP_PM_ADDR:   32-bit code (PAT check, PAE + NXE + LME + paging)
//   SMP_LONG_ADDR: 64-bit code (load stack, jump to smp_main)
//
// CRITICAL: The GDT table and descriptor MUST occupy these exact offsets.
// The `lgdt [0x8030]` instruction reads physical 0x8030 which contains the
// GDTR. Without this embedded GDT data, the AP loads garbage and triple-faults.
//
// Data area (at smp_trampoline_end, written by BSP via HHDM):
//   +0 u64: CR3 (PML4 physical address)
//   +8 u64: RSP (kernel stack top virtual address)
// ---------------------------------------------------------------------------
#[cfg(target_arch = "x86_64")]
global_asm!(
    r#"
.section .text
.code16

.global smp_trampoline
.global smp_trampoline_end

.set SMP_VAR_ADDR, 0x8000 + (smp_trampoline_end - smp_trampoline)
.set SMP_PM_ADDR, 0x8000 + (smp_trampoline_32 - smp_trampoline)
.set SMP_LONG_ADDR, 0x8000 + (smp_trampoline_64 - smp_trampoline)

smp_trampoline:
    cli
    cld
    # Jump over the GDT data : code continues at 0x8040.
    ljmp 0, 0x8040

# -------------------------------------------------------------------
# GDT : must be at physical offset 0x10 so that:
#   _gdt_table starts at 0x8010, _gdt (GDTR) is at 0x8030.
# -------------------------------------------------------------------
.align 16
_gdt_table:
    .long 0, 0                       # null  (selector 0)
    .long 0x0000ffff, 0x00af9b00     # code64, accessed (identity page is RX)
    .long 0x0000ffff, 0x00cf9300     # data, accessed
    .long 0x0000ffff, 0x00cf9b00     # code32, accessed
_gdt:
    .word _gdt - _gdt_table - 1      # limit = 31 (4 entries × 8 - 1)
    .long 0x8010                     # base  = 0x8010
    .long 0, 0                       # padding
.align 64

# -------------------------------------------------------------------
# Real-mode setup continues at 0x8040.
# -------------------------------------------------------------------
    xor ax, ax
    mov ds, ax
    lgdt [0x8030]                    # loads GDTR from 0x8030

    # Enter protected mode
    mov eax, cr0
    or eax, 1
    mov cr0, eax
    # Far JMP ptr16:16 (EA iw iw). LLVM's Intel-syntax `ljmp` rejects
    # this forward expression; data directives resolve it after layout.
    # The trampoline lives entirely in 0x8000..0x8fff, so IP fits in u16.
    .byte 0xea
    .word SMP_PM_ADDR
    .word 24                        # => code32 segment

.align 32
.code32
smp_trampoline_32:
    mov ax, 16
    mov ds, ax
    mov ss, ax

    # INIT does not reset PAT. Check the selectors used by BSP's new tables
    # before enabling paging; stop this AP explicitly if firmware disagrees.
    mov ecx, 0x277
    rdmsr
    and eax, 0xff0000ff
    cmp eax, 6                       # PAT[0] = WB, PAT[3] = UC
    je 3f
    mov al, 0x50                     # E9: P (PAT mismatch)
    out 0xe9, al
2:  hlt
    jmp 2b
3:
    mov eax, cr0
    or eax, 0x40000000              # CD = 1
    and eax, 0xdfffffff             # NW = 0
    mov cr0, eax
    wbinvd

    # Enable PAE + PSE + OSFXSR + OSXMMEXCPT
    # NOTE: do NOT force SMEP/SMAP (CR4 bits 20/21) here. qemu64 (and many
    # older hosts) lack them: `mov cr4` then raises #GP and the AP dies
    # silently (no IDT yet -> triple fault -> "waiting for APs" hang).
    # Feature-gated bits (OSXSAVE, SMEP, SMAP) are enabled conditionally
    # in Rust once CPUID detection has run.
    mov eax, cr4
    or eax, 0x630
    mov cr4, eax

    # Enable Long Mode (EFER.LME)
    mov ecx, 0xc0000080
    xor edx, edx
    rdmsr
    or eax, 0x901                    # NXE + LME + SCE
    wrmsr

    # Load kernel PML4 from data area
    mov eax, [SMP_VAR_ADDR]
    mov cr3, eax

    # Enable paging (activates long mode)
    mov eax, cr0
    and eax, 0x9FFFFFF3              # Clear CD/NW/EM/TS
    or eax, 0x80010002               # PG + WP + MP
    mov cr0, eax

    # Far JMP ptr16:32 (EA id iw), decoded while CS is still 32-bit.
    # Keep the destination symbolic when the preceding stub changes size.
    .byte 0xea
    .long SMP_LONG_ADDR
    .word 8                         # => code64 segment

.align 32
.code64
smp_trampoline_64:
    # Load kernel stack pointer
    mov rsp, [SMP_VAR_ADDR + 8]

    # Clear RFLAGS.IF only, preserve architectural default state.
    # pushfq reads RFLAGS; AND clears IF (bit 9); popfq restores.
    pushfq
    pop rax
    btr rax, 9             # IF = 0
    push rax
    popfq

    # Jump to smp_main (Rust)
    # NOTE: must use movabs (64-bit absolute), NOT RIP-relative lea.
    # The trampoline is copied to 0x8000, so RIP-relative would compute
    # an offset based on the wrong base address (the linker's original
    # placement in the kernel .text section, not 0x8000).
    movabs rax, offset smp_main
    jmp rax

.align 8
smp_trampoline_end:
"#
);

unsafe extern "C" {
    fn smp_trampoline();
    fn smp_trampoline_end();
}

/// Busy-wait for the given number of microseconds (very rough).
fn udelay(us: u32) {
    for _ in 0..us {
        io_wait();
    }
}

/// Identity-map the trampoline physical pages so the AP can execute the
/// trampoline code in real mode / protected mode before paging is enabled.
fn ensure_identity_mapping(phys_start: u64, length: usize) -> Result<(), &'static str> {
    let start = phys_start & !0xFFFu64;
    let end = (phys_start + length as u64 + 0xFFF) & !0xFFFu64;
    let flags = PageTableFlags::PRESENT;

    let mut addr = start;
    while addr < end {
        let virt = VirtAddr::new(addr);
        if let Some(mapped) = crate::memory::paging::translate(virt) {
            if mapped.as_u64() != addr {
                return Err("SMP: trampoline identity map collision");
            }
            crate::memory::paging::set_trampoline_execution(addr, true)?;
        } else {
            let page = Page::<Size4KiB>::containing_address(virt);
            let frame = PhysFrame::<Size4KiB>::containing_address(PhysAddr::new(addr));
            crate::memory::paging::map_page(page, frame, flags)?;
        }
        addr += 0x1000;
    }
    Ok(())
}

/// Copy the trampoline to physical address 0x8000 and write the data area.
///
/// Data area layout (at smp_trampoline_end):
///   +0: CR3 (PML4 physical address)
///   +8: RSP (kernel stack top virtual address)
///
/// After writing, performs WBINVD to flush the cache hierarchy to RAM.
/// This is essential on real hardware: the BSP writes the trampoline via
/// HHDM (WB cacheable), but the AP boots in real mode where the effective
/// memory type is determined by MTRRs. If MTRRs mark the region as UC, or
/// if platform firmware does not guarantee cache coherency, the AP would
/// read stale data from RAM without this flush.
fn copy_trampoline(cr3_phys: u64, stack_top_virt: u64) -> Result<(), &'static str> {
    let tramp_len = (smp_trampoline_end as *const u8 as usize)
        .saturating_sub(smp_trampoline as *const u8 as usize);

    if tramp_len + 16 > 4096 {
        return Err("SMP: trampoline exceeds its reserved page");
    }
    ensure_identity_mapping(TRAMPOLINE_PHYS_ADDR, tramp_len + 16)?;

    let tramp_virt = memory::phys_to_virt(TRAMPOLINE_PHYS_ADDR) as *mut u8;

    // SAFETY: trampoline destination is mapped and writable in HHDM.
    unsafe {
        core::ptr::copy_nonoverlapping(smp_trampoline as *const u8, tramp_virt, tramp_len);
        let data = tramp_virt.add(tramp_len) as *mut u64;
        // +0: CR3
        core::ptr::write_volatile(data, cr3_phys);
        // +8: RSP (stack top virtual address)
        core::ptr::write_volatile(data.add(1), stack_top_virt);

        // SFENCE + WBINVD: ensure all stores reach RAM before the AP starts.
        core::arch::asm!("sfence");
        core::arch::asm!("wbinvd");
    }
    Ok(())
}

/// Wait for ICR delivery to complete.
///
/// In x2APIC mode, the ICR MSR write is synchronous : the CPU blocks until
/// the IPI is dispatched, so there is nothing to wait for.  In xAPIC mode,
/// we poll the delivery-pending bit (ICR_LOW bit 12) until it clears.
fn wait_delivery() {
    // x2APIC: MSR write is synchronous, no pending bit to poll.
    if apic::is_x2apic_enabled() {
        return;
    }

    const DELIVERY_PENDING: u32 = 1 << 12;
    for i in 0..1_000_000 {
        let val = unsafe { apic::read_reg(apic::REG_ICR_LOW) };
        if val & DELIVERY_PENDING == 0 {
            return;
        }
        if i > 0 && i % 200_000 == 0 {
            crate::serial_println!(
                "[smp] wait_delivery: still pending (iter={}, icr={:#x})",
                i,
                val,
            );
        }
        core::hint::spin_loop();
    }
    let final_val = unsafe { apic::read_reg(apic::REG_ICR_LOW) };
    crate::serial_println!("[smp] wait_delivery: TIMEOUT, final icr={:#x}", final_val,);
    log::warn!("SMP: IPI delivery timeout");
}

/// Send an IPI and wait for delivery.
fn send_ipi(apic_id: u32, value: u32) {
    apic::send_ipi_raw(apic_id, value);
    wait_delivery();
}

/// Send INIT + SIPI×2 to an AP per Intel SDM Volume 3, Section 10.6.7.1.
///
/// The Intel SDM recommends sending SIPI twice (200 µs apart) so that if
/// the first SIPI is missed, the second one still catches the AP. Some
/// platforms (Redox-OS) succeed with a single SIPI, but sending two is
/// more robust and matches Linux/FreeBSD behaviour.
///
///   Step 1: INIT level-assert (0xC500) : wait ≥10 ms
///   Step 2: INIT level-de-assert (0x8500) : wait ≥200 µs
///   Step 3: SIPI (0x4608, vector=0x08 => trampoline at 0x8000) : 200 µs
///   Step 4: SIPI again : 200 µs
///
/// ICR values (xAPIC format, delivery-mode bit-fields):
///   INIT:  delivery=INIT(101), level=1(assert), trigger=1(level)
///   Deassert: delivery=INIT(101), level=0(de-assert), trigger=1(level)
///   SIPI:  delivery=STARTUP(110), level=1, vector=0x8
fn send_init_sipi(apic_id: u32) {
    crate::serial_println!("[smp] send_init_sipi: apic_id={}", apic_id);

    // Step 1: INIT level-assert
    crate::serial_println!("[smp]   INIT assert -> {}", apic_id);
    send_ipi(apic_id, 0xC500);
    udelay(10_000);

    // Step 2: INIT level-de-assert
    crate::serial_println!("[smp]   INIT de-assert -> {}", apic_id);
    send_ipi(apic_id, 0x8500);
    udelay(200);

    // Step 3: SIPI #1
    crate::serial_println!("[smp]   SIPI (1/2) -> {} (vector=0x8)", apic_id);
    send_ipi(apic_id, 0x0608);
    udelay(200);

    // Step 4: SIPI #2  (SDM recommends two SIPIs)
    crate::serial_println!("[smp]   SIPI (2/2) -> {} (vector=0x8)", apic_id);
    send_ipi(apic_id, 0x0608);
    udelay(200);

    crate::serial_println!("[smp]   INIT+SIPI complete for {}", apic_id);
}

/// Broadcast a halt command to all other CPUs.
///
/// Used during panic to stop the system and prevent log corruption.
pub fn broadcast_panic_halt() {
    if !apic::is_initialized() {
        return;
    }
    // SMI delivery mode (0b100 << 8), all-excluding-self shorthand (0b11 << 18).
    // Level bit = 0 (reserved for SMI mode per Intel SDM Vol. 3A Table 10-1).
    let icr_low = (0b11 << 18) | (0b100 << 8);
    apic::send_ipi_raw(0, icr_low);
}

/// Wait at a synchronization barrier until all expected CPUs arrive.
fn rendezvous_barrier() {
    let expected = BARRIER_TARGET.load(Ordering::Acquire);
    SYNC_BARRIER.fetch_add(1, Ordering::AcqRel);
    while SYNC_BARRIER.load(Ordering::Acquire) < expected {
        core::hint::spin_loop();
    }
    // BSP may have revoked the trampoline before publishing BARRIER_TARGET.
    // Each AP drops its formerly executable translation before normal work.
    unsafe {
        core::arch::asm!(
            "invlpg [{}]", in(reg) TRAMPOLINE_PHYS_ADDR,
            options(nostack, preserves_flags),
        );
    }
}

/// Boot Application Processors.
///
/// AP kernel stack frames are allocated from the buddy allocator during this
/// function and intentionally leaked (never freed). The allocation happens
/// once at boot while only the BSP is online, so there is no lock contention
/// and no risk of heap exhaustion under concurrent pressure.
///
/// If CPU hotplug is added in the future, stack allocation must be moved to
/// a sleepable context (e.g. `Mutex`-protected) to avoid heap allocation
/// under any spinlock that could be held during hot-add.
pub fn init() -> Result<usize, &'static str> {
    crate::serial_println!("[smp] init: entering SMP initialization");
    if !apic::is_initialized() {
        crate::serial_println!("[smp] init: ERROR - APIC not initialized");
        return Err("APIC not initialized");
    }

    BOOTED_CORES.store(1, Ordering::Release);
    SYNC_BARRIER.store(0, Ordering::Release);
    BARRIER_TARGET.store(0, Ordering::Release);

    let madt_info = madt::parse_madt().ok_or("MADT not available")?;
    let bsp_apic_id = apic::lapic_id();

    crate::serial_println!(
        "[smp] init: MADT parsed, {} local APICs, BSP apic_id={}",
        madt_info.local_apic_count,
        bsp_apic_id,
    );

    if madt_info.local_apic_count <= 1 {
        crate::serial_println!("[smp] init: single CPU system, returning");
        log::info!("SMP: single CPU system");
        return Ok(1);
    }

    let mut max_apic_id: usize = 0;
    for i in 0..madt_info.local_apic_count {
        if let Some(ref entry) = madt_info.local_apics[i] {
            crate::serial_println!(
                "[smp]   APIC[{}]: id={} proc={} flags={:#x} {}",
                i,
                entry.apic_id,
                entry.processor,
                entry.flags,
                if entry.flags & 1 == 0 {
                    "(DISABLED)"
                } else {
                    ""
                },
            );
            max_apic_id = max_apic_id.max(entry.apic_id as usize);
        }
    }

    let mut stack_tops: Vec<u64> = vec![0; max_apic_id + 1];
    let cr3_phys = crate::memory::paging::kernel_l4_phys().as_u64();
    let mut targets: Vec<u32> = Vec::new();
    let mut expected: usize = 1;

    for i in 0..madt_info.local_apic_count {
        let Some(ref entry) = madt_info.local_apics[i] else {
            continue;
        };

        let apic_id = entry.apic_id as u32;
        if apic_id == bsp_apic_id {
            continue;
        }

        // Allocate AP kernel stack from the buddy allocator.
        let stack_size = crate::process::task::Task::DEFAULT_STACK_SIZE;
        let pages = (stack_size + 4095) / 4096;
        let order = pages.next_power_of_two().trailing_zeros() as u8;
        let frame = crate::sync::with_irqs_disabled(|token| {
            crate::memory::allocate_phys_contiguous(token, order)
        })
        .map_err(|_| "SMP: failed to allocate AP stack from buddy")?;
        let stack_phys = frame.start_address.as_u64();
        let stack_virt = crate::memory::phys_to_virt(stack_phys);

        // Zero the stack.
        unsafe { core::ptr::write_bytes(stack_virt as *mut u8, 0, stack_size) };

        // Stack grows downward: top = base_virt + size.
        let stack_top = stack_virt.saturating_add(stack_size as u64);

        if apic_id as usize >= stack_tops.len() {
            log::warn!("SMP: APIC id {} out of stack array range", apic_id);
            continue;
        }

        stack_tops[apic_id as usize] = stack_top;

        let cpu_index =
            percpu::register_cpu(apic_id).ok_or("SMP: exceeded MAX_CPUS for per-CPU data")?;
        percpu::set_kernel_stack_top(cpu_index, stack_top);

        targets.push(apic_id);
        expected += 1;
    }

    // Copy trampoline and write data area (CR3 + first AP's stack top).
    // All APs share the same CR3; each gets its own stack from stack_tops.
    // We write the first target's stack top; subsequent APs will use their
    // own stack (set up in smp_main via percpu::kernel_stack_top).
    let first_stack_top = targets
        .first()
        .and_then(|id| stack_tops.get(*id as usize))
        .copied()
        .unwrap_or(0);
    copy_trampoline(cr3_phys, first_stack_top)?;
    crate::serial_println!(
        "[smp] init: trampoline at {:#x}, cr3={:#x}, stack={:#x}",
        TRAMPOLINE_PHYS_ADDR,
        cr3_phys,
        first_stack_top,
    );

    // Send INIT + single SIPI to each AP, one at a time.
    // Critical: each AP must consume the trampoline RSP before we overwrite
    // it for the next AP. We wait for AP_REACHED_RUST ack after each SIPI.
    crate::serial_println!(
        "[smp] init: sending INIT+SIPI to {} APs (sequential)",
        targets.len(),
    );
    for apic_id in &targets {
        AP_REACHED_RUST.store(usize::MAX, Ordering::Release);

        // Write this AP's stack pointer into the trampoline data area.
        if let Some(stack_top) = stack_tops.get(*apic_id as usize) {
            let tramp_len = (smp_trampoline_end as *const u8 as usize)
                .saturating_sub(smp_trampoline as *const u8 as usize);
            let data = memory::phys_to_virt(TRAMPOLINE_PHYS_ADDR) as *mut u64;
            unsafe {
                // +8: RSP
                core::ptr::write_volatile(data.add(tramp_len / 8 + 1), *stack_top);
            }
        }

        send_init_sipi(*apic_id);

        // Wait for this AP to reach Rust and consume the stack pointer
        // before we modify the slot for the next AP.
        let mut spin: u64 = 0;
        const ACK_TIMEOUT: u64 = 500_000_000;
        while AP_REACHED_RUST.load(Ordering::Acquire) != *apic_id as usize && spin < ACK_TIMEOUT {
            core::hint::spin_loop();
            spin = spin.saturating_add(1);
        }
        if spin >= ACK_TIMEOUT {
            crate::serial_println!(
                "[smp] init: WARNING AP {} did not ACK within timeout",
                apic_id
            );
        } else {
            crate::serial_println!("[smp] init: AP {} ACKed (rust reached)", apic_id);
        }
    }

    // Wait for APs to come online (they increment BOOTED_CORES in smp_main).
    let mut spins: u64 = 0;
    const MAX_SPINS: u64 = 200_000_000;
    crate::serial_println!("[smp] init: waiting for APs (expected={})...", expected,);
    while BOOTED_CORES.load(Ordering::Acquire) < expected && spins < MAX_SPINS {
        if spins > 0 && spins % 50_000_000 == 0 {
            crate::serial_println!(
                "[smp] init: waiting... online={} expected={} spins={}",
                BOOTED_CORES.load(Ordering::Acquire),
                expected,
                spins,
            );
        }
        core::hint::spin_loop();
        spins = spins.saturating_add(1);
    }
    let online = BOOTED_CORES.load(Ordering::Acquire);
    crate::serial_println!(
        "[smp] init: AP wait done: online={} expected={} spins={}",
        online,
        expected,
        spins,
    );
    if online < expected {
        crate::serial_println!(
            "[smp] init: WARNING only {}/{} APs online",
            online,
            expected,
        );
        log::warn!(
            "SMP: timeout waiting APs (online={} expected={}), continuing",
            online,
            expected
        );
    }

    log::info!("SMP: {} cores online (expected {})", online, expected);

    if online == expected {
        if let Err(error) =
            crate::memory::paging::set_trampoline_execution(TRAMPOLINE_PHYS_ADDR, false)
        {
            // A legacy boot path can provide a huge identity leaf. Never revoke
            // execution of unrelated addresses when retiring the trampoline.
            log::warn!("SMP: cannot retire trampoline execution: {}", error);
        }
    }

    // Publish barrier target so APs can proceed.
    BARRIER_TARGET.store(online, Ordering::Release);
    rendezvous_barrier();

    Ok(online)
}

/// First Rust function executed on APs after the trampoline.
///
/// The trampoline enables paging with the kernel's PML4, sets up the stack,
/// and jumps here. All virtual addresses are valid at this point.
#[unsafe(no_mangle)]
pub extern "C" fn smp_main() -> ! {
    // Read APIC ID first : needed to find our per-CPU state.
    // Use raw port output since serial mutex isn't initialized yet.
    let apic_id: u32;
    {
        use core::fmt::Write;
        let mut port = unsafe { uart_16550::SerialPort::new(0x3F8) };
        port.init();
        let _ = port.write_fmt(format_args!("[smp][ap] entered smp_main\n"));
        // CPUID leaf 1 EBX[31:24] = initial APIC ID.
        // rbx is LLVM-reserved, so save/restore it manually.
        let apic_id_raw: u32;
        unsafe {
            core::arch::asm!(
                "push rbx",
                "cpuid",
                "mov {val:e}, ebx",
                "pop rbx",
                val = out(reg) apic_id_raw,
                in("eax") 1u32,
                out("ecx") _,
                out("edx") _,
            );
        }
        apic_id = apic_id_raw >> 24;
        let _ = port.write_fmt(format_args!("[smp][ap] cpuid apic_id={}\n", apic_id));
    }

    // Signal BSP that this AP has reached Rust and consumed the trampoline RSP.
    // Must happen immediately : before any per-CPU init that might fail.
    AP_REACHED_RUST.store(apic_id as usize, Ordering::Release);

    let cpu_index = match percpu::cpu_index_by_apic(apic_id) {
        Some(idx) => idx,
        None => {
            {
                use core::fmt::Write;
                let mut port = unsafe { uart_16550::SerialPort::new(0x3F8) };
                let _ = port.write_fmt(format_args!(
                    "[smp][ap] ERROR: apic_id={} not registered, halting\n",
                    apic_id
                ));
            }
            loop {
                core::hint::spin_loop();
            }
        }
    };

    // ── Phase 1: per-CPU invariants (before any interrupt can fire) ──
    // Order matters: GS base must be set before any percpu access,
    // GDT before IDT (CS selector must be valid), TSS before any
    // ring-0 → ring-3 transition.

    crate::arch::x86_64::percpu::init_gs_base(cpu_index);

    crate::arch::x86_64::tss::init_cpu(cpu_index);
    crate::arch::x86_64::gdt::init_cpu(cpu_index);

    crate::arch::x86_64::syscall::init();
    crate::arch::x86_64::init_cpu_extensions();

    // Verify XCR0 was programmed correctly on this AP.
    if crate::arch::x86_64::cpuid::host_uses_xsave() {
        let expected = crate::arch::x86_64::cpuid::host_default_xcr0_fast();
        let actual = crate::arch::x86_64::xgetbv(0);
        if actual != expected {
            use core::fmt::Write;
            let mut port = unsafe { uart_16550::SerialPort::new(0x3F8) };
            let _ = port.write_fmt(format_args!(
                "[smp][ap] WARNING: XCR0 mismatch! expected={:#x} actual={:#x}\n",
                expected, actual
            ));
        }
    }

    if let Some(stack_top) = percpu::kernel_stack_top(cpu_index) {
        crate::arch::x86_64::tss::set_kernel_stack_for(cpu_index, x86_64::VirtAddr::new(stack_top));
    }

    // ── Phase 2: interrupt infrastructure ──
    // IDT is loaded LAST after all per-CPU state is established.
    // This prevents any exception/IPI from firing with incomplete state.
    idt::load();

    // Re-initialize Local APIC for this core.
    apic::init_ap();

    // ── Phase 3: signal online ──
    let _ = percpu::mark_online_by_apic(apic_id);
    BOOTED_CORES.fetch_add(1, Ordering::Release);

    {
        use core::fmt::Write;
        let mut port = unsafe { uart_16550::SerialPort::new(0x3F8) };
        let _ = port.write_fmt(format_args!(
            "[smp][ap] cpu_index={} online, waiting for barrier\n",
            cpu_index
        ));
    }

    // AP spins until BSP publishes the barrier target.
    while BARRIER_TARGET.load(Ordering::Acquire) == 0 {
        core::hint::spin_loop();
    }
    rendezvous_barrier();

    crate::serial_println!("[trace][ap] cpu_index={} entering scheduler", cpu_index);

    // Wait until BSP has finished scheduler initialization.
    while !AP_SCHED_GATE_OPEN.load(Ordering::Acquire) {
        core::hint::spin_loop();
    }

    // Start APIC timer on this CPU.
    timer::start_apic_timer_cached();

    // Start per-CPU scheduler (never returns).
    crate::process::scheduler::schedule_on_cpu(cpu_index)
}

/// Return the number of online CPUs.
pub fn cpu_count() -> usize {
    BOOTED_CORES.load(Ordering::Acquire)
}

/// Allow APs to start their local timer and enter the scheduler.
pub fn open_ap_scheduler_gate() {
    AP_SCHED_GATE_OPEN.store(true, Ordering::Release);
}