Fix `vec_deque::Drain` FIXME

2 files changed, 19 insertions, 31 deletions

diff --git a/library/alloc/src/collections/vec_deque/drain.rs b/library/alloc/src/collections/vec_deque/drain.rs
index 89feb361ddc..a102aaad452 100644
--- a/library/alloc/src/collections/vec_deque/drain.rs
+++ b/library/alloc/src/collections/vec_deque/drain.rs
@@ -57,31 +57,17 @@ impl<'a, T, A: Allocator> Drain<'a, T, A> {
     unsafe fn as_slices(&self) -> (*mut [T], *mut [T]) {
         unsafe {
             let deque = self.deque.as_ref();
-            // FIXME: This is doing almost exactly the same thing as the else branch in `VecDeque::slice_ranges`.
-            // Unfortunately, we can't just call `slice_ranges` here, as the deque's `len` is currently
-            // just `drain_start`, so the range check would (almost) always panic. Between temporarily
-            // adjusting the deques `len` to call `slice_ranges`, and just copy pasting the `slice_ranges`
-            // implementation, this seemed like the less hacky solution, though it might be good to
-            // find a better one in the future.
-
-            // because `self.remaining != 0`, we know that `self.idx < deque.original_len`, so it's a valid
-            // logical index.
-            let wrapped_start = deque.to_physical_idx(self.idx);
-
-            let head_len = deque.capacity() - wrapped_start;
-
-            let (a_range, b_range) = if head_len >= self.remaining {
-                (wrapped_start..wrapped_start + self.remaining, 0..0)
-            } else {
-                let tail_len = self.remaining - head_len;
-                (wrapped_start..deque.capacity(), 0..tail_len)
-            };
-
-            // SAFETY: the range `self.idx..self.idx+self.remaining` lies strictly inside
-            // the range `0..deque.original_len`. because of this, and because of the fact
-            // that we acquire `a_range` and `b_range` exactly like `slice_ranges` would,
+
+            let start = self.idx;
+            // We know that `self.idx + self.remaining <= deque.len <= usize::MAX`, so this won't overflow.
+            let end = start + self.remaining;
+
+            // SAFETY: the range `start..end` lies strictly inside
+            // the range `0..deque.original_len`. Because of this, and because
+            // we haven't touched the elements inside this range yet,
             // it's guaranteed that `a_range` and `b_range` represent valid ranges into
             // the deques buffer.
+            let (a_range, b_range) = deque.slice_ranges(start..end, end);
             (deque.buffer_range(a_range), deque.buffer_range(b_range))
         }
     }
diff --git a/library/alloc/src/collections/vec_deque/mod.rs b/library/alloc/src/collections/vec_deque/mod.rs
index c955db46d29..6d3e784c8b7 100644
--- a/library/alloc/src/collections/vec_deque/mod.rs
+++ b/library/alloc/src/collections/vec_deque/mod.rs
@@ -1147,7 +1147,7 @@ impl<T, A: Allocator> VecDeque<T, A> {
     #[inline]
     #[stable(feature = "deque_extras_15", since = "1.5.0")]
     pub fn as_slices(&self) -> (&[T], &[T]) {
-        let (a_range, b_range) = self.slice_ranges(..);
+        let (a_range, b_range) = self.slice_ranges(.., self.len);
         // SAFETY: `slice_ranges` always returns valid ranges into
         // the physical buffer.
         unsafe { (&*self.buffer_range(a_range), &*self.buffer_range(b_range)) }
@@ -1181,7 +1181,7 @@ impl<T, A: Allocator> VecDeque<T, A> {
     #[inline]
     #[stable(feature = "deque_extras_15", since = "1.5.0")]
     pub fn as_mut_slices(&mut self) -> (&mut [T], &mut [T]) {
-        let (a_range, b_range) = self.slice_ranges(..);
+        let (a_range, b_range) = self.slice_ranges(.., self.len);
         // SAFETY: `slice_ranges` always returns valid ranges into
         // the physical buffer.
         unsafe { (&mut *self.buffer_range(a_range), &mut *self.buffer_range(b_range)) }
@@ -1223,19 +1223,21 @@ impl<T, A: Allocator> VecDeque<T, A> {
 
     /// Given a range into the logical buffer of the deque, this function
     /// return two ranges into the physical buffer that correspond to
-    /// the given range.
-    fn slice_ranges<R>(&self, range: R) -> (Range<usize>, Range<usize>)
+    /// the given range. The `len` parameter should usually just be `self.len`;
+    /// the reason it's passed explicitly is that if the deque is wrapped in
+    /// a `Drain`, then `self.len` is not actually the length of the deque.
+    fn slice_ranges<R>(&self, range: R, len: usize) -> (Range<usize>, Range<usize>)
     where
         R: RangeBounds<usize>,
     {
-        let Range { start, end } = slice::range(range, ..self.len);
+        let Range { start, end } = slice::range(range, ..len);
         let len = end - start;
 
         if len == 0 {
             (0..0, 0..0)
         } else {
             // `slice::range` guarantees that `start <= end <= self.len`.
-            // because `len != 0`, we know that `start < end`, so `start < self.len`
+            // because `len != 0`, we know that `start < end`, so `start < len`
             // and the indexing is valid.
             let wrapped_start = self.to_physical_idx(start);
 
@@ -1281,7 +1283,7 @@ impl<T, A: Allocator> VecDeque<T, A> {
     where
         R: RangeBounds<usize>,
     {
-        let (a_range, b_range) = self.slice_ranges(range);
+        let (a_range, b_range) = self.slice_ranges(range, self.len);
         // SAFETY: The ranges returned by `slice_ranges`
         // are valid ranges into the physical buffer, so
         // it's ok to pass them to `buffer_range` and
@@ -1321,7 +1323,7 @@ impl<T, A: Allocator> VecDeque<T, A> {
     where
         R: RangeBounds<usize>,
     {
-        let (a_range, b_range) = self.slice_ranges(range);
+        let (a_range, b_range) = self.slice_ranges(range, self.len);
         // SAFETY: The ranges returned by `slice_ranges`
         // are valid ranges into the physical buffer, so
         // it's ok to pass them to `buffer_range` and