fix advance_by impl for vec_deque and add tests

rust-lang · Mar 27, 2023 · 4180793 · 4180793
1 parent 9cd9da2
commit 4180793
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 7 deletions.
diff --git a/library/alloc/src/collections/vec_deque/into_iter.rs b/library/alloc/src/collections/vec_deque/into_iter.rs
@@ -56,10 +56,10 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
 
     #[inline]
     fn advance_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
-        let rem = if self.inner.len < n {
-            let len = self.inner.len;
+        let len = self.inner.len;
+        let rem = if len < n {
             self.inner.clear();
-            len - n
+            n - len
         } else {
             self.inner.drain(..n);
             0
@@ -186,12 +186,12 @@ impl<T, A: Allocator> DoubleEndedIterator for IntoIter<T, A> {
     #[inline]
     fn advance_back_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
         let len = self.inner.len;
-        let rem = if len >= n {
-            self.inner.truncate(len - n);
-            0
-        } else {
+        let rem = if len < n {
             self.inner.clear();
             n - len
+        } else {
+            self.inner.truncate(len - n);
+            0
         };
         NonZeroUsize::new(rem).map_or(Ok(()), Err)
     }

diff --git a/library/alloc/tests/vec_deque.rs b/library/alloc/tests/vec_deque.rs
@@ -1,3 +1,4 @@
+use core::num::NonZeroUsize;
 use std::assert_matches::assert_matches;
 use std::collections::TryReserveErrorKind::*;
 use std::collections::{vec_deque::Drain, VecDeque};
@@ -426,6 +427,28 @@ fn test_into_iter() {
         assert_eq!(it.next(), Some(7));
         assert_eq!(it.size_hint(), (5, Some(5)));
     }
+
+    // advance_by
+    {
+        let mut d = VecDeque::new();
+        for i in 0..=4 {
+            d.push_back(i);
+        }
+        for i in 6..=8 {
+            d.push_front(i);
+        }
+
+        let mut it = d.into_iter();
+        assert_eq!(it.advance_by(1), Ok(()));
+        assert_eq!(it.next(), Some(7));
+        assert_eq!(it.advance_back_by(1), Ok(()));
+        assert_eq!(it.next_back(), Some(3));
+
+        let mut it = VecDeque::from(vec![1, 2, 3, 4, 5]).into_iter();
+        assert_eq!(it.advance_by(10), Err(NonZeroUsize::new(5).unwrap()));
+        let mut it = VecDeque::from(vec![1, 2, 3, 4, 5]).into_iter();
+        assert_eq!(it.advance_back_by(10), Err(NonZeroUsize::new(5).unwrap()));
+    }
 }
 
 #[test]