//! Arena submodule defining types allocated from a dedicated bump arena,
//! allowing many cheap allocations with fast bulk reclamation.
//!
//! ## Destruction and resource management
//!
//! Dropping the [`Arena`] reclaims the arena's memory in bulk. Destructors are
//! not run for arena allocations that are still live at that point. Therefore,
//! avoid storing types whose cleanup must reliably happen at arena release,
//! especially types that own memory allocations or external resources outside
//! the arena.
//!
//! [`ArenaNode`], [`ArenaVec`], and [`ArenaString`] are provided so commonly
//! used owned data can keep their storage inside the arena rather than in
//! separate global-heap allocations.

use core::borrow::Borrow;
use core::fmt::{Debug, Display, Formatter, Result};
use core::ops::{Deref, DerefMut};

use bumpalo::{Bump, boxed, collections};

use crate::lexer::{TokenPosition, TokenSpan};

/// Object that manages a separate memory space, which can be deallocated all
/// at once after use.
///
/// All allocations borrow the arena immutably.
///
/// Dropping the [`Arena`] does not run [`Drop`] for values allocated within it
/// (including values contained in [`ArenaNode`], [`ArenaVec`] and
/// [`ArenaString`]).
#[derive(Debug)]
pub struct Arena {
    bump: Bump,
}

impl Arena {
    /// Creates a new, empty arena.
    #[must_use]
    pub fn new() -> Self {
        Self { bump: Bump::new() }
    }

    /// Constructs an empty [`ArenaVec`] allocated in this arena.
    ///
    /// The returned vector borrows this arena and cannot outlive it.
    #[must_use]
    pub fn vec<T>(&self) -> ArenaVec<'_, T> {
        ArenaVec(collections::Vec::new_in(&self.bump))
    }

    /// Allocates a copy of `string` in this arena and returns
    /// an [`ArenaString`].
    #[must_use]
    pub fn string(&self, string: &str) -> ArenaString<'_> {
        ArenaString(collections::String::from_str_in(string, &self.bump))
    }

    /// Allocates `value` in this arena and attaches `span`.
    ///
    /// The returned node borrows the arena and cannot outlive it.
    /// If it is still live when the arena is dropped, its destructor is not run.
    #[must_use]
    pub fn alloc_node<T>(&self, value: T, span: TokenSpan) -> ArenaNode<'_, T> {
        ArenaNode {
            value: boxed::Box::new_in(value, &self.bump),
            span,
        }
    }

    /// Allocates `value` in this arena and attaches the span from `start` to
    /// `end`.
    ///
    /// The returned node borrows the arena and cannot outlive it.
    /// If it is still live when the arena is dropped, its destructor is not run.
    #[must_use]
    pub fn alloc_node_between<T>(
        &self,
        value: T,
        start: TokenPosition,
        end: TokenPosition,
    ) -> ArenaNode<'_, T> {
        self.alloc_node(value, TokenSpan::range(start, end))
    }

    /// Allocates `value` in this arena and attaches a span covering `at`.
    ///
    /// The returned node borrows the arena and cannot outlive it.
    /// If it is still live when the arena is dropped, its destructor is not run.
    #[must_use]
    pub fn alloc_node_at<T>(&self, value: T, at: TokenPosition) -> ArenaNode<'_, T> {
        self.alloc_node(value, TokenSpan::new(at))
    }
}

impl Default for Arena {
    fn default() -> Self {
        Self::new()
    }
}

/// An arena-allocated box with an attached source span.
///
/// Dropping the node normally runs `Drop` for the inner value.
/// Dropping the arena does not itself perform a separate destructor pass.
#[derive(Hash, PartialEq, Eq)]
pub struct ArenaNode<'arena, T> {
    value: boxed::Box<'arena, T>,
    span: TokenSpan,
}

impl<'arena, T> ArenaNode<'arena, T> {
    /// Creates a new [`ArenaNode`] by allocating `value` in `arena`.
    #[must_use]
    pub fn new_in(value: T, span: TokenSpan, arena: &'arena Arena) -> Self {
        Self {
            value: boxed::Box::new_in(value, &arena.bump),
            span,
        }
    }

