rapx/analysis/utils/
show_mir.rs

1use colorful::{Color, Colorful};
2use rustc_hir::def_id::DefId;
3use rustc_middle::mir::{
4    BasicBlockData, BasicBlocks, Body, LocalDecl, LocalDecls, Operand, Rvalue, Statement,
5    StatementKind, Terminator, TerminatorKind,
6};
7use rustc_middle::ty::{self, TyCtxt, TyKind};
8
9const NEXT_LINE: &str = "\n";
10const PADDING: &str = "    ";
11const EXPLAIN: &str = " @ ";
12
13// This trait is a wrapper towards std::Display or std::Debug, and is to resolve orphan restrictions.
14pub trait Display {
15    fn display(&self) -> String;
16}
17
18impl<'tcx> Display for Terminator<'tcx> {
19    fn display(&self) -> String {
20        let mut s = String::new();
21        s += &format!("{}{:?}{}", PADDING, self.kind, self.kind.display());
22        s
23    }
24}
25
26impl<'tcx> Display for TerminatorKind<'tcx> {
27    fn display(&self) -> String {
28        let mut s = String::new();
29        s += EXPLAIN;
30        match &self {
31            TerminatorKind::Goto { .. } => s += "Goto",
32            TerminatorKind::SwitchInt { .. } => s += "SwitchInt",
33            TerminatorKind::Return => s += "Return",
34            TerminatorKind::Unreachable => s += "Unreachable",
35            TerminatorKind::Drop { .. } => s += "Drop",
36            TerminatorKind::Assert { .. } => s += "Assert",
37            TerminatorKind::Yield { .. } => s += "Yield",
38            TerminatorKind::FalseEdge { .. } => s += "FalseEdge",
39            TerminatorKind::FalseUnwind { .. } => s += "FalseUnwind",
40            TerminatorKind::InlineAsm { .. } => s += "InlineAsm",
41            TerminatorKind::UnwindResume => s += "UnwindResume",
42            TerminatorKind::UnwindTerminate(..) => s += "UnwindTerminate",
43            TerminatorKind::CoroutineDrop => s += "CoroutineDrop",
44            TerminatorKind::Call { func, .. } => match func {
45                Operand::Constant(constant) => match constant.ty().kind() {
46                    ty::FnDef(id, ..) => {
47                        s += &format!("Call: FnDid: {}", id.index.as_usize()).as_str()
48                    }
49                    _ => (),
50                },
51                _ => (),
52            },
53            TerminatorKind::TailCall { .. } => todo!(),
54        };
55        s
56    }
57}
58
59impl<'tcx> Display for Statement<'tcx> {
60    fn display(&self) -> String {
61        let mut s = String::new();
62        s += &format!("{}{:?}{}", PADDING, self.kind, self.kind.display());
63        s
64    }
65}
66
67impl<'tcx> Display for StatementKind<'tcx> {
68    fn display(&self) -> String {
69        let mut s = String::new();
70        s += EXPLAIN;
71        match &self {
72            StatementKind::Assign(assign) => {
73                s += &format!("{:?}={:?}{}", assign.0, assign.1, assign.1.display());
74            }
75            StatementKind::FakeRead(..) => s += "FakeRead",
76            StatementKind::SetDiscriminant { .. } => s += "SetDiscriminant",
77            StatementKind::StorageLive(..) => s += "StorageLive",
78            StatementKind::StorageDead(..) => s += "StorageDead",
79            StatementKind::Retag(..) => s += "Retag",
80            StatementKind::AscribeUserType(..) => s += "AscribeUserType",
81            StatementKind::Coverage(..) => s += "Coverage",
82            StatementKind::Nop => s += "Nop",
83            StatementKind::PlaceMention(..) => s += "PlaceMention",
84            StatementKind::Intrinsic(..) => s += "Intrinsic",
85            StatementKind::ConstEvalCounter => s += "ConstEvalCounter",
86            _ => todo!(),
87        }
88        s
89    }
90}
91
92impl<'tcx> Display for Rvalue<'tcx> {
93    fn display(&self) -> String {
94        let mut s = String::new();
95        s += EXPLAIN;
96        match self {
97            Rvalue::Use(..) => s += "Use",
98            Rvalue::Repeat(..) => s += "Repeat",
99            Rvalue::Ref(..) => s += "Ref",
100            Rvalue::ThreadLocalRef(..) => s += "ThreadLocalRef",
101            Rvalue::Cast(..) => s += "Cast",
102            Rvalue::BinaryOp(..) => s += "BinaryOp",
103            Rvalue::NullaryOp(..) => s += "NullaryOp",
104            Rvalue::UnaryOp(..) => s += "UnaryOp",
105            Rvalue::Discriminant(..) => s += "Discriminant",
106            Rvalue::Aggregate(..) => s += "Aggregate",
107            Rvalue::ShallowInitBox(..) => s += "ShallowInitBox",
108            Rvalue::CopyForDeref(..) => s += "CopyForDeref",
109            Rvalue::RawPtr(_, _) => s += "RawPtr",
110            _ => todo!(),
111        }
112        s
113    }
114}
115
116impl<'tcx> Display for BasicBlocks<'tcx> {
117    fn display(&self) -> String {
118        let mut s = String::new();
119        for (index, bb) in self.iter().enumerate() {
120            s += &format!(
121                "bb {} {{{}{}}}{}",
122                index,
123                NEXT_LINE,
124                bb.display(),
125                NEXT_LINE
126            );
127        }
128        s
129    }
130}
131
132impl<'tcx> Display for BasicBlockData<'tcx> {
133    fn display(&self) -> String {
134        let mut s = String::new();
135        s += &format!("CleanUp: {}{}", self.is_cleanup, NEXT_LINE);
136        for stmt in self.statements.iter() {
137            s += &format!("{}{}", stmt.display(), NEXT_LINE);
138        }
139        s += &format!(
140            "{}{}",
141            self.terminator.clone().unwrap().display(),
142            NEXT_LINE
143        );
144        s
145    }
146}
147
148impl<'tcx> Display for LocalDecls<'tcx> {
149    fn display(&self) -> String {
150        let mut s = String::new();
151        for (index, ld) in self.iter().enumerate() {
152            s += &format!("_{}: {} {}", index, ld.display(), NEXT_LINE);
153        }
154        s
155    }
156}
157
158impl<'tcx> Display for LocalDecl<'tcx> {
159    fn display(&self) -> String {
160        let mut s = String::new();
161        s += &format!("{}{}", EXPLAIN, self.ty.kind().display());
162        s
163    }
164}
165
166impl<'tcx> Display for Body<'tcx> {
167    fn display(&self) -> String {
168        let mut s = String::new();
169        s += &self.local_decls.display();
170        s += &self.basic_blocks.display();
171        s
172    }
173}
174
175impl<'tcx> Display for TyKind<'tcx> {
176    fn display(&self) -> String {
177        let mut s = String::new();
178        s += &format!("{:?}", self);
179        s
180    }
181}
182
183impl Display for DefId {
184    fn display(&self) -> String {
185        format!("{:?}", self)
186    }
187}
188
189pub struct ShowMir<'tcx> {
190    pub tcx: TyCtxt<'tcx>,
191}
192
193// #[inline(always)]
194pub fn display_mir(did: DefId, body: &Body) {
195    rap_info!("{}", did.display().color(Color::LightRed));
196    rap_info!("{}", body.local_decls.display().color(Color::Green));
197    rap_info!(
198        "{}",
199        body.basic_blocks.display().color(Color::LightGoldenrod2a)
200    );
201}
202
203impl<'tcx> ShowMir<'tcx> {
204    pub fn new(tcx: TyCtxt<'tcx>) -> Self {
205        Self { tcx }
206    }
207
208    pub fn start(&mut self) {
209        rap_info!("Show MIR");
210        let mir_keys = self.tcx.mir_keys(());
211        for each_mir in mir_keys {
212            let def_id = each_mir.to_def_id();
213            let body = self.tcx.instance_mir(ty::InstanceKind::Item(def_id));
214            display_mir(def_id, body);
215        }
216    }
217
218    pub fn start_generate_dot(&mut self) {
219        rap_info!("Generate MIR DOT");
220        std::process::Command::new("mkdir")
221            .args(["MIR_dot_graph"])
222            .output()
223            .expect("Failed to create directory");
224
225        let mir_keys = self.tcx.mir_keys(());
226        for each_mir in mir_keys {
227            let def_id = each_mir.to_def_id();
228            let _ = crate::analysis::utils::fn_info::generate_mir_cfg_dot(self.tcx, def_id, &[]);
229        }
230    }
231}