rustc_utils/mir/
location_or_arg.rs

1use either::Either;
2use rustc_middle::mir::{Body, Local, Location, Place};
3
4use crate::PlaceExt;
5
6/// Used to represent dependencies of places.
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
8pub enum LocationOrArg {
9  Location(Location),
10  Arg(Local),
11}
12
13impl LocationOrArg {
14  pub fn from_place<'tcx>(place: Place<'tcx>, body: &Body<'tcx>) -> Option<Self> {
15    place
16      .is_arg(body)
17      .then_some(LocationOrArg::Arg(place.local))
18  }
19
20  pub fn to_string(self, body: &Body<'_>) -> String {
21    match self {
22      LocationOrArg::Arg(local) => format!("{local:?}"),
23      LocationOrArg::Location(location) => match body.stmt_at(location) {
24        Either::Left(stmt) => format!("{:?}", stmt.kind),
25        Either::Right(terminator) => format!("{:?}", terminator.kind),
26      },
27    }
28  }
29}
30
31impl From<Location> for LocationOrArg {
32  fn from(location: Location) -> Self {
33    LocationOrArg::Location(location)
34  }
35}
36
37impl From<Local> for LocationOrArg {
38  fn from(local: Local) -> Self {
39    LocationOrArg::Arg(local)
40  }
41}
42
43#[cfg(feature = "indexical")]
44pub mod index {
45  use indexical::{IndexedDomain, ToIndex, bitset::rustc::IndexSet, define_index_type};
46
47  use super::*;
48
49  define_index_type! {
50    pub struct LocationOrArgIndex for LocationOrArg = u32;
51  }
52
53  pub type LocationOrArgSet = IndexSet<LocationOrArg>;
54  pub type LocationOrArgDomain = IndexedDomain<LocationOrArg>;
55
56  pub struct CustomMarker;
57
58  impl ToIndex<LocationOrArg, CustomMarker> for Location {
59    fn to_index(self, domain: &IndexedDomain<LocationOrArg>) -> LocationOrArgIndex {
60      LocationOrArg::Location(self).to_index(domain)
61    }
62  }
63
64  impl ToIndex<LocationOrArg, CustomMarker> for Local {
65    fn to_index(self, domain: &IndexedDomain<LocationOrArg>) -> LocationOrArgIndex {
66      LocationOrArg::Arg(self).to_index(domain)
67    }
68  }
69
70  impl rustc_index::Idx for LocationOrArgIndex {
71    fn new(idx: usize) -> Self {
72      LocationOrArgIndex::new(idx)
73    }
74
75    fn index(self) -> usize {
76      self.index()
77    }
78  }
79}