feat: 实现 AmoOncePtr 并用于 HSM

Cargo.toml

@@ -15,10 +15,12 @@ keywords = ["riscv", "sbi", "rustsbi"]
 categories = ["os", "embedded", "hardware-support", "no-std"]
 edition = "2021"
 
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
 [dependencies]
 embedded-hal = "0.2.7"
 nb = "1.0"
-riscv = "0.7"
+riscv = "0.8"
 
 [features]
 default = []
@@ -27,8 +29,4 @@ guest = []
 
 [package.metadata.docs.rs]
 default-target = "riscv64imac-unknown-none-elf"
-targets = [
-    "riscv32imac-unknown-none-elf", "riscv64imac-unknown-none-elf",
-]
-
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+targets = ["riscv32imac-unknown-none-elf", "riscv64imac-unknown-none-elf"]

src/hsm.rs

@@ -42,7 +42,7 @@ use crate::ecall::SbiRet;
 /// - Always prefer most recent suspend state requested for higher topology group
 ///
 /// Ref: [Section 8, RISC-V Supervisor Binary Interface Specification](https://github.com/riscv-non-isa/riscv-sbi-doc/blob/master/riscv-sbi.adoc#8-hart-state-management-extension-eid-0x48534d-hsm)
-pub trait Hsm: Send {
+pub trait Hsm: Send + Sync {
     /// Request the SBI implementation to start executing the given hart at specified address in supervisor-mode.
     ///
     /// This call is asynchronous - more specifically, the `sbi_hart_start()` may return before target hart
@@ -190,15 +190,12 @@ pub trait Hsm: Send {
     }
 }
 
-use crate::util::OnceFatBox;
-use alloc::boxed::Box;
+use crate::util::AmoOncePtr;
 
-static HSM: OnceFatBox<dyn Hsm + Sync + 'static> = OnceFatBox::new();
+static HSM: AmoOncePtr<dyn Hsm> = AmoOncePtr::new();
 
-#[doc(hidden)] // use through a macro or a call from implementation
-pub fn init_hsm<T: Hsm + Sync + 'static>(hsm: T) {
-    let result = HSM.set(Box::new(hsm));
-    if result.is_err() {
+pub fn init_hsm(hsm: &'static dyn Hsm) {
+    if !HSM.try_call_once(hsm) {
         panic!("load sbi module when already loaded")
     }
 }

src/util.rs

@@ -191,9 +191,125 @@ impl<'a, T: ?Sized> Drop for AmoMutexGuard<'a, T> {
     fn drop(&mut self) {
         unsafe {
             asm!(
-                "amoswap.w.rl x0, x0, ({lock})", // release lock by storing 0
+                "amoswap.w.rl zero, zero, ({lock})", // release lock by storing 0
                 lock = in(reg) self.lock,
             );
         }
     }
 }
+
+/// 只使用 AMO 指令的一次初始化指针。
+pub struct AmoOncePtr<T: ?Sized> {
+    /// As atomic bool, to check if it is the first time to set `ptr`.
+    lock: UnsafeCell<u32>,
+    ptr: UnsafeCell<*const ()>,
+    meta: UnsafeCell<MaybeUninit<<T as Pointee>::Metadata>>,
+}
+
+unsafe impl<T: ?Sized + Send> Send for AmoOncePtr<T> {}
+unsafe impl<T: ?Sized + Send + Sync> Sync for AmoOncePtr<T> {}
+
+impl<T: ?Sized> AmoOncePtr<T> {
+    #[inline]
+    pub const fn new() -> Self {
+        Self {
+            lock: UnsafeCell::new(0),
+            ptr: UnsafeCell::new(core::ptr::null()),
+            meta: UnsafeCell::new(MaybeUninit::uninit()),
+        }
+    }
+
+    pub fn try_call_once(&self, r#ref: &'static T) -> bool {
+        let ptr = r#ref as *const T;
+        let locked: u32;
+        unsafe {
+            asm!(
+                "
+                lw           {locked}, ({lock})
+                bnez         {locked}, 1f
+                amoswap.w.aq {locked}, {one}, ({lock})
+                1: ",
+                lock   =  in(reg) self.lock.get(),
+                one    =  in(reg) 1,
+                locked = out(reg) locked,
+            );
+        }
+        if locked == 0 {
+            // 取得锁，初始化对象
+            unsafe {
+                // amo rl 保证先初始化 meta 后设置指针
+                (*self.meta.get()) = MaybeUninit::new(ptr::metadata(ptr));
+                #[cfg(target_pointer_width = "32")]
+                asm!(
+                    "amoswap.d.rl zero, {src}, ({dst})",
+                    src = in(reg) ptr as *const (),
+                    dst = in(reg) self.ptr.get(),
+                );
+                #[cfg(target_pointer_width = "64")]
+                asm!(
+                    "amoswap.d.rl zero, {src}, ({dst})",
+                    src = in(reg) ptr as *const (),
+                    dst = in(reg) self.ptr.get(),
+                );
+            }
+            true
+        } else {
+            // 未取得锁，对象已被初始化过
+            false
+        }
+    }
+
+    #[allow(unused)]
+    pub fn call_once(&self, r#ref: &'static T) -> Result<&T, &T> {
+        if self.try_call_once(r#ref) {
+            Ok(r#ref)
+        } else {
+            Err(self.wait())
+        }
+    }
+
+    pub fn wait(&self) -> &T {
+        loop {
+            // 反复读直到非空。
+            let ptr = unsafe { *self.ptr.get() };
+            if !ptr.is_null() {
+                return unsafe { self.build_ref_unchecked(ptr) };
+            }
+        }
+    }
+
+    pub fn get(&self) -> Option<&T> {
+        let ptr: *const ();
+        unsafe {
+            // 先获取指针。如果指针非空，元数据一定存在。
+            // FIXME AMO 设的值是否一定对 LD 可见？如果确定就不需要 AMO 读了。
+            #[cfg(target_pointer_width = "32")]
+            asm!(" lw          {dst}, ({src})
+                   bnez        {dst}, 1f
+                   amoadd.w.aq {dst}, zero, ({src})
+                1: ",
+                src =  in(reg) self.ptr.get(),
+                dst = out(reg) ptr,
+            );
+            #[cfg(target_pointer_width = "64")]
+            asm!(" ld          {dst}, ({src})
+                   bnez        {dst}, 1f
+                   amoadd.d.aq {dst}, zero, ({src})
+                1: ",
+                src =  in(reg) self.ptr.get(),
+                dst = out(reg) ptr,
+            );
+        }
+        if !ptr.is_null() {
+            Some(unsafe { self.build_ref_unchecked(ptr) })
+        } else {
+            None
+        }
+    }
+
+    /// 利用指针和元数据生成引用。需要保证传入的指针非空。如果能传入非空指针，meta 也一定存在。
+    #[inline]
+    unsafe fn build_ref_unchecked(&self, ptr: *const ()) -> &T {
+        &*ptr::from_raw_parts(ptr, (*self.meta.get()).assume_init())
+    }
+}