Skip to main content

ax_hal/
percpu.rs

1//! CPU-local data structures.
2
3pub use ax_plat::percpu::*;
4
5#[ax_percpu::def_percpu]
6static CURRENT_TASK_PTR: usize = 0;
7
8/// Gets the pointer to the current task with preemption-safety.
9///
10/// Preemption may be enabled when calling this function. This function will
11/// guarantee the correctness even the current task is preempted.
12#[inline]
13pub fn current_task_ptr<T>() -> *const T {
14    #[cfg(target_arch = "x86_64")]
15    unsafe {
16        // on x86, only one instruction is needed to read the per-CPU task pointer from `gs:[off]`.
17        CURRENT_TASK_PTR.read_current_raw() as _
18    }
19    #[cfg(any(
20        target_arch = "aarch64",
21        target_arch = "riscv32",
22        target_arch = "riscv64",
23        target_arch = "loongarch64"
24    ))]
25    unsafe {
26        // on RISC-V and LA64, reading `CURRENT_TASK_PTR` requires multiple instruction, so we disable local IRQs.
27        let _guard = ax_kernel_guard::IrqSave::new();
28        CURRENT_TASK_PTR.read_current_raw() as _
29    }
30}
31
32/// Sets the pointer to the current task with preemption-safety.
33///
34/// Preemption may be enabled when calling this function. This function will
35/// guarantee the correctness even the current task is preempted.
36///
37/// # Safety
38///
39/// The given `ptr` must be pointed to a valid task structure.
40#[inline]
41pub unsafe fn set_current_task_ptr<T>(ptr: *const T) {
42    #[cfg(target_arch = "x86_64")]
43    {
44        unsafe { CURRENT_TASK_PTR.write_current_raw(ptr as usize) }
45    }
46    #[cfg(any(
47        target_arch = "aarch64",
48        target_arch = "riscv32",
49        target_arch = "riscv64",
50        target_arch = "loongarch64"
51    ))]
52    {
53        let _guard = ax_kernel_guard::IrqSave::new();
54        unsafe { CURRENT_TASK_PTR.write_current_raw(ptr as usize) }
55    }
56}