refactor(prototyper): trap

`trap.rs` is complex enough to refactor.
First we can make all trap expcet boot use `fast_trap` crate,
second we can split this file into boot, handler and itself.

Signed-off-by: Woshiluo Luo <woshiluo.luo@outlook.com>

prototyper/src/main.rs

@@ -32,7 +32,7 @@ use crate::sbi::hart_context::NextStage;
 use crate::sbi::heap::sbi_heap_init;
 use crate::sbi::hsm::local_remote_hsm;
 use crate::sbi::ipi;
-use crate::sbi::trap::{self, trap_vec};
+use crate::sbi::trap;
 use crate::sbi::trap_stack;
 
 pub const R_RISCV_RELATIVE: usize = 3;
@@ -119,8 +119,8 @@ extern "C" fn rust_main(_hart_id: usize, opaque: usize, nonstandard_a2: usize) {
                 menvcfg::set_bits(menvcfg::CBIE_INVALIDATE | menvcfg::CBCFE | menvcfg::CBZE);
             }
         }
-        // Set up vectored trap handling.
-        mtvec::write(trap_vec as _, mtvec::TrapMode::Vectored);
+        // Set up trap handling.
+        mtvec::write(fast_trap::trap_entry as _, mtvec::TrapMode::Direct);
     }
 }
 
@@ -166,7 +166,7 @@ unsafe extern "C" fn start() -> ! {
         relocation_update = sym relocation_update,
         locate_stack = sym trap_stack::locate,
         main         = sym rust_main,
-        hart_boot    = sym trap::msoft,
+        hart_boot    = sym trap::boot::boot,
         options(noreturn)
     )
 }

prototyper/src/sbi/hsm.rs

@@ -244,7 +244,7 @@ impl rustsbi::Hsm for SbiHsm {
             }
             local_hsm().suspend();
             riscv::asm::wfi();
-            crate::trap::msoft_ipi_handler();
+            crate::sbi::trap::handler::msoft_ipi_handler();
             local_hsm().resume();
             SbiRet::success(0)
         } else {

prototyper/src/sbi/ipi.rs

@@ -4,7 +4,6 @@ use crate::riscv::current_hartid;
 use crate::sbi::extensions::{hart_extension_probe, Extension};
 use crate::sbi::hsm::remote_hsm;
 use crate::sbi::rfence;
-use crate::sbi::trap;
 use crate::sbi::trap_stack::ROOT_STACK;
 use alloc::boxed::Box;
 use core::sync::atomic::Ordering::Relaxed;
@@ -182,7 +181,7 @@ impl SbiIpi {
 
         // Wait for all fence operations to complete
         while !rfence::local_rfence().unwrap().is_sync() {
-            trap::rfence_single_handler();
+            rfence::rfence_single_handler();
         }
 
         SbiRet::success(0)

prototyper/src/sbi/rfence.rs

@@ -1,11 +1,12 @@
 use rustsbi::{HartMask, SbiRet};
 use spin::Mutex;
 
+use crate::cfg::{PAGE_SIZE, TLB_FLUSH_LIMIT};
 use crate::platform::PLATFORM;
 use crate::riscv::current_hartid;
 use crate::sbi::fifo::{Fifo, FifoError};
-use crate::sbi::trap;
 use crate::sbi::trap_stack::ROOT_STACK;
+use core::arch::asm;
 
 use core::sync::atomic::{AtomicU32, Ordering};
 
@@ -134,7 +135,7 @@ impl LocalRFenceCell<'_> {
                 Ok(_) => break,
                 Err(FifoError::Full) => {
                     drop(queue);
-                    trap::rfence_single_handler();
+                    rfence_single_handler();
                 }
                 Err(_) => panic!("Unable to push fence ops to fifo"),
             }
@@ -153,7 +154,7 @@ impl RemoteRFenceCell<'_> {
                 Ok(_) => return,
                 Err(FifoError::Full) => {
                     drop(queue);
-                    trap::rfence_single_handler();
+                    rfence_single_handler();
                 }
                 Err(_) => panic!("Unable to push fence ops to fifo"),
             }
@@ -253,3 +254,70 @@ impl rustsbi::Fence for SbiRFence {
         )
     }
 }
