smp.rs

//! ARM64 SMP (Symmetric Multi-Processing) support.
//!
//! This module handles bringing up secondary CPUs on ARM64 using PSCI
//! (Power State Coordination Interface). All supported hypervisors
//! (QEMU, Parallels, VMware) implement PSCI via HVC calls.
//!
//! Flow:
//! 1. CPU 0 probes CPUs 1..MAX_CPUS via `release_cpu()` (PSCI CPU_ON)
//! 2. PSCI firmware starts each target CPU at `secondary_cpu_entry` (boot.S)
//! 3. boot.S sets up stack, MMU, and calls `secondary_cpu_entry_rust()`
//! 4. Rust entry initializes per-CPU data, GIC, timer, creates idle thread

use core::sync::atomic::{AtomicBool, AtomicU64, Ordering};

/// Maximum number of CPUs supported.
pub const MAX_CPUS: usize = 8;

/// PSCI function IDs (SMCCC compliant).
const PSCI_CPU_ON_64: u64 = 0xC400_0003;
const PSCI_CPU_ON_32: u64 = 0x8400_0003;

extern "C" {
    /// Physical address of secondary_cpu_entry, stored in .rodata by boot.S.
    /// We cannot reference secondary_cpu_entry directly from Rust because it lives
    /// in .text.boot (low physical memory) while Rust code is in high-half virtual
    /// memory — the ~1 TiB gap exceeds the ADRP relocation range (+/- 4 GiB).
    static SECONDARY_CPU_ENTRY_PHYS: u64;

    /// Pointer to SMP_UART_PHYS (which lives in .bss.boot, low physical memory).
    /// Stored in .rodata so Rust can reach it via ADRP, then we dereference to
    /// get the actual address of the variable and write through it.
    static SMP_UART_PHYS_PTR: u64;

    /// Pointers to SMP_TTBR0_PHYS / SMP_TTBR1_PHYS / SMP_MAIR_PHYS / SMP_TCR_PHYS
    /// variables (in .bss.boot). Secondary CPUs read these to get the correct
    /// page table addresses and MMU configuration.
    static SMP_TTBR0_PTR: u64;
    static SMP_TTBR1_PTR: u64;
    static SMP_MAIR_PTR: u64;
    static SMP_TCR_PTR: u64;

    /// Pointer to SMP_STACK_BASE_PHYS (in .bss.boot). CPU 0 writes the
    /// physical base address of the per-CPU stack region here before PSCI CPU_ON.
    /// On QEMU/Parallels: 0x4100_0000; on VMware: 0x8100_0000.
    static SMP_STACK_BASE_PTR: u64;
}

/// Write CPU 0's actual MMU configuration to .bss.boot variables so
/// secondary CPUs can replicate the exact same setup.
///
/// Stores TTBR0, TTBR1, MAIR_EL1, and TCR_EL1 values.
/// On QEMU, these come from boot.S's setup_mmu. On Parallels, they come
/// from the UEFI loader's page_tables module (different TCR, different TTBRs).
/// Secondary CPUs in boot.S read these to configure MMU identically to CPU 0.
///
/// Includes cache clean + DSB so values are visible to secondary CPUs
/// which start with MMU off (uncached reads from physical memory).
///
/// Must be called before the first `release_cpu()` call.
pub fn set_smp_ttbrs() {
    // Helper: write a u64 to a .bss.boot variable via .rodata pointer indirection,
    // then clean the cache line so uncached readers (secondary CPUs) see it.
    unsafe fn write_bss_boot_var(ptr_ro: &u64, value: u64) {
        let var_phys = core::ptr::read_volatile(ptr_ro);
        let var_virt = (var_phys + 0xFFFF_0000_0000_0000u64) as *mut u64;
        core::ptr::write_volatile(var_virt, value);
        core::arch::asm!(
            "dc cvac, {addr}",
            addr = in(reg) var_virt,
            options(nostack),
        );
    }

    unsafe {
        let ttbr0: u64;
        let ttbr1: u64;
        let mair: u64;
        let tcr: u64;
        core::arch::asm!(
            "mrs {}, ttbr0_el1",
            "mrs {}, ttbr1_el1",
            "mrs {}, mair_el1",
            "mrs {}, tcr_el1",
            out(reg) ttbr0,
            out(reg) ttbr1,
            out(reg) mair,
            out(reg) tcr,
            options(nomem, nostack),
        );

        write_bss_boot_var(&SMP_TTBR0_PTR, ttbr0);
        write_bss_boot_var(&SMP_TTBR1_PTR, ttbr1);
        write_bss_boot_var(&SMP_MAIR_PTR, mair);
        write_bss_boot_var(&SMP_TCR_PTR, tcr);

        // Single DSB to ensure all cache cleans complete
        core::arch::asm!(
            "dsb ish",
            options(nostack),
        );
    }
}

