2 files changed, 30 insertions, 15 deletions

diff --git a/compiler/rustc_index/src/slice.rs b/compiler/rustc_index/src/slice.rs
index f17ea9e4b59..67ac805c2bf 100644
--- a/compiler/rustc_index/src/slice.rs
+++ b/compiler/rustc_index/src/slice.rs
@@ -1,6 +1,7 @@
 use std::fmt;
 use std::marker::PhantomData;
 use std::ops::{Index, IndexMut};
+use std::slice::GetDisjointMutError::*;
 use std::slice::{self, SliceIndex};
 
 use crate::{Idx, IndexVec, IntoSliceIdx};
@@ -65,6 +66,8 @@ impl<I: Idx, T> IndexSlice<I, T> {
 
     #[inline]
     pub fn iter_enumerated(&self) -> impl DoubleEndedIterator<Item = (I, &T)> + ExactSizeIterator {
+        // Allow the optimizer to elide the bounds checking when creating each index.
+        let _ = I::new(self.len());
         self.raw.iter().enumerate().map(|(n, t)| (I::new(n), t))
     }
 
@@ -72,6 +75,8 @@ impl<I: Idx, T> IndexSlice<I, T> {
     pub fn indices(
         &self,
     ) -> impl DoubleEndedIterator<Item = I> + ExactSizeIterator + Clone + 'static {
+        // Allow the optimizer to elide the bounds checking when creating each index.
+        let _ = I::new(self.len());
         (0..self.len()).map(|n| I::new(n))
     }
 
@@ -84,6 +89,8 @@ impl<I: Idx, T> IndexSlice<I, T> {
     pub fn iter_enumerated_mut(
         &mut self,
     ) -> impl DoubleEndedIterator<Item = (I, &mut T)> + ExactSizeIterator {
+        // Allow the optimizer to elide the bounds checking when creating each index.
+        let _ = I::new(self.len());
         self.raw.iter_mut().enumerate().map(|(n, t)| (I::new(n), t))
     }
 
@@ -115,32 +122,36 @@ impl<I: Idx, T> IndexSlice<I, T> {
 
     /// Returns mutable references to two distinct elements, `a` and `b`.
     ///
-    /// Panics if `a == b`.
+    /// Panics if `a == b` or if some of them are out of bounds.
     #[inline]
     pub fn pick2_mut(&mut self, a: I, b: I) -> (&mut T, &mut T) {
         let (ai, bi) = (a.index(), b.index());
-        assert!(ai != bi);
-
-        if ai < bi {
-            let (c1, c2) = self.raw.split_at_mut(bi);
-            (&mut c1[ai], &mut c2[0])
-        } else {
-            let (c2, c1) = self.pick2_mut(b, a);
-            (c1, c2)
+
+        match self.raw.get_disjoint_mut([ai, bi]) {
+            Ok([a, b]) => (a, b),
+            Err(OverlappingIndices) => panic!("Indices {ai:?} and {bi:?} are not disjoint!"),
+            Err(IndexOutOfBounds) => {
+                panic!("Some indices among ({ai:?}, {bi:?}) are out of bounds")
+            }
         }
     }
 
     /// Returns mutable references to three distinct elements.
     ///
-    /// Panics if the elements are not distinct.
+    /// Panics if the elements are not distinct or if some of them are out of bounds.
     #[inline]
     pub fn pick3_mut(&mut self, a: I, b: I, c: I) -> (&mut T, &mut T, &mut T) {
         let (ai, bi, ci) = (a.index(), b.index(), c.index());
-        assert!(ai != bi && bi != ci && ci != ai);
-        let len = self.raw.len();
-        assert!(ai < len && bi < len && ci < len);
-        let ptr = self.raw.as_mut_ptr();
-        unsafe { (&mut *ptr.add(ai), &mut *ptr.add(bi), &mut *ptr.add(ci)) }
+
+        match self.raw.get_disjoint_mut([ai, bi, ci]) {
+            Ok([a, b, c]) => (a, b, c),
+            Err(OverlappingIndices) => {
+                panic!("Indices {ai:?}, {bi:?} and {ci:?} are not disjoint!")
+            }
+            Err(IndexOutOfBounds) => {
+                panic!("Some indices among ({ai:?}, {bi:?}, {ci:?}) are out of bounds")
+            }
+        }
     }
 
     #[inline]
diff --git a/compiler/rustc_index/src/vec.rs b/compiler/rustc_index/src/vec.rs
index 7f3f3ead5f2..13f0dda180b 100644
--- a/compiler/rustc_index/src/vec.rs
+++ b/compiler/rustc_index/src/vec.rs
@@ -93,6 +93,8 @@ impl<I: Idx, T> IndexVec<I, T> {
     /// be allocated only once, with a capacity of at least `n`.)
     #[inline]
     pub fn from_fn_n(func: impl FnMut(I) -> T, n: usize) -> Self {
+        // Allow the optimizer to elide the bounds checking when creating each index.
+        let _ = I::new(n);
         IndexVec::from_raw((0..n).map(I::new).map(func).collect())
     }
 
@@ -128,6 +130,8 @@ impl<I: Idx, T> IndexVec<I, T> {
     pub fn into_iter_enumerated(
         self,
     ) -> impl DoubleEndedIterator<Item = (I, T)> + ExactSizeIterator {
+        // Allow the optimizer to elide the bounds checking when creating each index.
+        let _ = I::new(self.len());
         self.raw.into_iter().enumerate().map(|(n, t)| (I::new(n), t))
     }