use std::collections::HashMap;
use std::fs::File;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{sync_channel, Receiver, SyncSender};
use std::sync::Arc;
use std::time::Duration;
use ahash::RandomState;
use anyhow::{bail, Result};
use console::style;
use regex::Regex;
use serde::Serialize;
use serde_json::json;
use structopt::StructOpt;
use super::util::{expand_dirs, DataExecutor, DataInstance};
use crate::util;
#[derive(Debug, StructOpt, Clone)]
pub(crate) struct Opt {
#[structopt(parse(from_os_str))]
path: Vec<PathBuf>,
#[structopt(short = "p", long = "pattern", number_of_values = 1)]
pattern: Vec<String>,
#[structopt(long = "--with-locations")]
with_locations: bool,
#[structopt(long = "--locations-limit")]
locations_limit: Option<usize>,
#[structopt(short = "l", long = "limit")]
limit: Option<usize>,
#[structopt(long = "file-limit")]
file_limit: Option<usize>,
#[structopt(short = "j", long = "workers")]
workers: Option<usize>,
#[structopt(short = "o", long = "out")]
out: Option<PathBuf>,
#[structopt(short = "q", long = "quiet")]
quiet: bool,
#[structopt(long = "json")]
json: bool,
#[structopt(short = "f", long = "force")]
force: bool,
}
type MatchLocationsPerPath = HashMap<PathBuf, Vec<MatchLocation>, RandomState>;
pub(crate) fn main(mut opt: Opt) -> Result<()> {
opt.path = expand_dirs(&opt.path)?;
if opt.pattern.is_empty() {
bail!("At least one -p/--pattern regex term is required");
}
if let Some(file_limit) = opt.file_limit {
if file_limit == 0 {
bail!("File limit cannot be 0");
}
opt.path.truncate(file_limit);
}
if opt.path.is_empty() {
bail!("at least one path is required");
}
if opt.locations_limit.is_some() && !opt.with_locations {
bail!("--locations-limit is only valid when using --with-locations");
}
let mut counts: HashMap<String, Arc<AtomicUsize>, RandomState> =
HashMap::with_capacity_and_hasher(opt.pattern.len(), RandomState::new());
let mut patterns: HashMap<String, Regex, RandomState> =
HashMap::with_capacity_and_hasher(opt.pattern.len(), RandomState::new());
let mut match_locations: Option<HashMap<String, MatchLocationsPerPath, RandomState>> = None;
let mut tx: Option<SyncSender<(String, PathBuf, Vec<MatchLocation>)>> = None;
let mut rx: Option<Receiver<(String, PathBuf, Vec<MatchLocation>)>> = None;
if opt.with_locations {
let (tx_, rx_) = sync_channel::<(String, PathBuf, Vec<MatchLocation>)>(512_000);
tx = Some(tx_);
rx = Some(rx_);
match_locations = Some(HashMap::with_capacity_and_hasher(
opt.pattern.len(),
RandomState::new(),
));
}
for pattern in &opt.pattern {
counts.insert(pattern.to_string(), Arc::new(AtomicUsize::new(0)));
patterns.insert(pattern.to_string(), Regex::new(pattern)?);
if let Some(ref mut locations) = match_locations {
locations.insert(
pattern.to_string(),
HashMap::with_hasher(RandomState::new()),
);
}
}
let (mut out_file, out_path) = match get_output_file(&opt)? {
Some(out) => (Some(out.0), Some(out.1)),
None => (None, None),
};
let executor = DataExecutor::new(&opt.path, opt.workers, opt.limit, "Searching", opt.quiet)?;
for path in &opt.path {
let counts = counts.clone();
let patterns = patterns.clone();
let sync_match_locations = {
let tx = tx.clone();
let path = path.clone();
move |local_match_locations: Option<
HashMap<String, Vec<MatchLocation>, RandomState>,
>|
-> Result<()> {
if let Some(mut local_match_locations) = local_match_locations {
let tx = tx.as_ref().unwrap();
for (pattern, matches) in local_match_locations.drain() {
tx.send((pattern, path.clone(), matches))?;
}
}
Ok(())
}
};
let local_match_locations_factory = {
let opt = opt.clone();
move || -> Result<Option<HashMap<String, Vec<MatchLocation>, RandomState>>> {
let mut local_match_locations: Option<
HashMap<String, Vec<MatchLocation>, RandomState>,
> = None;
if opt.with_locations {
local_match_locations = Some(HashMap::with_capacity_and_hasher(
opt.pattern.len(),
RandomState::new(),
));
for pattern in &opt.pattern {
local_match_locations
.as_mut()
.unwrap()
.insert(pattern.into(), Vec::new());
}
}
Ok(local_match_locations)
}
};
executor.execute_with_callback(
path,
move |data: DataInstance,
_: &Path,
line_num: usize,
local_match_locations: &mut Option<
HashMap<String, Vec<MatchLocation>, RandomState>,
>|
-> Result<()> {
if let Some(text) = data.text {
for (pattern, regex) in &patterns {
let mut collect_matches = local_match_locations.is_some();
let mut submatches: Option<Vec<SubMatch>> = None;
if collect_matches {
submatches = Some(Vec::new());
}
for m in regex.find_iter(&text) {
let count =
counts.get(pattern).unwrap().fetch_add(1, Ordering::Relaxed);
if let Some(locations_limit) = opt.locations_limit {
if count >= locations_limit {
collect_matches = false;
}
}
if collect_matches {
submatches.as_mut().unwrap().push(SubMatch {
start_col: m.start(),
end_col: m.end(),
});
}
}
if let Some(submatches) = submatches {
if !submatches.is_empty() {
let locations = local_match_locations.as_mut().unwrap();
let match_location = MatchLocation {
text: text.clone(),
line_num,
submatches,
};
locations.get_mut(pattern).unwrap().push(match_location);
}
}
}
};
Ok(())
},
local_match_locations_factory,
sync_match_locations,
)?;
}
drop(tx);
while !executor.done() {
if let Some(ref rx) = rx {
while let Ok((pattern, path, matches)) = rx.recv_timeout(Duration::from_secs(1)) {
let matches_for_pattern: &mut HashMap<
PathBuf,
Vec<MatchLocation>,
ahash::RandomState,
> = match_locations
.as_mut()
.map(|m| m.get_mut(&pattern).unwrap())
.unwrap();
if !matches.is_empty() {
matches_for_pattern.insert(path.clone(), matches);
}
}
}
}
executor.join()?;
for (i, (pattern, count)) in counts.iter().enumerate() {
let count = count.load(Ordering::Relaxed);
let matches_for_pattern = match_locations.as_ref().map(|m| m.get(pattern).unwrap());
let json_out = if let Some(matches_for_pattern) = matches_for_pattern {
json!({
"pattern": pattern,
"count": count,
"matches": matches_for_pattern,
})
.to_string()
} else {
json!({
"pattern": pattern,
"count": count,
})
.to_string()
};
if opt.json {
println!("{json_out}");
} else if !opt.quiet {
println!(
"[{}/{}] {:?} (count = {})",
i + 1,
counts.len(),
style(pattern).cyan(),
count
);
if let Some(locations) = matches_for_pattern {
for (path, locs) in locations.iter() {
if locs.is_empty() {
continue;
}
for loc in locs.iter() {
println!(
" → [{}, line {}, {} matches] {:?}",
path.to_string_lossy(),
loc.line_num,
loc.submatches.len(),
loc.text,
);
}
}
}
}
if let Some(ref mut file) = out_file {
writeln!(file, "{json_out}")?;
}
}
if let Some(path) = out_path {
log::info!("Output written to {:?}", path);
}
Ok(())
}
fn get_output_file(opt: &Opt) -> Result<Option<(File, PathBuf)>> {
if let Some(path) = &opt.out {
if path.is_dir() {
bail!("-o/--out must be a valid file name, not a directory");
} else {
Ok(Some(util::get_output_file(path, opt.force)?))
}
} else {
Ok(None)
}
}
#[derive(Debug, Clone, Serialize)]
struct SubMatch {
start_col: usize,
end_col: usize,
}
#[derive(Debug, Clone, Serialize)]
struct MatchLocation {
text: String,
line_num: usize,
submatches: Vec<SubMatch>,
}