Rollup merge of #132187 - shahn:extend_more_tuples, r=dtolnay

Add Extend impls for tuples of arity 1 through 12

2 files changed, 142 insertions, 110 deletions

diff --git a/library/core/src/iter/traits/collect.rs b/library/core/src/iter/traits/collect.rs
index 2cf2ea58fd4..c3c7288e389 100644
--- a/library/core/src/iter/traits/collect.rs
+++ b/library/core/src/iter/traits/collect.rs
@@ -492,131 +492,150 @@ 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([(2, 3), (4, 5), (6, 7)]);
-    /// assert_eq!(tuple.0, [0, 2, 4, 6]);
-    /// assert_eq!(tuple.1, [1, 3, 5, 7]);
-    ///
-    /// // also allows for arbitrarily nested tuples as elements
-    /// let mut nested_tuple = (vec![1], (vec![2], vec![3]));
-    /// nested_tuple.extend([(4, (5, 6)), (7, (8, 9))]);
-    ///
-    /// let (a, (b, c)) = nested_tuple;
-    /// assert_eq!(a, [1, 4, 7]);
-    /// assert_eq!(b, [2, 5, 8]);
-    /// assert_eq!(c, [3, 6, 9]);
-    /// ```
-    fn extend<T: IntoIterator<Item = (A, B)>>(&mut self, into_iter: T) {
-        let (a, b) = self;
-        let iter = into_iter.into_iter();
-        SpecTupleExtend::extend(iter, a, b);
-    }
+macro_rules! spec_tuple_impl {
+    ( ($ty_name:ident, $var_name:ident, $extend_ty_name: ident, $trait_name:ident, $default_fn_name:ident, $cnt:tt), ) => {
+        spec_tuple_impl!($trait_name, $default_fn_name, #[doc(fake_variadic)] #[doc = "This trait is implemented for tuples up to twelve items long. The `impl`s for 1- and 3- through 12-ary tuples were stabilized after 2-tuples, in RUSTC_CURRENT_VERSION."] => ($ty_name, $var_name, $extend_ty_name, $cnt),);
+    };
+    ( ($ty_name:ident, $var_name:ident, $extend_ty_name: ident, $trait_name:ident, $default_fn_name:ident, $cnt:tt), $(($ty_names:ident, $var_names:ident,  $extend_ty_names:ident, $trait_names:ident, $default_fn_names:ident, $cnts:tt),)*) => {
 
-    fn extend_one(&mut self, item: (A, B)) {
-        self.0.extend_one(item.0);
-        self.1.extend_one(item.1);
-    }
+        spec_tuple_impl!($(($ty_names, $var_names, $extend_ty_names, $trait_names, $default_fn_names, $cnts),)*);
+        spec_tuple_impl!($trait_name, $default_fn_name, #[doc(hidden)] => ($ty_name, $var_name, $extend_ty_name, $cnt), $(($ty_names, $var_names, $extend_ty_names, $cnts),)*);
+    };
+    ($trait_name:ident, $default_fn_name:ident, #[$meta:meta] $(#[$doctext:meta])? => $(($ty_names:ident, $var_names:ident, $extend_ty_names:ident, $cnts:tt),)*) => {
+        #[$meta]
+        $(#[$doctext])?
+        #[stable(feature = "extend_for_tuple", since = "1.56.0")]
+        impl<$($ty_names,)* $($extend_ty_names,)*> Extend<($($ty_names,)*)> for ($($extend_ty_names,)*)
+        where
+            $($extend_ty_names: Extend<$ty_names>,)*
+        {
+            /// Allows to `extend` a tuple of collections that also implement `Extend`.
+            ///
+            /// See also: [`Iterator::unzip`]
+            ///
+            /// # Examples
+            /// ```
+            /// // Example given for a 2-tuple, but 1- through 12-tuples are supported
+            /// let mut tuple = (vec![0], vec![1]);
+            /// tuple.extend([(2, 3), (4, 5), (6, 7)]);
+            /// assert_eq!(tuple.0, [0, 2, 4, 6]);
+            /// assert_eq!(tuple.1, [1, 3, 5, 7]);
+            ///
+            /// // also allows for arbitrarily nested tuples as elements
+            /// let mut nested_tuple = (vec![1], (vec![2], vec![3]));
+            /// nested_tuple.extend([(4, (5, 6)), (7, (8, 9))]);
+            ///
+            /// let (a, (b, c)) = nested_tuple;
+            /// assert_eq!(a, [1, 4, 7]);
+            /// assert_eq!(b, [2, 5, 8]);
+            /// assert_eq!(c, [3, 6, 9]);
+            /// ```
+            fn extend<T: IntoIterator<Item = ($($ty_names,)*)>>(&mut self, into_iter: T) {
+                let ($($var_names,)*) = self;
+                let iter = into_iter.into_iter();
+                $trait_name::extend(iter, $($var_names,)*);
+            }
 
-    fn extend_reserve(&mut self, additional: usize) {
-        self.0.extend_reserve(additional);
-        self.1.extend_reserve(additional);
-    }
+            fn extend_one(&mut self, item: ($($ty_names,)*)) {
+                $(self.$cnts.extend_one(item.$cnts);)*
+            }
 
-    unsafe fn extend_one_unchecked(&mut self, item: (A, B)) {
-        // SAFETY: Those are our safety preconditions, and we correctly forward `extend_reserve`.
-        unsafe {
-            self.0.extend_one_unchecked(item.0);
-            self.1.extend_one_unchecked(item.1);
-        }
-    }
-}
+            fn extend_reserve(&mut self, additional: usize) {
+                $(self.$cnts.extend_reserve(additional);)*
+            }
 
-fn default_extend_tuple<A, B, ExtendA, ExtendB>(
-    iter: impl Iterator<Item = (A, B)>,
-    a: &mut ExtendA,
-    b: &mut ExtendB,
-) where
-    ExtendA: Extend<A>,
-    ExtendB: Extend<B>,
-{
-    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);
+            unsafe fn extend_one_unchecked(&mut self, item: ($($ty_names,)*)) {
+                // SAFETY: Those are our safety preconditions, and we correctly forward `extend_reserve`.
+                unsafe {
+                     $(self.$cnts.extend_one_unchecked(item.$cnts);)*
+                }
+            }
         }
-    }
 
-    let (lower_bound, _) = iter.size_hint();
-    if lower_bound > 0 {
-        a.extend_reserve(lower_bound);
-        b.extend_reserve(lower_bound);
-    }
+        trait $trait_name<$($ty_names),*> {
+            fn extend(self, $($var_names: &mut $ty_names,)*);
+        }
 
-    iter.fold((), extend(a, b));
-}
+        fn $default_fn_name<$($ty_names,)* $($extend_ty_names,)*>(
+            iter: impl Iterator<Item = ($($ty_names,)*)>,
+            $($var_names: &mut $extend_ty_names,)*
+        ) where
+            $($extend_ty_names: Extend<$ty_names>,)*
+        {
+            fn extend<'a, $($ty_names,)*>(
+                $($var_names: &'a mut impl Extend<$ty_names>,)*
+            ) -> impl FnMut((), ($($ty_names,)*)) + 'a {
+                #[allow(non_snake_case)]
+                move |(), ($($extend_ty_names,)*)| {
+                    $($var_names.extend_one($extend_ty_names);)*
+                }
+            }
 
-trait SpecTupleExtend<A, B> {
-    fn extend(self, a: &mut A, b: &mut B);
-}
+            let (lower_bound, _) = iter.size_hint();
+            if lower_bound > 0 {
+                $($var_names.extend_reserve(lower_bound);)*
+            }
 
-impl<A, B, ExtendA, ExtendB, Iter> SpecTupleExtend<ExtendA, ExtendB> for Iter
-where
-    ExtendA: Extend<A>,
-    ExtendB: Extend<B>,
-    Iter: Iterator<Item = (A, B)>,
-{
-    default fn extend(self, a: &mut ExtendA, b: &mut ExtendB) {
-        default_extend_tuple(self, a, b);
-    }
-}
+            iter.fold((), extend($($var_names,)*));
+        }
 
-impl<A, B, ExtendA, ExtendB, Iter> SpecTupleExtend<ExtendA, ExtendB> for Iter
-where
-    ExtendA: Extend<A>,
-    ExtendB: Extend<B>,
-    Iter: TrustedLen<Item = (A, B)>,
-{
-    fn extend(self, a: &mut ExtendA, b: &mut ExtendB) {
-        fn extend<'a, A, B>(
-            a: &'a mut impl Extend<A>,
-            b: &'a mut impl Extend<B>,
-        ) -> impl FnMut((), (A, B)) + 'a {
-            // SAFETY: We reserve enough space for the `size_hint`, and the iterator is `TrustedLen`
-            // so its `size_hint` is exact.
-            move |(), (t, u)| unsafe {
-                a.extend_one_unchecked(t);
-                b.extend_one_unchecked(u);
+        impl<$($ty_names,)* $($extend_ty_names,)* Iter> $trait_name<$($extend_ty_names),*> for Iter
+            where
+            $($extend_ty_names: Extend<$ty_names>,)*
+            Iter: Iterator<Item = ($($ty_names,)*)>,
+        {
+            default fn extend(self, $($var_names: &mut $extend_ty_names),*) {
+                $default_fn_name(self, $($var_names),*);
             }
         }
 
-        let (lower_bound, upper_bound) = self.size_hint();
+        impl<$($ty_names,)* $($extend_ty_names,)* Iter> $trait_name<$($extend_ty_names),*> for Iter
+            where
+            $($extend_ty_names: Extend<$ty_names>,)*
+            Iter: TrustedLen<Item = ($($ty_names,)*)>,
+        {
+            fn extend(self, $($var_names: &mut $extend_ty_names,)*) {
+                fn extend<'a, $($ty_names,)*>(
+                    $($var_names: &'a mut impl Extend<$ty_names>,)*
+                ) -> impl FnMut((), ($($ty_names,)*)) + 'a {
+                #[allow(non_snake_case)]
+                // SAFETY: We reserve enough space for the `size_hint`, and the iterator is `TrustedLen`
+                // so its `size_hint` is exact.
+                move |(), ($($extend_ty_names,)*)| unsafe {
+                    $($var_names.extend_one_unchecked($extend_ty_names);)*
+                }
+            }
 
-        if upper_bound.is_none() {
-            // We cannot reserve more than `usize::MAX` items, and this is likely to go out of memory anyway.
-            default_extend_tuple(self, a, b);
-            return;
-        }
+            let (lower_bound, upper_bound) = self.size_hint();
 
-        if lower_bound > 0 {
-            a.extend_reserve(lower_bound);
-            b.extend_reserve(lower_bound);
-        }
+            if upper_bound.is_none() {
+                // We cannot reserve more than `usize::MAX` items, and this is likely to go out of memory anyway.
+                $default_fn_name(self, $($var_names,)*);
+                return;
+            }
 
-        self.fold((), extend(a, b));
+            if lower_bound > 0 {
+                $($var_names.extend_reserve(lower_bound);)*
+            }
+
+            self.fold((), extend($($var_names,)*));
+        }
     }
+
+    };
 }
