rt-write-lock / src / shared.rs

use core::{
    cell::UnsafeCell,
    ptr::NonNull,
    sync::atomic::{self, AtomicU8, Ordering},
};

use alloc::boxed::Box;

#[derive(Clone, Copy, PartialEq, Eq)]
pub enum Ptr {
    Value1,
    Value2,
}

impl Ptr {
    fn other(self) -> Self {
        match self {
            Ptr::Value1 => Ptr::Value2,
            Ptr::Value2 => Ptr::Value1,
        }
    }
}

#[derive(Clone, Copy)]
struct State(u8);

impl State {
    /// - Unique
    /// - no read
    /// - new read allowed
    /// - should read value 0
    const INITIAL: u8 = 0b01000;

    const CURRENT_READ_MASK: u8 = 0b00110;
    const INV_CURRENT_READ_MASK: u8 = 0b11001;
    const UNIQUE_MASK: u8 = 0b00001;
    const INV_UNIQUE_MASK: u8 = 0b11110;

    const NEXT_READ: u8 = 0b10000;
    const INV_NEXT_READ: u8 = 0b01111;
    const READ_ALLOWED: u8 = 0b01000;
    const INV_READ_ALLOWED: u8 = 0b10111;

    fn new(value: u8) -> Self {
        // could also be assert unchecked
        debug_assert_eq!(
            value & 0b1110_0000,
            0,
            "only the lower 5 bits should be used. value: {value:b}"
        );
        debug_assert!(
            (value & Self::CURRENT_READ_MASK).count_ones() <= 1,
            "max 1 read. value: {value:b}"
        );
        Self(value)
    }

    // returns None if no read is allowed
    fn get_next_read(self) -> Option<Ptr> {
        // is reading allowed
        if self.0 & Self::READ_ALLOWED == 0 {
            None
        } else {
            // which value should be read
            if self.0 & Self::NEXT_READ == 0 {
                Some(Ptr::Value1)
            } else {
                Some(Ptr::Value2)
            }
        }
    }

    fn set_current_read(self, ptr: Ptr) -> Self {
        debug_assert!(
            self.0 & Self::CURRENT_READ_MASK == 0,
            "no read should be active. value: {:b}",
            self.0
        );

        let mask = match ptr {
            Ptr::Value1 => 0b10,
            Ptr::Value2 => 0b100,
        };

        Self(self.0 | mask)
    }

    // none if there is no current read
    fn get_current_read(self) -> Option<Ptr> {
        match self.0 & Self::CURRENT_READ_MASK {
            0b000 => None,
            0b010 => Some(Ptr::Value1),
            0b100 => Some(Ptr::Value2),
            _ => unreachable!(),
        }
    }

    fn is_unique(self) -> bool {
        self.0 & Self::UNIQUE_MASK == 0
    }

    fn disallow_read(self) -> Self {
        Self(self.0 & Self::INV_READ_ALLOWED)
    }
}

pub struct Shared<T> {
    /// # Bits from low to high
    /// | bit | meaning |
    /// |---|---|
    /// | 0 | 0: unique, 1: second object exists |
    /// | 1 | is value 1 being read (never both) |
    /// | 2 | is value 2 being read (never both) |
    /// | 3 | new read allowed |
    /// | 4 | which value should be read (0: value1, 1: value2) |
    state: AtomicU8,
    value_1: UnsafeCell<T>,
    value_2: UnsafeCell<T>,
}

impl<T: Clone> Shared<T> {
    pub fn new(value: T) -> NonNull<Shared<T>> {
        NonNull::new(Box::into_raw(Box::new(Shared {
            state: AtomicU8::new(State::INITIAL),
            value_1: UnsafeCell::new(value.clone()),
            value_2: UnsafeCell::new(value),
        })))
        .unwrap()
    }
}

impl<T: Default> Shared<T> {
    pub fn new_default() -> NonNull<Shared<T>> {
        NonNull::new(Box::into_raw(Box::new(Shared {
            state: AtomicU8::new(State::INITIAL),
            value_1: UnsafeCell::new(T::default()),
            value_2: UnsafeCell::new(T::default()),
        })))
        .unwrap()
    }
}

impl<T> Shared<T> {
    // returns a unsafecell to allow some lifetime stuff
    pub fn try_read(&self) -> Option<&UnsafeCell<T>> {
        let result = self
            .state
            .fetch_update(Ordering::Acquire, Ordering::Relaxed, |value| {
                let state = State::new(value);
                let read = state.get_next_read();
                read.map(|p| state.set_current_read(p).0)
            });
        match result {
            // read is allowed. could be unwrap unchecked
            Ok(v) => {
                // unwrap ok, because Ok returned. could also be unchecked
                let ptr = State::new(v).get_next_read().unwrap();
                // SAFETY: atomic was just locked and the read state was set.
                Some(match ptr {
                    Ptr::Value1 => &self.value_1,
                    Ptr::Value2 => &self.value_2,
                })
            }
            // read isn't allowed
            Err(_) => None,
        }
    }

