loongarch64, trap: rework to fix nested captured traps

Similar to 661a15940badf992d7836eb93c7ab293590c6ad4:
- Save previous {trap, task}_ctx and restore them afterwards.
- Set kernel tp when setting trap context user mode.
- Add the program counter with 4 bytes on breakpoints.

Signed-off-by: greatbridf <greatbridf@icloud.com>

crates/eonix_hal/src/arch/loongarch64/trap/mod.rs

@@ -278,11 +278,18 @@ impl TrapReturn for TrapContext {
         to_ctx.set_interrupt_enabled(false);
 
         unsafe {
+            let mut old_trap_ctx: usize;
+            let mut old_task_ctx: usize;
+
             asm!(
+                "csrrd {old_trap_ctx}, {CSR_CAPTURED_TRAP_CONTEXT_ADDR}",
+                "csrrd {old_task_ctx}, {CSR_CAPTURER_TASK_CONTEXT_ADDR}",
                 "csrwr {captured_trap_context}, {CSR_CAPTURED_TRAP_CONTEXT_ADDR}",
                 "csrwr {capturer_task_context}, {CSR_CAPTURER_TASK_CONTEXT_ADDR}",
                 captured_trap_context = inout(reg) &raw mut *self => _,
                 capturer_task_context = inout(reg) &raw mut capturer_ctx => _,
+                old_trap_ctx = out(reg) old_trap_ctx,
+                old_task_ctx = out(reg) old_task_ctx,
                 CSR_CAPTURED_TRAP_CONTEXT_ADDR = const CSR_CAPTURED_TRAP_CONTEXT_ADDR,
                 CSR_CAPTURER_TASK_CONTEXT_ADDR = const CSR_CAPTURER_TASK_CONTEXT_ADDR,
                 options(nomem, nostack, preserves_flags),
@@ -291,8 +298,10 @@ impl TrapReturn for TrapContext {
             TaskContext::switch(&mut capturer_ctx, &mut to_ctx);
 
             asm!(
-                "csrwr $zero, {CSR_CAPTURED_TRAP_CONTEXT_ADDR}",
-                "csrwr $zero, {CSR_CAPTURER_TASK_CONTEXT_ADDR}",
+                "csrwr {old_trap_ctx}, {CSR_CAPTURED_TRAP_CONTEXT_ADDR}",
+                "csrwr {old_task_ctx}, {CSR_CAPTURER_TASK_CONTEXT_ADDR}",
+                old_trap_ctx = inout(reg) old_trap_ctx,
+                old_task_ctx = inout(reg) old_task_ctx,
                 CSR_CAPTURED_TRAP_CONTEXT_ADDR = const CSR_CAPTURED_TRAP_CONTEXT_ADDR,
                 CSR_CAPTURER_TASK_CONTEXT_ADDR = const CSR_CAPTURER_TASK_CONTEXT_ADDR,
                 options(nomem, nostack, preserves_flags),

crates/eonix_hal/src/arch/loongarch64/trap/trap_context.rs

@@ -1,4 +1,4 @@
-use crate::processor::CPU;
+use crate::{arch::trap::CSR_KERNEL_TP, processor::CPU};
 use core::{arch::asm, mem::offset_of};
 use eonix_hal_traits::{
     fault::{Fault, PageFaultErrorCode},
@@ -226,7 +226,17 @@ impl RawTrapContext for TrapContext {
     fn set_user_mode(&mut self, user: bool) {
         match user {
             true => self.prmd |= 0x3,
-            false => self.prmd &= !0x3,
+            false => {
+                unsafe {
+                    asm!(
+                        "csrrd {tp}, {CSR_KERNEL_TP}",
+                        tp = out(reg) self.regs.tp,
+                        CSR_KERNEL_TP = const CSR_KERNEL_TP,
+                        options(nomem, nostack, preserves_flags),
+                    )
+                }
+                self.prmd &= !0x3;
+            }
         }
     }
 

src/driver/virtio/loongarch64.rs

@@ -3,13 +3,13 @@ use crate::kernel::{
     block::{make_device, BlockDevice},
     constants::EIO,
     pcie::{self, PCIDevice, PCIDriver, PciError, SegmentGroup},
+    task::block_on,
 };
 use alloc::sync::Arc;
 use core::sync::atomic::{AtomicUsize, Ordering};
 use eonix_hal::{fence::memory_barrier, mm::ArchPhysAccess};
 use eonix_log::println_warn;
 use eonix_mm::address::PhysAccess;
-use eonix_runtime::task::Task;
 use eonix_sync::Spin;
 use virtio_drivers::{
     device::blk::VirtIOBlk,
@@ -134,7 +134,7 @@ impl PCIDriver for VirtIODriver {
             Arc::new(Spin::new(virtio_block)),
         )?;
 
-        Task::block_on(block_device.partprobe()).map_err(|err| {
+        block_on(block_device.partprobe()).map_err(|err| {
             println_warn!(
                 "Failed to probe partitions for VirtIO Block device: {}",
                 err

src/kernel/task.rs

@@ -136,7 +136,11 @@ where
                     assert_preempt_enabled!("Blocking in stackful futures is not allowed.");
 
                     unsafe {
+                        #[cfg(target_arch = "riscv64")]
                         core::arch::asm!("ebreak");
+
+                        #[cfg(target_arch = "loongarch64")]
+                        core::arch::asm!("break 1");
                     }
                 }
             }
@@ -147,7 +151,11 @@ where
         }
 
         unsafe {
+            #[cfg(target_arch = "riscv64")]
             core::arch::asm!("ebreak");
+
+            #[cfg(target_arch = "loongarch64")]
+            core::arch::asm!("break 1");
         }
 
         unreachable!()
@@ -184,7 +192,11 @@ where
                         wait_for_wakeups().await;
                     }
 
+                    #[cfg(target_arch = "riscv64")]
                     trap_ctx.set_program_counter(trap_ctx.get_program_counter() + 2);
+
+                    #[cfg(target_arch = "loongarch64")]
+                    trap_ctx.set_program_counter(trap_ctx.get_program_counter() + 4);
                 } else {
                     default_fault_handler(fault, &mut trap_ctx)
                 }