Skip to main content

lean_ctx_patterns/
grep.rs

1// SPDX-License-Identifier: Apache-2.0
2
3use 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/// One rendered line: a match, or a context line around one.
12#[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    // Pass 1 — matches only (`path:line:content`).
29    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    // Pass 2 — context (`path-line-content`), resolved against the paths pass 1
45    // already proved exist (GH #1648).
46    //
47    // Doing it this way, rather than scanning for a `-digits-` delimiter, is
48    // what makes it unambiguous: both a path and the context text may contain
49    // hyphens and digits, but grep only emits context for a file it also
50    // matched in, so the path is always one we have already seen.
51    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        // Anything else — grep's `--` group separators, stray text — is
64        // dropped. It used to become a synthetic file with empty content,
65        // inflating both the file and the match count.
66    }
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        // The cap counts *matches*, and any context around a shown match comes
85        // with it — capping raw lines would silently drop matches once context
86        // was requested.
87        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            // grep's own convention: `:` marks a match, `-` marks context.
103            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
126/// Parse a match line: `path:line:content`.
127///
128/// The delimiter is the first colon *followed by digits and another colon* —
129/// not simply the first colon on the line (GH #1648). The old rule split
130/// `…/mod_1.py-2-    try:` at the colon inside the Python source, inventing a
131/// file named after the whole fragment.
132fn parse_match_line(line: &str) -> Option<(&str, usize, &str)> {
133    // Skip a Windows drive letter (e.g. "C:" at position 1).
134    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
151/// Split `prefix<D><digits><D>rest` at the first delimiter that is actually
152/// followed by a line number.
153fn 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
172/// Parse a context line: `path-line-content`, where `path` is one grep already
173/// reported a match in.
174///
175/// Anchoring on a known path rather than scanning for `-digits-` is what keeps
176/// this unambiguous: `src/v-1-x/mod.py-10-ctx` has three candidate delimiters
177/// and only one correct answer.
178fn 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    // --- GH #1648: context lines (`path-line-text`) are not matches ---
326
327    /// The reporter's repro, reduced: 40 files, one match each, `-C 2`.
328    /// GNU grep separates a match with `:` and context with `-`; the parser
329    /// took the FIRST colon anywhere on the line, so `…/mod_1.py-2-    try:`
330    /// became a *file* named `…/mod_1.py-2-    try` with empty content —
331    /// counted as both a file and a match.
332    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    /// The content the reporter lost. A context line whose text contains a
365    /// colon (Python, YAML, Go tags, prose) is exactly the case that broke.
366    #[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    /// `--` group separators are grep's, not data.
376    #[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    /// A path containing `-<digits>-` must not be mistaken for a context
387    /// delimiter. Context is resolved against paths already seen in match
388    /// lines, so the ambiguity never arises.
389    #[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}