ax_cpu/aarch64/
context.rs

1use core::{
2    arch::naked_asm,
3    fmt,
4    mem::{align_of, offset_of, size_of},
5    ptr::NonNull,
6};
7
8use ax_memory_addr::VirtAddr;
9use cpu_local::{ExecutionContextHeader, PreparedContextSwitch};
10
11use crate::{KernelTlsBase, TaskLocalState};
12
13/// Saved registers when a trap (exception) occurs.
14#[repr(C)]
15#[derive(Default, Clone, Copy)]
16pub struct TrapFrame {
17    /// General-purpose registers (X0..X30).
18    pub x: [u64; 31],
19    /// Exception Link Register (ELR_EL1).
20    pub elr: u64,
21    /// Saved Process Status Register (SPSR_EL1).
22    pub spsr: u64,
23
24    /// Stack pointer at the time of the exception.
25    /// Populated by SAVE_REGS as `sp_before_sub = sp_after_sub + trapframe_size`.
26    ///
27    /// Note: This field is read-only (saved by SAVE_REGS for inspection only).
28    /// The actual SP is restored by RESTORE_REGS via `add sp, sp, #trapframe_size`,
29    /// not from this field. Modifying this value will NOT affect the actual SP
30    /// after exception return.
31    pub sp: u64,
32}
33
34impl fmt::Debug for TrapFrame {
35    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36        writeln!(f, "TrapFrame: {{")?;
37        for (i, &reg) in self.x.iter().enumerate() {
38            writeln!(f, "    x{i}: {reg:#x},")?;
39        }
40        writeln!(f, "    elr: {:#x},", self.elr)?;
41        writeln!(f, "    spsr: {:#x},", self.spsr)?;
42        writeln!(f, "    sp: {:#x},", self.sp)?;
43        write!(f, "}}")?;
44        Ok(())
45    }
46}
47
48impl TrapFrame {
49    /// Returns the privilege domain represented by this register image.
50    pub const fn origin(&self) -> crate::TrapOrigin {
51        if self.spsr & 0b1_1111 == 0 {
52            crate::TrapOrigin::User
53        } else {
54            crate::TrapOrigin::Kernel
55        }
56    }
57
58    /// Gets the 0th syscall argument.
59    pub const fn arg0(&self) -> usize {
60        self.x[0] as _
61    }
62
63    /// Sets the 0th syscall argument.
64    pub const fn set_arg0(&mut self, a0: usize) {
65        self.x[0] = a0 as _;
66    }
67
68    /// Gets the 1st syscall argument.
69    pub const fn arg1(&self) -> usize {
70        self.x[1] as _
71    }
72
73    /// Sets the 1st syscall argument.
74    pub const fn set_arg1(&mut self, a1: usize) {
75        self.x[1] = a1 as _;
76    }
77
78    /// Gets the 2nd syscall argument.
79    pub const fn arg2(&self) -> usize {
80        self.x[2] as _
81    }
82
83    /// Sets the 2nd syscall argument.
84    pub const fn set_arg2(&mut self, a2: usize) {
85        self.x[2] = a2 as _;
86    }
87
88    /// Gets the 3rd syscall argument.
89    pub const fn arg3(&self) -> usize {
90        self.x[3] as _
91    }
92
93    /// Sets the 3rd syscall argument.
94    pub const fn set_arg3(&mut self, a3: usize) {
95        self.x[3] = a3 as _;
96    }
97
98    /// Gets the 4th syscall argument.
99    pub const fn arg4(&self) -> usize {
100        self.x[4] as _
101    }
102
103    /// Sets the 4th syscall argument.
104    pub const fn set_arg4(&mut self, a4: usize) {
105        self.x[4] = a4 as _;
106    }
107
108    /// Gets the 5th syscall argument.
109    pub const fn arg5(&self) -> usize {
110        self.x[5] as _
111    }
112
113    /// Sets the 5th syscall argument.
114    pub const fn set_arg5(&mut self, a5: usize) {
115        self.x[5] = a5 as _;
116    }
117
118    /// Gets the instruction pointer.
119    pub const fn ip(&self) -> usize {
120        self.elr as _
121    }
122
123    /// Sets the instruction pointer.
124    pub const fn set_ip(&mut self, pc: usize) {
125        self.elr = pc as _;
126    }
127
128    /// Get the syscall number.
129    pub const fn sysno(&self) -> usize {
130        self.x[8] as usize
131    }
132
133    /// Sets the syscall number.
134    pub const fn set_sysno(&mut self, sysno: usize) {
135        self.x[8] = sysno as _;
136    }
137
138    /// Gets the return value register.
139    pub const fn retval(&self) -> usize {
140        self.x[0] as _
141    }
142
143    /// Sets the return value register.
144    pub const fn set_retval(&mut self, r0: usize) {
145        self.x[0] = r0 as _;
146    }
147
148    /// Sets the return address.
149    pub const fn set_ra(&mut self, lr: usize) {
150        self.x[30] = lr as _;
151    }
152
153    /// Unwind the stack and get the backtrace.
154    pub fn backtrace(&self) -> axbacktrace::Backtrace {
155        axbacktrace::Backtrace::capture_trap(self.x[29] as _, self.elr as _, self.x[30] as _)
156    }
157}
158
159/// FP & SIMD registers.
160#[repr(C, align(16))]
161#[derive(Debug, Default)]
162pub struct FpState {
163    /// 128-bit SIMD & FP registers (V0..V31)
164    pub regs: [u128; 32],
165    /// Floating-point Control Register (FPCR)
166    pub fpcr: u32,
167    /// Floating-point Status Register (FPSR)
168    pub fpsr: u32,
169}
170
171#[cfg(feature = "fp-simd")]
172impl FpState {
173    /// Saves the current FP/SIMD states from CPU to this structure.
174    pub fn save(&mut self) {
175        unsafe { fpstate_save(self) }
176    }
177
178    /// Restores the FP/SIMD states from this structure to CPU.
179    pub fn restore(&self) {
180        unsafe { fpstate_restore(self) }
181    }
182}
183
184/// Saved hardware states of a task.
185///
186/// The context usually includes:
187///
188/// - Callee-saved registers
189/// - Stack pointer register
190/// - Thread pointer register (for kernel-space thread-local storage)
191/// - FP/SIMD registers
192///
193/// On context switch, current task saves its context from CPU to memory,
194/// and the next task restores its context from memory to CPU.
195#[allow(missing_docs)]
196#[repr(C)]
197#[derive(Debug, Default)]
198pub struct TaskContext {
199    sp: u64,
200    r19: u64,
201    r20: u64,
202    r21: u64,
203    r22: u64,
204    r23: u64,
205    r24: u64,
206    r25: u64,
207    r26: u64,
208    r27: u64,
209    r28: u64,
210    r29: u64,
211    lr: u64, // r30
212    /// Architecture-neutral current-header and kernel-TLS switch state.
213    task_local: TaskLocalState,
214    #[cfg(feature = "fp-simd")]
215    fp_state: FpState,
216}
217
218// `stp`/`ldp` address only the first member of each pair. Prove the paired
219// fields are adjacent and that the task TLS newtype has the register width.
220const _: () = {
221    assert!(size_of::<KernelTlsBase>() == size_of::<usize>());
222    assert!(align_of::<KernelTlsBase>() == align_of::<usize>());
223    assert!(offset_of!(TaskContext, sp) == 0);
224    assert!(offset_of!(TaskContext, r20) == offset_of!(TaskContext, r19) + size_of::<u64>());
225    assert!(offset_of!(TaskContext, r22) == offset_of!(TaskContext, r21) + size_of::<u64>());
226    assert!(offset_of!(TaskContext, r24) == offset_of!(TaskContext, r23) + size_of::<u64>());
227    assert!(offset_of!(TaskContext, r26) == offset_of!(TaskContext, r25) + size_of::<u64>());
228    assert!(offset_of!(TaskContext, r28) == offset_of!(TaskContext, r27) + size_of::<u64>());
229    assert!(offset_of!(TaskContext, lr) == offset_of!(TaskContext, r29) + size_of::<u64>());
230    assert!(offset_of!(TaskContext, task_local) == offset_of!(TaskContext, lr) + size_of::<u64>());
231};
232
233impl TaskContext {
234    /// Creates a dummy context for a new task.
235    ///
236    /// Note the context is not initialized, it will be filled by
237    /// [`switch_to_prepared`](Self::switch_to_prepared) (for initial tasks) and [`init`]
238    /// (for regular tasks) methods.
239    ///
240    /// [`init`]: TaskContext::init
241    pub fn new() -> Self {
242        Self::default()
243    }
244
245    /// Initializes the context for a new task, with the given entry point and
246    /// kernel stack.
247    pub fn init(&mut self, entry: usize, kstack_top: VirtAddr, kernel_tls: KernelTlsBase) {
248        self.sp = kstack_top.as_usize() as u64;
249        self.lr = entry as u64;
250        self.task_local.set_kernel_tls(kernel_tls);
251    }
252
253    /// Sets the pinned task-owned execution-context header.
254    pub fn set_context_header(&mut self, header: NonNull<ExecutionContextHeader>) {
255        self.task_local.set_context_header(header);
256    }
257
258    /// Returns the configured task-owned execution-context header.
259    pub const fn context_header(&self) -> Option<NonNull<ExecutionContextHeader>> {
260        self.task_local.context_header()
261    }
262
263    /// Completes FP/SIMD work before current-context publication.
264    pub fn prepare_switch_to(&mut self, _next_ctx: &Self) {
265        #[cfg(feature = "fp-simd")]
266        {
267            self.fp_state.save();
268            _next_ctx.fp_state.restore();
269        }
270    }
271
272    /// Performs only the final GPR/current/TLS transfer.
273    ///
274    /// # Safety
275    ///
276    /// Scheduling must be serialized, FP state prepared, and the next current
277    /// header published. No fallible Rust work may follow before this call.
278    #[inline(always)]
279    pub unsafe fn switch_to_prepared(
280        &mut self,
281        next_ctx: &Self,
282        prepared: PreparedContextSwitch<'_>,
283    ) {
284        unsafe { prepared.commit() };
285        unsafe { context_switch_raw(self, next_ctx) }
286    }
287}
288
289#[cfg(kernel_tls)]
290#[unsafe(naked)]
291unsafe extern "C" fn context_switch_raw(_current_task: &mut TaskContext, _next_task: &TaskContext) {
292    naked_asm!(
293        "
294        // save old context (callee-saved registers)
295        stp     x29, x30, [x0, {r29_offset}]
296        stp     x27, x28, [x0, {r27_offset}]
297        stp     x25, x26, [x0, {r25_offset}]
298        stp     x23, x24, [x0, {r23_offset}]
299        stp     x21, x22, [x0, {r21_offset}]
300        stp     x19, x20, [x0, {r19_offset}]
301        mov     x19, sp
302        str     x19, [x0, {sp_offset}]
303        mrs     x9, tpidr_el0
304        str     x9, [x0, {kernel_tls_offset}]
305
306        // restore new context
307        ldr     x9, [x1, {kernel_tls_offset}]
308        msr     tpidr_el0, x9
309        ldr     x19, [x1, {sp_offset}]
310        mov     sp, x19
311        ldp     x19, x20, [x1, {r19_offset}]
312        ldp     x21, x22, [x1, {r21_offset}]
313        ldp     x23, x24, [x1, {r23_offset}]
314        ldp     x25, x26, [x1, {r25_offset}]
315        ldp     x27, x28, [x1, {r27_offset}]
316        ldp     x29, x30, [x1, {r29_offset}]
317        ldr     x9, [x1, {context_header_offset}]
318        msr     sp_el0, x9
319
320        ret",
321        sp_offset = const offset_of!(TaskContext, sp),
322        r19_offset = const offset_of!(TaskContext, r19),
323        r21_offset = const offset_of!(TaskContext, r21),
324        r23_offset = const offset_of!(TaskContext, r23),
325        r25_offset = const offset_of!(TaskContext, r25),
326        r27_offset = const offset_of!(TaskContext, r27),
327        r29_offset = const offset_of!(TaskContext, r29),
328        context_header_offset = const offset_of!(TaskContext, task_local)
329            + offset_of!(TaskLocalState, context_header),
330        kernel_tls_offset = const offset_of!(TaskContext, task_local)
331            + offset_of!(TaskLocalState, kernel_tls),
332    )
333}
334
335#[cfg(not(kernel_tls))]
336#[unsafe(naked)]
337unsafe extern "C" fn context_switch_raw(_current_task: &mut TaskContext, _next_task: &TaskContext) {
338    naked_asm!(
339        "
340        // save old context (callee-saved registers)
341        stp     x29, x30, [x0, {r29_offset}]
342        stp     x27, x28, [x0, {r27_offset}]
343        stp     x25, x26, [x0, {r25_offset}]
344        stp     x23, x24, [x0, {r23_offset}]
345        stp     x21, x22, [x0, {r21_offset}]
346        stp     x19, x20, [x0, {r19_offset}]
347        mov     x19, sp
348        str     x19, [x0, {sp_offset}]
349
350        // LinuxCurrent keeps task identity in SP_EL0. TPIDR_EL0 remains
351        // userspace-owned and is never part of a kernel task switch.
352        ldr     x19, [x1, {sp_offset}]
353        mov     sp, x19
354        ldp     x19, x20, [x1, {r19_offset}]
355        ldp     x21, x22, [x1, {r21_offset}]
356        ldp     x23, x24, [x1, {r23_offset}]
357        ldp     x25, x26, [x1, {r25_offset}]
358        ldp     x27, x28, [x1, {r27_offset}]
359        ldp     x29, x30, [x1, {r29_offset}]
360        ldr     x9, [x1, {context_header_offset}]
361        msr     sp_el0, x9
362        ret",
363        sp_offset = const offset_of!(TaskContext, sp),
364        r19_offset = const offset_of!(TaskContext, r19),
365        r21_offset = const offset_of!(TaskContext, r21),
366        r23_offset = const offset_of!(TaskContext, r23),
367        r25_offset = const offset_of!(TaskContext, r25),
368        r27_offset = const offset_of!(TaskContext, r27),
369        r29_offset = const offset_of!(TaskContext, r29),
370        context_header_offset = const offset_of!(TaskContext, task_local)
371            + offset_of!(TaskLocalState, context_header),
372    )
373}
374
375#[unsafe(naked)]
376#[cfg(feature = "fp-simd")]
377unsafe extern "C" fn fpstate_save(state: &mut FpState) {
378    naked_asm!(
379        ".arch armv8
380        // save fp/neon context
381        mrs     x9, fpcr
382        mrs     x10, fpsr
383        stp     q0, q1, [x0, 0 * 16]
384        stp     q2, q3, [x0, 2 * 16]
385        stp     q4, q5, [x0, 4 * 16]
386        stp     q6, q7, [x0, 6 * 16]
387        stp     q8, q9, [x0, 8 * 16]
388        stp     q10, q11, [x0, 10 * 16]
389        stp     q12, q13, [x0, 12 * 16]
390        stp     q14, q15, [x0, 14 * 16]
391        stp     q16, q17, [x0, 16 * 16]
392        stp     q18, q19, [x0, 18 * 16]
393        stp     q20, q21, [x0, 20 * 16]
394        stp     q22, q23, [x0, 22 * 16]
395        stp     q24, q25, [x0, 24 * 16]
396        stp     q26, q27, [x0, 26 * 16]
397        stp     q28, q29, [x0, 28 * 16]
398        stp     q30, q31, [x0, 30 * 16]
399        str     x9, [x0, 64 *  8]
400        str     x10, [x0, 65 * 8]
401
402        isb
403        ret"
404    )
405}
406
407#[unsafe(naked)]
408#[cfg(feature = "fp-simd")]
409unsafe extern "C" fn fpstate_restore(state: &FpState) {
410    naked_asm!(
411        ".arch armv8
412        // restore fp/neon context
413        ldp     q0, q1, [x0, 0 * 16]
414        ldp     q2, q3, [x0, 2 * 16]
415        ldp     q4, q5, [x0, 4 * 16]
416        ldp     q6, q7, [x0, 6 * 16]
417        ldp     q8, q9, [x0, 8 * 16]
418        ldp     q10, q11, [x0, 10 * 16]
419        ldp     q12, q13, [x0, 12 * 16]
420        ldp     q14, q15, [x0, 14 * 16]
421        ldp     q16, q17, [x0, 16 * 16]
422        ldp     q18, q19, [x0, 18 * 16]
423        ldp     q20, q21, [x0, 20 * 16]
424        ldp     q22, q23, [x0, 22 * 16]
425        ldp     q24, q25, [x0, 24 * 16]
426        ldp     q26, q27, [x0, 26 * 16]
427        ldp     q28, q29, [x0, 28 * 16]
428        ldp     q30, q31, [x0, 30 * 16]
429        ldr     x9, [x0, 64 * 8]
430        ldr     x10, [x0, 65 * 8]
431        msr     fpcr, x9
432        msr     fpsr, x10
433
434        isb
435        ret"
436    )
437}
ax_cpu/aarch64/context.rs

ax_cpu/aarch64/
context.rs
