Add Iterator::intersperse_with

5 files changed, 192 insertions, 25 deletions

diff --git a/library/core/src/iter/adapters/intersperse.rs b/library/core/src/iter/adapters/intersperse.rs
index 36232672549..adae47459f7 100644
--- a/library/core/src/iter/adapters/intersperse.rs
+++ b/library/core/src/iter/adapters/intersperse.rs
@@ -40,37 +40,149 @@ where
         }
     }
 
-    fn fold<B, F>(mut self, init: B, mut f: F) -> B
+    fn fold<B, F>(self, init: B, f: F) -> B
     where
         Self: Sized,
         F: FnMut(B, Self::Item) -> B,
     {
-        let mut accum = init;
+        let separator = self.separator;
+        intersperse_fold(self.iter, init, f, move || separator.clone(), self.needs_sep)
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        intersperse_size_hint(&self.iter, self.needs_sep)
+    }
+}
+
+/// An iterator adapter that places a separator between all elements.
+#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
+pub struct IntersperseWith<I, G>
+where
+    I: Iterator,
+{
+    separator: G,
+    iter: Peekable<I>,
+    needs_sep: bool,
+}
+
+// FIXME This manual implementation is needed as #[derive] misplaces trait bounds,
+// requiring <I as Iterator>::Item to be Debug on the struct-definition, which is
+// not what we want.
+#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
+impl<I, G> crate::fmt::Debug for IntersperseWith<I, G>
+where
+    I: Iterator + crate::fmt::Debug,
+    I::Item: crate::fmt::Debug,
+    G: crate::fmt::Debug,
+{
+    fn fmt(&self, f: &mut crate::fmt::Formatter<'_>) -> crate::fmt::Result {
+        f.debug_struct("IntersperseWith")
+            .field("separator", &self.separator)
+            .field("iter", &self.iter)
+            .field("needs_sep", &self.needs_sep)
+            .finish()
+    }
+}
 
-        // Use `peek()` first to avoid calling `next()` on an empty iterator.
-        if !self.needs_sep || self.iter.peek().is_some() {
-            if let Some(x) = self.iter.next() {
-                accum = f(accum, x);
-            }
+// FIXME This manual implementation is needed as #[derive] misplaces trait bounds,
+// requiring <I as Iterator>::Item to be Clone on the struct-definition, which is
+// not what we want.
+#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
+impl<I, G> crate::clone::Clone for IntersperseWith<I, G>
+where
+    I: Iterator + crate::clone::Clone,
+    I::Item: crate::clone::Clone,
+    G: Clone,
+{
+    fn clone(&self) -> Self {
+        IntersperseWith {
+            separator: self.separator.clone(),
+            iter: self.iter.clone(),
+            needs_sep: self.needs_sep.clone(),
         }
+    }
+}
 
-        let element = &self.separator;
+impl<I, G> IntersperseWith<I, G>
+where
+    I: Iterator,
+    G: FnMut() -> I::Item,
+{
+    pub(in crate::iter) fn new(iter: I, separator: G) -> Self {
+        Self { iter: iter.peekable(), separator, needs_sep: false }
+    }
+}
 
-        self.iter.fold(accum, |mut accum, x| {
-            accum = f(accum, element.clone());
-            accum = f(accum, x);
-            accum
-        })
+#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
+impl<I, G> Iterator for IntersperseWith<I, G>
+where
+    I: Iterator,
+    G: FnMut() -> I::Item,
+{
+    type Item = I::Item;
+
+    #[inline]
+    fn next(&mut self) -> Option<I::Item> {
+        if self.needs_sep && self.iter.peek().is_some() {
+            self.needs_sep = false;
+            Some((self.separator)())
+        } else {
+            self.needs_sep = true;
+            self.iter.next()
+        }
+    }
+
+    fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        intersperse_fold(self.iter, init, f, self.separator, self.needs_sep)
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        let (lo, hi) = self.iter.size_hint();
-        let next_is_elem = !self.needs_sep;
-        let lo = lo.saturating_sub(next_is_elem as usize).saturating_add(lo);
-        let hi = match hi {
-            Some(hi) => hi.saturating_sub(next_is_elem as usize).checked_add(hi),
-            None => None,
-        };
-        (lo, hi)
+        intersperse_size_hint(&self.iter, self.needs_sep)
     }
 }
+
+fn intersperse_size_hint<I>(iter: &I, needs_sep: bool) -> (usize, Option<usize>)
+where
+    I: Iterator,
+{
+    let (lo, hi) = iter.size_hint();
+    let next_is_elem = !needs_sep;
+    let lo = lo.saturating_sub(next_is_elem as usize).saturating_add(lo);
+    let hi = match hi {
+        Some(hi) => hi.saturating_sub(next_is_elem as usize).checked_add(hi),
+        None => None,
+    };
+    (lo, hi)
+}
+
+fn intersperse_fold<I, B, F, G>(
+    mut iter: Peekable<I>,
+    init: B,
+    mut f: F,
+    mut separator: G,
+    needs_sep: bool,
+) -> B
+where
+    I: Iterator,
+    F: FnMut(B, I::Item) -> B,
+    G: FnMut() -> I::Item,
+{
+    let mut accum = init;
+
+    // Use `peek()` first to avoid calling `next()` on an empty iterator.
+    if !needs_sep || iter.peek().is_some() {
+        if let Some(x) = iter.next() {
+            accum = f(accum, x);
+        }
+    }
+
+    iter.fold(accum, |mut accum, x| {
+        accum = f(accum, separator());
+        accum = f(accum, x);
+        accum
+    })
+}
diff --git a/library/core/src/iter/adapters/mod.rs b/library/core/src/iter/adapters/mod.rs
index 7dfbf32cea7..41a7b13232a 100644
--- a/library/core/src/iter/adapters/mod.rs
+++ b/library/core/src/iter/adapters/mod.rs
@@ -43,7 +43,7 @@ pub use self::flatten::Flatten;
 pub use self::copied::Copied;
 
 #[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
