Merge pull request #2 from wheals/master
[rust-101.git] / 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 }