use rowan::{NodeOrToken, TextRange, TextSize};
use smol_str::{SmolStr, SmolStrBuilder};
use super::{AstNode, AstToken, child, child_token, children};
use crate::ast::tokens::ControlWord;
use crate::syntax::{SyntaxKind, SyntaxNode};
macro_rules! ast_node {
($(#[$meta:meta])* $name:ident, $kind:ident) => {
$(#[$meta])*
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct $name {
syntax: SyntaxNode,
}
impl AstNode for $name {
fn can_cast(kind: SyntaxKind) -> bool {
kind == SyntaxKind::$kind
}
fn cast(syntax: SyntaxNode) -> Option<Self> {
Self::can_cast(syntax.kind()).then_some(Self { syntax })
}
fn syntax(&self) -> &SyntaxNode {
&self.syntax
}
}
};
}
ast_node!(
Command, COMMAND
);
ast_node!(
Group, GROUP
);
ast_node!(
Optional, OPTIONAL
);
ast_node!(
NameGroup, NAME_GROUP
);
ast_node!(
Begin, BEGIN
);
ast_node!(
End, END
);
ast_node!(
Environment, ENVIRONMENT
);
ast_node!(
Conditional, CONDITIONAL
);
ast_node!(
ConditionalBranch, CONDITIONAL_BRANCH
);
impl Command {
pub fn control_word(&self) -> Option<ControlWord> {
self.syntax
.children_with_tokens()
.filter_map(NodeOrToken::into_token)
.find_map(ControlWord::cast)
}
pub fn name(&self) -> Option<SmolStr> {
self.control_word().map(|cw| SmolStr::new(cw.name()))
}
pub fn control_word_range(&self) -> Option<TextRange> {
self.control_word().map(|cw| cw.range())
}
pub fn nth_group(&self, n: usize) -> Option<Group> {
self.groups().nth(n)
}
pub fn groups(&self) -> impl Iterator<Item = Group> {
children::<Group>(&self.syntax)
}
pub fn optionals(&self) -> impl Iterator<Item = Optional> {
children::<Optional>(&self.syntax)
}
pub fn nth_group_text(&self, n: usize) -> Option<SmolStr> {
self.nth_group(n)?.inner_text()
}
pub fn nth_group_inner(&self, n: usize) -> Option<(TextRange, SmolStr)> {
self.nth_group(n)?.inner()
}
pub fn first_group_range(&self) -> TextRange {
match self.nth_group(0) {
Some(group) => TextRange::new(
self.syntax.text_range().start(),
group.syntax.text_range().end(),
),
None => self.syntax.text_range(),
}
}
}
impl Group {
pub fn inner_text(&self) -> Option<SmolStr> {
Some(flat_inner(&self.syntax)?.text)
}
pub fn inner(&self) -> Option<(TextRange, SmolStr)> {
let inner = flat_inner(&self.syntax)?;
let range = inner
.range
.unwrap_or_else(|| TextRange::empty(inner.empty_anchor));
Some((range, inner.text))
}
pub fn inner_source(&self) -> String {
inner_source_of(&self.syntax)
}
pub fn command(&self) -> Option<Command> {
child::<Command>(&self.syntax)
}
pub fn command_name(&self) -> Option<SmolStr> {
self.command()?.name()
}
}
pub(crate) fn inner_source_of(node: &SyntaxNode) -> String {
let mut text = String::new();
for element in node.descendants_with_tokens() {
if let NodeOrToken::Token(token) = element {
text.push_str(token.text());
}
}
let inner = text.strip_prefix('{').unwrap_or(&text);
inner.strip_suffix('}').unwrap_or(inner).to_string()
}
impl NameGroup {
pub fn text(&self) -> Option<String> {
Some(flat_inner(&self.syntax)?.text.to_string())
}
pub fn range(&self) -> Option<TextRange> {
flat_inner(&self.syntax)?.range
}
}
struct FlatInner {
text: SmolStr,
range: Option<TextRange>,
empty_anchor: TextSize,
}
fn flat_inner(node: &SyntaxNode) -> Option<FlatInner> {
let mut text = SmolStrBuilder::new();
let mut start = None;
let mut end = None;
let mut empty_anchor = node.text_range().start();
for element in node.children_with_tokens() {
match element {
NodeOrToken::Token(token) => match token.kind() {
SyntaxKind::L_BRACE => empty_anchor = token.text_range().end(),
SyntaxKind::R_BRACE => {}
SyntaxKind::HASH => return None,
_ => {
let token_range = token.text_range();
start.get_or_insert(token_range.start());
end = Some(token_range.end());
text.push_str(token.text());
}
},
NodeOrToken::Node(_) => return None,
}
}
Some(FlatInner {
text: text.finish(),
range: start
.zip(end)
.map(|(start, end)| TextRange::new(start, end)),
empty_anchor,
})
}
fn alias_delimiter_name(node: &SyntaxNode, keyword: &str) -> Option<String> {
let head = child_token::<ControlWord>(node)?;
let text = head.syntax().text();
(text != keyword)
.then(|| text.strip_prefix('\\'))
.flatten()
.filter(|name| !name.is_empty())
.map(str::to_owned)
}
impl Begin {
pub fn name_group(&self) -> Option<NameGroup> {
child::<NameGroup>(&self.syntax)
}
pub fn name(&self) -> Option<String> {
match self.name_group() {
Some(group) => group.text(),
None => alias_delimiter_name(&self.syntax, "\\begin"),
}
}
pub fn is_alias(&self) -> bool {
self.name_group().is_none() && alias_delimiter_name(&self.syntax, "\\begin").is_some()
}
pub fn name_range(&self) -> Option<TextRange> {
self.name_group()?.range()
}
}
impl End {
pub fn name_group(&self) -> Option<NameGroup> {
child::<NameGroup>(&self.syntax)
}
pub fn name(&self) -> Option<String> {
self.name_group()?.text()
}
pub fn name_range(&self) -> Option<TextRange> {
self.name_group()?.range()
}
}
impl Environment {
pub fn begin(&self) -> Option<Begin> {
child::<Begin>(&self.syntax)
}
pub fn end(&self) -> Option<End> {
child::<End>(&self.syntax)
}
pub fn name(&self) -> Option<String> {
self.begin()?.name()
}
}
impl Conditional {
pub fn branches(&self) -> impl Iterator<Item = ConditionalBranch> {
children::<ConditionalBranch>(&self.syntax)
}
pub fn closer(&self) -> Option<Command> {
self.syntax.last_child().and_then(Command::cast)
}
}
impl ConditionalBranch {
pub fn head(&self) -> Option<Command> {
self.syntax.first_child().and_then(Command::cast)
}
}