Rollup merge of #80567 - lukaslueg:intersperse_with, r=m-ou-se

Add Iterator::intersperse_with

This is a follow-up to #79479, tracking in #79524, as discussed https://github.com/rust-lang/rust/pull/79479#issuecomment-752671731.

~~Note that I had to manually implement `Clone` and `Debug` because `derive` insists on placing a `Clone`-bound on the struct-definition, which is too narrow. There is a long-standing issue # for this somewhere around here :-)~~

Also, note that I refactored the guts of `Intersperse` into private functions and re-used them in `IntersperseWith`, so I also went light on duplicating all the tests.

If this is suitable to be merged, the tracking issue should be updated, since it only mentions `intersperse`.

Happy New Year!

r? ``@m-ou-se``

5 files changed, 202 insertions, 27 deletions

diff --git a/library/core/src/iter/adapters/intersperse.rs b/library/core/src/iter/adapters/intersperse.rs
index 36232672549..1d01e9b5fb7 100644
--- a/library/core/src/iter/adapters/intersperse.rs
+++ b/library/core/src/iter/adapters/intersperse.rs
@@ -1,6 +1,9 @@
 use super::Peekable;
 
 /// An iterator adapter that places a separator between all elements.
+///
+/// This `struct` is created by [`Iterator::intersperse`]. See its documentation
+/// for more information.
 #[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
 #[derive(Debug, Clone)]
 pub struct Intersperse<I: Iterator>
@@ -40,37 +43,146 @@ 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.
+///
+/// This `struct` is created by [`Iterator::intersperse_with`]. See its
+/// documentation for more information.
+#[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,
+}
+
+#[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);
-            }
+#[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 0023de65d2b..83d339d8f40 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,
 };
@@ -571,6 +571,9 @@ pub trait Iterator {
 
     /// Places a copy of `separator` between all elements.
     ///
+    /// In case the separator does not implement [`Clone`] or needs to be
+    /// computed every time, use [`intersperse_with`].
+    ///
     /// # Examples
     ///
     /// Basic usage:
@@ -578,9 +581,12 @@ pub trait Iterator {
     /// ```
     /// #![feature(iter_intersperse)]
     ///
-    /// let hello = ["Hello", "World"].iter().copied().intersperse(" ").collect::<String>();
-    /// assert_eq!(hello, "Hello World");
+    /// let hello = ["Hello", "World", "!"].iter().copied().intersperse(" ").collect::<String>();
+    /// assert_eq!(hello, "Hello World !");
     /// ```
+    ///
+    /// [`Clone`]: crate::clone::Clone
+    /// [`intersperse_with`]: Iterator::intersperse_with
     #[inline]
     #[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")]
     fn intersperse(self, separator: Self::Item) -> Intersperse<Self>
@@ -591,6 +597,33 @@ 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 happy_emojis = [" ❤️ ", " 😀 "].iter().copied();
+    /// let separator = || happy_emojis.next().unwrap_or(" 🦀 ");
+    ///
+    /// let result = src.intersperse_with(separator).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]);
+}