scanner/
lib.rs

1use std::{
2    fmt,
3    io::{self, BufReader, Cursor},
4    str,
5};
6
7pub struct Scanner<R>
8where
9    R: io::BufRead,
10{
11    reader: R,
12    buf: String,
13    pos: usize,
14}
15
16impl Scanner<BufReader<io::StdinLock<'static>>> {
17    /// Creates a scanner that reads from standard input.
18    pub fn stdin_lock() -> Self {
19        Self {
20            reader: BufReader::new(io::stdin().lock()),
21            buf: String::new(),
22            pos: 0,
23        }
24    }
25}
26
27impl<T> Scanner<Cursor<T>>
28where
29    T: AsRef<[u8]>,
30{
31    /// Creates a scanner that reads from a string or byte slice.
32    pub fn cursor(inner: T) -> Self {
33        Self {
34            reader: Cursor::new(inner),
35            buf: String::new(),
36            pos: 0,
37        }
38    }
39}
40
41impl<R> Scanner<R>
42where
43    R: io::BufRead,
44{
45    /// Scans and parses the next token from the input.
46    ///
47    /// For more convenient input scanning with variable declarations, see [`scan!`].
48    ///
49    /// # Examples
50    ///
51    /// ```
52    /// use scanner::Scanner;
53    ///
54    /// let mut scanner = Scanner::cursor("-10 20");
55    ///
56    /// let x = scanner.scan::<i32>();
57    /// let y = scanner.scan::<i32>();
58    ///
59    /// assert_eq!(x, -10);
60    /// assert_eq!(y, 20);
61    /// ```
62    pub fn scan<T>(&mut self) -> T
63    where
64        T: str::FromStr,
65        T::Err: fmt::Debug,
66    {
67        // skip whitespace
68        loop {
69            match self.buf[self.pos..].find(|ch| !char::is_ascii_whitespace(&ch)) {
70                Some(j) => {
71                    self.pos += j;
72                    break;
73                }
74                None => {
75                    let num_bytes = self
76                        .reader
77                        .read_line(&mut self.buf)
78                        .unwrap_or_else(|_| panic!("invalid UTF-8"));
79                    assert!(num_bytes > 0, "reached EOF :(");
80                }
81            }
82        }
83
84        let rest = &self.buf[self.pos..];
85        let token_len = rest
86            .find(|ch| char::is_ascii_whitespace(&ch))
87            .unwrap_or(rest.len());
88        let value = rest[..token_len]
89            .parse()
90            .unwrap_or_else(|e| panic!("{:?}, attempt to read `{}`", e, rest));
91        self.pos += token_len;
92
93        value
94    }
95}
96
97/// Macro for convenient input scanning with variable declarations.
98///
99/// For direct token scanning, see [`Scanner::scan()`].
100///
101/// # Special Types
102///
103/// - `Usize1`: Converts 1-indexed input to 0-indexed (subtracts 1)
104/// - `Chars`: Reads a string and converts it to `Vec<char>`
105/// - `Bytes`: Reads a string and converts it to `Vec<u8>`
106///
107/// # Examples
108///
109/// ## Basic usage
110///
111/// ```
112/// use scanner::{Scanner, scan};
113///
114/// let mut scanner = Scanner::cursor("3 10\n1 2 3");
115///
116/// scan! {
117///     via scanner,
118///     (n, k): (usize, usize),
119///     a: [i32; n],
120/// };
121///
122/// assert_eq!((n, k), (3, 10));
123/// assert_eq!(a, vec![1, 2, 3]);
124/// ```
125///
126/// ## Special types
127///
128/// ```
129/// use scanner::{Scanner, scan};
130///
131/// let mut scanner = Scanner::cursor("5\n1 3\nabcde\nABCDE");
132///
133/// scan! {
134///     via scanner,
135///     n: usize,
136///     (l, r): (Usize1, Usize1),
137///     s: Chars,
138///     t: Bytes,
139/// };
140///
141/// assert_eq!((l, r), (0, 2));
142/// assert_eq!(s, vec!['a', 'b', 'c', 'd', 'e']);
143/// assert_eq!(t, vec![b'A', b'B', b'C', b'D', b'E']);
144/// ```
145#[macro_export]
146macro_rules! scan {
147    (via $scanner:expr, $($rest:tt)*) => {
148        $crate::scan!(@via [$scanner] @rest $($rest)*);
149    };
150
151    (@via [$via:expr] @rest) => {};
152    (@via [$via:expr] @rest ,) => {};
153
154    (@via [$via:expr] @rest mut $($rest:tt)*) => {
155        $crate::scan!(@via [$via] @mut [mut] @rest $($rest)*);
156    };
157    (@via [$via:expr] @rest $($rest:tt)*) => {
158        $crate::scan!(@via [$via] @mut [] @rest $($rest)*);
159    };
160
161    (@via [$via:expr] @mut [$($mut:tt)?] @rest $var:tt : $t:tt) => {
162        let $($mut)? $var = $crate::scan_inner!(via $via, $t);
163    };
164    (@via [$via:expr] @mut [$($mut:tt)?] @rest $var:tt : $t:tt , $($rest:tt)*) => {
165        $crate::scan!(@via [$via] @mut [$($mut)?] @rest $var : $t);
166        $crate::scan!(@via [$via] @rest $($rest)*);
167    };
168}
169
170#[doc(hidden)]
171#[macro_export]
172macro_rules! scan_inner {
173    // (i32, i32)
174    (via $scanner:expr, ( $($t:tt),* )) => {
175        ( $($crate::scan_inner!(via $scanner, $t)),* )
176    };
177
178    // [i32; n]
179    (via $scanner:expr, [ $t:tt ; $len:expr ]) => {
180        ::std::iter::repeat_with(|| $crate::scan_inner!(via $scanner, $t)).take($len).collect::<Vec<_>>()
181    };
182
183    (via $scanner:expr, Usize1) => {
184        $scanner.scan::<usize>().checked_sub(1).expect("Usize1: input was 0, expected >= 1")
185    };
186
187    (via $scanner:expr, Chars) => {
188        $scanner.scan::<String>().chars().collect::<Vec<_>>()
189    };
190
191    (via $scanner:expr, Bytes) => {
192        $scanner.scan::<String>().bytes().collect::<Vec<_>>()
193    };
194
195    // i32
196    (via $scanner:expr, $ty:ty) => {
197        $scanner.scan::<$ty>()
198    };
199}
200
201#[cfg(test)]
202mod tests {
203    use crate::Scanner;
204
205    #[test]
206    fn scan_test() {
207        let mut scanner = Scanner::cursor("42 123\n456\r\nABC");
208        assert_eq!(scanner.scan::<i32>(), 42);
209        assert_eq!(scanner.scan::<i32>(), 123);
210        assert_eq!(scanner.scan::<i32>(), 456);
211        assert_eq!(scanner.scan::<String>(), String::from("ABC"));
212    }
213
214    #[test]
215    fn scan_macro_test() {
216        let mut scanner = Scanner::cursor(
217            r#"
2183 10
2194
2201 2 3
221a 1
222b 2
223c 3
224d 4
225        "#,
226        );
227        scan! {
228            via scanner,
229            (n, k): (usize, usize),
230            mut q: usize,
231            a: [i32; n],
232            queries: [(char, i32); q],
233        };
234
235        assert_eq!((n, k), (3, 10));
236        assert_eq!(q, 4);
237        // test mutable
238        q += 1;
239        assert_eq!(q, 5);
240        assert_eq!(a, vec![1, 2, 3]);
241        assert_eq!(queries, vec![('a', 1), ('b', 2), ('c', 3), ('d', 4)]);
242    }
243
244    #[test]
245    fn test_special_type() {
246        let mut scanner = Scanner::cursor(
247            r#"
2481 10
249abc
250XYZ
251        "#,
252        );
253
254        scan! {
255            via scanner,
256            (l, r): (Usize1, Usize1),
257            s: Chars,
258            t: Bytes,
259        };
260
261        assert_eq!((l, r), (0, 9));
262        assert_eq!(s, vec!['a', 'b', 'c']);
263        assert_eq!(t, vec![b'X', b'Y', b'Z']);
264    }
265}
scanner/lib.rs

scanner/
lib.rs
