use super::{matcher::build_matcher, push_grep_match, reader};
use crate::{
agent::cancellation::AgentCancellation,
tools::workspace::{WorkspaceWalkOptions, WorkspaceWalker},
};
use crossbeam_channel as channel;
use grep_searcher::{BinaryDetection, Searcher, SearcherBuilder, Sink, SinkContext, SinkMatch};
use std::{
io,
path::{Path, PathBuf},
sync::{
Arc,
atomic::{AtomicBool, AtomicUsize, Ordering},
},
};
pub(super) const ORDERED_STREAMING_STOP_MAX_COUNT: usize = 64;
pub(super) const GREP_STREAM_WINDOW: usize = 512;
#[derive(Debug, Default, Clone)]
pub(super) struct GrepScan {
pub(super) matches: Vec<String>,
pub(super) files_scanned: usize,
pub(super) bytes_scanned: u64,
pub(super) truncated: bool,
pub(super) timed_out: bool,
pub(super) limit_reason: Option<&'static str>,
}
pub(super) struct GrepRequest<'a> {
pub(super) cwd: &'a Path,
pub(super) walker: &'a WorkspaceWalker,
pub(super) pattern: &'a str,
pub(super) root: &'a Path,
pub(super) raw_limit: Option<usize>,
pub(super) context: usize,
pub(super) cancellation: &'a AgentCancellation,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Strategy {
Parallel,
Sequential,
Windowed,
}
pub(super) fn run(request: GrepRequest<'_>) -> anyhow::Result<GrepScan> {
let matcher = Arc::new(build_matcher(request.pattern)?);
let files = files_for_request(request.walker, request.root, request.cancellation)?;
let limit = request.raw_limit.unwrap_or(100);
let strategy = select_strategy(request.raw_limit);
match strategy {
Strategy::Parallel => run_parallel(
&files,
request.cwd,
matcher,
limit,
request.context,
request.cancellation,
),
Strategy::Sequential => run_sequential(
&files,
request.cwd,
&matcher,
limit,
request.context,
request.cancellation,
),
Strategy::Windowed => run_windowed(
&files,
request.cwd,
matcher,
limit,
request.context,
request.cancellation,
),
}
}
fn files_for_request(
walker: &WorkspaceWalker,
root: &Path,
cancellation: &AgentCancellation,
) -> anyhow::Result<Vec<PathBuf>> {
let Ok(metadata) = root.symlink_metadata() else {
return Ok(Vec::new());
};
if metadata.file_type().is_file() {
return Ok(vec![root.to_path_buf()]);
}
if !metadata.file_type().is_dir() {
return Ok(Vec::new());
}
let walk = walker.walk(
WorkspaceWalkOptions {
root,
include_files: true,
include_dirs: false,
skip_dirs: &[".git", "target"],
cancel_interval: 32,
},
Some(cancellation),
)?;
let mut files = walk
.entries
.into_iter()
.filter(|entry| entry.file_type.is_file())
.map(|entry| entry.path)
.collect::<Vec<_>>();
files.sort();
Ok(files)
}
fn select_strategy(raw_limit: Option<usize>) -> Strategy {
match raw_limit {
None => Strategy::Parallel,
Some(1..=ORDERED_STREAMING_STOP_MAX_COUNT) => Strategy::Sequential,
Some(_) => Strategy::Windowed,
}
}
fn run_sequential(
files: &[PathBuf],
cwd: &Path,
matcher: &grep_regex::RegexMatcher,
limit: usize,
context: usize,
cancellation: &AgentCancellation,
) -> anyhow::Result<GrepScan> {
let mut scan = GrepScan::default();
for file in files {
cancellation.check()?;
if scan.matches.len() >= limit {
scan.truncated = true;
scan.limit_reason = Some("matches");
break;
}
let remaining = limit - scan.matches.len();
let result = scan_file(file, cwd, matcher, remaining, context, cancellation, None)?;
append_file_scan(&mut scan, result, limit);
if scan.matches.len() >= limit {
break;
}
}
Ok(scan)
}
fn run_windowed(
files: &[PathBuf],
cwd: &Path,
matcher: Arc<grep_regex::RegexMatcher>,
limit: usize,
context: usize,
cancellation: &AgentCancellation,
) -> anyhow::Result<GrepScan> {
let mut scan = GrepScan::default();
for window in files.chunks(GREP_STREAM_WINDOW) {
cancellation.check()?;
let stop = AtomicBool::new(false);
let mut results = scan_parallel_window(
window,
cwd,
matcher.clone(),
limit,
context,
cancellation,
ParallelStop {
stop: &stop,
accepted_matches: None,
},
)?;
results.sort_by_key(|(index, _)| *index);
for (_, result) in results {
append_file_scan(&mut scan, result, limit);
if scan.matches.len() >= limit {
return Ok(scan);
}
}
}
Ok(scan)
}
fn run_parallel(
files: &[PathBuf],
cwd: &Path,
matcher: Arc<grep_regex::RegexMatcher>,
limit: usize,
context: usize,
cancellation: &AgentCancellation,
) -> anyhow::Result<GrepScan> {
let stop = AtomicBool::new(false);
let accepted_matches = AtomicUsize::new(0);
let mut results = scan_parallel_window(
files,
cwd,
matcher,
limit,
context,
cancellation,
ParallelStop {
stop: &stop,
accepted_matches: Some(&accepted_matches),
},
)?;
let mut scan = GrepScan::default();
for (_, result) in results.drain(..) {
append_file_scan(&mut scan, result, limit);
if scan.matches.len() >= limit {
break;
}
}
Ok(scan)
}
struct ParallelStop<'a> {
stop: &'a AtomicBool,
accepted_matches: Option<&'a AtomicUsize>,
}
fn scan_parallel_window(
files: &[PathBuf],
cwd: &Path,
matcher: Arc<grep_regex::RegexMatcher>,
limit: usize,
context: usize,
cancellation: &AgentCancellation,
state: ParallelStop<'_>,
) -> anyhow::Result<Vec<(usize, FileScan)>> {
if files.is_empty() {
return Ok(Vec::new());
}
let workers = std::thread::available_parallelism()
.map_or(1, usize::from)
.min(files.len())
.max(1);
let chunk_size = files.len().div_ceil(workers);
let (tx, rx) = channel::bounded(files.len().clamp(1, GREP_STREAM_WINDOW));
let stop = state.stop;
let accepted_matches = state.accepted_matches;
std::thread::scope(|scope| {
for (chunk_index, chunk) in files.chunks(chunk_size).enumerate() {
let tx = tx.clone();
let matcher = matcher.clone();
scope.spawn(move || {
for (offset, file) in chunk.iter().enumerate() {
if stop.load(Ordering::Relaxed) {
break;
}
let index = chunk_index * chunk_size + offset;
let result = scan_file(
file,
cwd,
&matcher,
limit,
context,
cancellation,
Some(stop),
);
let should_stop = result.as_ref().is_ok_and(|scan| {
if let Some(accepted_matches) = accepted_matches {
if scan.matches.is_empty() {
return false;
}
let accepted = accepted_matches
.fetch_add(scan.matches.len(), Ordering::Relaxed)
+ scan.matches.len();
accepted >= limit
} else {
scan.limit_reached
}
});
if should_stop {
stop.store(true, Ordering::Relaxed);
}
if tx.send((index, result)).is_err() {
break;
}
if should_stop {
stop.store(true, Ordering::Relaxed);
break;
}
}
});
}
drop(tx);
let mut results = Vec::new();
let mut first_error = None;
for (index, result) in rx {
match result {
Ok(file_scan) => {
if file_scan.matches.len() >= limit {
stop.store(true, Ordering::Relaxed);
}
results.push((index, file_scan));
}
Err(error) => {
stop.store(true, Ordering::Relaxed);
first_error.get_or_insert(error);
}
}
}
if let Some(error) = first_error {
Err(error)
} else {
Ok(results)
}
})
}
fn append_file_scan(scan: &mut GrepScan, mut file: FileScan, limit: usize) {
scan.files_scanned += file.files_scanned;
scan.bytes_scanned = scan.bytes_scanned.saturating_add(file.bytes_scanned);
if file.read_truncated {
scan.truncated = true;
scan.limit_reason.get_or_insert("bytes");
}
let remaining = limit.saturating_sub(scan.matches.len());
if file.matches.len() > remaining {
file.matches.truncate(remaining);
file.limit_reached = true;
}
scan.matches.extend(file.matches);
if file.limit_reached || scan.matches.len() >= limit {
scan.truncated = true;
scan.limit_reason = Some("matches");
}
}
#[derive(Debug, Default)]
struct FileScan {
matches: Vec<String>,
files_scanned: usize,
bytes_scanned: u64,
read_truncated: bool,
limit_reached: bool,
}
fn scan_file(
file: &Path,
cwd: &Path,
matcher: &grep_regex::RegexMatcher,
limit: usize,
context: usize,
cancellation: &AgentCancellation,
stop: Option<&AtomicBool>,
) -> anyhow::Result<FileScan> {
cancellation.check()?;
let reader::ReadOutcome::Bytes {
bytes,
bytes_scanned,
truncated,
..
} = reader::read_file_bytes(file)?
else {
return Ok(FileScan::default());
};
let mut scan = FileScan {
files_scanned: 1,
bytes_scanned,
read_truncated: truncated,
..FileScan::default()
};
let mut searcher = SearcherBuilder::new()
.binary_detection(BinaryDetection::quit(b'\x00'))
.line_number(true)
.multi_line(false)
.before_context(context)
.after_context(context)
.build();
let mut collector = MatchCollector {
cwd,
file,
matches: Vec::new(),
limit,
limit_reached: false,
cancellation,
stop,
};
searcher
.search_slice(matcher, &bytes, &mut collector)
.map_err(|error| anyhow::anyhow!(error))?;
scan.matches = collector.matches;
scan.limit_reached = collector.limit_reached;
Ok(scan)
}
struct MatchCollector<'a> {
cwd: &'a Path,
file: &'a Path,
matches: Vec<String>,
limit: usize,
limit_reached: bool,
cancellation: &'a AgentCancellation,
stop: Option<&'a AtomicBool>,
}
impl MatchCollector<'_> {
fn push(&mut self, line_no: Option<u64>, bytes: &[u8]) -> io::Result<bool> {
self.cancellation
.check()
.map_err(|error| io::Error::new(io::ErrorKind::Interrupted, error.to_string()))?;
if self.stop.is_some_and(|stop| stop.load(Ordering::Relaxed)) {
return Ok(false);
}
let Ok(line) = std::str::from_utf8(bytes) else {
return Ok(false);
};
push_grep_match(
&mut self.matches,
self.cwd,
self.file,
line_no.unwrap_or(0) as usize,
line,
);
if self.matches.len() >= self.limit {
self.limit_reached = true;
if let Some(stop) = self.stop {
stop.store(true, Ordering::Relaxed);
}
return Ok(false);
}
Ok(true)
}
}
impl Sink for MatchCollector<'_> {
type Error = io::Error;
fn matched(&mut self, _searcher: &Searcher, mat: &SinkMatch<'_>) -> Result<bool, Self::Error> {
self.push(mat.line_number(), mat.bytes())
}
fn context(
&mut self,
_searcher: &Searcher,
context: &SinkContext<'_>,
) -> Result<bool, Self::Error> {
self.push(context.line_number(), context.bytes())
}
fn binary_data(
&mut self,
_searcher: &Searcher,
_binary_byte_offset: u64,
) -> Result<bool, Self::Error> {
Ok(false)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::tools::workspace::WorkspaceWalker;
use std::fs;
fn run_temp(temp: &tempfile::TempDir, limit: Option<usize>) -> GrepScan {
run(GrepRequest {
cwd: temp.path(),
walker: &WorkspaceWalker,
pattern: "needle",
root: temp.path(),
raw_limit: limit,
context: 0,
cancellation: &AgentCancellation::default(),
})
.unwrap()
}
#[test]
fn selector_chooses_required_strategies() {
assert_eq!(select_strategy(None), Strategy::Parallel);
assert_eq!(select_strategy(Some(1)), Strategy::Sequential);
assert_eq!(select_strategy(Some(64)), Strategy::Sequential);
assert_eq!(select_strategy(Some(65)), Strategy::Windowed);
}
#[test]
fn sequential_limit_one_stops_early_in_order() {
let temp = tempfile::TempDir::new().unwrap();
fs::write(temp.path().join("001.txt"), "needle\n").unwrap();
fs::write(temp.path().join("002.txt"), "needle\n").unwrap();
fs::write(temp.path().join("003.txt"), "needle\n").unwrap();
let scan = run_temp(&temp, Some(1));
assert_eq!(scan.matches, vec!["001.txt:1:needle"]);
assert_eq!(scan.files_scanned, 1);
assert!(scan.truncated);
assert_eq!(scan.limit_reason, Some("matches"));
}
#[test]
fn windowed_limit_sixty_five_preserves_order_across_boundary() {
let temp = tempfile::TempDir::new().unwrap();
for index in 0..520 {
fs::write(temp.path().join(format!("{index:03}.txt")), "needle\n").unwrap();
}
let scan = run_temp(&temp, Some(65));
assert_eq!(scan.matches.len(), 65);
assert_eq!(scan.matches.first().unwrap(), "000.txt:1:needle");
assert_eq!(scan.matches.last().unwrap(), "064.txt:1:needle");
assert!(scan.truncated);
}
#[test]
fn parallel_aggregates_results_for_no_limit() {
let temp = tempfile::TempDir::new().unwrap();
for index in 0..16 {
fs::write(temp.path().join(format!("{index:03}.txt")), "needle\n").unwrap();
}
let scan = run_temp(&temp, None);
assert_eq!(scan.matches.len(), 16);
for index in 0..16 {
assert!(
scan.matches
.iter()
.any(|line| line == &format!("{index:03}.txt:1:needle"))
);
}
}
#[test]
fn parallel_default_limit_sets_cumulative_stop() {
let temp = tempfile::TempDir::new().unwrap();
let total_files = std::thread::available_parallelism()
.map_or(1, usize::from)
.saturating_add(150)
.max(200);
for index in 0..total_files {
fs::write(temp.path().join(format!("{index:04}.txt")), "needle\n").unwrap();
}
let scan = run_temp(&temp, None);
assert!(scan.matches.len() <= 100);
assert!(scan.files_scanned < total_files);
assert!(scan.truncated);
assert_eq!(scan.limit_reason, Some("matches"));
}
#[test]
fn oversized_prefix_sets_bytes_limit_reason() {
let temp = tempfile::TempDir::new().unwrap();
let path = temp.path().join("huge.txt");
fs::write(&path, vec![b'a'; reader::MAX_FILE_BYTES as usize + 1]).unwrap();
let scan = run_temp(&temp, Some(10));
assert!(scan.truncated);
assert_eq!(scan.limit_reason, Some("bytes"));
assert_eq!(scan.bytes_scanned, reader::MAX_FILE_BYTES);
}
#[test]
fn binary_file_does_not_emit_partial_garbage() {
let temp = tempfile::TempDir::new().unwrap();
fs::write(temp.path().join("bin.dat"), b"needle\0needle\n").unwrap();
let scan = run_temp(&temp, Some(10));
assert!(scan.matches.is_empty());
}
}