iterext/
lib.rs

1use itertools::Itertools;
2use std::iter::FromIterator;
3
4pub type Padded<'a, T> = Box<dyn 'a + Iterator<Item = T>>;
5pub trait Pad<'a, T>
6where
7    T: Copy,
8{
9    /// Ensure that a stream has a length of a multiple of `group_size`.
10    ///
11    /// If `iter` ends at a length which is not a multiple of `group_size`,
12    /// instances of `padding` are copied into the stream until the length
13    /// is correct.
14    ///
15    /// This is a fused iterator.
16    fn pad(self, padding: T, group_size: usize) -> Padded<'a, T>;
17}
18
19impl<'a, I, T> Pad<'a, T> for I
20where
21    I: IntoIterator<Item = T>,
22    <I as IntoIterator>::IntoIter: 'a,
23    T: 'a + Copy,
24{
25    fn pad(self, padding: T, group_size: usize) -> Padded<'a, T> {
26        use itertools::EitherOrBoth::*;
27        Box::new(
28            self.into_iter()
29                .fuse()
30                .zip_longest(std::iter::repeat(padding))
31                .enumerate()
32                .take_while(move |(idx, eob)| match eob {
33                    Left(_) => unreachable!(),
34                    Both(_, _) => true,
35                    Right(_) => idx % group_size != 0,
36                })
37                .map(|(_, eob)| match eob {
38                    Left(_) => unreachable!(),
39                    Both(b, _) => b,
40                    Right(b) => b,
41                }),
42        )
43    }
44}
45
46pub trait Separate<'a, I, T, O>
47where
48    I: IntoIterator<Item = T>,
49    T: Copy,
50    O: FromIterator<T>,
51{
52    /// Separate a stream into groups, inserting a copy of T between each.
53    /// Then collect it into an appropriate container.
54    ///
55    /// This is a fused iterator.
56    fn separate(self, group_sep: T, group_size: usize) -> O;
57}
58
59impl<'a, I, T, O> Separate<'a, I, T, O> for I
60where
61    I: 'a + IntoIterator<Item = T>,
62    <I as IntoIterator>::IntoIter: 'a,
63    T: 'a + Copy + PartialEq,
64    O: FromIterator<T>,
65{
66    fn separate(self, group_sep: T, group_size: usize) -> O {
67        self.into_iter()
68            .fuse()
69            .chunks(group_size)
70            .into_iter()
71            .map(|chunk| {
72                let d: Box<dyn Iterator<Item = T>> = Box::new(chunk);
73                d
74            })
75            .interleave_shortest(std::iter::repeat(std::iter::once(group_sep)).map(|cyc| {
76                let d: Box<dyn Iterator<Item = T>> = Box::new(cyc);
77                d
78            }))
79            .flatten()
80            .with_position()
81            .filter_map(move |pos| {
82                use itertools::Position::*;
83                match pos {
84                    Only(c) => Some(c),
85                    First(c) => Some(c),
86                    Middle(c) => Some(c),
87                    Last(c) if c != group_sep => Some(c),
88                    _ => None,
89                }
90            })
91            .collect()
92    }
93}
94
95pub mod prelude {
96    pub use super::Pad;
97    pub use super::Separate;
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    const GROUP_SIZE: usize = 5;
105    const PAD_CHAR: u8 = b'X' - b'A' + 1;
106
107    /// Convert a text input into a numeric stream from 1..26 according to its chars.
108    ///
109    /// ASCII letters are uppercased, then assigned `A==1 .. Z==26`. All other chars
110    /// are discarded.
111    pub fn textbyte(text: &str) -> impl '_ + Iterator<Item = u8> {
112        text.chars()
113            .filter(char::is_ascii_alphabetic)
114            .map(|c| (c.to_ascii_uppercase() as u8) - b'A' + 1)
115    }
116
117    fn padding_impl(msg: &str, expect_len: usize) {
118        assert_eq!(
119            textbyte(msg)
120                .pad(PAD_CHAR, GROUP_SIZE)
121                .collect::<Vec<_>>()
122                .len(),
123            expect_len
124        );
125    }
126
127    #[test]
128    fn test_padding() {
129        padding_impl("a", 5);
130        padding_impl("abcde", 5);
131        padding_impl(".", 0);
132        padding_impl("abcdef", 10);
133        padding_impl("a.b.c.d", 5);
134        padding_impl("", 0);
135    }
136
137    fn padding_impl_2(msg: &str, expect: &[u8]) {
138        let have: Vec<u8> = textbyte(msg).pad(PAD_CHAR, GROUP_SIZE).collect();
139        assert_eq!(have, expect);
140    }
141
142    #[test]
143    fn test_padding_2() {
144        padding_impl_2("a", &[1, 24, 24, 24, 24]);
145        padding_impl_2("abcde", &[1, 2, 3, 4, 5]);
146        padding_impl_2(".", &[]);
147        padding_impl_2("abcdef", &[1, 2, 3, 4, 5, 6, 24, 24, 24, 24]);
148        padding_impl_2("a.b.c.d", &[1, 2, 3, 4, 24]);
149        padding_impl_2("", &[]);
150    }
151
152    #[test]
153    fn test_padding_chars() {
154        let have = "foo".chars().pad('X', 5).collect::<String>();
155        assert_eq!(have, "fooXX");
156    }
157
158    #[test]
159    fn test_separate() {
160        for (msg, expect) in &[
161            ("abc", "abc"),
162            ("zyx", "zyx"),
163            (
164                "abcdefghijklmnopqrstuvwxyz",
165                "abcde fghij klmno pqrst uvwxy z",
166            ),
167            (
168                "thequickbrownfoxjumpedoverthelazydog",
169                "thequ ickbr ownfo xjump edove rthel azydo g",
170            ),
171            (
172                "abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz",
173                "abcde fghij klmno pqrst uvwxy zabcd efghi jklmn opqrs tuvwx yz",
174            ),
175        ] {
176            let got: String = msg.chars().separate(' ', 5);
177            assert_eq!(&got, expect,);
178        }
179    }
180}