Skip to main content

brokk_bifrost_python/
clones.rs

1//! Python's clone-detection token and AST-signature normalization.
2//!
3//! `analyzer/python/clones.rs` in `brokk-bifrost-analysis` keeps the entry
4//! point: it reads the declaration's source through the analyzer and assembles
5//! the analysis-owned `CloneCandidateData`. Everything that knows what a Python
6//! token *is* lives here.
7
8use crate::declarations::parse_python_tree;
9#[cfg(any(test, feature = "test-support"))]
10use std::cell::Cell;
11use tree_sitter::Node;
12
13#[cfg(any(test, feature = "test-support"))]
14thread_local! {
15    static CLONE_PARSE_CALLS: Cell<usize> = const { Cell::new(0) };
16}
17
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct PythonClonePreparation {
20    pub normalized_tokens: Vec<String>,
21    pub ast_signature: String,
22}
23
24#[cfg(any(test, feature = "test-support"))]
25pub fn reset_python_clone_parse_count_for_test() {
26    CLONE_PARSE_CALLS.set(0);
27}
28
29#[cfg(any(test, feature = "test-support"))]
30pub fn python_clone_parse_count_for_test() -> usize {
31    CLONE_PARSE_CALLS.get()
32}
33
34fn parse_clone_tree(source: &str) -> Option<tree_sitter::Tree> {
35    #[cfg(any(test, feature = "test-support"))]
36    CLONE_PARSE_CALLS.set(CLONE_PARSE_CALLS.get() + 1);
37    parse_python_tree(source)
38}
39
40const PYTHON_CLONE_AST_IDENTIFIER_TYPES: &[&str] = &["identifier", "keyword_identifier"];
41const PYTHON_CLONE_AST_STRING_TYPES: &[&str] = &["string", "string_content", "interpolation"];
42const PYTHON_CLONE_AST_NUMBER_TYPES: &[&str] = &["integer", "float"];
43
44pub fn normalized_clone_tokens_python(source: &str) -> Vec<String> {
45    let Some(tree) = parse_clone_tree(source) else {
46        return Vec::new();
47    };
48    let mut out = Vec::new();
49    collect_normalized_leaf_tokens_python(tree.root_node(), source, &mut out);
50    out
51}
52
53fn collect_normalized_leaf_tokens_python(node: Node<'_>, source: &str, out: &mut Vec<String>) {
54    let mut stack = vec![node];
55    while let Some(node) = stack.pop() {
56        if node.named_child_count() == 0 {
57            let token = normalize_python_clone_leaf_token(node, source);
58            if !token.is_empty() {
59                out.push(token);
60            }
61        }
62        for index in (0..node.child_count()).rev() {
63            if let Some(child) = node.child(index) {
64                stack.push(child);
65            }
66        }
67    }
68}
69
70fn normalize_python_clone_leaf_token(node: Node<'_>, source: &str) -> String {
71    let kind = node.kind();
72    let token = source
73        .get(node.start_byte()..node.end_byte())
74        .unwrap_or("")
75        .trim();
76    if token.is_empty() || kind == "comment" {
77        return String::new();
78    }
79    if PYTHON_CLONE_AST_IDENTIFIER_TYPES.contains(&kind) {
80        return "ID".to_string();
81    }
82    if PYTHON_CLONE_AST_STRING_TYPES.contains(&kind) {
83        return "STR".to_string();
84    }
85    if PYTHON_CLONE_AST_NUMBER_TYPES.contains(&kind) {
86        return "NUM".to_string();
87    }
88    if kind == "true" || kind == "false" || token == "True" || token == "False" {
89        return "BOOL".to_string();
90    }
91    if token.chars().count() == 1 && token.chars().all(|ch| !ch.is_alphanumeric()) {
92        return format!("OP:{token}");
93    }
94    format!("T:{kind}")
95}
96
97pub fn build_python_clone_ast_signature(source: &str) -> String {
98    let Some(tree) = parse_clone_tree(source) else {
99        return String::new();
100    };
101    let mut labels = Vec::new();
102    collect_python_clone_ast_labels(tree.root_node(), source, &mut labels);
103    labels.join("|")
104}
105
106/// Parse and prepare one Python clone candidate, stopping before AST-label work
107/// when the normalized token threshold cannot be met.
108pub fn prepare_python_clone(
109    source: &str,
110    min_normalized_tokens: usize,
111) -> Option<PythonClonePreparation> {
112    let Some(tree) = parse_clone_tree(source) else {
113        return (min_normalized_tokens == 0).then(|| PythonClonePreparation {
114            normalized_tokens: Vec::new(),
115            ast_signature: String::new(),
116        });
117    };
118    let mut normalized_tokens = Vec::new();
119    collect_normalized_leaf_tokens_python(tree.root_node(), source, &mut normalized_tokens);
120    if normalized_tokens.len() < min_normalized_tokens {
121        return None;
122    }
123
124    let mut labels = Vec::new();
125    collect_python_clone_ast_labels(tree.root_node(), source, &mut labels);
126    Some(PythonClonePreparation {
127        normalized_tokens,
128        ast_signature: labels.join("|"),
129    })
130}
131
132fn collect_python_clone_ast_labels(node: Node<'_>, source: &str, out: &mut Vec<String>) {
133    let mut stack = vec![node];
134    while let Some(node) = stack.pop() {
135        out.push(normalize_python_clone_ast_label(node, source));
136        for index in (0..node.child_count()).rev() {
137            if let Some(child) = node.child(index) {
138                stack.push(child);
139            }
140        }
141    }
142}
143
144fn normalize_python_clone_ast_label(node: Node<'_>, source: &str) -> String {
145    let kind = node.kind();
146    let text = source
147        .get(node.start_byte()..node.end_byte())
148        .unwrap_or("")
149        .trim();
150    if PYTHON_CLONE_AST_IDENTIFIER_TYPES.contains(&kind) {
151        return "ID".to_string();
152    }
153    if PYTHON_CLONE_AST_STRING_TYPES.contains(&kind) {
154        return "STR".to_string();
155    }
156    if PYTHON_CLONE_AST_NUMBER_TYPES.contains(&kind) {
157        return "NUM".to_string();
158    }
159    if kind == "true" || kind == "false" || text == "True" || text == "False" {
160        return "BOOL".to_string();
161    }
162    format!("N:{kind}")
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168
169    fn legacy_tokens(node: Node<'_>, source: &str, out: &mut Vec<String>) {
170        if node.named_child_count() == 0 {
171            let token = normalize_python_clone_leaf_token(node, source);
172            if !token.is_empty() {
173                out.push(token);
174            }
175        }
176        for index in 0..node.child_count() {
177            if let Some(child) = node.child(index) {
178                legacy_tokens(child, source, out);
179            }
180        }
181    }
182
183    fn legacy_labels(node: Node<'_>, source: &str, out: &mut Vec<String>) {
184        out.push(normalize_python_clone_ast_label(node, source));
185        for index in 0..node.child_count() {
186            if let Some(child) = node.child(index) {
187                legacy_labels(child, source, out);
188            }
189        }
190    }
191
192    #[test]
193    fn combined_preparation_matches_the_existing_outputs_with_one_parse() {
194        let sources = [
195            r#"
196def compare(left, right):
197    # Keep comments in the AST signature, but not the token stream.
198    message = f"same: {left}"
199    ratio = 1.5 + 0x2A
200    return left == right and True
201"#,
202            "def broken(:\n    return False + 1\n",
203        ];
204
205        for source in sources {
206            let tree = parse_python_tree(source).unwrap();
207            if source.contains("broken") {
208                assert!(tree.root_node().has_error());
209            }
210            let mut expected_tokens = Vec::new();
211            legacy_tokens(tree.root_node(), source, &mut expected_tokens);
212            let mut expected_labels = Vec::new();
213            legacy_labels(tree.root_node(), source, &mut expected_labels);
214
215            reset_python_clone_parse_count_for_test();
216            let preparation = prepare_python_clone(source, 0).unwrap();
217
218            assert_eq!(preparation.normalized_tokens, expected_tokens);
219            assert_eq!(preparation.ast_signature, expected_labels.join("|"));
220            assert_eq!(python_clone_parse_count_for_test(), 1);
221            if source.contains("compare") {
222                for expected in ["STR", "NUM", "BOOL", "OP:="] {
223                    assert!(
224                        preparation
225                            .normalized_tokens
226                            .iter()
227                            .any(|token| token == expected),
228                        "missing normalized token {expected}"
229                    );
230                }
231                assert!(
232                    preparation
233                        .normalized_tokens
234                        .iter()
235                        .all(|token| token != "T:comment")
236                );
237                assert!(
238                    preparation
239                        .ast_signature
240                        .split('|')
241                        .any(|label| label == "N:comment")
242                );
243            }
244        }
245    }
246
247    #[test]
248    fn rejected_candidate_still_uses_one_parse() {
249        let token_count = prepare_python_clone("value = 1", 0)
250            .unwrap()
251            .normalized_tokens
252            .len();
253        reset_python_clone_parse_count_for_test();
254
255        assert!(prepare_python_clone("value = 1", token_count).is_some());
256        assert_eq!(python_clone_parse_count_for_test(), 1);
257        reset_python_clone_parse_count_for_test();
258
259        assert!(prepare_python_clone("value = 1", token_count + 1).is_none());
260        assert_eq!(python_clone_parse_count_for_test(), 1);
261    }
262
263    #[test]
264    fn deeply_nested_source_is_walked_without_recursion() {
265        let depth = 2_000;
266        let source = format!("value = {}0{}\n", "(".repeat(depth), ")".repeat(depth));
267
268        let preparation = prepare_python_clone(&source, 0).unwrap();
269
270        assert_eq!(
271            preparation
272                .normalized_tokens
273                .iter()
274                .filter(|token| token.as_str() == "OP:(")
275                .count(),
276            depth
277        );
278        assert_eq!(
279            preparation
280                .normalized_tokens
281                .iter()
282                .filter(|token| token.as_str() == "OP:)")
283                .count(),
284            depth
285        );
286        assert_eq!(preparation.ast_signature.matches("N:(").count(), depth);
287        assert_eq!(preparation.ast_signature.matches("N:)").count(), depth);
288    }
289}