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 Scan: Sized {
34    type Output;
35    fn scan<S: ScanSource>(source: &mut S) -> Option<Self::Output>;
36}
37pub trait MarkedScan: Sized {
38    type Output;
39    fn mscan<S: ScanSource>(self, source: &mut S) -> Option<Self::Output>;
40}
41
42pub trait ScanSource: Sized {
43    /// Reads a token under the source's input requirements.
44    /// Checked sources return `None` at EOF; unchecked sources require available input.
45    fn next_token(&mut self) -> Option<&str>;
46    /// Skips separators without consuming the next token.
47    #[inline]
48    fn skip_whitespace(&mut self) {}
49    #[inline]
50    fn read_u8(&mut self) -> Option<u8> {
51        self.next_token()?.parse().ok()
52    }
53    #[inline]
54    fn read_u16(&mut self) -> Option<u16> {
55        self.next_token()?.parse().ok()
56    }
57    #[inline]
58    fn read_u32(&mut self) -> Option<u32> {
59        self.next_token()?.parse().ok()
60    }
61    #[inline]
62    fn read_u64(&mut self) -> Option<u64> {
63        self.next_token()?.parse().ok()
64    }
65    #[inline]
66    fn read_u128(&mut self) -> Option<u128> {
67        self.next_token()?.parse().ok()
68    }
69    #[inline]
70    fn read_usize(&mut self) -> Option<usize> {
71        self.next_token()?.parse().ok()
72    }
73    #[inline]
74    fn read_i8(&mut self) -> Option<i8> {
75        self.next_token()?.parse().ok()
76    }
77    #[inline]
78    fn read_i16(&mut self) -> Option<i16> {
79        self.next_token()?.parse().ok()
80    }
81    #[inline]
82    fn read_i32(&mut self) -> Option<i32> {
83        self.next_token()?.parse().ok()
84    }
85    #[inline]
86    fn read_i64(&mut self) -> Option<i64> {
87        self.next_token()?.parse().ok()
88    }
89    #[inline]
90    fn read_i128(&mut self) -> Option<i128> {
91        self.next_token()?.parse().ok()
92    }
93    #[inline]
94    fn read_isize(&mut self) -> Option<isize> {
95        self.next_token()?.parse().ok()
96    }
97    /// Panics if reading fails.
98    #[inline]
99    fn scan<T: Scan>(&mut self) -> T::Output {
100        T::scan(self).expect("scan error")
101    }
102    /// Panics if reading fails.
103    #[inline]
104    fn mscan<T: MarkedScan>(&mut self, marker: T) -> T::Output {
105        marker.mscan(self).expect("scan error")
106    }
107    fn scan_vec<T: Scan>(&mut self, size: usize) -> Vec<T::Output> {
108        if size == 0 {
109            self.skip_whitespace();
110        }
111        (0..size).map(|_| self.scan::<T>()).collect()
112    }
113    #[inline]
114    fn iter<T: Scan>(&mut self) -> ScannerIter<'_, Self, T> {
115        ScannerIter {
116            inner: self,
117            _marker: PhantomData,
118        }
119    }
120}
121
122#[derive(Clone, Debug)]
123pub struct Scanner<'a, I: Iterator<Item = &'a str> = std::str::SplitAsciiWhitespace<'a>> {
124    iter: I,
125}
126impl<'a> Scanner<'a> {
127    pub fn new(s: &'a str) -> Self {
128        Self {
129            iter: s.split_ascii_whitespace(),
130        }
131    }
132}
133impl<'a, I: Iterator<Item = &'a str>> Scanner<'a, I> {
134    pub fn new_from_iter(iter: I) -> Self {
135        Self { iter }
136    }
137}
138impl<'a, I: Iterator<Item = &'a str>> ScanSource for Scanner<'a, I> {
139    fn next_token(&mut self) -> Option<&str> {
140        self.iter.next()
141    }
142}
143
144macro_rules! impl_scan {
145    ($($t:ty)*) => {$(
146        impl Scan for $t {
147            type Output = Self;
148            fn scan<I: ScanSource>(iter: &mut I) -> Option<Self> {
149                iter.next_token()?.parse::<$t>().ok()
150            }
151        })*
152    };
153}
154impl_scan!(char f32 f64 String);
155
156macro_rules! impl_integer_scan {
157    ($($t:ty => $read:ident),* $(,)?) => {$(
158        impl Scan for $t {
159            type Output = Self;
160            #[inline]
161            fn scan<S: ScanSource>(source: &mut S) -> Option<Self> {
162                source.$read()
163            }
164        }
165    )*};
166}
167impl_integer_scan!(
168    u8 => read_u8, u16 => read_u16, u32 => read_u32, u64 => read_u64,
169    u128 => read_u128, usize => read_usize, i8 => read_i8, i16 => read_i16,
170    i32 => read_i32, i64 => read_i64, i128 => read_i128, isize => read_isize,
171);
172
173macro_rules! impl_scan_tuple {
174    (@impl $($T:ident)*) => {
175        impl<$($T: Scan),*> Scan for ($($T,)*) {
176            type Output = ($(<$T as Scan>::Output,)*);
177            fn scan<It: ScanSource>(_iter: &mut It) -> Option<Self::Output> {
178                Some(($(<$T as Scan>::scan(_iter)?,)*))
179            }
180        }
181    };
182    (@inner $($T:ident)*,) => {
183        impl_scan_tuple!(@impl $($T)*);
184    };
185    (@inner $($T:ident)*, $U:ident $($Rest:ident)*) => {
186        impl_scan_tuple!(@impl $($T)*);
187        impl_scan_tuple!(@inner $($T)* $U, $($Rest)*);
188    };
189    ($($T:ident)*) => {
190        impl_scan_tuple!(@inner , $($T)*);
191    };
192}
193impl_scan_tuple!(A B C D E F G H I J K);
194
195pub struct ScannerIter<'a, S, T> {
196    inner: &'a mut S,
197    _marker: PhantomData<fn() -> T>,
198}
199impl<S: ScanSource, T: Scan> Iterator for ScannerIter<'_, S, T> {
200    type Item = T::Output;
201    #[inline]
202    fn next(&mut self) -> Option<Self::Item> {
203        T::scan(self.inner)
204    }
205}
206
207/// scan a value with Scanner
208///
209/// - `scan_value!(scanner, ELEMENT)`
210///
211/// ELEMENT :=
212/// - `$ty`: Scan
213/// - `&str`: borrowed token; the source stays borrowed until its last use
214/// - `@$expr`: MarkedScan
215/// - `$ty = $expr`: MarkedScan
216/// - `[ELEMENT; $expr]`: vector
217/// - `[ELEMENT; const $expr]`: array
218/// - `[ELEMENT]`: iterator
219/// - `[ELEMENT; iter $expr]`: iterator of the specified length
220/// - `($(ELEMENT)*,)`: tuple
221#[macro_export]
222macro_rules! scan_value {
223    (@repeat $scanner:expr, [$($t:tt)*] $len:expr)                             => { { $crate::scan_value!(@iter $scanner, [$($t)*] $len).collect::<Vec<_>>() } };
224    (@repeat $scanner:expr, [$($t:tt)*])                                       => { { ::std::iter::repeat_with(|| $crate::scan_value!(@inner $scanner, [] $($t)*)) } };
225    (@iter $scanner:expr, [$($t:tt)*] $len:expr)                               => {{ let size = $len; if size == 0 { $scanner.skip_whitespace(); } $crate::scan_value!(@repeat $scanner, [$($t)*]).take(size) }};
226    (@array $scanner:expr, [$($t:tt)*] $len:expr)                              => { { if $len == 0 { $scanner.skip_whitespace(); } $crate::array![|| $crate::scan_value!(@inner $scanner, [] $($t)*); $len] } };
227    (@tuple $scanner:expr, [$([$($args:tt)*])*])                               => { ($($($args)*,)*) };
228    (@sparen $scanner:expr, [] @$e:expr; $($t:tt)*)                            => { $crate::scan_value!(@sparen $scanner, [@$e] $($t)*) };
229    (@sparen $scanner:expr, [] ($($tt:tt)*); $($t:tt)*)                        => { $crate::scan_value!(@sparen $scanner, [($($tt)*)] $($t)*) };
230    (@sparen $scanner:expr, [] [$($tt:tt)*]; $($t:tt)*)                        => { $crate::scan_value!(@sparen $scanner, [[$($tt)*]] $($t)*) };
231    (@sparen $scanner:expr, [] $ty:ty = $e:expr; $($t:tt)*)                    => { $crate::scan_value!(@sparen $scanner, [$ty = $e] $($t)*) };
232    (@sparen $scanner:expr, [] $ty:ty; $($t:tt)*)                              => { $crate::scan_value!(@sparen $scanner, [$ty] $($t)*) };
233    (@sparen $scanner:expr, [] $($args:tt)*)                                   => { $crate::scan_value!(@repeat $scanner, [$($args)*]) };
234    (@sparen $scanner:expr, [$($args:tt)+] const $len:expr)                    => { $crate::scan_value!(@array $scanner, [$($args)+] $len) };
235    (@sparen $scanner:expr, [$($args:tt)+] iter $len:expr)                     => { $crate::scan_value!(@iter $scanner, [$($args)+] $len) };
236    (@sparen $scanner:expr, [$($args:tt)+] $len:expr)                          => { $crate::scan_value!(@repeat $scanner, [$($args)+] $len) };
237    (@$tag:ident $scanner:expr, [[$($args:tt)*]])                              => { $($args)* };
238    (@$tag:ident $scanner:expr, [$($args:tt)*] @$e:expr $(, $($t:tt)*)?)       => { $crate::scan_value!(@$tag $scanner, [$($args)* [$scanner.mscan($e)]] $(, $($t)*)?) };
239    (@$tag:ident $scanner:expr, [$($args:tt)*] ($($tuple:tt)*) $($t:tt)*)      => { $crate::scan_value!(@$tag $scanner, [$($args)* [$crate::scan_value!(@tuple $scanner, [] $($tuple)*)]] $($t)*) };
240    (@$tag:ident $scanner:expr, [$($args:tt)*] [$($tt:tt)*] $($t:tt)*)         => { $crate::scan_value!(@$tag $scanner, [$($args)* [$crate::scan_value!(@sparen $scanner, [] $($tt)*)]] $($t)*) };
241    (@$tag:ident $scanner:expr, [$($args:tt)*] &str $(, $($t:tt)*)?)           => { $crate::scan_value!(@$tag $scanner, [$($args)* [$scanner.next_token().expect("scan error")]] $(, $($t)*)?) };
242    (@$tag:ident $scanner:expr, [$($args:tt)*] $ty:ty = $e:expr $(, $($t:tt)*)?) => { $crate::scan_value!(@$tag $scanner, [$($args)* [{ let _tmp: $ty = $scanner.mscan($e); _tmp }]] $(, $($t)*)?) };
243    (@$tag:ident $scanner:expr, [$($args:tt)*] $ty:ty $(, $($t:tt)*)?)         => { $crate::scan_value!(@$tag $scanner, [$($args)* [$scanner.scan::<$ty>()]] $(, $($t)*)?) };
244    (@$tag:ident $scanner:expr, [$($args:tt)*] , $($t:tt)*)                    => { $crate::scan_value!(@$tag $scanner, [$($args)*] $($t)*) };
245    (@$tag:ident $scanner:expr, [$($args:tt)*])                                => { ::std::compile_error!(::std::stringify!($($args)*)) };
246    (src = $src:expr, $($t:tt)*)                                               => { { let mut __scanner = Scanner::new($src); $crate::scan_value!(@inner __scanner, [] $($t)*) } };
247    (iter = $iter:expr, $($t:tt)*)                                             => { { let mut __scanner = Scanner::new_from_iter($iter); $crate::scan_value!(@inner __scanner, [] $($t)*) } };
248    ($scanner:expr, $($t:tt)*)                                                 => { $crate::scan_value!(@inner $scanner, [] $($t)*) }
249}
250
251/// scan and bind values with Scanner
252///
253/// - `scan!(scanner, $($pat $(: ELEMENT)?),*)`
254#[macro_export]
255macro_rules! scan {
256    (@assert $p:pat) => {};
257    (@assert $($p:tt)*) => { ::std::compile_error!(::std::concat!("expected pattern, found `", ::std::stringify!($($p)*), "`")); };
258    (@pat $scanner:expr, [] [])                                                     => {};
259    (@pat $scanner:expr, [] [] , $($t:tt)*)                                         => { $crate::scan!(@pat $scanner, [] [] $($t)*) };
260    (@pat $scanner:expr, [$($p:tt)*] [] $x:ident $($t:tt)*)                         => { $crate::scan!(@pat $scanner, [$($p)* $x] [] $($t)*) };
261    (@pat $scanner:expr, [$($p:tt)*] [] :: $($t:tt)*)                               => { $crate::scan!(@pat $scanner, [$($p)* ::] [] $($t)*) };
262    (@pat $scanner:expr, [$($p:tt)*] [] & $($t:tt)*)                                => { $crate::scan!(@pat $scanner, [$($p)* &] [] $($t)*) };
263    (@pat $scanner:expr, [$($p:tt)*] [] ($($x:tt)*) $($t:tt)*)                      => { $crate::scan!(@pat $scanner, [$($p)* ($($x)*)] [] $($t)*) };
264    (@pat $scanner:expr, [$($p:tt)*] [] [$($x:tt)*] $($t:tt)*)                      => { $crate::scan!(@pat $scanner, [$($p)* [$($x)*]] [] $($t)*) };
265    (@pat $scanner:expr, [$($p:tt)*] [] {$($x:tt)*} $($t:tt)*)                      => { $crate::scan!(@pat $scanner, [$($p)* {$($x)*}] [] $($t)*) };
266    (@pat $scanner:expr, [$($p:tt)*] [] : $($t:tt)*)                                => { $crate::scan!(@ty  $scanner, [$($p)*] [] $($t)*) };
267    (@pat $scanner:expr, [$($p:tt)*] [] $($t:tt)*)                                  => { $crate::scan!(@let $scanner, [$($p)*] [usize] $($t)*) };
268    (@ty  $scanner:expr, [$($p:tt)*] [$($tt:tt)*] @$e:expr $(, $($t:tt)*)?)         => { $crate::scan!(@let $scanner, [$($p)*] [$($tt)* @$e] $(, $($t)*)?) };
269    (@ty  $scanner:expr, [$($p:tt)*] [$($tt:tt)*] ($($x:tt)*) $($t:tt)*)            => { $crate::scan!(@let $scanner, [$($p)*] [$($tt)* ($($x)*)] $($t)*) };
270    (@ty  $scanner:expr, [$($p:tt)*] [$($tt:tt)*] [$($x:tt)*] $($t:tt)*)            => { $crate::scan!(@let $scanner, [$($p)*] [$($tt)* [$($x)*]] $($t)*) };
271    (@ty  $scanner:expr, [$($p:tt)*] [$($tt:tt)*] &str $(, $($t:tt)*)?)             => { $crate::scan!(@let $scanner, [$($p)*] [$($tt)* &str] $(, $($t)*)?) };
272    (@ty  $scanner:expr, [$($p:tt)*] [$($tt:tt)*] $ty:ty = $e:expr $(, $($t:tt)*)?) => { $crate::scan!(@let $scanner, [$($p)*] [$($tt)* $ty = $e] $(, $($t)*)?) };
273    (@ty  $scanner:expr, [$($p:tt)*] [$($tt:tt)*] $ty:ty $(, $($t:tt)*)?)           => { $crate::scan!(@let $scanner, [$($p)*] [$($tt)* $ty] $(, $($t)*)?) };
274    (@let $scanner:expr, [$($p:tt)*] [$($tt:tt)*] $($t:tt)*) => {
275        $crate::scan!{@assert $($p)*}
276        let $($p)* = $crate::scan_value!($scanner, $($tt)*);
277        $crate::scan!(@pat $scanner, [] [] $($t)*)
278    };
279    (src = $src:expr, $($t:tt)*)   => { let mut __scanner = Scanner::new($src); $crate::scan!(@pat __scanner, [] [] $($t)*) };
280    (iter = $iter:expr, $($t:tt)*) => { let mut __scanner = Scanner::new_from_iter($iter); $crate::scan!(@pat __scanner, [] [] $($t)*) };
281    ($scanner:expr, $($t:tt)*) => { $crate::scan!(@pat $scanner, [] [] $($t)*) }
282}
283
284/// define enum scan rules
285///
286/// # Example
287/// ```rust
288/// # use competitive::{define_enum_scan, tools::{CharsWithBase, Scan, ScanSource, Scanner, Usize1}};
289/// define_enum_scan! {
290///   enum Query: u8 {
291///     0 => Noop,
292///     1 => Args { i: Usize1, s: char },
293///     9 => Complex { n: usize, c: [(usize, Vec<usize> = CharsWithBase('a')); n] },
294///   }
295/// }
296/// ```
297#[macro_export]
298macro_rules! define_enum_scan {
299    (@field_ty @repeat [$($t:tt)*] $($len:expr)?)                           => { Vec<$crate::define_enum_scan!(@field_ty $($t)*)> };
300    (@field_ty @array [$($t:tt)*] $len:expr)                                => { [$crate::define_enum_scan!(@field_ty $($t)*); $len] };
301    (@field_ty @tuple [$([$($args:tt)*])*])                                 => { ($( $($args)* ,)*) };
302    (@field_ty @sparen [] ($($tt:tt)*); $($t:tt)*)                          => { $crate::define_enum_scan!(@field_ty @sparen [($($tt)*)] $($t)*) };
303    (@field_ty @sparen [] [$($tt:tt)*]; $($t:tt)*)                          => { $crate::define_enum_scan!(@field_ty @sparen [[$($tt)*]] $($t)*) };
304    (@field_ty @sparen [] $ty:ty = $e:expr; $($t:tt)*)                      => { $crate::define_enum_scan!(@field_ty @sparen [$ty = $e] $($t)*) };
305    (@field_ty @sparen [] $ty:ty; $($t:tt)*)                                => { $crate::define_enum_scan!(@field_ty @sparen [$ty] $($t)*) };
306    (@field_ty @sparen [] $($args:tt)*)                                     => { $crate::define_enum_scan!(@field_ty @repeat [$($args)*]) };
307    (@field_ty @sparen [$($args:tt)+] const $len:expr)                      => { $crate::define_enum_scan!(@field_ty @array [$($args)+] $len) };
308    (@field_ty @sparen [$($args:tt)+] $len:expr)                            => { $crate::define_enum_scan!(@field_ty @repeat [$($args)+] $len) };
309    (@field_ty @$tag:ident [$($args:tt)*] ($($tuple:tt)*) $($t:tt)*)        => { $crate::define_enum_scan!(@field_ty @$tag [$($args)* [$crate::define_enum_scan!(@field_ty @tuple [] $($tuple)*)]] $($t)*) };
310    (@field_ty @$tag:ident [$($args:tt)*] [$($tt:tt)*] $($t:tt)*)           => { $crate::define_enum_scan!(@field_ty @$tag [$($args)* [$crate::define_enum_scan!(@field_ty @sparen [] $($tt)*)]] $($t)*) };
311    (@field_ty @$tag:ident [$($args:tt)*] $ty:ty = $e:expr $(, $($t:tt)*)?) => { $crate::define_enum_scan!(@field_ty @$tag [$($args)* [$ty]] $(, $($t)*)?) };
312    (@field_ty @$tag:ident [$($args:tt)*] $ty:ty $(, $($t:tt)*)?)           => { $crate::define_enum_scan!(@field_ty @$tag [$($args)* [<$ty as Scan>::Output]] $(, $($t)*)?) };
313    (@field_ty @$tag:ident [$($args:tt)*] , $($t:tt)*)                      => { $crate::define_enum_scan!(@field_ty @$tag [$($args)*] $($t)*) };
314    (@field_ty @$tag:ident [[$($args:tt)*]])                                => { $($args)* };
315    (@field_ty @$tag:ident [$($args:tt)*])                                  => { ::std::compile_error!(::std::stringify!($($args)*)) };
316    (@field_ty $($t:tt)*) => { $crate::define_enum_scan!(@field_ty @inner [] $($t)*) };
317
318    (@tag_expr raw, $iter:ident) => { ScanSource::next_token($iter)? };
319    (@tag_expr $d:ty, $iter:ident) => { <$d as Scan>::scan($iter)? };
320    (@variant ([$($attr:tt)*] $vis:vis $T:ident $d:tt) [$($vars:tt)*]) => { $crate::define_enum_scan! { @def $($attr)* $vis enum $T : $d { $($vars)* } } };
321    (@variant $ctx:tt [$($vars:tt)*] $p:pat => $v:ident { $($fs:tt)* } $($rest:tt)*) => { $crate::define_enum_scan! { @field   $ctx [$($vars)*] $p => $v [] $($fs)* ; $($rest)* } };
322    (@variant $ctx:tt [$($vars:tt)*] $p:pat => $v:ident $($rest:tt)*)                    => { $crate::define_enum_scan! { @variant $ctx [$($vars)* $p => $v ,] $($rest)* } };
323    (@variant $ctx:tt [$($vars:tt)*] , $($rest:tt)*)                                     => { $crate::define_enum_scan! { @variant $ctx [$($vars)*] $($rest)* } };
324    (@endfield $ctx:tt [$($vars:tt)*] $p:pat => $v:ident [$($fs:tt)*] [$f:ident : $($spec:tt)*] , $($rest:tt)*) => { $crate::define_enum_scan! { @field $ctx [$($vars)*] $p => $v [$($fs)* [$f : $($spec)*]] $($rest)* } };
325    (@endfield $ctx:tt [$($vars:tt)*] $p:pat => $v:ident [$($fs:tt)*] [$f:ident : $($spec:tt)*] ; $($rest:tt)*) => { $crate::define_enum_scan! { @variant $ctx [$($vars)* $p => $v { $($fs)* [$f : $($spec)*] } ,] $($rest)* } };
326    (@field $ctx:tt [$($vars:tt)*] $p:pat => $v:ident [$($fs:tt)*] ; $($rest:tt)*)                                  => { $crate::define_enum_scan! { @variant $ctx [$($vars)* $p => $v { $($fs)* } ,] $($rest)* } };
327    (@field $ctx:tt [$($vars:tt)*] $p:pat => $v:ident [$($fs:tt)*] $f:ident : ($($tuple:tt)*) $sep:tt $($rest:tt)*) => { $crate::define_enum_scan! { @endfield $ctx [$($vars)*] $p => $v [$($fs)*] [$f : ($($tuple)*)] $sep $($rest)* } };
328    (@field $ctx:tt [$($vars:tt)*] $p:pat => $v:ident [$($fs:tt)*] $f:ident : [$($x:tt)*] $sep:tt $($rest:tt)*)     => { $crate::define_enum_scan! { @endfield $ctx [$($vars)*] $p => $v [$($fs)*] [$f : [$($x)*]] $sep $($rest)* } };
329    (@field $ctx:tt [$($vars:tt)*] $p:pat => $v:ident [$($fs:tt)*] $f:ident : $ty:ty = $e:expr , $($rest:tt)*)      => { $crate::define_enum_scan! { @endfield $ctx [$($vars)*] $p => $v [$($fs)*] [$f : $ty = $e] , $($rest)* } };
330    (@field $ctx:tt [$($vars:tt)*] $p:pat => $v:ident [$($fs:tt)*] $f:ident : $ty:ty ; $($rest:tt)*)                => { $crate::define_enum_scan! { @endfield $ctx [$($vars)*] $p => $v [$($fs)*] [$f : $ty] ; $($rest)* } };
331    (@field $ctx:tt [$($vars:tt)*] $p:pat => $v:ident [$($fs:tt)*] $f:ident : $ty:ty = $e:expr ; $($rest:tt)*)      => { $crate::define_enum_scan! { @endfield $ctx [$($vars)*] $p => $v [$($fs)*] [$f : $ty = $e] ; $($rest)* } };
332    (@field $ctx:tt [$($vars:tt)*] $p:pat => $v:ident [$($fs:tt)*] $f:ident : $ty:ty , $($rest:tt)*)                => { $crate::define_enum_scan! { @endfield $ctx [$($vars)*] $p => $v [$($fs)*] [$f : $ty] , $($rest)* } };
333    (
334        @def
335        $(#[$attr:meta])*
336        $vis:vis enum $T:ident : $d:tt {
337            $( $p:pat => $v:ident $( { $( [$f:ident : $($spec:tt)*] )* } )?, )*
338        }
339    ) => {
340        $(#[$attr])*
341        $vis enum $T {
342            $( $v $( { $( $f : $crate::define_enum_scan!(@field_ty $($spec)*) ),* } )? ),*
343        }
344        impl Scan for $T {
345            type Output = Self;
346            fn scan<I: ScanSource>(iter: &mut I) -> Option<Self> {
347                let tag = $crate::define_enum_scan!(@tag_expr $d, iter);
348                match tag {
349                    $(
350                        $p => {
351                            $($(
352                                let $f = $crate::scan_value!((*iter), $($spec)* );
353                            )*)?
354                            Some($T::$v $( { $( $f ),* } )?)
355                        }
356                    ),*
357                    _ => None,
358                }
359            }
360        }
361    };
362    (
363        $(#[$attr:meta])*
364        $vis:vis enum $T:ident : raw {
365            $($body:tt)*
366        }
367    ) => {
368        $crate::define_enum_scan! { @variant ([$(#[$attr])*] $vis $T raw) [] $($body)* }
369    };
370    (
371        $(#[$attr:meta])*
372        $vis:vis enum $T:ident : $d:ty {
373            $($body:tt)*
374        }
375    ) => {
376        $crate::define_enum_scan! { @variant ([$(#[$attr])*] $vis $T $d) [] $($body)* }
377    };
378}
379
380#[derive(Debug, Copy, Clone)]
381pub enum Usize1 {}
382impl Scan for Usize1 {
383    type Output = usize;
384    fn scan<I: ScanSource>(iter: &mut I) -> Option<Self::Output> {
385        <usize as Scan>::scan(iter)?.checked_sub(1)
386    }
387}
388#[derive(Debug, Copy, Clone)]
389pub struct CharWithBase(pub char);
390impl MarkedScan for CharWithBase {
391    type Output = usize;
392    fn mscan<I: ScanSource>(self, iter: &mut I) -> Option<Self::Output> {
393        Some((<char as Scan>::scan(iter)? as u8 - self.0 as u8) as usize)
394    }
395}
396#[derive(Debug, Copy, Clone)]
397pub enum Chars {}
398impl Scan for Chars {
399    type Output = Vec<char>;
400    fn scan<I: ScanSource>(iter: &mut I) -> Option<Self::Output> {
401        Some(iter.next_token()?.chars().collect())
402    }
403}
404#[derive(Debug, Copy, Clone)]
405pub struct CharsWithBase(pub char);
406impl MarkedScan for CharsWithBase {
407    type Output = Vec<usize>;
408    fn mscan<I: ScanSource>(self, iter: &mut I) -> Option<Self::Output> {
409        Some(
410            iter.next_token()?
411                .chars()
412                .map(|c| (c as u8 - self.0 as u8) as usize)
413                .collect(),
414        )
415    }
416}
417#[derive(Debug, Copy, Clone)]
418pub enum Byte1 {}
419impl Scan for Byte1 {
420    type Output = u8;
421    fn scan<I: ScanSource>(iter: &mut I) -> Option<Self::Output> {
422        let bytes = iter.next_token()?.as_bytes();
423        assert_eq!(bytes.len(), 1);
424        Some(bytes[0])
425    }
426}
427#[derive(Debug, Copy, Clone)]
428pub struct ByteWithBase(pub u8);
429impl MarkedScan for ByteWithBase {
430    type Output = usize;
431    fn mscan<I: ScanSource>(self, iter: &mut I) -> Option<Self::Output> {
432        Some((<char as Scan>::scan(iter)? as u8 - self.0) as usize)
433    }
434}
435#[derive(Debug, Copy, Clone)]
436pub enum Bytes {}
437impl Scan for Bytes {
438    type Output = Vec<u8>;
439    fn scan<I: ScanSource>(iter: &mut I) -> Option<Self::Output> {
440        Some(iter.next_token()?.bytes().collect())
441    }
442}
443#[derive(Debug, Copy, Clone)]
444pub struct BytesWithBase(pub u8);
445impl MarkedScan for BytesWithBase {
446    type Output = Vec<usize>;
447    fn mscan<I: ScanSource>(self, iter: &mut I) -> Option<Self::Output> {
448        Some(
449            iter.next_token()?
450                .bytes()
451                .map(|c| (c - self.0) as usize)
452                .collect(),
453        )
454    }
455}
456#[derive(Debug, Copy, Clone)]
457pub struct Collect<T, B = Vec<<T as Scan>::Output>>
458where
459    T: Scan,
460    B: FromIterator<<T as Scan>::Output>,
461{
462    size: usize,
463    _marker: PhantomData<fn() -> (T, B)>,
464}
465impl<T, B> Collect<T, B>
466where
467    T: Scan,
468    B: FromIterator<<T as Scan>::Output>,
469{
470    pub fn new(size: usize) -> Self {
471        Self {
472            size,
473            _marker: PhantomData,
474        }
475    }
476}
477impl<T, B> MarkedScan for Collect<T, B>
478where
479    T: Scan,
480    B: FromIterator<<T as Scan>::Output>,
481{
482    type Output = B;
483    fn mscan<I: ScanSource>(self, iter: &mut I) -> Option<Self::Output> {
484        if self.size == 0 {
485            iter.skip_whitespace();
486        }
487        repeat_with(|| <T as Scan>::scan(iter))
488            .take(self.size)
489            .collect()
490    }
491}
492#[derive(Debug, Copy, Clone)]
493pub struct SizedCollect<T, B = Vec<<T as Scan>::Output>>
494where
495    T: Scan,
496    B: FromIterator<<T as Scan>::Output>,
497{
498    _marker: PhantomData<fn() -> (T, B)>,
499}
500impl<T, B> Scan for SizedCollect<T, B>
501where
502    T: Scan,
503    B: FromIterator<<T as Scan>::Output>,
504{
505    type Output = B;
506    fn scan<I: ScanSource>(iter: &mut I) -> Option<Self::Output> {
507        let size = usize::scan(iter)?;
508        if size == 0 {
509            iter.skip_whitespace();
510        }
511        repeat_with(|| <T as Scan>::scan(iter)).take(size).collect()
512    }
513}
514#[derive(Debug, Copy, Clone)]
515pub struct Splitted<T, P>
516where
517    T: Scan,
518{
519    pat: P,
520    _marker: PhantomData<fn() -> T>,
521}
522impl<T, P> Splitted<T, P>
523where
524    T: Scan,
525{
526    pub fn new(pat: P) -> Self {
527        Self {
528            pat,
529            _marker: PhantomData,
530        }
531    }
532}
533impl<T> MarkedScan for Splitted<T, char>
534where
535    T: Scan,
536{
537    type Output = Vec<<T as Scan>::Output>;
538    fn mscan<I: ScanSource>(self, iter: &mut I) -> Option<Self::Output> {
539        let mut iter = Scanner::new_from_iter(iter.next_token()?.split(self.pat));
540        Some(from_fn(|| <T as Scan>::scan(&mut iter)).collect())
541    }
542}
543impl<T> MarkedScan for Splitted<T, &str>
544where
545    T: Scan,
546{
547    type Output = Vec<<T as Scan>::Output>;
548    fn mscan<I: ScanSource>(self, iter: &mut I) -> Option<Self::Output> {
549        let mut iter = Scanner::new_from_iter(iter.next_token()?.split(self.pat));
550        Some(from_fn(|| <T as Scan>::scan(&mut iter)).collect())
551    }
552}
553impl<T, F> MarkedScan for F
554where
555    F: Fn(&str) -> Option<T>,
556{
557    type Output = T;
558    fn mscan<I: ScanSource>(self, iter: &mut I) -> Option<Self::Output> {
559        self(iter.next_token()?)
560    }
561}
562
563#[cfg(test)]
564mod tests {
565    use super::*;
566    use crate::tools::Xorshift;
567    use std::array;
568
569    #[test]
570    fn test_scan() {
571        let mut rng = Xorshift::default();
572        for _ in 0..1000 {
573            let values: [usize; 10] = array::from_fn(|_| rng.random(1..=9));
574            let text = format!(
575                "{} {} {} {} {} {} {} {} {}.{} {} {}",
576                values[0],
577                values[1],
578                values[2],
579                char::from(b'a' + values[3] as u8),
580                values[4],
581                values[5],
582                values[6],
583                values[7],
584                values[8],
585                values[9],
586                values[0],
587                values[1]
588            );
589            let mut s = Scanner::new(&text);
590            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]);
591            assert_eq!(x, values[0]);
592            assert_eq!(y, char::from(b'0' + values[1] as u8));
593            assert_eq!(z, values[2] - 1);
594            assert_eq!(a, values[3]);
595            assert_eq!(b, values[4..6]);
596            assert_eq!(c, (values[6], values[7]));
597            assert_eq!(d, values[8..10]);
598            assert_eq!(e, [values[0], values[1]]);
599            let text = format!("{}{} {}{}", values[0], values[1], values[2], values[3]);
600            scan!(src = &text, c: Vec<usize> = CharsWithBase('0'), d: [Vec<usize> = CharsWithBase('0'); 1]);
601            assert_eq!(c, values[..2]);
602            assert_eq!(d, vec![values[2..4].to_vec()]);
603            let text = rng.random(0..=usize::MAX).to_string();
604            let expected: usize = text.parse().unwrap();
605            scan!(src = &text, x);
606            assert_eq!(x, expected);
607            assert_eq!(scan_value!(src = &text, usize), expected);
608            scan!(iter = text.split_ascii_whitespace(), x);
609            assert_eq!(x, expected);
610            assert_eq!(
611                scan_value!(iter = text.split_ascii_whitespace(), usize),
612                expected
613            );
614        }
615    }
616
617    #[test]
618    fn test_define_enum_scan() {
619        define_enum_scan! {
620            enum Query: u8 {
621                0 => Noop,
622                1 => Args { i: Usize1, s: char },
623                9 => Complex { n: usize, c: [(usize, Vec<usize> = CharsWithBase('a')); n] },
624            }
625        }
626        let mut rng = Xorshift::default();
627        for _ in 0..1000 {
628            let tag = rng.random(0..3);
629            let n = rng.random(0..=16);
630            let i = rng.random(1..=1000usize);
631            let ch = char::from(rng.random(b'a'..=b'z'));
632            let rows: Vec<_> = (0..n)
633                .map(|_| {
634                    let len = rng.random(1..=10);
635                    (
636                        rng.random(0..=100usize),
637                        rng.random_iter(0..26usize).take(len).collect::<Vec<_>>(),
638                    )
639                })
640                .collect();
641            let text = match tag {
642                0 => "0".to_owned(),
643                1 => format!("1 {i} {ch}"),
644                _ => format!(
645                    "9 {n} {}",
646                    rows.iter()
647                        .map(|(a, b)| format!(
648                            "{} {}",
649                            a,
650                            b.iter()
651                                .map(|&c| char::from(b'a' + c as u8))
652                                .collect::<String>()
653                        ))
654                        .collect::<Vec<_>>()
655                        .join(" ")
656                ),
657            };
658            let mut s = Scanner::new(&text);
659            match (tag, s.scan::<Query>()) {
660                (0, Query::Noop) => {}
661                (1, Query::Args { i: actual, s }) => {
662                    assert_eq!(actual, i - 1);
663                    assert_eq!(s, ch);
664                }
665                (2, Query::Complex { n: actual, c }) => {
666                    assert_eq!(actual, n);
667                    assert_eq!(c, rows);
668                }
669                _ => panic!("unexpected query: {text}"),
670            }
671        }
672    }
673}
competitive/tools/scanner.rs

competitive/tools/
scanner.rs
