manual_share/

shared_box.rs

//! # Manually shared box
//! ```
//! use std::thread;
//! use manual_share::SharedBox;
//!
//! let b = Box::new(13);
//! let mut b = SharedBox::from_box(b);
//!
//! let br = b.borrow();
//! let j1 = thread::spawn(move || {
//!     println!("{}", br.get());
//!     br
//! });
//!
//! let br = b.borrow();
//! let j2 = thread::spawn(move || {
//!     println!("{}", br.get() + 1);
//!     br
//! });
//!
//! b.try_return(j1.join().unwrap()).unwrap();
//! b.try_return(j2.join().unwrap()).unwrap();
//!
//! let b = b.try_into_box().unwrap();
//! println!("{:?}", b);
//! ```

use std::ops::Deref;

/// A structure owning the original `Box`
/// which can be used to create multiple `SharedBoxRef` to send to other thread.
/// It uses a counter to record the number of `SharedBoxRef` that has been created and not given back.
///
/// Dropping `SharedBox` without returning all `SharedBoxRef` that it has created leaks its heap memory.
/// When **`panic-on-drop`** feature is enabled, it will panic:
/// ```should_panic
/// let r = {
///     let mut b = manual_share::SharedBox::new(0);
///     b.borrow()
///
///     // dropping SharedBox, panic here
/// };
/// println!("{}", r.get());
/// ```
///
/// When all `SharedBoxRef` have been returned back to `SharedBox`,
/// it will free and properly release its heap memory when `SharedBox` is dropped.
/// ```
/// let mut b = manual_share::SharedBox::new(0);
/// let r = b.borrow();
/// b.try_return(r).unwrap();
///
/// // Dropping SharedBox here will free the heap memory it owns. No panic will occur.
/// ```
#[derive(Debug)]
pub struct SharedBox<T: ?Sized> {
    borrow_count: usize,
    ptr: *mut T,
}

impl<T> SharedBox<T> {
    /// Create a `SharedBox` by creating a `Box` first,
    /// then convert it to a `SharedBox`.
    pub fn new(value: T) -> Self {
        let b = Box::new(value);
        Self::from_box(b)
    }
}

impl<T: ?Sized> SharedBox<T> {
    /// Create a `SharedBox` by consuming a `Box`.
    pub fn from_box(unique: Box<T>) -> Self {
        Self {
            borrow_count: 0,
            ptr: Box::into_raw(unique),
        }
    }

    /// Create a `SharedBoxRef` and increase the borrow count.
    /// # panics
    /// Panics when borrow count overflows `usize`.
    pub fn borrow(&mut self) -> SharedBoxRef<T> {
        self.borrow_count = self.borrow_count.checked_add(1).unwrap();
        SharedBoxRef { ptr: self.ptr }
    }

    /// Try to return back the `SharedBoxRef`.
    /// Returns `Err` if the `SharedBoxRef` does not originate from the same `SharedBox`.
    ///
    /// Decrease the borrow count if not error occurs.
    ///
    /// ```
    /// use manual_share::SharedBox;
    ///
    /// let mut b1 = SharedBox::from_box(Box::new(8));
    /// let r1 = b1.borrow();
    /// b1.try_return(r1).unwrap();
    ///
    /// let mut b2 = SharedBox::from_box(Box::new(9));
    /// let r2 = b2.borrow();
    ///
    /// // Giving SharedBoxRef to the wrong SharedBox returns Err.
    /// let r2 = b1.try_return(r2).unwrap_err();
    ///
    /// b2.try_return(r2).unwrap();
    /// ```
    pub fn try_return(&mut self, reference: SharedBoxRef<T>) -> Result<(), SharedBoxRef<T>> {
        if !core::ptr::eq(self.ptr, reference.ptr) {
            return Err(reference);
        }

        if size_of_val(unsafe { &*self.ptr }) == 0 {
            // ZST types can have multiple allocations to the same address, so we need to check for overflow.
            if let Some(new_count) = self.borrow_count.checked_sub(1) {
                self.borrow_count = new_count;
                let _ = core::mem::ManuallyDrop::new(reference);
                Ok(())
            } else {
                Err(reference)
            }
        } else {
            self.borrow_count -= 1;
            let _ = core::mem::ManuallyDrop::new(reference);
            Ok(())
        }
    }

