competitive_programming_rs/utils/
scanner.rs

1pub struct IO<R, W: std::io::Write>(R, std::io::BufWriter<W>);
2
3impl<R: std::io::Read, W: std::io::Write> IO<R, W> {
4    pub fn new(r: R, w: W) -> IO<R, W> {
5        IO(r, std::io::BufWriter::new(w))
6    }
7    pub fn write<S: ToString>(&mut self, s: S) {
8        use std::io::Write;
9        self.1.write_all(s.to_string().as_bytes()).unwrap();
10    }
11    pub fn read<T: std::str::FromStr>(&mut self) -> T {
12        use std::io::Read;
13        let buf = self
14            .0
15            .by_ref()
16            .bytes()
17            .map(|b| b.unwrap())
18            .skip_while(|&b| b == b' ' || b == b'\n' || b == b'\r' || b == b'\t')
19            .take_while(|&b| b != b' ' && b != b'\n' && b != b'\r' && b != b'\t')
20            .collect::<Vec<_>>();
21        unsafe { std::str::from_utf8_unchecked(&buf) }
22            .parse()
23            .ok()
24            .expect("Parse error.")
25    }
26    pub fn usize0(&mut self) -> usize {
27        self.read::<usize>() - 1
28    }
29    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
30        (0..n).map(|_| self.read()).collect()
31    }
32    pub fn chars(&mut self) -> Vec<char> {
33        self.read::<String>().chars().collect()
34    }
35}
36
37#[cfg(test)]
38mod tests {
39    use super::*;
40
41    #[test]
42    fn test_io() {
43        let input = br"10
44        1 2 3 4 5 6 7 8 9 10
45        abcde fghij
46        
47        3.14 -1592
48        no_empty_line";
49        let mut sc = IO::new(&input[..], Vec::new());
50
51        let n: usize = sc.read();
52        assert_eq!(n, 10);
53
54        let a: Vec<u64> = sc.vec(n);
55        assert_eq!(&a, &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
56
57        let s = sc.chars();
58        let t = sc.chars();
59        assert_eq!(&s, &['a', 'b', 'c', 'd', 'e']);
60        assert_eq!(&t, &['f', 'g', 'h', 'i', 'j']);
61
62        let f: f64 = sc.read();
63        assert_eq!(f, 3.14);
64
65        let neg: i64 = sc.read();
66        assert_eq!(neg, -1592);
67
68        let s = sc.read::<String>();
69        assert_eq!(&s, "no_empty_line");
70        sc.write(format!("1\n"));
71
72        let mut output = Vec::new();
73        {
74            let mut sc = IO::new(&b""[..], &mut output);
75            sc.write(format!("{}\n", 1));
76            sc.write(format!("{}\n", 2));
77            sc.write(format!("{}\n", 3));
78        }
79
80        let output = String::from_utf8(output).expect("Not UTF-8");
81        assert_eq!(&output, "1\n2\n3\n");
82    }
83}