Use atomic instructions to write once and mutex

No LR/SC instructions used

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "rustsbi"
 description = "Minimal RISC-V's SBI implementation library in Rust"
-version = "0.2.0-alpha.7"
+version = "0.2.0-alpha.8"
 authors = [
     "Luo Jia <[email protected]>",
     "Campbell He <[email protected]>",

src/legacy_stdio.rs

@@ -1,5 +1,7 @@
 //! 这个模块的两个宏应该公开
 //! 如果制造实例的时候，给定了stdout，那么就会打印到这个stdout里面
+use crate::util::AmoMutex;
+use alloc::boxed::Box;
 use embedded_hal::serial::{Read, Write};
 use nb::block;
 
@@ -59,10 +61,7 @@ where
     }
 }
 
-use alloc::boxed::Box;
-use spin::mutex::TicketMutex;
-
-static LEGACY_STDIO: TicketMutex<Option<Box<dyn LegacyStdio>>> = TicketMutex::new(None);
+static LEGACY_STDIO: AmoMutex<Option<Box<dyn LegacyStdio>>> = AmoMutex::new(None);
 
 #[doc(hidden)] // use through a macro
 pub fn init_legacy_stdio_embedded_hal<T: Read<u8> + Write<u8> + Send + 'static>(serial: T) {

src/lib.rs

@@ -137,7 +137,7 @@
 //! RustSBI provides implementations on common platforms in separate platform crates.
 
 #![no_std]
-#![feature(asm, asm_const)]
+#![feature(asm, asm_const, asm_sym)]
 #![feature(ptr_metadata)]
 
 extern crate alloc;

src/util.rs

@@ -8,6 +8,7 @@ use core::{
     fmt::{self, Debug},
     marker::PhantomData,
     mem::MaybeUninit,
+    ops::{Deref, DerefMut},
     ptr::{self, Pointee},
 };
 
@@ -80,16 +81,20 @@ impl<T: ?Sized> OnceFatBox<T> {
         let exchange = unsafe {
             let mut ans = Err(());
             asm!(
+                "li     {one}, 1",
                 "1: lw  {tmp}, ({lock})", // check if lock is held
+                // "call   {relax}", // spin loop hint
                 "bnez   {tmp}, 1b", // retry if held
                 "amoswap.w.aq {tmp}, {one}, ({lock})", // attempt to acquire lock
                 "bnez   {tmp}, 1b", // retry if held
                 lock = in(reg) self.lock.get(),
                 tmp = out(reg) _,
-                one = const 1,
+                one = out(reg) _,
+                // relax = sym pause,
+                options(nostack)
             );
             // critical section begin
-            if self.thin_ptr.get() == ptr::null_mut() {
+            if *self.thin_ptr.get() == ptr::null_mut() {
                 *self.thin_ptr.get() = data_address;
                 *(self.meta.as_ptr() as *mut _) = ptr::metadata(fat_ptr);
                 ans = Ok(())
@@ -109,4 +114,80 @@ impl<T: ?Sized> OnceFatBox<T> {
     }
 }
 
+// // FIXME: use `core::arch::riscv::pause()` after https://github.com/rust-lang/rust/pull/91548
+// #[inline]
+// fn pause() {
+//     unsafe { asm!(".word 0x0100000F", options(nomem, nostack)) }
+// }
+
 unsafe impl<T: Sync + Send + ?Sized> Sync for OnceFatBox<T> {}
+
+/// Use only amo instructions on mutex; no lr/sc instruction is used
+pub struct AmoMutex<T: ?Sized> {
+    lock: UnsafeCell<u8>,
+    data: UnsafeCell<T>,
+}
+
+pub struct AmoMutexGuard<'a, T: ?Sized> {
+    lock: *mut u8,
+    data: &'a mut T,
+}
+
+impl<T> AmoMutex<T> {
+    /// Create a new AmoMutex
+    pub const fn new(data: T) -> Self {
+        AmoMutex {
+            data: UnsafeCell::new(data),
+            lock: UnsafeCell::new(0),
+        }
+    }
+    /// Locks the mutex and returns a guard that permits access to the inner data.
+    pub fn lock(&self) -> AmoMutexGuard<T> {
+        unsafe {
+            asm!(
+                "li     {one}, 1",
+                "1: lw  {tmp}, ({lock})", // check if lock is held
+                // "call   {relax}", // spin loop hint
+                "bnez   {tmp}, 1b", // retry if held
+                "amoswap.w.aq {tmp}, {one}, ({lock})", // attempt to acquire lock
+                "bnez   {tmp}, 1b", // retry if held
+                lock = in(reg) self.lock.get(),
+                tmp = out(reg) _,
+                one = out(reg) _,
+                // relax = sym pause,
+                options(nostack)
+            );
+        }
+        AmoMutexGuard {
+            lock: self.lock.get(),
+            data: unsafe { &mut *self.data.get() },
+        }
+    }
+}
+
+unsafe impl<T: ?Sized + Send> Sync for AmoMutex<T> {}
+unsafe impl<T: ?Sized + Send> Send for AmoMutex<T> {}
+
+impl<'a, T: ?Sized> Deref for AmoMutexGuard<'a, T> {
+    type Target = T;
+    fn deref(&self) -> &T {
+        self.data
+    }
+}
+
+impl<'a, T: ?Sized> DerefMut for AmoMutexGuard<'a, T> {
+    fn deref_mut(&mut self) -> &mut T {
+        self.data
+    }
+}
+impl<'a, T: ?Sized> Drop for AmoMutexGuard<'a, T> {
+    /// The dropping of the mutex guard will release the lock it was created from.
+    fn drop(&mut self) {
+        unsafe {
+            asm!(
+                "amoswap.w.rl x0, x0, ({lock})", // release lock by storing 0
+                lock = in(reg) self.lock,
+            );
+        }
+    }
+}