use alloc::vec;
use alloc::vec::Vec;
use core::slice;
use span_lang::Span;
use token_lang::Token;
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Element<K> {
Node(Node<K>),
Token(Token<K>),
}
impl<K> Element<K> {
#[inline]
#[must_use]
pub fn span(&self) -> Span {
match self {
Element::Node(node) => node.span(),
Element::Token(token) => token.span(),
}
}
#[inline]
#[must_use]
pub fn kind(&self) -> &K {
match self {
Element::Node(node) => node.kind(),
Element::Token(token) => token.kind(),
}
}
#[inline]
#[must_use]
pub fn as_node(&self) -> Option<&Node<K>> {
match self {
Element::Node(node) => Some(node),
Element::Token(_) => None,
}
}
#[inline]
#[must_use]
pub fn as_token(&self) -> Option<&Token<K>> {
match self {
Element::Token(token) => Some(token),
Element::Node(_) => None,
}
}
#[inline]
#[must_use]
pub fn is_node(&self) -> bool {
matches!(self, Element::Node(_))
}
#[inline]
#[must_use]
pub fn is_token(&self) -> bool {
matches!(self, Element::Token(_))
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Node<K> {
kind: K,
span: Span,
children: Vec<Element<K>>,
}
impl<K> Node<K> {
#[must_use]
pub fn new(kind: K, children: Vec<Element<K>>) -> Self {
let span = cover(&children, Span::empty(0));
Self {
kind,
span,
children,
}
}
#[must_use]
pub(crate) fn with_span(kind: K, children: Vec<Element<K>>, empty: Span) -> Self {
let span = cover(&children, empty);
Self {
kind,
span,
children,
}
}
#[inline]
#[must_use]
pub fn kind(&self) -> &K {
&self.kind
}
#[inline]
#[must_use]
pub fn span(&self) -> Span {
self.span
}
#[inline]
#[must_use]
pub fn is_empty(&self) -> bool {
self.children.is_empty()
}
#[inline]
#[must_use]
pub fn len(&self) -> usize {
self.children.len()
}
#[inline]
pub fn children(&self) -> impl Iterator<Item = &Element<K>> {
self.children.iter()
}
#[inline]
pub fn child_nodes(&self) -> impl Iterator<Item = &Node<K>> {
self.children.iter().filter_map(Element::as_node)
}
#[inline]
pub fn child_tokens(&self) -> impl Iterator<Item = &Token<K>> {
self.children.iter().filter_map(Element::as_token)
}
#[inline]
pub fn descendants(&self) -> impl Iterator<Item = &Node<K>> {
Descendants {
root: Some(self),
stack: Vec::new(),
}
}
#[inline]
pub fn tokens(&self) -> impl Iterator<Item = &Token<K>> {
Tokens {
stack: vec![self.children.iter()],
}
}
#[inline]
#[must_use]
pub fn text<'s>(&self, source: &'s str) -> Option<&'s str> {
let start = self.span.start().to_usize();
let end = self.span.end().to_usize();
source.get(start..end)
}
}
impl<K> Drop for Node<K> {
fn drop(&mut self) {
if self.children.iter().all(Element::is_token) {
return;
}
let mut stack: Vec<Node<K>> = Vec::new();
drain_nodes(&mut self.children, &mut stack);
while let Some(mut node) = stack.pop() {
drain_nodes(&mut node.children, &mut stack);
}
}
}
fn drain_nodes<K>(children: &mut Vec<Element<K>>, stack: &mut Vec<Node<K>>) {
for element in children.drain(..) {
if let Element::Node(node) = element {
stack.push(node);
}
}
}
fn cover<K>(children: &[Element<K>], empty: Span) -> Span {
let mut iter = children.iter();
match iter.next() {
None => empty,
Some(first) => iter.fold(first.span(), |acc, child| acc.merge(child.span())),
}
}
struct Descendants<'a, K> {
root: Option<&'a Node<K>>,
stack: Vec<slice::Iter<'a, Element<K>>>,
}
impl<'a, K> Iterator for Descendants<'a, K> {
type Item = &'a Node<K>;
fn next(&mut self) -> Option<Self::Item> {
if let Some(root) = self.root.take() {
self.stack.push(root.children.iter());
return Some(root);
}
loop {
let top = self.stack.last_mut()?;
match top.next() {
None => {
let _ = self.stack.pop();
}
Some(Element::Node(node)) => {
self.stack.push(node.children.iter());
return Some(node);
}
Some(Element::Token(_)) => {}
}
}
}
}
struct Tokens<'a, K> {
stack: Vec<slice::Iter<'a, Element<K>>>,
}
impl<'a, K> Iterator for Tokens<'a, K> {
type Item = &'a Token<K>;
fn next(&mut self) -> Option<Self::Item> {
loop {
let top = self.stack.last_mut()?;
match top.next() {
None => {
let _ = self.stack.pop();
}
Some(Element::Node(node)) => {
self.stack.push(node.children.iter());
}
Some(Element::Token(token)) => return Some(token),
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use alloc::vec;
fn tok(kind: &'static str, lo: u32, hi: u32) -> Element<&'static str> {
Element::Token(Token::new(kind, Span::new(lo, hi)))
}
#[test]
fn test_new_covers_children_span() {
let n = Node::new("n", vec![tok("a", 2, 4), tok("b", 4, 9)]);
assert_eq!(n.span(), Span::new(2, 9));
}
#[test]
fn test_new_empty_node_has_empty_span() {
let n: Node<&str> = Node::new("n", vec![]);
assert_eq!(n.span(), Span::empty(0));
assert!(n.is_empty());
assert_eq!(n.len(), 0);
}
#[test]
fn test_children_iterators_split_nodes_and_tokens() {
let inner = Element::Node(Node::new("inner", vec![tok("x", 1, 2)]));
let n = Node::new("n", vec![tok("a", 0, 1), inner]);
assert_eq!(n.children().count(), 2);
assert_eq!(n.child_nodes().count(), 1);
assert_eq!(
n.child_tokens().map(|t| *t.kind()).collect::<Vec<_>>(),
["a"]
);
}
#[test]
fn test_descendants_preorder() {
let tree = Node::new(
"root",
vec![
Element::Node(Node::new("l", vec![tok("a", 0, 1)])),
Element::Node(Node::new("r", vec![tok("b", 1, 2)])),
],
);
let kinds: Vec<_> = tree.descendants().map(Node::kind).copied().collect();
assert_eq!(kinds, ["root", "l", "r"]);
}
#[test]
fn test_tokens_source_order_includes_all_leaves() {
let tree = Node::new(
"root",
vec![
tok("a", 0, 1),
Element::Node(Node::new("inner", vec![tok("b", 1, 2), tok("c", 2, 3)])),
tok("d", 3, 4),
],
);
let leaves: Vec<_> = tree.tokens().map(|t| *t.kind()).collect();
assert_eq!(leaves, ["a", "b", "c", "d"]);
}
#[test]
fn test_text_slices_source_and_rejects_out_of_bounds() {
let n = Node::new("n", vec![tok("a", 0, 2), tok("b", 2, 5)]);
assert_eq!(n.text("hello"), Some("hello"));
assert_eq!(n.text("hi"), None);
}
#[test]
fn test_element_accessors() {
let e = tok("a", 0, 1);
assert!(e.is_token());
assert!(!e.is_node());
assert_eq!(e.kind(), &"a");
assert_eq!(e.span(), Span::new(0, 1));
assert!(e.as_token().is_some());
assert!(e.as_node().is_none());
}
#[test]
fn test_deep_tree_drops_without_stack_overflow() {
let mut node = Node::new("leaf", vec![tok("t", 0, 1)]);
for _ in 0..200_000 {
node = Node::new("link", vec![Element::Node(node)]);
}
drop(node);
}
#[test]
fn test_deep_tree_tokens_and_descendants_are_iterative() {
let mut node = Node::new("leaf", vec![tok("t", 0, 1)]);
for _ in 0..100_000 {
node = Node::new("link", vec![Element::Node(node)]);
}
assert_eq!(node.tokens().count(), 1);
assert_eq!(node.descendants().count(), 100_001);
}
}