+
+/// Handles a single remote fence operation.
+#[inline]
+pub fn rfence_single_handler() {
+    let rfence_context = local_rfence().unwrap().get();
+    if let Some((ctx, id)) = rfence_context {
+        match ctx.op {
+            // Handle instruction fence
+            RFenceType::FenceI => unsafe {
+                asm!("fence.i");
+                remote_rfence(id).unwrap().sub();
+            },
+            // Handle virtual memory address fence
+            RFenceType::SFenceVma => {
+                // If the flush size is greater than the maximum limit then simply flush all
+                if (ctx.start_addr == 0 && ctx.size == 0)
+                    || (ctx.size == usize::MAX)
+                    || (ctx.size > TLB_FLUSH_LIMIT)
+                {
+                    unsafe {
+                        asm!("sfence.vma");
+                    }
+                } else {
+                    for offset in (0..ctx.size).step_by(PAGE_SIZE) {
+                        let addr = ctx.start_addr + offset;
+                        unsafe {
+                            asm!("sfence.vma {}", in(reg) addr);
+                        }
+                    }
+                }
+                remote_rfence(id).unwrap().sub();
+            }
+            // Handle virtual memory address fence with ASID
+            RFenceType::SFenceVmaAsid => {
+                let asid = ctx.asid;
+                // If the flush size is greater than the maximum limit then simply flush all
+                if (ctx.start_addr == 0 && ctx.size == 0)
+                    || (ctx.size == usize::MAX)
+                    || (ctx.size > TLB_FLUSH_LIMIT)
+                {
+                    unsafe {
+                        asm!("sfence.vma {}, {}", in(reg) 0, in(reg) asid);
+                    }
+                } else {
+                    for offset in (0..ctx.size).step_by(PAGE_SIZE) {
+                        let addr = ctx.start_addr + offset;
+                        unsafe {
+                            asm!("sfence.vma {}, {}", in(reg) addr, in(reg) asid);
+                        }
+                    }
+                }
+                remote_rfence(id).unwrap().sub();
+            }
+            rfencetype => {
+                error!("Unsupported RFence Type: {:?}!", rfencetype);
+            }
+        }
+    }
+}
+
+/// Process all pending remote fence operations.
+#[inline]
+pub fn rfence_handler() {
+    while !local_rfence().unwrap().is_empty() {
+        rfence_single_handler();
+    }
+}

prototyper/src/sbi/trap.rs