+
+spec_tuple_impl!(
+    (L, l, EL, TraitL, default_extend_tuple_l, 11),
+    (K, k, EK, TraitK, default_extend_tuple_k, 10),
+    (J, j, EJ, TraitJ, default_extend_tuple_j, 9),
+    (I, i, EI, TraitI, default_extend_tuple_i, 8),
+    (H, h, EH, TraitH, default_extend_tuple_h, 7),
+    (G, g, EG, TraitG, default_extend_tuple_g, 6),
+    (F, f, EF, TraitF, default_extend_tuple_f, 5),
+    (E, e, EE, TraitE, default_extend_tuple_e, 4),
+    (D, d, ED, TraitD, default_extend_tuple_d, 3),
+    (C, c, EC, TraitC, default_extend_tuple_c, 2),
+    (B, b, EB, TraitB, default_extend_tuple_b, 1),
+    (A, a, EA, TraitA, default_extend_tuple_a, 0),
+);
diff --git a/library/core/tests/iter/traits/iterator.rs b/library/core/tests/iter/traits/iterator.rs
index 93ef9c0812b..76f1e3319d4 100644
--- a/library/core/tests/iter/traits/iterator.rs
+++ b/library/core/tests/iter/traits/iterator.rs
@@ -617,6 +617,19 @@ fn test_next_chunk() {
     assert_eq!(it.next_chunk::<0>().unwrap(), []);
 }
 
+#[test]
+fn test_collect_into_tuples() {
+    let a = vec![(1, 2, 3), (4, 5, 6), (7, 8, 9)];
+    let b = vec![1, 4, 7];
+    let c = vec![2, 5, 8];
+    let d = vec![3, 6, 9];
+    let mut e = (Vec::new(), Vec::new(), Vec::new());
+    a.iter().cloned().collect_into(&mut e);
+    assert!(e.0 == b);
+    assert!(e.1 == c);
+    assert!(e.2 == d);
+}
+
 // just tests by whether or not this compiles
 fn _empty_impl_all_auto_traits<T>() {
     use std::panic::{RefUnwindSafe, UnwindSafe};