prealloc_ref_vec 0.1.1

#![no_std]
#![forbid(unsafe_code)]

//! Create temporary Vecs of references without allocating/deallocating (useful for
//! realtime code). This crate is implemented without any unsafe code, and can be
//! used without the standard library (#![no_std]).
//!
//! Note, if the maximum capacity is known at compile time and is small enough to
//! fit on the stack, then the same effect can be more easily achieved using
//! [`arrayvec`](https://crates.io/crates/arrayvec) instead.
//!
//! # Example
//! ```
//! # use prealloc_ref_vec::PreallocRefVec;
//! // Contruct the object to hold the allocation. Here we specify
//! // we want the data type to be slices of `f32` values.
//! //
//! // This object can be stored as a field in a struct for later use.
//! let capacity = 100;
//! let mut v_alloc: PreallocRefVec<[f32]> = PreallocRefVec::new(capacity);
//!
//! // ...
//!
//! let slice_1: [f32; 4] = [0.0; 4];
//! let slice_2: [f32; 4] = [0.0; 4];
//!
//! // Temporarily use the allocation to construct a Vec<&T>
//! // (Vec<&[f32]> in this case)
//! let mut tmp_vec = v_alloc.get_tmp();
//!
//! // The temporary vec is always cleared when created.
//! assert_eq!(tmp_vec.len(), 0);
//! // The temporary vec uses the internal allocation.
//! assert!(tmp_vec.capacity() >= capacity);
//!
//! tmp_vec.push(&slice_1);
//! tmp_vec.push(&slice_2);
//!
//! // Dropping the temporary vec returns the allocation to be used again.
//! drop(tmp_vec);
//!
//! // ...
//!
//! // Alternatively, a closure can be used to construct a temporary vec:
//! v_alloc.with_tmp(|tmp_vec| {
//!     assert_eq!(tmp_vec.len(), 0);
//!     assert!(tmp_vec.capacity() >= capacity);
//!     tmp_vec.push(&slice_1);
//!     tmp_vec.push(&slice_2);
//! });
//! ```

extern crate alloc;

use alloc::vec::Vec;
use core::ops::{Deref, DerefMut, Index, IndexMut};

/// A struct that is used to create a temporary `Vec` of references/unsized
/// data without allocating/deallocating. Useful for realtime code.
///
/// Note, if the maximum capacity is known at compile time and is small enough to
/// fit on the stack, then the same effect can be more easily achieved using
/// [`arrayvec`](https://crates.io/crates/arrayvec) instead.
///
/// # Example
/// ```
/// # use prealloc_ref_vec::PreallocRefVec;
/// // Contruct the object to hold the allocation. Here we specify
/// // we want the data type to be slices of `f32` values.
/// //
/// // This object can be stored as a field in a struct for later use.
/// let capacity = 100;
/// let mut v_alloc: PreallocRefVec<[f32]> = PreallocRefVec::new(capacity);
///
/// // ...
///
/// let slice_1: [f32; 4] = [0.0; 4];
/// let slice_2: [f32; 4] = [0.0; 4];
///
/// // Temporarily use the allocation to construct a Vec<&T>
/// // (Vec<&[f32]> in this case)
/// let mut tmp_vec = v_alloc.get_tmp();
///
/// // The temporary vec is always cleared when created.
/// assert_eq!(tmp_vec.len(), 0);
/// // The temporary vec uses the internal allocation.
/// assert!(tmp_vec.capacity() >= capacity);
///
/// tmp_vec.push(&slice_1);
/// tmp_vec.push(&slice_2);
///
/// // Dropping the temporary vec returns the allocation to be used again.
/// drop(tmp_vec);
///
/// // ...
///
/// // Alternatively, a closure can be used to construct a temporary vec:
/// v_alloc.with_tmp(|tmp_vec| {
///     assert_eq!(tmp_vec.len(), 0);
///     assert!(tmp_vec.capacity() >= capacity);
///     tmp_vec.push(&slice_1);
///     tmp_vec.push(&slice_2);
/// });
/// ```
pub struct PreallocRefVec<T: ?Sized + 'static> {
    buffer: Vec<&'static mut T>,
}

