sift_core/search/scan/
json.rs

1use std::io::Write;
2use std::sync::atomic::{AtomicBool, Ordering};
3use std::time::Instant;
4
5use grep_printer::{JSON, Stats as JsonStats};
6use grep_regex::RegexMatcher;
7use grep_searcher::Searcher;
8use rayon::prelude::*;
9
10use crate::Candidate;
11use crate::search::emit::result::{ChunkOutput, FileResult};
12use crate::search::emit::stats::SearchStats;
13use crate::search::output::SearchOutput;
14use crate::search::output::mode::OutputEmission;
15use crate::search::query::SearchQuery;
16
17struct NullWriter;
18
19impl std::io::Write for NullWriter {
20    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
21        Ok(buf.len())
22    }
23
24    fn flush(&mut self) -> std::io::Result<()> {
25        Ok(())
26    }
27}
28
29struct JsonWorker<'a> {
30    searcher: Searcher,
31    matcher: &'a RegexMatcher,
32    output: SearchOutput,
33}
34
35impl<'a> JsonWorker<'a> {
36    fn new(scan: &'a JsonScan<'a>) -> Self {
37        Self {
38            searcher: scan
39                .search
40                .build_searcher(true, scan.search.opts().max_results, true),
41            matcher: scan.matcher,
42            output: scan.output,
43        }
44    }
45
46    fn search_candidate(&mut self, candidate: &Candidate, stop: &AtomicBool) -> FileResult {
47        let quiet = self.output.emission == OutputEmission::Quiet;
48        if stop.load(Ordering::SeqCst) {
49            return FileResult {
50                output: ChunkOutput::empty(),
51                json_stats: None,
52                hit: None,
53            };
54        }
55        let (bytes, file_stats) = if quiet {
56            let mut json = JSON::new(NullWriter);
57            let mut sink = json.sink_with_path(self.matcher, candidate.abs_path());
58            let _ = self
59                .searcher
60                .search_path(self.matcher, candidate.abs_path(), &mut sink);
61            (Vec::new(), sink.stats().clone())
62        } else {
63            let mut json = JSON::new(Vec::new());
64            let file_stats = {
65                let mut sink = json.sink_with_path(self.matcher, candidate.abs_path());
66                let _ = self
67                    .searcher
68                    .search_path(self.matcher, candidate.abs_path(), &mut sink);
69                sink.stats().clone()
70            };
71            (json.into_inner(), file_stats)
72        };
73        let had_match = file_stats.matches() > 0;
74        if quiet && had_match {
75            stop.store(true, Ordering::SeqCst);
76        }
77        FileResult {
78            output: ChunkOutput {
79                bytes,
80                matched: had_match,
81                heading: false,
82            },
83            json_stats: Some(file_stats),
84            hit: None,
85        }
86    }
87}
88
89fn format_json_summary_line(
90    wall: std::time::Duration,
91    agg: &JsonStats,
92) -> Result<String, crate::search::SearchError> {
93    let stats_val = serde_json::to_value(agg)?;
94    let wall_secs = f64::from(wall.subsec_nanos()).mul_add(1e-9, wall.as_secs_f64());
95    let v = serde_json::json!({
96        "type": "summary",
97        "data": {
98            "elapsed_total": {
99                "secs": wall.as_secs(),
100                "nanos": wall.subsec_nanos(),
101                "human": format!("{wall_secs:0.6}s"),
102            },
103            "stats": stats_val,
104        }
105    });
106    Ok(serde_json::to_string(&v)?)
107}
108
109pub struct JsonScan<'a> {
110    search: &'a SearchQuery,
111    matcher: &'a RegexMatcher,
112    output: SearchOutput,
113    wall_start: Instant,
114}
115
116impl<'a> JsonScan<'a> {
117    pub const fn new(
118        search: &'a SearchQuery,
119        matcher: &'a RegexMatcher,
120        output: SearchOutput,
121        wall_start: Instant,
122    ) -> Self {
123        Self {
124            search,
125            matcher,
126            output,
127            wall_start,
128        }
129    }
130
131    /// # Errors
132    ///
133    /// Returns an error if scanning or writing output fails.
134    pub fn run(
135        &self,
136        candidates: &[Candidate],
137        stats: Option<&mut SearchStats>,
138    ) -> crate::Result<bool> {
139        let stop = AtomicBool::new(false);
140        let n = candidates.len();
141        let mut files = Vec::with_capacity(n);
142        candidates
143            .par_iter()
144            .map_init(
145                || JsonWorker::new(self),
146                |worker: &mut JsonWorker<'_>, candidate: &Candidate| {
147                    worker.search_candidate(candidate, &stop)
148                },
149            )
150            .collect_into_vec(&mut files);
151
152        let mut merged = JsonStats::new();
153        let mut outputs = Vec::with_capacity(files.len());
154        for f in files {
155            if let Some(st) = f.json_stats {
156                merged += &st;
157            }
158            outputs.push(f.output);
159        }
160        let any_match = ChunkOutput::flush_all(outputs, None)?;
161        let summary_line = format_json_summary_line(self.wall_start.elapsed(), &merged)?;
162        let summary_bytes = summary_line.len() as u64 + 1;
163        let mut stdout = std::io::stdout().lock();
164        stdout.write_all(summary_line.as_bytes())?;
165        stdout.write_all(b"\n")?;
166        if let Some(s) = stats {
167            s.fill_from_json(
168                &merged,
169                candidates.len(),
170                crate::Candidate::total_file_bytes(candidates),
171                self.wall_start.elapsed(),
172                summary_bytes,
173            );
174        }
175        Ok(any_match)
176    }
177}
sift_core/search/scan/json.rs

sift_core/search/scan/
json.rs