/// Set the UART physical address for secondary CPU boot debug output.
/// Must be called before `release_cpu()`.
///
/// Uses indirection through SMP_UART_PHYS_PTR because SMP_UART_PHYS lives in
/// .bss.boot (low physical memory) and direct ADRP from high-half Rust code
/// would overflow the +/-4GiB relocation range.
///
/// Includes cache clean + DSB so the value is visible to secondary CPUs
/// which start with MMU off (uncached reads from physical memory).
pub fn set_uart_phys(addr: u64) {
    unsafe {
        // SMP_UART_PHYS_PTR holds the physical address of SMP_UART_PHYS.
        // Add HHDM base to get the virtual address, then write through it.
        let phys = core::ptr::read_volatile(&SMP_UART_PHYS_PTR);
        let virt = phys + 0xFFFF_0000_0000_0000u64; // KERNEL_VIRT_BASE / HHDM
        let ptr = virt as *mut u64;
        core::ptr::write_volatile(ptr, addr);
        // Clean cache line to Point of Coherency so uncached reads see it
        core::arch::asm!(
            "dc cvac, {addr}",  // Clean by VA to PoC
            "dsb ish",          // Ensure completion
            addr = in(reg) ptr,
            options(nostack),
        );
    }
}

/// Set the physical base address of the per-CPU stack region.
/// Must be called before `release_cpu()`.
///
/// The stack base is `ram_base + 0x0100_0000` (16MB into RAM).
/// On QEMU/Parallels (ram at 0x40000000): 0x4100_0000.
/// On VMware (ram at 0x80000000): 0x8100_0000.
pub fn set_stack_base_phys(addr: u64) {
    unsafe {
        let phys = core::ptr::read_volatile(&SMP_STACK_BASE_PTR);
        let virt = phys + 0xFFFF_0000_0000_0000u64;
        let ptr = virt as *mut u64;
        core::ptr::write_volatile(ptr, addr);
        core::arch::asm!(
            "dc cvac, {addr}",
            "dsb ish",
            addr = in(reg) ptr,
            options(nostack),
        );
    }
}

/// Number of CPUs currently online (starts at 1 for the boot CPU).
static CPUS_ONLINE: AtomicU64 = AtomicU64::new(1);

/// Per-CPU online status flags.
static CPU_ONLINE: [AtomicBool; MAX_CPUS] = [
    AtomicBool::new(true),  // CPU 0 is online at boot
    AtomicBool::new(false),
    AtomicBool::new(false),
    AtomicBool::new(false),
    AtomicBool::new(false),
    AtomicBool::new(false),
    AtomicBool::new(false),
    AtomicBool::new(false),
];

/// Issue a PSCI CPU_ON call via HVC to start a secondary CPU.
///
/// Arguments:
/// - `target_cpu`: MPIDR of the target CPU (Aff0 = cpu_id for QEMU virt)
/// - `entry_point`: Physical address where the CPU starts executing
/// - `context_id`: Value passed in x0 to the new CPU (we use cpu_id)
///
/// Returns PSCI status: 0 = SUCCESS, negative = error.
fn psci_cpu_on(target_cpu: u64, entry_point: u64, context_id: u64) -> i64 {
    let ret: i64;
    unsafe {
        core::arch::asm!(
            "hvc #0",
            inout("x0") PSCI_CPU_ON_64 => ret,
            in("x1") target_cpu,
            in("x2") entry_point,
            in("x3") context_id,
            options(nomem, nostack),
        );
    }
    ret
}

