rt-write-lock / src / lib.rs

// std::mem::forget creates issues because you can run a lock function without running the unlock function
// i deal with this by tracking for Reader and Writer if they are locked currently. If a new locked is aquired
// while it's still locked it panics. If Reader or Writer are dropped while locked it also panics.
//
// This could be circumvented by leaking Reader or Writer, but this doesn't lead to issues as then it's impossible
// to aquire a new Reader or Writer, because then the count also wasn't decremented.
//
// This panics more often than UB would occur without this detection, but this exact specs is a implementation detail
// and users just shouldn't forget the Guards.

#![no_std]

extern crate alloc;

use core::{
    marker::PhantomData,
    ops::{Deref, DerefMut},
    ptr::NonNull,
};

mod shared;

use shared::{Ptr, Shared};

pub struct Reader<T> {
    shared: NonNull<Shared<T>>,
    locked: bool,
    _owned: PhantomData<T>,
}

impl<T> Reader<T> {
    pub fn try_read(&mut self) -> Option<ReadGuard<'_, T>> {
        if self.locked {
            self.locked = false; // don't also panic in the destructor
            panic!("ReadGuard was forgotten");
        }
        let shared = self.get_shared();
        let value = shared.try_read();
        // SAFETY: try_read returned Some to allowed to get a shared ref
        let value = value.map(|v| unsafe { &*v.get() });
        if let Some(value) = value {
            // only set locked if locking actually was successful
            self.locked = true;
            Some(ReadGuard {
                reader: self,
                value,
            })
        } else {
            None
        }
    }

    fn get_shared(&self) -> &Shared<T> {
        // SAFETY: the shared is always valid
        unsafe { self.shared.as_ref() }
    }

    pub fn get_writer(&mut self) -> Option<Writer<T>> {
        let shared = self.get_shared();
        if shared.try_create_other() {
            Some(Writer {
                shared: self.shared,
                locked: false,
                _owned: PhantomData,
            })
        } else {
            None
        }
    }
}

unsafe impl<T: Send + Sync> Send for Reader<T> {}

impl<T: Clone> Reader<T> {
    pub fn new(value: T) -> Self {
        Self {
            shared: Shared::new(value),
            _owned: PhantomData,
            locked: false,
        }
    }
}

impl<T: Default> Default for Reader<T> {
    fn default() -> Self {
        Self {
            shared: Shared::new_default(),
            _owned: PhantomData,
            locked: false,
        }
    }
}

impl<T> Drop for Reader<T> {
    fn drop(&mut self) {
        // SAFETY: ptr is valid and not used after this
        unsafe {
            Shared::drop(self.shared);
        }
        assert!(!self.locked, "ReadGuard was forgotten")
    }
}

pub struct ReadGuard<'a, T> {
    reader: &'a mut Reader<T>,
    value: &'a T,
}

impl<T> Deref for ReadGuard<'_, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.value
    }
}

impl<T> Drop for ReadGuard<'_, T> {
    fn drop(&mut self) {
        self.reader.locked = false;
        self.reader.get_shared().remove_current_read();
    }
}

pub struct Writer<T> {
    shared: NonNull<Shared<T>>,
    _owned: PhantomData<T>,
    locked: bool,
}

impl<T> Writer<T> {
    fn get_shared(&self) -> &Shared<T> {
        // SAFETY: the shared is always valid
        unsafe { self.shared.as_ref() }
    }

    pub fn write(&mut self) -> WriteGuard<'_, T> {
        if self.locked {
            self.locked = false; // don't also panic in the destructor
            panic!("WriteGuard was forgotten");
        }
        // locking here is always successfull, so we set it unconditionally
        self.locked = true;
        let shared = self.get_shared();
        // SAFETY: won't be called again while the mut ref is still live, as the write method needs a mut ref
        let pre_res = unsafe { shared.write() };
        // SAFETY: was just locked
        let value = unsafe { &mut *pre_res.0.get() };
        let ptr = pre_res.1;

        WriteGuard {
            writer: self,
            value,
            next_read: ptr,
        }
    }
}

impl<T> Drop for Writer<T> {
    fn drop(&mut self) {
        // SAFETY: ptr is valid and not used after this
        unsafe {
            Shared::drop(self.shared);
        }
        assert!(!self.locked, "WriteGuard was forgotten")
    }
}

unsafe impl<T: Send + Sync> Send for Writer<T> {}

/// The data should always be replaced completely. Incremental updates can lead to the reader seeing inconsistent data,
/// as the data structure uses two copies of the data internally and alternates between writing to both of them.
/// The copies will start to diverge when doing changes instead of replacing it.
///
/// I still offer DerefMut instead of just taking a new T, because this allows reusing allocations inside of T instead of
/// dropping and reallocing. This is important for realtime usecases.
pub struct WriteGuard<'a, T> {
    writer: &'a mut Writer<T>,
    value: &'a mut T,
    // if some: set the next_read bit to the ptr value for unlocking
    // if none: set the read_allowed bit. The next_read bit is already correct
    next_read: Option<Ptr>,
}

// i don't really want this. I would prefer to only write, but i need a mut ref to be able to write rt-safe.
// and for DerefMut i need Deref.
impl<T> Deref for WriteGuard<'_, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.value
    }
}

impl<T> DerefMut for WriteGuard<'_, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.value
    }
}

impl<T> Drop for WriteGuard<'_, T> {
    fn drop(&mut self) {
        self.writer.get_shared().unlock_write(self.next_read);
        self.writer.locked = false;
    }
}

#[cfg(test)]
mod internal_tests {
    use crate::Reader;

    #[test]
    fn drop_tests() {
        let mut reader: Reader<u32> = Reader::default();
        assert!(reader.get_shared().is_unique());
        let writer = reader.get_writer().unwrap();
        assert!(!writer.get_shared().is_unique());
        assert!(!reader.get_shared().is_unique());

        drop(writer);
        assert!(reader.get_shared().is_unique());

        let writer = reader.get_writer().unwrap();

        drop(reader);
        assert!(writer.get_shared().is_unique());
    }
}