rustc_utils/mir/
location_or_arg.rs1use either::Either;
2use rustc_middle::mir::{Body, Local, Location, Place};
3
4use crate::PlaceExt;
5
6#[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}