ax-cpu 0.9.0

//! Wrapper functions for assembly instructions.

use core::arch::asm;

use ax_memory_addr::{PAGE_SIZE_4K, PhysAddr, VirtAddr};
use loongArch64::register::{
    asid, crmd,
    ecfg::{self, LineBasedInterrupt},
    eentry, pgdh, pgdl,
};

const INVTLB_ADDR_GTRUE_OR_ASID: usize = 0x06;
const TLB_PAIR_SIZE: usize = PAGE_SIZE_4K * 2;

#[cfg(feature = "tls")]
use crate::KernelTlsBase;
#[cfg(feature = "uspace")]
use crate::{InstalledAddressSpace, InstalledAddressSpaceMode};

/// Returns the number of LoongArch ASIDs, including reserved ASID 0.
pub fn address_space_tag_capacity(_cpu_count: usize) -> u32 {
    // The generic installed identity stores a u16 tag. Preserve every ASID bit
    // representable by that contract instead of imposing an arbitrary 10-bit
    // software limit.
    let width = asid::read().asid_width().min(u16::BITS as usize);
    1u32.checked_shl(width as u32).unwrap_or(1).max(1)
}

#[cfg(feature = "uspace")]
fn flush_tlb_asid(tag: u16) {
    // op 0x4 invalidates every non-global entry matching the supplied ASID.
    unsafe {
        asm!(
            "dbar 0; invtlb 0x04, {asid}, $r0",
            asid = in(reg) usize::from(tag),
        )
    }
}

/// Installs one complete userspace identity into PGDL and CSR.ASID.
///
/// Incoming tagged contexts are invalidated before their root becomes usable;
/// the full-flush path installs ASID 0 and discards every local translation.
/// This is a conservative version of Linux's per-CPU ASID/version protocol:
/// tags may be globally allocated, but stale per-CPU state is never reused.
///
/// # Safety
///
/// The caller must own the current CPU with interrupts disabled and the root
/// must remain alive for the complete activation lease.
#[cfg(feature = "uspace")]
pub unsafe fn install_user_address_space(address_space: InstalledAddressSpace) {
    address_space.validate_architecture_support();
    let tagged = matches!(address_space.mode(), InstalledAddressSpaceMode::Tagged)
        && u32::from(address_space.hardware_tag()) < address_space_tag_capacity(1);
    if tagged {
        flush_tlb_asid(address_space.hardware_tag());
        // Linux writes PGDL before ASID so the new tag cannot name the old
        // page-table root. Scheduling is IRQ-disabled across both CSR writes.
        pgdl::set_base(address_space.root().as_usize() as _);
        asid::set_asid(usize::from(address_space.hardware_tag()));
    } else {
        pgdl::set_base(address_space.root().as_usize() as _);
        asid::set_asid(0);
        flush_tlb(None);
    }
}

core::arch::global_asm!(
    ".balign 16",
    ".global __axcpu_wait_for_irqs_disabled",
    ".global __axcpu_loongarch_idle_start",
    ".global __axcpu_loongarch_idle_exit",
    "__axcpu_wait_for_irqs_disabled:",
    "__axcpu_loongarch_idle_start:",
    "ori $t0, $zero, 4",
    "csrxchg $t0, $t0, 0x0",
    "idle 0",
    "__axcpu_loongarch_idle_exit:",
    "ret",
);

unsafe extern "C" {
    fn __axcpu_wait_for_irqs_disabled();
}

/// Allows the current CPU to respond to interrupts.
#[inline]
pub fn enable_irqs() {
    crmd::set_ie(true)
}

/// Makes the current CPU to ignore interrupts.
#[inline]
pub fn disable_irqs() {
    crmd::set_ie(false)
}

/// Returns whether the current CPU is allowed to respond to interrupts.
#[inline]
pub fn irqs_enabled() -> bool {
    crmd::read().ie()
}

/// Enables or disables the local timer interrupt line.
#[inline]
pub fn set_timer_irq_enabled(enabled: bool) {
    set_local_irq_line_enabled(LineBasedInterrupt::TIMER, enabled)
}

/// Enables or disables a local interrupt line.
#[inline]
fn set_local_irq_line_enabled(line: LineBasedInterrupt, enabled: bool) {
    let current = ecfg::read().lie();
    let new_value = if enabled {
        current | line
    } else {
        current & !line
    };
    ecfg::set_lie(new_value);
}

/// Relaxes the current CPU and waits for interrupts.
///
/// It must be called with interrupts enabled, otherwise it will never return.
#[inline]
pub fn wait_for_irqs() {
    unsafe { asm!("idle 0", options(nomem, nostack)) }
}

