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/// - `iter_print!(fast; writer, args...)`: use `FastOutput` and `FastPrint`
67/// - `@sep $expr`: set separator (default: `' '`)
68/// - `@ns`: alias for `@sep ""`
69/// - `@lf`: alias for `@sep '\n'`
70/// - `@sp`: alias for `@sep ' '`
71/// - `@fmt ($lit, $($expr),*)`: print `format!($lit, $($expr),*)`
72/// - `@flush`: flush writer (auto insert `!`)
73/// - `@it $expr`: print iterator
74/// - `@it1 $expr`: print iterator as 1-indexed
75/// - `@cw ($char $expr)`: print iterator as `(elem as u8 + $char as u8) as char`
76/// - `@bw ($byte $expr)`: print iterator as `(elem as u8 + $byte) as char`
77/// - `@it2d $expr`: print 2d-iterator
78/// - `@tup $expr`: print tuple (need to import [`IterPrint`])
79/// - `@ittup $expr`: print iterative tuple (need to import [`IterPrint`])
80/// - `$expr`: print expr
81/// - `{ args... }`: scoped
82/// - `;`: print `'\n'`
83/// - `!`: not print `'\n'` at the end
84#[macro_export]
85macro_rules! iter_print {
86    (@@fmt normal $writer:expr, $sep:expr, $is_head:expr, ($lit:literal $(, $e:expr)* $(,)?)) => {
87        if !$is_head {
88            ::std::write!($writer, "{}", $sep).expect("io error");
89        }
90        ::std::write!($writer, $lit, $($e),*).expect("io error");
91    };
92    (@@fmt fast $writer:expr, $sep:expr, $is_head:expr, ($lit:literal $(, $e:expr)* $(,)?)) => {{
93        use ::std::fmt::Write as _;
94        if !$is_head {
95            FastPrint::fast_print(&$sep, &mut $writer);
96        }
97        ::std::write!($writer, $lit, $($e),*).expect("io error");
98    }};
99    (@@item normal $writer:expr, $sep:expr, $is_head:expr, $e:expr) => {
100        $crate::iter_print!(@@fmt normal $writer, $sep, $is_head, ("{}", $e));
101    };
102    (@@item fast $writer:expr, $sep:expr, $is_head:expr, $e:expr) => {
103        if !$is_head {
104            FastPrint::fast_print(&$sep, &mut $writer);
105        }
106        FastPrint::fast_print(&$e, &mut $writer);
107    };
108    (@@line_feed normal $writer:expr $(,)?) => {
109        ::std::writeln!($writer).expect("io error");
110    };
111    (@@line_feed fast $writer:expr $(,)?) => {
112        $writer.byte(b'\n');
113    };
114    (@@flush normal $writer:expr) => {{
115        use ::std::io::Write as _;
116        $writer.flush().expect("io error");
117    }};
118    (@@flush fast $writer:expr) => {
119        $writer.flush();
120    };
121    (@@it fast $writer:expr, $sep:expr, $is_head:expr, $iter:expr) => {{
122        // Keep borrowed temporary sources alive through the loop.
123        match $iter.into_iter() {
124            mut iter => {
125                if let ::std::option::Option::Some(first) = iter.next() {
126                    {
127                        // Drop the first item before advancing the iterator.
128                        let item = first;
129                        $crate::iter_print!(@@item fast $writer, $sep, $is_head, item);
130                    }
131                    for item in iter {
132                        $crate::iter_print!(@@item fast $writer, $sep, false, item);
133                    }
134                }
135            }
136        }
137    }};
138    (@@it $mode:ident $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!(@@item $mode $writer, $sep, $is_head, item);
142        }
143        for item in iter {
144            $crate::iter_print!(@@item $mode $writer, $sep, false, item);
145        }
146    }};
147    (@@it1 $mode:ident $writer:expr, $sep:expr, $is_head:expr, $iter:expr) => {{
148        let mut iter = $iter.into_iter();
149        if let Some(item) = iter.next() {
150            $crate::iter_print!(@@item $mode $writer, $sep, $is_head, item + 1);
151        }
152        for item in iter {
153            $crate::iter_print!(@@item $mode $writer, $sep, false, item + 1);
154        }
155    }};
156    (@@cw $mode:ident $writer:expr, $sep:expr, $is_head:expr, ($ch:literal $iter:expr)) => {{
157        let mut iter = $iter.into_iter();
158        let b = $ch as u8;
159        if let Some(item) = iter.next() {
160            $crate::iter_print!(@@item $mode $writer, $sep, $is_head, (item as u8 + b) as char);
161        }
162        for item in iter {
163            $crate::iter_print!(@@item $mode $writer, $sep, false, (item as u8 + b) as char);
164        }
165    }};
166    (@@bw $mode:ident $writer:expr, $sep:expr, $is_head:expr, ($b:literal $iter:expr)) => {{
167        let mut iter = $iter.into_iter();
168        let b: u8 = $b;
169        if let Some(item) = iter.next() {
170            $crate::iter_print!(@@item $mode $writer, $sep, $is_head, (item as u8 + b) as char);
171        }
172        for item in iter {
173            $crate::iter_print!(@@item $mode $writer, $sep, false, (item as u8 + b) as char);
174        }
175    }};
176    (@@it2d fast $writer:expr, $sep:expr, $is_head:expr, $iter:expr) => {
177        for (i, item) in $iter.into_iter().enumerate() {
178            if i > 0 {
179                $crate::iter_print!(@@line_feed fast $writer);
180            }
181            $crate::iter_print!(@@it fast $writer, $sep, i > 0 || $is_head, item);
182        }
183    };
184    (@@it2d $mode:ident $writer:expr, $sep:expr, $is_head:expr, $iter:expr) => {
185        let mut iter = $iter.into_iter();
186        if let Some(item) = iter.next() {
187            $crate::iter_print!(@@it $mode $writer, $sep, $is_head, item);
188        }
189        for item in iter {
190            $crate::iter_print!(@@line_feed $mode $writer);
191            $crate::iter_print!(@@it $mode $writer, $sep, true, item);
192        }
193    };
194    (@@tup normal $writer:expr, $sep:expr, $is_head:expr, $tuple:expr) => {
195        IterPrint::iter_print($tuple, &mut $writer, $sep, $is_head).expect("io error");
196    };
197    (@@tup fast $writer:expr, $sep:expr, $is_head:expr, $tuple:expr) => {
198        FastIterPrint::fast_iter_print($tuple, &mut $writer, $sep, $is_head);
199    };
200    (@@ittup $mode:ident $writer:expr, $sep:expr, $is_head:expr, $iter:expr) => {
201        let mut iter = $iter.into_iter();
202        if let Some(item) = iter.next() {
203            $crate::iter_print!(@@tup $mode $writer, $sep, $is_head, item);
204        }
205        for item in iter {
206            $crate::iter_print!(@@line_feed $mode $writer);
207            $crate::iter_print!(@@tup $mode $writer, $sep, true, item);
208        }
209    };
210    (@@assert_tag item) => {};
211    (@@assert_tag it) => {};
212    (@@assert_tag it1) => {};
213    (@@assert_tag it2d) => {};
214    (@@assert_tag tup) => {};
215    (@@assert_tag ittup) => {};
216    (@@assert_tag $tag:ident) => {
217        ::std::compile_error!(::std::concat!("invalid tag in `iter_print!`: `", std::stringify!($tag), "`"));
218    };
219    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, @sep $e:expr, $($t:tt)*) => {
220        $crate::iter_print!(@@inner $mode $writer, $e, $is_head, $($t)*);
221    };
222    (@@inner fast $writer:expr, $sep:expr, $is_head:expr, @ns $($t:tt)*) => {
223        $crate::iter_print!(@@inner fast $writer, (), $is_head, $($t)*);
224    };
225    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, @ns $($t:tt)*) => {
226        $crate::iter_print!(@@inner $mode $writer, "", $is_head, $($t)*);
227    };
228    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, @lf $($t:tt)*) => {
229        $crate::iter_print!(@@inner $mode $writer, '\n', $is_head, $($t)*);
230    };
231    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, @sp $($t:tt)*) => {
232        $crate::iter_print!(@@inner $mode $writer, ' ', $is_head, $($t)*);
233    };
234    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, @flush $($t:tt)*) => {
235        $crate::iter_print!(@@flush $mode $writer);
236        $crate::iter_print!(@@inner $mode $writer, $sep, $is_head, ! $($t)*);
237    };
238    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, @fmt $arg:tt $($t:tt)*) => {
239        $crate::iter_print!(@@fmt $mode $writer, $sep, $is_head, $arg);
240        $crate::iter_print!(@@inner $mode $writer, $sep, $is_head, $($t)*);
241    };
242    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, @cw $arg:tt $($t:tt)*) => {
243        $crate::iter_print!(@@cw $mode $writer, $sep, $is_head, $arg);
244        $crate::iter_print!(@@inner $mode $writer, $sep, $is_head, $($t)*);
245    };
246    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, @bw $arg:tt $($t:tt)*) => {
247        $crate::iter_print!(@@bw $mode $writer, $sep, $is_head, $arg);
248        $crate::iter_print!(@@inner $mode $writer, $sep, $is_head, $($t)*);
249    };
250    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, @$tag:ident $e:expr, $($t:tt)*) => {
251        $crate::iter_print!(@@assert_tag $tag);
252        $crate::iter_print!(@@$tag $mode $writer, $sep, $is_head, $e);
253        $crate::iter_print!(@@inner $mode $writer, $sep, false, $($t)*);
254    };
255    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, @$tag:ident $e:expr; $($t:tt)*) => {
256        $crate::iter_print!(@@assert_tag $tag);
257        $crate::iter_print!(@@$tag $mode $writer, $sep, $is_head, $e);
258        $crate::iter_print!(@@line_feed $mode $writer);
259        $crate::iter_print!(@@inner $mode $writer, $sep, true, $($t)*);
260    };
261    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, @$tag:ident $e:expr) => {
262        $crate::iter_print!(@@assert_tag $tag);
263        $crate::iter_print!(@@$tag $mode $writer, $sep, $is_head, $e);
264        $crate::iter_print!(@@inner $mode $writer, $sep, false,);
265    };
266    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, @$tag:ident $($t:tt)*) => {
267        ::std::compile_error!(::std::concat!("invalid expr in `iter_print!`: `", std::stringify!($($t)*), "`"));
268    };
269    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, , $($t:tt)*) => {
270        $crate::iter_print!(@@inner $mode $writer, $sep, $is_head, $($t)*);
271    };
272    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, ; $($t:tt)*) => {
273        $crate::iter_print!(@@line_feed $mode $writer);
274        $crate::iter_print!(@@inner $mode $writer, $sep, $is_head, $($t)*);
275    };
276    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, ! $(,)?) => {};
277    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, ! $($t:tt)*) => {
278        $crate::iter_print!(@@inner $mode $writer, $sep, $is_head, $($t)*);
279    };
280    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr,) => {
281        $crate::iter_print!(@@line_feed $mode $writer);
282    };
283    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, { $($t:tt)* } $($rest:tt)*) => {
284        $crate::iter_print!(@@inner $mode $writer, $sep, $is_head, $($t)*, !);
285        $crate::iter_print!(@@inner $mode $writer, $sep, $is_head, $($rest)*);
286    };
287    (@@inner $mode:ident $writer:expr, $sep:expr, $is_head:expr, $($t:tt)*) => {
288        $crate::iter_print!(@@inner $mode $writer, $sep, $is_head, @item $($t)*);
289    };
290    (fast; $writer:expr, $($t:tt)*) => {{
291        $crate::iter_print!(@@inner fast $writer, ' ', true, $($t)*);
292    }};
293    ($writer:expr, $($t:tt)*) => {{
294        $crate::iter_print!(@@inner normal $writer, ' ', true, $($t)*);
295    }};
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301    use crate::tools::Xorshift;
302    use crate::tools::{FastIterPrint, FastOutput, FastPrint};
303    use std::array;
304    use std::iter;
305
306    #[test]
307    fn test_iter_print() {
308        use std::fmt::Write as _;
309        let mut rng = Xorshift::default();
310        for _ in 0..1000 {
311            let a: [i32; 23] = array::from_fn(|_| rng.random(-1000..=1000));
312            let decimal = rng.random(-1000..=1000) as f64 / 10.0;
313            let n = rng.random(0..=20);
314            let letters: Vec<u8> = rng.random_iter(0..26).take(n).collect();
315            let indices: Vec<usize> = rng.random_iter(0..100).take(n).collect();
316            let mut buf = Vec::new();
317            let mut fast_buf = Vec::new();
318            macro_rules! check {
319                ($writer:ident $(, $mode:ident)?) => {
320                    iter_print!(
321                        $($mode;)? $writer, a[0], a[1], @sep '.', a[2], a[3];
322                        a[4], a[5], @sp @it &a[6..10];
323                        @tup (a[0], a[1], a[2]); @flush a[3], @fmt ("{}?{}", a[4], decimal);
324                        { @ns @it &a[7..10]; @lf @it &a[10..13] },
325                        @it2d a[13..19].chunks(2);
326                        @ns @ittup a[19..23].chunks(2).map(|row| (row[0], row[1]));
327                        @flush,
328                        @bw (b'a' letters.iter().copied());
329                        @sp @it1 indices.iter().copied()
330                    );
331                };
332            }
333            check!(buf);
334            {
335                let mut out = FastOutput::new(&mut fast_buf);
336                check!(out, fast);
337            }
338            let text: Vec<_> = a.iter().map(ToString::to_string).collect();
339            let mut expected = String::new();
340            writeln!(expected, "{} {}.{}.{}", a[0], a[1], a[2], a[3]).unwrap();
341            writeln!(expected, "{}.{} {}", a[4], a[5], text[6..10].join(" ")).unwrap();
342            writeln!(expected, "{}", text[..3].join(" ")).unwrap();
343            writeln!(expected, "{} {}?{}", a[3], a[4], decimal).unwrap();
344            writeln!(expected, "{}", text[7..10].concat()).unwrap();
345            writeln!(
346                expected,
347                "{} {}",
348                text[10..13].join("\n"),
349                text[13..15].join(" ")
350            )
351            .unwrap();
352            for row in text[15..19].chunks(2) {
353                writeln!(expected, "{}", row.join(" ")).unwrap();
354            }
355            for row in text[19..23].chunks(2) {
356                writeln!(expected, "{}", row.concat()).unwrap();
357            }
358            writeln!(
359                expected,
360                "{}",
361                letters
362                    .iter()
363                    .map(|&x| char::from(b'a' + x))
364                    .collect::<String>()
365            )
366            .unwrap();
367            writeln!(
368                expected,
369                "{}",
370                indices
371                    .iter()
372                    .map(|x| (x + 1).to_string())
373                    .collect::<Vec<_>>()
374                    .join(" ")
375            )
376            .unwrap();
377            assert_eq!(buf, expected.as_bytes());
378            assert_eq!(fast_buf, expected.as_bytes());
379        }
380    }
381
382    #[test]
383    fn test_iter_print_iterators() {
384        for mask in 0..256 {
385            for n in 0..=4 {
386                let rows: Vec<Vec<_>> = (0..n)
387                    .map(|i| (0..(mask >> (i * 2)) & 3).map(|j| i * 10 + j).collect())
388                    .collect();
389                let expected = rows
390                    .iter()
391                    .map(|row| {
392                        row.iter()
393                            .map(ToString::to_string)
394                            .collect::<Vec<_>>()
395                            .join(":")
396                    })
397                    .collect::<Vec<_>>()
398                    .join("\n");
399                let mut normal = Vec::new();
400                iter_print!(normal, @sep ':', @it2d &rows, !);
401                let mut fast = Vec::new();
402                let mut calls = 0;
403                let mut separators = 0;
404                {
405                    let mut writer = FastOutput::new(&mut fast);
406                    iter_print!(fast; writer, @sep { separators += 1; ':' }, @it2d { calls += 1; &rows }, !);
407                }
408                assert_eq!(calls, 1);
409                assert_eq!(
410                    separators,
411                    rows.iter()
412                        .map(|row| row.len().saturating_sub(1))
413                        .sum::<usize>()
414                );
415                assert_eq!(normal, expected.as_bytes());
416                assert_eq!(fast, normal);
417            }
418        }
419    }
420
421    #[test]
422    fn test_fast_iter_print_order() {
423        use std::cell::Cell;
424
425        struct Item<'a> {
426            value: usize,
427            dropped: &'a Cell<usize>,
428        }
429        impl FastPrint for Item<'_> {
430            fn fast_print<W: std::io::Write>(&self, writer: &mut FastOutput<W>) {
431                writer.usize(self.value);
432            }
433        }
434        impl Drop for Item<'_> {
435            fn drop(&mut self) {
436                self.dropped.set(self.dropped.get() + 1);
437            }
438        }
439
440        for n in 0..=32 {
441            let dropped = Cell::new(0);
442            let mut next = 0;
443            let iter = iter::from_fn(|| {
444                let value = next;
445                assert_eq!(dropped.get(), value);
446                next += 1;
447                (value != n).then(|| Item {
448                    value,
449                    dropped: &dropped,
450                })
451            });
452            let mut output = Vec::new();
453            {
454                let mut writer = FastOutput::new(&mut output);
455                iter_print!(fast; writer, @it iter);
456            }
457            assert_eq!(next, n + 1);
458            let expected = format!(
459                "{}\n",
460                (0..n).map(|i| i.to_string()).collect::<Vec<_>>().join(" ")
461            );
462            assert_eq!(output, expected.as_bytes());
463            output.clear();
464            {
465                let mut writer = FastOutput::new(&mut output);
466                iter_print!(fast; writer, @it (0..n).collect::<Vec<_>>().iter());
467            }
468            assert_eq!(output, expected.as_bytes());
469        }
470    }
471}
competitive/tools/iter_print.rs

competitive/tools/
iter_print.rs
