CODEWITHSUNDEEP
rustiterator
Simplex
Simplex

Reputation: 1000

Iterator that repeats the first and last element of another Iterator?

Is there an Iterator, either in std or a maintained crate, that repeats the first and last element of another Iterator given Item is Clone?

Example:

let iter = [1, 2, 3].into_iter();

assert!(iter.repeat_first_and_last().eq([1, 1, 2, 3, 3]));

Upvotes: 0

Views: 537

Answers (3)

Simplex
Simplex

Reputation: 1000

I ended up rolling my own.

use std::{mem::take, iter::Fuse};

use tap::Tap;

struct Padded<T: Iterator> {
    inner: Fuse<T>,
    repeat_first: bool,
    repeat_last: bool,
    next: Option<T::Item>,
}

impl<T: Iterator> Padded<T>
where
    T::Item: Clone,
{
    fn new(inner: impl IntoIterator<IntoIter = T>) -> Self
    where
        T::Item: Clone,
    {
        let mut inner = inner.into_iter().fuse();
        let next = inner.next();
        
        Self { inner, next, repeat_first: true, repeat_last: true, }
    }
}

impl<T: Iterator> Iterator for Padded<T>
where
    T::Item: Clone,
{
    type Item = T::Item;

    fn next(&mut self) -> Option<Self::Item> {
        self.next.take().tap(|next| {
            if take(&mut self.repeat_first) {
                self.next = next.clone();
            } else {
                self.next = self.inner.next().or_else(|| {
                    next.as_ref().filter(|_| take(&mut self.repeat_last)).cloned()
                })
            }
        })
    }
}

fn main() {
    Padded::new([1, 2, 3]).for_each(|n| print!("{n}"));
}

Upvotes: 0

isaactfa
isaactfa

Reputation: 6651

You can also roll your own if you want:

use std::iter::Fuse;

pub struct FirstLastRepeat<I, T> {
    iter: Fuse<I>,
    head: Option<T>,
    tail: Option<T>,
    stash: Option<T>,
}

impl<I, T> FirstLastRepeat<I, T>
where
    I: Iterator<Item = T>,
{
    pub fn new(iter: I) -> Self {
        let mut iter = iter.fuse();
        let head = iter.next();
        Self {
            iter,
            head,
            tail: None,
            stash: None,
        }
    }
}

impl<I, T> Iterator for FirstLastRepeat<I, T>
where
    I: Iterator<Item = T>,
    T: Clone,
{
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        if let v @ Some(_) = self.head.take() {
            self.stash = v.clone();
            return v;
        }
        if let v @ Some(_) = self.tail.take() {
            return v;
        }
        let next = self.stash.take();
        match self.iter.next() {
            v @ Some(_) => self.stash = v,
            None => self.tail = next.clone(),
        }
        next
    }
}

#[test]
fn basic() {
    let iter = [1, 2, 3].into_iter();
    let flr = FirstLastRepeat::new(iter);

    assert!(flr.eq([1, 1, 2, 3, 3]));
}

The only weird edge case here is an original iterator with only one element. It's not clear what the right behaviour here is, but FristLastRepeat will yield it three times which doesn't seem correct in any case. That could be fixed though.

Upvotes: 1

Jonas Fassbender
Jonas Fassbender

Reputation: 2632

There's no method I know of in either std or itertools for that. You can combine multiple iterator methods to accomplish this, though:

use std::iter::repeat;

fn main() {
    let iter = [1, 2, 3, 4].into_iter();
    let len = iter.len();

    let res: Vec<u8> = iter
        .enumerate()
        .map(|(i, x)| repeat(x).take((i == 0 || i == len - 1).then(|| 2).unwrap_or(1)))
        .flatten()
        .collect();

    assert_eq!(res, vec![1, 1, 2, 3, 4, 4]);
}

Playground.

Note that repeating the last element twice requires us to know the size of the iterator beforehand, so the iterator type must implement ExactSizeIterator, or you have to provide some other means to get the length.

Upvotes: 0

Related Questions