use std::rc::Rc;

/// An iterator that allows to insert a given element `elem` every `interval` elements of the input
/// iterator.
///
/// The interleaved element is wrapped in a `Rc` smart pointer to avoid hard copies.
pub struct Intersperse<I: Iterator> {
    inner: std::iter::Peekable<I>, // required to have `peek(..)`
    elem: Rc<I::Item>,
    interval: u32,
    intersperse: u32,
}

impl<I: Iterator> Intersperse<I> {
    pub fn new(inner: I, elem: I::Item, interval: u32) -> Intersperse<I> {
        Intersperse {
            inner: inner.peekable(),
            elem: Rc::new(elem),
            interval: interval + 1, // `+ 1` needed to align semantics and implementation
            intersperse: 0,
        }
    }
}

impl<I: Iterator> Iterator for Intersperse<I> {
    type Item = Rc<I::Item>;

    fn next(&mut self) -> Option<Self::Item> {
        // Counter rotation
        self.intersperse = (self.intersperse + 1) % self.interval;
        // Insert a copy of the element only when counter goes to 0 and we haven't reached the end
        // of the inner iterator.
        if self.intersperse == 0 && self.inner.peek().is_some() {
            Some(Rc::clone(&self.elem))
        } else {
            self.inner.next()
                .map(|x| Rc::new(x)) // Go from Option<I::Item> to Option<Rc<I::Item>>
        }
    }
}

/// Extend `Iterator`s with `intersperse` function.
pub trait IntersperseIterator: Iterator {
    fn intersperse(self, elem: Self::Item, interval: u32) -> Intersperse<Self>
    where
        Self: Sized,
    {
        Intersperse::new(self, elem, interval)
    }
}

impl<I: Iterator> IntersperseIterator for I {}