use std::io::{self, Write};
use std::path::Path;
use std::time::Instant;
use grep_matcher::{Match, Matcher};
use grep_searcher::{
Searcher, Sink, SinkContext, SinkContextKind, SinkFinish, SinkMatch,
};
use serde_json as json;
use crate::counter::CounterWriter;
use crate::jsont;
use crate::stats::Stats;
use crate::util::find_iter_at_in_context;
#[derive(Debug, Clone)]
struct Config {
pretty: bool,
max_matches: Option<u64>,
always_begin_end: bool,
}
impl Default for Config {
fn default() -> Config {
Config { pretty: false, max_matches: None, always_begin_end: false }
}
}
#[derive(Clone, Debug)]
pub struct JSONBuilder {
config: Config,
}
impl JSONBuilder {
pub fn new() -> JSONBuilder {
JSONBuilder { config: Config::default() }
}
pub fn build<W: io::Write>(&self, wtr: W) -> JSON<W> {
JSON {
config: self.config.clone(),
wtr: CounterWriter::new(wtr),
matches: vec![],
}
}
pub fn pretty(&mut self, yes: bool) -> &mut JSONBuilder {
self.config.pretty = yes;
self
}
pub fn max_matches(&mut self, limit: Option<u64>) -> &mut JSONBuilder {
self.config.max_matches = limit;
self
}
pub fn always_begin_end(&mut self, yes: bool) -> &mut JSONBuilder {
self.config.always_begin_end = yes;
self
}
}
#[derive(Debug)]
pub struct JSON<W> {
config: Config,
wtr: CounterWriter<W>,
matches: Vec<Match>,
}
impl<W: io::Write> JSON<W> {
pub fn new(wtr: W) -> JSON<W> {
JSONBuilder::new().build(wtr)
}
pub fn sink<'s, M: Matcher>(
&'s mut self,
matcher: M,
) -> JSONSink<'static, 's, M, W> {
JSONSink {
matcher: matcher,
json: self,
path: None,
start_time: Instant::now(),
match_count: 0,
after_context_remaining: 0,
binary_byte_offset: None,
begin_printed: false,
stats: Stats::new(),
}
}
pub fn sink_with_path<'p, 's, M, P>(
&'s mut self,
matcher: M,
path: &'p P,
) -> JSONSink<'p, 's, M, W>
where
M: Matcher,
P: ?Sized + AsRef<Path>,
{
JSONSink {
matcher: matcher,
json: self,
path: Some(path.as_ref()),
start_time: Instant::now(),
match_count: 0,
after_context_remaining: 0,
binary_byte_offset: None,
begin_printed: false,
stats: Stats::new(),
}
}
fn write_message(
&mut self,
message: &jsont::Message<'_>,
) -> io::Result<()> {
if self.config.pretty {
json::to_writer_pretty(&mut self.wtr, message)?;
} else {
json::to_writer(&mut self.wtr, message)?;
}
self.wtr.write(&[b'\n'])?;
Ok(())
}
}
impl<W> JSON<W> {
pub fn has_written(&self) -> bool {
self.wtr.total_count() > 0
}
pub fn get_mut(&mut self) -> &mut W {
self.wtr.get_mut()
}
pub fn into_inner(self) -> W {
self.wtr.into_inner()
}
}
#[derive(Debug)]
pub struct JSONSink<'p, 's, M: Matcher, W> {
matcher: M,
json: &'s mut JSON<W>,
path: Option<&'p Path>,
start_time: Instant,
match_count: u64,
after_context_remaining: u64,
binary_byte_offset: Option<u64>,
begin_printed: bool,
stats: Stats,
}
impl<'p, 's, M: Matcher, W: io::Write> JSONSink<'p, 's, M, W> {
pub fn has_match(&self) -> bool {
self.match_count > 0
}
pub fn match_count(&self) -> u64 {
self.match_count
}
pub fn binary_byte_offset(&self) -> Option<u64> {
self.binary_byte_offset
}
pub fn stats(&self) -> &Stats {
&self.stats
}
fn record_matches(
&mut self,
searcher: &Searcher,
bytes: &[u8],
range: std::ops::Range<usize>,
) -> io::Result<()> {
self.json.matches.clear();
let matches = &mut self.json.matches;
find_iter_at_in_context(
searcher,
&self.matcher,
bytes,
range.clone(),
|m| {
let (s, e) = (m.start() - range.start, m.end() - range.start);
matches.push(Match::new(s, e));
true
},
)?;
if !matches.is_empty()
&& matches.last().unwrap().is_empty()
&& matches.last().unwrap().start() >= bytes.len()
{
matches.pop().unwrap();
}
Ok(())
}
fn should_quit(&self) -> bool {
let limit = match self.json.config.max_matches {
None => return false,
Some(limit) => limit,
};
if self.match_count < limit {
return false;
}
self.after_context_remaining == 0
}
fn match_more_than_limit(&self) -> bool {
let limit = match self.json.config.max_matches {
None => return false,
Some(limit) => limit,
};
self.match_count > limit
}
fn write_begin_message(&mut self) -> io::Result<()> {
if self.begin_printed {
return Ok(());
}
let msg = jsont::Message::Begin(jsont::Begin { path: self.path });
self.json.write_message(&msg)?;
self.begin_printed = true;
Ok(())
}
}
impl<'p, 's, M: Matcher, W: io::Write> Sink for JSONSink<'p, 's, M, W> {
type Error = io::Error;
fn matched(
&mut self,
searcher: &Searcher,
mat: &SinkMatch<'_>,
) -> Result<bool, io::Error> {
self.write_begin_message()?;
self.match_count += 1;
if self.match_more_than_limit() {
self.after_context_remaining =
self.after_context_remaining.saturating_sub(1);
} else {
self.after_context_remaining = searcher.after_context() as u64;
}
self.record_matches(
searcher,
mat.buffer(),
mat.bytes_range_in_buffer(),
)?;
self.stats.add_matches(self.json.matches.len() as u64);
self.stats.add_matched_lines(mat.lines().count() as u64);
let submatches = SubMatches::new(mat.bytes(), &self.json.matches);
let msg = jsont::Message::Match(jsont::Match {
path: self.path,
lines: mat.bytes(),
line_number: mat.line_number(),
absolute_offset: mat.absolute_byte_offset(),
submatches: submatches.as_slice(),
});
self.json.write_message(&msg)?;
Ok(!self.should_quit())
}
fn context(
&mut self,
searcher: &Searcher,
ctx: &SinkContext<'_>,
) -> Result<bool, io::Error> {
self.write_begin_message()?;
self.json.matches.clear();
if ctx.kind() == &SinkContextKind::After {
self.after_context_remaining =
self.after_context_remaining.saturating_sub(1);
}
let submatches = if searcher.invert_match() {
self.record_matches(searcher, ctx.bytes(), 0..ctx.bytes().len())?;
SubMatches::new(ctx.bytes(), &self.json.matches)
} else {
SubMatches::empty()
};
let msg = jsont::Message::Context(jsont::Context {
path: self.path,
lines: ctx.bytes(),
line_number: ctx.line_number(),
absolute_offset: ctx.absolute_byte_offset(),
submatches: submatches.as_slice(),
});
self.json.write_message(&msg)?;
Ok(!self.should_quit())
}
fn begin(&mut self, _searcher: &Searcher) -> Result<bool, io::Error> {
self.json.wtr.reset_count();
self.start_time = Instant::now();
self.match_count = 0;
self.after_context_remaining = 0;
self.binary_byte_offset = None;
if self.json.config.max_matches == Some(0) {
return Ok(false);
}
if !self.json.config.always_begin_end {
return Ok(true);
}
self.write_begin_message()?;
Ok(true)
}
fn finish(
&mut self,
_searcher: &Searcher,
finish: &SinkFinish,
) -> Result<(), io::Error> {
if !self.begin_printed {
return Ok(());
}
self.binary_byte_offset = finish.binary_byte_offset();
self.stats.add_elapsed(self.start_time.elapsed());
self.stats.add_searches(1);
if self.match_count > 0 {
self.stats.add_searches_with_match(1);
}
self.stats.add_bytes_searched(finish.byte_count());
self.stats.add_bytes_printed(self.json.wtr.count());
let msg = jsont::Message::End(jsont::End {
path: self.path,
binary_offset: finish.binary_byte_offset(),
stats: self.stats.clone(),
});
self.json.write_message(&msg)?;
Ok(())
}
}
enum SubMatches<'a> {
Empty,
Small([jsont::SubMatch<'a>; 1]),
Big(Vec<jsont::SubMatch<'a>>),
}
impl<'a> SubMatches<'a> {
fn new(bytes: &'a [u8], matches: &[Match]) -> SubMatches<'a> {
if matches.len() == 1 {
let mat = matches[0];
SubMatches::Small([jsont::SubMatch {
m: &bytes[mat],
start: mat.start(),
end: mat.end(),
}])
} else {
let mut match_ranges = vec![];
for &mat in matches {
match_ranges.push(jsont::SubMatch {
m: &bytes[mat],
start: mat.start(),
end: mat.end(),
});
}
SubMatches::Big(match_ranges)
}
}
fn empty() -> SubMatches<'static> {
SubMatches::Empty
}
fn as_slice(&self) -> &[jsont::SubMatch<'_>] {
match *self {
SubMatches::Empty => &[],
SubMatches::Small(ref x) => x,
SubMatches::Big(ref x) => x,
}
}
}
#[cfg(test)]
mod tests {
use grep_matcher::LineTerminator;
use grep_regex::{RegexMatcher, RegexMatcherBuilder};
use grep_searcher::SearcherBuilder;
use super::{JSONBuilder, JSON};
const SHERLOCK: &'static [u8] = b"\
For the Doctor Watsons of this world, as opposed to the Sherlock
Holmeses, success in the province of detective work must always
be, to a very large extent, the result of luck. Sherlock Holmes
can extract a clew from a wisp of straw or a flake of cigar ash;
but Doctor Watson has to have it taken out for him and dusted,
and exhibited clearly, with a label attached.
";
fn printer_contents(printer: &mut JSON<Vec<u8>>) -> String {
String::from_utf8(printer.get_mut().to_owned()).unwrap()
}
#[test]
fn binary_detection() {
use grep_searcher::BinaryDetection;
const BINARY: &'static [u8] = b"\
For the Doctor Watsons of this world, as opposed to the Sherlock
Holmeses, success in the province of detective work must always
be, to a very large extent, the result of luck. Sherlock Holmes
can extract a clew \x00 from a wisp of straw or a flake of cigar ash;
but Doctor Watson has to have it taken out for him and dusted,
and exhibited clearly, with a label attached.\
";
let matcher = RegexMatcher::new(r"Watson").unwrap();
let mut printer = JSONBuilder::new().build(vec![]);
SearcherBuilder::new()
.binary_detection(BinaryDetection::quit(b'\x00'))
.heap_limit(Some(80))
.build()
.search_reader(&matcher, BINARY, printer.sink(&matcher))
.unwrap();
let got = printer_contents(&mut printer);
assert_eq!(got.lines().count(), 3);
let last = got.lines().last().unwrap();
assert!(last.contains(r#""binary_offset":212,"#));
}
#[test]
fn max_matches() {
let matcher = RegexMatcher::new(r"Watson").unwrap();
let mut printer =
JSONBuilder::new().max_matches(Some(1)).build(vec![]);
SearcherBuilder::new()
.build()
.search_reader(&matcher, SHERLOCK, printer.sink(&matcher))
.unwrap();
let got = printer_contents(&mut printer);
assert_eq!(got.lines().count(), 3);
}
#[test]
fn max_matches_after_context() {
let haystack = "\
a
b
c
d
e
d
e
d
e
d
e
";
let matcher = RegexMatcher::new(r"d").unwrap();
let mut printer =
JSONBuilder::new().max_matches(Some(1)).build(vec![]);
SearcherBuilder::new()
.after_context(2)
.build()
.search_reader(
&matcher,
haystack.as_bytes(),
printer.sink(&matcher),
)
.unwrap();
let got = printer_contents(&mut printer);
assert_eq!(got.lines().count(), 5);
}
#[test]
fn no_match() {
let matcher = RegexMatcher::new(r"DOES NOT MATCH").unwrap();
let mut printer = JSONBuilder::new().build(vec![]);
SearcherBuilder::new()
.build()
.search_reader(&matcher, SHERLOCK, printer.sink(&matcher))
.unwrap();
let got = printer_contents(&mut printer);
assert!(got.is_empty());
}
#[test]
fn always_begin_end_no_match() {
let matcher = RegexMatcher::new(r"DOES NOT MATCH").unwrap();
let mut printer =
JSONBuilder::new().always_begin_end(true).build(vec![]);
SearcherBuilder::new()
.build()
.search_reader(&matcher, SHERLOCK, printer.sink(&matcher))
.unwrap();
let got = printer_contents(&mut printer);
assert_eq!(got.lines().count(), 2);
assert!(got.contains("begin") && got.contains("end"));
}
#[test]
fn missing_crlf() {
let haystack = "test\r\n".as_bytes();
let matcher = RegexMatcherBuilder::new().build("test").unwrap();
let mut printer = JSONBuilder::new().build(vec![]);
SearcherBuilder::new()
.build()
.search_reader(&matcher, haystack, printer.sink(&matcher))
.unwrap();
let got = printer_contents(&mut printer);
assert_eq!(got.lines().count(), 3);
assert!(
got.lines().nth(1).unwrap().contains(r"test\r\n"),
r"missing 'test\r\n' in '{}'",
got.lines().nth(1).unwrap(),
);
let matcher =
RegexMatcherBuilder::new().crlf(true).build("test").unwrap();
let mut printer = JSONBuilder::new().build(vec![]);
SearcherBuilder::new()
.line_terminator(LineTerminator::crlf())
.build()
.search_reader(&matcher, haystack, printer.sink(&matcher))
.unwrap();
let got = printer_contents(&mut printer);
assert_eq!(got.lines().count(), 3);
assert!(
got.lines().nth(1).unwrap().contains(r"test\r\n"),
r"missing 'test\r\n' in '{}'",
got.lines().nth(1).unwrap(),
);
}
}