Skip to main content

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::{
46    IndexedDomain, ToIndex, bitset::rustc::ArcIndexSet, define_index_type,
47  };
48
49  use super::*;
50
51  define_index_type! {
52    pub struct LocationOrArgIndex for LocationOrArg = u32;
53  }
54
55  pub type LocationOrArgSet = ArcIndexSet<LocationOrArg>;
56  pub type LocationOrArgDomain = IndexedDomain<LocationOrArg>;
57
58  pub struct CustomMarker;
59
60  impl ToIndex<LocationOrArg, CustomMarker> for Location {
61    fn to_index(self, domain: &IndexedDomain<LocationOrArg>) -> LocationOrArgIndex {
62      LocationOrArg::Location(self).to_index(domain)
63    }
64  }
65
66  impl ToIndex<LocationOrArg, CustomMarker> for Local {
67    fn to_index(self, domain: &IndexedDomain<LocationOrArg>) -> LocationOrArgIndex {
68      LocationOrArg::Arg(self).to_index(domain)
69    }
70  }
71
72  impl rustc_index::Idx for LocationOrArgIndex {
73    fn new(idx: usize) -> Self {
74      LocationOrArgIndex::new(idx)
75    }
76
77    fn index(self) -> usize {
78      self.index()
79    }
80  }
81}