/// Waits for an interrupt after the caller masks local IRQ delivery.
///
/// LoongArch requires `CRMD.IE` to be set when `IDLE` executes. The assembly
/// window updates `CRMD.IE` immediately before `IDLE`; the trap path
/// fast-forwards an interrupt inside that window to its exit label. The
/// function returns with local IRQs enabled.
#[inline]
pub fn wait_for_irqs_disabled() {
    debug_assert!(!irqs_enabled());
    // SAFETY: the caller has masked local IRQ delivery. The assembly routine
    // only updates CRMD.IE, executes IDLE, and returns after the trap path has
    // preserved or fast-forwarded its continuation.
    unsafe { __axcpu_wait_for_irqs_disabled() }
}

/// Halt the current CPU.
#[inline]
pub fn halt() {
    disable_irqs();
    unsafe { loongArch64::asm::idle() }
}

/// Reads the current page table root register for user space (`PGDL`).
///
/// Returns the physical address of the page table root.
#[inline]
pub fn read_user_page_table() -> PhysAddr {
    PhysAddr::from(pgdl::read().base())
}

/// Reads the current page table root register for kernel space (`PGDH`).
///
/// Returns the physical address of the page table root.
#[inline]
pub fn read_kernel_page_table() -> PhysAddr {
    PhysAddr::from(pgdh::read().base())
}

/// Writes the register to update the current page table root for user space
/// (`PGDL`).
///
/// Note that the TLB is **NOT** flushed after this operation.
///
/// # Safety
///
/// This function is unsafe as it changes the virtual memory address space.
pub unsafe fn write_user_page_table(root_paddr: PhysAddr) {
    pgdl::set_base(root_paddr.as_usize() as _);
}

/// Writes the register to update the current page table root for kernel space
/// (`PGDH`).
///
/// Note that the TLB is **NOT** flushed after this operation.
///
/// # Safety
///
/// This function is unsafe as it changes the virtual memory address space.
pub unsafe fn write_kernel_page_table(root_paddr: PhysAddr) {
    pgdh::set_base(root_paddr.as_usize());
}

/// Flushes the entire instruction cache.
/// See <https://elixir.bootlin.com/linux/v6.6/source/arch/loongarch/mm/cache.c#L38>
#[inline]
pub fn flush_icache_all() {
    unsafe { asm!("ibar 0") };
}

/// Flushes the TLB.
///
/// If `vaddr` is [`None`], flushes the entire TLB. Otherwise, flushes the TLB
/// entry that maps the given virtual address.
#[inline]
pub fn flush_tlb(vaddr: Option<VirtAddr>) {
    unsafe {
        if let Some(vaddr) = vaddr {
            // <https://loongson.github.io/LoongArch-Documentation/LoongArch-Vol1-EN.html#_dbar>
            //
            // Only after all previous load/store access operations are completely
            // executed, the DBAR 0 instruction can be executed; and only after the
            // execution of DBAR 0 is completed, all subsequent load/store access
            // operations can be executed.
            //
            // <https://loongson.github.io/LoongArch-Documentation/LoongArch-Vol1-EN.html#_invtlb>
            //
            // formats: invtlb op, asid, addr
            //
            // LoongArch TLB entries map an even/odd page pair. Match Linux's
            // local_flush_tlb_one() by invalidating from the pair base and by
            // including global entries; op 0x5 would leave a stale kernel
            // translation installed. Supplying the live ASID also makes the
            // same operation valid for current user-space mappings.
            let pair_base = vaddr.as_usize() & !(TLB_PAIR_SIZE - 1);
            let current_asid = asid::read().asid();
            asm!(
                "dbar 0; invtlb {op}, {asid}, {addr}",
                op = const INVTLB_ADDR_GTRUE_OR_ASID,
                asid = in(reg) current_asid,
                addr = in(reg) pair_base,
            );
        } else {
            // op 0x0: Clear all page table entries
            asm!("dbar 0; invtlb 0x00, $r0, $r0");
        }
    }
}

/// Makes a page-table entry installed by the local page-fault handler visible
/// before retrying the faulting instruction.
///
/// The software refill path may cache an invalid entry for a missing page-table
/// level. Invalidate that local entry so the retry refills from the newly
/// installed PTE.
#[inline]
pub fn update_mmu_cache(vaddr: VirtAddr) {
    flush_tlb(Some(vaddr));
}

