competitive/tools/
iter_print.rs

1use std::{
2    fmt::Display,
3    io::{Error, Write},
4};
5
6pub trait IterPrint {
7    fn iter_print<W, S>(self, writer: &mut W, sep: S, is_head: bool) -> Result<(), Error>
8    where
9        W: Write,
10        S: Display;
11}
12macro_rules! impl_iter_print_tuple {
13    (@impl $($A:ident $a:ident)?, $($B:ident $b:ident)*) => {
14        impl<$($A,)? $($B),*> IterPrint for ($($A,)? $($B),*)
15        where
16            $($A: Display,)? $($B: Display),*
17        {
18            #[allow(unused_variables)]
19            fn iter_print<W, S>(self, writer: &mut W, sep: S, is_head: bool) -> Result<(), Error>
20            where
21                W: Write,
22                S: Display
23            {
24                let ($($a,)? $($b,)*) = self;
25                $(
26                    if is_head {
27                        ::std::write!(writer, "{}", $a)?;
28                    } else {
29                        ::std::write!(writer, "{}{}", sep, $a)?;
30                    }
31                )?
32                $( ::std::write!(writer, "{}{}", sep, $b)?; )*
33                Ok(())
34            }
35        }
36    };
37    (@inc , , $C:ident $c:ident $($D:ident $d:ident)*) => {
38        impl_iter_print_tuple!(@impl ,);
39        impl_iter_print_tuple!(@inc $C $c, , $($D $d)*);
40    };
41    (@inc $A:ident $a:ident, $($B:ident $b:ident)*, $C:ident $c:ident $($D:ident $d:ident)*) => {
42        impl_iter_print_tuple!(@impl $A $a, $($B $b)*);
43        impl_iter_print_tuple!(@inc $A $a, $($B $b)* $C $c, $($D $d)*);
44    };
45    (@inc $A:ident $a:ident, $($B:ident $b:ident)*,) => {
46        impl_iter_print_tuple!(@impl $A $a, $($B $b)*);
47    };
48    ($($t:tt)*) => {
49        impl_iter_print_tuple!(@inc , , $($t)*);
50    };
51}
52impl_iter_print_tuple!(A a B b C c D d E e F f G g H h I i J j K k);
53
54/// Print expressions with a separator.
55/// - `iter_print!(writer, args...)`
56/// - `@sep $expr`: set separator (default: `' '`)
57/// - `@ns`: alias for `@sep ""`
58/// - `@lf`: alias for `@sep '\n'`
59/// - `@sp`: alias for `@sep ' '`
60/// - `@fmt ($lit, $($expr),*)`: print `format!($lit, $($expr),*)`
61/// - `@flush`: flush writer (auto insert `!`)
62/// - `@it $expr`: print iterator
63/// - `@it1 $expr`: print iterator as 1-indexed
64/// - `@cw ($char $expr)`: print iterator as `(elem as u8 + $char as u8) as char`
65/// - `@bw ($byte $expr)`: print iterator as `(elem as u8 + $byte) as char`
66/// - `@it2d $expr`: print 2d-iterator
67/// - `@tup $expr`: print tuple (need to import [`IterPrint`])
68/// - `@ittup $expr`: print iterative tuple (need to import [`IterPrint`])
69/// - `$expr`: print expr
70/// - `{ args... }`: scoped
71/// - `;`: print `'\n'`
72/// - `!`: not print `'\n'` at the end
73#[macro_export]
74macro_rules! iter_print {
75    (@@fmt $writer:expr, $sep:expr, $is_head:expr, ($lit:literal $(, $e:expr)* $(,)?)) => {
76        if !$is_head {
77            ::std::write!($writer, "{}", $sep).expect("io error");
78        }
79        ::std::write!($writer, $lit, $($e),*).expect("io error");
80    };
81    (@@item $writer:expr, $sep:expr, $is_head:expr, $e:expr) => {
82        $crate::iter_print!(@@fmt $writer, $sep, $is_head, ("{}", $e));
83    };
84    (@@line_feed $writer:expr $(,)?) => {
85        ::std::writeln!($writer).expect("io error");
86    };
87    (@@it $writer:expr, $sep:expr, $is_head:expr, $iter:expr) => {{
88        let mut iter = $iter.into_iter();
89        if let Some(item) = iter.next() {
90            $crate::iter_print!(@@item $writer, $sep, $is_head, item);
91        }
92        for item in iter {
93            $crate::iter_print!(@@item $writer, $sep, false, item);
94        }
95    }};
96    (@@it1 $writer:expr, $sep:expr, $is_head:expr, $iter:expr) => {{
97        let mut iter = $iter.into_iter();
98        if let Some(item) = iter.next() {
99            $crate::iter_print!(@@item $writer, $sep, $is_head, item + 1);
100        }
101        for item in iter {
102            $crate::iter_print!(@@item $writer, $sep, false, item + 1);
103        }
104    }};
105    (@@cw $writer:expr, $sep:expr, $is_head:expr, ($ch:literal $iter:expr)) => {{
106        let mut iter = $iter.into_iter();
107        let b = $ch as u8;
108        if let Some(item) = iter.next() {
109            $crate::iter_print!(@@item $writer, $sep, $is_head, (item as u8 + b) as char);
110        }
111        for item in iter {
112            $crate::iter_print!(@@item $writer, $sep, false, (item as u8 + b) as char);
113        }
114    }};
115    (@@bw $writer:expr, $sep:expr, $is_head:expr, ($b:literal $iter:expr)) => {{
116        let mut iter = $iter.into_iter();
117        let b: u8 = $b;
118        if let Some(item) = iter.next() {
119            $crate::iter_print!(@@item $writer, $sep, $is_head, (item as u8 + b) as char);
120        }
121        for item in iter {
122            $crate::iter_print!(@@item $writer, $sep, false, (item as u8 + b) as char);
123        }
124    }};
125    (@@it2d $writer:expr, $sep:expr, $is_head:expr, $iter:expr) => {
126        let mut iter = $iter.into_iter();
127        if let Some(item) = iter.next() {
128            $crate::iter_print!(@@it $writer, $sep, $is_head, item);
129        }
130        for item in iter {
131            $crate::iter_print!(@@line_feed $writer);
132            $crate::iter_print!(@@it $writer, $sep, true, item);
133        }
134    };
135    (@@tup $writer:expr, $sep:expr, $is_head:expr, $tuple:expr) => {
136        IterPrint::iter_print($tuple, &mut $writer, $sep, $is_head).expect("io error");
137    };
138    (@@ittup $writer:expr, $sep:expr, $is_head:expr, $iter:expr) => {
139        let mut iter = $iter.into_iter();
140        if let Some(item) = iter.next() {
141            $crate::iter_print!(@@tup $writer, $sep, $is_head, item);
142        }
143        for item in iter {
144            $crate::iter_print!(@@line_feed $writer);
145            $crate::iter_print!(@@tup $writer, $sep, true, item);
146        }
147    };
148    (@@assert_tag item) => {};
149    (@@assert_tag it) => {};
150    (@@assert_tag it1) => {};
151    (@@assert_tag it2d) => {};
152    (@@assert_tag tup) => {};
153    (@@assert_tag ittup) => {};
154    (@@assert_tag $tag:ident) => {
155        ::std::compile_error!(::std::concat!("invalid tag in `iter_print!`: `", std::stringify!($tag), "`"));
156    };
157    (@@inner $writer:expr, $sep:expr, $is_head:expr, @sep $e:expr, $($t:tt)*) => {
158        $crate::iter_print!(@@inner $writer, $e, $is_head, $($t)*);
159    };
160    (@@inner $writer:expr, $sep:expr, $is_head:expr, @ns $($t:tt)*) => {
161        $crate::iter_print!(@@inner $writer, "", $is_head, $($t)*);
162    };
163    (@@inner $writer:expr, $sep:expr, $is_head:expr, @lf $($t:tt)*) => {
164        $crate::iter_print!(@@inner $writer, '\n', $is_head, $($t)*);
165    };
166    (@@inner $writer:expr, $sep:expr, $is_head:expr, @sp $($t:tt)*) => {
167        $crate::iter_print!(@@inner $writer, ' ', $is_head, $($t)*);
168    };
169    (@@inner $writer:expr, $sep:expr, $is_head:expr, @flush $($t:tt)*) => {
170        $writer.flush().expect("io error");
171        $crate::iter_print!(@@inner $writer, $sep, $is_head, ! $($t)*);
172    };
173    (@@inner $writer:expr, $sep:expr, $is_head:expr, @fmt $arg:tt $($t:tt)*) => {
174        $crate::iter_print!(@@fmt $writer, $sep, $is_head, $arg);
175        $crate::iter_print!(@@inner $writer, $sep, $is_head, $($t)*);
176    };
177    (@@inner $writer:expr, $sep:expr, $is_head:expr, @cw $arg:tt $($t:tt)*) => {
178        $crate::iter_print!(@@cw $writer, $sep, $is_head, $arg);
179        $crate::iter_print!(@@inner $writer, $sep, $is_head, $($t)*);
180    };
181    (@@inner $writer:expr, $sep:expr, $is_head:expr, @bw $arg:tt $($t:tt)*) => {
182        $crate::iter_print!(@@bw $writer, $sep, $is_head, $arg);
183        $crate::iter_print!(@@inner $writer, $sep, $is_head, $($t)*);
184    };
185    (@@inner $writer:expr, $sep:expr, $is_head:expr, @$tag:ident $e:expr, $($t:tt)*) => {
186        $crate::iter_print!(@@assert_tag $tag);
187        $crate::iter_print!(@@$tag $writer, $sep, $is_head, $e);
188        $crate::iter_print!(@@inner $writer, $sep, false, $($t)*);
189    };
190    (@@inner $writer:expr, $sep:expr, $is_head:expr, @$tag:ident $e:expr; $($t:tt)*) => {
191        $crate::iter_print!(@@assert_tag $tag);
192        $crate::iter_print!(@@$tag $writer, $sep, $is_head, $e);
193        $crate::iter_print!(@@line_feed $writer);
194        $crate::iter_print!(@@inner $writer, $sep, true, $($t)*);
195    };
196    (@@inner $writer:expr, $sep:expr, $is_head:expr, @$tag:ident $e:expr) => {
197        $crate::iter_print!(@@assert_tag $tag);
198        $crate::iter_print!(@@$tag $writer, $sep, $is_head, $e);
199        $crate::iter_print!(@@inner $writer, $sep, false,);
200    };
201    (@@inner $writer:expr, $sep:expr, $is_head:expr, @$tag:ident $($t:tt)*) => {
202        ::std::compile_error!(::std::concat!("invalid expr in `iter_print!`: `", std::stringify!($($t)*), "`"));
203    };
204    (@@inner $writer:expr, $sep:expr, $is_head:expr, , $($t:tt)*) => {
205        $crate::iter_print!(@@inner $writer, $sep, $is_head, $($t)*);
206    };
207    (@@inner $writer:expr, $sep:expr, $is_head:expr, ; $($t:tt)*) => {
208        $crate::iter_print!(@@line_feed $writer);
209        $crate::iter_print!(@@inner $writer, $sep, $is_head, $($t)*);
210    };
211    (@@inner $writer:expr, $sep:expr, $is_head:expr, ! $(,)?) => {};
212    (@@inner $writer:expr, $sep:expr, $is_head:expr, ! $($t:tt)*) => {
213        $crate::iter_print!(@@inner $writer, $sep, $is_head, $($t)*);
214    };
215    (@@inner $writer:expr, $sep:expr, $is_head:expr,) => {
216        $crate::iter_print!(@@line_feed $writer);
217    };
218    (@@inner $writer:expr, $sep:expr, $is_head:expr, { $($t:tt)* } $($rest:tt)*) => {
219        $crate::iter_print!(@@inner $writer, $sep, $is_head, $($t)*, !);
220        $crate::iter_print!(@@inner $writer, $sep, $is_head, $($rest)*);
221    };
222    (@@inner $writer:expr, $sep:expr, $is_head:expr, $($t:tt)*) => {
223        $crate::iter_print!(@@inner $writer, $sep, $is_head, @item $($t)*);
224    };
225    ($writer:expr, $($t:tt)*) => {{
226        $crate::iter_print!(@@inner $writer, ' ', true, $($t)*);
227    }};
228}
229
230#[cfg(test)]
231mod tests {
232    use super::*;
233    use crate::iter_print;
234
235    #[test]
236    fn test_iter_print() {
237        let mut buf = Vec::new();
238        iter_print!(
239            buf, 1, 2, @sep '.', 3, 4; 5, 6, @sp @it 7..=10;
240            @tup (1, 2, 3); @flush 4, @fmt ("{}?{}", 5, 6.7);
241            { @ns @it 8..=10; @lf @it 11..=13 },
242            @it2d (0..3).map(|i| (14..=15).map(move |j| i * 2 + j));
243            @ns @ittup (0..2).map(|i| (i * 2 + 20, i * 2 + 21));
244            @flush,
245            @bw (b'a' [0, 1, 2].iter().cloned());
246            @sp @it1 (0..2)
247        );
248        let expected = r#"1 2.3.4
2495.6 7 8 9 10
2501 2 3
2514 5?6.7
2528910
25311
25412
25513 14 15
25616 17
25718 19
2582021
2592223
260abc
2611 2
262"#;
263        assert_eq!(expected, String::from_utf8_lossy(&buf));
264    }
265}
competitive/tools/iter_print.rs

competitive/tools/
iter_print.rs
