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::{
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}