@@ -1,537 +0,0 @@
-use core::arch::asm;
-use fast_trap::{trap_entry, FastContext, FastResult};
-use riscv::register::{
-    mcause::{self, Exception as E, Trap as T},
-    mepc, mie, mstatus, mtval, satp, sstatus,
-};
-use rustsbi::RustSBI;
-
-use crate::cfg::{PAGE_SIZE, TLB_FLUSH_LIMIT};
-use crate::platform::PLATFORM;
-use crate::riscv::csr::{CSR_TIME, CSR_TIMEH};
-use crate::riscv::current_hartid;
-use crate::sbi::console;
-use crate::sbi::hsm::local_hsm;
-use crate::sbi::ipi;
-use crate::sbi::rfence::{self, local_rfence, RFenceType};
-
-/// Trap vector table entry point. Maps different trap types to their handlers.
-#[naked]
-pub(crate) unsafe extern "C" fn trap_vec() {
-    asm!(
-        ".align 2",
-        ".option push",
-        ".option norvc",
-        "j {default}", // exception
-        "j {default}", // supervisor software
-        "j {default}", // reserved
-        "j {msoft} ",  // machine    software
-        "j {default}", // reserved
-        "j {default}", // supervisor timer
-        "j {default}", // reserved
-        "j {mtimer}",  // machine    timer
-        "j {default}", // reserved
-        "j {default}", // supervisor external
-        "j {default}", // reserved
-        "j {default}", // machine    external
-        ".option pop",
-        default = sym trap_entry,
-        msoft   = sym msoft,
-        mtimer  = sym mtimer,
-        options(noreturn)
-    )
-}
-
-/// Machine timer interrupt handler.
-/// Saves context, clears mtimecmp, sets STIP bit, and restores context.
-///
-/// # Safety
-///
-/// This is a naked function that directly manipulates registers and stack.
-#[naked]
-unsafe extern "C" fn mtimer() {
-    asm!(
-        // Switch stacks: sp <-> mscratch
-        "   csrrw sp, mscratch, sp",
-        // Save registers to stack
-        "   addi   sp, sp, -30*8",
-        "   sd     ra, 0*8(sp)
-            sd      gp, 2*8(sp)
-            sd      tp, 3*8(sp)
-            sd      t0, 4*8(sp)
-            sd      t1, 5*8(sp)
-            sd      t2, 6*8(sp)
-            sd      s0, 7*8(sp)
-            sd      s1, 8*8(sp)
-            sd      a0, 9*8(sp)
-            sd      a1, 10*8(sp)
-            sd      a2, 11*8(sp)
-            sd      a3, 12*8(sp)
-            sd      a4, 13*8(sp)
-            sd      a5, 14*8(sp)
-            sd      a6, 15*8(sp)
-            sd      a7, 16*8(sp)
-            sd      s2, 17*8(sp)
-            sd      s3, 18*8(sp)
-            sd      s4, 19*8(sp)
-            sd      s5, 20*8(sp)
-            sd      s6, 21*8(sp)
-            sd      s7, 22*8(sp)
-            sd      s8, 23*8(sp)
-            sd      s9, 24*8(sp)
-            sd     s10, 25*8(sp)
-            sd     s11, 26*8(sp)
-            sd      t3, 27*8(sp)
-            sd      t4, 28*8(sp)
-            sd      t5, 29*8(sp)
-            sd      t6, 1*8(sp)",
-        // Clear machine timer compare register
-        "    call  {clear_mtime}",
-        // Set supervisor timer interrupt pending bit
-        "   li    a0, {mip_stip}
-            csrrs zero, mip, a0
-        ",
-        // Restore registers from stack
-        "   ld     ra, 0*8(sp)
-            ld      gp, 2*8(sp)
-            ld      tp, 3*8(sp)
-            ld      t0, 4*8(sp)
-            ld      t1, 5*8(sp)
-            ld      t2, 6*8(sp)
-            ld      s0, 7*8(sp)
-            ld      s1, 8*8(sp)
-            ld      a0, 9*8(sp)
-            ld      a1, 10*8(sp)
-            ld      a2, 11*8(sp)
-            ld      a3, 12*8(sp)
-            ld      a4, 13*8(sp)
-            ld      a5, 14*8(sp)
-            ld      a6, 15*8(sp)
-            ld      a7, 16*8(sp)
-            ld      s2, 17*8(sp)
-            ld      s3, 18*8(sp)
-            ld      s4, 19*8(sp)
-            ld      s5, 20*8(sp)
-            ld      s6, 21*8(sp)
-            ld      s7, 22*8(sp)
-            ld      s8, 23*8(sp)
-            ld      s9, 24*8(sp)
-            ld     s10, 25*8(sp)
-            ld     s11, 26*8(sp)
-            ld      t3, 27*8(sp)
-            ld      t4, 28*8(sp)
-            ld      t5, 29*8(sp)
-            ld      t6, 1*8(sp)",
-        "   addi   sp, sp, 30*8",
-        // Switch stacks back: sp <-> mscratch
-        "   csrrw sp, mscratch, sp",
-        // Return from machine mode
-        "   mret",
-        mip_stip    = const 1 << 5,
-        clear_mtime = sym ipi::clear_mtime,
-        options(noreturn)
-    )
-}
-
-/// Machine software interrupt handler.
-///
-/// Handles inter-processor interrupts.
-#[naked]
-pub unsafe extern "C" fn msoft() -> ! {
-    asm!(
-        ".align 2",
-        // Switch stacks
-        "csrrw  sp, mscratch, sp",
-        // Allocate stack space
-        "addi   sp, sp, -32*8",
-        // Save registers
-        "sd     ra, 0*8(sp)
-        sd      gp, 2*8(sp)
-        sd      tp, 3*8(sp)
-        sd      t0, 4*8(sp)
-        sd      t1, 5*8(sp)
-        sd      t2, 6*8(sp)
-        sd      s0, 7*8(sp)
-        sd      s1, 8*8(sp)
-        sd      a0, 9*8(sp)
-        sd      a1, 10*8(sp)
-        sd      a2, 11*8(sp)
-        sd      a3, 12*8(sp)
-        sd      a4, 13*8(sp)
-        sd      a5, 14*8(sp)
-        sd      a6, 15*8(sp)
-        sd      a7, 16*8(sp)
-        sd      s2, 17*8(sp)
-        sd      s3, 18*8(sp)
-        sd      s4, 19*8(sp)
-        sd      s5, 20*8(sp)
-        sd      s6, 21*8(sp)
-        sd      s7, 22*8(sp)
-        sd      s8, 23*8(sp)
-        sd      s9, 24*8(sp)
-        sd     s10, 25*8(sp)
-        sd     s11, 26*8(sp)
-        sd      t3, 27*8(sp)
-        sd      t4, 28*8(sp)
-        sd      t5, 29*8(sp)
-        sd      t6, 30*8(sp)",
-        // Save mepc and mscratch
-        "csrr   t0, mepc
-        sd      t0, 31*8(sp)",
-        "csrr   t2, mscratch",
-        "sd     t2, 1*8(sp)",
-        // Call handler with context pointer
-        "mv     a0, sp",
-        "call   {msoft_handler}",
-        // Restore mepc
-        "ld     t0, 31*8(sp)
-        csrw    mepc, t0",
-        // Restore registers
-        "ld     ra, 0*8(sp)
-        ld      gp, 2*8(sp)
-        ld      tp, 3*8(sp)
-        ld      t0, 4*8(sp)
-        ld      t1, 5*8(sp)
-        ld      t2, 6*8(sp)
-        ld      s0, 7*8(sp)
-        ld      s1, 8*8(sp)
-        ld      a0, 9*8(sp)
-        ld      a1, 10*8(sp)
-        ld      a2, 11*8(sp)
-        ld      a3, 12*8(sp)
-        ld      a4, 13*8(sp)
-        ld      a5, 14*8(sp)
-        ld      a6, 15*8(sp)
-        ld      a7, 16*8(sp)
-        ld      s2, 17*8(sp)
-        ld      s3, 18*8(sp)
-        ld      s4, 19*8(sp)
-        ld      s5, 20*8(sp)
-        ld      s6, 21*8(sp)
-        ld      s7, 22*8(sp)
-        ld      s8, 23*8(sp)
-        ld      s9, 24*8(sp)
-        ld     s10, 25*8(sp)
-        ld     s11, 26*8(sp)
-        ld      t3, 27*8(sp)
-        ld      t4, 28*8(sp)
-        ld      t5, 29*8(sp)
-        ld      t6, 30*8(sp)",
-        // Restore stack pointer
-        "addi   sp, sp, 32*8",
-        // Switch stacks back
-        "csrrw  sp, mscratch, sp",
-        // Return from machine mode
-        "mret",
-        msoft_handler = sym msoft_handler,
-        options(noreturn)
-    );
-}
-
-/// Machine software interrupt handler implementation.
-///
-/// Handles HSM (Hart State Management) and RFence operations.
-pub extern "C" fn msoft_handler(ctx: &mut SupervisorContext) {
-    #[inline(always)]
-    fn boot(ctx: &mut SupervisorContext, start_addr: usize, opaque: usize) {
-        unsafe {
-            sstatus::clear_sie();
-            satp::write(0);
-        }
-        ctx.a0 = current_hartid();
-        ctx.a1 = opaque;
-        ctx.mepc = start_addr;
-    }
-
-    match local_hsm().start() {
-        // Handle HSM Start
-        Ok(next_stage) => {
-            ipi::clear_msip();
-            unsafe {
-                mstatus::set_mpie();
-                mstatus::set_mpp(next_stage.next_mode);
-                mie::set_msoft();
-                mie::set_mtimer();
-            }
-            boot(ctx, next_stage.start_addr, next_stage.opaque);
-        }
-        // Handle HSM Stop
-        Err(rustsbi::spec::hsm::HART_STOP) => {
-            ipi::clear_msip();
-            unsafe {
-                mie::set_msoft();
-            }
-            riscv::asm::wfi();
-        }
-        // Handle RFence
-        _ => {
-            msoft_ipi_handler();
-        }
-    }
-}
-
-/// Handles a single remote fence operation.
-pub fn rfence_single_handler() {
-    let rfence_context = local_rfence().unwrap().get();
-    if let Some((ctx, id)) = rfence_context {
-        match ctx.op {
-            // Handle instruction fence
-            RFenceType::FenceI => unsafe {
-                asm!("fence.i");
-                rfence::remote_rfence(id).unwrap().sub();
-            },
-            // Handle virtual memory address fence
-            RFenceType::SFenceVma => {
-                // If the flush size is greater than the maximum limit then simply flush all
-                if (ctx.start_addr == 0 && ctx.size == 0)
-                    || (ctx.size == usize::MAX)
-                    || (ctx.size > TLB_FLUSH_LIMIT)
-                {
-                    unsafe {
-                        asm!("sfence.vma");
-                    }
-                } else {
-                    for offset in (0..ctx.size).step_by(PAGE_SIZE) {
-                        let addr = ctx.start_addr + offset;
-                        unsafe {
-                            asm!("sfence.vma {}", in(reg) addr);
-                        }
-                    }
-                }
-                rfence::remote_rfence(id).unwrap().sub();
-            }
-            // Handle virtual memory address fence with ASID
-            RFenceType::SFenceVmaAsid => {
-                let asid = ctx.asid;
-                // If the flush size is greater than the maximum limit then simply flush all
-                if (ctx.start_addr == 0 && ctx.size == 0)
-                    || (ctx.size == usize::MAX)
-                    || (ctx.size > TLB_FLUSH_LIMIT)
-                {
-                    unsafe {
-                        asm!("sfence.vma {}, {}", in(reg) 0, in(reg) asid);
-                    }
-                } else {
-                    for offset in (0..ctx.size).step_by(PAGE_SIZE) {
-                        let addr = ctx.start_addr + offset;
-                        unsafe {
-                            asm!("sfence.vma {}, {}", in(reg) addr, in(reg) asid);
-                        }
-                    }
-                }
-                rfence::remote_rfence(id).unwrap().sub();
-            }
-            rfencetype => {
-                error!("Unsupported RFence Type: {:?}!", rfencetype);
-            }
-        }
-    }
-}
-
-/// Process all pending remote fence operations.
-pub fn rfence_handler() {
-    while !local_rfence().unwrap().is_empty() {
-        rfence_single_handler();
-    }
-}
-
-/// Handle machine software inter-processor interrupts.
-pub fn msoft_ipi_handler() {
-    use ipi::get_and_reset_ipi_type;
-    ipi::clear_msip();
-    let ipi_type = get_and_reset_ipi_type();
-    // Handle supervisor software interrupt
-    if (ipi_type & ipi::IPI_TYPE_SSOFT) != 0 {
-        unsafe {
-            riscv::register::mip::set_ssoft();
-        }
-    }
-    // Handle fence operation
-    if (ipi_type & ipi::IPI_TYPE_FENCE) != 0 {
-        rfence_handler();
-    }
-}
-
-/// Fast trap handler for SBI calls and illegal instructions.
-pub extern "C" fn fast_handler(
-    mut ctx: FastContext,
-    a1: usize,
-    a2: usize,
-    a3: usize,
-    a4: usize,
-    a5: usize,
-    a6: usize,
-    a7: usize,
-) -> FastResult {
-    #[inline]
-    fn resume(mut ctx: FastContext, start_addr: usize, opaque: usize) -> FastResult {
-        unsafe {
-            sstatus::clear_sie();
-            satp::write(0);
-        }
-        ctx.regs().a[0] = current_hartid();
-        ctx.regs().a[1] = opaque;
-        ctx.regs().pc = start_addr;
-        ctx.call(2)
-    }
-    match mcause::read().cause() {
-        // Handle SBI calls
-        T::Exception(E::SupervisorEnvCall) => {
-            use sbi_spec::{base, hsm, legacy};
-            let mut ret = unsafe {
-                PLATFORM
-                    .sbi
-                    .handle_ecall(a7, a6, [ctx.a0(), a1, a2, a3, a4, a5])
-            };
-            if ret.is_ok() {
-                match (a7, a6) {
-                    // Handle non-retentive suspend
-                    (hsm::EID_HSM, hsm::HART_SUSPEND)
-                        if matches!(ctx.a0() as u32, hsm::suspend_type::NON_RETENTIVE) =>
-                    {
-                        return resume(ctx, a1, a2);
-                    }
-                    // Handle legacy console probe
-                    (base::EID_BASE, base::PROBE_EXTENSION)
-                        if matches!(
-                            ctx.a0(),
-                            legacy::LEGACY_CONSOLE_PUTCHAR | legacy::LEGACY_CONSOLE_GETCHAR
-                        ) =>
-                    {
-                        ret.value = 1;
-                    }
-                    _ => {}
-                }
-            } else {
-                match a7 {
-                    legacy::LEGACY_CONSOLE_PUTCHAR => {
-                        ret.error = console::putchar(ctx.a0());
-                        ret.value = a1;
-                    }
-                    legacy::LEGACY_CONSOLE_GETCHAR => {
-                        ret.error = console::getchar();
-                        ret.value = a1;
-                    }
-                    _ => {}
-                }
-            }
-            ctx.regs().a = [ret.error, ret.value, a2, a3, a4, a5, a6, a7];
-            mepc::write(mepc::read() + 4);
-            ctx.restore()
-        }
-        // Handle illegal instructions
-        T::Exception(E::IllegalInstruction) => {
-            if mstatus::read().mpp() == mstatus::MPP::Machine {
-                panic!("Cannot handle illegal instruction exception from M-MODE");
-            }
-
-            ctx.regs().a = [ctx.a0(), a1, a2, a3, a4, a5, a6, a7];
-            if !illegal_instruction_handler(&mut ctx) {
-                delegate();
-            }
-            ctx.restore()
-        }
-        // Handle other traps
-        trap => {
-            error!("-----------------------------");
-            error!("trap:    {trap:?}");
-            error!("mepc:    {:#018x}", mepc::read());
-            error!("mtval:   {:#018x}", mtval::read());
-            error!("-----------------------------");
-            panic!("Stopped with unsupported trap")
-        }
-    }
-}
-
-/// Delegate trap handling to supervisor mode.
-#[inline]
-fn delegate() {
-    use riscv::register::{mcause, mepc, mtval, scause, sepc, sstatus, stval, stvec};
-    unsafe {
-        // TODO: 当支持中断嵌套时，需要从ctx里获取mpec。当前ctx.reg().pc与mepc不一致
-        sepc::write(mepc::read());
-        scause::write(mcause::read().bits());
-        stval::write(mtval::read());
-        sstatus::clear_sie();
-        if mstatus::read().mpp() == mstatus::MPP::Supervisor {
-            sstatus::set_spp(sstatus::SPP::Supervisor);
-        } else {
-            sstatus::set_spp(sstatus::SPP::User);
-        }
-        mstatus::set_mpp(mstatus::MPP::Supervisor);
-        mepc::write(stvec::read().address());
-    }
-}
-
-/// Handle illegal instructions, particularly CSR access.
-#[inline]
-fn illegal_instruction_handler(ctx: &mut FastContext) -> bool {
-    use riscv::register::{mepc, mtval};
-    use riscv_decode::{decode, Instruction};
-
-    let inst = decode(mtval::read() as u32);
-    match inst {
-        Ok(Instruction::Csrrs(csr)) => match csr.csr() {
-            CSR_TIME => {
-                assert!(
-                    10 <= csr.rd() && csr.rd() <= 17,
-                    "Unsupported CSR rd: {}",
-                    csr.rd()
-                );
-                ctx.regs().a[(csr.rd() - 10) as usize] =
-                    unsafe { PLATFORM.sbi.ipi.as_ref() }.unwrap().get_time();
-            }
-            CSR_TIMEH => {
-                assert!(
-                    10 <= csr.rd() && csr.rd() <= 17,
-                    "Unsupported CSR rd: {}",
-                    csr.rd()
-                );
-                ctx.regs().a[(csr.rd() - 10) as usize] =
-                    unsafe { PLATFORM.sbi.ipi.as_ref() }.unwrap().get_timeh();
-            }
-            _ => return false,
-        },
-        _ => return false,
-    }
-    mepc::write(mepc::read() + 4);
-    true
-}
-
-/// Supervisor context structure containing saved register state.
-#[derive(Debug)]
-#[repr(C)]
-pub struct SupervisorContext {
-    pub ra: usize, // 0
-    pub sp: usize,
-    pub gp: usize,
-    pub tp: usize,
-    pub t0: usize,
-    pub t1: usize,
-    pub t2: usize,
-    pub s0: usize,
-    pub s1: usize,
-    pub a0: usize,
-    pub a1: usize,
-    pub a2: usize,
-    pub a3: usize,
-    pub a4: usize,
-    pub a5: usize,
-    pub a6: usize,
-    pub a7: usize,
-    pub s2: usize,
-    pub s3: usize,
-    pub s4: usize,
-    pub s5: usize,
-    pub s6: usize,
-    pub s7: usize,
-    pub s8: usize,
-    pub s9: usize,
-    pub s10: usize,
-    pub s11: usize,
-    pub t3: usize,
-    pub t4: usize,
-    pub t5: usize,
-    pub t6: usize,   // 30
-    pub mepc: usize, // 31
-}

