From c5d6a0dd96ad9825751ac6189fa632f521a16db2 Mon Sep 17 00:00:00 2001
From: Tim Vermeulen <tvermeulen@me.com>
Date: Fri, 2 Oct 2020 05:26:36 +0200
Subject: [PATCH] Implement iter::Chain::{advance_by, advance_back_by}

---
 library/core/src/iter/adapters/chain.rs | 76 +++++++++++++++++++++----
 1 file changed, 64 insertions(+), 12 deletions(-)

diff --git a/library/core/src/iter/adapters/chain.rs b/library/core/src/iter/adapters/chain.rs
index 13c6a75d58b..dadd7bfb0f7 100644
--- a/library/core/src/iter/adapters/chain.rs
+++ b/library/core/src/iter/adapters/chain.rs
@@ -115,16 +115,42 @@ where
     }
 
     #[inline]
-    fn nth(&mut self, mut n: usize) -> Option<A::Item> {
+    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
+        let mut rem = n;
+
         if let Some(ref mut a) = self.a {
-            while let Some(x) = a.next() {
-                if n == 0 {
-                    return Some(x);
-                }
-                n -= 1;
+            match a.advance_by(rem) {
+                Ok(()) => return Ok(()),
+                Err(k) => rem -= k,
             }
             self.a = None;
         }
+
+        if let Some(ref mut b) = self.b {
+            match b.advance_by(rem) {
+                Ok(()) => return Ok(()),
+                Err(k) => rem -= k,
+            }
+            // we don't fuse the second iterator
+        }
+
+        if rem == 0 { Ok(()) } else { Err(n - rem) }
+    }
+
+    #[inline]
+    fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
+        if let Some(ref mut a) = self.a {
+            match a.advance_by(n) {
+                Ok(()) => match a.next() {
+                    None => n = 0,
+                    x => return x,
+                },
+                Err(k) => n -= k,
+            }
+
+            self.a = None;
+        }
+
         maybe!(self.b.nth(n))
     }
 
@@ -191,16 +217,42 @@ where
     }
 
     #[inline]
-    fn nth_back(&mut self, mut n: usize) -> Option<A::Item> {
+    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
+        let mut rem = n;
+
         if let Some(ref mut b) = self.b {
-            while let Some(x) = b.next_back() {
-                if n == 0 {
-                    return Some(x);
-                }
-                n -= 1;
+            match b.advance_back_by(rem) {
+                Ok(()) => return Ok(()),
+                Err(k) => rem -= k,
             }
             self.b = None;
         }
+
+        if let Some(ref mut a) = self.a {
+            match a.advance_back_by(rem) {
+                Ok(()) => return Ok(()),
+                Err(k) => rem -= k,
+            }
+            // we don't fuse the second iterator
+        }
+
+        if rem == 0 { Ok(()) } else { Err(n - rem) }
+    }
+
+    #[inline]
+    fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> {
+        if let Some(ref mut b) = self.b {
+            match b.advance_back_by(n) {
+                Ok(()) => match b.next_back() {
+                    None => n = 0,
+                    x => return x,
+                },
+                Err(k) => n -= k,
+            }
+
+            self.b = None;
+        }
+
         maybe!(self.a.nth_back(n))
     }