slice-arena 1.0.0

use std::ffi::OsStr;
use std::path::Path;
use std::sync::Mutex;

/// Thread-safe memory arena.
pub struct SliceArena<T = u8> {
    arenas: Mutex<Vec<Vec<T>>>,
}

impl<T: Clone> SliceArena<T> {
    /// New arena with ~64KB of initial capacity (which will grow exponentially as needed).
    ///
    /// It works with any cloneable element type, but `SliceArena<u8>` is the default (it's like `Vec<u8>`).
    #[inline]
    pub fn new() -> Self {
        Self::with_capacity(((1<<16) / std::mem::size_of::<T>()).next_power_of_two())
    }

    /// New arena with memory reserved for the initial number of elements.
    ///
    /// Arena grows expomnentially, so once this capacity is exhausted,
    /// a double of that size will be allocated.
    ///
    /// Capacity should be much much larger than sizes of slices pushed.
    /// Pushing a slice larger than the capacity leaves previous capacity unused.
    #[inline]
    pub fn with_capacity(capacity: usize) -> Self {
        let mut arenas = Vec::with_capacity(1.max(8.min(capacity)));
        arenas.push(Vec::with_capacity(capacity));
        Self {
            arenas: Mutex::new(arenas),
        }
    }

    /// Copy a slice into the arena. Returns arena-allocated slice.
    pub fn push<'arena>(&'arena self, slice: &[T]) -> &'arena mut [T] {
        let mut a = self.arenas.lock().unwrap();
        let slice_len = slice.len();
        let latest = &mut a[0];
        if latest.capacity() >= latest.len() + slice_len {
            let start = latest.len();
            let _start_ptr = latest.as_ptr();
            latest.extend_from_slice(slice);
            debug_assert_eq!(_start_ptr, latest.as_ptr(), "No realloc");
            unsafe {
                std::mem::transmute::<&mut [T], &'arena mut [T]>(&mut latest[start..start + slice_len])
            }
        } else {
            let new_capacity = slice_len.max(latest.capacity()*2);
            a.push(Vec::with_capacity(new_capacity));
            let last_idx = a.len()-1;
            a.swap(0, last_idx); // for convenience, 0th one should be the latest
            let latest = &mut a[0];
            let _start_ptr = latest.as_ptr();
            debug_assert!(latest.len() == 0);
            latest.extend_from_slice(slice);
            debug_assert!(latest.len() == slice_len);
            debug_assert_eq!(_start_ptr, latest.as_ptr(), "No realloc");
            unsafe {
                std::mem::transmute::<&mut [T], &'arena mut [T]>(&mut latest[0..slice_len])
            }
        }
    }
}

impl SliceArena<u8> {
    /// Store a string slice in the arena.
    #[inline]
    pub fn push_str<'arena>(&'arena self, s: impl AsRef<str>) -> &'arena str {
        let bytes = self.push(s.as_ref().as_bytes());
        unsafe {
            std::str::from_utf8_unchecked(bytes)
        }
    }

    /// Store a Unix `OsStr` in the arena. This function does not exist on Windows.
    #[inline]
    #[cfg(unix)]
    pub fn push_os_str<'arena>(&'arena self, s: &OsStr) -> &'arena OsStr {
        use std::os::unix::ffi::OsStrExt;

        let bytes = self.push(s.as_bytes());
        OsStr::from_bytes(bytes)
    }

    /// Store a Unix `Path` in the arena. This function does not exist on Windows.
    #[inline]
    #[cfg(unix)]
    pub fn push_path<'arena>(&'arena self, str: impl AsRef<Path>) -> &'arena Path {
        let bytes = self.push_os_str(str.as_ref().as_os_str());
        Path::new(bytes)
    }
}


#[test]
fn test_non_copy() {
    let arena = SliceArena::with_capacity(0);
    let a = arena.push(&[String::from("a")]);
    let b = arena.push(&[String::from("b"), String::from("c")]);
    a[0].push_str("1");
    b[0].push_str("2");
    b[1].push_str("3");
    assert_eq!(a[0], "a1");
    assert_eq!(b[0], "b2");
    assert_eq!(b[1], "c3");
}

#[test]
fn test_send_and_sync() {
    fn is_sync<S: Sync>(_: S) {}

    let arena = SliceArena::new();
    arena.push_str("hi");
    std::thread::spawn(move || {
        arena.push_str("bye");
        is_sync(arena);
    }).join().unwrap();
}

#[test]
#[cfg(unix)]
fn test_path() {
    let arena = SliceArena::with_capacity(9999);
    let p = Path::new("some path");
    assert_eq!(arena.push(&[p]), &[p]);
    assert_eq!(arena.push(&[p]), &[p]);
    let p = Path::new("some other path");
    assert_eq!(arena.push(&[p]), &[p]);
}

#[test]
fn kick_tires() {
    let arena = SliceArena::with_capacity(4);
    let h1 = arena.push(b"hello world");
    let h2 = arena.push(b"hello world again");
    assert_eq!(b"hello world", h1);
    assert_eq!(b"hello world again", h2);

    let arena = SliceArena::with_capacity(0);
    assert_eq!(&[0u8; 0], arena.push(&[]));
    assert_eq!(&[0u8; 0], arena.push(&[]));
    let mut grows = String::new();
    let mut in_arena = Vec::new();
    for i in 0..1000 {
        grows.push_str(&format!("and {} ", i));
        let res = arena.push_str(&grows);
        assert_eq!(res, grows);
        in_arena.push(res);
    }
    let h1 = arena.push(b"hello world");
    let h2 = arena.push(b"hello world again");
    assert_eq!(b"hello world", h1);
    assert_eq!(b"hello world again", h2);

    let mut grows = String::new();
    for i in 0..1000 {
        grows.push_str(&format!("and {} ", i));
        assert_eq!(in_arena[i], grows);
    }
}