Merge pull request #29 from chenzongyao200127/main

feat(prototyper): Improve code quality and enhance safety features

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

@@ -9,61 +9,92 @@ use crate::sbi::rfence;
 use crate::sbi::trap;
 use crate::sbi::trap_stack::ROOT_STACK;
 
+/// 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_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) => {
+                    drop(queue);
                     trap::rfence_single_handler();
-                    continue;
                 }
-                _ => 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) => {
+                    drop(queue);
                     trap::rfence_single_handler();
-                    continue;
                 }
-                _ => 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();
         }
     }
 }
 
+/// 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_single_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_single_handler() {
     }
 }
 
+/// 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
+/// 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,
             )