impl<T: ?Sized + 'static> PreallocRefVec<T> {
    /// Construct a new [`PreallocRefVec`] with the given allocated capacity.
    ///
    /// # Panics
    /// Panics if the capacity exceeds `isize::MAX` bytes.
    pub fn new(capacity: usize) -> Self {
        let mut buffer = Vec::new();
        buffer.reserve_exact(capacity);

        Self { buffer }
    }

    /// Create a temporary `Vec` of immutable references/unsized data.
    ///
    /// # Example
    /// ```
    /// # use prealloc_ref_vec::PreallocRefVec;
    /// // Contruct the object to hold the allocation. Here we specify
    /// // we want the data type to be slices of `f32` values.
    /// //
    /// // This object can be stored as a field in a struct for later use.
    /// let capacity = 100;
    /// let mut v_alloc: PreallocRefVec<[f32]> = PreallocRefVec::new(capacity);
    ///
    /// // ...
    ///
    /// let slice_1: [f32; 4] = [0.0; 4];
    /// let slice_2: [f32; 4] = [0.0; 4];
    ///
    /// // Temporarily use the allocation to construct a Vec<&T>
    /// // (Vec<&[f32]> in this case)
    /// let mut tmp_vec = v_alloc.get_tmp();
    ///
    /// // The temporary vec is always cleared when created.
    /// assert_eq!(tmp_vec.len(), 0);
    /// // The temporary vec uses the internal allocation.
    /// assert!(tmp_vec.capacity() >= capacity);
    ///
    /// tmp_vec.push(&slice_1);
    /// tmp_vec.push(&slice_2);
    ///
    /// // Dropping the temporary vec returns the allocation to be used again.
    /// drop(tmp_vec);
    /// ```
    pub fn get_tmp<'a>(&'a mut self) -> TmpRefVec<'a, T> {
        let mut tmp: Vec<&'static mut T> = Vec::new();
        core::mem::swap(&mut self.buffer, &mut tmp);
        tmp.clear();

        TmpRefVec {
            data: tmp.into_iter().map(|_| unreachable!()).collect(),
            buffer: &mut self.buffer,
        }
    }

    /// Create a temporary `Vec` of mutable references/unsized data.
    ///
    /// # Example
    /// ```
    /// # use prealloc_ref_vec::PreallocRefVec;
    /// // Contruct the object to hold the allocation. Here we specify
    /// // we want the data type to be slices of `f32` values.
    /// //
    /// // This object can be stored as a field in a struct for later use.
    /// let capacity = 100;
    /// let mut v_alloc: PreallocRefVec<[f32]> = PreallocRefVec::new(capacity);
    ///
    /// // ...
    ///
    /// let mut slice_1: [f32; 4] = [0.0; 4];
    /// let mut slice_2: [f32; 4] = [0.0; 4];
    ///
    /// // Temporarily use the allocation to construct a Vec<&mut T>
    /// // (Vec<&mut [f32]> in this case)
    /// let mut tmp_vec = v_alloc.get_tmp_mut();
    ///
    /// // The temporary vec is always cleared when created.
    /// assert_eq!(tmp_vec.len(), 0);
    /// // The temporary vec uses the internal allocation.
    /// assert!(tmp_vec.capacity() >= capacity);
    ///
    /// tmp_vec.push(&mut slice_1);
    /// tmp_vec.push(&mut slice_2);
    ///
    /// // Dropping the temporary vec returns the allocation to be used again.
    /// drop(tmp_vec);
    /// ```
    pub fn get_tmp_mut<'a>(&'a mut self) -> TmpRefMutVec<'a, T> {
        let mut tmp: Vec<&'static mut T> = Vec::new();
        core::mem::swap(&mut self.buffer, &mut tmp);
        tmp.clear();

        TmpRefMutVec {
            data: tmp.into_iter().map(|_| unreachable!()).collect(),
            buffer: &mut self.buffer,
        }
    }

    /// Get a temporary `Vec` of immutable references/unsized data.
    ///
    /// # Example
    /// ```
    /// # use prealloc_ref_vec::PreallocRefVec;
    /// // Contruct the object to hold the allocation. Here we specify
    /// // we want the data type to be slices of `f32` values.
    /// //
    /// // This object can be stored as a field in a struct for later use.
    /// let capacity = 100;
    /// let mut v_alloc: PreallocRefVec<[f32]> = PreallocRefVec::new(capacity);
    ///
    /// // ...
    ///
    /// let slice_1: [f32; 4] = [0.0; 4];
    /// let slice_2: [f32; 4] = [0.0; 4];
    ///
    /// v_alloc.with_tmp(|tmp_vec| {
    ///     // The temporary vec is always cleared when created.
    ///     assert_eq!(tmp_vec.len(), 0);
    ///     // The temporary vec uses the internal allocation.
    ///     assert!(tmp_vec.capacity() >= capacity);
    ///
    ///     tmp_vec.push(&slice_1);
    ///     tmp_vec.push(&slice_2);
    /// });
    /// ```
    pub fn with_tmp<'a>(&'a mut self, f: impl FnOnce(&mut Vec<&'a T>)) {
        let mut v = self.get_tmp();
        (f)(&mut v.data);
    }

    /// Get a temporary `Vec` of mutable references/unsized data.
    ///
    /// # Example
    /// ```
    /// # use prealloc_ref_vec::PreallocRefVec;
    /// // Contruct the object to hold the allocation. Here we specify
    /// // we want the data type to be slices of `f32` values.
    /// //
    /// // This object can be stored as a field in a struct for later use.
    /// let capacity = 100;
    /// let mut v_alloc: PreallocRefVec<[f32]> = PreallocRefVec::new(capacity);
    ///
    /// // ...
    ///
    /// let mut slice_1: [f32; 4] = [0.0; 4];
    /// let mut slice_2: [f32; 4] = [0.0; 4];
    ///
    /// v_alloc.with_tmp_mut(|tmp_vec| {
    ///     // The temporary vec is always cleared when created.
    ///     assert_eq!(tmp_vec.len(), 0);
    ///     // The temporary vec uses the internal allocation.
    ///     assert!(tmp_vec.capacity() >= capacity);
    ///
    ///     tmp_vec.push(&mut slice_1);
    ///     tmp_vec.push(&mut slice_2);
    /// });
    /// ```
    pub fn with_tmp_mut<'a>(&'a mut self, f: impl FnOnce(&mut Vec<&'a mut T>)) {
        let mut v = self.get_tmp_mut();
        (f)(&mut v.data);
    }

    /// The allocated capacity
    pub const fn capacity(&self) -> usize {
        self.buffer.capacity()
    }
}

