use crate::lang::{self, Grammar, Lang};
use std::{collections::HashSet, ops::Range};
use theme::HighlightKind;
use tree_sitter::{InputEdit, Parser, Point, Query, QueryCursor, StreamingIterator as _, Tree};
pub struct Document {
parser: Parser,
query: Query,
kinds: Vec<Option<HighlightKind>>,
tree: Tree,
source: String,
}
impl Document {
pub fn new(tag: &str, source: impl Into<String>) -> Option<Self> {
Self::with_lang(lang::resolve(tag)?, source)
}
pub fn with_lang(lang: &'static Lang, source: impl Into<String>) -> Option<Self> {
if !lang.injections.is_empty() {
return None;
}
let language = match &lang.grammar {
Grammar::Native(grammar) => tree_sitter::Language::from(*grammar),
Grammar::Wasm(_) => return None,
};
let query = Query::new(&language, lang.query).ok()?;
let kinds = query
.capture_names()
.iter()
.map(|name| recognized(name).map(lang::kind_of))
.collect();
let mut parser = Parser::new();
parser.set_language(&language).ok()?;
let source = source.into();
let tree = parser.parse(&source, None)?;
Some(Self {
parser,
query,
kinds,
tree,
source,
})
}
pub fn source(&self) -> &str {
&self.source
}
pub fn edit(&mut self, range: Range<usize>, new_end: usize, source: &str) -> Vec<Range<usize>> {
let edit = InputEdit {
start_byte: range.start,
old_end_byte: range.end,
new_end_byte: new_end,
start_position: point_at(&self.source, range.start),
old_end_position: point_at(&self.source, range.end),
new_end_position: point_at(source, new_end),
};
self.tree.edit(&edit);
self.source.clear();
self.source.push_str(source);
let Some(tree) = self.parser.parse(&self.source, Some(&self.tree)) else {
return std::iter::once(0..self.source.len()).collect();
};
let mut changed: Vec<Range<usize>> = self
.tree
.changed_ranges(&tree)
.map(|range| range.start_byte..range.end_byte)
.collect();
changed.push(range.start..new_end);
self.tree = tree;
merge(changed)
}
pub fn spans(&self) -> Vec<(Range<usize>, HighlightKind)> {
self.spans_in(0..self.source.len())
}
pub fn spans_in(&self, range: Range<usize>) -> Vec<(Range<usize>, HighlightKind)> {
let mut cursor = QueryCursor::new();
cursor.set_byte_range(range.clone());
let mut captures =
cursor.captures(&self.query, self.tree.root_node(), self.source.as_bytes());
let mut seen = HashSet::new();
let mut claimed: Vec<(Range<usize>, Option<HighlightKind>)> = Vec::new();
while let Some((found, index)) = captures.next() {
let capture = found.captures()[*index];
if seen.insert(capture.node.id()) {
claimed.push((
capture.node.byte_range(),
self.kinds[capture.index as usize],
));
}
}
claimed.sort_by(|a, b| a.0.start.cmp(&b.0.start).then(b.0.end.cmp(&a.0.end)));
flatten(
claimed
.into_iter()
.filter_map(|(span, kind)| Some((span, kind?))),
range,
)
}
}
fn flatten(
captures: impl Iterator<Item = (Range<usize>, HighlightKind)>,
within: Range<usize>,
) -> Vec<(Range<usize>, HighlightKind)> {
let mut spans = Vec::new();
let mut stack: Vec<(usize, HighlightKind)> = Vec::new();
let mut at = within.start;
let mut paint = |to: usize, stack: &[(usize, HighlightKind)], at: &mut usize| {
let to = to.min(within.end);
if to > *at {
if let Some((_, kind)) = stack.last() {
spans.push((*at..to, *kind));
}
*at = to;
}
};
for (span, kind) in captures {
while let Some(&(end, _)) = stack.last() {
if end > span.start {
break;
}
paint(end, &stack, &mut at);
stack.pop();
}
paint(span.start, &stack, &mut at);
stack.push((span.end, kind));
}
while let Some(&(end, _)) = stack.last() {
paint(end, &stack, &mut at);
stack.pop();
}
spans
}
fn point_at(text: &str, offset: usize) -> Point {
let before = &text.as_bytes()[..offset.min(text.len())];
let row = before.iter().filter(|byte| **byte == b'\n').count();
let column = before
.iter()
.rposition(|byte| *byte == b'\n')
.map_or(before.len(), |newline| before.len() - newline - 1);
Point { row, column }
}
fn merge(mut ranges: Vec<Range<usize>>) -> Vec<Range<usize>> {
ranges.sort_by_key(|range| range.start);
let mut merged: Vec<Range<usize>> = Vec::new();
for range in ranges {
match merged.last_mut() {
Some(last) if range.start <= last.end => last.end = last.end.max(range.end),
_ => merged.push(range),
}
}
merged
}
fn recognized(capture: &str) -> Option<&'static str> {
let parts: Vec<&str> = capture.split('.').collect();
let mut best: Option<(&'static str, usize)> = None;
for name in lang::NAMES {
let len = name.split('.').count();
if name.split('.').all(|part| parts.contains(&part))
&& best.is_none_or(|(_, longest)| len > longest)
{
best = Some((name, len));
}
}
best.map(|(name, _)| name)
}