//! `KeyedArena` is a bump allocator which allows for random access to its allocations.
//!
//! # Rationale
//! [`bumpalo`](https://docs.rs/bumpalo) is a popular Rust crate which implements a bump allocator.
//! The function signature for allocation a slice in `bumpalo` looks something like:
//!
//! ```text
//! fn alloc_slice_default<'a, T: Default>(&'a self, n: usize) -> &'a [T]
//! ```
//!
//! But, now if we want to store the arena alongside references to values allocated into it,
//! we have a problem
//!
//! ```compile_fail
//! struct OurFunStruct {
//!     arena: bumpalo::Bump,
//!     things: &'????? [usize], // what lifetime do we want?
//! }
//! ```
//!
//! The "right" solution to this problem would be to have a self-referential type: a type which
//! would expose the lifetime of the arena to be used for its own fields. Unfortunately, this
//! isn't possible in Rust (except through the use of `async`/`await`).
//!
//! Since we can't do the "right" thing, we'll work around it! And that's the point of `KeyedArena`
//!
//! # `KeyedArena`
//! To avoid the need for self-referential types, `KeyedArena` allocation outputs a key, which can
//! be used to retrieve the allocation later.
//! ```
//! # use keyed_arena::{AllocationKey, KeyedArena};
//! struct OurFunStruct {
//!     arena: KeyedArena,
//!     things: AllocationKey<usize>,
//! }
//! impl OurFunStruct {
//!     pub fn new(n: usize) -> Self {
//!         let arena = KeyedArena::default();
//!         let things = arena.alloc_slice_fill_with(n, |i| i);
//!         // Do stuff with things
//!         let key = things.key();
//!         std::mem::drop(things);
//!         Self { arena, things: key }
//!     }
//!     pub fn step2(&mut self) -> usize {
//!         let things = self.arena.borrow_mut(self.things);
//!         things.iter().copied().sum()
//!     }
//! }
//! ```
use crossbeam_queue::ArrayQueue;
use std::{
    any::TypeId,
    cell::RefCell,
    marker::PhantomData,
    ops::{Deref, DerefMut},
    sync::{Arc, Weak},
};

/// A key that can be used to retrieve a previous allocation of a slice of type `T`.
#[derive(Clone, Copy)]
pub struct AllocationKey<T: 'static + Copy + Sized + Send> {
    idx: usize,
    len: usize,
    phantom: PhantomData<&'static T>,
}

/// A mutable allocation of a slice of type `T`, borrowed from a `KeyedArena`.
///
/// This type is _not_ `Sync`.
/// ```compile_fail
/// fn foo<'a>(x: keyed_arena::BorrowedAllocation<'a, i32>) -> impl std::marker::Sync + 'a { x }
/// ```
pub struct BorrowedAllocation<'a, T: 'static + Copy + Sized + Send> {
    // T must be 'static since we capture its TypeId. It must be Copy since we won't call its
    // destructor.
    // We _could_ support more types in the future.
    arena: &'a KeyedArena,
    start: *mut T,
    key: AllocationKey<T>,
}
impl<'a, T: 'static + Copy + Sized + Send> BorrowedAllocation<'a, T> {
    /// Get a key that can be used to retrieve this allocation in the future.
    pub fn key(&self) -> AllocationKey<T> {
        self.key
    }
}
impl<'a, T: 'static + Copy + Sized + Send> Deref for BorrowedAllocation<'a, T> {
    type Target = [T];

    fn deref(&self) -> &Self::Target {
        unsafe { std::slice::from_raw_parts(self.start as *const T, self.key.len) }
    }
}
impl<'a, T: 'static + Copy + Sized + Send> DerefMut for BorrowedAllocation<'a, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { std::slice::from_raw_parts_mut(self.start, self.key.len) }
    }
}
impl<'a, T: 'static + Copy + Sized + Send> std::ops::Drop for BorrowedAllocation<'a, T> {
    fn drop(&mut self) {
        let borrowed = &mut self.arena.allocations.borrow_mut()[self.key.idx].borrowed;
        assert!(*borrowed);
        *borrowed = false;
    }
}

#[derive(Debug, Clone, Copy)]
pub struct KeyedArenaStats {
    pub num_allocations: usize,
    pub allocated_bytes: usize,
}

struct RawAllocation {
    type_id: TypeId,
    start: *const u8,
    len: usize,
    borrowed: bool,
}
unsafe impl Send for RawAllocation {}

