brokk_bifrost_python/
clones.rs1use 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
106pub 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}