use std::collections::HashMap;
use crate::tokens::count_tokens;
fn normalize_shell_tokens(text: &str) -> String {
text.split_whitespace().collect::<Vec<_>>().join(" ")
}
#[derive(Clone, Copy)]
struct Entry<'a> {
line_num: usize,
content: &'a str,
is_match: bool,
}
pub fn compress(output: &str) -> Option<String> {
let lines: Vec<&str> = output.lines().collect();
if lines.len() < 3 {
return None;
}
let mut by_file: HashMap<&str, Vec<Entry<'_>>> = HashMap::new();
let mut total_matches = 0usize;
for line in &lines {
if let Some((file, line_num, content)) = parse_match_line(line) {
total_matches += 1;
by_file.entry(file).or_default().push(Entry {
line_num,
content,
is_match: true,
});
}
}
if total_matches == 0 {
return None;
}
let known: Vec<&str> = by_file.keys().copied().collect();
for line in &lines {
if parse_match_line(line).is_some() {
continue;
}
if let Some((file, line_num, content)) = parse_context_line(line, &known) {
by_file.entry(file).or_default().push(Entry {
line_num,
content,
is_match: false,
});
}
}
let max_matches_per_file = if total_matches > 200 { 5 } else { 10 };
let mut result = format!("{total_matches} matches in {}F:\n", by_file.len());
let mut sorted_files: Vec<_> = by_file.iter_mut().collect();
sorted_files.sort_by(|a, b| {
let am = a.1.iter().filter(|e| e.is_match).count();
let bm = b.1.iter().filter(|e| e.is_match).count();
bm.cmp(&am).then_with(|| a.0.cmp(b.0))
});
for (file, entries) in &mut sorted_files {
entries.sort_by_key(|e| (e.line_num, !e.is_match));
let match_count = entries.iter().filter(|e| e.is_match).count();
let short = shorten_path(file);
result.push_str(&format!("\n{short} ({match_count}):"));
let mut shown_matches = 0usize;
for entry in entries.iter() {
if entry.is_match {
if shown_matches == max_matches_per_file {
break;
}
shown_matches += 1;
}
let trimmed = entry.content.trim();
let short_content = if trimmed.len() > 120 {
let truncated: String = trimmed.chars().take(119).collect();
format!("{truncated}…")
} else {
trimmed.to_string()
};
let sep = if entry.is_match { ':' } else { '-' };
if entry.line_num > 0 {
result.push_str(&format!("\n {}{sep} {short_content}", entry.line_num));
} else {
result.push_str(&format!("\n {short_content}"));
}
}
if match_count > shown_matches {
result.push_str(&format!("\n ... +{} more", match_count - shown_matches));
}
}
let out_n = normalize_shell_tokens(output);
let res_n = normalize_shell_tokens(&result);
let ct_r = count_tokens(&res_n);
let ct_o = count_tokens(&out_n);
if ct_r >= ct_o && !(ct_r == ct_o && res_n.len() < out_n.len()) {
return None;
}
Some(result)
}
fn parse_match_line(line: &str) -> Option<(&str, usize, &str)> {
let start = if line.len() >= 2
&& line.as_bytes()[0].is_ascii_alphabetic()
&& line.as_bytes()[1] == b':'
{
2
} else {
0
};
let (file, line_num, content) = split_numbered(line, ':', start)?;
if file.contains('/') || file.contains('\\') || file.contains('.') {
Some((file, line_num, content))
} else {
None
}
}
fn split_numbered(line: &str, delim: char, start: usize) -> Option<(&str, usize, &str)> {
let bytes = line.as_bytes();
let mut from = start;
while let Some(rel) = line[from..].find(delim) {
let at = from + rel;
let after = at + delim.len_utf8();
let digits = bytes[after..]
.iter()
.take_while(|b| b.is_ascii_digit())
.count();
if digits > 0 && bytes.get(after + digits).copied() == Some(delim as u8) {
let num = line[after..after + digits].parse().ok()?;
return Some((&line[..at], num, &line[after + digits + 1..]));
}
from = at + delim.len_utf8();
}
None
}
fn parse_context_line<'a>(line: &'a str, known: &[&'a str]) -> Option<(&'a str, usize, &'a str)> {
for file in known {
let Some(rest) = line.strip_prefix(*file).and_then(|r| r.strip_prefix('-')) else {
continue;
};
let digits = rest
.as_bytes()
.iter()
.take_while(|b| b.is_ascii_digit())
.count();
if digits > 0 && rest.as_bytes().get(digits).copied() == Some(b'-') {
let num = rest[..digits].parse().ok()?;
return Some((file, num, &rest[digits + 1..]));
}
}
None
}
fn shorten_path(path: &str) -> &str {
path.strip_prefix("./").unwrap_or(path)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn small_grep_output_is_not_claimed_without_matches() {
assert!(compress("hello\nworld").is_none());
}
#[test]
fn small_grep_output_still_compresses() {
let output = (0..20)
.map(|i| format!("src/main.rs:{i}: let x = {i};"))
.collect::<Vec<_>>()
.join("\n");
let result = compress(&output);
assert!(result.is_some());
let compressed = result.unwrap();
assert!(
compressed.contains("20 matches in 1F:"),
"should group by file: {compressed}"
);
assert!(
count_tokens(&compressed) < count_tokens(&output),
"should compress: {} vs {}",
count_tokens(&compressed),
count_tokens(&output)
);
}
#[test]
fn large_output_reduces_per_file_lines() {
let mut lines = Vec::new();
for i in 0..250 {
lines.push(format!("src/a.rs:{i}: line content {i}"));
}
let output = lines.join("\n");
let result = compress(&output).unwrap();
assert!(
result.contains("... +245 more"),
"should show +more for large output: {result}"
);
}
#[test]
fn non_grep_output_returns_none() {
let output = "no file:line pattern here\njust regular text\nmore text\nand more";
assert!(compress(output).is_none());
}
#[test]
fn tiny_grep_output_returns_none_if_inflation() {
let output = "a.rs:1:x\nb.rs:2:y\nc.rs:3:z\n";
let result = compress(output);
if let Some(ref compressed) = result {
assert!(
count_tokens(compressed) < count_tokens(output),
"must never inflate: compressed={} vs original={}",
count_tokens(compressed),
count_tokens(output)
);
}
}
#[test]
fn multi_file_many_matches_compresses_well() {
let mut lines = Vec::new();
for i in 0..50 {
lines.push(format!(
"src/models/user.rs:{}: pub fn method_{i}() {{}}",
i + 1
));
}
for i in 0..30 {
lines.push(format!(
"src/controllers/auth.rs:{}: let val = method_{i}();",
i + 1
));
}
let output = lines.join("\n");
let result = compress(&output).expect("80 matches should compress");
assert!(
count_tokens(&result) < count_tokens(&output),
"must compress: {} vs {}",
count_tokens(&result),
count_tokens(&output)
);
assert!(result.contains("80 matches in 2F:"));
assert!(result.contains("src/models/user.rs (50):"));
assert!(result.contains("src/controllers/auth.rs (30):"));
}
#[test]
fn many_single_match_files_falls_back_to_none() {
let lines: Vec<String> = (1..=30)
.map(|i| format!("src/file{i}.rs:42: fn search_result()"))
.collect();
let output = lines.join("\n");
let result = compress(&output);
if let Some(ref c) = result {
assert!(
count_tokens(c) < count_tokens(&output),
"if claimed, must be shorter in tokens: {} vs {}",
count_tokens(c),
count_tokens(&output)
);
}
}
#[test]
fn never_returns_inflated_output() {
for count in [3, 5, 10, 15, 25, 50] {
let lines: Vec<String> = (0..count).map(|i| format!("f{i}.rs:{i}:x")).collect();
let output = lines.join("\n");
if let Some(ref c) = compress(&output) {
assert!(
count_tokens(c) < count_tokens(&output),
"count={count}: inflated {} vs {}",
count_tokens(c),
count_tokens(&output)
);
}
}
}
fn repro_with_context(files: usize) -> String {
let mut out = Vec::new();
for i in 1..=files {
let f = format!("/tmp/leanctx-grep-repro/mod_{i}.py");
out.push(format!("{f}-1-def handler_{i}():"));
out.push(format!("{f}-2- try:"));
out.push(format!("{f}:3: target_pattern_here()"));
out.push(format!("{f}-4- except Exception:"));
out.push(format!("{f}-5- pass"));
out.push("--".to_string());
}
out.join("\n")
}
#[test]
fn context_lines_do_not_become_files_or_matches() {
let result = compress(&repro_with_context(40)).expect("should compress");
assert!(
result.starts_with("40 matches in 40F:"),
"40 real matches in 40 files, not one per context line: {}",
result.lines().next().unwrap_or("")
);
assert!(
!result.contains("-2-"),
"a context line must not appear as a file header: {result}"
);
assert!(
!result.contains("try (1)"),
"text must not be truncated at a colon into a path: {result}"
);
}
#[test]
fn context_text_survives_intact() {
let result = compress(&repro_with_context(40)).expect("should compress");
assert!(
result.contains("try:"),
"the context line's own text, colon included, must survive: {result}"
);
}
#[test]
fn group_separators_are_ignored() {
let out = (1..=20)
.map(|i| format!("a/b.rs-{}-ctx\na/b.rs:{i}:hit\n--", i * 10))
.collect::<Vec<_>>()
.join("\n");
let result = compress(&out).expect("should compress");
assert!(result.starts_with("20 matches in 1F:"), "{result}");
}
#[test]
fn a_hyphenated_path_is_not_split_as_context() {
let out = (1..=20)
.map(|i| format!("src/v-1-x/mod.py-{}-ctx\nsrc/v-1-x/mod.py:{i}:hit", i * 10))
.collect::<Vec<_>>()
.join("\n");
let result = compress(&out).expect("should compress");
assert!(
result.starts_with("20 matches in 1F:"),
"one file, not one per hyphen: {result}"
);
}
}