Polars Buffers Are Immutable: The Refcounted Storage Behind Every Column

/// [`Buffer`] is a contiguous memory region that can be shared across
/// thread boundaries.
///
/// The easiest way to think about [`Buffer<T>`] is being equivalent to
/// a `Arc<Vec<T>>`, with the following differences:
/// * slicing and cloning is `O(1)`.
/// * it supports external allocated memory
///
/// The easiest way to create one is to use its implementation of `From<Vec<T>>`.
///
/// # Examples
/// ```
/// use polars_buffer::Buffer;
///
/// let mut buffer: Buffer<u32> = vec![1, 2, 3].into();
/// assert_eq!(buffer.as_ref(), [1, 2, 3].as_ref());
///
/// // it supports copy-on-write semantics (i.e. back to a `Vec`)
/// let vec: Vec<u32> = buffer.into_mut().right().unwrap();
/// assert_eq!(vec, vec![1, 2, 3]);
///
/// // cloning and slicing is `O(1)` (data is shared)
/// let mut buffer: Buffer<u32> = vec![1, 2, 3].into();
/// let mut sliced = buffer.clone();
/// sliced.slice(1, 1);
/// assert_eq!(sliced.as_ref(), [2].as_ref());
/// // but cloning forbids getting mut since `slice` and `buffer` now share data
/// assert_eq!(buffer.get_mut_slice(), None);
/// ```

pub struct Buffer<T> {
    /// The internal byte buffer.
    storage: SharedStorage<T>,

    /// A pointer into the buffer where our data starts.
    ptr: *const T,

    // The length of the buffer.
    length: usize,
}

    /// Returns whether underlying data is sliced.
    /// If sliced the [`Buffer`] is backed by
    /// more data than the length of `Self`.
    pub fn is_sliced(&self) -> bool {
        self.storage.len() != self.length
    }

impl<T> Clone for Buffer<T> {
    #[inline(always)]
    fn clone(&self) -> Self {
        Self {
            storage: self.storage.clone(),
            ptr: self.ptr,
            length: self.length,
        }
    }
}

    #[inline]
    pub fn slice_in_place<R: RangeBounds<usize>>(&mut self, range: R) {
        unsafe {
            let Range { start, end } = check_range(range, ..self.len());
            self.ptr = self.ptr.add(start);
            self.length = end - start;
        }
    }

enum BackingStorage {
    Vec {
        original_capacity: usize, // Elements, not bytes.
        vtable: &'static VecVTable,
    },
    ForeignOwner(Box<dyn Any + Send + 'static>),

    /// Backed by some external method which we do not need to take care of,
    /// but we still should refcount and drop the SharedStorageInner.
    External,

    /// Both the backing storage and the SharedStorageInner are leaked, no
    /// refcounting is done. This technically should be a flag on
    /// SharedStorageInner instead of being here, but that would add 8 more
    /// bytes to SharedStorageInner, so here it is.
    Leaked,
}

struct SharedStorageInner<T> {
    ref_count: AtomicU64,
    ptr: *mut T,
    length_in_bytes: usize,
    backing: BackingStorage,
    // https://github.com/rust-lang/rfcs/blob/master/text/0769-sound-generic-drop.md#phantom-data
    phantom: PhantomData<T>,
}

// Allows us to transmute between types while also keeping the original
// stats and drop method of the Vec around.
struct VecVTable {
    size: usize,
    align: usize,
    drop_buffer: unsafe fn(*mut (), usize),
}

impl VecVTable {
    const fn new<T>() -> Self {
        unsafe fn drop_buffer<T>(ptr: *mut (), cap: usize) {
            unsafe { drop(Vec::from_raw_parts(ptr.cast::<T>(), 0, cap)) }
        }

        Self {
            size: size_of::<T>(),
            align: align_of::<T>(),
            drop_buffer: drop_buffer::<T>,
        }
    }

            // We may only go back to a Vec if we originally came from a Vec
            // where the desired size/align matches the original.
            let BackingStorage::Vec {
                original_capacity,
                vtable,
            } = &mut inner.backing
            else {
                return None;
            };

            if vtable.size != size_of::<T>() || vtable.align != align_of::<T>() {
                return None;
            }

        // The pointer must be properly aligned for U.
        // We can skip this check if the alignment of U divides the alignment of T.
        if !align_of::<T>().is_multiple_of(align_of::<U>()) && !inner.ptr.cast::<U>().is_aligned() {
            return Err(self);
        }

impl<T> Clone for SharedStorage<T> {
    fn clone(&self) -> Self {
        let inner = self.inner();
        if !matches!(inner.backing, BackingStorage::Leaked) {
            // Ordering semantics copied from Arc<T>.
            inner.ref_count.fetch_add(1, Ordering::Relaxed);
        }
        Self {
            inner: self.inner,
            phantom: PhantomData,
        }
    }
}

impl<T> Drop for SharedStorage<T> {
    fn drop(&mut self) {
        let inner = self.inner();
        if matches!(inner.backing, BackingStorage::Leaked) {
            return;
        }

        // Ordering semantics copied from Arc<T>.
        if inner.ref_count.fetch_sub(1, Ordering::Release) == 1 {
            std::sync::atomic::fence(Ordering::Acquire);
            unsafe {
                self.drop_slow();
            }
        }
    }
}

impl<T> Drop for SharedStorageInner<T> {
    fn drop(&mut self) {
        match core::mem::replace(&mut self.backing, BackingStorage::External) {
            BackingStorage::ForeignOwner(o) => drop(o),
            BackingStorage::Vec {
                original_capacity,
                vtable,
            } => unsafe {
                // Drop the elements in our slice.
                if std::mem::needs_drop::<T>() {
                    core::ptr::drop_in_place(core::ptr::slice_from_raw_parts_mut(
                        self.ptr,
                        self.length_in_bytes / size_of::<T>(),
                    ));
                }

                // Free the buffer.
                if original_capacity > 0 {
                    (vtable.drop_buffer)(self.ptr.cast(), original_capacity);
                }
            },
            BackingStorage::External | BackingStorage::Leaked => {},
        }
    }
}