    /// Returns a mutable reference to the token span covered by this node.
    #[must_use]
    pub const fn span_mut(&mut self) -> &mut TokenSpan {
        &mut self.span
    }

    /// Returns the token span covered by this node.
    #[must_use]
    pub const fn span(&self) -> &TokenSpan {
        &self.span
    }
}

impl<T> Deref for ArenaNode<'_, T> {
    type Target = T;

    fn deref(&self) -> &T {
        &self.value
    }
}

impl<T> DerefMut for ArenaNode<'_, T> {
    fn deref_mut(&mut self) -> &mut T {
        &mut self.value
    }
}

impl<T: Debug> Debug for ArenaNode<'_, T> {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
        f.debug_struct("ArenaNode")
            .field("inner", &**self)
            .field("span", self.span())
            .finish()
    }
}

/// Version of [`Vec`] whose backing storage lives in the arena.
///
/// Elements are dropped when the `ArenaVec` itself is dropped normally.
/// Capacity growth may leave old buffers in the arena until the whole arena
/// is reclaimed.
#[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct ArenaVec<'arena, T>(collections::Vec<'arena, T>);

impl<'arena, T> ArenaVec<'arena, T> {
    /// Creates an empty `ArenaVec` allocated in `arena`.
    #[must_use]
    pub fn new_in(arena: &'arena Arena) -> Self {
        Self(collections::Vec::new_in(&arena.bump))
    }

    /// Appends an element to the end of the vector.
    ///
    /// Growth is backed by the arena; increasing capacity allocates new space
    /// in the arena and never frees previous blocks.
    pub fn push(&mut self, value: T) {
        self.0.push(value);
    }

    /// Reserves capacity for at least `additional` more elements.
    ///
    /// The collection may reserve more space to avoid frequent reallocations.
    /// If growth requires a new allocation, the previous buffer remains in the
    /// arena until the arena is reclaimed.
    pub fn reserve(&mut self, additional: usize) {
        self.0.reserve(additional);
    }

    /// Extends the vector with the contents of `items`.
    ///
    /// Growth may allocate a new buffer in the arena and leave the previous
    /// buffer in place until the arena is reclaimed.
    pub fn extend<I: IntoIterator<Item = T>>(&mut self, items: I) {
        self.0.extend(items);
    }
}

impl<T> Deref for ArenaVec<'_, T> {
    type Target = [T];

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

impl<T> DerefMut for ArenaVec<'_, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl<'iter, T> IntoIterator for &'iter ArenaVec<'_, T> {
    type Item = &'iter T;
    type IntoIter = core::slice::Iter<'iter, T>;
    fn into_iter(self) -> Self::IntoIter {
        self.0.iter()
    }
}

impl<'iter, T> IntoIterator for &'iter mut ArenaVec<'_, T> {
    type Item = &'iter mut T;
    type IntoIter = core::slice::IterMut<'iter, T>;
    fn into_iter(self) -> Self::IntoIter {
        self.0.iter_mut()
    }
}

/// Version of [`String`] whose backing storage lives in the arena.
#[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct ArenaString<'arena>(collections::String<'arena>);

impl<'arena> ArenaString<'arena> {
    /// Allocates a copy of `string` in `arena` and returns an [`ArenaString`].
    #[must_use]
    pub fn from_str_in(text: &str, arena: &'arena Arena) -> Self {
        Self(collections::String::from_str_in(text, &arena.bump))
    }
}

impl Deref for ArenaString<'_> {
    type Target = str;

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

impl AsRef<str> for ArenaString<'_> {
    fn as_ref(&self) -> &str {
        &self.0
    }
}

impl Borrow<str> for ArenaString<'_> {
    fn borrow(&self) -> &str {
        &self.0
    }
}

impl Display for ArenaString<'_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
        Display::fmt(&self.0, f)
    }
}
