Rollup merge of #136610 - Jarcho:range_idx, r=Noratrieb

Allow `IndexSlice` to be indexed by ranges.

This comes with some annoyances as the index type can no longer inferred from indexing expressions. The biggest offender for this is `IndexVec::from_fn_n(|idx| ..., n)` where the index type won't be inferred from the call site or any index expressions inside the closure.

My main use case for this is mapping a `Place` to `Range<Idx>` for value tracking where the range represents all the values the place contains.

12 files changed, 144 insertions, 42 deletions

diff --git a/compiler/rustc_codegen_ssa/src/mir/mod.rs b/compiler/rustc_codegen_ssa/src/mir/mod.rs
index ba28720afec..0758e5d0456 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 1654867739f..b38b09d60eb 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/arg_matrix.rs b/compiler/rustc_hir_typeck/src/fn_ctxt/arg_matrix.rs
index 358bc389bd1..f6298adf2eb 100644
--- a/compiler/rustc_hir_typeck/src/fn_ctxt/arg_matrix.rs
+++ b/compiler/rustc_hir_typeck/src/fn_ctxt/arg_matrix.rs
@@ -22,6 +22,12 @@ impl ExpectedIdx {
     }
 }
 
+impl ProvidedIdx {
+    pub(crate) fn to_expected_idx(self) -> ExpectedIdx {
+        ExpectedIdx::from_u32(self.as_u32())
+    }
+}
+
 // An issue that might be found in the compatibility matrix
 #[derive(Debug)]
 enum Issue {
diff --git a/compiler/rustc_hir_typeck/src/fn_ctxt/checks.rs b/compiler/rustc_hir_typeck/src/fn_ctxt/checks.rs
index 120a43576e8..edd740d8d8f 100644
--- a/compiler/rustc_hir_typeck/src/fn_ctxt/checks.rs
+++ b/compiler/rustc_hir_typeck/src/fn_ctxt/checks.rs
@@ -775,7 +775,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
 
         // First, check if we just need to wrap some arguments in a tuple.
         if let Some((mismatch_idx, terr)) =
-            compatibility_diagonal.iter().enumerate().find_map(|(i, c)| {
+            compatibility_diagonal.iter_enumerated().find_map(|(i, c)| {
                 if let Compatibility::Incompatible(Some(terr)) = c {
                     Some((i, *terr))
                 } else {
@@ -787,24 +787,25 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             // Do we have as many extra provided arguments as the tuple's length?
             // If so, we might have just forgotten to wrap some args in a tuple.
             if let Some(ty::Tuple(tys)) =
-                formal_and_expected_inputs.get(mismatch_idx.into()).map(|tys| tys.1.kind())
+                formal_and_expected_inputs.get(mismatch_idx.to_expected_idx()).map(|tys| tys.1.kind())
                 // If the tuple is unit, we're not actually wrapping any arguments.
                 && !tys.is_empty()
                 && 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_as_tuple = Ty::new_tup_from_iter(
-                    tcx,
-                    provided_arg_tys.iter().map(|(ty, _)| *ty).skip(mismatch_idx).take(tys.len()),
-                );
+                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 =
+                    Ty::new_tup_from_iter(tcx, provided_args_to_tuple.iter().map(|&(ty, _)| ty));
 
                 let mut satisfied = true;
                 // Check if the newly wrapped tuple + rest of the arguments are compatible.
                 for ((_, expected_ty), provided_ty) in std::iter::zip(
-                    formal_and_expected_inputs.iter().skip(mismatch_idx),
-                    [provided_as_tuple].into_iter().chain(
-                        provided_arg_tys.iter().map(|(ty, _)| *ty).skip(mismatch_idx + tys.len()),
-                    ),
+                    formal_and_expected_inputs[mismatch_idx.to_expected_idx()..].iter(),
+                    [provided_as_tuple]
+                        .into_iter()
+                        .chain(provided_args_after_tuple.iter().map(|&(ty, _)| ty)),
                 ) {
                     if !self.may_coerce(provided_ty, *expected_ty) {
                         satisfied = false;
@@ -816,10 +817,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 // Take some care with spans, so we don't suggest wrapping a macro's
                 // innards in parenthesis, for example.
                 if satisfied
-                    && let Some((_, lo)) =
-                        provided_arg_tys.get(ProvidedIdx::from_usize(mismatch_idx))
-                    && let Some((_, hi)) =
-                        provided_arg_tys.get(ProvidedIdx::from_usize(mismatch_idx + tys.len() - 1))
+                    && let &[(_, hi @ lo)] | &[(_, lo), .., (_, hi)] = provided_args_to_tuple
                 {
                     let mut err;
                     if tys.len() == 1 {
@@ -827,9 +825,9 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                         // so don't do anything special here.
                         err = self.err_ctxt().report_and_explain_type_error(
                             mk_trace(
-                                *lo,
-                                formal_and_expected_inputs[mismatch_idx.into()],
-                                provided_arg_tys[mismatch_idx.into()].0,
+                                lo,
+                                formal_and_expected_inputs[mismatch_idx.to_expected_idx()],
+                                provided_arg_tys[mismatch_idx].0,
                             ),
                             self.param_env,
                             terr,
@@ -868,7 +866,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                         callee_ty,
                         call_expr,
                         None,
-                        Some(mismatch_idx),
+                        Some(mismatch_idx.as_usize()),
                         &matched_inputs,
                         &formal_and_expected_inputs,
                         is_method,
@@ -2615,7 +2613,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 }
 
                 let expected_display_type = self
-                    .resolve_vars_if_possible(formal_and_expected_inputs[idx.into()].1)
+                    .resolve_vars_if_possible(formal_and_expected_inputs[idx].1)
                     .sort_string(self.tcx);
                 let label = if idxs_matched == params_with_generics.len() - 1 {
                     format!(
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 4636f294f13..f17ea9e4b59 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`.
 ///
@@ -97,13 +98,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`.
@@ -184,19 +191,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/check_pointers.rs b/compiler/rustc_mir_transform/src/check_pointers.rs
index d693f739180..2d04b621935 100644
--- a/compiler/rustc_mir_transform/src/check_pointers.rs
+++ b/compiler/rustc_mir_transform/src/check_pointers.rs
@@ -71,8 +71,7 @@ pub(crate) fn check_pointers<'tcx, F>(
     // statements/blocks after. Iterating or visiting the MIR in order would require updating
     // our current location after every insertion. By iterating backwards, we dodge this issue:
     // The only Locations that an insertion changes have already been handled.
-    for block in (0..basic_blocks.len()).rev() {
-        let block = block.into();
+    for block in basic_blocks.indices().rev() {
         for statement_index in (0..basic_blocks[block].statements.len()).rev() {
             let location = Location { block, statement_index };
             let statement = &basic_blocks[block].statements[statement_index];
diff --git a/compiler/rustc_mir_transform/src/coverage/counters.rs b/compiler/rustc_mir_transform/src/coverage/counters.rs
index 039f346495b..5568d42ab8f 100644
--- a/compiler/rustc_mir_transform/src/coverage/counters.rs
+++ b/compiler/rustc_mir_transform/src/coverage/counters.rs
@@ -50,7 +50,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 09384defea8..dcc7c5b91d7 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 2f8a3050199..29521d42720 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()?;