/// A bump allocating arena where allocations can be indexed with keys.
pub struct KeyedArena {
    arena: bumpalo::Bump,
    allocations: RefCell<Vec<RawAllocation>>,
}
impl Default for KeyedArena {
    fn default() -> Self {
        Self::with_capacity(0, 0)
    }
}
impl KeyedArena {
    pub fn with_capacity(arena_size: usize, num_allocations: usize) -> Self {
        KeyedArena {
            arena: bumpalo::Bump::with_capacity(arena_size),
            allocations: RefCell::new(Vec::with_capacity(num_allocations)),
        }
    }

    pub fn stats(&self) -> KeyedArenaStats {
        KeyedArenaStats {
            num_allocations: self.allocations.borrow().len(),
            allocated_bytes: self.arena.allocated_bytes(),
        }
    }

    /// Allocate a slice, `&[T]` of length `len`, generated by invoking `f(0), f(1), ...`
    pub fn alloc_slice_fill_with<T: 'static + Copy + Sized + Send, F: FnMut(usize) -> T>(
        &self,
        len: usize,
        f: F,
    ) -> BorrowedAllocation<T> {
        let slice = self.arena.alloc_slice_fill_with(len, f);
        assert_eq!(slice.len(), len);
        // borrow this _after_ the slice is built, in case it also wants to allocate.
        let mut allocations = self.allocations.borrow_mut();
        let idx = allocations.len();
        allocations.push(RawAllocation {
            type_id: TypeId::of::<T>(),
            start: slice.as_mut_ptr() as *mut u8,
            len,
            borrowed: true,
        });
        BorrowedAllocation {
            arena: self,
            start: slice.as_mut_ptr(),
            key: AllocationKey {
                idx,
                len,
                phantom: PhantomData,
            },
        }
    }

    /// Retrieve a previous allocation from this arena.
    /// # Panic
    /// Calling this function will panic if:
    /// 1. the `key` wasn't created by this allocator, or this allocator has been `reset()` since
    ///    the key was created. Note that a panic is not guaranteed in this situation.
    /// 2. `key` already has a `BorrowedAllocation` that hasn't been dropped.
    pub fn borrow_mut<T: 'static + Copy + Sized + Send>(
        &self,
        key: AllocationKey<T>,
    ) -> BorrowedAllocation<T> {
        let mut allocations = self.allocations.borrow_mut();
        let allocation = &mut allocations[key.idx];
        assert_eq!(allocation.type_id, TypeId::of::<T>());
        assert_eq!(allocation.len, key.len);
        assert!(!allocation.borrowed);
        allocation.borrowed = true;
        BorrowedAllocation {
            arena: self,
            start: allocation.start as *mut T,
            key,
        }
    }

    pub fn reset(&mut self) {
        self.arena.reset();
        self.allocations.get_mut().clear();
    }
}

/// A thread-safe object pool to cache `KeyedArena` values.
pub struct KeyedArenaPool {
    arena_size: usize,
    num_allocations: usize,
    queue: Arc<ArrayQueue<KeyedArena>>,
}
impl KeyedArenaPool {
    /// Construct a new `KeyedArenaPool`. Fresh arenas will be allocated with the given `arena_size`
    /// and `num_allocations`. The pool will hold at most `capacity` arenas.
    pub fn new(arena_size: usize, num_allocations: usize, capacity: usize) -> Self {
        KeyedArenaPool {
            arena_size,
            num_allocations,
            queue: Arc::new(ArrayQueue::new(capacity)),
        }
    }

    /// Get a fresh `KeyedArena` from the pool.
    pub fn get(&self) -> KeyedArenaFromPool {
        KeyedArenaFromPool {
            queue: Arc::downgrade(&self.queue),
            arena: Some(self.queue.pop().unwrap_or_else(|| {
                KeyedArena::with_capacity(self.arena_size, self.num_allocations)
            })),
        }
    }
}

/// A wrapper for a `KeyedArena` which will return it to a `KeyedArenaPool` on drop.
pub struct KeyedArenaFromPool {
    // This is an Option<> to allow us to take() it in the drop handler.
    arena: Option<KeyedArena>,
    queue: Weak<ArrayQueue<KeyedArena>>,
}
impl Deref for KeyedArenaFromPool {
    type Target = KeyedArena;

    #[inline]
    fn deref(&self) -> &Self::Target {
        self.arena.as_ref().unwrap()
    }
}
impl DerefMut for KeyedArenaFromPool {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.arena.as_mut().unwrap()
    }
}
impl std::ops::Drop for KeyedArenaFromPool {
    fn drop(&mut self) {
        let mut arena = self.arena.take().unwrap();
        arena.reset();
        // If adding the arena fails, then just free it.
        if let Some(queue) = self.queue.upgrade() {
            let _ = queue.push(arena);
        }
    }
}
