Allow SliceIndex to be indexed by ranges.

10 files changed, 121 insertions, 22 deletions

diff --git a/compiler/rustc_codegen_ssa/src/mir/mod.rs b/compiler/rustc_codegen_ssa/src/mir/mod.rs
index 3a896071bc6..3cbec337d6b 100644
--- a/compiler/rustc_codegen_ssa/src/mir/mod.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/mod.rs
@@ -3,7 +3,7 @@ use std::iter;
 use rustc_index::IndexVec;
 use rustc_index::bit_set::DenseBitSet;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
-use rustc_middle::mir::{UnwindTerminateReason, traversal};
+use rustc_middle::mir::{Local, UnwindTerminateReason, traversal};
 use rustc_middle::ty::layout::{FnAbiOf, HasTyCtxt, HasTypingEnv, TyAndLayout};
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt, TypeFoldable, TypeVisitableExt};
 use rustc_middle::{bug, mir, span_bug};
@@ -240,7 +240,7 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     let local_values = {
         let args = arg_local_refs(&mut start_bx, &mut fx, &memory_locals);
 
-        let mut allocate_local = |local| {
+        let mut allocate_local = |local: Local| {
             let decl = &mir.local_decls[local];
             let layout = start_bx.layout_of(fx.monomorphize(decl.ty));
             assert!(!layout.ty.has_erasable_regions());
diff --git a/compiler/rustc_data_structures/src/sorted_map/index_map.rs b/compiler/rustc_data_structures/src/sorted_map/index_map.rs
index e9a5fb51975..278b3a9270c 100644
--- a/compiler/rustc_data_structures/src/sorted_map/index_map.rs
+++ b/compiler/rustc_data_structures/src/sorted_map/index_map.rs
@@ -147,7 +147,7 @@ impl<I: Idx, K: Ord, V> FromIterator<(K, V)> for SortedIndexMultiMap<I, K, V> {
     where
         J: IntoIterator<Item = (K, V)>,
     {
-        let items = IndexVec::from_iter(iter);
+        let items = IndexVec::<I, _>::from_iter(iter);
         let mut idx_sorted_by_item_key: Vec<_> = items.indices().collect();
 
         // `sort_by_key` is stable, so insertion order is preserved for duplicate items.
diff --git a/compiler/rustc_hir_typeck/src/fn_ctxt/checks.rs b/compiler/rustc_hir_typeck/src/fn_ctxt/checks.rs
index 29ee120a675..10185117b9c 100644
--- a/compiler/rustc_hir_typeck/src/fn_ctxt/checks.rs
+++ b/compiler/rustc_hir_typeck/src/fn_ctxt/checks.rs
@@ -791,7 +791,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 && provided_arg_tys.len() == formal_and_expected_inputs.len() - 1 + tys.len()
             {
                 // Wrap up the N provided arguments starting at this position in a tuple.
-                let provided_args_to_tuple = &provided_arg_tys.raw[mismatch_idx.idx()..];
+                let provided_args_to_tuple = &provided_arg_tys[mismatch_idx..];
                 let (provided_args_to_tuple, provided_args_after_tuple) =
                     provided_args_to_tuple.split_at(tys.len());
                 let provided_as_tuple =
diff --git a/compiler/rustc_index/src/idx.rs b/compiler/rustc_index/src/idx.rs
index b85160540d8..33f406e2113 100644
--- a/compiler/rustc_index/src/idx.rs
+++ b/compiler/rustc_index/src/idx.rs
@@ -1,5 +1,7 @@
 use std::fmt::Debug;
 use std::hash::Hash;
+use std::ops;
+use std::slice::SliceIndex;
 
 /// Represents some newtyped `usize` wrapper.
 ///
@@ -43,3 +45,92 @@ impl Idx for u32 {
         self as usize
     }
 }
+
+/// Helper trait for indexing operations with a custom index type.
+pub trait IntoSliceIdx<I, T: ?Sized> {
+    type Output: SliceIndex<T>;
+    fn into_slice_idx(self) -> Self::Output;
+}
+
+impl<I: Idx, T> IntoSliceIdx<I, [T]> for I {
+    type Output = usize;
+    #[inline]
+    fn into_slice_idx(self) -> Self::Output {
+        self.index()
+    }
+}
+
+impl<I, T> IntoSliceIdx<I, [T]> for ops::RangeFull {
+    type Output = ops::RangeFull;
+    #[inline]
+    fn into_slice_idx(self) -> Self::Output {
+        self
+    }
+}
+
+impl<I: Idx, T> IntoSliceIdx<I, [T]> for ops::Range<I> {
+    type Output = ops::Range<usize>;
+    #[inline]
+    fn into_slice_idx(self) -> Self::Output {
+        ops::Range { start: self.start.index(), end: self.end.index() }
+    }
+}
+
+impl<I: Idx, T> IntoSliceIdx<I, [T]> for ops::RangeFrom<I> {
+    type Output = ops::RangeFrom<usize>;
+    #[inline]
+    fn into_slice_idx(self) -> Self::Output {
+        ops::RangeFrom { start: self.start.index() }
+    }
+}
+
+impl<I: Idx, T> IntoSliceIdx<I, [T]> for ops::RangeTo<I> {
+    type Output = ops::RangeTo<usize>;
+    #[inline]
+    fn into_slice_idx(self) -> Self::Output {
+        ..self.end.index()
+    }
+}
+
+impl<I: Idx, T> IntoSliceIdx<I, [T]> for ops::RangeInclusive<I> {
+    type Output = ops::RangeInclusive<usize>;
+    #[inline]
+    fn into_slice_idx(self) -> Self::Output {
+        ops::RangeInclusive::new(self.start().index(), self.end().index())
+    }
+}
+
+impl<I: Idx, T> IntoSliceIdx<I, [T]> for ops::RangeToInclusive<I> {
+    type Output = ops::RangeToInclusive<usize>;
+    #[inline]
+    fn into_slice_idx(self) -> Self::Output {
+        ..=self.end.index()
+    }
+}
+
+#[cfg(feature = "nightly")]
+impl<I: Idx, T> IntoSliceIdx<I, [T]> for core::range::Range<I> {
+    type Output = core::range::Range<usize>;
+    #[inline]
+    fn into_slice_idx(self) -> Self::Output {
+        core::range::Range { start: self.start.index(), end: self.end.index() }
+    }
+}
+
+#[cfg(feature = "nightly")]
+impl<I: Idx, T> IntoSliceIdx<I, [T]> for core::range::RangeFrom<I> {
+    type Output = core::range::RangeFrom<usize>;
+    #[inline]
+    fn into_slice_idx(self) -> Self::Output {
+        core::range::RangeFrom { start: self.start.index() }
+    }
+}
+
+#[cfg(feature = "nightly")]
+impl<I: Idx, T> IntoSliceIdx<I, [T]> for core::range::RangeInclusive<I> {
+    type Output = core::range::RangeInclusive<usize>;
+    #[inline]
+    fn into_slice_idx(self) -> Self::Output {
+        core::range::RangeInclusive { start: self.start.index(), end: self.end.index() }
+    }
+}
diff --git a/compiler/rustc_index/src/lib.rs b/compiler/rustc_index/src/lib.rs
index cae55230b06..3441a5f65c7 100644
--- a/compiler/rustc_index/src/lib.rs
+++ b/compiler/rustc_index/src/lib.rs
@@ -2,6 +2,7 @@
 #![cfg_attr(all(feature = "nightly", test), feature(stmt_expr_attributes))]
 #![cfg_attr(feature = "nightly", allow(internal_features))]
 #![cfg_attr(feature = "nightly", feature(extend_one, step_trait, test))]
+#![cfg_attr(feature = "nightly", feature(new_range_api))]
 #![cfg_attr(feature = "nightly", feature(new_zeroed_alloc))]
 #![warn(unreachable_pub)]
 // tidy-alphabetical-end
@@ -14,7 +15,7 @@ mod idx;
 mod slice;
 mod vec;
 
-pub use idx::Idx;
+pub use idx::{Idx, IntoSliceIdx};
 pub use rustc_index_macros::newtype_index;
 pub use slice::IndexSlice;
 #[doc(no_inline)]
diff --git a/compiler/rustc_index/src/slice.rs b/compiler/rustc_index/src/slice.rs
index 956d32c9a69..cc8080ee697 100644
--- a/compiler/rustc_index/src/slice.rs
+++ b/compiler/rustc_index/src/slice.rs
@@ -1,8 +1,9 @@
+use std::fmt;
 use std::marker::PhantomData;
 use std::ops::{Index, IndexMut};
-use std::{fmt, slice};
+use std::slice::{self, SliceIndex};
 
-use crate::{Idx, IndexVec};
+use crate::{Idx, IndexVec, IntoSliceIdx};
 
 /// A view into contiguous `T`s, indexed by `I` rather than by `usize`.
 ///
@@ -99,13 +100,19 @@ impl<I: Idx, T> IndexSlice<I, T> {
     }
 
     #[inline]
-    pub fn get(&self, index: I) -> Option<&T> {
-        self.raw.get(index.index())
+    pub fn get<R: IntoSliceIdx<I, [T]>>(
+        &self,
+        index: R,
+    ) -> Option<&<R::Output as SliceIndex<[T]>>::Output> {
+        self.raw.get(index.into_slice_idx())
     }
 
     #[inline]
-    pub fn get_mut(&mut self, index: I) -> Option<&mut T> {
-        self.raw.get_mut(index.index())
+    pub fn get_mut<R: IntoSliceIdx<I, [T]>>(
+        &mut self,
+        index: R,
+    ) -> Option<&mut <R::Output as SliceIndex<[T]>>::Output> {
+        self.raw.get_mut(index.into_slice_idx())
     }
 
     /// Returns mutable references to two distinct elements, `a` and `b`.
@@ -186,19 +193,19 @@ impl<I: Idx, T: fmt::Debug> fmt::Debug for IndexSlice<I, T> {
     }
 }
 
-impl<I: Idx, T> Index<I> for IndexSlice<I, T> {
-    type Output = T;
+impl<I: Idx, T, R: IntoSliceIdx<I, [T]>> Index<R> for IndexSlice<I, T> {
+    type Output = <R::Output as SliceIndex<[T]>>::Output;
 
     #[inline]
-    fn index(&self, index: I) -> &T {
-        &self.raw[index.index()]
+    fn index(&self, index: R) -> &Self::Output {
+        &self.raw[index.into_slice_idx()]
     }
 }
 
-impl<I: Idx, T> IndexMut<I> for IndexSlice<I, T> {
+impl<I: Idx, T, R: IntoSliceIdx<I, [T]>> IndexMut<R> for IndexSlice<I, T> {
     #[inline]
-    fn index_mut(&mut self, index: I) -> &mut T {
-        &mut self.raw[index.index()]
+    fn index_mut(&mut self, index: R) -> &mut Self::Output {
+        &mut self.raw[index.into_slice_idx()]
     }
 }
 
diff --git a/compiler/rustc_infer/src/infer/lexical_region_resolve/mod.rs b/compiler/rustc_infer/src/infer/lexical_region_resolve/mod.rs
index e454a88e847..03c4614af13 100644
--- a/compiler/rustc_infer/src/infer/lexical_region_resolve/mod.rs
+++ b/compiler/rustc_infer/src/infer/lexical_region_resolve/mod.rs
@@ -160,7 +160,7 @@ impl<'cx, 'tcx> LexicalResolver<'cx, 'tcx> {
     /// empty region. The `expansion` phase will grow this larger.
     fn construct_var_data(&self) -> LexicalRegionResolutions<'tcx> {
         LexicalRegionResolutions {
-            values: IndexVec::from_fn_n(
+            values: IndexVec::<RegionVid, _>::from_fn_n(
                 |vid| {
                     let vid_universe = self.var_infos[vid].universe;
                     VarValue::Empty(vid_universe)
diff --git a/compiler/rustc_mir_transform/src/coverage/counters.rs b/compiler/rustc_mir_transform/src/coverage/counters.rs
index adb99a75a9e..fddcd341a99 100644
--- a/compiler/rustc_mir_transform/src/coverage/counters.rs
+++ b/compiler/rustc_mir_transform/src/coverage/counters.rs
@@ -47,7 +47,7 @@ fn make_node_flow_priority_list(
     // A "reloop" node has exactly one out-edge, which jumps back to the top
     // of an enclosing loop. Reloop nodes are typically visited more times
     // than loop-exit nodes, so try to avoid giving them physical counters.
-    let is_reloop_node = IndexVec::from_fn_n(
+    let is_reloop_node = IndexVec::<BasicCoverageBlock, _>::from_fn_n(
         |node| match graph.successors[node].as_slice() {
             &[succ] => graph.dominates(succ, node),
             _ => false,
diff --git a/compiler/rustc_mir_transform/src/coverage/graph.rs b/compiler/rustc_mir_transform/src/coverage/graph.rs
index 392b54c8d81..d7db0140f3d 100644
--- a/compiler/rustc_mir_transform/src/coverage/graph.rs
+++ b/compiler/rustc_mir_transform/src/coverage/graph.rs
@@ -42,7 +42,7 @@ impl CoverageGraph {
         // `SwitchInt` to have multiple targets to the same destination `BasicBlock`, so
         // de-duplication is required. This is done without reordering the successors.
 
-        let successors = IndexVec::from_fn_n(
+        let successors = IndexVec::<BasicCoverageBlock, _>::from_fn_n(
             |bcb| {
                 let mut seen_bcbs = FxHashSet::default();
                 let terminator = mir_body[bcbs[bcb].last_bb()].terminator();
diff --git a/compiler/rustc_mir_transform/src/gvn.rs b/compiler/rustc_mir_transform/src/gvn.rs
index d2ffd26f0a0..4128e8b09d0 100644
--- a/compiler/rustc_mir_transform/src/gvn.rs
+++ b/compiler/rustc_mir_transform/src/gvn.rs
@@ -1259,7 +1259,7 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
 
         let layout = self.ecx.layout_of(lhs_ty).ok()?;
 
-        let as_bits = |value| {
+        let as_bits = |value: VnIndex| {
             let constant = self.evaluated[value].as_ref()?;
             if layout.backend_repr.is_scalar() {
                 let scalar = self.ecx.read_scalar(constant).discard_err()?;