    /// Try to convert `Self` into a `Box` if all borrowed `SharedBoxRef` has been given back.
    ///
    /// Note that the returned error type is `Self`, dropping it may cause panic.
    ///
    /// ```
    /// use manual_share::SharedBox;
    ///
    /// let b = Box::new(0);
    /// let mut b = SharedBox::from_box(b);
    ///
    /// let r = b.borrow();
    ///
    /// // Try to convert to Box without returning all SharedBoxRef returns Err.
    /// let mut b = b.try_into_box().unwrap_err();
    ///
    /// b.try_return(r).unwrap();
    ///
    /// let b = b.try_into_box().unwrap();
    /// assert_eq!(b, Box::new(0));
    /// ```
    ///
    pub fn try_into_box(self) -> Result<Box<T>, Self> {
        if self.borrow_count > 0 {
            Err(self)
        } else {
            let r = core::mem::ManuallyDrop::new(self);
            Ok(unsafe { Box::from_raw(r.ptr) })
        }
    }
    /// Directly get a reference to the value inside the `SharedBox`.
    /// This use rust built-in lifetime check to ensure the reference is valid as long as the `SharedBox` is alive,
    /// and has no runtime overhead.
    pub fn get(&self) -> &T {
        // SAFETY:
        // The pointer is valid as long as the SharedBox is alive.
        // All other references can only get immutable reference.
        unsafe { &*self.ptr }
    }

    /// Get a mutable reference if borrow count is 0.
    /// ```
    /// use manual_share::SharedBox;
    ///
    /// let mut b = SharedBox::new(0);
    /// let r = b.borrow();
    ///
    /// assert!(b.get_mut().is_none());
    ///
    /// b.try_return(r).unwrap();
    ///
    /// let r = b.get_mut().unwrap();
    /// *r += 1;
    ///
    /// assert_eq!(*b, 1);
    /// ```
    pub fn get_mut(&mut self) -> Option<&mut T> {
        if self.borrow_count == 0 {
            // SAFETY:
            // The pointer is valid as long as the SharedBox is alive.
            // And there is no other references.
            let r = unsafe { &mut *self.ptr };
            Some(r)
        } else {
            None
        }
    }
}

impl<T: ?Sized> Deref for SharedBox<T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        self.get()
    }
}

/// See <https://users.rust-lang.org/t/built-a-crate-to-safely-share-box-and-vec-manually/141138/25?u=newdino> for why it require `Sync`.
unsafe impl<T: Send + Sync> Send for SharedBox<T> {}

unsafe impl<T: Sync> Sync for SharedBox<T> {}

impl<T: ?Sized> Drop for SharedBox<T> {
    fn drop(&mut self) {
        #[cfg(feature = "panic-on-drop")]
        {
            // Let user deal with other panics.
            #[cfg(feature = "do-not-panic-when-panicking")]
            if std::thread::panicking() {
                return;
            }

            if self.borrow_count > 0 {
                panic!("Dropping a SharedBox without giving back all SharedBoxRef")
            }
        }
        // Only drops when there are no outstanding SharedBoxRef values to prevent use-after-free.
        if self.borrow_count == 0 {
            unsafe {
                drop(Box::from_raw(self.ptr));
            }
        }
    }
}

/// A Reference to `SharedBox` that can be sent to other threads.
///
/// Dropping a `SharedBoxRef` leaks the heap memory it points to.
/// When **`panic-on-drop`** feature is enabled, dropping it will panic:
/// ```should_panic
/// let mut b = manual_share::SharedBox::new(1);
/// b.borrow();
///
/// // forget SharedBox to make sure the panic is not caused by dropping it first.
/// std::mem::forget(b);
///
/// // panic here due to dropping SharedBoxRef
/// ```
///
/// Use `SharedBox::try_return` to consume it without causing panic.
/// ```
/// let mut b = manual_share::SharedBox::new(1);
/// let r = b.borrow();
/// b.try_return(r).unwrap();
/// ```
#[derive(Debug)]
pub struct SharedBoxRef<T: ?Sized> {
    ptr: *const T,
}

