Add indexical support, refactor LocationOrArg

crates/rustc_utils/Cargo.toml

@@ -15,6 +15,8 @@ serde = ["dep:serde"]
 test = ["dep:textwrap"]
 graphviz = ["dep:regex"]
 ts-rs = ["dep:ts-rs"]
+indexical = ["dep:indexical"]
+default = ["indexical"]
 
 [dependencies]
 anyhow = "1"
@@ -25,6 +27,7 @@ serde = {version = "1", features = ["derive"], optional = true}
 textwrap = {version = "0.16", optional = true}
 regex = {version = "1", optional = true}
 ts-rs = {version = "6", optional = true}
+indexical = {version = "0.1", default-features = false, features = ["rustc"], optional = true}
 
 [dev-dependencies]
 rustc_utils = {path = ".", features = ["test"]}

crates/rustc_utils/src/mir/location_or_arg.rs

@@ -0,0 +1,86 @@
+use either::Either;
+use rustc_middle::mir::{Body, Local, Location, Place};
+
+use crate::PlaceExt;
+
+/// Used to represent dependencies of places.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum LocationOrArg {
+  Location(Location),
+  Arg(Local),
+}
+
+impl LocationOrArg {
+  pub fn from_place<'tcx>(place: Place<'tcx>, body: &Body<'tcx>) -> Option<Self> {
+    place
+      .is_arg(body)
+      .then_some(LocationOrArg::Arg(place.local))
+  }
+
+  pub fn to_string(self, body: &Body<'_>) -> String {
+    match self {
+      LocationOrArg::Arg(local) => format!("{local:?}"),
+      LocationOrArg::Location(location) => match body.stmt_at(location) {
+        Either::Left(stmt) => format!("{:?}", stmt.kind),
+        Either::Right(terminator) => format!("{:?}", terminator.kind),
+      },
+    }
+  }
+}
+
+impl From<Location> for LocationOrArg {
+  fn from(location: Location) -> Self {
+    LocationOrArg::Location(location)
+  }
+}
+
+impl From<Local> for LocationOrArg {
+  fn from(local: Local) -> Self {
+    LocationOrArg::Arg(local)
+  }
+}
+
+#[cfg(feature = "indexical")]
+pub mod index {
+  use indexical::{
+    impls::RustcIndexSet, index_vec::define_index_type, IndexedDomain, IndexedValue,
+    ToIndex,
+  };
+
+  use super::*;
+
+  define_index_type! {
+    pub struct LocationOrArgIndex = u32;
+  }
+
+  impl IndexedValue for LocationOrArg {
+    type Index = LocationOrArgIndex;
+  }
+
+  pub type LocationOrArgSet = RustcIndexSet<LocationOrArg>;
+  pub type LocationOrArgDomain = IndexedDomain<LocationOrArg>;
+
+  pub struct CustomMarker;
+
+  impl ToIndex<LocationOrArg, CustomMarker> for Location {
+    fn to_index(&self, domain: &IndexedDomain<LocationOrArg>) -> LocationOrArgIndex {
+      LocationOrArg::Location(*self).to_index(domain)
+    }
+  }
+
+  impl ToIndex<LocationOrArg, CustomMarker> for Local {
+    fn to_index(&self, domain: &IndexedDomain<LocationOrArg>) -> LocationOrArgIndex {
+      LocationOrArg::Arg(*self).to_index(domain)
+    }
+  }
+
+  impl rustc_index::Idx for LocationOrArgIndex {
+    fn new(idx: usize) -> Self {
+      LocationOrArgIndex::new(idx)
+    }
+
+    fn index(self) -> usize {
+      self.index()
+    }
+  }
+}

crates/rustc_utils/src/mir/mod.rs

@@ -4,6 +4,7 @@ pub mod adt_def;
 pub mod body;
 pub mod borrowck_facts;
 pub mod control_dependencies;
+pub mod location_or_arg;
 pub mod mutability;
 pub mod operand;
 pub mod place;

crates/rustc_utils/src/source_map/spanner/mir_span.rs

@@ -12,8 +12,8 @@ use rustc_middle::mir::{
 use rustc_span::SpanData;
 use smallvec::{smallvec, SmallVec};
 
-use super::{location_or_arg::LocationOrArg, Spanner};
-use crate::{BodyExt, PlaceExt, SpanExt};
+use super::Spanner;
+use crate::{mir::location_or_arg::LocationOrArg, BodyExt, PlaceExt, SpanExt};
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct MirSpannedPlace<'tcx> {

crates/rustc_utils/src/source_map/spanner/mod.rs

@@ -11,15 +11,14 @@ use rustc_middle::{
 };
 use rustc_span::{source_map::Spanned, Span, SpanData};
 
-pub use self::{hir_span::EnclosingHirSpans, location_or_arg::LocationOrArg};
+pub use self::hir_span::EnclosingHirSpans;
 use self::{
   mir_span::{MirSpanCollector, MirSpannedPlace},
   span_tree::SpanTree,
 };
-use crate::{BodyExt, SpanDataExt, SpanExt};
+use crate::{mir::location_or_arg::LocationOrArg, BodyExt, SpanDataExt, SpanExt};
 
 mod hir_span;
-mod location_or_arg;
 mod mir_span;
 mod span_tree;