//! Vectors allocated in arenas, with small per-vector overhead.

use std::marker::PhantomData;

use std::mem::MaybeUninit;

use std::ops::Range;

/// A vector of `T` stored within a `BumpArena`.

///

/// This is something like a normal `Vec`, except that all accesses

/// and updates require a separate borrow of the `BumpArena`. This, in

/// turn, makes the Vec itself very compact: only three `u32`s (12

/// bytes). The `BumpSlice` variant is only two `u32`s (8 bytes) and

/// is sufficient to reconstruct a slice, but not grow the vector.

///

/// The `BumpVec` does *not* implement `Clone` or `Copy`; it

/// represents unique ownership of a range of indices in the arena. If

/// dropped, those indices will be unavailable until the arena is

/// freed. This is "fine" (it is normally how arena allocation

/// works). To explicitly free and make available for some

/// allocations, a very rudimentary reuse mechanism exists via

/// `BumpVec::free(arena)`. (The allocation path opportunistically

/// checks the first range on the freelist, and can carve off a piece

/// of it if larger than needed, but it does not attempt to traverse

/// the entire freelist; this is a compromise between bump-allocation

/// speed and memory efficiency, which also influences speed through

/// cached-memory reuse.)

///

/// The type `T` should not have a `Drop` implementation. This

/// typically means that it does not own any boxed memory,

/// sub-collections, or other resources. This is important for the

/// efficiency of the data structure (otherwise, to call `Drop` impls,

/// the arena needs to track which indices are live or dead; the

/// BumpVec itself cannot do the drop because it does not retain a

/// reference to the arena). Note that placing a `T` with a `Drop`

/// impl in the arena is still *safe*, because leaking (that is, never

/// calling `Drop::drop()`) is safe. It is merely less efficient, and

/// so should be avoided if possible.

#[derive(Debug)]

pub struct BumpVec<T> {

    base: u32,

    len: u32,

    cap: u32,

    _phantom: PhantomData<T>,

}

/// A slice in an arena: like a `BumpVec`, but has a fixed size that

/// cannot grow. The size of this struct is one 32-bit word smaller

/// than `BumpVec`. It is copyable/cloneable because it will never be

/// freed.

#[derive(Debug, Clone, Copy)]

pub struct BumpSlice<T> {

    base: u32,

    len: u32,

    _phantom: PhantomData<T>,

}

#[derive(Default)]

pub struct BumpArena<T> {

    vec: Vec<MaybeUninit<T>>,

    freelist: Vec<Range<u32>>,

}

impl<T> BumpArena<T> {

    /// Create a new arena into which one can allocate `BumpVec`s.

    pub fn new() -> Self {

        Self {

            vec: vec![],

            freelist: vec![],

        }

    }

    /// Create a new arena, pre-allocating space for `cap` total `T`

    /// elements.

    pub fn arena_with_capacity(cap: usize) -> Self {

        Self {

            vec: Vec::with_capacity(cap),

            freelist: Vec::with_capacity(cap / 16),

        }

    }

    /// Create a new `BumpVec` with the given pre-allocated capacity

    /// and zero length.

    pub fn vec_with_capacity(&mut self, cap: usize) -> BumpVec<T> {

        let cap = u32::try_from(cap).unwrap();

        if let Some(range) = self.maybe_freelist_alloc(cap) {

            BumpVec {

                base: range.start,

                len: 0,

                cap,

                _phantom: PhantomData,

            }

        } else {

            let base = self.vec.len() as u32;

            for _ in 0..cap {

                self.vec.push(MaybeUninit::uninit());

            }

            BumpVec {

                base,

                len: 0,

                cap,

                _phantom: PhantomData,

            }

        }

    }

    /// Create a new `BumpVec` with a single element. The capacity is

    /// also only one element; growing the vector further will require

    /// a reallocation.

    pub fn single(&mut self, t: T) -> BumpVec<T> {

        let mut vec = self.vec_with_capacity(1);

        unsafe {

            self.write_into_index(vec.base, t);

        }

        vec.len = 1;

        vec

    }

    /// Create a new `BumpVec` with the sequence from an iterator.

    pub fn from_iter<I: Iterator<Item = T>>(&mut self, i: I) -> BumpVec<T> {

        let base = self.vec.len() as u32;

        self.vec.extend(i.map(|item| MaybeUninit::new(item)));

        let len = self.vec.len() as u32 - base;

        BumpVec {

            base,

            len,

            cap: len,

            _phantom: PhantomData,

        }

    }