/// Writes the Exception Entry Base Address register (`EENTRY`).
///
/// It also set the Exception Configuration register (`ECFG`) to `VS=0`.
///
/// - ECFG: <https://loongson.github.io/LoongArch-Documentation/LoongArch-Vol1-EN.html#exception-configuration>
/// - EENTRY: <https://loongson.github.io/LoongArch-Documentation/LoongArch-Vol1-EN.html#exception-entry-base-address>
///
/// # Safety
///
/// This function is unsafe as it changes the exception handling behavior of the
/// current CPU.
#[inline]
pub unsafe fn write_exception_entry_base(eentry: usize) {
    ecfg::set_vs(0);
    eentry::set_eentry(eentry);
}

/// Writes the Page Walk Controller registers (`PWCL` and `PWCH`).
///
/// # Safety
///
/// This function is unsafe as it changes the page walk configuration such as
/// levels and starting bits.
///
/// - `PWCL`: <https://loongson.github.io/LoongArch-Documentation/LoongArch-Vol1-EN.html#page-walk-controller-for-lower-half-address-space>
/// - `PWCH`: <https://loongson.github.io/LoongArch-Documentation/LoongArch-Vol1-EN.html#page-walk-controller-for-higher-half-address-space>
#[inline]
pub unsafe fn write_pwc(pwcl: u32, pwch: u32) {
    unsafe {
        asm!(
            include_asm_macros!(),
            "csrwr {}, LA_CSR_PWCL",
            "csrwr {}, LA_CSR_PWCH",
            in(reg) pwcl,
            in(reg) pwch
        )
    }
}

/// Reads the current kernel task's TLS base from `$tp`.
///
/// This register follows the execution context across CPUs. It is distinct
/// from the CPU-local base kept in `$r21`.
#[inline]
#[cfg(feature = "tls")]
pub fn read_thread_pointer() -> KernelTlsBase {
    let address;
    unsafe { asm!("move {}, $tp", out(reg) address) };
    KernelTlsBase::new(address)
}

/// Writes the current kernel task's TLS base to `$tp`.
///
/// This register follows the execution context across CPUs. It is distinct
/// from the CPU-local base kept in `$r21`.
///
/// # Safety
///
/// The caller must ensure `kernel_tls` belongs to the execution context that
/// is becoming current and that no Rust code observes a half-completed context
/// switch.
#[inline]
#[cfg(feature = "tls")]
pub unsafe fn write_thread_pointer(kernel_tls: KernelTlsBase) {
    unsafe { asm!("move $tp, {}", in(reg) kernel_tls.as_usize()) }
}

/// Enables floating-point instructions by setting `EUEN.FPE`.
///
/// - `EUEN`: <https://loongson.github.io/LoongArch-Documentation/LoongArch-Vol1-EN.html#extended-component-unit-enable>
#[inline]
pub fn enable_fp() {
    loongArch64::register::euen::set_fpe(true);
}

/// Enables LSX extension by setting `EUEN.LSX`.
///
/// - `EUEN`: <https://loongson.github.io/LoongArch-Documentation/LoongArch-Vol1-EN.html#extended-component-unit-enable>
pub fn enable_lsx() {
    loongArch64::register::euen::set_sxe(true);
}

/// Enables LASX extension by setting `EUEN.ASXE`.
///
/// - `EUEN`: <https://loongson.github.io/LoongArch-Documentation/LoongArch-Vol1-EN.html#extended-component-unit-enable>
pub fn enable_lasx() {
    loongArch64::register::euen::set_asxe(true);
}

#[cfg(feature = "uspace")]
core::arch::global_asm!(
    include_asm_macros!(),
    include_str!("user_copy.S"),
    include_str!("user_atomic.S"),
);

#[cfg(feature = "uspace")]
unsafe extern "C" {
    /// Copies data from source to destination, where addresses may be in user
    /// space. Equivalent to memcpy.
    ///
    /// # Safety
    /// This function is unsafe because it performs raw memory operations.
    ///
    /// # Returns
    /// Returns the number of bytes not copied. This means 0 indicates success,
    /// while a value > 0 indicates failure.
    pub fn user_copy(dst: *mut u8, src: *const u8, size: usize) -> usize;
}

/// Lock-free EL0/user access probe. No hardware address-translation probe is
/// wired up on this architecture yet, so always report a present-page probe miss
/// and let the caller take the locked slow path (correctness preserved).
///
/// # Safety
///
/// No precondition — this stub reads nothing and always returns `false`. It is
/// `unsafe` only to share the signature of the aarch64 EL1 probe (which requires
/// IRQs-off), so callers can use one `unsafe` block across all targets.
#[cfg(feature = "uspace")]
#[inline]
pub unsafe fn user_access_ok_page(_vaddr: usize, _access: crate::UserAccessType) -> bool {
    false
}