lean_ctx_patterns/
grep.rs1use std::collections::HashMap;
4
5use crate::tokens::count_tokens;
6
7fn normalize_shell_tokens(text: &str) -> String {
8 text.split_whitespace().collect::<Vec<_>>().join(" ")
9}
10
11#[derive(Clone, Copy)]
13struct Entry<'a> {
14 line_num: usize,
15 content: &'a str,
16 is_match: bool,
17}
18
19pub fn compress(output: &str) -> Option<String> {
20 let lines: Vec<&str> = output.lines().collect();
21 if lines.len() < 3 {
22 return None;
23 }
24
25 let mut by_file: HashMap<&str, Vec<Entry<'_>>> = HashMap::new();
26 let mut total_matches = 0usize;
27
28 for line in &lines {
30 if let Some((file, line_num, content)) = parse_match_line(line) {
31 total_matches += 1;
32 by_file.entry(file).or_default().push(Entry {
33 line_num,
34 content,
35 is_match: true,
36 });
37 }
38 }
39
40 if total_matches == 0 {
41 return None;
42 }
43
44 let known: Vec<&str> = by_file.keys().copied().collect();
52 for line in &lines {
53 if parse_match_line(line).is_some() {
54 continue;
55 }
56 if let Some((file, line_num, content)) = parse_context_line(line, &known) {
57 by_file.entry(file).or_default().push(Entry {
58 line_num,
59 content,
60 is_match: false,
61 });
62 }
63 }
67
68 let max_matches_per_file = if total_matches > 200 { 5 } else { 10 };
69
70 let mut result = format!("{total_matches} matches in {}F:\n", by_file.len());
71 let mut sorted_files: Vec<_> = by_file.iter_mut().collect();
72 sorted_files.sort_by(|a, b| {
73 let am = a.1.iter().filter(|e| e.is_match).count();
74 let bm = b.1.iter().filter(|e| e.is_match).count();
75 bm.cmp(&am).then_with(|| a.0.cmp(b.0))
76 });
77
78 for (file, entries) in &mut sorted_files {
79 entries.sort_by_key(|e| (e.line_num, !e.is_match));
80 let match_count = entries.iter().filter(|e| e.is_match).count();
81 let short = shorten_path(file);
82 result.push_str(&format!("\n{short} ({match_count}):"));
83
84 let mut shown_matches = 0usize;
88 for entry in entries.iter() {
89 if entry.is_match {
90 if shown_matches == max_matches_per_file {
91 break;
92 }
93 shown_matches += 1;
94 }
95 let trimmed = entry.content.trim();
96 let short_content = if trimmed.len() > 120 {
97 let truncated: String = trimmed.chars().take(119).collect();
98 format!("{truncated}…")
99 } else {
100 trimmed.to_string()
101 };
102 let sep = if entry.is_match { ':' } else { '-' };
104 if entry.line_num > 0 {
105 result.push_str(&format!("\n {}{sep} {short_content}", entry.line_num));
106 } else {
107 result.push_str(&format!("\n {short_content}"));
108 }
109 }
110 if match_count > shown_matches {
111 result.push_str(&format!("\n ... +{} more", match_count - shown_matches));
112 }
113 }
114
115 let out_n = normalize_shell_tokens(output);
116 let res_n = normalize_shell_tokens(&result);
117 let ct_r = count_tokens(&res_n);
118 let ct_o = count_tokens(&out_n);
119 if ct_r >= ct_o && !(ct_r == ct_o && res_n.len() < out_n.len()) {
120 return None;
121 }
122
123 Some(result)
124}
125
126fn parse_match_line(line: &str) -> Option<(&str, usize, &str)> {
133 let start = if line.len() >= 2
135 && line.as_bytes()[0].is_ascii_alphabetic()
136 && line.as_bytes()[1] == b':'
137 {
138 2
139 } else {
140 0
141 };
142
143 let (file, line_num, content) = split_numbered(line, ':', start)?;
144 if file.contains('/') || file.contains('\\') || file.contains('.') {
145 Some((file, line_num, content))
146 } else {
147 None
148 }
149}
150
151fn split_numbered(line: &str, delim: char, start: usize) -> Option<(&str, usize, &str)> {
154 let bytes = line.as_bytes();
155 let mut from = start;
156 while let Some(rel) = line[from..].find(delim) {
157 let at = from + rel;
158 let after = at + delim.len_utf8();
159 let digits = bytes[after..]
160 .iter()
161 .take_while(|b| b.is_ascii_digit())
162 .count();
163 if digits > 0 && bytes.get(after + digits).copied() == Some(delim as u8) {
164 let num = line[after..after + digits].parse().ok()?;
165 return Some((&line[..at], num, &line[after + digits + 1..]));
166 }
167 from = at + delim.len_utf8();
168 }
169 None
170}
171
172fn parse_context_line<'a>(line: &'a str, known: &[&'a str]) -> Option<(&'a str, usize, &'a str)> {
179 for file in known {
180 let Some(rest) = line.strip_prefix(*file).and_then(|r| r.strip_prefix('-')) else {
181 continue;
182 };
183 let digits = rest
184 .as_bytes()
185 .iter()
186 .take_while(|b| b.is_ascii_digit())
187 .count();
188 if digits > 0 && rest.as_bytes().get(digits).copied() == Some(b'-') {
189 let num = rest[..digits].parse().ok()?;
190 return Some((file, num, &rest[digits + 1..]));
191 }
192 }
193 None
194}
195
196fn shorten_path(path: &str) -> &str {
197 path.strip_prefix("./").unwrap_or(path)
198}
199
200#[cfg(test)]
201mod tests {
202 use super::*;
203
204 #[test]
205 fn small_grep_output_is_not_claimed_without_matches() {
206 assert!(compress("hello\nworld").is_none());
207 }
208
209 #[test]
210 fn small_grep_output_still_compresses() {
211 let output = (0..20)
212 .map(|i| format!("src/main.rs:{i}: let x = {i};"))
213 .collect::<Vec<_>>()
214 .join("\n");
215 let result = compress(&output);
216 assert!(result.is_some());
217 let compressed = result.unwrap();
218 assert!(
219 compressed.contains("20 matches in 1F:"),
220 "should group by file: {compressed}"
221 );
222 assert!(
223 count_tokens(&compressed) < count_tokens(&output),
224 "should compress: {} vs {}",
225 count_tokens(&compressed),
226 count_tokens(&output)
227 );
228 }
229
230 #[test]
231 fn large_output_reduces_per_file_lines() {
232 let mut lines = Vec::new();
233 for i in 0..250 {
234 lines.push(format!("src/a.rs:{i}: line content {i}"));
235 }
236 let output = lines.join("\n");
237 let result = compress(&output).unwrap();
238 assert!(
239 result.contains("... +245 more"),
240 "should show +more for large output: {result}"
241 );
242 }
243
244 #[test]
245 fn non_grep_output_returns_none() {
246 let output = "no file:line pattern here\njust regular text\nmore text\nand more";
247 assert!(compress(output).is_none());
248 }
249
250 #[test]
251 fn tiny_grep_output_returns_none_if_inflation() {
252 let output = "a.rs:1:x\nb.rs:2:y\nc.rs:3:z\n";
253 let result = compress(output);
254 if let Some(ref compressed) = result {
255 assert!(
256 count_tokens(compressed) < count_tokens(output),
257 "must never inflate: compressed={} vs original={}",
258 count_tokens(compressed),
259 count_tokens(output)
260 );
261 }
262 }
263
264 #[test]
265 fn multi_file_many_matches_compresses_well() {
266 let mut lines = Vec::new();
267 for i in 0..50 {
268 lines.push(format!(
269 "src/models/user.rs:{}: pub fn method_{i}() {{}}",
270 i + 1
271 ));
272 }
273 for i in 0..30 {
274 lines.push(format!(
275 "src/controllers/auth.rs:{}: let val = method_{i}();",
276 i + 1
277 ));
278 }
279 let output = lines.join("\n");
280 let result = compress(&output).expect("80 matches should compress");
281 assert!(
282 count_tokens(&result) < count_tokens(&output),
283 "must compress: {} vs {}",
284 count_tokens(&result),
285 count_tokens(&output)
286 );
287 assert!(result.contains("80 matches in 2F:"));
288 assert!(result.contains("src/models/user.rs (50):"));
289 assert!(result.contains("src/controllers/auth.rs (30):"));
290 }
291
292 #[test]
293 fn many_single_match_files_falls_back_to_none() {
294 let lines: Vec<String> = (1..=30)
295 .map(|i| format!("src/file{i}.rs:42: fn search_result()"))
296 .collect();
297 let output = lines.join("\n");
298 let result = compress(&output);
299 if let Some(ref c) = result {
300 assert!(
301 count_tokens(c) < count_tokens(&output),
302 "if claimed, must be shorter in tokens: {} vs {}",
303 count_tokens(c),
304 count_tokens(&output)
305 );
306 }
307 }
308
309 #[test]
310 fn never_returns_inflated_output() {
311 for count in [3, 5, 10, 15, 25, 50] {
312 let lines: Vec<String> = (0..count).map(|i| format!("f{i}.rs:{i}:x")).collect();
313 let output = lines.join("\n");
314 if let Some(ref c) = compress(&output) {
315 assert!(
316 count_tokens(c) < count_tokens(&output),
317 "count={count}: inflated {} vs {}",
318 count_tokens(c),
319 count_tokens(&output)
320 );
321 }
322 }
323 }
324
325 fn repro_with_context(files: usize) -> String {
333 let mut out = Vec::new();
334 for i in 1..=files {
335 let f = format!("/tmp/leanctx-grep-repro/mod_{i}.py");
336 out.push(format!("{f}-1-def handler_{i}():"));
337 out.push(format!("{f}-2- try:"));
338 out.push(format!("{f}:3: target_pattern_here()"));
339 out.push(format!("{f}-4- except Exception:"));
340 out.push(format!("{f}-5- pass"));
341 out.push("--".to_string());
342 }
343 out.join("\n")
344 }
345
346 #[test]
347 fn context_lines_do_not_become_files_or_matches() {
348 let result = compress(&repro_with_context(40)).expect("should compress");
349 assert!(
350 result.starts_with("40 matches in 40F:"),
351 "40 real matches in 40 files, not one per context line: {}",
352 result.lines().next().unwrap_or("")
353 );
354 assert!(
355 !result.contains("-2-"),
356 "a context line must not appear as a file header: {result}"
357 );
358 assert!(
359 !result.contains("try (1)"),
360 "text must not be truncated at a colon into a path: {result}"
361 );
362 }
363
364 #[test]
367 fn context_text_survives_intact() {
368 let result = compress(&repro_with_context(40)).expect("should compress");
369 assert!(
370 result.contains("try:"),
371 "the context line's own text, colon included, must survive: {result}"
372 );
373 }
374
375 #[test]
377 fn group_separators_are_ignored() {
378 let out = (1..=20)
379 .map(|i| format!("a/b.rs-{}-ctx\na/b.rs:{i}:hit\n--", i * 10))
380 .collect::<Vec<_>>()
381 .join("\n");
382 let result = compress(&out).expect("should compress");
383 assert!(result.starts_with("20 matches in 1F:"), "{result}");
384 }
385
386 #[test]
390 fn a_hyphenated_path_is_not_split_as_context() {
391 let out = (1..=20)
392 .map(|i| format!("src/v-1-x/mod.py-{}-ctx\nsrc/v-1-x/mod.py:{i}:hit", i * 10))
393 .collect::<Vec<_>>()
394 .join("\n");
395 let result = compress(&out).expect("should compress");
396 assert!(
397 result.starts_with("20 matches in 1F:"),
398 "one file, not one per hyphen: {result}"
399 );
400 }
401}