prototyper/src/sbi/trap/boot.rs

@@ -0,0 +1,85 @@
+use crate::riscv::current_hartid;
+use crate::sbi::hsm::local_hsm;
+use crate::sbi::ipi;
+use core::arch::asm;
+use riscv::register::{mie, mstatus, satp, sstatus};
+
+/// Boot Function.
+/// After boot, this flow will never back again,
+/// so we can store a0, a1 and mepc only.
+#[naked]
+pub unsafe extern "C" fn boot() -> ! {
+    asm!(
+        ".align 2",
+        // Switch stacks
+        "csrrw  sp, mscratch, sp",
+        // Allocate stack space
+        "addi   sp, sp, -3*8",
+        // Call handler with context pointer
+        "mv     a0, sp",
+        "call   {boot_handler}",
+        // Restore mepc
+        "ld     t0, 0*8(sp)
+        csrw    mepc, t0",
+        // Restore registers
+        "
+        ld      a0, 1*8(sp)
+        ld      a1, 2*8(sp)",
+        // Restore stack pointer
+        "addi   sp, sp, 3*8",
+        // Switch stacks back
+        "csrrw  sp, mscratch, sp",
+        // Return from machine mode
+        "mret",
+        boot_handler = sym boot_handler,
+        options(noreturn)
+    );
+}
+
+/// Boot Handler.
+pub extern "C" fn boot_handler(ctx: &mut BootContext) {
+    #[inline(always)]
+    fn boot(ctx: &mut BootContext, start_addr: usize, opaque: usize) {
+        unsafe {
+            sstatus::clear_sie();
+            satp::write(0);
+        }
+        ctx.a0 = current_hartid();
+        ctx.a1 = opaque;
+        ctx.mepc = start_addr;
+    }
+
+    match local_hsm().start() {
+        // Handle HSM Start
+        Ok(next_stage) => {
+            ipi::clear_msip();
+            unsafe {
+                mstatus::set_mpie();
+                mstatus::set_mpp(next_stage.next_mode);
+                mie::set_msoft();
+                mie::set_mtimer();
+            }
+            boot(ctx, next_stage.start_addr, next_stage.opaque);
+        }
+        // Handle HSM Stop
+        Err(rustsbi::spec::hsm::HART_STOP) => {
+            ipi::clear_msip();
+            unsafe {
+                mie::set_msoft();
+            }
+            riscv::asm::wfi();
+        }
+        _ => {
+            unreachable!("Boot stage hsm should be start or stop.");
+        }
+    }
+}
+
+/// Boot context structure containing saved register state.
+#[derive(Debug)]
+#[repr(C)]
+pub struct BootContext {
+    pub mepc: usize, // 0
+    pub a0: usize,
+    pub a1: usize, // 2
+}

