use crate::lang::detect_language;
use crate::parser::get_parser;
use std::collections::HashSet;
use std::path::Path;
use tree_sitter::{Node, Tree};
#[derive(Debug)]
pub struct TreeContext {
lines: Vec<String>,
tree: Option<Tree>,
pub lois: HashSet<usize>,
context_lines: HashSet<usize>,
}
impl TreeContext {
pub fn new(content: &str, path: &Path) -> Self {
let lines: Vec<String> = content.lines().map(String::from).collect();
let tree = detect_language(path)
.and_then(get_parser)
.and_then(|mut parser| parser.parse(content, None));
TreeContext {
lines,
tree,
lois: HashSet::new(),
context_lines: HashSet::new(),
}
}
pub fn reset_lois(&mut self) {
self.lois.clear();
self.context_lines.clear();
}
pub fn add_lines_of_interest(&mut self, lois: &[i32]) {
for &line in lois {
if line > 0 {
self.lois.insert((line - 1) as usize);
}
}
}
pub fn add_context(&mut self) {
if let Some(tree) = &self.tree {
let root = tree.root_node();
let lois: Vec<usize> = self.lois.iter().copied().collect();
let mut new_context = HashSet::new();
for loi in lois {
collect_scope_context(&root, loi, &mut new_context);
}
self.context_lines.extend(new_context);
}
}
pub fn format(&self) -> String {
if self.lines.is_empty() {
return String::new();
}
let mut show_lines: Vec<usize> = self
.lois
.iter()
.chain(self.context_lines.iter())
.copied()
.filter(|&l| l < self.lines.len())
.collect();
show_lines.sort_unstable();
show_lines.dedup();
if show_lines.is_empty() {
return String::new();
}
let mut output = String::new();
let mut last_line: Option<usize> = None;
for line_idx in show_lines {
if let Some(last) = last_line
&& line_idx > last + 1
{
output.push_str("⋮\n");
}
output.push_str(&self.lines[line_idx]);
output.push('\n');
last_line = Some(line_idx);
}
output
}
}
fn collect_scope_context(root: &Node, loi: usize, context_lines: &mut HashSet<usize>) {
let point = tree_sitter::Point::new(loi, 0);
let mut node = root.descendant_for_point_range(point, point);
while let Some(n) = node {
if is_scope_node(&n) {
let start_line = n.start_position().row;
context_lines.insert(start_line);
}
node = n.parent();
}
}
fn is_scope_node(node: &Node) -> bool {
let kind = node.kind();
matches!(
kind,
"function_definition"
| "function_declaration"
| "function_item"
| "method_definition"
| "method_declaration"
| "arrow_function"
| "function_expression"
| "generator_function"
| "generator_function_declaration"
| "class_definition"
| "class_declaration"
| "struct_item"
| "enum_item"
| "trait_item"
| "interface_declaration"
| "impl_item"
| "module"
| "mod_item"
| "namespace_definition"
| "if_statement"
| "if_expression"
| "for_statement"
| "for_expression"
| "while_statement"
| "while_expression"
| "match_expression"
| "try_statement"
)
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::Path;
#[test]
fn tree_context_basic() {
let content = r#"fn main() {
println!("hello");
}
fn helper() {
something();
}
"#;
let mut ctx = TreeContext::new(content, Path::new("test.rs"));
ctx.add_lines_of_interest(&[2]); ctx.add_context();
let output = ctx.format();
assert!(output.contains("fn main()"));
assert!(output.contains("println"));
}
#[test]
fn tree_context_ellipsis() {
let content = "line1\nline2\nline3\nline4\nline5\n";
let mut ctx = TreeContext::new(content, Path::new("test.txt"));
ctx.lois.insert(0); ctx.lois.insert(4);
let output = ctx.format();
assert!(output.contains("line1"));
assert!(output.contains("⋮")); assert!(output.contains("line5"));
}
#[test]
fn tree_context_negative_line() {
let content = "fn foo() {}\n";
let mut ctx = TreeContext::new(content, Path::new("test.rs"));
ctx.add_lines_of_interest(&[-1, 1]); ctx.add_context();
let output = ctx.format();
assert!(output.contains("fn foo()"));
}
#[test]
fn tree_context_empty() {
let ctx = TreeContext::new("", Path::new("test.rs"));
let output = ctx.format();
assert!(output.is_empty());
}
#[test]
fn tree_context_reset() {
let content = "fn foo() {}\n";
let mut ctx = TreeContext::new(content, Path::new("test.rs"));
ctx.add_lines_of_interest(&[1]);
ctx.add_context();
ctx.reset_lois();
let output = ctx.format();
assert!(output.is_empty());
}
}