Rollup merge of #85835 - Seppel3210:master, r=yaahc

Implement Extend<(A, B)> for (Extend<A>, Extend<B>)

I oriented myself at the implementation of `Iterator::unzip` and also rewrote the impl in terms of `(A, B)::extend` after that.

Since (A, B) now also implements Extend we could also mention in the documentation of unzip that it can do "nested unzipping" (you could unzip `Iterator<Item=(A, (B, C))>` into `(Vec<A>, (Vec<B>, Vec<C>))` for example) but I'm not sure of that so I'm asking here 🙂

(P.S. I saw a couple of people asking if there is an unzip3 but there isn't. So this could be a way to get equivalent functionality)

2 files changed, 69 insertions, 22 deletions

diff --git a/library/core/src/iter/traits/collect.rs b/library/core/src/iter/traits/collect.rs
index 7f87ead6fee..aa91346851f 100644
--- a/library/core/src/iter/traits/collect.rs
+++ b/library/core/src/iter/traits/collect.rs
@@ -360,3 +360,61 @@ impl Extend<()> for () {
     }
     fn extend_one(&mut self, _item: ()) {}
 }
+
+#[stable(feature = "extend_for_tuple", since = "1.56.0")]
+impl<A, B, ExtendA, ExtendB> Extend<(A, B)> for (ExtendA, ExtendB)
+where
+    ExtendA: Extend<A>,
+    ExtendB: Extend<B>,
+{
+    /// Allows to `extend` a tuple of collections that also implement `Extend`.
+    ///
+    /// See also: [`Iterator::unzip`]
+    ///
+    /// # Examples
+    /// ```
+    /// let mut tuple = (vec![0], vec![1]);
+    /// tuple.extend(vec![(2, 3), (4, 5), (6, 7)]);
+    /// assert_eq!(tuple.0, vec![0, 2, 4, 6]);
+    /// assert_eq!(tuple.1, vec![1, 3, 5, 7]);
+    ///
+    /// // also allows for arbitrarily nested tuples
+    /// let mut nested_tuple = (vec![(1, -1)], vec![(2, -2)]);
+    /// nested_tuple.extend(vec![((3, -3), (4, -4)), ((5, -5), (6, -6))]);
+    ///
+    /// assert_eq!(nested_tuple.0, vec![(1, -1), (3, -3), (5, -5)]);
+    /// assert_eq!(nested_tuple.1, vec![(2, -2), (4, -4), (6, -6)]);
+    /// ```
+    fn extend<T: IntoIterator<Item = (A, B)>>(&mut self, into_iter: T) {
+        let (a, b) = self;
+        let iter = into_iter.into_iter();
+
+        fn extend<'a, A, B>(
+            a: &'a mut impl Extend<A>,
+            b: &'a mut impl Extend<B>,
+        ) -> impl FnMut((), (A, B)) + 'a {
+            move |(), (t, u)| {
+                a.extend_one(t);
+                b.extend_one(u);
+            }
+        }
+
+        let (lower_bound, _) = iter.size_hint();
+        if lower_bound > 0 {
+            a.extend_reserve(lower_bound);
+            b.extend_reserve(lower_bound);
+        }
+
+        iter.fold((), extend(a, b));
+    }
+
+    fn extend_one(&mut self, item: (A, B)) {
+        self.0.extend_one(item.0);
+        self.1.extend_one(item.1);
+    }
+
+    fn extend_reserve(&mut self, additional: usize) {
+        self.0.extend_reserve(additional);
+        self.1.extend_reserve(additional);
+    }
+}
diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs
index 6b24d33bebc..524d8f857e2 100644
--- a/library/core/src/iter/traits/iterator.rs
+++ b/library/core/src/iter/traits/iterator.rs
@@ -2841,6 +2841,14 @@ pub trait Iterator {
     ///
     /// assert_eq!(left, [1, 3]);
     /// assert_eq!(right, [2, 4]);
+    ///
+    /// // you can also unzip multiple nested tuples at once
+    /// let a = [(1, (2, 3)), (4, (5, 6))];
+    ///
+    /// let (x, (y, z)): (Vec<_>, (Vec<_>, Vec<_>)) = a.iter().cloned().unzip();
+    /// assert_eq!(x, [1, 4]);
+    /// assert_eq!(y, [2, 5]);
+    /// assert_eq!(z, [3, 6]);
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     fn unzip<A, B, FromA, FromB>(self) -> (FromA, FromB)
@@ -2849,28 +2857,9 @@ pub trait Iterator {
         FromB: Default + Extend<B>,
         Self: Sized + Iterator<Item = (A, B)>,
     {
-        fn extend<'a, A, B>(
-            ts: &'a mut impl Extend<A>,
-            us: &'a mut impl Extend<B>,
-        ) -> impl FnMut((), (A, B)) + 'a {
-            move |(), (t, u)| {
-                ts.extend_one(t);
-                us.extend_one(u);
-            }
-        }
-
-        let mut ts: FromA = Default::default();
-        let mut us: FromB = Default::default();
-
-        let (lower_bound, _) = self.size_hint();
-        if lower_bound > 0 {
-            ts.extend_reserve(lower_bound);
-            us.extend_reserve(lower_bound);
-        }
-
-        self.fold((), extend(&mut ts, &mut us));
-
-        (ts, us)
+        let mut unzipped: (FromA, FromB) = Default::default();
+        unzipped.extend(self);
+        unzipped
     }
 
     /// Creates an iterator which copies all of its elements.