prototyper/src/sbi/trap/handler.rs

@@ -0,0 +1,181 @@
+use fast_trap::{FastContext, FastResult};
+use riscv::register::{mepc, mie, mstatus, satp, sstatus};
+use rustsbi::RustSBI;
+
+use crate::platform::PLATFORM;
+use crate::riscv::csr::{CSR_TIME, CSR_TIMEH};
+use crate::riscv::current_hartid;
+use crate::sbi::console;
+use crate::sbi::hsm::local_hsm;
+use crate::sbi::ipi;
+use crate::sbi::rfence;
+
+#[inline]
+pub fn switch(mut ctx: FastContext, start_addr: usize, opaque: usize) -> FastResult {
+    unsafe {
+        sstatus::clear_sie();
+        satp::write(0);
+    }
+
+    ctx.regs().a[0] = current_hartid();
+    ctx.regs().a[1] = opaque;
+    ctx.regs().pc = start_addr;
+    ctx.call(2)
+}
+
+/// Handle machine software inter-processor interrupts.
+#[inline]
+pub fn msoft_ipi_handler() {
+    use ipi::get_and_reset_ipi_type;
+    ipi::clear_msip();
+    let ipi_type = get_and_reset_ipi_type();
+    // Handle supervisor software interrupt
+    if (ipi_type & ipi::IPI_TYPE_SSOFT) != 0 {
+        unsafe {
+            riscv::register::mip::set_ssoft();
+        }
+    }
+    // Handle fence operation
+    if (ipi_type & ipi::IPI_TYPE_FENCE) != 0 {
+        rfence::rfence_handler();
+    }
+}
+
+#[inline]
+pub fn msoft_handler(ctx: FastContext) -> FastResult {
+    match local_hsm().start() {
+        // Handle HSM Start
+        Ok(next_stage) => {
+            ipi::clear_msip();
+            unsafe {
+                mstatus::set_mpie();
+                mstatus::set_mpp(next_stage.next_mode);
+                mie::set_msoft();
+                mie::set_mtimer();
+            }
+            switch(ctx, next_stage.start_addr, next_stage.opaque)
+        }
+        // Handle HSM Stop
+        Err(rustsbi::spec::hsm::HART_STOP) => {
+            ipi::clear_msip();
+            unsafe {
+                mie::set_msoft();
+            }
+            riscv::asm::wfi();
+            ctx.restore()
+        }
+        // Handle RFence
+        _ => {
+            msoft_ipi_handler();
+            ctx.restore()
+        }
+    }
+}
+
+#[inline]
+pub fn sbi_call_handler(
+    mut ctx: FastContext,
+    a1: usize,
+    a2: usize,
+    a3: usize,
+    a4: usize,
+    a5: usize,
+    a6: usize,
+    a7: usize,
+) -> FastResult {
+    use sbi_spec::{base, hsm, legacy};
+    let mut ret = unsafe {
+        PLATFORM
+            .sbi
+            .handle_ecall(a7, a6, [ctx.a0(), a1, a2, a3, a4, a5])
+    };
+    if ret.is_ok() {
+        match (a7, a6) {
+            // Handle non-retentive suspend
+            (hsm::EID_HSM, hsm::HART_SUSPEND)
+                if matches!(ctx.a0() as u32, hsm::suspend_type::NON_RETENTIVE) =>
+            {
+                return switch(ctx, a1, a2);
+            }
+            // Handle legacy console probe
+            (base::EID_BASE, base::PROBE_EXTENSION)
+                if matches!(
+                    ctx.a0(),
+                    legacy::LEGACY_CONSOLE_PUTCHAR | legacy::LEGACY_CONSOLE_GETCHAR
+                ) =>
+            {
+                ret.value = 1;
+            }
+            _ => {}
+        }
+    } else {
+        match a7 {
+            legacy::LEGACY_CONSOLE_PUTCHAR => {
+                ret.error = console::putchar(ctx.a0());
+                ret.value = a1;
+            }
+            legacy::LEGACY_CONSOLE_GETCHAR => {
+                ret.error = console::getchar();
+                ret.value = a1;
+            }
+            _ => {}
+        }
+    }
+    ctx.regs().a = [ret.error, ret.value, a2, a3, a4, a5, a6, a7];
+    mepc::write(mepc::read() + 4);
+    ctx.restore()
+}
+
+/// Delegate trap handling to supervisor mode.
+#[inline]
+pub fn delegate(ctx: &mut FastContext) {
+    use riscv::register::{mcause, mepc, mtval, scause, sepc, sstatus, stval, stvec};
+    unsafe {
+        sepc::write(ctx.regs().pc);
+        scause::write(mcause::read().bits());
+        stval::write(mtval::read());
+        sstatus::clear_sie();
+        if mstatus::read().mpp() == mstatus::MPP::Supervisor {
+            sstatus::set_spp(sstatus::SPP::Supervisor);
+        } else {
+            sstatus::set_spp(sstatus::SPP::User);
+        }
+        mstatus::set_mpp(mstatus::MPP::Supervisor);
+        mepc::write(stvec::read().address());
+    }
+}
+
+/// Handle illegal instructions, particularly CSR access.
+#[inline]
+pub fn illegal_instruction_handler(ctx: &mut FastContext) -> bool {
+    use riscv::register::{mepc, mtval};
+    use riscv_decode::{decode, Instruction};
+
+    let inst = decode(mtval::read() as u32);
+    match inst {
+        Ok(Instruction::Csrrs(csr)) => match csr.csr() {
+            CSR_TIME => {
+                assert!(
+                    10 <= csr.rd() && csr.rd() <= 17,
+                    "Unsupported CSR rd: {}",
+                    csr.rd()
+                );
+                ctx.regs().a[(csr.rd() - 10) as usize] =
+                    unsafe { PLATFORM.sbi.ipi.as_ref() }.unwrap().get_time();
+            }
+            CSR_TIMEH => {
+                assert!(
+                    10 <= csr.rd() && csr.rd() <= 17,
+                    "Unsupported CSR rd: {}",
+                    csr.rd()
+                );
+                ctx.regs().a[(csr.rd() - 10) as usize] =
+                    unsafe { PLATFORM.sbi.ipi.as_ref() }.unwrap().get_timeh();
+            }
+            _ => return false,
+        },
+        _ => return false,
+    }
+    mepc::write(mepc::read() + 4);
+    true
+}

