rapx/helpers/
show_mir.rs

1use crate::compat::FxHashSet;
2use crate::def_id::is_drop_fn;
3use crate::helpers::draw_dot::render_dot_string;
4use crate::helpers::name::get_cleaned_def_path_name;
5use colorful::{Color, Colorful};
6use rustc_hir::def_id::DefId;
7use rustc_middle::mir::{
8    BasicBlockData, BasicBlocks, Body, LocalDecl, LocalDecls, Operand, Rvalue, Statement,
9    StatementKind, Terminator, TerminatorKind,
10};
11use rustc_middle::ty::{self, TyCtxt, TyKind};
12
13const NEXT_LINE: &str = "\n";
14const PADDING: &str = "    ";
15const EXPLAIN: &str = " @ ";
16
17// This trait is a wrapper towards std::Display or std::Debug, and is to resolve orphan restrictions.
18pub trait Display {
19    fn display(&self) -> String;
20}
21
22impl<'tcx> Display for Terminator<'tcx> {
23    fn display(&self) -> String {
24        let mut s = String::new();
25        s += &format!("{}{:?}{}", PADDING, self.kind, self.kind.display());
26        s
27    }
28}
29
30impl<'tcx> Display for TerminatorKind<'tcx> {
31    fn display(&self) -> String {
32        let mut s = String::new();
33        s += EXPLAIN;
34        match &self {
35            TerminatorKind::Goto { .. } => s += "Goto",
36            TerminatorKind::SwitchInt { .. } => s += "SwitchInt",
37            TerminatorKind::Return => s += "Return",
38            TerminatorKind::Unreachable => s += "Unreachable",
39            TerminatorKind::Drop { .. } => s += "Drop",
40            TerminatorKind::Assert { .. } => s += "Assert",
41            TerminatorKind::Yield { .. } => s += "Yield",
42            TerminatorKind::FalseEdge { .. } => s += "FalseEdge",
43            TerminatorKind::FalseUnwind { .. } => s += "FalseUnwind",
44            TerminatorKind::InlineAsm { .. } => s += "InlineAsm",
45            TerminatorKind::UnwindResume => s += "UnwindResume",
46            TerminatorKind::UnwindTerminate(..) => s += "UnwindTerminate",
47            TerminatorKind::CoroutineDrop => s += "CoroutineDrop",
48            TerminatorKind::Call { func, .. } => match func {
49                Operand::Constant(constant) => match constant.ty().kind() {
50                    ty::FnDef(id, ..) => {
51                        s += &format!("Call: FnDid: {}", id.index.as_usize()).as_str()
52                    }
53                    _ => (),
54                },
55                _ => (),
56            },
57            TerminatorKind::TailCall { .. } => todo!(),
58        };
59        s
60    }
61}
62
63impl<'tcx> Display for Statement<'tcx> {
64    fn display(&self) -> String {
65        let mut s = String::new();
66        s += &format!("{}{:?}{}", PADDING, self.kind, self.kind.display());
67        s
68    }
69}
70
71impl<'tcx> Display for StatementKind<'tcx> {
72    fn display(&self) -> String {
73        let mut s = String::new();
74        s += EXPLAIN;
75        match &self {
76            StatementKind::Assign(assign) => {
77                s += &format!("{:?}={:?}{}", assign.0, assign.1, assign.1.display());
78            }
79            StatementKind::FakeRead(..) => s += "FakeRead",
80            StatementKind::SetDiscriminant { .. } => s += "SetDiscriminant",
81            StatementKind::StorageLive(..) => s += "StorageLive",
82            StatementKind::StorageDead(..) => s += "StorageDead",
83            #[cfg(not(rapx_rustc_ge_198))]
84            StatementKind::Retag(..) => s += "Retag",
85            StatementKind::AscribeUserType(..) => s += "AscribeUserType",
86            StatementKind::Coverage(..) => s += "Coverage",
87            StatementKind::Nop => s += "Nop",
88            StatementKind::PlaceMention(..) => s += "PlaceMention",
89            StatementKind::Intrinsic(..) => s += "Intrinsic",
90            StatementKind::ConstEvalCounter => s += "ConstEvalCounter",
91            _ => todo!(),
92        }
93        s
94    }
95}
96
97impl<'tcx> Display for Rvalue<'tcx> {
98    fn display(&self) -> String {
99        let mut s = String::new();
100        s += EXPLAIN;
101        match self {
102            Rvalue::Use(..) => s += "Use",
103            Rvalue::Repeat(..) => s += "Repeat",
104            Rvalue::Ref(..) => s += "Ref",
105            Rvalue::ThreadLocalRef(..) => s += "ThreadLocalRef",
106            Rvalue::Cast(..) => s += "Cast",
107            Rvalue::BinaryOp(..) => s += "BinaryOp",
108            #[cfg(not(rapx_rustc_ge_196))]
109            Rvalue::NullaryOp(..) => s += "NullaryOp",
110            Rvalue::UnaryOp(..) => s += "UnaryOp",
111            Rvalue::Discriminant(..) => s += "Discriminant",
112            Rvalue::Aggregate(..) => s += "Aggregate",
113            #[cfg(not(rapx_rustc_ge_196))]
114            Rvalue::ShallowInitBox(..) => s += "ShallowInitBox",
115            Rvalue::CopyForDeref(..) => s += "CopyForDeref",
116            Rvalue::RawPtr(_, _) => s += "RawPtr",
117            _ => todo!(),
118        }
119        s
120    }
121}
122
123impl<'tcx> Display for BasicBlocks<'tcx> {
124    fn display(&self) -> String {
125        let mut s = String::new();
126        for (index, bb) in self.iter().enumerate() {
127            s += &format!(
128                "bb {} {{{}{}}}{}",
129                index,
130                NEXT_LINE,
131                bb.display(),
132                NEXT_LINE
133            );
134        }
135        s
136    }
137}
138
139impl<'tcx> Display for BasicBlockData<'tcx> {
140    fn display(&self) -> String {
141        let mut s = String::new();
142        s += &format!("CleanUp: {}{}", self.is_cleanup, NEXT_LINE);
143        for stmt in self.statements.iter() {
144            s += &format!("{}{}", stmt.display(), NEXT_LINE);
145        }
146        s += &format!(
147            "{}{}",
148            self.terminator.clone().unwrap().display(),
149            NEXT_LINE
150        );
151        s
152    }
153}
154
155impl<'tcx> Display for LocalDecls<'tcx> {
156    fn display(&self) -> String {
157        let mut s = String::new();
158        for (index, ld) in self.iter().enumerate() {
159            s += &format!("_{}: {} {}", index, ld.display(), NEXT_LINE);
160        }
161        s
162    }
163}
164
165impl<'tcx> Display for LocalDecl<'tcx> {
166    fn display(&self) -> String {
167        let mut s = String::new();
168        s += &format!("{}{}", EXPLAIN, self.ty.kind().display());
169        s
170    }
171}
172
173impl<'tcx> Display for Body<'tcx> {
174    fn display(&self) -> String {
175        let mut s = String::new();
176        s += &self.local_decls.display();
177        s += &self.basic_blocks.display();
178        s
179    }
180}
181
182impl<'tcx> Display for TyKind<'tcx> {
183    fn display(&self) -> String {
184        let mut s = String::new();
185        s += &format!("{:?}", self);
186        s
187    }
188}
189
190impl Display for DefId {
191    fn display(&self) -> String {
192        format!("{:?}", self)
193    }
194}
195
196pub struct ShowMir<'tcx> {
197    pub tcx: TyCtxt<'tcx>,
198}
199
200// #[inline(always)]
201pub fn display_mir(did: DefId, body: &Body) {
202    rap_info!("{}", did.display().color(Color::LightRed));
203    rap_info!("{}", body.local_decls.display().color(Color::Green));
204    rap_info!(
205        "{}",
206        body.basic_blocks.display().color(Color::LightGoldenrod2a)
207    );
208}
209
210impl<'tcx> ShowMir<'tcx> {
211    pub fn new(tcx: TyCtxt<'tcx>) -> Self {
212        Self { tcx }
213    }
214
215    pub fn start(&mut self) {
216        rap_info!("Show MIR");
217        let mir_keys = self.tcx.mir_keys(());
218        for each_mir in mir_keys {
219            let def_id = each_mir.to_def_id();
220            let body = self.tcx.instance_mir(ty::InstanceKind::Item(def_id));
221            display_mir(def_id, body);
222        }
223    }
224
225    pub fn start_generate_dot(&mut self) {
226        rap_info!("Generate MIR DOT");
227        std::process::Command::new("mkdir")
228            .args(["MIR_dot_graph"])
229            .output()
230            .expect("Failed to create directory");
231
232        let mir_keys = self.tcx.mir_keys(());
233        for each_mir in mir_keys {
234            let def_id = each_mir.to_def_id();
235            let _ = generate_mir_cfg_dot(self.tcx, def_id, &Vec::new());
236        }
237    }
238}
239
240pub fn generate_mir_cfg_dot<'tcx>(
241    tcx: TyCtxt<'tcx>,
242    def_id: DefId,
243    alias_sets: &Vec<FxHashSet<usize>>,
244) -> Result<(), std::io::Error> {
245    let mir = tcx.optimized_mir(def_id);
246    let mut dot_content = String::new();
247    let alias_info_str = format!("Alias Sets: {:?}", alias_sets);
248
249    dot_content.push_str(&format!(
250        "digraph mir_cfg_{} {{\n",
251        get_cleaned_def_path_name(tcx, def_id)
252    ));
253    dot_content.push_str(&format!(
254        "    label = \"MIR CFG for {}\\n{}\\n\";\n",
255        tcx.def_path_str(def_id),
256        alias_info_str.replace("\"", "\\\"")
257    ));
258    dot_content.push_str("    labelloc = \"t\";\n");
259    dot_content.push_str("    node [shape=box, fontname=\"Courier\", align=\"left\"];\n\n");
260
261    for (bb_index, bb_data) in mir.basic_blocks.iter_enumerated() {
262        let mut lines: Vec<String> = bb_data
263            .statements
264            .iter()
265            .map(|stmt| format!("{:?}", stmt))
266            .collect();
267        let mut node_style = String::new();
268
269        if let Some(terminator) = &bb_data.terminator {
270            let mut is_drop_related = false;
271            match &terminator.kind {
272                TerminatorKind::Drop { .. } => is_drop_related = true,
273                TerminatorKind::Call { func, .. } => {
274                    if let Operand::Constant(c) = func
275                        && let ty::FnDef(def_id, _) = *c.ty().kind()
276                        && is_drop_fn(def_id)
277                    {
278                        is_drop_related = true;
279                    }
280                }
281                _ => {}
282            }
283            if is_drop_related {
284                node_style = ", style=\"filled\", fillcolor=\"#ffdddd\", color=\"red\"".to_string();
285            }
286            lines.push(format!("{:?}", terminator.kind));
287        } else {
288            lines.push("(no terminator)".to_string());
289        }
290
291        let label_content = lines.join("\\l");
292        let node_label = format!("BB{}:\\l{}\\l", bb_index.index(), label_content);
293        dot_content.push_str(&format!(
294            "    BB{} [label=\"{}\"{}];\n",
295            bb_index.index(),
296            node_label.replace("\"", "\\\""),
297            node_style
298        ));
299
300        if let Some(terminator) = &bb_data.terminator {
301            for target in terminator.successors() {
302                dot_content.push_str(&format!(
303                    "    BB{} -> BB{} [label=\"\"];\n",
304                    bb_index.index(),
305                    target.index(),
306                ));
307            }
308        }
309    }
310    dot_content.push_str("}\n");
311    let name = get_cleaned_def_path_name(tcx, def_id);
312    render_dot_string(name, dot_content);
313    Ok(())
314}