use std::cell::RefCell;
use std::ops::ControlFlow;
use tree_sitter::{Node, Parser, Tree};
thread_local! {
static R_PARSER: RefCell<Parser> = RefCell::new({
let mut parser = Parser::new();
parser
.set_language(&tree_sitter_r::LANGUAGE.into())
.expect("Failed to set tree-sitter-r language");
parser
});
}
pub fn parse_r(source: &str) -> Option<Tree> {
R_PARSER.with(|parser| parser.borrow_mut().parse(source.as_bytes(), None))
}
pub fn is_atomic_node(kind: &str) -> bool {
matches!(
kind,
"string"
| "comment"
| "integer"
| "float"
| "complex"
| "true"
| "false"
| "null"
| "inf"
| "nan"
| "na"
| "dots"
| "dot_dot_i"
| "special" )
}
pub fn walk_atomic_tokens<'tree>(
root: Node<'tree>,
mut visit: impl FnMut(Node<'tree>) -> ControlFlow<()>,
) {
let mut cursor = root.walk();
loop {
let node = cursor.node();
if is_atomic_node(node.kind()) || node.child_count() == 0 {
if visit(node).is_break() {
return;
}
} else if cursor.goto_first_child() {
continue;
}
while !cursor.goto_next_sibling() {
if !cursor.goto_parent() {
return;
}
}
}
}
fn meaningful_token_range(node: Node<'_>, source: &[u8]) -> Option<(usize, usize)> {
let start = node.start_byte();
let end = node.end_byte();
if start >= end {
return None;
}
if is_atomic_node(node.kind()) {
return Some((start, end));
}
let text = std::str::from_utf8(&source[start..end]).ok()?;
if text.chars().all(char::is_whitespace) {
return None;
}
Some((start, end))
}
pub fn collect_atomic_tokens(root: Node<'_>, source: &[u8]) -> Vec<(usize, usize)> {
let mut tokens = Vec::new();
walk_atomic_tokens(root, |node| {
tokens.extend(meaningful_token_range(node, source));
ControlFlow::Continue(())
});
tokens.sort_by_key(|(start, _)| *start);
tokens
}
pub fn first_atomic_token(root: Node<'_>, source: &[u8]) -> Option<(usize, usize)> {
let mut first = None;
walk_atomic_tokens(root, |node| {
first = meaningful_token_range(node, source);
if first.is_some() {
ControlFlow::Break(())
} else {
ControlFlow::Continue(())
}
});
first
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_r_simple() {
let tree = parse_r("x <- 1");
assert!(tree.is_some());
let tree = tree.unwrap();
assert_eq!(tree.root_node().kind(), "program");
}
#[test]
fn test_parse_r_empty() {
let tree = parse_r("");
assert!(tree.is_some());
}
#[test]
fn test_parse_r_complex() {
let tree = parse_r("function(x) { x |> filter(y > 0) %>% select(-z) }");
assert!(tree.is_some());
}
#[test]
fn test_is_atomic_node() {
for kind in [
"string",
"comment",
"integer",
"float",
"complex",
"true",
"false",
"null",
"na",
"special",
"dots",
"dot_dot_i",
] {
assert!(is_atomic_node(kind), "{kind} should be atomic");
}
for kind in ["identifier", "program", "call", "binary_operator"] {
assert!(!is_atomic_node(kind), "{kind} should not be atomic");
}
}
fn tokens_of(source: &str) -> Vec<&str> {
let tree = parse_r(source).unwrap();
collect_atomic_tokens(tree.root_node(), source.as_bytes())
.into_iter()
.map(|(start, end)| &source[start..end])
.collect()
}
#[test]
fn test_collect_atomic_tokens_keeps_atomic_nodes_whole() {
assert_eq!(
tokens_of(r#"x |> f("a b", 1L) # note"#),
["x", "|>", "f", "(", r#""a b""#, ",", "1L", ")", "# note"]
);
assert_eq!(tokens_of("a %in% b"), ["a", "%in%", "b"]);
assert_eq!(tokens_of(""), Vec::<&str>::new());
}
#[test]
fn test_collect_atomic_tokens_skips_empty_missing_nodes() {
assert_eq!(tokens_of("f("), ["f", "("]);
}
#[test]
fn test_first_atomic_token_and_early_exit() {
let source = " x <- 1";
let tree = parse_r(source).unwrap();
let first = first_atomic_token(tree.root_node(), source.as_bytes());
assert_eq!(first, Some((2, 3)));
let mut visited = 0;
walk_atomic_tokens(tree.root_node(), |_| {
visited += 1;
ControlFlow::Break(())
});
assert_eq!(visited, 1);
let empty = parse_r("").unwrap();
assert_eq!(first_atomic_token(empty.root_node(), b""), None);
}
}