How can I intersperse a rust iterator with a value every n items?

Question

I have an iterator of characters, and I want to add a newline every N characters:

let iter = "abcdefghijklmnopqrstuvwxyz".chars();
let iter_with_newlines = todo!();
let string: String = iter_with_newlines.collect();
assert_eq("abcdefghij\nklmnopqrst\nuvwxyz", string);

So basically, I want to intersperse the iterator with a newline every n characters. How can I do this?

Some Ideas I had

It would be great if I could do something like this, where chunks would be a method to make Iterator<T> into Iterator<Iterator<T>: iter.chunks(10).intersperse('\n').flatten()

It would also be cool if I could do something like this: iter.chunks.intersperseEvery(10, '\n'), where intersperseEvery is a method that would only intersperse the value every n items.

Answer 1

You can do it without temporary allocation using enumerate and flat_map:

use either::Either;

fn main() {
    let iter = "abcdefghijklmnopqrstuvwxyz".chars();
    let iter_with_newlines = iter
        .enumerate()
        .flat_map(|(i, c)| {
            if i % 10 == 0 {
                Either::Left(['\n', c].into_iter())
            } else {
                Either::Right(std::iter::once(c))
            }
        })
        .skip(1); // The above code add a newline in first position -> skip it
    let string: String = iter_with_newlines.collect();
    assert_eq!("abcdefghij\nklmnopqrst\nuvwxyz", string);
}

Playground

Answer 2

Here's what I ended up doing:

// src/intersperse_sparse.rs

use core::iter::Peekable;

/// An iterator adaptor to insert a particular value
/// every n elements of the adapted iterator.
///
/// Iterator element type is `I::Item`
pub struct IntersperseSparse<I>
where
    I: Iterator,
    I::Item: Clone,
{
    iter: Peekable<I>,
    step_length: usize,
    index: usize,
    separator: I::Item,
}

impl<I> IntersperseSparse<I>
where
    I: Iterator,
    I::Item: Clone,
{
    #[allow(unused)] // Although this function isn't explicitly exported, it is called in the default implementation of the IntersperseSparseAdapter, which is exported.
    fn new(iter: I, step_length: usize, separator: I::Item) -> Self {
        if step_length == 0 {
            panic!("Chunk size cannot be 0!")
        }
        Self {
            iter: iter.peekable(),
            step_length,
            separator,
            index: 0,
        }
    }
}

impl<I> Iterator for IntersperseSparse<I>
where
    I: Iterator,
    I::Item: Clone,
{
    type Item = I::Item;
    fn next(&mut self) -> Option<Self::Item> {
        if self.index == self.step_length && self.iter.peek().is_some() {
            self.index = 0;
            Some(self.separator.clone())
        } else {
            self.index += 1;
            self.iter.next()
        }
    }
}

/// An iterator adaptor to insert a particular value created by a function
/// every n elements of the adapted iterator.
///
/// Iterator element type is `I::Item`
pub struct IntersperseSparseWith<I, G>
where
    I: Iterator,
    G: FnMut() -> I::Item,
{
    iter: Peekable<I>,
    step_length: usize,
    index: usize,
    separator_closure: G,
}

impl<I, G> IntersperseSparseWith<I, G>
where
    I: Iterator,
    G: FnMut() -> I::Item,
{
    #[allow(unused)] // Although this function isn't explicitly exported, it is called in the default implementation of the IntersperseSparseAdapter, which is exported.
    fn new(iter: I, step_length: usize, separator_closure: G) -> Self {
        if step_length == 0 {
            panic!("Chunk size cannot be 0!")
        }
        Self {
            iter: iter.peekable(),
            step_length,
            separator_closure,
            index: 0,
        }
    }
}

impl<I, G> Iterator for IntersperseSparseWith<I, G>
where
    I: Iterator,
    G: FnMut() -> I::Item,
{
    type Item = I::Item;
    fn next(&mut self) -> Option<Self::Item> {
        if self.index == self.step_length && self.iter.peek().is_some() {
            self.index = 0;
            Some((self.separator_closure)())
        } else {
            self.index += 1;
            self.iter.next()
        }
    }
}

/// Import this trait to use the `iter.intersperse_sparse(n, item)` and `iter.intersperse_sparse(n, ||item)` on all iterators.
pub trait IntersperseSparseAdapter: Iterator {
    fn intersperse_sparse(self, chunk_size: usize, separator: Self::Item) -> IntersperseSparse<Self>
    where
        Self: Sized,
        Self::Item: Clone,
    {
        IntersperseSparse::new(self, chunk_size, separator)
    }

    fn intersperse_sparse_with<G>(
        self,
        chunk_size: usize,
        separator_closure: G,
    ) -> IntersperseSparseWith<Self, G>
    where
        Self: Sized,
        G: FnMut() -> Self::Item,
    {
        IntersperseSparseWith::new(self, chunk_size, separator_closure)
    }
}

impl<I> IntersperseSparseAdapter for I where I: Iterator {}

To use it:

// src/main.rs

mod intersperse_sparse;
use intersperse_sparse::IntersperseSparseAdapter;

fn main() {
    let string = "abcdefg";
    let new_string: String = string.chars().intersperse_sparse(3, '\n').collect();
    assert_eq!(new_string, "abc\ndef\ng");
}

Answer 3

If you don't particularly care about performance, you can use chunks from itertools, collect the chunks into Vecs, and then intersperse your element as a single-element Vec, just to flatten the whole thing finally.

use itertools::Itertools;
iter
    .chunks(3)
    .into_iter()
    .map(|chunk| chunk.collect::<Vec<_>>())
    .intersperse(vec![','])
    .flat_map(|chunk| chunk.into_iter())
    .collect::<String>();

Playground

Other than that, consider writing your own iterator extension trait, just like itertools is one?

Answer 4

Build an Iterator with from_fn:

let mut iter = "abcdefghijklmnopqrstuvwxyz".chars().peekable();
let mut count = 0;
let iter_with_newlines = std::iter::from_fn(move || match iter.peek() {
    Some(_) => {
        if count < 10 {
            count += 1;
            iter.next()
        } else {
            count = 0;
            Some('\n')
        }
    }
    None => None,
});
assert_eq!(
    "abcdefghij\nklmnopqrst\nuvwxyz",
    iter_with_newlines.collect::<String>()
);

Playground