use std::collections::HashSet;
use reqwest::Client;
use tree_sitter::{Node, Parser, Point};
use crate::config::Config;
use crate::diff::{parse_valid_lines, split_diff_sections};
use crate::providers::{PrMeta, Provider};
fn hunk_trailing(line: &str) -> Option<&str> {
let rest = line.strip_prefix("@@")?;
let idx = rest.find("@@")?;
let trailing = rest[idx + 2..].trim();
(!trailing.is_empty()).then_some(trailing)
}
fn plusplus_path(line: &str) -> Option<String> {
let rest = line.strip_prefix("+++ ")?;
let p = rest.trim();
let p = p.strip_prefix("b/").unwrap_or(p);
(p != "/dev/null").then(|| p.to_string())
}
fn hunk_regions(diff: &str) -> Vec<(String, Vec<String>)> {
let mut out: Vec<(String, Vec<String>)> = Vec::new();
let mut cur: Option<String> = None;
for line in diff.lines() {
if let Some(path) = line.strip_prefix("+++ ").and_then(|_| plusplus_path(line)) {
cur = Some(path);
} else if line.starts_with("@@") {
let (Some(path), Some(sig)) = (cur.as_ref(), hunk_trailing(line)) else {
continue;
};
let sig = sig.to_string();
match out.iter_mut().find(|(p, _)| p == path) {
Some((_, sigs)) => {
if !sigs.contains(&sig) {
sigs.push(sig);
}
}
None => out.push((path.clone(), vec![sig])),
}
}
}
out
}
fn format_regions(regions: &[(String, Vec<String>)]) -> String {
let mut s = String::from("Changed regions:");
for (path, sigs) in regions {
s.push_str(&format!("\n- {}: {}", path, sigs.join("; ")));
}
s
}
pub fn hunk_context(diff: &str) -> String {
let regions = hunk_regions(diff);
if regions.is_empty() {
return String::new();
}
format_regions(®ions)
}
struct Sym {
label: &'static str,
name: String,
start: u64,
end: u64,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum Lang {
Rust,
TypeScript,
Tsx,
JavaScript,
Python,
Go,
}
fn language_for_path(path: &str) -> Option<Lang> {
let ext = path.rsplit('.').next()?.to_ascii_lowercase();
Some(match ext.as_str() {
"rs" => Lang::Rust,
"ts" | "mts" | "cts" => Lang::TypeScript,
"tsx" => Lang::Tsx,
"js" | "jsx" | "mjs" | "cjs" => Lang::JavaScript,
"py" | "pyi" => Lang::Python,
"go" => Lang::Go,
_ => return None,
})
}
impl Lang {
fn ts_language(self) -> tree_sitter::Language {
match self {
Lang::Rust => tree_sitter_rust::LANGUAGE.into(),
Lang::TypeScript => tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
Lang::Tsx => tree_sitter_typescript::LANGUAGE_TSX.into(),
Lang::JavaScript => tree_sitter_javascript::LANGUAGE.into(),
Lang::Python => tree_sitter_python::LANGUAGE.into(),
Lang::Go => tree_sitter_go::LANGUAGE.into(),
}
}
fn def_label(self, kind: &str) -> Option<&'static str> {
match self {
Lang::Rust => match kind {
"function_item" => Some("fn"),
"struct_item" => Some("struct"),
"enum_item" => Some("enum"),
"union_item" => Some("union"),
"trait_item" => Some("trait"),
"impl_item" => Some("impl"),
"mod_item" => Some("mod"),
"type_item" => Some("type"),
"const_item" => Some("const"),
"static_item" => Some("static"),
"macro_definition" => Some("macro"),
_ => None,
},
Lang::TypeScript | Lang::Tsx | Lang::JavaScript => match kind {
"function_declaration" | "generator_function_declaration" => Some("function"),
"method_definition" => Some("method"),
"class_declaration" | "abstract_class_declaration" => Some("class"),
"interface_declaration" => Some("interface"),
"type_alias_declaration" => Some("type"),
"enum_declaration" => Some("enum"),
_ => None,
},
Lang::Python => match kind {
"function_definition" => Some("function"),
"class_definition" => Some("class"),
_ => None,
},
Lang::Go => match kind {
"function_declaration" | "method_declaration" => Some("func"),
"type_spec" => Some("type"),
_ => None,
},
}
}
}
fn unwrap_export(lang: Lang, node: Node) -> Node {
if matches!(lang, Lang::TypeScript | Lang::Tsx | Lang::JavaScript)
&& node.kind() == "export_statement"
{
let mut cursor = node.walk();
for child in node.named_children(&mut cursor) {
if lang.def_label(child.kind()).is_some() {
return child;
}
}
}
node
}
fn node_name(lang: Lang, node: Node, src: &[u8]) -> Option<String> {
if let Some(n) = node.child_by_field_name("name") {
if let Ok(t) = n.utf8_text(src) {
return Some(t.to_string());
}
}
if lang == Lang::Rust && node.kind() == "impl_item" {
if let Some(t) = node.child_by_field_name("type") {
if let Ok(s) = t.utf8_text(src) {
return Some(s.to_string());
}
}
}
None
}
fn symbols_for_file(lang: Lang, source: &str, changed: &HashSet<u64>) -> Vec<Sym> {
let mut parser = Parser::new();
if parser.set_language(&lang.ts_language()).is_err() {
return Vec::new();
}
let tree = match parser.parse(source, None) {
Some(t) => t,
None => return Vec::new(),
};
let src = source.as_bytes();
let root = tree.root_node();
let mut lines: Vec<u64> = changed.iter().copied().collect();
lines.sort_unstable();
let mut seen: HashSet<(&'static str, String, u64, u64)> = HashSet::new();
let mut out: Vec<Sym> = Vec::new();
for line in lines {
if line == 0 {
continue;
}
let pt = Point {
row: (line - 1) as usize,
column: 0,
};
let mut node = match root.descendant_for_point_range(pt, pt) {
Some(n) => n,
None => continue,
};
loop {
let def = unwrap_export(lang, node);
if let Some(label) = lang.def_label(def.kind()) {
if let Some(name) = node_name(lang, def, src) {
let start = def.start_position().row as u64 + 1;
let end = def.end_position().row as u64 + 1;
if seen.insert((label, name.clone(), start, end)) {
out.push(Sym {
label,
name,
start,
end,
});
}
break;
}
}
match node.parent() {
Some(p) => node = p,
None => break,
}
}
}
out
}
fn format_symbols(path: &str, syms: &[Sym]) -> String {
let parts: Vec<String> = syms
.iter()
.map(|s| {
format!(
"{} {} (lines {}\u{2013}{})",
s.label, s.name, s.start, s.end
)
})
.collect();
format!("- {}: {}", path, parts.join(", "))
}
fn diff_file_order(diff: &str) -> Vec<String> {
let mut seen: HashSet<String> = HashSet::new();
let mut order: Vec<String> = Vec::new();
for (path, _) in split_diff_sections(diff) {
if !path.is_empty() && seen.insert(path.clone()) {
order.push(path);
}
}
order
}
pub async fn structural_context(
provider: &Provider,
client: &Client,
cfg: &Config,
repo: &str,
meta: &PrMeta,
diff: &str,
) -> String {
let head = match meta.head_sha.as_deref() {
Some(s) if !s.is_empty() => s,
_ => return hunk_context(diff),
};
let valid = parse_valid_lines(diff);
let mut covered: HashSet<String> = HashSet::new();
let mut ts_blocks: Vec<String> = Vec::new();
let mut attempted = 0usize;
for path in diff_file_order(diff) {
if attempted >= cfg.structural_max_files {
break;
}
let Some(lang) = language_for_path(&path) else {
continue;
};
let Some(lines) = valid.get(&path).filter(|s| !s.is_empty()) else {
continue;
};
attempted += 1;
let content = match provider
.get_file_contents(client, cfg, repo, head, &path)
.await
{
Ok(Some(c)) if c.len() <= 400_000 => c,
_ => continue,
};
let syms = symbols_for_file(lang, &content, lines);
if !syms.is_empty() {
covered.insert(path.clone());
ts_blocks.push(format_symbols(&path, &syms));
}
}
let mut out = String::new();
if !ts_blocks.is_empty() {
out.push_str("Changed symbols:\n");
out.push_str(&ts_blocks.join("\n"));
}
let regions: Vec<(String, Vec<String>)> = hunk_regions(diff)
.into_iter()
.filter(|(p, _)| !covered.contains(p))
.collect();
if !regions.is_empty() {
if !out.is_empty() {
out.push_str("\n\n");
}
out.push_str(&format_regions(®ions));
}
out
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn hunk_context_extracts_trailing_signature() {
let d = "+++ b/orders.go\n\
@@ -10,7 +10,8 @@ func processOrder(o Order) {\n\
+added\n";
let ctx = hunk_context(d);
assert!(ctx.starts_with("Changed regions:"));
assert!(ctx.contains("orders.go"));
assert!(ctx.contains("func processOrder(o Order) {"));
}
#[test]
fn hunk_context_empty_when_no_trailing() {
let d = "+++ b/a.rs\n@@ -1,2 +1,3 @@\n+added\n";
assert!(hunk_context(d).is_empty());
}
#[test]
fn hunk_context_dedups_and_groups_per_file() {
let d = "+++ b/a.ts\n\
@@ -1,3 +1,4 @@ class Foo {\n\
+a\n\
@@ -20,3 +21,4 @@ class Foo {\n\
+b\n\
+++ b/b.ts\n\
@@ -1,3 +1,4 @@ function bar() {\n\
+c\n";
let ctx = hunk_context(d);
assert_eq!(ctx.matches("class Foo {").count(), 1);
assert!(ctx.contains("a.ts: class Foo {"));
assert!(ctx.contains("b.ts: function bar() {"));
}
#[test]
fn hunk_context_skips_dev_null() {
let d = "+++ /dev/null\n@@ -1,1 +0,0 @@ func gone() {\n-gone\n";
assert!(hunk_context(d).is_empty());
}
#[test]
fn ts_finds_enclosing_rust_fn_and_impl() {
let src = "\
struct Point { x: i32 }
impl Point {
fn clamp(v: i32) -> i32 {
v
}
}
";
let changed: HashSet<u64> = [5u64].into_iter().collect();
let syms = symbols_for_file(Lang::Rust, src, &changed);
assert_eq!(syms.len(), 1, "smallest enclosing def only");
assert_eq!(syms[0].label, "fn");
assert_eq!(syms[0].name, "clamp");
}
#[test]
fn ts_reports_outer_def_for_line_outside_inner() {
let src = "\
impl Point {
fn clamp(v: i32) -> i32 {
v
}
}
";
let changed: HashSet<u64> = [1u64].into_iter().collect();
let syms = symbols_for_file(Lang::Rust, src, &changed);
assert_eq!(syms.len(), 1);
assert_eq!(syms[0].label, "impl");
assert_eq!(syms[0].name, "Point");
}
#[test]
fn ts_finds_typescript_class_and_method() {
let src = "\
export class OrderService {
processOrder(o: Order) {
return o;
}
}
";
let changed: HashSet<u64> = [1u64, 3u64].into_iter().collect();
let syms = symbols_for_file(Lang::TypeScript, src, &changed);
let names: Vec<_> = syms.iter().map(|s| (s.label, s.name.as_str())).collect();
assert!(names.contains(&("class", "OrderService")));
assert!(names.contains(&("method", "processOrder")));
}
#[test]
fn ts_finds_python_function() {
let src = "\
def process(order):
total = 0
return total
";
let changed: HashSet<u64> = [2u64].into_iter().collect();
let syms = symbols_for_file(Lang::Python, src, &changed);
assert_eq!(syms.len(), 1);
assert_eq!(syms[0].label, "function");
assert_eq!(syms[0].name, "process");
}
#[test]
fn ts_finds_go_func() {
let src = "\
package main
func processOrder(o Order) int {
return 0
}
";
let changed: HashSet<u64> = [4u64].into_iter().collect();
let syms = symbols_for_file(Lang::Go, src, &changed);
assert_eq!(syms.len(), 1);
assert_eq!(syms[0].label, "func");
assert_eq!(syms[0].name, "processOrder");
}
#[test]
fn language_for_path_maps_extensions() {
assert!(matches!(language_for_path("a/b.rs"), Some(Lang::Rust)));
assert!(matches!(
language_for_path("a/b.ts"),
Some(Lang::TypeScript)
));
assert!(matches!(language_for_path("a/b.tsx"), Some(Lang::Tsx)));
assert!(matches!(
language_for_path("a/b.jsx"),
Some(Lang::JavaScript)
));
assert!(matches!(language_for_path("a/b.py"), Some(Lang::Python)));
assert!(matches!(language_for_path("a/b.go"), Some(Lang::Go)));
assert!(language_for_path("a/b.md").is_none());
assert!(language_for_path("Makefile").is_none());
}
#[test]
fn format_symbols_uses_en_dash_span() {
let syms = vec![Sym {
label: "fn",
name: "clamp".to_string(),
start: 5,
end: 9,
}];
let s = format_symbols("src/util.rs", &syms);
assert_eq!(s, "- src/util.rs: fn clamp (lines 5\u{2013}9)");
}
}