input/
lib.rs

1use std::{
2    any::type_name,
3    fmt::Debug,
4    io::{BufRead, BufReader, Stdin},
5    marker::PhantomData,
6    str::FromStr,
7    sync::{Mutex, OnceLock},
8};
9
10pub static STDIN_SOURCE: OnceLock<Mutex<AutoSource<BufReader<Stdin>>>> =
11    OnceLock::new();
12
13pub trait Source<R: BufRead> {
14    fn next_token(&mut self) -> Option<String>;
15    fn next_token_unwrap(&mut self) -> String { self.next_token().unwrap() }
16}
17
18impl<R: BufRead, S: Source<R>> Source<R> for &'_ mut S {
19    fn next_token(&mut self) -> Option<String> { (*self).next_token() }
20}
21
22pub type AutoSource<R> = OnceSource<R>;
23
24pub struct OnceSource<R: BufRead> {
25    tokens: std::vec::IntoIter<String>,
26    _phantom: PhantomData<R>,
27}
28
29impl<R: BufRead> OnceSource<R> {
30    pub fn new(mut source: R) -> Self {
31        let mut context = "".to_owned();
32        source.read_to_string(&mut context).unwrap();
33        let tokens: Vec<_> =
34            context.split_whitespace().map(|s| s.to_owned()).collect();
35        Self {
36            tokens: tokens.into_iter(),
37            _phantom: PhantomData,
38        }
39    }
40}
41
42impl<R: BufRead> Source<R> for OnceSource<R> {
43    fn next_token(&mut self) -> Option<String> { self.tokens.next() }
44}
45
46impl<'a> From<&'a str> for OnceSource<BufReader<&'a [u8]>> {
47    fn from(s: &'a str) -> Self {
48        OnceSource::new(BufReader::new(s.as_bytes()))
49    }
50}
51
52pub trait Readable {
53    type Output;
54    fn read<R: BufRead, S: Source<R>>(source: &mut S) -> Self::Output;
55}
56
57impl<T: FromStr> Readable for T
58where
59    T::Err: Debug,
60{
61    type Output = T;
62    fn read<R: BufRead, S: Source<R>>(source: &mut S) -> T {
63        let token = source.next_token_unwrap();
64        match token.parse() {
65            Ok(v) => v,
66            Err(e) => panic!(
67                "`{input}` `{ty}` `{err:?}`",
68                input = token,
69                ty = type_name::<T>(),
70                err = e
71            ),
72        }
73    }
74}
75
76#[macro_export]
77macro_rules! scan {
78    // terminator
79    (@from [$source:expr] @rest) => {};
80
81    // parse mutability
82    (@from [$source:expr] @rest mut $($rest:tt)*) => {
83        $crate::scan! {
84            @from [$source]
85            @mut [mut]
86            @rest $($rest)*
87        }
88    };
89    (@from [$source:expr] @rest $($rest:tt)*) => {
90        $crate::scan! {
91            @from [$source]
92            @mut []
93            @rest $($rest)*
94        }
95    };
96
97    // parse variable pattern
98    (@from [$source:expr] @mut [$($mut:tt)?] @rest $var:tt: $($rest:tt)*) => {
99        $crate::scan! {
100            @from [$source]
101            @mut [$($mut)?]
102            @var $var
103            @kind []
104            @rest $($rest)*
105        }
106    };
107
108    // parse kind (type)
109    (@from [$source:expr] @mut [$($mut:tt)?] @var $var:tt @kind [$($kind:tt)*] @rest) => {
110        let $($mut)? $var = $crate::read_value!(@source [$source] @kind [$($kind)*] @hint [$var]);
111    };
112    (@from [$source:expr] @mut [$($mut:tt)?] @var $var:tt @kind [$($kind:tt)*] @rest, $($rest:tt)*) => {
113        $crate::scan!(@from [$source] @mut [$($mut)?] @var $var @kind [$($kind)*] @rest);
114        $crate::scan!(@from [$source] @rest $($rest)*);
115    };
116    (@from [$source:expr] @mut [$($mut:tt)?] @var $var:tt @kind [$($kind:tt)*] @rest $tt:tt $($rest:tt)*) => {
117        $crate::scan!(@from [$source] @mut [$($mut)?] @var $var @kind [$($kind)* $tt] @rest $($rest)*);
118    };
119
120    (from $source:expr, $($rest:tt)*) => {
121        #[allow(unused_variables, unused_mut)]
122        let mut s = $source;
123        $crate::scan! {
124            @from [&mut s]
125            @rest $($rest)*
126        }
127    };
128    ($($rest:tt)*) => {
129        let mut locked_stdin = $crate::STDIN_SOURCE.get_or_init(|| {
130            ::std::sync::Mutex::new($crate::AutoSource::new(::std::io::BufReader::new(::std::io::stdin())))
131        }).lock().unwrap();
132        $crate::scan! {
133            @from [&mut *locked_stdin]
134            @rest $($rest)*
135        }
136        drop(locked_stdin);
137    };
138}
139
140#[macro_export]
141macro_rules! read_value {
142    // array and variable-length array
143    (@source [$source:expr] @kind [[$($kind:tt)*]] @hint [$var:tt]) => {
144        $crate::read_value!(@vec @source [$source] @kind [] @hint [$var] @rest $($kind)*)
145    };
146    (@vec @source [$source:expr] @kind [$($kind:tt)*] @hint [$var:tt] @rest) => {{
147        let len = <usize as $crate::Readable>::read($source);
148        $crate::read_value!(@source [$source] @kind [[$($kind)*; len]] @hint [$var])
149    }};
150    (@vec @source [$source:expr] @kind [$($kind:tt)*] @hint[$var:tt] @rest ; const $($rest:tt)*) => {
151        $crate::read_value!(@array @source [$source] @kind [$($kind)*] @hint [$var] @len [$($rest)*])
152    };
153    (@vec @source [$source:expr] @kind [$($kind:tt)*] @hint[$var:tt] @rest ; const $($rest:tt)*) => {
154        $crate::read_value!(@array @source [$source] @kind [$($kind)*] @hint[$var] @len [$($rest)*])
155    };
156    (@vec @source [$source:expr] @kind [$($kind:tt)*] @hint[$var:tt] @rest ; _) => {
157        $crate::read_value!(@array @source [$source] @kind [$($kind)*] @hint [$var] @len [_])
158    };
159    (@vec @source [$source:expr] @kind [$($kind:tt)*] @hint[$var:tt] @rest ; $($rest:tt)*) => {
160        $crate::read_value!(@vec @source [$source] @kind [$($kind)*] @hint [$var] @len [$($rest)*])
161    };
162    (@vec @source [$source:expr] @kind [$($kind:tt)*] @hint[$var:tt] @rest $tt:tt $($rest:tt)*) => {
163        $crate::read_value!(@vec @source [$source] @kind [$($kind)* $tt] @hint [$var] @rest $($rest)*)
164    };
165    (@vec @source [$source:expr] @kind [$($kind:tt)*] @hint[$var:tt] @len [$($len:tt)*]) => {{
166        let len = $($len)*;
167        (0..len)
168            .map(|_| $crate::read_value!(@source [$source] @kind [$($kind)*] @hint [$var]))
169            .collect::<Vec<_>>()
170    }};
171    (@array @source [$source:expr] @kind [$($kind:tt)*] @hint[$var:tt] @len [_]) => {{
172        const LEN: usize = {
173            const fn zero_array<const N: usize>() -> [(); N] { [(); N] }
174            let $var = zero_array();
175            let a = $var;
176            a.len()
177        };
178        let mut tmp = [Default::default(); LEN];
179        for i in 0..LEN {
180            tmp[i] = $crate::read_value!(@source [$source] @kind [$($kind)*] @hint [$var])
181        }
182        tmp
183    }};
184    (@array @source [$source:expr] @kind [$($kind:tt)*] @hint[$var:tt] @len [$($len:tt)*]) => {{
185        const LEN: usize = $($len)*;
186        let mut tmp = [Default::default(); LEN];
187        for i in 0..LEN {
188            tmp[i] = $crate::read_value!(@source [$source] @kind [$($kind)*] @hint [$var])
189        }
190        tmp
191    }};
192
193    // tuple
194    (@source [$source:expr] @kind [($($kinds:tt)*)] @hint [$var:tt]) => {
195        $crate::read_value!(@tuple @source [$source] @kinds [] @current [] @hint [$var] @rest $($kinds)*)
196    };
197    (@tuple @source [$source:expr] @kinds [$([$($kind:tt)*])*] @current [] @hint [$var:tt] @rest) => {
198        (
199            $($crate::read_value!(@source [$source] @kind [$($kind)*] @hint [$var]),)*
200        )
201    };
202    (@tuple @source [$source:expr] @kinds [$($kinds:tt)*] @current [$($curr:tt)*] @hint [$var:tt] @rest) => {
203        $crate::read_value!(@tuple @source [$source] @kinds [$($kinds)* [$($curr)*]] @current [] @hint [$var] @rest)
204    };
205    (@tuple @source [$source:expr] @kinds [$($kinds:tt)*] @current [$($curr:tt)*] @hint [$var:tt] @rest, $($rest:tt)*) => {
206        $crate::read_value!(@tuple @source [$source] @kinds [$($kinds)* [$($curr)*]] @current [] @hint [$var] @rest $($rest)*)
207    };
208    (@tuple @source [$source:expr] @kinds [$($kinds:tt)*] @current [$($curr:tt)*] @hint [$var:tt] @rest $tt:tt $($rest:tt)*) => {
209        $crate::read_value!(@tuple @source [$source] @kinds [$($kinds)*] @current [$($curr)* $tt] @hint [$var] @rest $($rest)*)
210    };
211
212    // unreachable
213    (@source [$source:expr] @kind [] @hint [$var:tt]) => {
214        compile_error!("Reached unreachable statement while parsing macro input.")
215    };
216
217    // normal other
218    (@source [$source:expr] @kind [$kind:ty] @hint [$var:tt]) => {
219        <$kind as $crate::Readable>::read($source)
220    };
221}
222
223#[test]
224fn sanity_check() {
225    // primitives
226    let src = AutoSource::from("1 2 3 4");
227    scan! {
228        from src,
229        int: u32,
230        frac: f64,
231        ch: char,
232        string: String,
233    }
234    assert_eq!((int, frac, ch, string), (1, 2.0, '3', "4".to_owned()));
235
236    // lists
237    let src = AutoSource::from("2 1 2 3 1 2 3 1 2 3 1 2 3");
238    scan! {
239        from src,
240        n: usize,
241        vec_n: [u32; n],
242        vec: [u32],
243        array: [u32; const 3],
244        [un, pac, ked]: [u32; _],
245    }
246    assert_eq!((vec_n, vec, array), (vec![1, 2], vec![1, 2, 3], [1, 2, 3]));
247    assert_eq!([un, pac, ked], [1, 2, 3]);
248
249    // tuples
250    let src = AutoSource::from("1 2 3 4 5 6 7 8");
251    scan! {
252        from src,
253        a: (i32, i32, (i32, i32), i32),
254        (b, c): (i32, i32,),
255        (d,): (i32,),
256    }
257    assert_eq!((a, b, c, d), ((1, 2, (3, 4), 5), 6, 7, 8));
258}
input/lib.rs

input/
lib.rs