    /// The Option is None if the reader is blocked by this write. This can happen if the reader is keeping
    /// a readlock for a long time. If it is None the writelock doesn't need to set the next_read bit when
    /// unlocking. It only needs to remove the no_read bit.
    ///
    /// If this is Some the reader isn't blocked. The writelock needs to set the next_read bit to the Ptr value
    pub unsafe fn write(&self) -> (&UnsafeCell<T>, Option<Ptr>) {
        let result = self
            .state
            .fetch_update(Ordering::Acquire, Ordering::Acquire, |v| {
                let state = State::new(v);
                let current_read = state.get_current_read();
                // this can't be None because if WriteLock was forgotten instead of dropped the outer API layer panics,
                // before calling this function
                let next_read = unsafe { state.get_next_read().unwrap_unchecked() };

                match (current_read, next_read) {
                    (Some(c), n) if c != n => Some(state.disallow_read().0),
                    _ => {
                        // one of the following situations:
                        // - no read currently
                        // - should read and current read are the same
                        // both mean that the reader isn't behind, so no need to block it
                        //
                        // This is the branch that should be taken in most of the cases, so this fetch_update should be
                        // a single atomic load most of the time
                        None
                    }
                }
            });

        // this basically duplicates the logic from the closure. I am not sure if this is a good idea, but i want to
        // keep the closure as small as possible to not make the cmp exchange loop harder than necessary
        let pre_res = match result {
            Ok(v) => {
                let state = State::new(v);
                // SAFETY: when locking for writing reading is always allowed. This operates on the state before the
                // cmp exchange loop, so this is still true.
                let next_read = unsafe { state.get_next_read().unwrap_unchecked() };

                (next_read, None)
            }
            Err(v) => {
                let state = State::new(v);
                // SAFETY: when locking for writing reading is always allowed. This operates on the state before the
                // cmp exchange loop, so this is still true.
                let next_read = unsafe { state.get_next_read().unwrap_unchecked() };
                let write = next_read.other();

                (write, Some(write))
            }
        };

        let value = match pre_res.0 {
            Ptr::Value1 => &self.value_1,
            Ptr::Value2 => &self.value_2,
        };

        (value, pre_res.1)
    }

    pub fn remove_current_read(&self) {
        // relaxed because modifivations are not possible from the read side
        let old_state = self
            .state
            .fetch_and(State::INV_CURRENT_READ_MASK, Ordering::Relaxed);
        debug_assert!(
            (old_state & State::CURRENT_READ_MASK).count_ones() == 1,
            "don't call this method if there is no active read. value: {old_state:b}"
        )
    }

    pub fn unlock_write(&self, read_state: Option<Ptr>) {
        match read_state {
            // set the next read bit to the correct value
            // value1 so remove the bit
            Some(Ptr::Value1) => self
                .state
                .fetch_and(State::INV_NEXT_READ, Ordering::Release),
            // value2 so set the bit
            Some(Ptr::Value2) => self.state.fetch_or(State::NEXT_READ, Ordering::Release),
            // only set the allow read bit
            None => self.state.fetch_or(State::READ_ALLOWED, Ordering::Release),
        };
    }

    /// SAFETY: this has to be valid and gets invalidated by the function
    pub unsafe fn drop(this: NonNull<Self>) {
        // SAFETY: function preconditions
        let old_state = unsafe { this.as_ref() }
            .state
            .fetch_and(State::INV_UNIQUE_MASK, Ordering::Release);

        if State::new(old_state).is_unique() {
            // see std::arc
            atomic::fence(Ordering::Acquire);

            // SAFETY: was alloced with box and is unique
            core::mem::drop(unsafe { Box::from_raw(this.as_ptr()) })
        }
    }

    /// this allows then duplicating the pointer to it
    pub fn try_create_other(&self) -> bool {
        // relaxed, because it is singlethreaded if this succeeds
        let old_state = State::new(self.state.fetch_or(State::UNIQUE_MASK, Ordering::Relaxed));
        // if it was unique before creating another is successfull
        old_state.is_unique()
    }

    #[cfg(test)]
    pub fn is_unique(&self) -> bool {
        State::new(self.state.load(Ordering::Relaxed)).is_unique()
    }
}

#[cfg(test)]
mod state_tests {
    use crate::shared::{Ptr, State};

    #[test]
    fn unique() {
        let state = State::new(State::INITIAL);
        assert!(state.is_unique());
        let non_unique = State::new(state.0 | State::UNIQUE_MASK);
        assert!(!non_unique.is_unique());
        let unique_again = State::new(non_unique.0 & State::INV_UNIQUE_MASK);
        assert!(unique_again.is_unique());
    }

    #[test]
    fn read_state() {
        let state = State::new(State::INITIAL);
        assert!(state.get_current_read().is_none());
        // ptr 1 read
        let read1 = state.set_current_read(Ptr::Value1);
        assert!(read1.get_current_read() == Some(Ptr::Value1));
        // ptr 2 read
        let read2 = state.set_current_read(Ptr::Value2);
        assert!(read2.get_current_read() == Some(Ptr::Value2));
    }
}