mod pattern; use walkdir::WalkDir; #[derive(Debug)] struct Args { targets: Vec, num_threads: u8, } impl Args { fn parse() -> Args { let matches = clap::App::new("rubygrep") .version("0.1") .author("Getty Ritter ") .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::>() } 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 p = pattern::Pattern { node_type: pattern::NodeType::Send, name: Some("ActiveRecord".to_owned()), children: vec![], rest: false, }; let (work_send, work_recv) = crossbeam::channel::unbounded(); let (result_send, result_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.into_path()).expect("Unable to send work from producer thread"); } } } }); // produce a set of threads which can grab work to be done let workers: Vec> = (0..num_threads).map(|id| { let receiver = work_recv.clone(); let sender = result_send.clone(); let pat = p.clone(); std::thread::spawn(move || { while let Ok(ref msg) = receiver.recv() { use std::io::Read; let mut buf = Vec::new(); { let mut f = std::fs::File::open(msg).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 { let matches = pat.find_matches(&*ast); if matches.len() > 0 { sender.send(format!("{:?}: found matches: {:?}", msg, matches)).unwrap(); } } else { // sender.send(format!("Unable to parse {:?}", msg)).unwrap(); } } }) }).collect(); drop(result_send); while let Ok(msg) = result_recv.recv() { println!("{}", msg); } // join all the threads producer.join().expect("Producer thread panicked!"); for w in workers { w.join().expect("Worker thread panicked!"); } }