prototyper/src/sbi/trap/mod.rs

@@ -0,0 +1,71 @@
+pub mod boot;
+pub mod handler;
+
+use fast_trap::{FastContext, FastResult};
+use riscv::register::{
+    mcause::{self, Exception as E, Interrupt, Trap as T},
+    mepc, mip, mstatus, mtval,
+};
+
+/// Fast trap handler for all trap.
+pub extern "C" fn fast_handler(
+    mut ctx: FastContext,
+    a1: usize,
+    a2: usize,
+    a3: usize,
+    a4: usize,
+    a5: usize,
+    a6: usize,
+    a7: usize,
+) -> FastResult {
+    // Save mepc into context
+    ctx.regs().pc = mepc::read();
+
+    let save_regs = |ctx: &mut FastContext| {
+        ctx.regs().a = [ctx.a0(), a1, a2, a3, a4, a5, a6, a7];
+    };
+
+    match mcause::read().cause() {
+        // Handle Msoft
+        T::Interrupt(Interrupt::MachineSoft) => {
+            save_regs(&mut ctx);
+            handler::msoft_handler(ctx)
+        }
+        // Handle MTimer
+        T::Interrupt(Interrupt::MachineTimer) => {
+            use crate::sbi::ipi;
+
+            ipi::clear_mtime();
+            unsafe {
+                mip::clear_stimer();
+            }
+            save_regs(&mut ctx);
+            ctx.restore()
+        }
+        // Handle SBI calls
+        T::Exception(E::SupervisorEnvCall) => {
+            handler::sbi_call_handler(ctx, a1, a2, a3, a4, a5, a6, a7)
+        }
+        // Handle illegal instructions
+        T::Exception(E::IllegalInstruction) => {
+            if mstatus::read().mpp() == mstatus::MPP::Machine {
+                panic!("Cannot handle illegal instruction exception from M-MODE");
+            }
+
+            save_regs(&mut ctx);
+            if !handler::illegal_instruction_handler(&mut ctx) {
+                handler::delegate(&mut ctx);
+            }
+            ctx.restore()
+        }
+        // Handle other traps
+        trap => {
+            error!("-----------------------------");
+            error!("trap:    {trap:?}");
+            error!("mepc:    {:#018x}", mepc::read());
+            error!("mtval:   {:#018x}", mtval::read());
+            error!("-----------------------------");
+            panic!("Stopped with unsupported trap")
+        }
+    }
+}