impl<T: ?Sized + 'static> From<Vec<&'static mut T>> for PreallocRefVec<T> {
    fn from(buffer: Vec<&'static mut T>) -> Self {
        Self { buffer }
    }
}

impl<T: ?Sized + 'static> From<PreallocRefVec<T>> for Vec<&'static mut T> {
    fn from(vec: PreallocRefVec<T>) -> Self {
        vec.buffer
    }
}

/// A temporary `Vec` of immutable references/unsized data created using a
/// [`PreallocRefVec`]. Useful for realtime code.
///
/// This struct can be automatically dereferenced into a `Vec<&T>`.
///
/// # Example
/// ```
/// # use prealloc_ref_vec::PreallocRefVec;
/// // Contruct the object to hold the allocation. Here we specify
/// // we want the data type to be slices of `f32` values.
/// //
/// // This object can be stored as a field in a struct for later use.
/// let capacity = 100;
/// let mut v_alloc: PreallocRefVec<[f32]> = PreallocRefVec::new(capacity);
///
/// // ...
///
/// let slice_1: [f32; 4] = [0.0; 4];
/// let slice_2: [f32; 4] = [0.0; 4];
///
/// // Temporarily use the allocation to construct a Vec<&T>
/// // (Vec<&[f32]> in this case)
/// let mut tmp_vec = v_alloc.get_tmp();
///
/// // The temporary vec is always cleared when created.
/// assert_eq!(tmp_vec.len(), 0);
/// // The temporary vec uses the internal allocation.
/// assert!(tmp_vec.capacity() >= capacity);
///
/// tmp_vec.push(&slice_1);
/// tmp_vec.push(&slice_2);
///
/// // Dropping the temporary vec returns the allocation to be used again.
/// drop(tmp_vec);
/// ```
pub struct TmpRefVec<'a, T: ?Sized + 'static> {
    /// A temporary `Vec` of immutable references/unsized data created using a
    /// [`PreallocRefVec`].
    pub data: Vec<&'a T>,
    buffer: &'a mut Vec<&'static mut T>,
}

impl<'a, T: ?Sized + 'static> Deref for TmpRefVec<'a, T> {
    type Target = Vec<&'a T>;

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

impl<'a, T: ?Sized + 'static> DerefMut for TmpRefVec<'a, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.data
    }
}

impl<'a, T: ?Sized + 'static> Index<usize> for TmpRefVec<'a, T> {
    type Output = T;

    fn index(&self, index: usize) -> &Self::Output {
        &self.data[index]
    }
}

impl<'a, T: ?Sized + 'static> Drop for TmpRefVec<'a, T> {
    fn drop(&mut self) {
        let mut tmp: Vec<&'a T> = Vec::new();
        core::mem::swap(&mut self.data, &mut tmp);
        tmp.clear();
        *self.buffer = tmp.into_iter().map(|_| unreachable!()).collect();
    }
}