    /// Get a [`View`] for a specific set of bytes.
    pub fn push_value_into_buffer(&mut self, bytes: &[u8]) -> View {
        assert!(bytes.len() <= u32::MAX as usize);

        if bytes.len() <= View::MAX_INLINE_SIZE as usize {
            View::new_inline(bytes)
        } else {
            self.total_buffer_len += bytes.len();

            // We want to make sure that we never have to memcopy between buffers. So if the
            // current buffer is not large enough, create a new buffer that is large enough and try
            // to anticipate the larger size.
            if self.in_progress_buffer.len().saturating_add(bytes.len())
                > usize::min(
                    BINVIEW_ARROW_BUFFER_LEN_LIMIT,
                    self.in_progress_buffer.capacity(),
                )
            {
                const {
                    assert!(MAX_EXP_BLOCK_SIZE < BINVIEW_ARROW_BUFFER_LEN_LIMIT);
                }

                // Allocate a new buffer and flush the old buffer.
                let new_capacity = usize::max(
                    bytes.len(),
                    (self.in_progress_buffer.capacity() * 2)
                        .clamp(DEFAULT_BLOCK_SIZE, MAX_EXP_BLOCK_SIZE),
                );

                let in_progress = Vec::with_capacity(new_capacity);
                let flushed = std::mem::replace(&mut self.in_progress_buffer, in_progress);
                if !flushed.is_empty() {
                    self.completed_buffers.push(flushed.into())
                }
            }

#[repr(C, align(4096))]
#[derive(Copy, Clone)]
struct Aligned([u8; 4096]);

// We intentionally leak 8MiB of zeroed memory once so we don't have to
// refcount it.
const GLOBAL_ZERO_SIZE: usize = 8 * 1024 * 1024;
static GLOBAL_ZEROES: LazyLock<SharedStorage<Aligned>> = LazyLock::new(|| {
    const { assert!(GLOBAL_ZERO_SIZE.is_multiple_of(size_of::<Aligned>())) };
    let chunks = GLOBAL_ZERO_SIZE / size_of::<Aligned>();
    let v = vec![Aligned([0; _]); chunks];
    let mut ss = SharedStorage::from_vec(v);
    ss.leak();
    ss
});

    /// Creates an empty SharedStorage.
    pub const fn empty() -> Self {
        assert!(align_of::<T>() <= 1 << 30);
        static INNER: SharedStorageInner<()> = SharedStorageInner {
            ref_count: AtomicU64::new(1),
            ptr: core::ptr::without_provenance_mut(1 << 30), // Very overaligned for any T.
            length_in_bytes: 0,
            backing: BackingStorage::Leaked,
            phantom: PhantomData,
        };

/// A representation of a type as raw bytes with the same alignment as the original type.
pub trait AlignedBytes: Pod + Default + Eq {
    const ALIGNMENT: usize;
    const SIZE: usize;
    const SIZE_ALIGNMENT_PAIR: PrimitiveSizeAlignmentPair;

    type Unaligned: AsRef<[u8]>
        + AsMut<[u8]>
        + std::ops::Index<usize, Output = u8>
        + std::ops::IndexMut<usize, Output = u8>
        + for<'a> TryFrom<&'a [u8]>
        + std::fmt::Debug
        + Default
        + IntoIterator<Item = u8>
        + Pod;

        /// Bytes with a size and alignment.
        ///
        /// This is used to reduce the monomorphizations for routines that solely rely on the size
        /// and alignment of types.
        #[derive(Debug, Copy, Clone, PartialEq, Eq, Default, Pod, Zeroable)]
        #[repr(C, align($alignment))]
        pub struct $name([u8; $size]);

impl_aligned_bytes! {
    (Bytes1Alignment1, 1, 1, S1A1, [u8, i8], u8),
    (Bytes2Alignment2, 2, 2, S2A2, [u16, i16, pf16], u16),
    (Bytes4Alignment4, 4, 4, S4A4, [u32, i32, f32], u32),
    (Bytes8Alignment8, 8, 8, S8A8, [u64, i64, f64], u64),
    (Bytes8Alignment4, 8, 4, S8A4, [days_ms]),
    (Bytes12Alignment4, 12, 4, S12A4, [[u32; 3]]),
    (Bytes16Alignment4, 16, 4, S16A4, [View]),
    (Bytes16Alignment8, 16, 8, S16A8, [months_days_ns]),
    (Bytes16Alignment16, 16, 16, S16A16, [u128, i128]),
    (Bytes32Alignment16, 32, 16, S32A16, [i256]),
}

    /// Returns a mutable reference to its underlying [`Vec`], if possible.
    ///
    /// This operation returns [`Either::Right`] iff this [`Buffer`]:
    /// * has no alive clones
    /// * has not been imported from the C data interface (FFI)
    #[inline]
    pub fn into_mut(mut self) -> Either<Self, Vec<T>> {
        // We lose information if the data is sliced.
        if self.is_sliced() {
            return Either::Left(self);
        }
        match self.storage.try_into_vec() {
            Ok(v) => Either::Right(v),
            Err(slf) => {
                self.storage = slf;
                Either::Left(self)
            },
        }
    }

impl<T: Clone> Buffer<T> {
    pub fn to_vec(self) -> Vec<T> {
        match self.into_mut() {
            Either::Right(v) => v,
            Either::Left(same) => same.as_slice().to_vec(),
        }
    }
}