Skip to main content

sift_core/search/scan/
standard.rs

1use std::io::{self, Write};
2use std::path::PathBuf;
3use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
4
5use grep_matcher::{Captures, Matcher};
6use grep_regex::RegexMatcher;
7use grep_searcher::{Searcher, Sink, SinkContext, SinkMatch};
8use rayon::prelude::*;
9
10use crate::Candidate;
11use crate::search::emit::format::{ANSI_LINE, ANSI_PATH, ANSI_RESET};
12use crate::search::emit::result::{ChunkOutput, FileResult};
13use crate::search::emit::stats::TextStatsCounters;
14use crate::search::output::SearchOutput;
15use crate::search::output::format::ColumnAction;
16use crate::search::output::mode::OutputEmission;
17use crate::search::output::style::{
18    FilenameMode, LineStyleFlags, SearchRecordStyle, SearchSeparators,
19};
20use crate::search::query::SearchQuery;
21use crate::search::request::SearchCollection;
22
23#[derive(Clone, Copy)]
24pub struct SinkConfig {
25    pub before_context: usize,
26    pub after_context: usize,
27}
28
29pub struct StandardSink<'a> {
30    matcher: &'a RegexMatcher,
31    output: SearchOutput,
32    show_line_numbers: bool,
33    display_path: String,
34    bytes: &'a mut Vec<u8>,
35    separators: &'a SearchSeparators,
36    matched: bool,
37    match_count: usize,
38    replace: Option<&'a str>,
39    trim: bool,
40}
41
42impl<'a> StandardSink<'a> {
43    pub const fn new(
44        matcher: &'a RegexMatcher,
45        output: SearchOutput,
46        display_path: String,
47        bytes: &'a mut Vec<u8>,
48        separators: &'a SearchSeparators,
49        replace: Option<&'a str>,
50        config: SinkConfig,
51    ) -> Self {
52        Self {
53            matcher,
54            output,
55            show_line_numbers: output.lines.line_number()
56                || config.before_context > 0
57                || config.after_context > 0,
58            display_path,
59            bytes,
60            separators,
61            matched: false,
62            match_count: 0,
63            replace,
64            trim: output.lines.trim(),
65        }
66    }
67}
68
69impl Sink for StandardSink<'_> {
70    type Error = io::Error;
71
72    fn matched(&mut self, searcher: &Searcher, mat: &SinkMatch<'_>) -> Result<bool, Self::Error> {
73        std::hint::black_box(searcher);
74        self.matched = true;
75        self.match_count += 1;
76
77        if self.output.emission == OutputEmission::Quiet {
78            return Ok(true);
79        }
80
81        if matches!(
82            self.output.mode,
83            crate::search::output::mode::SearchMode::OnlyMatching
84        ) {
85            return Ok(self.handle_only_matching(mat));
86        }
87
88        let line_bytes = mat.bytes();
89        if self.handle_max_columns(line_bytes, mat)? {
90            return Ok(true);
91        }
92
93        let col = self.compute_first_column(mat);
94        self.write_prefix(
95            mat.line_number(),
96            false,
97            col,
98            if self.output.lines.byte_offset() {
99                Some(mat.absolute_byte_offset())
100            } else {
101                None
102            },
103        )?;
104        self.write_line_content(mat.bytes())
105    }
106
107    fn context(&mut self, searcher: &Searcher, ctx: &SinkContext<'_>) -> Result<bool, Self::Error> {
108        std::hint::black_box(searcher);
109        if self.output.emission == OutputEmission::Quiet {
110            return Ok(true);
111        }
112        if matches!(
113            self.output.mode,
114            crate::search::output::mode::SearchMode::OnlyMatching
115        ) {
116            return Ok(true);
117        }
118        let ctx_bytes = ctx.bytes();
119        let columns = self.output.lines.columns;
120        match columns.map(|c| c.classify(ctx_bytes)) {
121            Some(ColumnAction::Omit) => return Ok(true),
122            Some(ColumnAction::Preview) => {
123                let limit = columns.unwrap();
124                self.write_prefix(
125                    ctx.line_number(),
126                    true,
127                    None,
128                    if self.output.lines.byte_offset() {
129                        Some(ctx.absolute_byte_offset())
130                    } else {
131                        None
132                    },
133                )?;
134                let truncated = limit.truncate(ctx_bytes);
135                self.bytes.write_all(&truncated)?;
136                return Ok(true);
137            }
138            Some(ColumnAction::Normal) | None => {}
139        }
140        self.write_prefix(
141            ctx.line_number(),
142            true,
143            None,
144            if self.output.lines.byte_offset() {
145                Some(ctx.absolute_byte_offset())
146            } else {
147                None
148            },
149        )?;
150        let line_bytes = ctx.bytes();
151        if self.trim {
152            let s = String::from_utf8_lossy(line_bytes);
153            let trimmed = s.trim_start();
154            self.bytes.write_all(trimmed.as_bytes())?;
155            if !trimmed.ends_with('\n') {
156                self.bytes.write_all(b"\n")?;
157            }
158        } else {
159            self.bytes.write_all(line_bytes)?;
160            if !line_bytes.ends_with(b"\n") {
161                self.bytes.write_all(b"\n")?;
162            }
163        }
164        Ok(true)
165    }
166
167    fn context_break(&mut self, searcher: &Searcher) -> Result<bool, Self::Error> {
168        std::hint::black_box(searcher);
169        if self.output.emission == OutputEmission::Quiet {
170            return Ok(true);
171        }
172        if matches!(
173            self.output.mode,
174            crate::search::output::mode::SearchMode::OnlyMatching
175        ) {
176            return Ok(true);
177        }
178        if let Some(ref sep) = self.separators.context_separator {
179            self.bytes.write_all(sep)?;
180            self.bytes.write_all(b"\n")?;
181        }
182        Ok(true)
183    }
184}
185
186impl StandardSink<'_> {
187    fn handle_only_matching(&mut self, mat: &SinkMatch<'_>) -> bool {
188        let show_column = self.output.lines.flags.contains(LineStyleFlags::COLUMN);
189        let line_number = mat.line_number();
190        let line = mat.bytes();
191        let byte_offset = mat.absolute_byte_offset();
192        let _ = self.matcher.find_iter(line, |m: grep_matcher::Match| {
193            let col = if show_column {
194                Some(m.start() + 1)
195            } else {
196                None
197            };
198            let matched_slice = &line[m.start()..m.end()];
199            let text = if let Some(rep) = self.replace {
200                apply_replace(self.matcher, matched_slice, rep)
201            } else {
202                String::from_utf8_lossy(matched_slice).into_owned()
203            };
204            let _ = self.write_prefix(
205                line_number,
206                false,
207                col,
208                if self.output.lines.byte_offset() {
209                    Some(byte_offset)
210                } else {
211                    None
212                },
213            );
214            let _ = self.bytes.write_all(text.as_bytes());
215            let _ = self.bytes.write_all(b"\n");
216            true
217        });
218        true
219    }
220
221    fn handle_max_columns(
222        &mut self,
223        line_bytes: &[u8],
224        mat: &SinkMatch<'_>,
225    ) -> Result<bool, io::Error> {
226        let columns = self.output.lines.columns;
227        match columns.map(|c| c.classify(line_bytes)) {
228            Some(ColumnAction::Omit) => return Ok(true),
229            Some(ColumnAction::Preview) => {
230                let limit = columns.unwrap();
231                self.write_prefix(
232                    mat.line_number(),
233                    false,
234                    None,
235                    if self.output.lines.byte_offset() {
236                        Some(mat.absolute_byte_offset())
237                    } else {
238                        None
239                    },
240                )?;
241                let truncated = limit.truncate(line_bytes);
242                self.bytes.write_all(&truncated)?;
243                return Ok(true);
244            }
245            Some(ColumnAction::Normal) | None => {}
246        }
247        Ok(false)
248    }
249
250    fn compute_first_column(&self, mat: &SinkMatch<'_>) -> Option<usize> {
251        if !self.output.lines.flags.contains(LineStyleFlags::COLUMN) {
252            return None;
253        }
254        let line = mat.bytes();
255        let mut first_col = None;
256        let _ = self.matcher.find_iter(line, |m: grep_matcher::Match| {
257            first_col = Some(m.start() + 1);
258            false
259        });
260        first_col
261    }
262
263    fn write_prefix(
264        &mut self,
265        line_number: Option<u64>,
266        is_context_line: bool,
267        column: Option<usize>,
268        byte_offset: Option<u64>,
269    ) -> io::Result<()> {
270        let color = self.output.records.should_color();
271        let print_filename = self.output.lines.filename_mode != FilenameMode::Never;
272        let field_sep = if is_context_line {
273            &self.separators.field_context_separator
274        } else {
275            &self.separators.field_match_separator
276        };
277        if print_filename {
278            if color {
279                self.bytes.extend_from_slice(ANSI_PATH);
280            }
281            write!(self.bytes, "{}", self.display_path)?;
282            if color {
283                self.bytes.extend_from_slice(ANSI_RESET);
284            }
285            self.bytes.extend_from_slice(field_sep);
286        }
287        if self.show_line_numbers {
288            if color {
289                self.bytes.extend_from_slice(ANSI_LINE);
290            }
291            write!(self.bytes, "{}", line_number.unwrap_or(0))?;
292            if color {
293                self.bytes.extend_from_slice(ANSI_RESET);
294            }
295            self.bytes.extend_from_slice(field_sep);
296        }
297        if let Some(col) = column {
298            if color {
299                self.bytes.extend_from_slice(ANSI_LINE);
300            }
301            write!(self.bytes, "{col}")?;
302            if color {
303                self.bytes.extend_from_slice(ANSI_RESET);
304            }
305            self.bytes.extend_from_slice(field_sep);
306        }
307        if let Some(offset) = byte_offset {
308            if color {
309                self.bytes.extend_from_slice(ANSI_LINE);
310            }
311            let sep = if is_context_line { '-' } else { ':' };
312            write!(self.bytes, "{offset}{sep}")?;
313            if color {
314                self.bytes.extend_from_slice(ANSI_RESET);
315            }
316        }
317        Ok(())
318    }
319
320    fn write_line_content(&mut self, line_bytes: &[u8]) -> Result<bool, io::Error> {
321        if let Some(rep) = self.replace {
322            let text = apply_replace(self.matcher, line_bytes, rep);
323            self.bytes.write_all(text.as_bytes())?;
324            if !text.ends_with('\n') {
325                self.bytes.write_all(b"\n")?;
326            }
327        } else if self.trim {
328            let trimmed = String::from_utf8_lossy(line_bytes);
329            let trimmed = trimmed.trim_start();
330            self.bytes.write_all(trimmed.as_bytes())?;
331            if !trimmed.ends_with('\n') {
332                self.bytes.write_all(b"\n")?;
333            }
334        } else {
335            self.bytes.write_all(line_bytes)?;
336            if !line_bytes.ends_with(b"\n") {
337                self.bytes.write_all(b"\n")?;
338            }
339        }
340        Ok(true)
341    }
342}
343
344struct StandardWorker<'a> {
345    matcher: &'a RegexMatcher,
346    searcher: Searcher,
347    output: SearchOutput,
348    separators: &'a SearchSeparators,
349    bytes: Vec<u8>,
350    match_counter: Option<&'a AtomicUsize>,
351    files_with_matches: Option<&'a AtomicUsize>,
352    replace: Option<String>,
353    sink_config: SinkConfig,
354    records: SearchRecordStyle,
355    lines_flags: LineStyleFlags,
356    filename_mode: FilenameMode,
357    path_display: crate::search::output::style::PathDisplay,
358    path_separator: Option<u8>,
359    emission: OutputEmission,
360    collect_hits: bool,
361}
362
363impl<'a> StandardWorker<'a> {
364    fn new(scan: &'a StandardScan<'a>, collect: SearchCollection) -> Self {
365        Self {
366            searcher: scan.search.build_searcher(
367                scan.output.lines.line_number(),
368                scan.search.opts().max_results,
369                true,
370            ),
371            matcher: scan.matcher,
372            output: scan.output,
373            separators: scan.separators,
374            bytes: Vec::new(),
375            match_counter: scan.counters.primary(),
376            files_with_matches: scan.counters.files_with_matches(),
377            replace: scan.search.opts().replace.clone(),
378            sink_config: SinkConfig {
379                before_context: scan.search.opts().before_context,
380                after_context: scan.search.opts().after_context,
381            },
382            records: scan.output.records,
383            lines_flags: scan.output.lines.flags,
384            filename_mode: scan.output.lines.filename_mode,
385            path_display: scan.output.lines.path_display,
386            path_separator: scan.output.records.path_separator,
387            emission: scan.output.emission,
388            collect_hits: collect.hits,
389        }
390    }
391
392    fn search_candidate(&mut self, candidate: &Candidate, stop: &AtomicBool) -> FileResult {
393        self.bytes.clear();
394        if stop.load(Ordering::SeqCst) {
395            return FileResult {
396                output: ChunkOutput::empty(),
397                json_stats: None,
398                hit: None,
399            };
400        }
401
402        let matched = {
403            let heading = self.lines_flags.contains(LineStyleFlags::HEADING)
404                && self.filename_mode != FilenameMode::Never;
405            let mut sink_output = self.output;
406            if heading {
407                sink_output.lines.filename_mode = FilenameMode::Never;
408            }
409            let display = candidate.display_path(self.path_display, self.path_separator);
410            let mut sink = StandardSink::new(
411                self.matcher,
412                sink_output,
413                display,
414                &mut self.bytes,
415                self.separators,
416                self.replace.as_deref(),
417                self.sink_config,
418            );
419            let _ = self
420                .searcher
421                .search_path(self.matcher, candidate.abs_path(), &mut sink);
422            let n = sink.match_count;
423            if let Some(c) = self.match_counter {
424                c.fetch_add(n, Ordering::Relaxed);
425            }
426            if sink.matched
427                && let Some(c) = self.files_with_matches
428            {
429                c.fetch_add(1, Ordering::Relaxed);
430            }
431            sink.matched
432        };
433
434        if self.emission == OutputEmission::Quiet && matched {
435            stop.store(true, Ordering::SeqCst);
436        }
437
438        FileResult {
439            output: ChunkOutput {
440                bytes: if matched
441                    && self.lines_flags.contains(LineStyleFlags::HEADING)
442                    && self.filename_mode != FilenameMode::Never
443                    && self.emission != OutputEmission::Quiet
444                {
445                    let mut out = Vec::new();
446                    if self.records.should_color() {
447                        out.extend_from_slice(ANSI_PATH);
448                    }
449                    let display = candidate.display_path(self.path_display, self.path_separator);
450                    let _ = write!(out, "{display}");
451                    if self.records.should_color() {
452                        out.extend_from_slice(ANSI_RESET);
453                    }
454                    self.records.terminator.write_to(&mut out);
455                    out.extend(std::mem::take(&mut self.bytes));
456                    out
457                } else {
458                    std::mem::take(&mut self.bytes)
459                },
460                matched,
461                heading: matched
462                    && self.lines_flags.contains(LineStyleFlags::HEADING)
463                    && self.filename_mode != FilenameMode::Never
464                    && self.emission != OutputEmission::Quiet,
465            },
466            json_stats: None,
467            hit: (matched && self.collect_hits).then(|| candidate.rel_path().to_path_buf()),
468        }
469    }
470}
471
472pub struct StandardScan<'a> {
473    search: &'a SearchQuery,
474    matcher: &'a RegexMatcher,
475    output: SearchOutput,
476    separators: &'a SearchSeparators,
477    counters: &'a TextStatsCounters,
478}
479
480impl<'a> StandardScan<'a> {
481    pub const fn new(
482        search: &'a SearchQuery,
483        matcher: &'a RegexMatcher,
484        output: SearchOutput,
485        separators: &'a SearchSeparators,
486        counters: &'a TextStatsCounters,
487    ) -> Self {
488        Self {
489            search,
490            matcher,
491            output,
492            separators,
493            counters,
494        }
495    }
496
497    /// # Errors
498    ///
499    /// Returns an error if scanning or writing output fails.
500    pub fn run(
501        &self,
502        candidates: &[Candidate],
503        collect: SearchCollection,
504    ) -> crate::Result<(bool, Vec<PathBuf>)> {
505        let stop = AtomicBool::new(false);
506        let n = candidates.len();
507        let mut files = Vec::with_capacity(n);
508        candidates
509            .par_iter()
510            .map_init(
511                || StandardWorker::new(self, collect),
512                |worker: &mut StandardWorker<'_>, candidate: &Candidate| {
513                    worker.search_candidate(candidate, &stop)
514                },
515            )
516            .collect_into_vec(&mut files);
517        let mut hits = Vec::new();
518        let mut outputs = Vec::with_capacity(files.len());
519        for file in files {
520            if collect.hits
521                && let Some(hit) = file.hit
522            {
523                hits.push(hit);
524            }
525            outputs.push(file.output);
526        }
527        let any_match = ChunkOutput::flush_all(outputs, self.counters.bytes_printed())?;
528        Ok((any_match, hits))
529    }
530}
531
532fn apply_replace(matcher: &RegexMatcher, line: &[u8], replacement: &str) -> String {
533    let Ok(mut caps) = matcher.new_captures() else {
534        return String::from_utf8_lossy(line).into_owned();
535    };
536    let mut dst = Vec::new();
537    let _ = matcher.replace_with_captures(line, &mut caps, &mut dst, |caps, dst| {
538        caps.interpolate(
539            |name| matcher.capture_index(name),
540            line,
541            replacement.as_bytes(),
542            dst,
543        );
544        true
545    });
546    String::from_utf8_lossy(&dst).into_owned()
547}