[rust] Beer Song

1 files changed, 52 insertions, 0 deletions

diff --git a/beer-song/src/intersperse.rs b/beer-song/src/intersperse.rs
new file mode 100644
index 0000000..51580e6
--- /dev/null
+++ b/beer-song/src/intersperse.rs
@@ -0,0 +1,52 @@
+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 {}