feat: Improve code quality and enhance safety features

This commit includes several improvements to the codebase, focusing on adding detailed comments, refactoring the FIFO queue, and enhancing the functionality and safety of RFenceCell.

- Added comprehensive comments to enhance code readability and maintainability across the project.
- Refactored the FIFO (First In First Out) queue to optimize its logic structure, improving overall performance and efficiency.
- Renamed `RFenceCTX` to `RFenceContext` to adhere to Rust naming conventions, promoting better code consistency.
- Introduced `drop(queue)` before `trap::rfence_single_handler()` to mitigate deadlock risks by ensuring mutexes are released during full queue scenarios, enhancing thread safety.
- Implemented input address validation in `remote_sfence` functions to ensure the validity of address ranges

Signed-off-by: Zongyao Chen <solar1s@163.com>

prototyper/Cargo.toml

@@ -15,7 +15,7 @@ panic-halt = "0.2.0"
 riscv = "0.11.1"
 rustsbi = { version = "0.4.0", features = ["machine"] }
 sbi-spec = { version = "0.0.7", features = ["legacy"] }
-serde = { version = "1.0.202", default-features = false, features = ["derive"]}
+serde = { version = "1.0.202", default-features = false, features = ["derive"] }
 serde-device-tree = { git = "https://github.com/rustsbi/serde-device-tree", default-features = false }
 sifive-test-device = "0.0.0"
 spin = "0.9.8"

prototyper/src/dt.rs

@@ -4,47 +4,67 @@ use serde_device_tree::{
     Dtb, DtbPtr,
 };
 
+/// Root device tree structure containing system information.
 #[derive(Deserialize)]
 pub struct Tree<'a> {
+    /// Optional model name string.
     pub model: Option<StrSeq<'a>>,
+    /// Chosen node containing boot parameters.
     pub chosen: Chosen<'a>,
+    /// CPU information.
     pub cpus: Cpus<'a>,
+    /// System-on-chip components.
     pub soc: Soc<'a>,
 }
 
+/// Chosen node containing boot parameters.
 #[derive(Deserialize)]
 #[serde(rename_all = "kebab-case")]
 pub struct Chosen<'a> {
+    /// Path to stdout device.
     pub stdout_path: StrSeq<'a>,
 }
 
+/// CPU information container.
 #[derive(Deserialize)]
 #[serde(rename_all = "kebab-case")]
 pub struct Cpus<'a> {
+    /// Sequence of CPU nodes.
     pub cpu: NodeSeq<'a>,
 }
 
+/// Individual CPU node information.
 #[derive(Deserialize, Debug)]
 pub struct Cpu<'a> {
+    /// RISC-V ISA extensions supported by this CPU.
     #[serde(rename = "riscv,isa-extensions")]
     pub isa: Option<StrSeq<'a>>,
+    /// CPU register information.
     pub reg: Reg<'a>,
 }
 
+/// System-on-chip components.
 #[derive(Deserialize, Debug)]
 pub struct Soc<'a> {
+    /// Serial (UART) device nodes.
     pub serial: Option<NodeSeq<'a>>,
+    /// Test device nodes.
     pub test: Option<NodeSeq<'a>>,
+    /// CLINT (Core Local Interruptor) nodes.
     pub clint: Option<NodeSeq<'a>>,
 }
 
+/// Generic device node information.
 #[allow(unused)]
 #[derive(Deserialize, Debug)]
 pub struct Device<'a> {
+    /// Device register information.
     pub reg: Reg<'a>,
 }
 
+/// Errors that can occur during device tree parsing.
 pub enum ParseDeviceTreeError {
+    /// Invalid device tree format.
     Format,
 }
 

prototyper/src/fail.rs

@@ -8,6 +8,7 @@ use crate::platform::dynamic;
 #[cfg(not(feature = "payload"))]
 use riscv::register::mstatus;
 
