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

competitive/tools/
iter_print.rs