    /// Append two `BumpVec`s, returning a new one. Consumes both

    /// vectors. This will use the capacity at the end of `a` if

    /// possible to move `b`'s elements into place; otherwise it will

    /// need to allocate new space.

    pub fn append(&mut self, a: BumpVec<T>, b: BumpVec<T>) -> BumpVec<T> {

        if (a.cap - a.len) >= b.len {

            self.append_into_cap(a, b)

        } else {

            self.append_into_new(a, b)

        }

    }

    /// Helper: read the `T` out of a given arena index. After

    /// reading, that index becomes uninitialized.

    unsafe fn read_out_of_index(&self, index: u32) -> T {

        // Note that we don't actually *track* uninitialized status

        // (and this is fine because we will never `Drop` and we never

        // allow a `BumpVec` to refer to an uninitialized index, so

        // the bits are effectively dead). We simply read the bits out

        // and return them.

        self.vec[index as usize].as_ptr().read()

    }

    /// Helper: write a `T` into the given arena index. Index must

    /// have been uninitialized previously.

    unsafe fn write_into_index(&mut self, index: u32, t: T) {

        self.vec[index as usize].as_mut_ptr().write(t);

    }

    /// Helper: move a `T` from one index to another. Old index

    /// becomes uninitialized and new index must have previously been

    /// uninitialized.

    unsafe fn move_item(&mut self, from: u32, to: u32) {

        let item = self.read_out_of_index(from);

        self.write_into_index(to, item);

    }

    /// Helper: push a `T` onto the end of the arena, growing its

    /// storage. The `T` to push is read out of another index, and

    /// that index subsequently becomes uninitialized.

    unsafe fn push_item(&mut self, from: u32) -> u32 {

        let index = self.vec.len() as u32;

        let item = self.read_out_of_index(from);

        self.vec.push(MaybeUninit::new(item));

        index

    }

    /// Helper: append `b` into the capacity at the end of `a`.

    fn append_into_cap(&mut self, mut a: BumpVec<T>, b: BumpVec<T>) -> BumpVec<T> {

        debug_assert!(a.cap - a.len >= b.len);

        for i in 0..b.len {

            // Safety: initially, the indices in `b` are initialized;

            // the indices in `a`'s cap, beyond its length, are

            // uninitialized. We move the initialized contents from

            // `b` to the tail beyond `a`, and we consume `b` (so it

            // no longer exists), and we update `a`'s length to cover

            // the initialized contents in their new location.

            unsafe {

                self.move_item(b.base + i, a.base + a.len + i);

            }

        }

        a.len += b.len;

        b.free(self);

        a

    }

    /// Helper: return a range of indices that are available

    /// (uninitialized) according to the freelist for `len` elements,

    /// if possible.

    fn maybe_freelist_alloc(&mut self, len: u32) -> Option<Range<u32>> {

        if let Some(entry) = self.freelist.last_mut() {

            if entry.len() >= len as usize {

                let base = entry.start;

                entry.start += len;

                if entry.start == entry.end {

                    self.freelist.pop();

                }

                return Some(base..(base + len));

            }

        }

        None

    }

    /// Helper: append `a` and `b` into a completely new allocation.

    fn append_into_new(&mut self, a: BumpVec<T>, b: BumpVec<T>) -> BumpVec<T> {

        // New capacity: round up to a power of two.

        let len = a.len + b.len;

        let cap = round_up_power_of_two(len);

        if let Some(range) = self.maybe_freelist_alloc(cap) {

            for i in 0..a.len {

                // Safety: the indices in `a` must be initialized. We read

                // out the item and copy it to a new index; the old index

                // is no longer covered by a BumpVec, because we consume

                // `a`.

                unsafe {

                    self.move_item(a.base + i, range.start + i);

                }

            }

            for i in 0..b.len {

                // Safety: the indices in `b` must be initialized. We read

                // out the item and copy it to a new index; the old index

                // is no longer covered by a BumpVec, because we consume

                // `b`.

                unsafe {

                    self.move_item(b.base + i, range.start + a.len + i);

                }

            }

            a.free(self);

            b.free(self);

            BumpVec {

                base: range.start,

                len,

                cap,

                _phantom: PhantomData,

            }

        } else {

            self.vec.reserve(cap as usize);

            let base = self.vec.len() as u32;

            for i in 0..a.len {

                // Safety: the indices in `a` must be initialized. We read

                // out the item and copy it to a new index; the old index

                // is no longer covered by a BumpVec, because we consume

                // `a`.

                unsafe {

                    self.push_item(a.base + i);

                }

            }

            for i in 0..b.len {

                // Safety: the indices in `b` must be initialized. We read

                // out the item and copy it to a new index; the old index

                // is no longer covered by a BumpVec, because we consume

                // `b`.

                unsafe {

                    self.push_item(b.base + i);

                }

            }

            let len = self.vec.len() as u32 - base;

            for _ in len..cap {

                self.vec.push(MaybeUninit::uninit());

            }

            a.free(self);

            b.free(self);

            BumpVec {

                base,

                len,

                cap,

                _phantom: PhantomData,

            }

        }

    }

