rubygrep/src/main.rs

mod pattern;

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 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<std::thread::JoinHandle<()>> = (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!");
    }
}