use std::cell::UnsafeCell;
use std::sync::atomic::{spin_loop_hint, AtomicUsize, Ordering};
pub struct SpinRwLock<T> {
lock: AtomicUsize,
data: UnsafeCell<T>,
}
impl<T> SpinRwLock<T> {
#[allow(dead_code)]
pub fn new(data: T) -> SpinRwLock<T> {
SpinRwLock {
lock: AtomicUsize::new(0),
data: UnsafeCell::new(data),
}
}
pub fn write(&self) -> RwLockWriteGuard<'_, T> {
loop {
if let Some(guard) = self.try_write() {
return guard;
} else {
spin_loop_hint();
}
}
}
pub fn read(&self) -> RwLockReadGuard<'_, T> {
loop {
if let Some(guard) = self.try_read() {
return guard;
} else {
spin_loop_hint();
}
}
}
pub fn try_write(&self) -> Option<RwLockWriteGuard<'_, T>> {
if self.lock.compare_and_swap(0, 1, Ordering::Acquire) == 0 {
Some(RwLockWriteGuard { lock: self })
} else {
None
}
}
pub fn try_read(&self) -> Option<RwLockReadGuard<'_, T>> {
self.lock.fetch_add(2, Ordering::Release);
if (self.lock.load(Ordering::Acquire) & 1) == 0 {
Some(RwLockReadGuard { lock: self })
} else {
self.lock.fetch_sub(2, Ordering::Release);
None
}
}
fn release_write(&self) {
self.lock.fetch_sub(1, Ordering::Release);
}
fn release_read(&self) {
self.lock.fetch_sub(2, Ordering::Release);
}
pub(crate) fn as_ptr(&self) -> *mut T {
self.data.get()
}
}
pub struct RwLockReadGuard<'a, T> {
lock: &'a SpinRwLock<T>,
}
impl<'a, T> std::ops::Deref for RwLockReadGuard<'a, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
unsafe { self.lock.data.get().as_ref().unwrap() }
}
}
impl<'a, T> Drop for RwLockReadGuard<'a, T> {
fn drop(&mut self) {
self.lock.release_read();
}
}
pub struct RwLockWriteGuard<'a, T> {
lock: &'a SpinRwLock<T>,
}
impl<'a, T> std::ops::Deref for RwLockWriteGuard<'a, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
unsafe { self.lock.data.get().as_ref().unwrap() }
}
}
impl<'a, T> std::ops::DerefMut for RwLockWriteGuard<'a, T> {
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe { self.lock.data.get().as_mut().unwrap() }
}
}
impl<'a, T> Drop for RwLockWriteGuard<'a, T> {
fn drop(&mut self) {
self.lock.release_write();
}
}