Rollup merge of #81306 - SkiFire13:fuse-flatten, r=cuviper

Fuse inner iterator in FlattenCompat and improve related tests Fixes #81248
rust-lang · Jan 29, 2021 · 94e093a · 94e093a
2 parents 788036d + 5aa625b
commit 94e093a
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 2 deletions.
diff --git a/library/core/src/iter/adapters/flatten.rs b/library/core/src/iter/adapters/flatten.rs
@@ -265,7 +265,13 @@ where
                 }
             }
             match self.iter.next() {
-                None => return self.backiter.as_mut()?.next(),
+                None => match self.backiter.as_mut()?.next() {
+                    None => {
+                        self.backiter = None;
+                        return None;
+                    }
+                    elt @ Some(_) => return elt,
+                },
                 Some(inner) => self.frontiter = Some(inner.into_iter()),
             }
         }
@@ -353,7 +359,13 @@ where
                 }
             }
             match self.iter.next_back() {
-                None => return self.frontiter.as_mut()?.next_back(),
+                None => match self.frontiter.as_mut()?.next_back() {
+                    None => {
+                        self.frontiter = None;
+                        return None;
+                    }
+                    elt @ Some(_) => return elt,
+                },
                 next => self.backiter = next.map(IntoIterator::into_iter),
             }
         }

diff --git a/library/core/tests/iter/adapters/flatten.rs b/library/core/tests/iter/adapters/flatten.rs
@@ -64,6 +64,14 @@ fn test_flatten_non_fused_outer() {
     assert_eq!(iter.next_back(), Some(1));
     assert_eq!(iter.next(), Some(0));
     assert_eq!(iter.next(), None);
+    assert_eq!(iter.next(), None);
+
+    let mut iter = NonFused::new(once(0..2)).flatten();
+
+    assert_eq!(iter.next(), Some(0));
+    assert_eq!(iter.next_back(), Some(1));
+    assert_eq!(iter.next_back(), None);
+    assert_eq!(iter.next_back(), None);
 }
 
 #[test]
@@ -74,6 +82,15 @@ fn test_flatten_non_fused_inner() {
     assert_eq!(iter.next(), Some(0));
     assert_eq!(iter.next(), Some(1));
     assert_eq!(iter.next(), None);
+    assert_eq!(iter.next(), None);
+
+    let mut iter = once(0..1).chain(once(1..3)).flat_map(NonFused::new);
+
+    assert_eq!(iter.next(), Some(0));
+    assert_eq!(iter.next_back(), Some(2));
+    assert_eq!(iter.next_back(), Some(1));
+    assert_eq!(iter.next_back(), None);
+    assert_eq!(iter.next_back(), None);
 }
 
 #[test]