From 435c45ee86713256209bb37a912e4850aa267d22 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ga=C3=ABtan=20Cassiers?= <gaetan.cassiers@uclouvain.be>
Date: Thu, 4 Feb 2021 18:28:28 +0100
Subject: [PATCH] Add a chain! macro.

The chain! macro creates an iterator running multiple iterators
sequentially.
---
 src/lib.rs         | 34 ++++++++++++++++++++++++++++++++++
 tests/quick.rs     |  6 ++++++
 tests/test_core.rs | 20 ++++++++++++++++++++
 3 files changed, 60 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index 4db8f5d40..eaf670d88 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -346,6 +346,40 @@ macro_rules! izip {
     };
 }
 
+#[macro_export]
+/// Create an iterator running multiple iterators sequentially.
+///
+/// This is a version of the standard ``.chain()`` that's supporting more than
+/// two iterators. `chain!` takes `IntoIterator` arguments.
+/// Alternatively, this is an alternative to the standard ``.flatten()`` for a
+/// fixed number of iterators of distinct sizes.
+///
+/// **Note:** The result of this macro is in the general case an iterator
+/// composed of repeated `.chain()`.
+/// The special case of one arguments produce $a.into_iter().
+///
+///
+/// ```
+/// # use itertools::chain;
+/// #
+/// # fn main() {
+///
+/// // chain three sequences
+/// let chained: Vec<i32> = chain!(0..=3, 4..6, vec![6, 7]).collect();
+/// assert_eq!(chained, (0..=7).collect::<Vec<i32>>());
+/// # }
+/// ```
+macro_rules! chain {
+    ( $first:expr $(, $rest:expr )* $(,)*) => {
+        core::iter::IntoIterator::into_iter($first)
+            $(
+                .chain(
+                    core::iter::IntoIterator::into_iter($rest)
+                )
+            )*
+    };
+}
+
 /// An [`Iterator`] blanket implementation that provides extra adaptors and
 /// methods.
 ///
diff --git a/tests/quick.rs b/tests/quick.rs
index b4ae576af..b9b87f2d0 100644
--- a/tests/quick.rs
+++ b/tests/quick.rs
@@ -15,6 +15,7 @@ use itertools::{
     EitherOrBoth,
     iproduct,
     izip,
+    chain,
 };
 use itertools::free::{
     cloned,
@@ -526,6 +527,11 @@ quickcheck! {
         correct_size_hint(izip!(filt, b.clone(), c.clone())) &&
             exact_size(izip!(a, b, c))
     }
+
+    fn size_chain_macro(a: Iter<i16, Exact>, b: Iter<i16, Exact>, c: Iter<i16, Exact>) -> bool {
+        correct_size_hint(chain!(a.clone(), b.clone(), c.clone()))
+    }
+
     fn equal_kmerge(a: Vec<i16>, b: Vec<i16>, c: Vec<i16>) -> bool {
         use itertools::free::kmerge;
         let mut sa = a.clone();
diff --git a/tests/test_core.rs b/tests/test_core.rs
index 5861653da..bfff9a005 100644
--- a/tests/test_core.rs
+++ b/tests/test_core.rs
@@ -13,6 +13,7 @@ use crate::it::multizip;
 use crate::it::free::put_back;
 use crate::it::iproduct;
 use crate::it::izip;
+use crate::it::chain;
 
 #[test]
 fn product2() {
@@ -87,6 +88,25 @@ fn multizip3() {
     }
 }
 
+#[test]
+fn chain_macro() {
+    let mut chain = chain!(2..3);
+    assert!(chain.next() == Some(2));
+    assert!(chain.next().is_none());
+
+    let mut chain = chain!(0..2, 2..3, 3..5i8);
+    for i in 0..5i8 {
+        assert_eq!(Some(i), chain.next());
+    }
+    assert!(chain.next().is_none());
+}
+
+#[test]
+fn chain2() {
+    let _ = chain!(1.., 2..);
+    let _ = chain!(1.., 2.., );
+}
+
 #[test]
 fn write_to() {
     let xs = [7, 9, 8];