/// `SharedBoxRef` is like `&T`, which only requires `T: Sync` to implement `Send`.
unsafe impl<T: Sync> Send for SharedBoxRef<T> {}
unsafe impl<T: Sync> Sync for SharedBoxRef<T> {}

impl<T: ?Sized> SharedBoxRef<T> {
    /// Example usage:
    /// ```
    /// let mut b = manual_share::SharedBox::new(42);
    /// let r = b.borrow();
    /// let value = *r.get();
    /// assert_eq!(value, 42);
    ///
    /// b.try_return(r).unwrap();
    /// ```
    ///
    /// The reference got from this method has the same lifetime of the `SharedBoxRef`,
    /// which means it will be invalidated after `SharedBoxRef` is given back to `SharedBox`:
    /// ```compile_fail
    /// let mut b = manual_share::SharedBox::new(42);
    /// let br = b.borrow();
    /// let r = br.get();
    ///
    /// b.try_return(br).unwrap();
    ///
    /// // r is no longer valid here.
    /// println!("{}", r);
    /// ```
    pub fn get(&self) -> &T {
        unsafe { &*self.ptr }
    }
}

impl<T: ?Sized> Drop for SharedBoxRef<T> {
    fn drop(&mut self) {
        #[cfg(feature = "panic-on-drop")]
        {
            #[cfg(feature = "do-not-panic-when-panicking")]
            // Let user deal with other panics.
            if std::thread::panicking() {
                return;
            }

            panic!("SharedBoxRef should not be dropped. Use SharedBox::try_return to consume it.");
        }
    }
}

impl<T: ?Sized> Deref for SharedBoxRef<T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        self.get()
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn zst() {
        let mut b1 = SharedBox::new(());
        let mut b2 = SharedBox::new(());

        let r11 = b1.borrow();
        let r12 = b1.borrow();

        let r2 = b2.borrow();

        b1.try_return(r2).unwrap();

        b1.try_return(r11).unwrap();
        let r12 = b1.try_return(r12).unwrap_err();

        b2.try_return(r12).unwrap();
    }
    #[test]
    fn zst_dyn_trait() {
        // When comparing wide pointers, both the address and the metadata are tested for equality.
        // It should become impossible to return a ZST ref with one vtable to another SharedBox.

        trait T1 {
            fn f(&self) -> u32;
        }

        struct A;
        impl T1 for A {
            fn f(&self) -> u32 {
                0
            }
        }

        struct B;
        impl T1 for B {
            fn f(&self) -> u32 {
                1
            }
        }

        let b1: Box<dyn T1> = Box::new(A);
        let b2: Box<dyn T1> = Box::new(B);

        let mut b1 = SharedBox::from_box(b1);
        let mut b2 = SharedBox::from_box(b2);

        let r1 = b1.borrow();
        let r2 = b2.borrow();

        assert_eq!(r1.get().f(), 0);
        assert_eq!(r2.get().f(), 1);

        let r2 = b1.try_return(r2).unwrap_err();
        let r1 = b2.try_return(r1).unwrap_err();

        assert!(b1.try_return(r1).is_ok());
        assert!(b2.try_return(r2).is_ok());
    }

    #[test]
    fn zst_slice() {
        // When comparing wide pointers, both the address and the metadata are tested for equality.
        // It should become impossible to return a ZST slice with one length to SharedBox with another length.

        let b1: Box<[()]> = Box::new([(); 1]);
        let b2: Box<[()]> = Box::new([(); 2]);

        let mut b1 = SharedBox::from_box(b1);
        let mut b2 = SharedBox::from_box(b2);

        let r1 = b1.borrow();
        let r2 = b2.borrow();

        let r2 = b1.try_return(r2).unwrap_err();
        let r1 = b2.try_return(r1).unwrap_err();

        assert!(b1.try_return(r1).is_ok());
        assert!(b2.try_return(r2).is_ok());
    }
}