//! Arena submodule defining types that exist in their own memory space and
//! allow multiple cheap allocations (both performance- and fragmentation-wise).
//!
//! ## Memory safety
//!
//! Dropping the [`Arena`] frees all its memory at once and does not run
//! [`Drop`] for values allocated within it. Avoid storing types that implement
//! [`Drop`] or own external resources inside [`ArenaNode`], [`ArenaVec`], or
//! [`ArenaString`]. If you must, arrange an explicit "drain/drop" pass before
//! the arena is dropped.

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

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

use crate::ast::AstSpan;
use crate::lexer::TokenLocation;

/// 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`]).
///
/// This arena is not thread-safe (`!Send`, `!Sync`). Values borrow the arena
/// and therefore cannot be sent across threads independently.
#[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 with the given `span`,
    /// returning an [`ArenaNode`].
    ///
    /// The node's storage borrows this arena and cannot outlive it.
    ///
    /// Note: `T`'s [`Drop`] is not run when the arena is dropped.
    #[must_use]
    pub fn alloc<T>(&self, value: T, span: AstSpan) -> ArenaNode<'_, T> {
        ArenaNode {
            inner: boxed::Box::new_in(value, &self.bump),
            span,
        }
    }

    pub fn alloc_between<T>(
        &self,
        value: T,
        from: TokenLocation,
        to: TokenLocation,
    ) -> ArenaNode<'_, T> {
        self.alloc(value, AstSpan { from, to })
    }

    pub fn alloc_at<T>(&self, value: T, at: TokenLocation) -> ArenaNode<'_, T> {
        self.alloc(value, AstSpan { from: at, to: at })
    }
}

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

/// An arena-allocated box with an attached source span.
///
/// Equality and hashing take into account both the contained `T` and the `span`
/// (when `T: Eq + Hash`).
///
/// Note: `T`'s [`Drop`] is not run when the arena is dropped.
#[derive(Hash, PartialEq, Eq)]
pub struct ArenaNode<'arena, T> {
    /// Value allocated in the arena; this node owns it.
    inner: boxed::Box<'arena, T>,
    /// Token range covered by the value.
    span: AstSpan,
}

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

    /// Creates a new [`ArenaNode`] for an AST node that spans a single token.
    pub fn from_token_location(
        value: T,
        token_location: crate::lexer::TokenLocation,
        arena: &'arena Arena,
    ) -> Self {
        Self {
            inner: boxed::Box::new_in(value, &arena.bump),
            span: AstSpan {
                from: token_location,
                to: token_location,
            },
        }
    }

    pub fn span_mut(&mut self) -> &mut AstSpan {
        &mut self.span
    }

    pub fn extend_to(&mut self, to: TokenLocation) {
        self.span.to = to;
    }

    pub fn extend_from(&mut self, from: TokenLocation) {
        self.span.from = from;
    }

    /// Returns the token span covered by this node.
    pub fn span(&self) -> &AstSpan {
        &self.span
    }
}

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

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

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

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

/// Version of [`Vec`] that can be safely used inside a memory arena.
///
/// Elements do not have their destructors run when the arena is dropped.
///
/// This type dereferences to `[T]` and supports iteration by reference
/// (`&ArenaVec` and `&mut ArenaVec` implement [`IntoIterator`]).
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
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)
    }

    pub fn reserve(&mut self, additional: usize) {
        self.0.reserve(additional)
    }
    pub fn extend<I: IntoIterator<Item = T>>(&mut self, it: I) {
        self.0.extend(it)
    }
}

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

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

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

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

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

/// Version of [`String`] that can be safely used inside a memory arena.
///
/// This type dereferences to [`str`] and implements [`AsRef<str>`] and
/// [`core::borrow::Borrow<str>`] for ergonomic use with APIs expecting string
/// slices.
///
/// The string borrows the arena and cannot outlive it. Dropping the arena
/// frees its memory without running `Drop` for the string contents.
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
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(string: &str, arena: &'arena Arena) -> Self {
        Self(collections::String::from_str_in(string, &arena.bump))
    }
}

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

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

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

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

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