    /// Returns the size of the backing `Vec`.

    pub fn size(&self) -> usize {

        self.vec.len()

    }

}

fn round_up_power_of_two(x: u32) -> u32 {

    debug_assert!(x > 0);

    debug_assert!(x < 0x8000_0000);

    let log2 = 32 - (x - 1).leading_zeros();

    1 << log2

}

impl<T> BumpVec<T> {

    /// Returns a slice view of this `BumpVec`, given a borrow of the

    /// arena.

    pub fn as_slice<'a>(&'a self, arena: &'a BumpArena<T>) -> &'a [T] {

        let maybe_uninit_slice =

            &arena.vec[(self.base as usize)..((self.base + self.len) as usize)];

        // Safety: the index range we represent must be initialized.

        unsafe { std::mem::transmute(maybe_uninit_slice) }

    }

    /// Returns a mutable slice view of this `BumpVec`, given a

    /// mutable borrow of the arena.

    pub fn as_mut_slice<'a>(&'a mut self, arena: &'a mut BumpArena<T>) -> &'a mut [T] {

        let maybe_uninit_slice =

            &mut arena.vec[(self.base as usize)..((self.base + self.len) as usize)];

        // Safety: the index range we represent must be initialized.

        unsafe { std::mem::transmute(maybe_uninit_slice) }

    }

    /// Returns the length of this vector. Does not require access to

    /// the arena.

    pub fn len(&self) -> usize {

        self.len as usize

    }

    /// Returns the capacity of this vector. Does not require access

    /// to the arena.

    pub fn cap(&self) -> usize {

        self.cap as usize

    }

    /// Reserve `extra_len` capacity at the end of the vector,

    /// reallocating if necessary.

    pub fn reserve(&mut self, extra_len: usize, arena: &mut BumpArena<T>) {

        let extra_len = u32::try_from(extra_len).unwrap();

        if self.cap - self.len < extra_len {

            if self.base + self.cap == arena.vec.len() as u32 {

                for _ in 0..extra_len {

                    arena.vec.push(MaybeUninit::uninit());

                }

                self.cap += extra_len;

            } else {

                let new_cap = self.cap + extra_len;

                let new = arena.vec_with_capacity(new_cap as usize);

                unsafe {

                    for i in 0..self.len {

                        arena.move_item(self.base + i, new.base + i);

                    }

                }

                self.base = new.base;

                self.cap = new.cap;

            }

        }

    }

    /// Push an item, growing the capacity if needed.

    pub fn push(&mut self, t: T, arena: &mut BumpArena<T>) {

        if self.cap > self.len {

            unsafe {

                arena.write_into_index(self.base + self.len, t);

            }

            self.len += 1;

        } else if (self.base + self.cap) as usize == arena.vec.len() {

            arena.vec.push(MaybeUninit::new(t));

            self.cap += 1;

            self.len += 1;

        } else {

            let new_cap = round_up_power_of_two(self.cap + 1);

            let extra = new_cap - self.cap;

            self.reserve(extra as usize, arena);

            unsafe {

                arena.write_into_index(self.base + self.len, t);

            }

            self.len += 1;

        }

    }

    /// Clone, if `T` is cloneable.

    pub fn clone(&self, arena: &mut BumpArena<T>) -> BumpVec<T>

    where

        T: Clone,

    {

        let mut new = arena.vec_with_capacity(self.len as usize);

        for i in 0..self.len {

            let item = self.as_slice(arena)[i as usize].clone();

            new.push(item, arena);

        }

        new

    }

    /// Truncate the length to a smaller-or-equal length.

    pub fn truncate(&mut self, len: usize) {

        let len = len as u32;

        assert!(len <= self.len);

        self.len = len;

    }

    /// Consume the BumpVec and return its indices to a free pool in

    /// the arena.

    pub fn free(self, arena: &mut BumpArena<T>) {

        arena.freelist.push(self.base..(self.base + self.cap));

    }

