_diffctx/
signatures.rs

1use std::sync::Arc;
2
3use rustc_hash::FxHashSet;
4
5use crate::config::graph_filtering::GRAPH_FILTERING;
6use crate::types::{Fragment, FragmentId, FragmentKind};
7
8fn is_signature_eligible(kind: FragmentKind) -> bool {
9    // Variable covers TS/JS arrow-function bindings (`const f = (...) => {...}`);
10    // without a stub variant a large changed arrow function that misses the core
11    // budget vanishes from the output entirely (#106).
12    matches!(
13        kind,
14        FragmentKind::Function
15            | FragmentKind::Class
16            | FragmentKind::Struct
17            | FragmentKind::Interface
18            | FragmentKind::Enum
19            | FragmentKind::Variable
20    )
21}
22
23fn signature_kind(kind: FragmentKind) -> FragmentKind {
24    match kind {
25        FragmentKind::Function => FragmentKind::FunctionSignature,
26        FragmentKind::Class => FragmentKind::ClassSignature,
27        FragmentKind::Struct => FragmentKind::StructSignature,
28        FragmentKind::Interface => FragmentKind::InterfaceSignature,
29        FragmentKind::Enum => FragmentKind::EnumSignature,
30        _ => FragmentKind::FunctionSignature,
31    }
32}
33
34fn count_brackets_outside_strings(line: &str) -> (i32, i32, i32, i32) {
35    let mut open_parens = 0i32;
36    let mut close_parens = 0i32;
37    let mut open_braces = 0i32;
38    let mut close_braces = 0i32;
39    let mut in_string: Option<char> = None;
40    let mut escaped = false;
41
42    for ch in line.chars() {
43        if let Some(quote) = in_string {
44            if escaped {
45                escaped = false;
46            } else if ch == '\\' {
47                escaped = true;
48            } else if ch == quote {
49                in_string = None;
50            }
51            continue;
52        }
53        match ch {
54            '\'' | '"' | '`' => {
55                in_string = Some(ch);
56                escaped = false;
57            }
58            '(' => open_parens += 1,
59            ')' => close_parens += 1,
60            '{' => open_braces += 1,
61            '}' => close_braces += 1,
62            _ => {}
63        }
64    }
65
66    (open_parens, close_parens, open_braces, close_braces)
67}
68
69fn decorator_prefix_len(lines: &[&str]) -> usize {
70    let mut i = 0;
71    let mut paren_depth = 0i32;
72    while i < lines.len() {
73        let trimmed = lines[i].trim_start();
74        let starts_decorator = trimmed.starts_with('@') || trimmed.starts_with("#[");
75        if paren_depth <= 0 && !starts_decorator {
76            break;
77        }
78        let (op, cp, _, _) = count_brackets_outside_strings(lines[i]);
79        paren_depth += op - cp;
80        i += 1;
81    }
82    if i >= lines.len() { 0 } else { i }
83}
84
85fn find_signature_end(lines: &[&str]) -> usize {
86    let mut paren_depth = 0i32;
87    let mut seen_open_paren = false;
88
89    for (i, line) in lines.iter().enumerate() {
90        let (op, cp, ob, cb) = count_brackets_outside_strings(line);
91        paren_depth += op - cp;
92        if op > 0 {
93            seen_open_paren = true;
94        }
95        // A body-opening brace only ends the signature once we are outside the
96        // parameter list. Braces inside parameter defaults or annotations
97        // (e.g. Python `def f(x={}):`) appear while `paren_depth > 0` and must
98        // not truncate the signature mid-parameter-list.
99        if paren_depth <= 0 && ob - cb > 0 {
100            return i + 1;
101        }
102        if seen_open_paren && paren_depth <= 0 {
103            return i + 1;
104        }
105    }
106
107    2.min(lines.len())
108}
109
110pub fn generate_signature_variants(fragments: &[Fragment]) -> Vec<Fragment> {
111    let mut signatures: Vec<Fragment> = Vec::new();
112    let mut seen: FxHashSet<FragmentId> = FxHashSet::default();
113
114    for frag in fragments {
115        if !is_signature_eligible(frag.kind) {
116            continue;
117        }
118        if frag.line_count() < GRAPH_FILTERING.min_lines_for_signature {
119            continue;
120        }
121        let lines: Vec<&str> = frag.content.lines().collect();
122        let decorators = decorator_prefix_len(&lines);
123        let sig_end = decorators + find_signature_end(&lines[decorators..]);
124        let sig_content: String = lines[..sig_end].join("\n");
125        let sig_end_line = frag.start_line() + sig_end as u32 - 1;
126        let sig_id = FragmentId::new(frag.id.path.clone(), frag.start_line(), sig_end_line);
127
128        if seen.contains(&sig_id) {
129            continue;
130        }
131        seen.insert(sig_id.clone());
132
133        signatures.push(Fragment {
134            id: sig_id,
135            kind: signature_kind(frag.kind),
136            content: Arc::from(sig_content),
137            identifiers: frag.identifiers.clone(),
138            token_count: 0,
139            symbol_name: frag.symbol_name.clone(),
140        });
141    }
142
143    signatures
144}
_diffctx/signatures.rs

_diffctx/
signatures.rs
