123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106 |
- use walkdir::WalkDir;
- #[derive(Debug)]
- struct Args {
- targets: Vec<String>,
- num_threads: u8,
- }
- impl Args {
- fn parse() -> Args {
- let matches = clap::App::new("rubygrep")
- .version("0.1")
- .author("Getty Ritter <rubygrep@infinitenegativeutility.com>")
- .about("Search Ruby source trees")
- .arg(clap::Arg::with_name("threads")
- .short("j")
- .value_name("THREADS")
- .help("Spawn the specified number of worker threads")
- .takes_value(true))
- .arg(clap::Arg::with_name("targets")
- .help("The Ruby sources and directories to search")
- .multiple(true))
- .get_matches();
- let targets = if let Some(values) = matches.values_of("targets") {
- values.map(|x| x.to_owned()).collect::<Vec<String>>()
- } else {
- vec![".".to_owned()]
- };
- let num_threads = if let Some(val) = matches.value_of("threads") {
- match val.parse() {
- Ok(x) => x,
- Err(_) => {
- panic!("Invalid number of threads: {}", val);
- }
- }
- } else {
- 8
- };
- Args { targets, num_threads }
- }
- }
- fn is_ruby_source(path: &std::path::Path) -> bool {
- if let Some(ext) = path.extension() {
- return ext == "rb";
- }
- false
- }
- fn main() {
- let args = Args::parse();
- println!("Got args: {:?}", args);
- let Args {
- targets,
- num_threads,
- } = args;
- let (work_send, work_recv) = crossbeam::channel::unbounded();
- // produce a thread which populates the channel with work
- let producer = std::thread::spawn(move || {
- for target in targets {
- for entry in WalkDir::new(target).into_iter().filter_map(|e| e.ok()) {
- if is_ruby_source(entry.path()) {
- work_send.send(entry).expect("Unable to send work from producer thread");
- }
- }
- }
- });
- // produce a set of threads which can grab work to be done
- let workers: Vec<std::thread::JoinHandle<()>> = (0..num_threads).map(|id| {
- let receiver = work_recv.clone();
- std::thread::spawn(move || {
- while let Ok(msg) = receiver.recv() {
- use std::io::Read;
- // println!("thread {} processing {:?}", id, msg);
- let mut buf = Vec::new();
- {
- let mut f = std::fs::File::open(msg.path()).expect("Unable to read file");
- f.read_to_end(&mut buf).expect("Unable to read file");
- }
- let parser = lib_ruby_parser::Parser::new(buf, std::default::Default::default());
- let lib_ruby_parser::ParserResult { ast, diagnostics, .. } = parser.do_parse();
- if let Some(_ast) = ast {
- // println!("Got ast: {:?}", ast);
- } else {
- println!("Unable to parse {:?}", msg.path());
- for d in diagnostics {
- println!(" - {:?}", d);
- }
- }
- }
- println!("thread {} done", id);
- })
- }).collect();
- // join all the threads
- producer.join().expect("Producer thread panicked!");
- for w in workers {
- w.join().expect("Worker thread panicked!");
- }
- }
|