linedance/
lib.rs

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<dyn Fn() -> io::Result<Box<dyn BufRead>>>>,
9    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}