Rollup merge of #106276 - Sp00ph:unify_slice_ranges, r=the8472

Fix `vec_deque::Drain` FIXME

In my original `VecDeque` rewrite, I didn't use `VecDeque::slice_ranges` in `Drain::as_slices`, even though that's basically the exact use case for `slice_ranges`. The reason for this was that a `VecDeque` wrapped in a `Drain` actually has its length set to `drain_start`, so that there's no potential use after free if you `mem::forget` the `Drain`. I modified `slice_ranges` to accept an explicit `len` parameter instead, which it now uses to bounds check the given range. This way, `Drain::as_slices` can use `slice_ranges` internally instead of having to basically just copy paste the `slice_ranges` code. Since `slice_ranges` is just an internal helper function, this shouldn't change the user facing behavior in any way.

2 files changed, 31 insertions, 36 deletions

diff --git a/library/alloc/src/collections/vec_deque/drain.rs b/library/alloc/src/collections/vec_deque/drain.rs
index 89feb361ddc..0be274a3822 100644
--- a/library/alloc/src/collections/vec_deque/drain.rs
+++ b/library/alloc/src/collections/vec_deque/drain.rs
@@ -52,36 +52,22 @@ impl<'a, T, A: Allocator> Drain<'a, T, A> {
         }
     }
 
-    // Only returns pointers to the slices, as that's
-    // all we need to drop them. May only be called if `self.remaining != 0`.
+    // Only returns pointers to the slices, as that's all we need
+    // to drop them. May only be called if `self.remaining != 0`.
     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,
-            // it's guaranteed that `a_range` and `b_range` represent valid ranges into
-            // the deques buffer.
+
+            // We know that `self.idx + self.remaining <= deque.len <= usize::MAX`, so this won't overflow.
+            let logical_remaining_range = self.idx..self.idx + self.remaining;
+
+            // SAFETY: `logical_remaining_range` represents the
+            // range into the logical buffer of elements that
+            // haven't been drained yet, so they're all initialized,
+            // and `slice::range(start..end, end) == start..end`,
+            // so the preconditions for `slice_ranges` are met.
+            let (a_range, b_range) =
+                deque.slice_ranges(logical_remaining_range.clone(), logical_remaining_range.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 8317ac431a5..48e907e402c 100644
--- a/library/alloc/src/collections/vec_deque/mod.rs
+++ b/library/alloc/src/collections/vec_deque/mod.rs
@@ -1156,7 +1156,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)) }
@@ -1190,7 +1190,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)) }
@@ -1232,19 +1232,28 @@ 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.
+    ///
+    /// # Safety
+    ///
+    /// This function is always safe to call. For the resulting ranges to be valid
+    /// ranges into the physical buffer, the caller must ensure that the result of
+    /// calling `slice::range(range, ..len)` represents a valid range into the
+    /// logical buffer, and that all elements in that range are initialized.
+    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`
+            // `slice::range` guarantees that `start <= end <= 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);
 
@@ -1290,7 +1299,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
@@ -1330,7 +1339,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