1use std::env;
2use std::fs::File;
3use std::io::{self, BufRead, BufReader};
4
5type MakeReader = dyn Fn() -> io::Result<Box<dyn BufRead>>;
6
7struct FileInput {
8 readers: Vec<Box<MakeReader>>,
10 current_reader_idx: usize,
11 current_reader: Option<Box<dyn BufRead>>,
12 buffer: String,
13}
14
15impl FileInput {
16 pub fn new(parameter_name: &str) -> io::Result<Self> {
17 let readers = Self::parse_args(parameter_name)?;
18 Ok(FileInput {
19 readers,
20 current_reader_idx: 0,
21 current_reader: None,
22 buffer: String::new(),
23 })
24 }
25
26 fn parse_args(parameter_name: &str) -> io::Result<Vec<Box<MakeReader>>> {
27 let args: Vec<String> = env::args().collect();
28 let file_args = Self::extract_file_args(parameter_name, &args);
29 if file_args.is_empty() {
30 Ok(vec![Box::new(|| Ok(Box::new(BufReader::new(io::stdin())) as Box<dyn BufRead>))])
31 } else {
32 Ok(file_args.into_iter().map(|arg| {
33 Box::new(move || {
34 let file = File::open(&arg)?;
35 Ok(Box::new(BufReader::new(file)) as Box<dyn BufRead>)
36 }) as Box<dyn Fn() -> io::Result<Box<dyn BufRead>>>
37 }).collect())
38 }
39 }
40
41 fn extract_file_args(parameter_name: &str, args: &[String]) -> Vec<String> {
42 let mut file_args = Vec::new();
43 let mut reading_files = false;
44 for arg in args.iter().skip(1) {
45 if *arg == format!("--{parameter_name}") {
46 reading_files = true;
47 } else if reading_files && arg.starts_with('-') {
48 break;
49 } else if reading_files {
50 file_args.push(arg.clone());
51 }
52 }
53 file_args
54 }
55}
56
57impl Iterator for FileInput {
58 type Item = io::Result<String>;
59
60 fn next(&mut self) -> Option<Self::Item> {
61 loop {
62 if self.current_reader_idx >= self.readers.len() {
63 return None;
64 }
65
66 if self.current_reader.is_none() {
67 match self.readers[self.current_reader_idx]() {
68 Ok(reader) => self.current_reader = Some(reader),
69 Err(e) => return Some(Err(e)),
70 }
71 }
72
73 self.buffer.clear();
74 let reader = match self.current_reader.as_mut() {
75 Some(reader) => reader,
76 None => return Some(Err(io::Error::new(io::ErrorKind::Other, "reader should be Some"))),
77 };
78
79 match reader.read_line(&mut self.buffer) {
80 Ok(0) => {
81 self.current_reader_idx += 1;
82 self.current_reader = None;
83 continue;
84 }
85 Ok(_) => return Some(Ok(self.buffer.trim_end().to_string())),
86 Err(e) => return Some(Err(e)),
87 }
88 }
89 }
90}
91
92pub fn input() -> io::Result<impl Iterator<Item = io::Result<String>>> {
93 FileInput::new("files")
94}