use normalize_facts::{Extractor, Symbol};
use normalize_languages::parsers::{grammar_loader, parse_with_grammar};
use normalize_languages::{Language, support_for_path};
use std::path::Path;
use streaming_iterator::StreamingIterator;
pub use normalize_languages::ContainerBody;
#[derive(Debug)]
#[allow(dead_code)] pub struct SymbolLocation {
pub name: String,
pub kind: String,
pub start_byte: usize,
pub end_byte: usize,
pub start_line: usize,
pub end_line: usize,
pub indent: String,
}
pub fn line_to_byte(content: &str, line: usize) -> usize {
if line <= 1 {
return 0;
}
let target = line - 1; let mut newlines_seen = 0usize;
let mut i = 0usize;
while i < content.len() {
let b = content.as_bytes()[i];
if b == b'\n' {
newlines_seen += 1;
if newlines_seen == target {
return (i + 1).min(content.len());
}
}
let ch_len = content[i..]
.chars()
.next()
.map(|c| c.len_utf8())
.unwrap_or(1);
i += ch_len;
}
content.len()
}
pub struct Editor {}
impl Default for Editor {
fn default() -> Self {
Self::new()
}
}
impl Editor {
pub fn new() -> Self {
Self {}
}
pub fn find_symbol(
&self,
path: &Path,
content: &str,
name: &str,
case_insensitive: bool,
) -> Option<SymbolLocation> {
let extractor = Extractor::new();
let result = extractor.extract(path, content);
fn search_symbols(
symbols: &[Symbol],
name: &str,
content: &str,
case_insensitive: bool,
) -> Option<SymbolLocation> {
for sym in symbols {
let matches = if case_insensitive {
sym.name.eq_ignore_ascii_case(name)
} else {
sym.name == name
};
if matches {
let start_byte = line_to_byte(content, sym.start_line);
let end_byte = line_to_byte(content, sym.end_line + 1);
return Some(SymbolLocation {
name: sym.name.clone(),
kind: sym.kind.as_str().to_string(),
start_byte,
end_byte,
start_line: sym.start_line,
end_line: sym.end_line,
indent: String::new(),
});
}
if let Some(loc) = search_symbols(&sym.children, name, content, case_insensitive) {
return Some(loc);
}
}
None
}
search_symbols(&result.symbols, name, content, case_insensitive)
}
pub fn delete_symbol(&self, content: &str, loc: &SymbolLocation) -> String {
let mut result = String::new();
let line_start = content[..loc.start_byte]
.rfind('\n')
.map(|i| i + 1)
.unwrap_or(0);
let mut end_byte = loc.end_byte;
if end_byte < content.len() && content.as_bytes()[end_byte] == b'\n' {
end_byte += 1;
}
let has_blank_before =
line_start >= 2 && &content[line_start.saturating_sub(2)..line_start] == "\n\n";
if has_blank_before {
while end_byte < content.len() && content.as_bytes()[end_byte] == b'\n' {
end_byte += 1;
if end_byte < content.len() && content.as_bytes()[end_byte] != b'\n' {
break;
}
}
}
result.push_str(&content[..line_start]);
result.push_str(&content[end_byte..]);
result
}
pub fn replace_symbol(&self, content: &str, loc: &SymbolLocation, new_content: &str) -> String {
let mut result = String::new();
let indented = self.apply_indent(new_content, &loc.indent);
result.push_str(&content[..loc.start_byte]);
result.push_str(&indented);
result.push_str(&content[loc.end_byte..]);
result
}
fn count_blank_lines_before(&self, content: &str, pos: usize) -> usize {
let mut count = 0usize;
let mut i = pos;
while i > 0 {
i -= 1;
if content.as_bytes()[i] == b'\n' {
count += 1;
} else if !content.as_bytes()[i].is_ascii_whitespace() {
break;
}
}
count.saturating_sub(1) }
fn count_blank_lines_after(&self, content: &str, pos: usize) -> usize {
let mut count = 0;
let mut i = pos;
if i < content.len() && content.as_bytes()[i] == b'\n' {
i += 1;
}
while i < content.len() {
if content.as_bytes()[i] == b'\n' {
count += 1;
i += 1;
} else if content.as_bytes()[i].is_ascii_whitespace() {
i += 1;
} else {
break;
}
}
count
}
pub fn insert_before(&self, content: &str, loc: &SymbolLocation, new_content: &str) -> String {
let mut result = String::new();
let line_start = content[..loc.start_byte]
.rfind('\n')
.map(|i| i + 1)
.unwrap_or(0);
let blank_lines = self.count_blank_lines_before(content, line_start);
let spacing = "\n".repeat(blank_lines.max(1) + 1);
let indented = self.apply_indent(new_content, &loc.indent);
result.push_str(&content[..line_start]);
result.push_str(&indented);
result.push_str(&spacing);
result.push_str(&content[line_start..]);
result
}
pub fn insert_after(&self, content: &str, loc: &SymbolLocation, new_content: &str) -> String {
let mut result = String::new();
let indented = self.apply_indent(new_content, &loc.indent);
let end_pos = if loc.end_byte < content.len() && content.as_bytes()[loc.end_byte] == b'\n' {
loc.end_byte + 1
} else {
loc.end_byte
};
let blank_lines = self.count_blank_lines_after(content, loc.end_byte);
let spacing = "\n".repeat(blank_lines.max(1));
let mut next_content_pos = end_pos;
while next_content_pos < content.len() && content.as_bytes()[next_content_pos] == b'\n' {
next_content_pos += 1;
}
result.push_str(&content[..end_pos]);
result.push_str(&spacing);
result.push_str(&indented);
if next_content_pos < content.len() {
result.push_str(&"\n".repeat(blank_lines.max(1) + 1));
result.push_str(&content[next_content_pos..]);
} else {
result.push('\n');
}
result
}
pub fn prepend_to_file(&self, content: &str, new_content: &str) -> String {
let mut result = String::new();
result.push_str(new_content);
if !new_content.ends_with('\n') {
result.push('\n');
}
result.push_str(content);
result
}
pub fn append_to_file(&self, content: &str, new_content: &str) -> String {
let mut result = String::new();
result.push_str(content);
if !content.ends_with('\n') {
result.push('\n');
}
result.push_str(new_content);
if !new_content.ends_with('\n') {
result.push('\n');
}
result
}
pub fn find_container_body(
&self,
path: &Path,
content: &str,
name: &str,
) -> Option<ContainerBody> {
let support = support_for_path(path)?;
let grammar = support.grammar_name();
let tree = parse_with_grammar(grammar, content)?;
let loader = grammar_loader();
let tags_scm = loader.get_tags(grammar)?;
let ts_lang = loader.get(grammar).ok()?;
let tags_query = tree_sitter::Query::new(&ts_lang, &tags_scm).ok()?;
find_container_body_via_tags(&tree, &tags_query, content, name, support)
}
pub fn prepend_to_container(
&self,
content: &str,
body: &ContainerBody,
new_content: &str,
) -> String {
let mut result = String::new();
let indented = self.apply_indent(new_content, &body.inner_indent);
result.push_str(&content[..body.content_start]);
result.push_str(&indented);
result.push('\n');
if !body.is_empty {
result.push('\n');
}
result.push_str(&content[body.content_start..]);
result
}
pub fn append_to_container(
&self,
content: &str,
body: &ContainerBody,
new_content: &str,
) -> String {
let mut result = String::new();
let indented = self.apply_indent(new_content, &body.inner_indent);
let mut end_pos = body.content_end;
while end_pos > 0
&& content
.as_bytes()
.get(end_pos - 1)
.map(|&b| b == b'\n' || b == b' ')
== Some(true)
{
end_pos -= 1;
}
result.push_str(&content[..end_pos]);
if !body.is_empty {
result.push_str("\n\n");
} else {
result.push('\n');
}
result.push_str(&indented);
result.push('\n');
result.push_str(&content[body.content_end..]);
result
}
pub fn rename_identifier_in_line(
&self,
content: &str,
line_no: usize,
old_name: &str,
new_name: &str,
) -> Option<String> {
let (line_start, line_end) = line_byte_range(content, line_no)?;
let line = &content[line_start..line_end];
let new_line = replace_all_words(line, old_name, new_name);
if new_line == line {
return None;
}
let mut result = String::with_capacity(content.len() + new_name.len() * 4);
result.push_str(&content[..line_start]);
result.push_str(&new_line);
result.push_str(&content[line_end..]);
Some(result)
}
pub fn apply_indent(&self, content: &str, indent: &str) -> String {
content
.lines()
.enumerate()
.map(|(i, line)| {
if i == 0 {
format!("{}{}", indent, line)
} else if line.is_empty() {
line.to_string()
} else {
format!("{}{}", indent, line)
}
})
.collect::<Vec<_>>()
.join("\n")
}
}
fn line_byte_range(content: &str, line_no: usize) -> Option<(usize, usize)> {
if line_no == 0 {
return None;
}
let mut start = 0usize;
let mut current_line = 1usize;
for (i, c) in content.char_indices() {
if current_line == line_no {
let end = content[i..]
.find('\n')
.map(|n| i + n)
.unwrap_or(content.len());
return Some((start, end));
}
if c == '\n' {
current_line += 1;
start = i + 1;
}
}
if current_line == line_no {
Some((start, content.len()))
} else {
None
}
}
pub fn replace_all_words(text: &str, old: &str, new_word: &str) -> String {
if old.is_empty() {
return text.to_string();
}
let bytes = text.as_bytes();
let mut result = String::with_capacity(text.len());
let mut offset = 0;
loop {
match text[offset..].find(old) {
None => {
result.push_str(&text[offset..]);
break;
}
Some(pos) => {
let abs = offset + pos;
let before_ok = abs == 0 || {
let b = bytes[abs - 1];
!b.is_ascii_alphanumeric() && b != b'_'
};
let after = abs + old.len();
let after_ok = after >= bytes.len() || {
let b = bytes[after];
!b.is_ascii_alphanumeric() && b != b'_'
};
if before_ok && after_ok {
result.push_str(&text[offset..abs]);
result.push_str(new_word);
offset = after;
} else {
let next = text[abs..]
.chars()
.next()
.map(|c| c.len_utf8())
.unwrap_or(1);
result.push_str(&text[offset..abs + next]);
offset = abs + next;
}
}
}
}
result
}
fn find_container_body_via_tags(
tree: &tree_sitter::Tree,
tags_query: &tree_sitter::Query,
content: &str,
name: &str,
support: &dyn Language,
) -> Option<ContainerBody> {
let capture_names = tags_query.capture_names();
let root = tree.root_node();
let mut qcursor = tree_sitter::QueryCursor::new();
let mut matches = qcursor.matches(tags_query, root, content.as_bytes());
while let Some(m) = matches.next() {
for capture in m.captures {
let cn = capture_names[capture.index as usize];
if !matches!(
cn,
"definition.class" | "definition.module" | "definition.interface"
) {
continue;
}
let node = capture.node;
let container_name = support.node_name(&node, content)?;
if container_name != name {
continue;
}
let body_node = support.container_body(&node)?;
let start_byte = node.start_byte();
let line_start = content[..start_byte]
.rfind('\n')
.map(|i| i + 1)
.unwrap_or(0);
let container_indent: String = content[line_start..start_byte]
.chars()
.take_while(|c| c.is_whitespace())
.collect();
let inner_indent = format!("{} ", container_indent);
if let Some(body) = support.analyze_container_body(&body_node, content, &inner_indent) {
return Some(body);
}
}
}
None
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::PathBuf;
#[test]
fn test_find_python_function() {
let editor = Editor::new();
let content = r#"
def foo():
pass
def bar():
return 42
"#;
let loc = editor.find_symbol(&PathBuf::from("test.py"), content, "bar", false);
assert!(loc.is_some());
let loc = loc.unwrap();
assert_eq!(loc.name, "bar");
assert_eq!(loc.kind, "function");
}
#[test]
fn test_delete_symbol() {
let editor = Editor::new();
let content = "def foo():\n pass\n\ndef bar():\n return 42\n";
let loc = editor
.find_symbol(&PathBuf::from("test.py"), content, "bar", false)
.unwrap();
let result = editor.delete_symbol(content, &loc);
assert!(!result.contains("bar"));
assert!(result.contains("foo"));
}
#[test]
fn test_insert_before() {
let editor = Editor::new();
let content = "def foo():\n pass\n\ndef bar():\n return 42\n";
let loc = editor
.find_symbol(&PathBuf::from("test.py"), content, "bar", false)
.unwrap();
let result = editor.insert_before(content, &loc, "def baz():\n pass");
assert!(result.contains("baz"));
assert!(result.find("baz").unwrap() < result.find("bar").unwrap());
}
#[test]
fn test_prepend_to_python_class() {
let editor = Editor::new();
let content = r#"class Foo:
"""Docstring."""
def first(self):
pass
"#;
let body = editor
.find_container_body(&PathBuf::from("test.py"), content, "Foo")
.unwrap();
let result =
editor.prepend_to_container(content, &body, "def new_method(self):\n return 1");
assert!(result.contains("new_method"));
let docstring_pos = result.find("Docstring").unwrap();
let new_method_pos = result.find("new_method").unwrap();
let first_pos = result.find("first").unwrap();
assert!(docstring_pos < new_method_pos);
assert!(new_method_pos < first_pos);
}
#[test]
fn test_append_to_python_class() {
let editor = Editor::new();
let content = r#"class Foo:
def first(self):
pass
def second(self):
return 42
"#;
let body = editor
.find_container_body(&PathBuf::from("test.py"), content, "Foo")
.unwrap();
let result = editor.append_to_container(content, &body, "def last(self):\n return 99");
assert!(result.contains("last"));
let second_pos = result.find("second").unwrap();
let last_pos = result.find("last").unwrap();
assert!(second_pos < last_pos);
}
#[test]
fn test_prepend_to_rust_impl() {
let editor = Editor::new();
let content = r#"impl Foo {
fn first(&self) -> i32 {
1
}
}
"#;
let body = editor
.find_container_body(&PathBuf::from("test.rs"), content, "Foo")
.unwrap();
let result =
editor.prepend_to_container(content, &body, "fn new() -> Self {\n Self {}\n}");
assert!(result.contains("new"));
let new_pos = result.find("new").unwrap();
let first_pos = result.find("first").unwrap();
assert!(new_pos < first_pos);
}
#[test]
fn test_append_to_rust_impl() {
let editor = Editor::new();
let content = r#"impl Foo {
fn first(&self) -> i32 {
1
}
}
"#;
let body = editor
.find_container_body(&PathBuf::from("test.rs"), content, "Foo")
.unwrap();
let result =
editor.append_to_container(content, &body, "fn last(&self) -> i32 {\n 99\n}");
assert!(result.contains("last"));
let first_pos = result.find("first").unwrap();
let last_pos = result.find("last").unwrap();
assert!(first_pos < last_pos);
assert!(result.contains("}"));
}
}