    /// Freeze the capacity of this BumpVec, turning it into a slice,

    /// for a smaller struct (8 bytes rather than 12). Once this

    /// exists, it is copyable, because the slice will never be freed.

    pub fn freeze(self, arena: &mut BumpArena<T>) -> BumpSlice<T> {

        if self.cap > self.len {

            arena

                .freelist

                .push((self.base + self.len)..(self.base + self.cap));

        }

        BumpSlice {

            base: self.base,

            len: self.len,

            _phantom: PhantomData,

        }

    }

}

impl<T> BumpSlice<T> {

    /// Returns a slice view of the `BumpSlice`, given a borrow of the

    /// arena.

    pub fn as_slice<'a>(&'a self, arena: &'a BumpArena<T>) -> &'a [T] {

        let maybe_uninit_slice =

            &arena.vec[(self.base as usize)..((self.base + self.len) as usize)];

        // Safety: the index range we represent must be initialized.

        unsafe { std::mem::transmute(maybe_uninit_slice) }

    }

    /// Returns a mutable slice view of the `BumpSlice`, given a

    /// mutable borrow of the arena.

    pub fn as_mut_slice<'a>(&'a mut self, arena: &'a mut BumpArena<T>) -> &'a mut [T] {

        let maybe_uninit_slice =

            &mut arena.vec[(self.base as usize)..((self.base + self.len) as usize)];

        // Safety: the index range we represent must be initialized.

        unsafe { std::mem::transmute(maybe_uninit_slice) }

    }

    /// Returns the length of the `BumpSlice`.

    pub fn len(&self) -> usize {

        self.len as usize

    }

}

impl<T> std::default::Default for BumpVec<T> {

    fn default() -> Self {

        BumpVec {

            base: 0,

            len: 0,

            cap: 0,

            _phantom: PhantomData,

        }

    }

}

impl<T> std::default::Default for BumpSlice<T> {

    fn default() -> Self {

        BumpSlice {

            base: 0,

            len: 0,

            _phantom: PhantomData,

        }

    }

}

#[cfg(test)]

mod test {

    use super::*;

    #[test]

    fn test_round_up() {

        assert_eq!(1, round_up_power_of_two(1));

        assert_eq!(2, round_up_power_of_two(2));

        assert_eq!(4, round_up_power_of_two(3));

        assert_eq!(4, round_up_power_of_two(4));

        assert_eq!(32, round_up_power_of_two(24));

        assert_eq!(0x8000_0000, round_up_power_of_two(0x7fff_ffff));

    }

    #[test]

    fn test_basic() {

        let mut arena: BumpArena<u32> = BumpArena::new();

        let a = arena.single(1);

        let b = arena.single(2);

        let c = arena.single(3);

        let ab = arena.append(a, b);

        assert_eq!(ab.as_slice(&arena), &[1, 2]);

        assert_eq!(ab.cap(), 2);

        let abc = arena.append(ab, c);

        assert_eq!(abc.len(), 3);

        assert_eq!(abc.cap(), 4);

        assert_eq!(abc.as_slice(&arena), &[1, 2, 3]);

        assert_eq!(arena.size(), 9);

        let mut d = arena.single(4);

        // Should have reused the freelist.

        assert_eq!(arena.size(), 9);

        assert_eq!(d.len(), 1);

        assert_eq!(d.cap(), 1);

        assert_eq!(d.as_slice(&arena), &[4]);

        d.as_mut_slice(&mut arena)[0] = 5;

        assert_eq!(d.as_slice(&arena), &[5]);

        abc.free(&mut arena);

        let d2 = d.clone(&mut arena);

        let dd = arena.append(d, d2);

        // Should have reused the freelist.

        assert_eq!(arena.size(), 9);

        assert_eq!(dd.as_slice(&arena), &[5, 5]);

        let mut e = arena.from_iter([10, 11, 12].into_iter());

        e.push(13, &mut arena);

        assert_eq!(arena.size(), 13);

        e.reserve(4, &mut arena);

        assert_eq!(arena.size(), 17);

        let _f = arena.from_iter([1, 2, 3, 4, 5, 6, 7, 8].into_iter());

        assert_eq!(arena.size(), 25);

        e.reserve(8, &mut arena);

        assert_eq!(e.cap(), 16);

        assert_eq!(e.as_slice(&arena), &[10, 11, 12, 13]);

        // `e` must have been copied now that `f` is at the end of the

        // arena.

        assert_eq!(arena.size(), 41);

    }

}