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, I: Iterator<Item = &'a str> = std::str::SplitAsciiWhitespace<'a>> {
43    iter: I,
44}
45impl<'a> Scanner<'a> {
46    pub fn new(s: &'a str) -> Self {
47        let iter = s.split_ascii_whitespace();
48        Self { iter }
49    }
50}
51impl<'a, I: Iterator<Item = &'a str>> Scanner<'a, I> {
52    pub fn new_from_iter(iter: I) -> Self {
53        Self { iter }
54    }
55    pub fn scan<T>(&mut self) -> <T as IterScan>::Output
56    where
57        T: IterScan,
58    {
59        <T as IterScan>::scan(&mut self.iter).expect("scan error")
60    }
61    pub fn mscan<T>(&mut self, marker: T) -> <T as MarkedIterScan>::Output
62    where
63        T: MarkedIterScan,
64    {
65        marker.mscan(&mut self.iter).expect("scan error")
66    }
67    pub fn scan_vec<T>(&mut self, size: usize) -> Vec<<T as IterScan>::Output>
68    where
69        T: IterScan,
70    {
71        (0..size)
72            .map(|_| <T as IterScan>::scan(&mut self.iter).expect("scan error"))
73            .collect()
74    }
75    #[inline]
76    pub fn iter<'b, T>(&'b mut self) -> ScannerIter<'a, 'b, I, T>
77    where
78        T: IterScan,
79    {
80        ScannerIter {
81            inner: self,
82            _marker: std::marker::PhantomData,
83        }
84    }
85}
86
87macro_rules! impl_iter_scan {
88    ($($t:ty)*) => {$(
89        impl IterScan for $t {
90            type Output = Self;
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            fn scan<'a, It: Iterator<Item = &'a str>>(_iter: &mut It) -> Option<Self::Output> {
104                Some(($(<$T as IterScan>::scan(_iter)?,)*))
105            }
106        }
107    };
108    (@inner $($T:ident)*,) => {
109        impl_iter_scan_tuple!(@impl $($T)*);
110    };
111    (@inner $($T:ident)*, $U:ident $($Rest:ident)*) => {
112        impl_iter_scan_tuple!(@impl $($T)*);
113        impl_iter_scan_tuple!(@inner $($T)* $U, $($Rest)*);
114    };
115    ($($T:ident)*) => {
116        impl_iter_scan_tuple!(@inner , $($T)*);
117    };
118}
119impl_iter_scan_tuple!(A B C D E F G H I J K);
120
121pub struct ScannerIter<'a, 'b, I: Iterator<Item = &'a str>, T> {
122    inner: &'b mut Scanner<'a, I>,
123    _marker: std::marker::PhantomData<fn() -> T>,
124}
125impl<'a, I, T> Iterator for ScannerIter<'a, '_, I, T>
126where
127    I: Iterator<Item = &'a str>,
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/// - `$ty = $expr`: MarkedIterScan
145/// - `[ELEMENT; $expr]`: vector
146/// - `[ELEMENT; const $expr]`: array
147/// - `[ELEMENT]`: iterator
148/// - `($(ELEMENT)*,)`: tuple
149#[macro_export]
150macro_rules! scan_value {
151    (@repeat $scanner:expr, [$($t:tt)*] $($len:expr)?)                           => { ::std::iter::repeat_with(|| $crate::scan_value!(@inner $scanner, [] $($t)*)) $(.take($len).collect::<Vec<_>>())? };
152    (@array $scanner:expr, [$($t:tt)*] $len:expr)                                => { $crate::array![|| $crate::scan_value!(@inner $scanner, [] $($t)*); $len] };
153    (@tuple $scanner:expr, [$([$($args:tt)*])*])                                 => { ($($($args)*,)*) };
154    (@sparen $scanner:expr, [] @$e:expr; $($t:tt)*)                              => { $crate::scan_value!(@sparen $scanner, [@$e] $($t)*) };
155    (@sparen $scanner:expr, [] ($($tt:tt)*); $($t:tt)*)                          => { $crate::scan_value!(@sparen $scanner, [($($tt)*)] $($t)*) };
156    (@sparen $scanner:expr, [] [$($tt:tt)*]; $($t:tt)*)                          => { $crate::scan_value!(@sparen $scanner, [[$($tt)*]] $($t)*) };
157    (@sparen $scanner:expr, [] $ty:ty = $e:expr; $($t:tt)*)                      => { $crate::scan_value!(@sparen $scanner, [$ty = $e] $($t)*) };
158    (@sparen $scanner:expr, [] $ty:ty; $($t:tt)*)                                => { $crate::scan_value!(@sparen $scanner, [$ty] $($t)*) };
159    (@sparen $scanner:expr, [] $($args:tt)*)                                     => { $crate::scan_value!(@repeat $scanner, [$($args)*]) };
160    (@sparen $scanner:expr, [$($args:tt)+] const $len:expr)                      => { $crate::scan_value!(@array $scanner, [$($args)+] $len) };
161    (@sparen $scanner:expr, [$($args:tt)+] $len:expr)                            => { $crate::scan_value!(@repeat $scanner, [$($args)+] $len) };
162    (@$tag:ident $scanner:expr, [[$($args:tt)*]])                                => { $($args)* };
163    (@$tag:ident $scanner:expr, [$($args:tt)*] @$e:expr $(, $($t:tt)*)?)         => { $crate::scan_value!(@$tag $scanner, [$($args)* [$scanner.mscan($e)]] $(, $($t)*)?) };
164    (@$tag:ident $scanner:expr, [$($args:tt)*] ($($tuple:tt)*) $($t:tt)*)        => { $crate::scan_value!(@$tag $scanner, [$($args)* [$crate::scan_value!(@tuple $scanner, [] $($tuple)*)]] $($t)*) };
165    (@$tag:ident $scanner:expr, [$($args:tt)*] [$($tt:tt)*] $($t:tt)*)           => { $crate::scan_value!(@$tag $scanner, [$($args)* [$crate::scan_value!(@sparen $scanner, [] $($tt)*)]] $($t)*) };
166    (@$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)*)?) };
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    (iter = $iter:expr, $($t:tt)*)                                               => { { let mut __scanner = Scanner::new_from_iter($iter); $crate::scan_value!(@inner __scanner, [] $($t)*) } };
172    ($scanner:expr, $($t:tt)*)                                                   => { $crate::scan_value!(@inner $scanner, [] $($t)*) }
173}
174
175/// scan and bind values with Scanner
176///
177/// - `scan!(scanner, $($pat $(: ELEMENT)?),*)`
178#[macro_export]
179macro_rules! scan {
180    (@assert $p:pat) => {};
181    (@assert $($p:tt)*) => { ::std::compile_error!(::std::concat!("expected pattern, found `", ::std::stringify!($($p)*), "`")); };
182    (@pat $scanner:expr, [] [])                                                     => {};
183    (@pat $scanner:expr, [] [] , $($t:tt)*)                                         => { $crate::scan!(@pat $scanner, [] [] $($t)*) };
184    (@pat $scanner:expr, [$($p:tt)*] [] $x:ident $($t:tt)*)                         => { $crate::scan!(@pat $scanner, [$($p)* $x] [] $($t)*) };
185    (@pat $scanner:expr, [$($p:tt)*] [] :: $($t:tt)*)                               => { $crate::scan!(@pat $scanner, [$($p)* ::] [] $($t)*) };
186    (@pat $scanner:expr, [$($p:tt)*] [] & $($t:tt)*)                                => { $crate::scan!(@pat $scanner, [$($p)* &] [] $($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)*] [] {$($x:tt)*} $($t:tt)*)                      => { $crate::scan!(@pat $scanner, [$($p)* {$($x)*}] [] $($t)*) };
190    (@pat $scanner:expr, [$($p:tt)*] [] : $($t:tt)*)                                => { $crate::scan!(@ty  $scanner, [$($p)*] [] $($t)*) };
191    (@pat $scanner:expr, [$($p:tt)*] [] $($t:tt)*)                                  => { $crate::scan!(@let $scanner, [$($p)*] [usize] $($t)*) };
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 = $e:expr $(, $($t:tt)*)?) => { $crate::scan!(@let $scanner, [$($p)*] [$($tt)* $ty = $e] $(, $($t)*)?) };
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    (iter = $iter:expr, $($t:tt)*) => { let mut __scanner = Scanner::new_from_iter($iter); $crate::scan!(@pat __scanner, [] [] $($t)*) };
204    ($scanner:expr, $($t:tt)*) => { $crate::scan!(@pat $scanner, [] [] $($t)*) }
205}
206
207/// define enum scan rules
208///
209/// # Example
210/// ```rust
211/// # use competitive::{define_enum_scan, tools::{CharsWithBase, IterScan, Scanner, Usize1}};
212/// define_enum_scan! {
213///   enum Query: u8 {
214///     0 => Noop,
215///     1 => Args { i: Usize1, s: char },
216///     9 => Complex { n: usize, c: [(usize, Vec<usize> = CharsWithBase('a')); n] },
217///   }
218/// }
219/// ```
220#[macro_export]
221macro_rules! define_enum_scan {
222    (@field_ty @repeat [$($t:tt)*] $($len:expr)?)                           => { Vec<$crate::define_enum_scan!(@field_ty $($t)*)> };
223    (@field_ty @array [$($t:tt)*] $len:expr)                                => { [$crate::define_enum_scan!(@field_ty $($t)*); $len] };
224    (@field_ty @tuple [$([$($args:tt)*])*])                                 => { ($( $($args)* ,)*) };
225    (@field_ty @sparen [] ($($tt:tt)*); $($t:tt)*)                          => { $crate::define_enum_scan!(@field_ty @sparen [($($tt)*)] $($t)*) };
226    (@field_ty @sparen [] [$($tt:tt)*]; $($t:tt)*)                          => { $crate::define_enum_scan!(@field_ty @sparen [[$($tt)*]] $($t)*) };
227    (@field_ty @sparen [] $ty:ty = $e:expr; $($t:tt)*)                      => { $crate::define_enum_scan!(@field_ty @sparen [$ty = $e] $($t)*) };
228    (@field_ty @sparen [] $ty:ty; $($t:tt)*)                                => { $crate::define_enum_scan!(@field_ty @sparen [$ty] $($t)*) };
229    (@field_ty @sparen [] $($args:tt)*)                                     => { $crate::define_enum_scan!(@field_ty @repeat [$($args)*]) };
230    (@field_ty @sparen [$($args:tt)+] const $len:expr)                      => { $crate::define_enum_scan!(@field_ty @array [$($args)+] $len) };
231    (@field_ty @sparen [$($args:tt)+] $len:expr)                            => { $crate::define_enum_scan!(@field_ty @repeat [$($args)+] $len) };
232    (@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)*) };
233    (@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)*) };
234    (@field_ty @$tag:ident [$($args:tt)*] $ty:ty = $e:expr $(, $($t:tt)*)?) => { $crate::define_enum_scan!(@field_ty @$tag [$($args)* [$ty]] $(, $($t)*)?) };
235    (@field_ty @$tag:ident [$($args:tt)*] $ty:ty $(, $($t:tt)*)?)           => { $crate::define_enum_scan!(@field_ty @$tag [$($args)* [<$ty as IterScan>::Output]] $(, $($t)*)?) };
236    (@field_ty @$tag:ident [$($args:tt)*] , $($t:tt)*)                      => { $crate::define_enum_scan!(@field_ty @$tag [$($args)*] $($t)*) };
237    (@field_ty @$tag:ident [[$($args:tt)*]])                                => { $($args)* };
238    (@field_ty @$tag:ident [$($args:tt)*])                                  => { ::std::compile_error!(::std::stringify!($($args)*)) };
239    (@field_ty $($t:tt)*) => { $crate::define_enum_scan!(@field_ty @inner [] $($t)*) };
240
241    (@tag_expr raw, $iter:ident) => { $iter.next()? };
242    (@tag_expr $d:ty, $iter:ident) => { <$d as IterScan>::scan($iter)? };
243    (@variant ([$($attr:tt)*] $vis:vis $T:ident $d:tt) [$($vars:tt)*]) => { $crate::define_enum_scan! { @def $($attr)* $vis enum $T : $d { $($vars)* } } };
244    (@variant $ctx:tt [$($vars:tt)*] $p:pat => $v:ident { $($fs:tt)* } $($rest:tt)*) => { $crate::define_enum_scan! { @field   $ctx [$($vars)*] $p => $v [] $($fs)* ; $($rest)* } };
245    (@variant $ctx:tt [$($vars:tt)*] $p:pat => $v:ident $($rest:tt)*)                    => { $crate::define_enum_scan! { @variant $ctx [$($vars)* $p => $v ,] $($rest)* } };
246    (@variant $ctx:tt [$($vars:tt)*] , $($rest:tt)*)                                     => { $crate::define_enum_scan! { @variant $ctx [$($vars)*] $($rest)* } };
247    (@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)* } };
248    (@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)* } };
249    (@field $ctx:tt [$($vars:tt)*] $p:pat => $v:ident [$($fs:tt)*] ; $($rest:tt)*)                                  => { $crate::define_enum_scan! { @variant $ctx [$($vars)* $p => $v { $($fs)* } ,] $($rest)* } };
250    (@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)* } };
251    (@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)* } };
252    (@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)* } };
253    (@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)* } };
254    (@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)* } };
255    (@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)* } };
256    (
257        @def
258        $(#[$attr:meta])*
259        $vis:vis enum $T:ident : $d:tt {
260            $( $p:pat => $v:ident $( { $( [$f:ident : $($spec:tt)*] )* } )?, )*
261        }
262    ) => {
263        $(#[$attr])*
264        $vis enum $T {
265            $( $v $( { $( $f : $crate::define_enum_scan!(@field_ty $($spec)*) ),* } )? ),*
266        }
267        impl IterScan for $T {
268            type Output = Self;
269            fn scan<'a, I: Iterator<Item = &'a str>>(iter: &mut I) -> Option<Self> {
270                let tag = $crate::define_enum_scan!(@tag_expr $d, iter);
271                match tag {
272                    $(
273                        $p => {
274                            $($(
275                                let $f = $crate::scan_value!(iter = &mut *iter, $($spec)* );
276                            )*)?
277                            Some($T::$v $( { $( $f ),* } )?)
278                        }
279                    ),*
280                    _ => None,
281                }
282            }
283        }
284    };
285    (
286        $(#[$attr:meta])*
287        $vis:vis enum $T:ident : raw {
288            $($body:tt)*
289        }
290    ) => {
291        $crate::define_enum_scan! { @variant ([$(#[$attr])*] $vis $T raw) [] $($body)* }
292    };
293    (
294        $(#[$attr:meta])*
295        $vis:vis enum $T:ident : $d:ty {
296            $($body:tt)*
297        }
298    ) => {
299        $crate::define_enum_scan! { @variant ([$(#[$attr])*] $vis $T $d) [] $($body)* }
300    };
301}
302
303#[derive(Debug, Copy, Clone)]
304pub enum Usize1 {}
305impl IterScan for Usize1 {
306    type Output = usize;
307    fn scan<'a, I: Iterator<Item = &'a str>>(iter: &mut I) -> Option<Self::Output> {
308        <usize as IterScan>::scan(iter)?.checked_sub(1)
309    }
310}
311#[derive(Debug, Copy, Clone)]
312pub struct CharWithBase(pub char);
313impl MarkedIterScan for CharWithBase {
314    type Output = usize;
315    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output> {
316        Some((<char as IterScan>::scan(iter)? as u8 - self.0 as u8) as usize)
317    }
318}
319#[derive(Debug, Copy, Clone)]
320pub enum Chars {}
321impl IterScan for Chars {
322    type Output = Vec<char>;
323    fn scan<'a, I: Iterator<Item = &'a str>>(iter: &mut I) -> Option<Self::Output> {
324        Some(iter.next()?.chars().collect())
325    }
326}
327#[derive(Debug, Copy, Clone)]
328pub struct CharsWithBase(pub char);
329impl MarkedIterScan for CharsWithBase {
330    type Output = Vec<usize>;
331    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output> {
332        Some(
333            iter.next()?
334                .chars()
335                .map(|c| (c as u8 - self.0 as u8) as usize)
336                .collect(),
337        )
338    }
339}
340#[derive(Debug, Copy, Clone)]
341pub enum Byte1 {}
342impl IterScan for Byte1 {
343    type Output = u8;
344    fn scan<'a, I: Iterator<Item = &'a str>>(iter: &mut I) -> Option<Self::Output> {
345        let bytes = iter.next()?.as_bytes();
346        assert_eq!(bytes.len(), 1);
347        Some(bytes[0])
348    }
349}
350#[derive(Debug, Copy, Clone)]
351pub struct ByteWithBase(pub u8);
352impl MarkedIterScan for ByteWithBase {
353    type Output = usize;
354    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output> {
355        Some((<char as IterScan>::scan(iter)? as u8 - self.0) as usize)
356    }
357}
358#[derive(Debug, Copy, Clone)]
359pub enum Bytes {}
360impl IterScan for Bytes {
361    type Output = Vec<u8>;
362    fn scan<'a, I: Iterator<Item = &'a str>>(iter: &mut I) -> Option<Self::Output> {
363        Some(iter.next()?.bytes().collect())
364    }
365}
366#[derive(Debug, Copy, Clone)]
367pub struct BytesWithBase(pub u8);
368impl MarkedIterScan for BytesWithBase {
369    type Output = Vec<usize>;
370    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output> {
371        Some(
372            iter.next()?
373                .bytes()
374                .map(|c| (c - self.0) as usize)
375                .collect(),
376        )
377    }
378}
379#[derive(Debug, Copy, Clone)]
380pub struct Collect<T, B = Vec<<T as IterScan>::Output>>
381where
382    T: IterScan,
383    B: FromIterator<<T as IterScan>::Output>,
384{
385    size: usize,
386    _marker: PhantomData<fn() -> (T, B)>,
387}
388impl<T, B> Collect<T, B>
389where
390    T: IterScan,
391    B: FromIterator<<T as IterScan>::Output>,
392{
393    pub fn new(size: usize) -> Self {
394        Self {
395            size,
396            _marker: PhantomData,
397        }
398    }
399}
400impl<T, B> MarkedIterScan for Collect<T, B>
401where
402    T: IterScan,
403    B: FromIterator<<T as IterScan>::Output>,
404{
405    type Output = B;
406    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output> {
407        repeat_with(|| <T as IterScan>::scan(iter))
408            .take(self.size)
409            .collect()
410    }
411}
412#[derive(Debug, Copy, Clone)]
413pub struct SizedCollect<T, B = Vec<<T as IterScan>::Output>>
414where
415    T: IterScan,
416    B: FromIterator<<T as IterScan>::Output>,
417{
418    _marker: PhantomData<fn() -> (T, B)>,
419}
420impl<T, B> IterScan for SizedCollect<T, B>
421where
422    T: IterScan,
423    B: FromIterator<<T as IterScan>::Output>,
424{
425    type Output = B;
426    fn scan<'a, I: Iterator<Item = &'a str>>(iter: &mut I) -> Option<Self::Output> {
427        let size = usize::scan(iter)?;
428        repeat_with(|| <T as IterScan>::scan(iter))
429            .take(size)
430            .collect()
431    }
432}
433#[derive(Debug, Copy, Clone)]
434pub struct Splitted<T, P>
435where
436    T: IterScan,
437{
438    pat: P,
439    _marker: PhantomData<fn() -> T>,
440}
441impl<T, P> Splitted<T, P>
442where
443    T: IterScan,
444{
445    pub fn new(pat: P) -> Self {
446        Self {
447            pat,
448            _marker: PhantomData,
449        }
450    }
451}
452impl<T> MarkedIterScan for Splitted<T, char>
453where
454    T: IterScan,
455{
456    type Output = Vec<<T as IterScan>::Output>;
457    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output> {
458        let mut iter = iter.next()?.split(self.pat);
459        Some(from_fn(|| <T as IterScan>::scan(&mut iter)).collect())
460    }
461}
462impl<T> MarkedIterScan for Splitted<T, &str>
463where
464    T: IterScan,
465{
466    type Output = Vec<<T as IterScan>::Output>;
467    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output> {
468        let mut iter = iter.next()?.split(self.pat);
469        Some(from_fn(|| <T as IterScan>::scan(&mut iter)).collect())
470    }
471}
472impl<T, F> MarkedIterScan for F
473where
474    F: Fn(&str) -> Option<T>,
475{
476    type Output = T;
477    fn mscan<'a, I: Iterator<Item = &'a str>>(self, iter: &mut I) -> Option<Self::Output> {
478        self(iter.next()?)
479    }
480}
481
482#[cfg(test)]
483mod tests {
484    use super::*;
485
486    #[test]
487    fn test_scan() {
488        let mut s = Scanner::new("1 2 3 a 1 2 1 1 1.1 2 3");
489        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]);
490        assert_eq!(x, 1);
491        assert_eq!(y, '2');
492        assert_eq!(z, 2);
493        assert_eq!(a, 0);
494        assert_eq!(b, vec![1, 2]);
495        assert_eq!(c, (1, 1));
496        assert_eq!(d, vec![1, 1]);
497        assert_eq!(e, [2, 3]);
498
499        scan!(src = "12 34", c: Vec<usize> = CharsWithBase('0'), d: [Vec<usize> = CharsWithBase('0'); 1]);
500        assert_eq!(c, vec![1, 2]);
501        assert_eq!(d, vec![vec![3, 4]]);
502
503        scan!(src = "1", x);
504        assert_eq!(x, 1);
505        assert_eq!(scan_value!(src = "1", usize), 1);
506
507        scan!(iter = "1".split_ascii_whitespace(), x);
508        assert_eq!(x, 1);
509        assert_eq!(scan_value!(iter = "1".split_ascii_whitespace(), usize), 1);
510    }
511
512    #[test]
513    fn test_define_enum_scan() {
514        define_enum_scan! {
515            enum Query: u8 {
516                0 => Noop,
517                1 => Args { i: Usize1, s: char },
518                9 => Complex { n: usize, c: [(usize, Vec<usize> = CharsWithBase('a')); n] },
519            }
520        }
521
522        let mut s = Scanner::new("0   1 2 a  9 2 3 ab 2 ab");
523        scan!(s, q1: Query, q2: Query, q3: Query);
524        match q1 {
525            Query::Noop => {}
526            _ => panic!("unexpected"),
527        }
528        match q2 {
529            Query::Args { i, s } => {
530                assert_eq!(i, 1);
531                assert_eq!(s, 'a');
532            }
533            _ => panic!("unexpected"),
534        }
535        match q3 {
536            Query::Complex { n, c } => {
537                assert_eq!(n, 2);
538                assert_eq!(c, vec![(3, vec![0, 1]), (2, vec![0, 1])]);
539            }
540            _ => panic!("unexpected"),
541        }
542    }
543}
competitive/tools/scanner.rs

competitive/tools/
scanner.rs
