Use amoswap based lock in once fat box

src/lib.rs

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

src/util.rs

@@ -4,16 +4,17 @@
 
 use alloc::boxed::Box;
 use core::{
+    cell::UnsafeCell,
     fmt::{self, Debug},
     marker::PhantomData,
     mem::MaybeUninit,
     ptr::{self, Pointee},
-    sync::atomic::{AtomicPtr, Ordering},
 };
 
 /// A thread-safe fat pointer cell which can be written to only once.
 pub struct OnceFatBox<T: ?Sized> {
-    thin_ptr: AtomicPtr<()>,
+    thin_ptr: UnsafeCell<*mut ()>,
+    lock: UnsafeCell<u8>,
     meta: MaybeUninit<<T as Pointee>::Metadata>,
     _marker: PhantomData<Option<Box<T>>>,
 }
@@ -41,7 +42,7 @@ where
 {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("OnceFatBox")
-            .field("data_address", &self.thin_ptr)
+            .field("address", &self.thin_ptr)
             .field("meta", &self.meta)
             .finish()
     }
@@ -51,7 +52,8 @@ impl<T: ?Sized> OnceFatBox<T> {
     /// Creates a new empty cell.
     pub const fn new() -> OnceFatBox<T> {
         OnceFatBox {
-            thin_ptr: AtomicPtr::new(ptr::null_mut()),
+            thin_ptr: UnsafeCell::new(ptr::null_mut()),
+            lock: UnsafeCell::new(0),
             meta: MaybeUninit::uninit(), // value meaning ignored when thin_ptr is null
             _marker: PhantomData,
         }
@@ -59,12 +61,12 @@ impl<T: ?Sized> OnceFatBox<T> {
 
     /// Gets a reference to the underlying value.
     pub fn get(&self) -> Option<&T> {
-        let data_address = self.thin_ptr.load(Ordering::Acquire);
+        let data_address = self.thin_ptr.get();
         if data_address.is_null() {
             return None;
         }
         let metadata = unsafe { self.meta.assume_init() };
-        let fat_ptr: *const T = ptr::from_raw_parts(data_address, metadata);
+        let fat_ptr: *const T = ptr::from_raw_parts(unsafe { *data_address }, metadata);
         Some(unsafe { &*fat_ptr })
     }
 
@@ -74,21 +76,35 @@ impl<T: ?Sized> OnceFatBox<T> {
     pub fn set(&self, value: Box<T>) -> Result<(), Box<T>> {
         let fat_ptr = Box::into_raw(value);
         let data_address = fat_ptr as *mut ();
-        let exchange = self.thin_ptr.compare_exchange(
-            ptr::null_mut(),
-            data_address,
-            Ordering::AcqRel,
-            Ordering::Acquire,
-        );
+        // compare-exchange using amo operations
+        let exchange = unsafe {
+            let mut ans = Err(());
+            asm!(
+                "1: lw  {tmp}, ({lock})", // check if lock is held
+                "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,
+            );
+            // critical section begin
+            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(())
+            }
+            // critical section end
+            asm!(
+                "amoswap.w.rl x0, x0, ({lock})", // release lock by storing 0
+                lock = in(reg) self.lock.get(),
+            );
+            ans
+        };
         if let Err(_) = exchange {
             let value = unsafe { Box::from_raw(fat_ptr) };
             return Err(value);
         }
-        // 对once_cell来说这样做是对的，因为其它的线程失败后，不会再更改元数据了。
-        // 如果其它的线程仍然需要更改元数据，就不能这样做。
-        unsafe {
-            *(self.meta.as_ptr() as *mut _) = ptr::metadata(fat_ptr);
-        }
         Ok(())
     }
 }