+/// Handles device tree format parsing errors by logging and resetting.
 #[cold]
 pub fn device_tree_format(err: dt::ParseDeviceTreeError) -> Dtb {
     match err {
@@ -16,12 +17,14 @@ pub fn device_tree_format(err: dt::ParseDeviceTreeError) -> Dtb {
     reset::fail()
 }
 
+/// Handles device tree deserialization errors by logging and resetting.
 #[cold]
 pub fn device_tree_deserialize<'a>(err: serde_device_tree::error::Error) -> Tree<'a> {
     error!("Device tree deserialization error: {:?}", err);
     reset::fail()
 }
 
+/// Handles invalid dynamic information data by logging details and resetting.
 #[cold]
 #[cfg(not(feature = "payload"))]
 pub fn invalid_dynamic_data(err: dynamic::DynamicError) -> (mstatus::MPP, usize) {
@@ -45,6 +48,7 @@ pub fn invalid_dynamic_data(err: dynamic::DynamicError) -> (mstatus::MPP, usize)
     reset::fail()
 }
 
+/// Handles case where dynamic information is not available by logging details and resetting.
 #[cold]
 #[cfg(not(feature = "payload"))]
 pub fn no_dynamic_info_available(err: dynamic::DynamicReadError) -> dynamic::DynamicInfo {
@@ -74,6 +78,9 @@ pub fn no_dynamic_info_available(err: dynamic::DynamicReadError) -> dynamic::Dyn
     reset::fail()
 }
 
+/// Fallback function that returns default dynamic info with boot_hart set to MAX.
+///
+/// Used when dynamic info read fails but execution should continue.
 #[cold]
 #[cfg(not(feature = "payload"))]
 pub fn use_lottery(_err: dynamic::DynamicReadError) -> dynamic::DynamicInfo {

prototyper/src/main.rs

@@ -39,43 +39,44 @@ pub const R_RISCV_RELATIVE: usize = 3;
 
 #[no_mangle]
 extern "C" fn rust_main(_hart_id: usize, opaque: usize, nonstandard_a2: usize) {
-    // parse dynamic information
+    // Track whether SBI is initialized and ready.
     static SBI_READY: AtomicBool = AtomicBool::new(false);
 
     let boot_hart_info = platform::get_boot_hart(opaque, nonstandard_a2);
-    // boot hart task entry
+    // boot hart task entry.
     if boot_hart_info.is_boot_hart {
         let fdt_addr = boot_hart_info.fdt_address;
 
         // 1. Init FDT
-        // parse the device tree
-        // TODO: shoule remove `fail:device_tree_format`
+        // parse the device tree.
+        // TODO: shoule remove `fail:device_tree_format`.
         let dtb = dt::parse_device_tree(fdt_addr).unwrap_or_else(fail::device_tree_format);
         let dtb = dtb.share();
 
-        // TODO: should remove `fail:device_tree_deserialize`
+        // TODO: should remove `fail:device_tree_deserialize`.
         let tree =
             serde_device_tree::from_raw_mut(&dtb).unwrap_or_else(fail::device_tree_deserialize);
 
         // 2. Init device
-        // TODO: The device base address should be find in a better way
+        // TODO: The device base address should be find in a better way.
         let console_base = tree.soc.serial.unwrap().iter().next().unwrap();
         let clint_device = tree.soc.clint.unwrap().iter().next().unwrap();
         let cpu_num = tree.cpus.cpu.len();
         let console_base_address = console_base.at();
         let ipi_base_address = clint_device.at();
 
-        // Set reset device if found it
+        // Initialize reset device if present.
         if let Some(test) = tree.soc.test {
             let reset_device = test.iter().next().unwrap();
             let reset_base_address = reset_device.at();
             board::reset_dev_init(usize::from_str_radix(reset_base_address, 16).unwrap());
         }
 
+        // Initialize console and IPI devices.
         board::console_dev_init(usize::from_str_radix(console_base_address, 16).unwrap());
         board::ipi_dev_init(usize::from_str_radix(ipi_base_address, 16).unwrap());
 
-        // 3. Init SBI
+        // 3. Init the SBI implementation
         unsafe {
             SBI_IMPL = MaybeUninit::new(SBI {
                 console: Some(SbiConsole::new(&UART)),
@@ -85,12 +86,14 @@ extern "C" fn rust_main(_hart_id: usize, opaque: usize, nonstandard_a2: usize) {
                 rfence: Some(SbiRFence),
             });
         }
-        // 设置陷入栈
+
+        // Setup trap handling.
         trap_stack::prepare_for_trap();
         extensions::init(&tree.cpus.cpu);
         SBI_READY.swap(true, Ordering::AcqRel);
+
         // 4. Init Logger
-        logger::Logger::init();
+        logger::Logger::init().unwrap();
 
         info!("RustSBI version {}", rustsbi::VERSION);
         rustsbi::LOGO.lines().for_each(|line| info!("{}", line));
@@ -120,9 +123,11 @@ extern "C" fn rust_main(_hart_id: usize, opaque: usize, nonstandard_a2: usize) {
             pmpaddr1::write(usize::MAX >> 2);
         }
 
+        // Get boot information and prepare for kernel entry.
         let boot_info = platform::get_boot_info(nonstandard_a2);
         let (mpp, next_addr) = (boot_info.mpp, boot_info.next_address);
-        // 设置内核入口
+
+        // Start kernel.
         local_remote_hsm().start(NextStage {
             start_addr: next_addr,
             next_mode: mpp,
@@ -136,7 +141,9 @@ extern "C" fn rust_main(_hart_id: usize, opaque: usize, nonstandard_a2: usize) {
             mpp
         );
     } else {
-        // TODO: PMP configuration needs to be obtained through the memory range in the device tree
+        // Non-boot hart initialization path.
+
+        // TODO: PMP configuration needs to be obtained through the memory range in the device tree.
         use riscv::register::*;
         unsafe {
             pmpcfg0::set_pmp(0, Range::OFF, Permission::NONE, false);
@@ -145,23 +152,29 @@ extern "C" fn rust_main(_hart_id: usize, opaque: usize, nonstandard_a2: usize) {
             pmpaddr1::write(usize::MAX >> 2);
         }
 
-        // 设置陷入栈
+        // Setup trap handling.
         trap_stack::prepare_for_trap();
 
-        // waiting for sbi ready
+        // Wait for boot hart to complete SBI initialization.
         while !SBI_READY.load(Ordering::Relaxed) {
             core::hint::spin_loop()
         }
     }
 
+    // Clear all pending IPIs.
     ipi::clear_all();
+
+    // Configure CSRs and trap handling.
     unsafe {
+        // Delegate all interrupts and exceptions to supervisor mode.
         asm!("csrw mideleg,    {}", in(reg) !0);
         asm!("csrw medeleg,    {}", in(reg) !0);
         asm!("csrw mcounteren, {}", in(reg) !0);
         use riscv::register::{medeleg, mtvec};
+        // Keep supervisor environment calls and illegal instructions in M-mode.
         medeleg::clear_supervisor_env_call();
         medeleg::clear_illegal_instruction();
+        // Configure environment features based on available extensions.
         if hart_extension_probe(current_hartid(), Extension::Sstc) {
             menvcfg::set_bits(
                 menvcfg::STCE | menvcfg::CBIE_INVALIDATE | menvcfg::CBCFE | menvcfg::CBZE,
@@ -169,6 +182,7 @@ extern "C" fn rust_main(_hart_id: usize, opaque: usize, nonstandard_a2: usize) {
         } else {
             menvcfg::set_bits(menvcfg::CBIE_INVALIDATE | menvcfg::CBCFE | menvcfg::CBZE);
         }
+        // Set up vectored trap handling.
         mtvec::write(trap_vec as _, mtvec::TrapMode::Vectored);
     }
 }
@@ -178,39 +192,39 @@ extern "C" fn rust_main(_hart_id: usize, opaque: usize, nonstandard_a2: usize) {
 #[export_name = "_start"]
 unsafe extern "C" fn start() -> ! {
     core::arch::asm!(
-        // 1. Turn off interrupt
+        // 1. Turn off interrupt.
         "   csrw    mie, zero",
-        // 2. Initialize programming langauge runtime
-        // only clear bss if hartid matches preferred boot hart id
+        // 2. Initialize programming langauge runtime.
+        // only clear bss if hartid matches preferred boot hart id.
         "   csrr    t0, mhartid",
         "   bne     t0, zero, 4f",
         "   call    {relocation_update}",
         "1:",
-        // 3. Hart 0 clear bss segment
+        // 3. Hart 0 clear bss segment.
         "   lla     t0, sbss
             lla     t1, ebss
          2: bgeu    t0, t1, 3f
             sd      zero, 0(t0)
             addi    t0, t0, 8
             j       2b",
-        "3: ", // Hart 0 set bss ready signal
+        "3: ", // Hart 0 set bss ready signal.
         "   lla     t0, 6f
             li      t1, 1
             amoadd.w t0, t1, 0(t0)
             j       5f",
-        "4:", // Other harts are waiting for bss ready signal
+        "4:", // Other harts are waiting for bss ready signal.
         "   li      t1, 1
             lla     t0, 6f
             lw      t0, 0(t0)
             bne     t0, t1, 4b", 
         "5:",
-         // 4. Prepare stack for each hart
+         // 4. Prepare stack for each hart.
         "   call    {locate_stack}",
         "   call    {main}",
         "   csrw    mscratch, sp",
         "   j       {hart_boot}",
         "  .balign  4",
-        "6:",  // bss ready signal
+        "6:",  // bss ready signal.
         "  .word    0",
         relocation_update = sym relocation_update,
         locate_stack = sym trap_stack::locate,
@@ -220,15 +234,16 @@ unsafe extern "C" fn start() -> ! {
     )
 }
 
+// Handle relocations for position-independent code
 #[naked]
 unsafe extern "C" fn relocation_update() {
     asm!(
-        // Get load offset
+        // Get load offset.
         "   li t0, {START_ADDRESS}",
         "   lla t1, .text.entry",
         "   sub t2, t1, t0",
 
-        // Foreach rela.dyn and update relocation
+        // Foreach rela.dyn and update relocation.
         "   lla t0, __rel_dyn_start",
         "   lla t1, __rel_dyn_end",
         "   li  t3, {R_RISCV_RELATIVE}",

prototyper/src/platform/dynamic.rs

@@ -9,20 +9,33 @@ use crate::riscv_spec::current_hartid;
 
 use riscv::register::mstatus;
 
+/// Gets boot hart information based on opaque and nonstandard_a2 parameters.
+///
+/// Returns a BootHart struct containing FDT address and whether this is the boot hart.
 pub fn get_boot_hart(opaque: usize, nonstandard_a2: usize) -> BootHart {
+    // Track whether this is the first hart to boot
     static GENESIS: AtomicBool = AtomicBool::new(true);
+
     let info = read_paddr(nonstandard_a2).unwrap_or_else(fail::use_lottery);
+
+    // Determine if this is the boot hart based on hart ID
     let is_boot_hart = if info.boot_hart == usize::MAX {
+        // If boot_hart is MAX, use atomic bool to determine first hart
         GENESIS.swap(false, Ordering::AcqRel)
     } else {
+        // Otherwise check if current hart matches designated boot hart
         current_hartid() == info.boot_hart
     };
+
     BootHart {
         fdt_address: opaque,
         is_boot_hart,
     }
 }
 
+/// Gets boot information from nonstandard_a2 parameter.
+///
+/// Returns BootInfo containing next stage address and privilege mode.
 pub fn get_boot_info(nonstandard_a2: usize) -> BootInfo {
     let dynamic_info = read_paddr(nonstandard_a2).unwrap_or_else(fail::no_dynamic_info_available);
     let (mpp, next_addr) = mpp_next_addr(&dynamic_info).unwrap_or_else(fail::invalid_dynamic_data);
@@ -58,26 +71,31 @@ const NEXT_ADDR_VALID_ADDRESSES: Range<usize> = 0x80000000..0x90000000;
 pub(crate) const MAGIC: usize = 0x4942534f;
 const SUPPORTED_VERSION: Range<usize> = 2..3;
 
+/// Error type for dynamic info read failures.
 pub struct DynamicReadError {
     pub bad_paddr: Option<usize>,
     pub bad_magic: Option<usize>,
     pub bad_version: Option<usize>,
 }
 
-// TODO unconstrained lifetime
+// TODO: unconstrained lifetime
+/// Reads dynamic info from physical address.
+///
+/// Returns Result containing DynamicInfo or error details.
 pub fn read_paddr(paddr: usize) -> Result<DynamicInfo, DynamicReadError> {
     let mut error = DynamicReadError {
         bad_paddr: None,
         bad_magic: None,
         bad_version: None,
     };
-    // check pointer before dereference
+    // check pointer before dereference.
     if DYNAMIC_INFO_INVALID_ADDRESSES == paddr {
         error.bad_paddr = Some(paddr);
         return Err(error);
     }
     let ans = unsafe { *(paddr as *const DynamicInfo) };
 
+    // Validate magic number and version.
     if ans.magic != MAGIC {
         error.bad_magic = Some(ans.magic);
     }
@@ -90,12 +108,16 @@ pub fn read_paddr(paddr: usize) -> Result<DynamicInfo, DynamicReadError> {
     Ok(ans)
 }
 
+/// Error type for dynamic info validation failures.
 pub struct DynamicError<'a> {
     pub invalid_mpp: bool,
     pub invalid_next_addr: bool,
     pub bad_info: &'a DynamicInfo,
 }
 
+/// Validates and extracts privilege mode and next address from dynamic info.
+///
+/// Returns Result containing tuple of (MPP, next_addr) or error details.
 pub fn mpp_next_addr(info: &DynamicInfo) -> Result<(mstatus::MPP, usize), DynamicError> {
     let mut error = DynamicError {
         invalid_mpp: false,

prototyper/src/riscv_spec.rs

@@ -1,40 +1,60 @@
 #![allow(unused)]
 
+/// CSR addresses for timer registers.
+///
+/// Time value (lower 32 bits).
 pub const CSR_TIME: u32 = 0xc01;
+/// Time value (upper 32 bits).
 pub const CSR_TIMEH: u32 = 0xc81;
+/// Supervisor timer compare value.
 pub const CSR_STIMECMP: u32 = 0x14D;
 
+/// Machine environment configuration register (menvcfg) bit fields.
 pub mod menvcfg {
     use core::arch::asm;
 
+    /// Fence of I/O implies memory.
     pub const FIOM: usize = 0x1 << 0;
+    /// Cache block invalidate - flush.
     pub const CBIE_FLUSH: usize = 0x01 << 4;
+    /// Cache block invalidate - invalidate.
     pub const CBIE_INVALIDATE: usize = 0x11 << 4;
+    /// Cache block clean for enclave.
     pub const CBCFE: usize = 0x1 << 6;
+    /// Cache block zero for enclave.
     pub const CBZE: usize = 0x1 << 7;
+    /// Page-based memory types enable.
     pub const PBMTE: usize = 0x1 << 62;
+    /// Supervisor timer counter enable.
     pub const STCE: usize = 0x1 << 63;
 
+    /// Sets the STCE bit to enable supervisor timer counter.
     #[inline(always)]
     pub fn set_stce() {
         set_bits(STCE);
     }
 
+    /// Sets specified bits in menvcfg register.
     pub fn set_bits(option: usize) {
         let mut bits: usize;
         unsafe {
+            // Read current `menvcfg` value.
             asm!("csrr {}, menvcfg", out(reg) bits, options(nomem));
         }
+        // Set requested bits
         bits |= option;
         unsafe {
+            // Write back updated value
             asm!("csrw menvcfg, {}", in(reg) bits, options(nomem));
         }
     }
 }
 
+/// Supervisor timer compare register operations.
 pub mod stimecmp {
     use core::arch::asm;
 
+    /// Sets the supervisor timer compare value.
     pub fn set(value: u64) {
         unsafe {
             asm!("csrrw zero, stimecmp, {}", in(reg) value, options(nomem));
@@ -42,6 +62,7 @@ pub mod stimecmp {
     }
 }
 
+/// Returns the current hart (hardware thread) ID.
 #[inline]
 pub fn current_hartid() -> usize {
     riscv::register::mhartid::read()

prototyper/src/sbi/console.rs

@@ -3,58 +3,92 @@ use core::fmt::{self, Write};
 use rustsbi::{Console, Physical, SbiRet};
 use spin::Mutex;
 
+/// A trait that must be implemented by console devices to provide basic I/O functionality.
 pub trait ConsoleDevice {
+    /// Reads bytes from the console into the provided buffer.
+    ///
+    /// # Returns
+    /// The number of bytes that were successfully read.
     fn read(&self, buf: &mut [u8]) -> usize;
+
+    /// Writes bytes from the provided buffer to the console.
+    ///
+    /// # Returns
+    /// The number of bytes that were successfully written.
     fn write(&self, buf: &[u8]) -> usize;
 }
 
+/// An implementation of the SBI console interface that wraps a console device.
+///
+/// This provides a safe interface for interacting with console hardware through the
+/// SBI specification.
 pub struct SbiConsole<'a, T: ConsoleDevice> {
     inner: &'a Mutex<T>,
 }
 
 impl<'a, T: ConsoleDevice> SbiConsole<'a, T> {
+    /// Creates a new SBI console that wraps the provided locked console device.
+    ///
+    /// # Arguments
+    /// * `inner` - A mutex containing the console device implementation
+    #[inline]
     pub fn new(inner: &'a Mutex<T>) -> Self {
         Self { inner }
     }
 
+    /// Writes a single character to the console.
+    ///
+    /// # Arguments
+    /// * `c` - The character to write, as a usize
+    ///
+    /// # Returns
+    /// Always returns 0 to indicate success
     #[inline]
     pub fn putchar(&mut self, c: usize) -> usize {
         self.write_char(c as u8 as char).unwrap();
         0
     }
 
+    /// Reads a single character from the console.
+    ///
+    /// This method will block until a character is available to be read.
+    ///
+    /// # Returns
+    /// The read character as a usize
     #[inline]
     pub fn getchar(&self) -> usize {
         let mut c = 0u8;
         let console = self.inner.lock();
-        loop {
-            if console.read(core::slice::from_mut(&mut c)) == 1 {
-                break;
-            }
+        // Block until we successfully read 1 byte
+        while console.read(core::slice::from_mut(&mut c)) != 1 {
+            core::hint::spin_loop();
         }
-        c as _
+        c as usize
     }
 }
 
 impl<'a, T: ConsoleDevice> Console for SbiConsole<'a, T> {
+    /// Write a physical memory buffer to the console.
     #[inline]
     fn write(&self, bytes: Physical<&[u8]>) -> SbiRet {
-        // TODO verify valid memory range for a `Physical` slice.
+        // TODO: verify valid memory range for a `Physical` slice.
         let start = bytes.phys_addr_lo();
         let buf = unsafe { core::slice::from_raw_parts(start as *const u8, bytes.num_bytes()) };
-        let bytes_num: usize = self.inner.lock().write(buf);
-        SbiRet::success(bytes_num)
+        let bytes_written = self.inner.lock().write(buf);
+        SbiRet::success(bytes_written)
     }
 
+    /// Read from console into a physical memory buffer.
     #[inline]
     fn read(&self, bytes: Physical<&mut [u8]>) -> SbiRet {
-        // TODO verify valid memory range for a `Physical` slice.
+        // TODO: verify valid memory range for a `Physical` slice.
         let start = bytes.phys_addr_lo();
         let buf = unsafe { core::slice::from_raw_parts_mut(start as *mut u8, bytes.num_bytes()) };
-        let bytes_num = self.inner.lock().read(buf);
-        SbiRet::success(bytes_num)
+        let bytes_read = self.inner.lock().read(buf);
+        SbiRet::success(bytes_read)
     }
 
+    /// Write a single byte to the console.
     #[inline]
     fn write_byte(&self, byte: u8) -> SbiRet {
         self.inner.lock().write(&[byte]);
@@ -63,10 +97,12 @@ impl<'a, T: ConsoleDevice> Console for SbiConsole<'a, T> {
 }
 
 impl<'a, T: ConsoleDevice> fmt::Write for SbiConsole<'a, T> {
+    /// Implement Write trait for string formatting.
     #[inline]
     fn write_str(&mut self, s: &str) -> fmt::Result {
         let mut bytes = s.as_bytes();
         let console = self.inner.lock();
+        // Write all bytes in chunks
         while !bytes.is_empty() {
             let count = console.write(bytes);
             bytes = &bytes[count..];
@@ -75,6 +111,7 @@ impl<'a, T: ConsoleDevice> fmt::Write for SbiConsole<'a, T> {
     }
 }
 
+/// Global function to write a character to the console.
 #[inline]
 pub fn putchar(c: usize) -> usize {
     unsafe { SBI_IMPL.assume_init_mut() }
@@ -84,6 +121,7 @@ pub fn putchar(c: usize) -> usize {
         .putchar(c)
 }
 
+/// Global function to read a character from the console.
 #[inline]
 pub fn getchar() -> usize {
     unsafe { SBI_IMPL.assume_init_mut() }

prototyper/src/sbi/extensions.rs

@@ -10,7 +10,7 @@ pub enum Extension {
 
 impl Extension {
     const COUNT: usize = 1;
-    const ITER: [Self;Extension::COUNT] = [Extension::Sstc];
+    const ITER: [Self; Extension::COUNT] = [Extension::Sstc];
 
     pub fn as_str(&self) -> &'static str {
         match self {
@@ -28,7 +28,8 @@ pub fn hart_extension_probe(hart_id: usize, ext: Extension) -> bool {
     unsafe {
         ROOT_STACK
             .get_mut(hart_id)
-            .map(|x| x.hart_context().extensions.0[ext.index()]).unwrap()
+            .map(|x| x.hart_context().extensions.0[ext.index()])
+            .unwrap()
     }
 }
 
@@ -38,7 +39,7 @@ pub fn init(cpus: &NodeSeq) {
     for cpu_iter in cpus.iter() {
         let cpu = cpu_iter.deserialize::<Cpu>();
         let hart_id = cpu.reg.iter().next().unwrap().0.start;
-        let mut hart_exts = [false;Extension::COUNT];
+        let mut hart_exts = [false; Extension::COUNT];
         let isa = cpu.isa.unwrap();
         Extension::ITER.iter().for_each(|ext| {
             hart_exts[ext.index()] = isa.iter().any(|e| e == ext.as_str());
@@ -47,21 +48,22 @@ pub fn init(cpus: &NodeSeq) {
         unsafe {
             ROOT_STACK
                 .get_mut(hart_id)
-                .map(|stack| stack.hart_context().extensions = HartExtensions(hart_exts)).unwrap()
+                .map(|stack| stack.hart_context().extensions = HartExtensions(hart_exts))
+                .unwrap()
         }
     }
 }
 
-#[cfg(feature = "nemu")] 
+#[cfg(feature = "nemu")]
 pub fn init(cpus: &NodeSeq) {
     for hart_id in 0..cpus.len() {
-        let mut hart_exts = [false;Extension::COUNT];
+        let mut hart_exts = [false; Extension::COUNT];
         hart_exts[Extension::Sstc.index()] = true;
         unsafe {
             ROOT_STACK
                 .get_mut(hart_id)
-                .map(|stack| stack.hart_context().extensions = HartExtensions(hart_exts)).unwrap()
+                .map(|stack| stack.hart_context().extensions = HartExtensions(hart_exts))
+                .unwrap()
         }
     }
 }
-    

prototyper/src/sbi/fifo.rs

@@ -1,46 +1,54 @@
 use core::mem::MaybeUninit;
 
+/// Size of the FIFO buffer.
 const FIFO_SIZE: usize = 16;
 
+#[derive(Debug)]
 pub enum FifoError {
     Empty,
     Full,
 }
 
+/// A fixed-size FIFO (First In First Out) queue implementation.
 pub struct Fifo<T: Copy + Clone> {
     data: [MaybeUninit<T>; FIFO_SIZE],
-    size: usize,
-    avil: usize,
+    head: usize,
     tail: usize,
+    count: usize,
 }
 
 impl<T: Copy + Clone> Fifo<T> {
-    pub fn new() -> Fifo<T> {
-        let data: [MaybeUninit<T>; FIFO_SIZE] = [const { MaybeUninit::uninit() }; FIFO_SIZE];
-        Fifo {
+    #[inline]
+    pub fn new() -> Self {
+        // Initialize array with uninitialized values
+        let data = [MaybeUninit::uninit(); FIFO_SIZE];
+        Self {
             data,
-            size: FIFO_SIZE,
-            avil: 0,
+            head: 0,
             tail: 0,
+            count: 0,
         }
     }
+
+    #[inline]
     pub fn is_full(&self) -> bool {
-        self.avil == self.size
+        self.count == FIFO_SIZE
     }
+
+    #[inline]
     pub fn is_empty(&self) -> bool {
-        self.avil == 0
+        self.count == 0
     }
 
-    pub fn push(&mut self, new_element: T) -> Result<(), FifoError> {
+    pub fn push(&mut self, element: T) -> Result<(), FifoError> {
         if self.is_full() {
             return Err(FifoError::Full);
         }
-        self.data[self.tail].write(new_element);
-        self.tail += 1;
-        self.avil += 1;
-        if self.tail >= self.size {
-            self.tail -= self.size;
-        }
+
+        // Write element and update tail position
+        self.data[self.tail].write(element);
+        self.tail = (self.tail + 1) % FIFO_SIZE;
+        self.count += 1;
 
         Ok(())
     }
@@ -49,16 +57,14 @@ impl<T: Copy + Clone> Fifo<T> {
         if self.is_empty() {
             return Err(FifoError::Empty);
         }
-        let raw_head = self.tail as isize - self.avil as isize;
-        let head = if raw_head < 0 {
-            raw_head + self.size as isize
-        } else {
-            raw_head
-        } as usize;
 
-        self.avil -= 1;
-        let result = *unsafe { self.data[head].assume_init_ref() };
-        unsafe { self.data[head].assume_init_drop() }
-        Ok(result)
+        // unsafe: Take ownership of element at head
+        let element = unsafe { self.data[self.head].assume_init_read() };
+
+        // Update head position
+        self.head = (self.head + 1) % FIFO_SIZE;
+        self.count -= 1;
+
+        Ok(element)
     }
 }

prototyper/src/sbi/hart_context.rs

@@ -1,37 +1,47 @@
+use crate::sbi::extensions::HartExtensions;
+use crate::sbi::hsm::HsmCell;
+use crate::sbi::rfence::RFenceCell;
 use core::ptr::NonNull;
 use core::sync::atomic::AtomicU8;
 use fast_trap::FlowContext;
 use riscv::register::mstatus;
-use crate::sbi::hsm::HsmCell;
-use crate::sbi::rfence::RFenceCell;
-use crate::sbi::extensions::HartExtensions;
 
+/// Context for managing hart (hardware thread) state and operations.
 pub(crate) struct HartContext {
-    /// trap context
+    /// Trap context for handling exceptions and interrupts.
     trap: FlowContext,
+    /// Hart state management cell containing next stage boot info.
     pub hsm: HsmCell<NextStage>,
+    /// Remote fence synchronization cell.
     pub rfence: RFenceCell,
+    /// Type of inter-processor interrupt pending.
     pub ipi_type: AtomicU8,
+    /// Supported hart extensions.
     pub extensions: HartExtensions,
 }
 
 impl HartContext {
+    /// Initialize the hart context by creating new HSM and RFence cells
     #[inline]
     pub fn init(&mut self) {
         self.hsm = HsmCell::new();
         self.rfence = RFenceCell::new();
     }
 
+    /// Get a non-null pointer to the trap context.
     #[inline]
     pub fn context_ptr(&mut self) -> NonNull<FlowContext> {
         unsafe { NonNull::new_unchecked(&mut self.trap) }
     }
 }
 
-/// Next Stage boot info
+/// Information needed to boot into the next execution stage.
 #[derive(Debug)]
 pub struct NextStage {
+    /// Starting address to jump to.
     pub start_addr: usize,
+    /// Opaque value passed to next stage.
     pub opaque: usize,
+    /// Privilege mode for next stage.
     pub next_mode: mstatus::MPP,
 }

prototyper/src/sbi/hsm.rs

@@ -11,17 +11,19 @@ use crate::riscv_spec::current_hartid;
 use crate::sbi::hart_context::NextStage;
 use crate::sbi::trap_stack::ROOT_STACK;
 
+/// Special state indicating a hart is in the process of starting.
 const HART_STATE_START_PENDING_EXT: usize = usize::MAX;
 
 type HsmState = AtomicUsize;
 
+/// Cell for managing hart state and shared data between harts.
 pub(crate) struct HsmCell<T> {
     status: HsmState,
     inner: UnsafeCell<Option<T>>,
 }
 
 impl<T> HsmCell<T> {
-    /// 创建一个新的共享对象。
+    /// Creates a new HsmCell with STOPPED state and no inner data.
     pub const fn new() -> Self {
         Self {
             status: HsmState::new(hart_state::STOPPED),
@@ -29,34 +31,37 @@ impl<T> HsmCell<T> {
         }
     }
 
-    /// 从当前硬件线程的状态中获取线程间共享对象。
+    /// Gets a local view of this cell for the current hart.
     ///
     /// # Safety
     ///
-    /// 用户需要确保对象属于当前硬件线程。
+    /// Caller must ensure this cell belongs to the current hart.
     #[inline]
     pub unsafe fn local(&self) -> LocalHsmCell<'_, T> {
         LocalHsmCell(self)
     }
 
-    /// 取出共享对象。
+    /// Gets a remote view of this cell for accessing from other harts.
     #[inline]
     pub fn remote(&self) -> RemoteHsmCell<'_, T> {
         RemoteHsmCell(self)
     }
 }
 
-/// 当前硬件线程的共享对象。
+/// View of HsmCell for operations on the current hart.
 pub struct LocalHsmCell<'a, T>(&'a HsmCell<T>);
 
-/// 任意硬件线程的共享对象。
+/// View of HsmCell for operations from other harts.
 pub struct RemoteHsmCell<'a, T>(&'a HsmCell<T>);
 
+// Mark HsmCell as safe to share between threads
 unsafe impl<T: Send> Sync for HsmCell<T> {}
 unsafe impl<T: Send> Send for HsmCell<T> {}
 
 impl<T> LocalHsmCell<'_, T> {
-    /// 从启动挂起状态的硬件线程取出共享数据，并将其状态设置为启动，如果成功返回取出的数据，否则返回当前状态。
+    /// Attempts to transition hart from START_PENDING to STARTED state.
+    ///
+    /// Returns inner data if successful, otherwise returns current state.
     #[inline]
     pub fn start(&self) -> Result<T, usize> {
         loop {
@@ -73,14 +78,14 @@ impl<T> LocalHsmCell<'_, T> {
         }
     }
 
-    /// 关闭。
+    /// Transitions hart to STOPPED state.
     #[allow(unused)]
     #[inline]
     pub fn stop(&self) {
         self.0.status.store(hart_state::STOPPED, Ordering::Release)
     }
 
-    /// 关闭。
+    /// Transitions hart to SUSPENDED state.
     #[allow(unused)]
     #[inline]
     pub fn suspend(&self) {
@@ -89,7 +94,7 @@ impl<T> LocalHsmCell<'_, T> {
             .store(hart_state::SUSPENDED, Ordering::Relaxed)
     }
 
-    /// 关闭。
+    /// Transitions hart to STARTED state.
     #[allow(unused)]
     #[inline]
     pub fn resume(&self) {
@@ -98,7 +103,9 @@ impl<T> LocalHsmCell<'_, T> {
 }
 
 impl<T: core::fmt::Debug> RemoteHsmCell<'_, T> {
-    /// 向关闭状态的硬件线程传入共享数据，并将其状态设置为启动挂起，返回是否放入成功。
+    /// Attempts to start a stopped hart by providing startup data.
+    ///
+    /// Returns true if successful, false if hart was not in STOPPED state.
     #[inline]
     pub fn start(&self, t: T) -> bool {
         if self
@@ -122,7 +129,7 @@ impl<T: core::fmt::Debug> RemoteHsmCell<'_, T> {
         }
     }
 
-    /// 取出当前状态。
+    /// Gets the current state of the hart.
     #[allow(unused)]
     #[inline]
     pub fn sbi_get_status(&self) -> usize {
@@ -132,7 +139,7 @@ impl<T: core::fmt::Debug> RemoteHsmCell<'_, T> {
         }
     }
 
-    /// 判断这个 HART 能否接收 IPI。
+    /// Checks if hart can receive IPIs (must be STARTED or SUSPENDED).
     #[allow(unused)]
     #[inline]
     pub fn allow_ipi(&self) -> bool {
@@ -143,7 +150,7 @@ impl<T: core::fmt::Debug> RemoteHsmCell<'_, T> {
     }
 }
 
-/// 获取此 hart 的 local hsm 对象。
+/// Gets the local HSM cell for the current hart.
 pub(crate) fn local_hsm() -> LocalHsmCell<'static, NextStage> {
     unsafe {
         ROOT_STACK
@@ -154,7 +161,7 @@ pub(crate) fn local_hsm() -> LocalHsmCell<'static, NextStage> {
     }
 }
 
-/// 获取此 hart 的 remote hsm 对象。
+/// Gets a remote view of the current hart's HSM cell.
 pub(crate) fn local_remote_hsm() -> RemoteHsmCell<'static, NextStage> {
     unsafe {
         ROOT_STACK
@@ -165,7 +172,7 @@ pub(crate) fn local_remote_hsm() -> RemoteHsmCell<'static, NextStage> {
     }
 }
 
-/// 获取任意 hart 的 remote hsm 对象。
+/// Gets a remote view of any hart's HSM cell.
 #[allow(unused)]
 pub(crate) fn remote_hsm(hart_id: usize) -> Option<RemoteHsmCell<'static, NextStage>> {
     unsafe {
@@ -175,10 +182,11 @@ pub(crate) fn remote_hsm(hart_id: usize) -> Option<RemoteHsmCell<'static, NextSt
     }
 }
 
-/// HSM
+/// Implementation of SBI HSM (Hart State Management) extension.
 pub(crate) struct SbiHsm;
 
 impl rustsbi::Hsm for SbiHsm {
+    /// Starts execution on a stopped hart.
     fn hart_start(&self, hartid: usize, start_addr: usize, opaque: usize) -> SbiRet {
         match remote_hsm(hartid) {
             Some(remote) => {
@@ -201,6 +209,7 @@ impl rustsbi::Hsm for SbiHsm {
         }
     }
 
+    /// Stops execution on the current hart.
     #[inline]
     fn hart_stop(&self) -> SbiRet {
         local_hsm().stop();
@@ -216,6 +225,7 @@ impl rustsbi::Hsm for SbiHsm {
         SbiRet::success(0)
     }
 
+    /// Gets the current state of a hart.
     #[inline]
     fn hart_get_status(&self, hartid: usize) -> SbiRet {
         match remote_hsm(hartid) {
@@ -224,6 +234,7 @@ impl rustsbi::Hsm for SbiHsm {
         }
     }
 
+    /// Suspends execution on the current hart.
     fn hart_suspend(&self, suspend_type: u32, _resume_addr: usize, _opaque: usize) -> SbiRet {
         use rustsbi::spec::hsm::suspend_type::{NON_RETENTIVE, RETENTIVE};
         if matches!(suspend_type, NON_RETENTIVE | RETENTIVE) {

prototyper/src/sbi/ipi.rs

@@ -3,67 +3,98 @@ use rustsbi::SbiRet;
 
 use crate::board::SBI_IMPL;
 use crate::riscv_spec::{current_hartid, stimecmp};
+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 crate::sbi::extensions::{hart_extension_probe, Extension};
 
+/// IPI type for supervisor software interrupt.
 pub(crate) const IPI_TYPE_SSOFT: u8 = 1 << 0;
+/// IPI type for memory fence operations.
 pub(crate) const IPI_TYPE_FENCE: u8 = 1 << 1;
 
+/// Trait defining interface for inter-processor interrupt device
 #[allow(unused)]
 pub trait IpiDevice {
+    /// Read machine time value.
     fn read_mtime(&self) -> u64;
+    /// Write machine time value.
     fn write_mtime(&self, val: u64);
+    /// Read machine timer compare value for given hart.
     fn read_mtimecmp(&self, hart_idx: usize) -> u64;
+    /// Write machine timer compare value for given hart.
     fn write_mtimecmp(&self, hart_idx: usize, val: u64);
+    /// Read machine software interrupt pending bit for given hart.
     fn read_msip(&self, hart_idx: usize) -> bool;
+    /// Set machine software interrupt pending bit for given hart.
     fn set_msip(&self, hart_idx: usize);
+    /// Clear machine software interrupt pending bit for given hart.
     fn clear_msip(&self, hart_idx: usize);
 }
 
+/// SBI IPI implementation.
 pub struct SbiIpi<'a, T: IpiDevice> {
+    /// Reference to atomic pointer to IPI device.
     pub ipi_dev: &'a AtomicPtr<T>,
+    /// Maximum hart ID in the system
     pub max_hart_id: usize,
 }
 
 impl<'a, T: IpiDevice> rustsbi::Timer for SbiIpi<'a, T> {
+    /// Set timer value for current hart.
     #[inline]
     fn set_timer(&self, stime_value: u64) {
-        if hart_extension_probe(current_hartid(),Extension::Sstc) {
+        let hart_id = current_hartid();
+        let uses_sstc = hart_extension_probe(hart_id, Extension::Sstc);
+
+        // Set timer value based on extension support.
+        if uses_sstc {
             stimecmp::set(stime_value);
-            unsafe {
-                riscv::register::mie::set_mtimer();
-            }
         } else {
-            self.write_mtimecmp(current_hartid(), stime_value);
+            self.write_mtimecmp(hart_id, stime_value);
             unsafe {
                 riscv::register::mip::clear_stimer();
-                riscv::register::mie::set_mtimer();
             }
         }
+        // Enable machine timer interrupt.
+        unsafe {
+            riscv::register::mie::set_mtimer();
+        }
     }
 }
 
 impl<'a, T: IpiDevice> rustsbi::Ipi for SbiIpi<'a, T> {
+    /// Send IPI to specified harts.
     #[inline]
     fn send_ipi(&self, hart_mask: rustsbi::HartMask) -> SbiRet {
+        let ipi_dev = unsafe { &*self.ipi_dev.load(Relaxed) };
+
         for hart_id in 0..=self.max_hart_id {
-            if hart_mask.has_bit(hart_id) && remote_hsm(hart_id).unwrap().allow_ipi() {
-                let old_ipi_type = set_ipi_type(hart_id, IPI_TYPE_SSOFT);
-                if old_ipi_type == 0 {
-                    unsafe {
-                        (*self.ipi_dev.load(Relaxed)).set_msip(hart_id);
-                    }
-                }
+            if !hart_mask.has_bit(hart_id) {
+                continue;
+            }
+
+            let Some(hsm) = remote_hsm(hart_id) else {
+                continue;
+            };
+
+            if !hsm.allow_ipi() {
+                continue;
+            }
+
+            if set_ipi_type(hart_id, IPI_TYPE_SSOFT) == 0 {
+                ipi_dev.set_msip(hart_id);
             }
         }
+
         SbiRet::success(0)
     }
 }
 
 impl<'a, T: IpiDevice> SbiIpi<'a, T> {
+    /// Create new SBI IPI instance.
+    #[inline]
     pub fn new(ipi_dev: &'a AtomicPtr<T>, max_hart_id: usize) -> Self {
         Self {
             ipi_dev,
@@ -71,67 +102,93 @@ impl<'a, T: IpiDevice> SbiIpi<'a, T> {
         }
     }
 
+    /// Send IPI for remote fence operation.
     pub fn send_ipi_by_fence(
         &self,
         hart_mask: rustsbi::HartMask,
-        ctx: rfence::RFenceCTX,
+        ctx: rfence::RFenceContext,
     ) -> SbiRet {
+        let current_hart = current_hartid();
+        let ipi_dev = unsafe { &*self.ipi_dev.load(Relaxed) };
+
+        // Send fence operations to target harts
         for hart_id in 0..=self.max_hart_id {
-            if hart_mask.has_bit(hart_id) && remote_hsm(hart_id).unwrap().allow_ipi() {
-                rfence::remote_rfence(hart_id).unwrap().set(ctx);
-                rfence::local_rfence().unwrap().add();
-                if hart_id == current_hartid() {
-                    continue;
+            if !hart_mask.has_bit(hart_id) {
+                continue;
+            }
+
+            let Some(hsm) = remote_hsm(hart_id) else {
+                continue;
+            };
+
+            if !hsm.allow_ipi() {
+                continue;
+            }
+
+            if let Some(remote) = rfence::remote_rfence(hart_id) {
+                remote.set(ctx);
+                if let Some(local) = rfence::local_rfence() {
+                    local.add();
                 }
-                let old_ipi_type = set_ipi_type(hart_id, IPI_TYPE_FENCE);
-                if old_ipi_type == 0 {
-                    unsafe {
-                        (*self.ipi_dev.load(Relaxed)).set_msip(hart_id);
+                if hart_id != current_hart {
+                    let old_ipi_type = set_ipi_type(hart_id, IPI_TYPE_FENCE);
+                    if old_ipi_type == 0 {
+                        ipi_dev.set_msip(hart_id);
                     }
                 }
             }
         }
+
+        // Wait for all fence operations to complete
         while !rfence::local_rfence().unwrap().is_sync() {
-            trap::rfence_signle_handler();
+            trap::rfence_single_handler();
         }
+
         SbiRet::success(0)
     }
 
+    /// Get lower 32 bits of machine time.
     #[inline]
     pub fn get_time(&self) -> usize {
         unsafe { (*self.ipi_dev.load(Relaxed)).read_mtime() as usize }
     }
 
+    /// Get upper 32 bits of machine time.
     #[inline]
     pub fn get_timeh(&self) -> usize {
         unsafe { ((*self.ipi_dev.load(Relaxed)).read_mtime() >> 32) as usize }
     }
 
+    /// Set machine software interrupt pending for hart.
     #[inline]
     pub fn set_msip(&self, hart_idx: usize) {
         unsafe { (*self.ipi_dev.load(Relaxed)).set_msip(hart_idx) }
     }
 
+    /// Clear machine software interrupt pending for hart.
     #[inline]
     pub fn clear_msip(&self, hart_idx: usize) {
         unsafe { (*self.ipi_dev.load(Relaxed)).clear_msip(hart_idx) }
     }
 
+    /// Write machine timer compare value for hart.
     #[inline]
     pub fn write_mtimecmp(&self, hart_idx: usize, val: u64) {
         unsafe { (*self.ipi_dev.load(Relaxed)).write_mtimecmp(hart_idx, val) }
     }
 
+    /// Clear all pending interrupts for current hart.
     #[inline]
     pub fn clear(&self) {
         let hart_id = current_hartid();
-        unsafe {
-            (*self.ipi_dev.load(Relaxed)).clear_msip(hart_id);
-            (*self.ipi_dev.load(Relaxed)).write_mtimecmp(hart_id, u64::MAX);
-        }
+        // Load ipi_dev once instead of twice
+        let ipi_dev = unsafe { &*self.ipi_dev.load(Relaxed) };
+        ipi_dev.clear_msip(hart_id);
+        ipi_dev.write_mtimecmp(hart_id, u64::MAX);
     }
 }
 
+/// Set IPI type for specified hart.
 pub fn set_ipi_type(hart_id: usize, event_id: u8) -> u8 {
     unsafe {
         ROOT_STACK
@@ -142,6 +199,7 @@ pub fn set_ipi_type(hart_id: usize, event_id: u8) -> u8 {
     }
 }
 
+/// Get and reset IPI type for current hart.
 pub fn get_and_reset_ipi_type() -> u8 {
     unsafe {
         ROOT_STACK
@@ -152,26 +210,29 @@ pub fn get_and_reset_ipi_type() -> u8 {
     }
 }
 
+/// Clear machine software interrupt pending for current hart.
+#[inline]
 pub fn clear_msip() {
-    unsafe { SBI_IMPL.assume_init_ref() }
-        .ipi
-        .as_ref()
-        .unwrap()
-        .clear_msip(current_hartid())
+    match unsafe { SBI_IMPL.as_ptr().as_ref().and_then(|sbi| sbi.ipi.as_ref()) } {
+        Some(ipi) => ipi.clear_msip(current_hartid()),
+        None => error!("SBI or IPI device not initialized"),
+    }
 }
 
+/// Clear machine timer interrupt for current hart.
+#[inline]
 pub fn clear_mtime() {
-    unsafe { SBI_IMPL.assume_init_ref() }
-        .ipi
-        .as_ref()
-        .unwrap()
-        .write_mtimecmp(current_hartid(), u64::MAX)
+    match unsafe { SBI_IMPL.as_ptr().as_ref().and_then(|sbi| sbi.ipi.as_ref()) } {
+        Some(ipi) => ipi.write_mtimecmp(current_hartid(), u64::MAX),
+        None => error!("SBI or IPI device not initialized"),
+    }
 }
 
+/// Clear all pending interrupts for current hart.
+#[inline]
 pub fn clear_all() {
-    unsafe { SBI_IMPL.assume_init_ref() }
-        .ipi
-        .as_ref()
-        .unwrap()
-        .clear();
-}
+    match unsafe { SBI_IMPL.as_ptr().as_ref().and_then(|sbi| sbi.ipi.as_ref()) } {
+        Some(ipi) => ipi.clear(),
+        None => error!("SBI or IPI device not initialized"),
+    }
+}

prototyper/src/sbi/logger.rs

@@ -1,39 +1,55 @@
 use core::str::FromStr;
 use log::{Level, LevelFilter};
 
+/// Simple logger implementation for RustSBI that supports colored output.
 pub struct Logger;
 
 impl Logger {
-    pub fn init() {
-        log::set_max_level(
-            option_env!("RUST_LOG")
-                .and_then(|s| LevelFilter::from_str(s).ok())
-                .unwrap_or(LevelFilter::Info),
-        );
-        log::set_logger(&Logger).unwrap();
+    /// Initialize the logger with log level from RUST_LOG env var or default to Info.
+    pub fn init() -> Result<(), log::SetLoggerError> {
+        // Set max log level from RUST_LOG env var if present, otherwise use Info
+        let max_level = option_env!("RUST_LOG")
+            .and_then(|s| LevelFilter::from_str(s).ok())
+            .unwrap_or(LevelFilter::Info);
+
+        log::set_max_level(max_level);
+        log::set_logger(&Logger)
     }
 }
 
 impl log::Log for Logger {
+    // Always enable logging for all log levels
     #[inline]
     fn enabled(&self, _metadata: &log::Metadata) -> bool {
         true
     }
 
+    // Log messages with color-coded levels
     #[inline]
     fn log(&self, record: &log::Record) {
-        let color_code: u8 = match record.level() {
-            Level::Error => 31,
-            Level::Warn => 93,
-            Level::Info => 32,
-            Level::Debug => 36,
-            Level::Trace => 90,
+        // ANSI color codes for different log levels
+        const ERROR_COLOR: u8 = 31; // Red
+        const WARN_COLOR: u8 = 93; // Bright yellow
+        const INFO_COLOR: u8 = 32; // Green
+        const DEBUG_COLOR: u8 = 36; // Cyan
+        const TRACE_COLOR: u8 = 90; // Bright black
+
+        let color_code = match record.level() {
+            Level::Error => ERROR_COLOR,
+            Level::Warn => WARN_COLOR,
+            Level::Info => INFO_COLOR,
+            Level::Debug => DEBUG_COLOR,
+            Level::Trace => TRACE_COLOR,
         };
+
         println!(
             "\x1b[1;37m[RustSBI] \x1b[1;{color_code}m{:^5}\x1b[0m - {}",
             record.level(),
             record.args(),
         );
     }
+
+    // No-op flush since we use println! which is already line-buffered
+    #[inline]
     fn flush(&self) {}
 }

prototyper/src/sbi/mod.rs

@@ -6,12 +6,12 @@ pub mod ipi;
 pub mod reset;
 pub mod rfence;
 
+pub mod extensions;
 pub mod fifo;
 pub mod hart_context;
 pub mod logger;
 pub mod trap;
 pub mod trap_stack;
-pub mod extensions;
 
 use console::{ConsoleDevice, SbiConsole};
 use hsm::SbiHsm;

prototyper/src/sbi/rfence.rs

@@ -9,34 +9,52 @@ use crate::sbi::trap_stack::ROOT_STACK;
 
 use core::sync::atomic::{AtomicU32, Ordering};
 
+/// Cell for managing remote fence operations between harts.
 pub(crate) struct RFenceCell {
-    queue: Mutex<Fifo<(RFenceCTX, usize)>>, // (ctx, source_hart_id)
+    // Queue of fence operations with source hart ID
+    queue: Mutex<Fifo<(RFenceContext, usize)>>,
+    // Counter for tracking pending synchronization operations
     wait_sync_count: AtomicU32,
 }
 
+/// Context information for a remote fence operation.
 #[repr(C)]
 #[derive(Clone, Copy, Debug)]
-pub struct RFenceCTX {
+pub struct RFenceContext {
+    /// Start address of memory region to fence.
     pub start_addr: usize,
+    /// Size of memory region to fence.
     pub size: usize,
+    /// Address space ID.
     pub asid: usize,
+    /// Virtual machine ID.
     pub vmid: usize,
+    /// Type of fence operation.
     pub op: RFenceType,
 }
 
+/// Types of remote fence operations supported.
 #[allow(unused)]
 #[derive(Clone, Copy, Debug)]
 pub enum RFenceType {
+    /// Instruction fence.
     FenceI,
+    /// Supervisor fence for virtual memory.
     SFenceVma,
+    /// Supervisor fence for virtual memory with ASID.
     SFenceVmaAsid,
+    /// Hypervisor fence for guest virtual memory with VMID.
     HFenceGvmaVmid,
+    /// Hypervisor fence for guest virtual memory.
     HFenceGvma,
+    /// Hypervisor fence for virtual machine virtual memory with ASID.
     HFenceVvmaAsid,
+    /// Hypervisor fence for virtual machine virtual memory.
     HFenceVvma,
 }
 
 impl RFenceCell {
+    /// Creates a new RFenceCell with empty queue and zero sync count.
     pub fn new() -> Self {
         Self {
             queue: Mutex::new(Fifo::new()),
@@ -44,28 +62,30 @@ impl RFenceCell {
         }
     }
 
+    /// Gets a local view of this fence cell for the current hart.
     #[inline]
     pub fn local(&self) -> LocalRFenceCell<'_> {
         LocalRFenceCell(self)
     }
 
-    /// 取出共享对象。
+    /// Gets a remote view of this fence cell for accessing from other harts.
     #[inline]
     pub fn remote(&self) -> RemoteRFenceCell<'_> {
         RemoteRFenceCell(self)
     }
 }
 
+// Mark RFenceCell as safe to share between threads
 unsafe impl Sync for RFenceCell {}
 unsafe impl Send for RFenceCell {}
 
-/// 当前硬件线程的rfence上下文。
+/// View of RFenceCell for operations on the current hart.
 pub struct LocalRFenceCell<'a>(&'a RFenceCell);
 
-/// 任意硬件线程的rfence上下文。
+/// View of RFenceCell for operations from other harts.
 pub struct RemoteRFenceCell<'a>(&'a RFenceCell);
 
-/// 获取当前hart的rfence上下文
+/// Gets the local fence context for the current hart.
 pub(crate) fn local_rfence() -> Option<LocalRFenceCell<'static>> {
     unsafe {
         ROOT_STACK
@@ -74,7 +94,7 @@ pub(crate) fn local_rfence() -> Option<LocalRFenceCell<'static>> {
     }
 }
 
-/// 获取任意hart的rfence上下文
+/// Gets the remote fence context for a specific hart.
 pub(crate) fn remote_rfence(hart_id: usize) -> Option<RemoteRFenceCell<'static>> {
     unsafe {
         ROOT_STACK
@@ -85,31 +105,38 @@ pub(crate) fn remote_rfence(hart_id: usize) -> Option<RemoteRFenceCell<'static>>
 
 #[allow(unused)]
 impl LocalRFenceCell<'_> {
+    /// Checks if all synchronization operations are complete.
     pub fn is_sync(&self) -> bool {
         self.0.wait_sync_count.load(Ordering::Relaxed) == 0
     }
+
+    /// Increments the synchronization counter.
     pub fn add(&self) {
         self.0.wait_sync_count.fetch_add(1, Ordering::Relaxed);
     }
 
+    /// Checks if the operation queue is empty.
     pub fn is_empty(&self) -> bool {
         self.0.queue.lock().is_empty()
     }
-    pub fn get(&self) -> Option<(RFenceCTX, usize)> {
-        match self.0.queue.lock().pop() {
-            Ok(res) => Some(res),
-            Err(_) => None,
-        }
+
+    /// Gets the next fence operation from the queue.
+    pub fn get(&self) -> Option<(RFenceContext, usize)> {
+        self.0.queue.lock().pop().ok()
     }
-    pub fn set(&self, ctx: RFenceCTX) {
+
+    /// Adds a fence operation to the queue, retrying if full.
+    pub fn set(&self, ctx: RFenceContext) {
+        let hart_id = current_hartid();
         loop {
-            match self.0.queue.lock().push((ctx, current_hartid())) {
+            let mut queue = self.0.queue.lock();
+            match queue.push((ctx, hart_id)) {
                 Ok(_) => break,
                 Err(FifoError::Full) => {
-                    trap::rfence_signle_handler();
-                    continue;
+                    drop(queue);
+                    trap::rfence_single_handler();
                 }
-                _ => panic!("Unable to push fence ops to fifo"),
+                Err(_) => panic!("Unable to push fence ops to fifo"),
             }
         }
     }
@@ -117,29 +144,54 @@ impl LocalRFenceCell<'_> {
 
 #[allow(unused)]
 impl RemoteRFenceCell<'_> {
-    pub fn set(&self, ctx: RFenceCTX) {
-        // TODO: maybe deadlock
+    /// Adds a fence operation to the queue from a remote hart.
+    pub fn set(&self, ctx: RFenceContext) {
+        let hart_id = current_hartid();
         loop {
-            match self.0.queue.lock().push((ctx, current_hartid())) {
-                Ok(_) => break,
+            let mut queue = self.0.queue.lock();
+            match queue.push((ctx, hart_id)) {
+                Ok(_) => return,
                 Err(FifoError::Full) => {
-                    trap::rfence_signle_handler();
-                    continue;
+                    drop(queue);
+                    trap::rfence_single_handler();
                 }
-                _ => panic!("Unable to push fence ops to fifo"),
+                Err(_) => panic!("Unable to push fence ops to fifo"),
             }
         }
     }
 
+    /// Decrements the synchronization counter.
     pub fn sub(&self) {
         self.0.wait_sync_count.fetch_sub(1, Ordering::Relaxed);
     }
 }
 
-/// RFENCE
+/// Implementation of RISC-V remote fence operations.
 pub(crate) struct SbiRFence;
 
-fn remote_fence_process(rfence_ctx: RFenceCTX, hart_mask: HartMask) -> SbiRet {
+/// Validates address range for fence operations
+#[inline(always)]
+fn validate_address_range(start_addr: usize, size: usize) -> Result<usize, SbiRet> {
+    // Check page alignment using bitwise AND instead of modulo
+    if start_addr & 0xFFF != 0 {
+        return Err(SbiRet::invalid_address());
+    }
+
+    // Avoid checked_add by checking for overflow directly
+    if size > usize::MAX - start_addr {
+        return Err(SbiRet::invalid_address());
+    }
+
+    let end_addr = start_addr + size;
+    if end_addr > usize::MAX {
+        Ok(usize::MAX)
+    } else {
+        Ok(size)
+    }
+}
+
+/// Processes a remote fence operation by sending IPI to target harts.
+fn remote_fence_process(rfence_ctx: RFenceContext, hart_mask: HartMask) -> SbiRet {
     let sbi_ret = unsafe { SBI_IMPL.assume_init_mut() }
         .ipi
         .as_ref()
@@ -150,9 +202,10 @@ fn remote_fence_process(rfence_ctx: RFenceCTX, hart_mask: HartMask) -> SbiRet {
 }
 
 impl rustsbi::Fence for SbiRFence {
+    /// Remote instruction fence for specified harts.
     fn remote_fence_i(&self, hart_mask: HartMask) -> SbiRet {
         remote_fence_process(
-            RFenceCTX {
+            RFenceContext {
                 start_addr: 0,
                 size: 0,
                 asid: 0,
@@ -163,14 +216,15 @@ impl rustsbi::Fence for SbiRFence {
         )
     }
 
+    /// Remote supervisor fence for virtual memory on specified harts.
     fn remote_sfence_vma(&self, hart_mask: HartMask, start_addr: usize, size: usize) -> SbiRet {
-        // TODO: return SBI_ERR_INVALID_ADDRESS, when start_addr or size is not valid.
-        let flush_size = match start_addr.checked_add(size) {
-            None => usize::MAX,
-            Some(_) => size,
+        let flush_size = match validate_address_range(start_addr, size) {
+            Ok(size) => size,
+            Err(e) => return e,
         };
+
         remote_fence_process(
-            RFenceCTX {
+            RFenceContext {
                 start_addr,
                 size: flush_size,
                 asid: 0,
@@ -181,6 +235,7 @@ impl rustsbi::Fence for SbiRFence {
         )
     }
 
+    /// Remote supervisor fence for virtual memory with ASID on specified harts.
     fn remote_sfence_vma_asid(
         &self,
         hart_mask: HartMask,
@@ -188,13 +243,13 @@ impl rustsbi::Fence for SbiRFence {
         size: usize,
         asid: usize,
     ) -> SbiRet {
-        // TODO: return SBI_ERR_INVALID_ADDRESS, when start_addr or size is not valid.
-        let flush_size = match start_addr.checked_add(size) {
-            None => usize::MAX,
-            Some(_) => size,
+        let flush_size = match validate_address_range(start_addr, size) {
+            Ok(size) => size,
+            Err(e) => return e,
         };
+
         remote_fence_process(
-            RFenceCTX {
+            RFenceContext {
                 start_addr,
                 size: flush_size,
                 asid,

prototyper/src/sbi/trap.rs

@@ -13,10 +13,12 @@ use crate::sbi::hsm::local_hsm;
 use crate::sbi::ipi;
 use crate::sbi::rfence::{self, local_rfence, RFenceType};
 
+// Constants for page and TLB management
 const PAGE_SIZE: usize = 4096;
 // TODO: `TLB_FLUSH_LIMIT` is a platform-dependent parameter
 const TLB_FLUSH_LIMIT: usize = 4 * PAGE_SIZE;
 
+/// Trap vector table entry point. Maps different trap types to their handlers.
 #[naked]
 pub(crate) unsafe extern "C" fn trap_vec() {
     asm!(
@@ -43,19 +45,18 @@ pub(crate) unsafe extern "C" fn trap_vec() {
     )
 }
 
-/// machine timer 中断代理
+/// 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!(
-        // 换栈：
-        // sp      : M sp
-        // mscratch: S sp
+        // 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)
@@ -87,13 +88,13 @@ unsafe extern "C" fn mtimer() {
             sd      t4, 28*8(sp)
             sd      t5, 29*8(sp)
             sd      t6, 1*8(sp)",
-        // 清除 mtimecmp
+        // Clear machine timer compare register
         "    call  {clear_mtime}",
-        // 设置 stip
+        // 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)
@@ -125,11 +126,9 @@ unsafe extern "C" fn mtimer() {
             ld      t5, 29*8(sp)
             ld      t6, 1*8(sp)",
         "   addi   sp, sp, 30*8",
-        // 换栈：
-        // sp      : S sp
-        // mscratch: M sp
+        // 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,
@@ -137,14 +136,18 @@ unsafe extern "C" fn mtimer() {
     )
 }
 
-/// machine soft 中断代理
+/// 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)
@@ -175,14 +178,18 @@ pub unsafe extern "C" fn msoft() -> ! {
         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_hanlder}",
+        // 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)
@@ -213,14 +220,20 @@ pub unsafe extern "C" fn msoft() -> ! {
         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_hanlder = sym msoft_hanlder,
         options(noreturn)
     );
 }
 
+/// Machine software interrupt handler implementation.
+///
+/// Handles HSM (Hart State Management) and RFence operations.
 pub extern "C" fn msoft_hanlder(ctx: &mut SupervisorContext) {
     #[inline(always)]
     fn boot(ctx: &mut SupervisorContext, start_addr: usize, opaque: usize) {
@@ -234,7 +247,7 @@ pub extern "C" fn msoft_hanlder(ctx: &mut SupervisorContext) {
     }
 
     match local_hsm().start() {
-        // HSM Start
+        // Handle HSM Start
         Ok(next_stage) => {
             ipi::clear_msip();
             unsafe {
@@ -245,6 +258,7 @@ pub extern "C" fn msoft_hanlder(ctx: &mut SupervisorContext) {
             }
             boot(ctx, next_stage.start_addr, next_stage.opaque);
         }
+        // Handle HSM Stop
         Err(rustsbi::spec::hsm::HART_STOP) => {
             ipi::clear_msip();
             unsafe {
@@ -252,21 +266,24 @@ pub extern "C" fn msoft_hanlder(ctx: &mut SupervisorContext) {
             }
             riscv::asm::wfi();
         }
-        // RFence
+        // Handle RFence
         _ => {
             msoft_ipi_handler();
         }
     }
 }
 
-pub fn rfence_signle_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)
@@ -286,6 +303,7 @@ pub fn rfence_signle_handler() {
                 }
                 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
@@ -313,27 +331,31 @@ pub fn rfence_signle_handler() {
     }
 }
 
+/// Process all pending remote fence operations.
 pub fn rfence_handler() {
     while !local_rfence().unwrap().is_empty() {
-        rfence_signle_handler();
+        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
+/// Fast trap handler for SBI calls and illegal instructions.
 pub extern "C" fn fast_handler(
     mut ctx: FastContext,
     a1: usize,
@@ -356,7 +378,7 @@ pub extern "C" fn fast_handler(
         ctx.call(2)
     }
     match mcause::read().cause() {
-        // SBI call
+        // Handle SBI calls
         T::Exception(E::SupervisorEnvCall) => {
             use sbi_spec::{base, hsm, legacy};
             let mut ret = unsafe { SBI_IMPL.assume_init_ref() }.handle_ecall(
@@ -366,13 +388,13 @@ pub extern "C" fn fast_handler(
             );
             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);
                     }
-                    // legacy console 探测
+                    // Handle legacy console probe
                     (base::EID_BASE, base::PROBE_EXTENSION)
                         if matches!(
                             ctx.a0(),
@@ -400,6 +422,7 @@ pub extern "C" fn fast_handler(
             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");
@@ -411,7 +434,7 @@ pub extern "C" fn fast_handler(
             }
             ctx.restore()
         }
-        // 其他陷入
+        // Handle other traps
         trap => {
             println!(
                 "
@@ -429,6 +452,7 @@ pub extern "C" fn fast_handler(
     }
 }
 
+/// Delegate trap handling to supervisor mode.
 #[inline]
 fn delegate() {
     use riscv::register::{mcause, mepc, mtval, scause, sepc, sstatus, stval, stvec};
@@ -448,6 +472,7 @@ fn delegate() {
     }
 }
 
+/// Handle illegal instructions, particularly CSR access.
 #[inline]
 fn illegal_instruction_handler(ctx: &mut FastContext) -> bool {
     use riscv::register::{mepc, mtval};
@@ -488,6 +513,7 @@ fn illegal_instruction_handler(ctx: &mut FastContext) -> bool {
     true
 }
 
+/// Supervisor context structure containing saved register state.
 #[derive(Debug)]
 #[repr(C)]
 pub struct SupervisorContext {

prototyper/src/sbi/trap_stack.rs

@@ -1,28 +1,33 @@
-use core::mem::forget;
-use fast_trap::FreeTrapStack;
 use crate::riscv_spec::current_hartid;
 use crate::sbi::hart_context::HartContext;
 use crate::sbi::trap::fast_handler;
+use core::mem::forget;
+use fast_trap::FreeTrapStack;
 
+/// Stack size per hart (hardware thread) in bytes.
 const LEN_STACK_PER_HART: usize = 16 * 1024;
+/// Maximum number of supported harts.
 pub const NUM_HART_MAX: usize = 8;
-/// 栈空间。
+
+/// Root stack array for all harts, placed in uninitialized BSS section.
 #[link_section = ".bss.uninit"]
 pub(crate) static mut ROOT_STACK: [Stack; NUM_HART_MAX] = [Stack::ZERO; NUM_HART_MAX];
 
-/// 定位每个 hart 的栈。
+/// Locates and initializes stack for each hart.
+///
+/// This is a naked function that sets up the stack pointer based on hart ID.
 #[naked]
 pub(crate) unsafe extern "C" fn locate() {
     core::arch::asm!(
-        "   la   sp, {stack}
-            li   t0, {per_hart_stack_size}
-            csrr t1, mhartid
-            addi t1, t1,  1
-         1: add  sp, sp, t0
-            addi t1, t1, -1
-            bnez t1, 1b
-            call t1, {move_stack}
-            ret
+        "   la   sp, {stack}            // Load stack base address
+            li   t0, {per_hart_stack_size} // Load stack size per hart
+            csrr t1, mhartid            // Get current hart ID
+            addi t1, t1,  1             // Add 1 to hart ID
+         1: add  sp, sp, t0             // Calculate stack pointer
+            addi t1, t1, -1             // Decrement counter
+            bnez t1, 1b                 // Loop if not zero
+            call t1, {move_stack}       // Call stack reuse function
+            ret                         // Return
         ",
         per_hart_stack_size = const LEN_STACK_PER_HART,
         stack               =   sym ROOT_STACK,
@@ -31,7 +36,7 @@ pub(crate) unsafe extern "C" fn locate() {
     )
 }
 
-/// 预备陷入栈。
+/// Prepares trap stack for current hart
 pub(crate) fn prepare_for_trap() {
     unsafe {
         ROOT_STACK
@@ -40,33 +45,42 @@ pub(crate) fn prepare_for_trap() {
     };
 }
 
-/// 类型化栈。
+/// Stack type for each hart.
+///
+/// Memory layout:
+/// - Bottom: HartContext struct.
+/// - Middle: Stack space for the hart.
+/// - Top: Trap handling space.
 ///
-/// 每个硬件线程拥有一个满足这样条件的内存块。
-/// 这个内存块的底部放着硬件线程状态 [`HartContext`]，顶部用于陷入处理，中间是这个硬件线程的栈空间。
-/// 不需要 M 态线程，每个硬件线程只有这一个栈。
+/// Each hart has a single stack that contains both its context and working space.
 #[repr(C, align(128))]
 pub(crate) struct Stack([u8; LEN_STACK_PER_HART]);
 
 impl Stack {
-    /// 零初始化以避免加载。
     const ZERO: Self = Self([0; LEN_STACK_PER_HART]);
 
-    /// 从栈上取出硬件线程状态。
+    /// Gets mutable reference to hart context at bottom of stack.
     #[inline]
     pub fn hart_context(&mut self) -> &mut HartContext {
         unsafe { &mut *self.0.as_mut_ptr().cast() }
     }
 
+    /// Initializes stack for trap handling.
+    /// - Sets up hart context.
+    /// - Creates and loads FreeTrapStack with the stack range.
     fn load_as_stack(&'static mut self) {
         let hart = self.hart_context();
         let context_ptr = hart.context_ptr();
         hart.init();
+
+        // Get stack memory range.
         let range = self.0.as_ptr_range();
+
+        // Create and load trap stack, forgetting it to avoid drop
         forget(
             FreeTrapStack::new(
                 range.start as usize..range.end as usize,
-                |_| {},
+                |_| {}, // Empty callback
                 context_ptr,
                 fast_handler,
             )