ax-hal 0.5.13

//! Architecture-specific KCOV trampolines.
//!
//! These are the `__sanitizer_cov_trace_pc` entry points that the compiler
//! injects into every instrumented basic block. The recursion guard lives
//! in the per-CPU `.percpu` section so that tracing on one CPU never blocks
//! coverage collection on another under SMP.  The guard is checked and set
//! entirely in naked asm *before* any C-ABI call so that LLVM KCOV
//! instrumentation of the callee cannot re-enter this function (infinite
//! recursion guard).

#![cfg(feature = "starry-kcov")]

use core::arch::naked_asm;

use ax_percpu::def_percpu;

/// Per-CPU recursion guard: nonzero while inside `kcov_trace_pc_impl` on this
/// CPU.  Lives in the `.percpu` section so each CPU has its own copy.
///
/// Checked and set by the architecture-specific naked trampolines before
/// calling `kcov_trace_pc_impl`.  The underlying static is named
/// `__PERCPU_IN_KCOV_TRACE` (generated by `#[def_percpu]`) and is referenced
/// directly from asm via `sym __PERCPU_IN_KCOV_TRACE`.
#[def_percpu]
static IN_KCOV_TRACE: u8 = 0;

/// Safety gate: blocks the trace handler until a thread explicitly enables
/// KCOV via the `KCOV_ENABLE` ioctl (which happens from userspace, long
/// after boot).  Without this check, instrumented edges during early boot
/// call `kcov_trace_pc_impl` → `ax_task::current()` before the scheduler
/// has set the per-CPU task pointer (that happens at the end of
/// `primary_init`, well after the first instrumented code runs).
///
/// `ax_task::current()` then panics with "current task is uninitialized".
/// The panic handler tries `ax_println!`, but UART is not ready yet either:
/// the platform init runs `init_trap` before `console::init_early`, so the
/// UART `LazyInit` is still empty.  That is a double-panic, which rustc's
/// abort guard turns into a `ud2` → #UD (unhandled) → #DF (unhandled) →
/// triple fault → CPU reset → Seabios "Booting from ROM…" → boot again →
/// same crash → infinite reset loop.
///
/// Setting this flag to 1 in `KCOV_ENABLE` is safe because by the time
/// userspace can issue ioctls the boot is complete and all affected
/// subsystems are live.
///
/// This makes sure initialization of per-cpu is before IN_KCOV_TRACE is resolved.
#[used]
pub static mut KCOV_GLOBAL_GATE: u8 = 0;

// Provided by the kernel crate (`starry_kernel::kcov`).
unsafe extern "C" {
    fn kcov_trace_pc_impl(pc: u64);
}

// ---------------------------------------------------------------------------
// x86_64  –  return address at [rsp]  (System V AMD64 ABI)
// per-CPU via gs-segment base
// ---------------------------------------------------------------------------

#[cfg(target_arch = "x86_64")]
#[unsafe(no_mangle)]
#[unsafe(naked)]
/// # Safety
///
/// Called by the compiler as a coverage instrumentation hook. The standard
/// x86_64 System V C ABI applies: the return address is on the stack.
/// Must only be called by instrumented code when KCOV is enabled.
pub unsafe extern "C" fn __sanitizer_cov_trace_pc() {
    naked_asm!(
        "cmp byte ptr [rip + {global}], 0",
        "je 2f",
        "cmp byte ptr gs:[offset {guard}], 0",
        "jne 1f",
        "mov byte ptr gs:[offset {guard}], 1",
        "mov rdi, [rsp]",
        "call {impl}",
        "mov byte ptr gs:[offset {guard}], 0",
        "1:",
        "2:",
        "ret",
        global = sym KCOV_GLOBAL_GATE,
        guard = sym __PERCPU_IN_KCOV_TRACE,
        impl = sym kcov_trace_pc_impl,
    );
}

// ---------------------------------------------------------------------------
// aarch64  –  return address in x30 (LR)
// per-CPU via TPIDR_EL1 + symbol VMA
// ---------------------------------------------------------------------------

