Implement iter::Chain::{advance_by, advance_back_by}

Author: timvermeulen

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))
     }