/// PSCI CPU_ON with 32-bit function ID via HVC.
fn psci_cpu_on_32(target_cpu: u64, entry_point: u64, context_id: u64) -> i64 {
    let ret: i64;
    unsafe {
        core::arch::asm!(
            "hvc #0",
            inout("x0") PSCI_CPU_ON_32 => ret,
            in("x1") target_cpu,
            in("x2") entry_point,
            in("x3") context_id,
            options(nomem, nostack),
        );
    }
    ret
}

/// PSCI CPU_ON with 64-bit function ID via SMC (EL3 firmware conduit).
fn psci_cpu_on_smc(target_cpu: u64, entry_point: u64, context_id: u64) -> i64 {
    let ret: i64;
    unsafe {
        core::arch::asm!(
            "smc #0",
            inout("x0") PSCI_CPU_ON_64 => ret,
            in("x1") target_cpu,
            in("x2") entry_point,
            in("x3") context_id,
            options(nomem, nostack),
        );
    }
    ret
}

/// Release a secondary CPU using PSCI CPU_ON.
///
/// The CPU will start executing at `secondary_cpu_entry` in boot.S,
/// which sets up the stack and MMU, then calls `secondary_cpu_entry_rust(cpu_id)`.
///
/// Returns the PSCI result: 0 = success, negative = error
/// (-2 = INVALID_PARAMS, -9 = NOT_PRESENT, etc.)
pub fn release_cpu(cpu_id: usize) -> i64 {
    if cpu_id == 0 || cpu_id >= MAX_CPUS {
        return -2; // INVALID_PARAMS
    }

    // Get the physical address of the secondary entry point in boot.S
    let entry_phys = unsafe { core::ptr::read_volatile(&SECONDARY_CPU_ENTRY_PHYS) };

    // MPIDR: Aff0 = cpu_id, all other affinity fields = 0
    // This is the standard layout for ARM virt machines (QEMU, Parallels, VMware)
    let target_mpidr = cpu_id as u64;

    // Context ID: pass cpu_id so the new CPU knows who it is
    let context_id = cpu_id as u64;

    // Try 64-bit PSCI CPU_ON via HVC first (standard for ARM64 hypervisors)
    let mut ret = psci_cpu_on(target_mpidr, entry_phys, context_id);

    // If 64-bit failed, try 32-bit function ID (some hypervisors only support this)
    if ret != 0 {
        crate::serial_println!("[smp] CPU {}: HVC64 failed (ret={}), trying HVC32...", cpu_id, ret);
        ret = psci_cpu_on_32(target_mpidr, entry_phys, context_id);
    }

    // Note: SMC conduit not attempted — on VMware (EL1 guest, no EL3),
    // SMC would trap to EL2 and likely fault. HVC is the correct conduit.

    if ret != 0 {
        crate::serial_println!("[smp] PSCI CPU_ON failed for CPU {}: ret={} (MPIDR={:#x} entry={:#x})",
            cpu_id, ret, target_mpidr, entry_phys);
    }

    ret
}

/// Get the number of CPUs currently online.
pub fn cpus_online() -> u64 {
    CPUS_ONLINE.load(Ordering::Acquire)
}

/// Check if a specific CPU is online.
#[allow(dead_code)]
pub fn is_cpu_online(cpu_id: usize) -> bool {
    if cpu_id >= MAX_CPUS {
        return false;
    }
    CPU_ONLINE[cpu_id].load(Ordering::Acquire)
}

/// Raw UART output for secondary CPUs (no locks, no allocations).
#[inline(always)]
fn raw_uart_char(c: u8) {
    let addr = crate::platform_config::uart_virt() as *mut u8;
    unsafe {
        core::ptr::write_volatile(addr, c);
    }
}