-pub use self::intersperse::Intersperse;
+pub use self::intersperse::{Intersperse, IntersperseWith};
 
 #[unstable(feature = "iter_map_while", reason = "recently added", issue = "68537")]
 pub use self::map_while::MapWhile;
diff --git a/library/core/src/iter/mod.rs b/library/core/src/iter/mod.rs
index 569de719d03..c57ba2bf626 100644
--- a/library/core/src/iter/mod.rs
+++ b/library/core/src/iter/mod.rs
@@ -395,8 +395,6 @@ pub use self::adapters::Cloned;
 pub use self::adapters::Copied;
 #[stable(feature = "iterator_flatten", since = "1.29.0")]
 pub use self::adapters::Flatten;
-#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
-pub use self::adapters::Intersperse;
 #[unstable(feature = "iter_map_while", reason = "recently added", issue = "68537")]
 pub use self::adapters::MapWhile;
 #[unstable(feature = "inplace_iteration", issue = "none")]
@@ -410,6 +408,8 @@ pub use self::adapters::{
     Chain, Cycle, Enumerate, Filter, FilterMap, FlatMap, Fuse, Inspect, Map, Peekable, Rev, Scan,
     Skip, SkipWhile, Take, TakeWhile, Zip,
 };
+#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
+pub use self::adapters::{Intersperse, IntersperseWith};
 
 pub(crate) use self::adapters::process_results;
 
diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs
index 633175702d8..91d7a47907a 100644
--- a/library/core/src/iter/traits/iterator.rs
+++ b/library/core/src/iter/traits/iterator.rs
@@ -8,7 +8,7 @@ use crate::ops::{Add, ControlFlow, Try};
 use super::super::TrustedRandomAccess;
 use super::super::{Chain, Cloned, Copied, Cycle, Enumerate, Filter, FilterMap, Fuse};
 use super::super::{FlatMap, Flatten};
-use super::super::{FromIterator, Intersperse, Product, Sum, Zip};
+use super::super::{FromIterator, Intersperse, IntersperseWith, Product, Sum, Zip};
 use super::super::{
     Inspect, Map, MapWhile, Peekable, Rev, Scan, Skip, SkipWhile, StepBy, Take, TakeWhile,
 };
@@ -591,6 +591,31 @@ pub trait Iterator {
         Intersperse::new(self, separator)
     }
 
+    /// Places an element generated by `separator` between all elements.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// #![feature(iter_intersperse)]
+    ///
+    /// let src = ["Hello", "to", "all", "people"].iter().copied();
+    /// let mut separator = [" ❤️ ", " 😀 "].iter().copied().cycle();
+    ///
+    /// let result = src.intersperse_with(|| separator.next().unwrap()).collect::<String>();
+    /// assert_eq!(result, "Hello ❤️ to 😀 all ❤️ people");
+    /// ```
+    #[inline]
+    #[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
+    fn intersperse_with<G>(self, separator: G) -> IntersperseWith<Self, G>
+    where
+        Self: Sized,
+        G: FnMut() -> Self::Item,
+    {
+        IntersperseWith::new(self, separator)
+    }
+
     /// Takes a closure and creates an iterator which calls that closure on each
     /// element.
     ///
diff --git a/library/core/tests/iter.rs b/library/core/tests/iter.rs
index 7376e7848ef..691767edea6 100644
--- a/library/core/tests/iter.rs
+++ b/library/core/tests/iter.rs
@@ -3508,6 +3508,12 @@ pub fn extend_for_unit() {
 
 #[test]
 fn test_intersperse() {
+    let v = std::iter::empty().intersperse(0u32).collect::<Vec<_>>();
+    assert_eq!(v, vec![]);
+
+    let v = std::iter::once(1).intersperse(0).collect::<Vec<_>>();
+    assert_eq!(v, vec![1]);
+
     let xs = ["a", "", "b", "c"];
     let v: Vec<&str> = xs.iter().map(|x| x.clone()).intersperse(", ").collect();
     let text: String = v.concat();
@@ -3520,6 +3526,9 @@ fn test_intersperse() {
 
 #[test]
 fn test_intersperse_size_hint() {
+    let iter = std::iter::empty::<i32>().intersperse(0);
+    assert_eq!(iter.size_hint(), (0, Some(0)));
+
     let xs = ["a", "", "b", "c"];
     let mut iter = xs.iter().map(|x| x.clone()).intersperse(", ");
     assert_eq!(iter.size_hint(), (7, Some(7)));
@@ -3587,3 +3596,24 @@ fn test_try_fold_specialization_intersperse_err() {
     iter.try_for_each(|item| if item == "b" { None } else { Some(()) });
     assert_eq!(iter.next(), None);
 }
+
+#[test]
+fn test_intersperse_with() {
+    #[derive(PartialEq, Debug)]
+    struct NotClone {
+        u: u32,
+    }
+    let r = vec![NotClone { u: 0 }, NotClone { u: 1 }]
+        .into_iter()
+        .intersperse_with(|| NotClone { u: 2 })
+        .collect::<Vec<_>>();
+    assert_eq!(r, vec![NotClone { u: 0 }, NotClone { u: 2 }, NotClone { u: 1 }]);
+
+    let mut ctr = 100;
+    let separator = || {
+        ctr *= 2;
+        ctr
+    };
+    let r = (0..3).intersperse_with(separator).collect::<Vec<_>>();
+    assert_eq!(r, vec![0, 200, 1, 400, 2]);
+}