From 5a8a47283574a8fd11f81fbe059ecd0361b5969d Mon Sep 17 00:00:00 2001
From: Barak Ugav <barakugav@gmail.com>
Date: Sun, 1 Sep 2024 12:41:02 +0300
Subject: [PATCH] Add `zip_clones`, zips an iterator with clones of a value

Similar to `iter.zip(repeat_n(val, n))' but does not require knowing n
---
 src/lib.rs        | 39 ++++++++++++++++++++++++
 src/zip_clones.rs | 77 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 116 insertions(+)
 create mode 100644 src/zip_clones.rs

diff --git a/src/lib.rs b/src/lib.rs
index a4bb5e46c..0bb3c9ecf 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -142,6 +142,7 @@ pub mod structs {
     #[cfg(feature = "use_std")]
     pub use crate::unique_impl::{Unique, UniqueBy};
     pub use crate::with_position::WithPosition;
+    pub use crate::zip_clones::ZipClones;
     pub use crate::zip_eq_impl::ZipEq;
     pub use crate::zip_longest::ZipLongest;
     pub use crate::ziptuple::Zip;
@@ -233,6 +234,7 @@ mod tuple_impl;
 mod unique_impl;
 mod unziptuple;
 mod with_position;
+mod zip_clones;
 mod zip_eq_impl;
 mod zip_longest;
 mod ziptuple;
@@ -617,6 +619,43 @@ pub trait Itertools: Iterator {
         zip_eq(self, other)
     }
 
+    /// Create an iterator which iterates over this iterator paired with clones of a given value.
+    ///
+    /// If the iterator has `n` elements, the zipped value will be cloned `n-1` times. This function
+    /// is useful when the zipped value is expensive to clone and you want to avoid cloning it `n` times,
+    /// using the trivial following code:
+    /// ```rust
+    /// let it = [0, 1, 2, 3, 4].iter();
+    /// let zipped = "expensive-to-clone".to_string();
+    /// for a in it {
+    ///     let b = zipped.clone();
+    ///     // do something that consumes the expensive zipped value
+    ///     drop((a, b));
+    /// }
+    /// ```
+    /// Instead, you can use `zip_clones`:
+    /// ```rust
+    /// use itertools::Itertools;
+    /// let it = [0, 1, 2, 3, 4].iter();
+    /// let zipped = "expensive-to-clone".to_string();
+    /// for (a, b) in it.zip_clones(zipped) {
+    ///     // do something that consumes the expensive zipped value
+    ///     drop((a, b));
+    /// }
+    /// ```
+    ///
+    /// The [`repeat_n()`](crate::repeat_n) function can be used to create from a zipped value
+    /// an iterator that also clones the value `n-1` times, but it require to know the number of
+    /// elements in the iterator in advance.
+    #[inline]
+    fn zip_clones<T>(self, zipped: T) -> ZipClones<Self, T>
+    where
+        Self: Sized,
+        T: Clone,
+    {
+        zip_clones::zip_clones(self, zipped)
+    }
+
     /// A “meta iterator adaptor”. Its closure receives a reference to the
     /// iterator and may pick off as many elements as it likes, to produce the
     /// next iterator element.
diff --git a/src/zip_clones.rs b/src/zip_clones.rs
new file mode 100644
index 000000000..c59ed0618
--- /dev/null
+++ b/src/zip_clones.rs
@@ -0,0 +1,77 @@
+use std::iter::Peekable;
+
+/// An iterator which iterate over an iterator and clones of a value.
+///
+///
+/// See [`.zip_clones()`](crate::Itertools::zip_clones) for more information.
+pub struct ZipClones<I, T>
+where
+    I: Iterator,
+{
+    iter: Peekable<I>,
+    cloned: Option<T>,
+}
+
+/// Zips an iterator with clones of a value.
+///
+/// [`IntoIterator`] enabled version of [`Itertools::zip_clones`](crate::Itertools::zip_clones).
+///
+/// ```
+/// use itertools::Itertools;
+///
+/// let data = [1, 2, 3, 4, 5];
+/// let zipped = "expensive-to-clone".to_string();
+/// for (a, b) in data.iter().zip_clones(zipped) {
+///     // do something that consumes the expensive zipped value
+///     drop((a, b));
+/// }
+/// ```
+pub fn zip_clones<I, T>(i: I, zipped: T) -> ZipClones<I::IntoIter, T>
+where
+    I: IntoIterator,
+    T: Clone,
+{
+    ZipClones {
+        iter: i.into_iter().peekable(),
+        cloned: Some(zipped),
+    }
+}
+
+impl<I: Iterator, T: Clone> Iterator for ZipClones<I, T> {
+    type Item = (I::Item, T);
+    fn next(&mut self) -> Option<Self::Item> {
+        // let cur = self.next.take()?;
+        let cur = self.iter.next()?;
+        let zipped = if self.iter.peek().is_some() {
+            self.cloned.clone()
+        } else {
+            self.cloned.take()
+        };
+        // Safety: the zipped field is Some as long as the iterator is not exhausted
+        let zipped = unsafe { zipped.unwrap_unchecked() };
+        Some((cur, zipped))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::Itertools;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+
+    #[test]
+    fn test_zip_clones() {
+        static ZIPPED_CLONES_COUNTER: AtomicUsize = AtomicUsize::new(0);
+        struct Zipped {}
+        impl Clone for Zipped {
+            fn clone(&self) -> Self {
+                ZIPPED_CLONES_COUNTER.fetch_add(1, Ordering::SeqCst);
+                Zipped {}
+            }
+        }
+
+        ZIPPED_CLONES_COUNTER.store(0, Ordering::SeqCst);
+        let iter_len = [1, 2, 3, 4, 5, 6].iter().zip_clones(Zipped {}).count();
+        assert_eq!(iter_len, 6);
+        assert_eq!(ZIPPED_CLONES_COUNTER.load(Ordering::SeqCst), 5);
+    }
+}