#[cfg(target_arch = "aarch64")]
#[unsafe(no_mangle)]
#[unsafe(naked)]
/// # Safety
///
/// Called by the compiler as a coverage instrumentation hook. The standard
/// AAPCS64 ABI applies: `x30` holds the return address. Must only be called
/// by instrumented code when KCOV is enabled.
pub unsafe extern "C" fn __sanitizer_cov_trace_pc() {
    naked_asm!(
        "adrp x16, {global}",
        "ldrb w17, [x16, #:lo12:{global}]",
        "cbz w17, 2f",
        "mrs x16, tpidr_el1",
        "movz x17, #:abs_g0_nc:{guard}",
        "add x16, x16, x17",
        "ldrb w17, [x16]",
        "cbnz w17, 1f",
        "mov w17, #1",
        "strb w17, [x16]",
        "str x30, [sp, #-16]!",
        "mov x0, x30",
        "bl {impl}",
        "ldr x30, [sp], #16",
        // bl {impl} clobbered x16/x17 — recalculate per-CPU address
        "mrs x16, tpidr_el1",
        "movz x17, #:abs_g0_nc:{guard}",
        "add x16, x16, x17",
        "strb wzr, [x16]",
        "1:",
        "2:",
        "ret",
        global = sym KCOV_GLOBAL_GATE,
        guard = sym __PERCPU_IN_KCOV_TRACE,
        impl = sym kcov_trace_pc_impl,
    );
}

// ---------------------------------------------------------------------------
// riscv64  –  return address in ra (x1)
// per-CPU via gp + symbol offset
// ---------------------------------------------------------------------------

#[cfg(target_arch = "riscv64")]
#[unsafe(no_mangle)]
#[unsafe(naked)]
/// # Safety
///
/// Called by the compiler as a coverage instrumentation hook. The standard
/// RISC-V calling convention applies: `ra` holds the return address. Must
/// only be called by instrumented code when KCOV is enabled.
pub unsafe extern "C" fn __sanitizer_cov_trace_pc() {
    naked_asm!(
        "1: auipc t0, %pcrel_hi({global})",
        "addi t0, t0, %pcrel_lo(1b)",
        "lb t1, 0(t0)",
        "beqz t1, 2f",
        "lui t0, %hi({guard})",
        "add t0, t0, gp",
        "addi t0, t0, %lo({guard})",
        "lb t1, 0(t0)",
        "bnez t1, 1f",
        "li t1, 1",
        "sb t1, 0(t0)",
        "addi sp, sp, -16",
        "sd ra, 0(sp)",
        "mv a0, ra",
        "call {impl}",
        "ld ra, 0(sp)",
        "addi sp, sp, 16",
        "lui t0, %hi({guard})",
        "add t0, t0, gp",
        "addi t0, t0, %lo({guard})",
        "sb zero, 0(t0)",
        "1:",
        "2:",
        "ret",
        global = sym KCOV_GLOBAL_GATE,
        guard = sym __PERCPU_IN_KCOV_TRACE,
        impl = sym kcov_trace_pc_impl,
    );
}

// ---------------------------------------------------------------------------
// loongarch64  –  return address in ra ($r1)
// per-CPU via $r21 + symbol offset
// ---------------------------------------------------------------------------

#[cfg(target_arch = "loongarch64")]
#[unsafe(no_mangle)]
#[unsafe(naked)]
/// # Safety
///
/// Called by the compiler as a coverage instrumentation hook. The standard
/// LoongArch calling convention applies: `ra` holds the return address. Must
/// only be called by instrumented code when KCOV is enabled.
pub unsafe extern "C" fn __sanitizer_cov_trace_pc() {
    naked_asm!(
        "pcalau12i $t0, %pc_hi20({global})",
        "addi.d $t0, $t0, %pc_lo12({global})",
        "ld.b $t1, $t0, 0",
        "beqz $t1, 2f",
        "lu12i.w $t0, %abs_hi20({guard})",
        "ori $t0, $t0, %abs_lo12({guard})",
        "add.d $t0, $t0, $r21",
        "ld.b $t1, $t0, 0",
        "bnez $t1, 1f",
        "ori $t1, $zero, 1",
        "st.b $t1, $t0, 0",
        "addi.d $sp, $sp, -16",
        "st.d $ra, $sp, 0",
        "ori $a0, $ra, 0",
        "bl {impl}",
        "ld.d $ra, $sp, 0",
        "addi.d $sp, $sp, 16",
        "lu12i.w $t0, %abs_hi20({guard})",
        "ori $t0, $t0, %abs_lo12({guard})",
        "add.d $t0, $t0, $r21",
        "st.b $zero, $t0, 0",
        "1:",
        "2:",
        "jirl $zero, $ra, 0",
        global = sym KCOV_GLOBAL_GATE,
        guard = sym __PERCPU_IN_KCOV_TRACE,
        impl = sym kcov_trace_pc_impl,
    );
}