competitive/tools/
scanner.rs

1use std::{
2    iter::{FromIterator, from_fn, repeat_with},
3    marker::PhantomData,
4};
5
6pub fn read_stdin_all() -> String {
7    use std::io::Read as _;
8    let mut s = String::new();
9    std::io::stdin().read_to_string(&mut s).expect("io error");
10    s
11}
12pub fn read_stdin_all_unchecked() -> String {
13    use std::io::Read as _;
14    let mut buf = Vec::new();
15    std::io::stdin().read_to_end(&mut buf).expect("io error");
16    unsafe { String::from_utf8_unchecked(buf) }
17}
18pub fn read_all(mut reader: impl std::io::Read) -> String {
19    let mut s = String::new();
20    reader.read_to_string(&mut s).expect("io error");
21    s
22}
23pub fn read_all_unchecked(mut reader: impl std::io::Read) -> String {
24    let mut buf = Vec::new();
25    reader.read_to_end(&mut buf).expect("io error");
26    unsafe { String::from_utf8_unchecked(buf) }
27}
28pub fn read_stdin_line() -> String {
29    let mut s = String::new();
30    std::io::stdin().read_line(&mut s).expect("io error");
31    s
32}
33pub trait IterScan: Sized {
34    type Output;
35    fn scan<'a, I: Iterator<Item = &'a str>>(iter: &mut I) -> Option<Self::Output>;
36}
37pub trait MarkedIterScan: Sized {
38    type Output;
39    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output>;
40}
41#[derive(Clone, Debug)]
42pub struct Scanner<'a> {
43    iter: std::str::SplitAsciiWhitespace<'a>,
44}
45impl<'a> Scanner<'a> {
46    #[inline]
47    pub fn new(s: &'a str) -> Self {
48        let iter = s.split_ascii_whitespace();
49        Self { iter }
50    }
51    #[inline]
52    pub fn scan<T>(&mut self) -> <T as IterScan>::Output
53    where
54        T: IterScan,
55    {
56        <T as IterScan>::scan(&mut self.iter).expect("scan error")
57    }
58    #[inline]
59    pub fn mscan<T>(&mut self, marker: T) -> <T as MarkedIterScan>::Output
60    where
61        T: MarkedIterScan,
62    {
63        marker.mscan(&mut self.iter).expect("scan error")
64    }
65    #[inline]
66    pub fn scan_vec<T>(&mut self, size: usize) -> Vec<<T as IterScan>::Output>
67    where
68        T: IterScan,
69    {
70        (0..size)
71            .map(|_| <T as IterScan>::scan(&mut self.iter).expect("scan error"))
72            .collect()
73    }
74    #[inline]
75    pub fn iter<'b, T>(&'b mut self) -> ScannerIter<'a, 'b, T>
76    where
77        T: IterScan,
78    {
79        ScannerIter {
80            inner: self,
81            _marker: std::marker::PhantomData,
82        }
83    }
84}
85
86macro_rules! impl_iter_scan {
87    ($($t:ty)*) => {$(
88        impl IterScan for $t {
89            type Output = Self;
90            #[inline]
91            fn scan<'a, I: Iterator<Item = &'a str>>(iter: &mut I) -> Option<Self> {
92                iter.next()?.parse::<$t>().ok()
93            }
94        })*
95    };
96}
97impl_iter_scan!(char u8 u16 u32 u64 usize i8 i16 i32 i64 isize f32 f64 u128 i128 String);
98
99macro_rules! impl_iter_scan_tuple {
100    (@impl $($T:ident)*) => {
101        impl<$($T: IterScan),*> IterScan for ($($T,)*) {
102            type Output = ($(<$T as IterScan>::Output,)*);
103            #[inline]
104            fn scan<'a, It: Iterator<Item = &'a str>>(_iter: &mut It) -> Option<Self::Output> {
105                Some(($(<$T as IterScan>::scan(_iter)?,)*))
106            }
107        }
108    };
109    (@inner $($T:ident)*,) => {
110        impl_iter_scan_tuple!(@impl $($T)*);
111    };
112    (@inner $($T:ident)*, $U:ident $($Rest:ident)*) => {
113        impl_iter_scan_tuple!(@impl $($T)*);
114        impl_iter_scan_tuple!(@inner $($T)* $U, $($Rest)*);
115    };
116    ($($T:ident)*) => {
117        impl_iter_scan_tuple!(@inner , $($T)*);
118    };
119}
120impl_iter_scan_tuple!(A B C D E F G H I J K);
121
122pub struct ScannerIter<'a, 'b, T> {
123    inner: &'b mut Scanner<'a>,
124    _marker: std::marker::PhantomData<fn() -> T>,
125}
126impl<T> Iterator for ScannerIter<'_, '_, T>
127where
128    T: IterScan,
129{
130    type Item = <T as IterScan>::Output;
131    #[inline]
132    fn next(&mut self) -> Option<Self::Item> {
133        <T as IterScan>::scan(&mut self.inner.iter)
134    }
135}
136
137/// scan a value with Scanner
138///
139/// - `scan_value!(scanner, ELEMENT)`
140///
141/// ELEMENT :=
142/// - `$ty`: IterScan
143/// - `@$expr`: MarkedIterScan
144/// - `[ELEMENT; $expr]`: vector
145/// - `[ELEMENT; const $expr]`: array
146/// - `[ELEMENT]`: iterator
147/// - `($(ELEMENT)*,)`: tuple
148#[macro_export]
149macro_rules! scan_value {
150    (@repeat $scanner:expr, [$($t:tt)*] $($len:expr)?)                                    => { ::std::iter::repeat_with(|| $crate::scan_value!(@inner $scanner, [] $($t)*)) $(.take($len).collect::<Vec<_>>())? };
151    (@array $scanner:expr, [$($t:tt)*] $len:expr)                                         => { $crate::array![|| $crate::scan_value!(@inner $scanner, [] $($t)*); $len] };
152    (@tuple $scanner:expr, [$([$($args:tt)*])*])                                          => { ($($($args)*,)*) };
153    (@$tag:ident $scanner:expr, [[$($args:tt)*]])                                         => { $($args)* };
154    (@$tag:ident $scanner:expr, [$($args:tt)*] @$e:expr)                                  => { $crate::scan_value!(@$tag $scanner, [$($args)* [$scanner.mscan($e)]]) };
155    (@$tag:ident $scanner:expr, [$($args:tt)*] @$e:expr, $($t:tt)*)                       => { $crate::scan_value!(@$tag $scanner, [$($args)* [$scanner.mscan($e)]] $($t)*) };
156    (@$tag:ident $scanner:expr, [$($args:tt)*] ($($tuple:tt)*) $($t:tt)*)                 => { $crate::scan_value!(@$tag $scanner, [$($args)* [$crate::scan_value!(@tuple $scanner, [] $($tuple)*)]] $($t)*) };
157    (@$tag:ident $scanner:expr, [$($args:tt)*] [@$e:expr; const $len:expr] $($t:tt)*)     => { $crate::scan_value!(@$tag $scanner, [$($args)* [$crate::scan_value!(@array $scanner, [@$e] $len)]] $($t)*) };
158    (@$tag:ident $scanner:expr, [$($args:tt)*] [@$e:expr; $len:expr] $($t:tt)*)           => { $crate::scan_value!(@$tag $scanner, [$($args)* [$crate::scan_value!(@repeat $scanner, [@$e] $len)]] $($t)*) };
159    (@$tag:ident $scanner:expr, [$($args:tt)*] [[$($tt:tt)*]; const $len:expr] $($t:tt)*) => { $crate::scan_value!(@$tag $scanner, [$($args)* [$crate::scan_value!(@array $scanner, [[$($tt)*]] $len)]] $($t)*) };
160    (@$tag:ident $scanner:expr, [$($args:tt)*] [[$($tt:tt)*]; $len:expr] $($t:tt)*)       => { $crate::scan_value!(@$tag $scanner, [$($args)* [$crate::scan_value!(@repeat $scanner, [[$($tt)*]] $len)]] $($t)*) };
161    (@$tag:ident $scanner:expr, [$($args:tt)*] [($($tt:tt)*); const $len:expr] $($t:tt)*) => { $crate::scan_value!(@$tag $scanner, [$($args)* [$crate::scan_value!(@array $scanner, [($($tt)*)] $len)]] $($t)*) };
162    (@$tag:ident $scanner:expr, [$($args:tt)*] [($($tt:tt)*); $len:expr] $($t:tt)*)       => { $crate::scan_value!(@$tag $scanner, [$($args)* [$crate::scan_value!(@repeat $scanner, [($($tt)*)] $len)]] $($t)*) };
163    (@$tag:ident $scanner:expr, [$($args:tt)*] [$ty:ty; const $len:expr] $($t:tt)*)       => { $crate::scan_value!(@$tag $scanner, [$($args)* [$crate::scan_value!(@array $scanner, [$ty] $len)]] $($t)*) };
164    (@$tag:ident $scanner:expr, [$($args:tt)*] [$ty:ty; $len:expr] $($t:tt)*)             => { $crate::scan_value!(@$tag $scanner, [$($args)* [$crate::scan_value!(@repeat $scanner, [$ty] $len)]] $($t)*) };
165    (@$tag:ident $scanner:expr, [$($args:tt)*] [$($tt:tt)*] $($t:tt)*)                    => { $crate::scan_value!(@$tag $scanner, [$($args)* [$crate::scan_value!(@repeat $scanner, [$($tt)*])]] $($t)*) };
166    (@$tag:ident $scanner:expr, [$($args:tt)*] $ty:ty)                                    => { $crate::scan_value!(@$tag $scanner, [$($args)* [$scanner.scan::<$ty>()]]) };
167    (@$tag:ident $scanner:expr, [$($args:tt)*] $ty:ty, $($t:tt)*)                         => { $crate::scan_value!(@$tag $scanner, [$($args)* [$scanner.scan::<$ty>()]] $($t)*) };
168    (@$tag:ident $scanner:expr, [$($args:tt)*] , $($t:tt)*)                               => { $crate::scan_value!(@$tag $scanner, [$($args)*] $($t)*) };
169    (@$tag:ident $scanner:expr, [$($args:tt)*])                                           => { ::std::compile_error!(::std::stringify!($($args)*)) };
170    (src = $src:expr, $($t:tt)*)                                                          => { { let mut __scanner = Scanner::new($src); $crate::scan_value!(@inner __scanner, [] $($t)*) } };
171    ($scanner:expr, $($t:tt)*)                                                            => { $crate::scan_value!(@inner $scanner, [] $($t)*) }
172}
173
174/// scan and bind values with Scanner
175///
176/// - `scan!(scanner, $($pat $(: ELEMENT)?),*)`
177#[macro_export]
178macro_rules! scan {
179    (@assert $p:pat) => {};
180    (@assert $($p:tt)*) => { ::std::compile_error!(::std::concat!("expected pattern, found `", ::std::stringify!($($p)*), "`")); };
181    (@pat $scanner:expr, [] [])                                          => {};
182    (@pat $scanner:expr, [] [] , $($t:tt)*)                              => { $crate::scan!(@pat $scanner, [] [] $($t)*) };
183    (@pat $scanner:expr, [$($p:tt)*] [] $x:ident $($t:tt)*)              => { $crate::scan!(@pat $scanner, [$($p)* $x] [] $($t)*) };
184    (@pat $scanner:expr, [$($p:tt)*] [] :: $($t:tt)*)                    => { $crate::scan!(@pat $scanner, [$($p)* ::] [] $($t)*) };
185    (@pat $scanner:expr, [$($p:tt)*] [] & $($t:tt)*)                     => { $crate::scan!(@pat $scanner, [$($p)* &] [] $($t)*) };
186    (@pat $scanner:expr, [$($p:tt)*] [] ($($x:tt)*) $($t:tt)*)           => { $crate::scan!(@pat $scanner, [$($p)* ($($x)*)] [] $($t)*) };
187    (@pat $scanner:expr, [$($p:tt)*] [] [$($x:tt)*] $($t:tt)*)           => { $crate::scan!(@pat $scanner, [$($p)* [$($x)*]] [] $($t)*) };
188    (@pat $scanner:expr, [$($p:tt)*] [] {$($x:tt)*} $($t:tt)*)           => { $crate::scan!(@pat $scanner, [$($p)* {$($x)*}] [] $($t)*) };
189    (@pat $scanner:expr, [$($p:tt)*] [] : $($t:tt)*)                     => { $crate::scan!(@ty  $scanner, [$($p)*] [] $($t)*) };
190    (@pat $scanner:expr, [$($p:tt)*] [] $($t:tt)*)                       => { $crate::scan!(@let $scanner, [$($p)*] [usize] $($t)*) };
191    (@ty  $scanner:expr, [$($p:tt)*] [$($tt:tt)*] @$e:expr)              => { $crate::scan!(@let $scanner, [$($p)*] [$($tt)* @$e]) };
192    (@ty  $scanner:expr, [$($p:tt)*] [$($tt:tt)*] @$e:expr, $($t:tt)*)   => { $crate::scan!(@let $scanner, [$($p)*] [$($tt)* @$e], $($t)*) };
193    (@ty  $scanner:expr, [$($p:tt)*] [$($tt:tt)*] ($($x:tt)*) $($t:tt)*) => { $crate::scan!(@let $scanner, [$($p)*] [$($tt)* ($($x)*)] $($t)*) };
194    (@ty  $scanner:expr, [$($p:tt)*] [$($tt:tt)*] [$($x:tt)*] $($t:tt)*) => { $crate::scan!(@let $scanner, [$($p)*] [$($tt)* [$($x)*]] $($t)*) };
195    (@ty  $scanner:expr, [$($p:tt)*] [$($tt:tt)*] $ty:ty)                => { $crate::scan!(@let $scanner, [$($p)*] [$($tt)* $ty]) };
196    (@ty  $scanner:expr, [$($p:tt)*] [$($tt:tt)*] $ty:ty, $($t:tt)*)     => { $crate::scan!(@let $scanner, [$($p)*] [$($tt)* $ty], $($t)*) };
197    (@let $scanner:expr, [$($p:tt)*] [$($tt:tt)*] $($t:tt)*) => {
198        $crate::scan!{@assert $($p)*}
199        let $($p)* = $crate::scan_value!($scanner, $($tt)*);
200        $crate::scan!(@pat $scanner, [] [] $($t)*)
201    };
202    (src = $src:expr, $($t:tt)*) => { let mut __scanner = Scanner::new($src); $crate::scan!(@pat __scanner, [] [] $($t)*) };
203    ($scanner:expr, $($t:tt)*) => { $crate::scan!(@pat $scanner, [] [] $($t)*) }
204}
205
206#[derive(Debug, Copy, Clone)]
207pub enum Usize1 {}
208impl IterScan for Usize1 {
209    type Output = usize;
210    #[inline]
211    fn scan<'a, I: Iterator<Item = &'a str>>(iter: &mut I) -> Option<Self::Output> {
212        <usize as IterScan>::scan(iter)?.checked_sub(1)
213    }
214}
215#[derive(Debug, Copy, Clone)]
216pub struct CharWithBase(pub char);
217impl MarkedIterScan for CharWithBase {
218    type Output = usize;
219    #[inline]
220    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output> {
221        Some((<char as IterScan>::scan(iter)? as u8 - self.0 as u8) as usize)
222    }
223}
224#[derive(Debug, Copy, Clone)]
225pub enum Chars {}
226impl IterScan for Chars {
227    type Output = Vec<char>;
228    #[inline]
229    fn scan<'a, I: Iterator<Item = &'a str>>(iter: &mut I) -> Option<Self::Output> {
230        Some(iter.next()?.chars().collect())
231    }
232}
233#[derive(Debug, Copy, Clone)]
234pub struct CharsWithBase(pub char);
235impl MarkedIterScan for CharsWithBase {
236    type Output = Vec<usize>;
237    #[inline]
238    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output> {
239        Some(
240            iter.next()?
241                .chars()
242                .map(|c| (c as u8 - self.0 as u8) as usize)
243                .collect(),
244        )
245    }
246}
247#[derive(Debug, Copy, Clone)]
248pub enum Byte1 {}
249impl IterScan for Byte1 {
250    type Output = u8;
251    #[inline]
252    fn scan<'a, I: Iterator<Item = &'a str>>(iter: &mut I) -> Option<Self::Output> {
253        let bytes = iter.next()?.as_bytes();
254        assert_eq!(bytes.len(), 1);
255        Some(bytes[0])
256    }
257}
258#[derive(Debug, Copy, Clone)]
259pub struct ByteWithBase(pub u8);
260impl MarkedIterScan for ByteWithBase {
261    type Output = usize;
262    #[inline]
263    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output> {
264        Some((<char as IterScan>::scan(iter)? as u8 - self.0) as usize)
265    }
266}
267#[derive(Debug, Copy, Clone)]
268pub enum Bytes {}
269impl IterScan for Bytes {
270    type Output = Vec<u8>;
271    #[inline]
272    fn scan<'a, I: Iterator<Item = &'a str>>(iter: &mut I) -> Option<Self::Output> {
273        Some(iter.next()?.bytes().collect())
274    }
275}
276#[derive(Debug, Copy, Clone)]
277pub struct BytesWithBase(pub u8);
278impl MarkedIterScan for BytesWithBase {
279    type Output = Vec<usize>;
280    #[inline]
281    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output> {
282        Some(
283            iter.next()?
284                .bytes()
285                .map(|c| (c - self.0) as usize)
286                .collect(),
287        )
288    }
289}
290#[derive(Debug, Copy, Clone)]
291pub struct Collect<T, B = Vec<<T as IterScan>::Output>>
292where
293    T: IterScan,
294    B: FromIterator<<T as IterScan>::Output>,
295{
296    size: usize,
297    _marker: PhantomData<fn() -> (T, B)>,
298}
299impl<T, B> Collect<T, B>
300where
301    T: IterScan,
302    B: FromIterator<<T as IterScan>::Output>,
303{
304    pub fn new(size: usize) -> Self {
305        Self {
306            size,
307            _marker: PhantomData,
308        }
309    }
310}
311impl<T, B> MarkedIterScan for Collect<T, B>
312where
313    T: IterScan,
314    B: FromIterator<<T as IterScan>::Output>,
315{
316    type Output = B;
317    #[inline]
318    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output> {
319        repeat_with(|| <T as IterScan>::scan(iter))
320            .take(self.size)
321            .collect()
322    }
323}
324#[derive(Debug, Copy, Clone)]
325pub struct SizedCollect<T, B = Vec<<T as IterScan>::Output>>
326where
327    T: IterScan,
328    B: FromIterator<<T as IterScan>::Output>,
329{
330    _marker: PhantomData<fn() -> (T, B)>,
331}
332impl<T, B> IterScan for SizedCollect<T, B>
333where
334    T: IterScan,
335    B: FromIterator<<T as IterScan>::Output>,
336{
337    type Output = B;
338    #[inline]
339    fn scan<'a, I: Iterator<Item = &'a str>>(iter: &mut I) -> Option<Self::Output> {
340        let size = usize::scan(iter)?;
341        repeat_with(|| <T as IterScan>::scan(iter))
342            .take(size)
343            .collect()
344    }
345}
346#[derive(Debug, Copy, Clone)]
347pub struct Splitted<T, P>
348where
349    T: IterScan,
350{
351    pat: P,
352    _marker: PhantomData<fn() -> T>,
353}
354impl<T, P> Splitted<T, P>
355where
356    T: IterScan,
357{
358    pub fn new(pat: P) -> Self {
359        Self {
360            pat,
361            _marker: PhantomData,
362        }
363    }
364}
365impl<T> MarkedIterScan for Splitted<T, char>
366where
367    T: IterScan,
368{
369    type Output = Vec<<T as IterScan>::Output>;
370    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output> {
371        let mut iter = iter.next()?.split(self.pat);
372        Some(from_fn(|| <T as IterScan>::scan(&mut iter)).collect())
373    }
374}
375impl<T> MarkedIterScan for Splitted<T, &str>
376where
377    T: IterScan,
378{
379    type Output = Vec<<T as IterScan>::Output>;
380    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output> {
381        let mut iter = iter.next()?.split(self.pat);
382        Some(from_fn(|| <T as IterScan>::scan(&mut iter)).collect())
383    }
384}
385impl<T, F> MarkedIterScan for F
386where
387    F: Fn(&str) -> Option<T>,
388{
389    type Output = T;
390    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output> {
391        self(iter.next()?)
392    }
393}
394
395#[test]
396fn test_scan() {
397    use crate::scan;
398    let mut s = Scanner::new("1 2 3 a 1 2 1 1 1.1 2 3");
399    scan!(s, x, y: char, z: Usize1, a: @CharWithBase('a'), b: [usize; 2], c: (usize, @CharWithBase('0')), d: @Splitted::<usize, _>::new('.'), e: [usize; const 2]);
400    assert_eq!(x, 1);
401    assert_eq!(y, '2');
402    assert_eq!(z, 2);
403    assert_eq!(a, 0);
404    assert_eq!(b, vec![1, 2]);
405    assert_eq!(c, (1, 1));
406    assert_eq!(d, vec![1, 1]);
407    assert_eq!(e, [2, 3]);
408
409    scan!(src = "1", x);
410    assert_eq!(x, 1);
411    assert_eq!(scan_value!(src = "1", usize), 1);
412}
competitive/tools/scanner.rs

competitive/tools/
scanner.rs
