use std::ops::Index;

use ruff_index::{newtype_index, IndexVec};
use ruff_python_ast::Expr;

/// Id uniquely identifying an expression in a program.
///
/// Using a `u32` is sufficient because Ruff only supports parsing documents with a size of max
/// `u32::max` and it is impossible to have more nodes than characters in the file. We use a
/// `NonZeroU32` to take advantage of memory layout optimizations.
#[newtype_index]
#[derive(Ord, PartialOrd)]
pub struct ExpressionId;

/// An [`Expr`] AST node in a program, along with a pointer to its parent expression (if any).
#[derive(Debug)]
struct ExpressionWithParent<'a> {
    /// A pointer to the AST node.
    node: &'a Expr,
    /// The ID of the parent of this node, if any.
    parent: Option<ExpressionId>,
}

/// The nodes of a program indexed by [`ExpressionId`]
#[derive(Debug, Default)]
pub struct Expressions<'a>(IndexVec<ExpressionId, ExpressionWithParent<'a>>);

impl<'a> Expressions<'a> {
    /// Inserts a new expression into the node tree and returns its unique id.
    pub(crate) fn insert(&mut self, node: &'a Expr, parent: Option<ExpressionId>) -> ExpressionId {
        self.0.push(ExpressionWithParent { node, parent })
    }

    /// Return the [`ExpressionId`] of the parent node.
    #[inline]
    pub fn parent_id(&self, node_id: ExpressionId) -> Option<ExpressionId> {
        self.0[node_id].parent
    }

    /// Returns an iterator over all [`ExpressionId`] ancestors, starting from the given [`ExpressionId`].
    pub(crate) fn ancestor_ids(
        &self,
        node_id: ExpressionId,
    ) -> impl Iterator<Item = ExpressionId> + '_ {
        std::iter::successors(Some(node_id), |&node_id| self.0[node_id].parent)
    }
}

impl<'a> Index<ExpressionId> for Expressions<'a> {
    type Output = &'a Expr;

    #[inline]
    fn index(&self, index: ExpressionId) -> &Self::Output {
        &self.0[index].node
    }
}