1
0
mirror of https://github.com/BurntSushi/ripgrep.git synced 2024-12-12 19:18:24 +02:00
ripgrep/src/main.rs

281 lines
6.9 KiB
Rust
Raw Normal View History

2016-03-11 03:48:44 +02:00
#![allow(dead_code, unused_variables)]
2016-02-27 18:07:26 +02:00
extern crate crossbeam;
2016-02-27 18:07:26 +02:00
extern crate docopt;
extern crate env_logger;
2016-06-20 22:53:48 +02:00
extern crate grep;
2016-09-05 23:36:41 +02:00
#[cfg(windows)]
extern crate kernel32;
2016-09-04 03:48:23 +02:00
#[macro_use]
extern crate lazy_static;
2016-09-05 23:36:41 +02:00
extern crate libc;
#[macro_use]
extern crate log;
2016-03-11 03:48:44 +02:00
extern crate memchr;
extern crate memmap;
2016-08-26 03:44:37 +02:00
extern crate num_cpus;
2016-02-27 18:07:26 +02:00
extern crate regex;
2016-03-11 03:48:44 +02:00
extern crate regex_syntax as syntax;
2016-02-27 18:07:26 +02:00
extern crate rustc_serialize;
2016-09-05 23:36:41 +02:00
extern crate term;
extern crate thread_local;
2016-08-26 03:44:37 +02:00
extern crate walkdir;
2016-09-05 23:36:41 +02:00
#[cfg(windows)]
extern crate winapi;
2016-02-27 18:07:26 +02:00
use std::error::Error;
use std::fs::File;
2016-09-05 23:36:41 +02:00
use std::io;
2016-09-05 16:15:13 +02:00
use std::path::Path;
2016-02-27 18:07:26 +02:00
use std::process;
use std::result;
use std::sync::{Arc, Mutex};
use std::thread;
2016-02-27 18:07:26 +02:00
use crossbeam::sync::chase_lev::{self, Steal, Stealer};
use grep::Grep;
use memmap::{Mmap, Protection};
use term::Terminal;
2016-09-05 16:15:13 +02:00
use walkdir::DirEntry;
2016-08-26 03:44:37 +02:00
use args::Args;
use out::{NoColorTerminal, Out, OutBuffer};
use printer::Printer;
use search::InputBuffer;
2016-08-26 03:44:37 +02:00
macro_rules! errored {
($($tt:tt)*) => {
return Err(From::from(format!($($tt)*)));
}
}
2016-02-27 18:07:26 +02:00
2016-08-26 03:44:37 +02:00
macro_rules! eprintln {
($($tt:tt)*) => {{
use std::io::Write;
let _ = writeln!(&mut ::std::io::stderr(), $($tt)*);
}}
}
mod args;
mod gitignore;
2016-08-26 03:44:37 +02:00
mod glob;
mod ignore;
mod out;
mod printer;
mod search;
mod search_buffer;
2016-09-05 23:36:41 +02:00
mod sys;
2016-09-06 00:22:12 +02:00
mod terminal;
mod types;
mod walk;
2016-03-11 03:48:44 +02:00
pub type Result<T> = result::Result<T, Box<Error + Send + Sync>>;
2016-02-27 18:07:26 +02:00
fn main() {
match Args::parse().and_then(run) {
Ok(count) if count == 0 => process::exit(1),
Ok(count) => process::exit(0),
2016-02-27 18:07:26 +02:00
Err(err) => {
2016-09-05 23:36:41 +02:00
eprintln!("{}", err);
2016-02-27 18:07:26 +02:00
process::exit(1);
}
}
}
fn run(args: Args) -> Result<u64> {
if args.files() {
return run_files(args);
}
if args.type_list() {
return run_types(args);
}
let args = Arc::new(args);
let out = Arc::new(Mutex::new(args.out()));
let outbuf = args.outbuf();
let mut workers = vec![];
let mut workq = {
let (workq, stealer) = chase_lev::deque();
for _ in 0..args.threads() {
let worker = Worker {
args: args.clone(),
2016-09-04 03:48:23 +02:00
out: out.clone(),
chan_work: stealer.clone(),
inpbuf: args.input_buffer(),
outbuf: Some(outbuf.clone()),
2016-09-07 01:33:19 +02:00
grep: args.grep(),
2016-09-05 16:15:13 +02:00
match_count: 0,
};
workers.push(thread::spawn(move || worker.run()));
}
workq
};
for p in args.paths() {
if p == Path::new("-") {
workq.push(Work::Stdin)
} else {
2016-09-05 23:36:41 +02:00
for ent in try!(args.walker(p)) {
2016-09-05 16:15:13 +02:00
workq.push(Work::File(ent));
}
2016-08-05 06:10:58 +02:00
}
2016-03-11 03:48:44 +02:00
}
for _ in 0..workers.len() {
workq.push(Work::Quit);
}
let mut match_count = 0;
for worker in workers {
match_count += worker.join().unwrap();
}
Ok(match_count)
}
2016-08-26 03:44:37 +02:00
fn run_files(args: Args) -> Result<u64> {
let term = NoColorTerminal::new(io::BufWriter::new(io::stdout()));
let mut printer = args.printer(term);
let mut file_count = 0;
for p in args.paths() {
if p == Path::new("-") {
printer.path(&Path::new("<stdin>"));
file_count += 1;
2016-09-04 03:48:23 +02:00
} else {
2016-09-05 23:36:41 +02:00
for ent in try!(args.walker(p)) {
2016-09-05 16:15:13 +02:00
printer.path(ent.path());
file_count += 1;
}
2016-09-04 03:48:23 +02:00
}
}
Ok(file_count)
}
2016-09-04 03:48:23 +02:00
fn run_types(args: Args) -> Result<u64> {
let term = NoColorTerminal::new(io::BufWriter::new(io::stdout()));
let mut printer = args.printer(term);
let mut ty_count = 0;
for def in args.type_defs() {
printer.type_def(def);
ty_count += 1;
2016-09-04 03:48:23 +02:00
}
Ok(ty_count)
}
enum Work {
Stdin,
2016-09-05 16:15:13 +02:00
File(DirEntry),
Quit,
}
2016-09-05 16:15:13 +02:00
enum WorkReady {
Stdin,
File(DirEntry, File),
}
struct Worker {
args: Arc<Args>,
out: Arc<Mutex<Out>>,
chan_work: Stealer<Work>,
inpbuf: InputBuffer,
outbuf: Option<OutBuffer>,
grep: Grep,
2016-09-05 16:15:13 +02:00
match_count: u64,
}
impl Worker {
fn run(mut self) -> u64 {
2016-09-05 16:15:13 +02:00
self.match_count = 0;
loop {
2016-09-05 16:15:13 +02:00
let work = match self.chan_work.steal() {
Steal::Empty | Steal::Abort => continue,
Steal::Data(Work::Quit) => break,
2016-09-05 16:15:13 +02:00
Steal::Data(Work::Stdin) => WorkReady::Stdin,
Steal::Data(Work::File(ent)) => {
match File::open(ent.path()) {
Ok(file) => WorkReady::File(ent, file),
Err(err) => {
2016-09-05 16:15:13 +02:00
eprintln!("{}: {}", ent.path().display(), err);
continue;
}
}
}
};
let mut outbuf = self.outbuf.take().unwrap();
outbuf.clear();
let mut printer = self.args.printer(outbuf);
2016-09-05 16:15:13 +02:00
self.do_work(&mut printer, work);
let outbuf = printer.into_inner();
if !outbuf.get_ref().is_empty() {
let mut out = self.out.lock().unwrap();
out.write(&outbuf);
}
self.outbuf = Some(outbuf);
}
2016-09-05 16:15:13 +02:00
self.match_count
}
fn do_work<W: Send + Terminal>(
2016-09-05 16:15:13 +02:00
&mut self,
printer: &mut Printer<W>,
work: WorkReady,
) {
let result = match work {
WorkReady::Stdin => {
let stdin = io::stdin();
let stdin = stdin.lock();
self.search(printer, &Path::new("<stdin>"), stdin)
}
WorkReady::File(ent, file) => {
let mut path = ent.path();
if let Ok(p) = path.strip_prefix("./") {
path = p;
}
if self.args.mmap() {
self.search_mmap(printer, path, &file)
} else {
self.search(printer, path, file)
}
2016-09-05 16:15:13 +02:00
}
};
match result {
Ok(count) => {
self.match_count += count;
}
Err(err) => {
eprintln!("{}", err);
}
}
}
fn search<R: io::Read, W: Send + Terminal>(
&mut self,
printer: &mut Printer<W>,
path: &Path,
rdr: R,
) -> Result<u64> {
self.args.searcher(
&mut self.inpbuf,
printer,
&self.grep,
path,
rdr,
).run().map_err(From::from)
2016-09-04 03:48:23 +02:00
}
fn search_mmap<W: Send + Terminal>(
&mut self,
printer: &mut Printer<W>,
path: &Path,
file: &File,
) -> Result<u64> {
if try!(file.metadata()).len() == 0 {
// Opening a memory map with an empty file results in an error.
return Ok(0);
}
let mmap = try!(Mmap::open(file, Protection::Read));
Ok(self.args.searcher_buffer(
printer,
&self.grep,
path,
unsafe { mmap.as_slice() },
).run())
}
2016-09-04 03:48:23 +02:00
}