/// Secondary CPU entry point.
///
/// Called from boot.S after PSCI CPU_ON starts the CPU and boot.S
/// sets up the stack, MMU, and exception vectors.
///
/// Initializes per-CPU data, GIC CPU interface, timer, and creates
/// this CPU's idle thread. After initialization, enters the idle loop
/// and participates in scheduling.
#[no_mangle]
pub extern "C" fn secondary_cpu_entry_rust(cpu_id: u64) -> ! {
    // Emit raw UART character to signal this CPU is alive
    raw_uart_char(b'0' + cpu_id as u8);

    // Initialize per-CPU data (sets TPIDR_EL1 for this CPU)
    crate::per_cpu_aarch64::init_cpu(cpu_id as usize);

    // Set kernel stack top for this CPU.
    // boot.S sets SP to SMP_STACK_BASE_PHYS + (cpu_id+1)*0x200000 (physical),
    // then adds KERNEL_VIRT_BASE after enabling MMU.
    // This value is critical: when a user thread runs on this CPU and an
    // exception occurs, the kernel needs to switch to this stack.
    let stack_base = super::constants::percpu_stack_region_base();
    const STACK_SIZE: u64 = 0x20_0000; // 2MB per CPU
    let kernel_stack_top = stack_base + (cpu_id + 1) * STACK_SIZE;
    crate::per_cpu_aarch64::set_kernel_stack_top(kernel_stack_top);

    // Initialize GIC CPU interface (GICC registers are banked per-CPU)
    super::gic::init_cpu_interface_secondary();

    // Initialize timer for this CPU (arm virtual timer, enable PPI 27)
    super::timer_interrupt::init_secondary();

    // Create and register this CPU's idle thread with the scheduler.
    // This must happen before enabling interrupts — the scheduler needs
    // an idle thread for this CPU before timer interrupts fire.
    create_and_register_idle_thread(cpu_id as usize);

    // Enable interrupts so this CPU can handle timer ticks
    unsafe {
        super::cpu::enable_interrupts();
    }

    // Mark this CPU as online (after all init is complete)
    if (cpu_id as usize) < MAX_CPUS {
        CPU_ONLINE[cpu_id as usize].store(true, Ordering::Release);
    }
    CPUS_ONLINE.fetch_add(1, Ordering::Release);

    // Idle loop — wait for interrupts, handle timer, participate in scheduling
    loop {
        unsafe {
            core::arch::asm!("wfi", options(nomem, nostack));
        }
    }
}

/// Create an idle thread for a secondary CPU and register it with the scheduler.
fn create_and_register_idle_thread(cpu_id: usize) {
    use alloc::boxed::Box;
    use alloc::format;
    use crate::task::thread::{Thread, ThreadState, ThreadPrivilege};
    use crate::memory::arch_stub::VirtAddr;

    // Boot stack addresses — must match boot.S layout.
    let stack_base = super::constants::percpu_stack_region_base();
    const STACK_SIZE: u64 = 0x20_0000; // 2MB per CPU
    let boot_stack_top = VirtAddr::new(stack_base + ((cpu_id as u64) + 1) * STACK_SIZE);
    let boot_stack_bottom = VirtAddr::new(stack_base + (cpu_id as u64) * STACK_SIZE);
    let dummy_tls = VirtAddr::zero();

    let mut idle_task = Box::new(Thread::new(
        format!("swapper/{}", cpu_id),
        idle_thread_fn,
        boot_stack_top,
        boot_stack_bottom,
        dummy_tls,
        ThreadPrivilege::Kernel,
    ));

    // CRITICAL: Set kernel_stack_top to this CPU's boot stack. Without this,
    // setup_idle_return_arm64 falls back to the per-CPU kernel_stack_top from
    // the last dispatched thread, causing the idle loop to run on that thread's
    // kernel stack and corrupt its SVC frame.
    idle_task.kernel_stack_top = Some(boot_stack_top);

    // Mark as running and already started (this CPU is already executing)
    idle_task.state = ThreadState::Running;
    idle_task.has_started = true;

    // Set per-CPU current thread pointer
    let idle_task_ptr = &*idle_task as *const _ as *mut crate::task::thread::Thread;
    crate::per_cpu_aarch64::set_current_thread(idle_task_ptr);

    // Register with the global scheduler
    crate::task::scheduler::register_cpu_idle_thread(cpu_id, idle_task);
}

/// Idle thread function for secondary CPUs.
fn idle_thread_fn() {
    loop {
        unsafe {
            core::arch::asm!("wfi", options(nomem, nostack));
        }
    }
}