solutions/src/rgrep.rs

   1 use std::io::prelude::*;
   2 use std::{io, fs, thread, process, cmp};
   3 use std::sync::mpsc::{sync_channel, SyncSender, Receiver};
   4 use std::sync::Arc;
   5
   6 #[derive(Clone,Copy)]
   7 enum OutputMode {
   8     Print,
   9     SortAndPrint,
  10     Count,
  11 }
  12 use self::OutputMode::*;
  13
  14 struct Options {
  15     files: Vec<String>,
  16     pattern: String,
  17     output_mode: OutputMode,
  18 }
  19
  20 struct Line {
  21     data: String,
  22     file: usize,
  23     line: usize,
  24 }
  25
  26 impl PartialEq for Line {
  27     fn eq(&self, other: &Line) -> bool {
  28         self.data.eq(&other.data)
  29     }
  30 }
  31 impl PartialOrd for Line {
  32     fn partial_cmp(&self, other: &Line) -> Option<cmp::Ordering> {
  33         self.data.partial_cmp(&other.data)
  34     }
  35 }
  36
  37 fn read_files(options: Arc<Options>, out_channel: SyncSender<Line>) {
  38     for (fileidx, file) in options.files.iter().enumerate() {
  39         let file = fs::File::open(file).unwrap();
  40         let file = io::BufReader::new(file);
  41         for (lineidx, line) in file.lines().enumerate() {
  42             let line = Line { data: line.unwrap(), file: fileidx, line: lineidx };
  43             out_channel.send(line).unwrap();
  44         }
  45     }
  46 }
  47
  48 fn filter_lines(options: Arc<Options>, in_channel: Receiver<Line>, out_channel: SyncSender<Line>) {
  49     for line in in_channel.iter() {
  50         if line.data.contains(&options.pattern) {
  51             out_channel.send(line).unwrap();
  52         }
  53     }
  54 }
  55
  56 fn sort<T: PartialOrd>(data: &mut [T]) {
  57     if data.len() < 2 { return; }
  58
  59     let mut lpos = 1;
  60     let mut rpos = data.len();
  61     // Invariant: pivot is data[0]; (0,lpos) is <= pivot; [rpos,len) is >= pivot; lpos < rpos
  62     loop {
  63         while lpos < rpos && data[lpos] <= data[0] {
  64             lpos += 1;
  65         }
  66         while rpos > lpos && data[rpos-1] >= data[0] {
  67             rpos -= 1;
  68         }
  69         if rpos == lpos {
  70             break;
  71         }
  72
  73         data.swap(lpos, rpos-1);
  74     }
  75
  76     data.swap(0, lpos-1); // put pivot in the right place
  77
  78     let (part1, part2) = data.split_at_mut(lpos);
  79     sort(&mut part1[..lpos-1]);
  80     sort(part2);
  81 }
  82
  83 fn output_lines(options: Arc<Options>, in_channel: Receiver<Line>) {
  84     match options.output_mode {
  85         Print => {
  86             for line in in_channel.iter() {
  87                 println!("{}:{}: {}", options.files[line.file], line.line, line.data);
  88             }
  89         },
  90         Count => {
  91             let count = in_channel.iter().count();
  92             println!("{} hits for {}.", count, options.pattern);
  93         },
  94         SortAndPrint => {
  95             let mut data: Vec<Line> = in_channel.iter().collect();
  96             sort(&mut data[..]);
  97             for line in data.iter() {
  98                 println!("{}:{}: {}", options.files[line.file], line.line, line.data);
  99             }
 100         }
 101     }
 102 }
 103
 104 static USAGE: &'static str = "
 105 Usage: rgrep [-c] [-s] <pattern> <file>...
 106
 107 Options:
 108     -c, --count  Count number of matching lines (rather than printing them).
 109     -s, --sort   Sort the lines before printing.
 110 ";
 111
 112 fn get_options() -> Options {
 113     use docopt::Docopt;
 114
 115     // Parse argv and exit the program with an error message if it fails.
 116     let args = Docopt::new(USAGE).and_then(|d| d.parse()).unwrap_or_else(|e| e.exit());
 117     let count = args.get_bool("-c");
 118     let sort = args.get_bool("-s");
 119     let pattern = args.get_str("<pattern>");
 120     let files = args.get_vec("<file>");
 121     if count && sort {
 122         println!("Setting both '-c' and '-s' at the same time does not make any sense.");
 123         process::exit(1);
 124     }
 125
 126     // We need to make the strings owned to construct the `Options` instance.
 127     Options {
 128         files: files.iter().map(|file| file.to_string()).collect(),
 129         pattern: pattern.to_string(),
 130         output_mode: if count { Count } else if sort { SortAndPrint } else { Print },
 131     }
 132 }
 133
 134 fn run(options: Options) {
 135     let options = Arc::new(options);
 136
 137     // This sets up the chain of threads. Use `sync_channel` with buffer-size of 16 to avoid needlessly filling RAM.
 138     let (line_sender, line_receiver) = sync_channel(16);
 139     let (filtered_sender, filtered_receiver) = sync_channel(16);
 140
 141     let options1 = options.clone();
 142     let handle1 = thread::spawn(move || read_files(options1, line_sender));
 143     let options2 = options.clone();
 144     let handle2 = thread::spawn(move || filter_lines(options2, line_receiver, filtered_sender));
 145     let options3 = options.clone();
 146     let handle3 = thread::spawn(move || output_lines(options3, filtered_receiver));
 147     handle1.join().unwrap();
 148     handle2.join().unwrap();
 149     handle3.join().unwrap();
 150 }
 151
 152 pub fn main() {
 153     run(get_options());
 154 }