use oxc_allocator::{Address, GetAddress};
use oxc_ast::AstKind;
use oxc_data_structures::stack::Stack;
use oxc_syntax::node::NodeId;
pub struct AncestryStack<'a> {
stack: Stack<AstKind<'a>>,
}
type AncestorKinds<'s, 'a> = std::iter::Copied<std::iter::Rev<std::slice::Iter<'s, AstKind<'a>>>>;
impl Default for AncestryStack<'_> {
fn default() -> Self {
Self { stack: Stack::with_capacity(64) }
}
}
impl<'a> AncestryStack<'a> {
#[inline]
pub(crate) fn push(&mut self, kind: AstKind<'a>) {
self.stack.push(kind);
}
#[inline]
pub(crate) fn pop(&mut self) {
self.stack.pop();
}
#[inline]
pub(crate) fn current_node_id(&self) -> NodeId {
self.stack.last().map_or(NodeId::ROOT, AstKind::node_id)
}
#[inline]
pub(crate) fn parent_node_id(&self) -> NodeId {
let stack = self.stack.as_slice();
let parent_index = stack.len().saturating_sub(2);
stack.get(parent_index).map_or(NodeId::ROOT, AstKind::node_id)
}
#[inline]
pub fn current_kind(&self) -> AstKind<'a> {
*self.stack.last().expect("ancestry stack is empty")
}
#[inline]
pub fn current_address(&self) -> Address {
self.current_kind().address()
}
#[inline]
pub fn parent_kind(&self) -> AstKind<'a> {
let stack = self.stack.as_slice();
let parent_index = stack.len().saturating_sub(2);
stack[parent_index]
}
#[inline]
pub fn ancestor_kinds(&self) -> AncestorKinds<'_, 'a> {
let stack = self.stack.as_slice();
let upto = stack.len().saturating_sub(1);
stack[..upto].iter().rev().copied()
}
#[inline]
pub fn ancestor_kinds_from(&self, from: Option<NodeId>) -> AncestorKinds<'_, 'a> {
let stack = self.stack.as_slice();
let end = match from {
None => stack.len(),
Some(node_id) => {
stack.iter().rposition(|kind| kind.node_id() == node_id).map_or(0, |i| i + 1)
}
};
stack[..end].iter().rev().copied()
}
#[inline]
pub fn find_kind_by_node_id(&self, node_id: NodeId) -> AstKind<'a> {
self.stack
.as_slice()
.iter()
.rev()
.find(|kind| kind.node_id() == node_id)
.copied()
.expect("node id not found on the ancestry stack")
}
}
#[derive(Clone, Copy)]
pub enum Ancestry<'a, 'n> {
Nodes { nodes: &'n super::AstNodes<'a>, current_node_id: NodeId },
Stack(&'n AncestryStack<'a>),
}
impl<'a, 'n> Ancestry<'a, 'n> {
#[inline]
pub fn current_kind(self) -> AstKind<'a> {
match self {
Ancestry::Nodes { nodes, current_node_id } => nodes.kind(current_node_id),
Ancestry::Stack(stack) => stack.current_kind(),
}
}
#[inline]
pub fn current_address(self) -> Address {
self.current_kind().address()
}
#[inline]
pub fn parent_kind(self) -> AstKind<'a> {
match self {
Ancestry::Nodes { nodes, current_node_id } => nodes.parent_kind(current_node_id),
Ancestry::Stack(stack) => stack.parent_kind(),
}
}
#[inline]
pub fn parent_node_id(self) -> NodeId {
match self {
Ancestry::Nodes { nodes, current_node_id } => nodes.parent_id(current_node_id),
Ancestry::Stack(stack) => stack.parent_node_id(),
}
}
#[inline]
pub fn ancestor_kinds(self) -> impl Iterator<Item = AstKind<'a>> + Clone + 'n {
match self {
Ancestry::Nodes { nodes, current_node_id } => {
itertools::Either::Left(nodes.ancestor_kinds(current_node_id))
}
Ancestry::Stack(stack) => itertools::Either::Right(stack.ancestor_kinds()),
}
}
#[inline]
pub fn ancestor_kinds_from(
self,
from: Option<NodeId>,
) -> impl Iterator<Item = AstKind<'a>> + Clone + 'n {
match self {
Ancestry::Nodes { nodes, current_node_id } => {
let start = from.unwrap_or(current_node_id);
itertools::Either::Left(
std::iter::once(nodes.kind(start)).chain(nodes.ancestor_kinds(start)),
)
}
Ancestry::Stack(stack) => itertools::Either::Right(stack.ancestor_kinds_from(from)),
}
}
#[inline]
pub fn find_kind_by_node_id(self, node_id: NodeId) -> AstKind<'a> {
match self {
Ancestry::Nodes { nodes, .. } => nodes.kind(node_id),
Ancestry::Stack(stack) => stack.find_kind_by_node_id(node_id),
}
}
}