/// A temporary `Vec` of mutable references/unsized data created using a
/// [`PreallocRefVec`]. Useful for realtime code.
///
/// This struct can be automatically dereferenced into a `Vec<&mut T>`.
///
/// # Example
/// ```
/// # use prealloc_ref_vec::PreallocRefVec;
/// // Contruct the object to hold the allocation. Here we specify
/// // we want the data type to be slices of `f32` values.
/// //
/// // This object can be stored as a field in a struct for later use.
/// let capacity = 100;
/// let mut v_alloc: PreallocRefVec<[f32]> = PreallocRefVec::new(capacity);
///
/// // ...
///
/// let mut slice_1: [f32; 4] = [0.0; 4];
/// let mut slice_2: [f32; 4] = [0.0; 4];
///
/// // Temporarily use the allocation to construct a Vec<&mut T>
/// // (Vec<&mut [f32]> in this case)
/// let mut tmp_vec = v_alloc.get_tmp_mut();
///
/// // The temporary vec is always cleared when created.
/// assert_eq!(tmp_vec.len(), 0);
/// // The temporary vec uses the internal allocation.
/// assert!(tmp_vec.capacity() >= capacity);
///
/// tmp_vec.push(&mut slice_1);
/// tmp_vec.push(&mut slice_2);
///
/// // Dropping the temporary vec returns the allocation to be used again.
/// drop(tmp_vec);
/// ```
pub struct TmpRefMutVec<'a, T: ?Sized + 'static> {
    /// A temporary `Vec` of mutable references/unsized data created using a
    /// [`PreallocRefVec`].
    pub data: Vec<&'a mut T>,
    buffer: &'a mut Vec<&'static mut T>,
}

impl<'a, T: ?Sized + 'static> Deref for TmpRefMutVec<'a, T> {
    type Target = Vec<&'a mut T>;

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

impl<'a, T: ?Sized + 'static> DerefMut for TmpRefMutVec<'a, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.data
    }
}

impl<'a, T: ?Sized + 'static> Index<usize> for TmpRefMutVec<'a, T> {
    type Output = T;

    fn index(&self, index: usize) -> &Self::Output {
        &self.data[index]
    }
}

impl<'a, T: ?Sized + 'static> IndexMut<usize> for TmpRefMutVec<'a, T> {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        &mut self.data[index]
    }
}

impl<'a, T: ?Sized + 'static> Drop for TmpRefMutVec<'a, T> {
    fn drop(&mut self) {
        let mut tmp: Vec<&'a mut T> = Vec::new();
        core::mem::swap(&mut self.data, &mut tmp);
        tmp.clear();
        *self.buffer = tmp.into_iter().map(|_| unreachable!()).collect();
    }
}

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

    use assert_no_alloc::*;

    #[global_allocator]
    static A: AllocDisabler = AllocDisabler;

    #[test]
    fn test() {
        const CAPACITY: usize = 10;

        let mut v: PreallocRefVec<[f32]> = PreallocRefVec::new(CAPACITY);

        assert_no_alloc(|| test_inner(&mut v, CAPACITY));
    }

    // Make sure the compiler isn't doing any fancy allocation optimizations
    #[inline(never)]
    fn test_inner(v: &mut PreallocRefVec<[f32]>, capacity: usize) {
        let mut data: [f32; 4] = [0.0; 4];
        let v_ptr = v.buffer.as_ptr();

        assert!(v.capacity() >= capacity);
        assert_eq!(v.buffer.len(), 0);

        {
            let mut r = v.get_tmp();

            assert!(r.capacity() >= capacity);
            assert_eq!(r.data.len(), 0);
            assert!(core::ptr::addr_eq(v_ptr, r.data.as_ptr()));

            r.data.push(&data);
        }

        assert!(v.capacity() >= capacity);
        assert_eq!(v.buffer.len(), 0);
        assert!(core::ptr::addr_eq(v_ptr, v.buffer.as_ptr()));

        {
            let mut r = v.get_tmp_mut();

            assert!(r.capacity() >= capacity);
            assert_eq!(r.data.len(), 0);
            assert!(core::ptr::addr_eq(v_ptr, r.data.as_ptr()));

            r.data.push(&mut data);
        }

        assert!(v.capacity() >= capacity);
        assert_eq!(v.buffer.len(), 0);
        assert!(core::ptr::addr_eq(v_ptr, v.buffer.as_ptr()));
    }
}