rapx/analysis/core/alias_analysis/mop/
graph.rs

1use crate::{
2    analysis::core::alias_analysis::mop::{types::*, MopAAResult},
3    rap_debug,
4    utils::source::*,
5};
6use rustc_data_structures::fx::{FxHashMap, FxHashSet};
7use rustc_middle::{
8    mir::{
9        BasicBlock, Const, Operand, Place, Rvalue, StatementKind, Terminator, TerminatorKind,
10        UnwindAction,
11    },
12    ty::{TyCtxt, TypingEnv},
13};
14use rustc_span::{def_id::DefId, Span};
15use std::{cell::RefCell, cmp::min, vec::Vec};
16
17#[derive(PartialEq, Debug, Copy, Clone)]
18pub enum AssignType {
19    Copy,
20    Move,
21    InitBox,
22    Variant,
23}
24
25//self-defined assignments structure.
26#[derive(Debug, Clone)]
27pub struct Assignment<'tcx> {
28    pub lv: Place<'tcx>,
29    pub rv: Place<'tcx>,
30    pub atype: AssignType,
31    pub span: Span,
32}
33
34impl<'tcx> Assignment<'tcx> {
35    pub fn new(
36        lv: Place<'tcx>,
37        rv: Place<'tcx>,
38        atype: AssignType,
39        span: Span,
40    ) -> Assignment<'tcx> {
41        Assignment {
42            lv,
43            rv,
44            atype,
45            span,
46        }
47    }
48}
49
50/*
51 * Self-defined basicblock structure;
52 * Used both for the original CFG and after SCC.
53 */
54
55#[derive(Debug, Clone)]
56pub struct BlockNode<'tcx> {
57    pub index: usize,
58    pub next: FxHashSet<usize>,
59    pub assignments: Vec<Assignment<'tcx>>,
60    pub calls: Vec<Terminator<'tcx>>,
61    //store the index of the basic blocks as a SCC node.
62    pub scc_sub_blocks: Vec<usize>,
63    //store const values defined in this block, i.e., which id has what value;
64    pub const_value: Vec<(usize, usize)>,
65    //store switch stmts in current block for the path filtering in path-sensitive analysis.
66    pub switch_stmts: Vec<Terminator<'tcx>>,
67    pub modified_value: FxHashSet<usize>,
68    // (SwitchInt target, enum index) -> outside nodes.
69    pub scc_outer: RefCell<Option<FxHashMap<(usize, usize), Vec<usize>>>>,
70}
71
72impl<'tcx> BlockNode<'tcx> {
73    pub fn new(index: usize) -> BlockNode<'tcx> {
74        BlockNode {
75            index,
76            next: FxHashSet::<usize>::default(),
77            assignments: Vec::<Assignment<'tcx>>::new(),
78            calls: Vec::<Terminator<'tcx>>::new(),
79            scc_sub_blocks: Vec::<usize>::new(),
80            const_value: Vec::<(usize, usize)>::new(),
81            switch_stmts: Vec::<Terminator<'tcx>>::new(),
82            modified_value: FxHashSet::<usize>::default(),
83            scc_outer: RefCell::new(None),
84        }
85    }
86
87    pub fn add_next(&mut self, index: usize) {
88        self.next.insert(index);
89    }
90}
91
92#[derive(Debug, Clone)]
93pub struct ValueNode {
94    pub index: usize, // node index
95    pub local: usize, // location?
96    pub need_drop: bool,
97    pub may_drop: bool,
98    pub kind: TyKind,
99    pub father: usize,
100    pub field_id: usize, // the field id of its father node.
101    pub fields: FxHashMap<usize, usize>,
102}
103
104impl ValueNode {
105    pub fn new(index: usize, local: usize, need_drop: bool, may_drop: bool) -> Self {
106        ValueNode {
107            index,
108            local,
109            need_drop,
110            father: local,
111            field_id: usize::MAX,
112            may_drop,
113            kind: TyKind::Adt,
114            fields: FxHashMap::default(),
115        }
116    }
117
118    pub fn is_tuple(&self) -> bool {
119        self.kind == TyKind::Tuple
120    }
121
122    pub fn is_ptr(&self) -> bool {
123        self.kind == TyKind::RawPtr || self.kind == TyKind::Ref
124    }
125
126    pub fn is_ref(&self) -> bool {
127        self.kind == TyKind::Ref
128    }
129
130    pub fn is_corner_case(&self) -> bool {
131        self.kind == TyKind::CornerCase
132    }
133}
134
135pub struct MopGraph<'tcx> {
136    pub def_id: DefId,
137    pub tcx: TyCtxt<'tcx>,
138    pub span: Span,
139    // contains all varibles (including fields) as values.
140    pub values: Vec<ValueNode>,
141    // contains all blocks in the CFG
142    pub blocks: Vec<BlockNode<'tcx>>,
143    pub arg_size: usize,
144    // we shrink a SCC into a node and use a scc node to represent the SCC.
145    pub scc_indices: Vec<usize>,
146    // record the constant value during safedrop checking, i.e., which id has what value.
147    pub constant: FxHashMap<usize, usize>,
148    // contains the return results for inter-procedure analysis.
149    pub ret_alias: MopAAResult,
150    // a threhold to avoid path explosion.
151    pub visit_times: usize,
152    pub alias_set: Vec<usize>,
153    pub child_scc: FxHashMap<
154        usize,
155        (
156            BlockNode<'tcx>,
157            rustc_middle::mir::SwitchTargets,
158            FxHashSet<usize>,
159        ),
160    >,
161    pub disc_map: FxHashMap<usize, usize>,
162    pub terms: Vec<TerminatorKind<'tcx>>,
163}
164
165impl<'tcx> MopGraph<'tcx> {
166    pub fn new(tcx: TyCtxt<'tcx>, def_id: DefId) -> MopGraph<'tcx> {
167        let fn_name = get_fn_name(tcx, def_id);
168        rap_debug!("Building a new MoP graph for: {:?}", fn_name);
169        // handle variables
170        let body = tcx.optimized_mir(def_id);
171        let locals = &body.local_decls;
172        let arg_size = body.arg_count;
173        let mut values = Vec::<ValueNode>::new();
174        let mut alias = Vec::<usize>::new();
175        let ty_env = TypingEnv::post_analysis(tcx, def_id);
176        for (local, local_decl) in locals.iter_enumerated() {
177            let need_drop = local_decl.ty.needs_drop(tcx, ty_env); // the type is drop
178            let may_drop = !is_not_drop(tcx, local_decl.ty);
179            let mut node = ValueNode::new(
180                local.as_usize(),
181                local.as_usize(),
182                need_drop,
183                need_drop || may_drop,
184            );
185            node.kind = kind(local_decl.ty);
186            alias.push(values.len());
187            values.push(node);
188        }
189
190        let basicblocks = &body.basic_blocks;
191        let mut blocks = Vec::<BlockNode<'tcx>>::new();
192        let mut scc_indices = Vec::<usize>::new();
193        let mut disc_map = FxHashMap::default();
194        let mut terms = Vec::new();
195
196        // handle each basicblock
197        for i in 0..basicblocks.len() {
198            scc_indices.push(i);
199            let iter = BasicBlock::from(i);
200            let terminator = basicblocks[iter].terminator.clone().unwrap();
201            let mut cur_bb = BlockNode::new(i);
202
203            // handle general statements
204            for stmt in &basicblocks[iter].statements {
205                /* Assign is a tuple defined as Assign(Box<(Place<'tcx>, Rvalue<'tcx>)>) */
206                let span = stmt.source_info.span;
207                if let StatementKind::Assign(ref assign) = stmt.kind {
208                    let lv_local = assign.0.local.as_usize(); // assign.0 is a Place
209                    let lv = assign.0;
210                    cur_bb.modified_value.insert(lv_local);
211                    match assign.1 {
212                        // assign.1 is a Rvalue
213                        Rvalue::Use(ref x) => {
214                            match x {
215                                Operand::Copy(ref p) => {
216                                    let rv_local = p.local.as_usize();
217                                    if values[lv_local].may_drop && values[rv_local].may_drop {
218                                        let rv = *p;
219                                        let assign =
220                                            Assignment::new(lv, rv, AssignType::Copy, span);
221                                        cur_bb.assignments.push(assign);
222                                    }
223                                }
224                                Operand::Move(ref p) => {
225                                    let rv_local = p.local.as_usize();
226                                    if values[lv_local].may_drop && values[rv_local].may_drop {
227                                        let rv = *p;
228                                        let assign =
229                                            Assignment::new(lv, rv, AssignType::Move, span);
230                                        cur_bb.assignments.push(assign);
231                                    }
232                                }
233                                Operand::Constant(ref constant) => {
234                                    /* We should check the correctness due to the update of rustc
235                                     * https://doc.rust-lang.org/beta/nightly-rustc/rustc_middle/mir/enum.Const.html
236                                     */
237                                    match constant.const_ {
238                                        Const::Ty(_ty, const_value) => {
239                                            if let Some(val) = const_value.try_to_target_usize(tcx)
240                                            {
241                                                cur_bb.const_value.push((lv_local, val as usize));
242                                            }
243                                        }
244                                        Const::Unevaluated(_const_value, _ty) => {}
245                                        Const::Val(const_value, _ty) => {
246                                            //if let Some(val) = const_value.try_to_target_usize(tcx) {
247                                            if let Some(scalar) = const_value.try_to_scalar_int() {
248                                                let val = scalar.to_uint(scalar.size());
249                                                cur_bb.const_value.push((lv_local, val as usize));
250                                            }
251                                        }
252                                    }
253                                }
254                            }
255                        }
256                        Rvalue::Ref(_, _, ref p)
257                        | Rvalue::RawPtr(_, ref p)
258                        | Rvalue::CopyForDeref(ref p) => {
259                            let rv_local = p.local.as_usize();
260                            if values[lv_local].may_drop && values[rv_local].may_drop {
261                                let rv = *p;
262                                let assign = Assignment::new(lv, rv, AssignType::Copy, span);
263                                cur_bb.assignments.push(assign);
264                            }
265                        }
266                        Rvalue::ShallowInitBox(ref x, _) => {
267                            /*
268                             * Original ShllowInitBox is a two-level pointer: lvl0 -> lvl1 -> lvl2
269                             * Since our alias analysis does not consider multi-level pointer,
270                             * We simplify it as: lvl0
271                             */
272                            if !values[lv_local].fields.contains_key(&0) {
273                                let mut lvl0 = ValueNode::new(values.len(), lv_local, false, true);
274                                lvl0.field_id = 0;
275                                values[lv_local].fields.insert(0, lvl0.index);
276                                alias.push(values.len());
277                                values.push(lvl0);
278                            }
279                            match x {
280                                Operand::Copy(ref p) | Operand::Move(ref p) => {
281                                    let rv_local = p.local.as_usize();
282                                    if values[lv_local].may_drop && values[rv_local].may_drop {
283                                        let rv = *p;
284                                        let assign =
285                                            Assignment::new(lv, rv, AssignType::InitBox, span);
286                                        cur_bb.assignments.push(assign);
287                                    }
288                                }
289                                Operand::Constant(_) => {}
290                            }
291                        }
292                        Rvalue::Cast(_, ref x, _) => match x {
293                            Operand::Copy(ref p) => {
294                                let rv_local = p.local.as_usize();
295                                if values[lv_local].may_drop && values[rv_local].may_drop {
296                                    let rv = *p;
297                                    let assign = Assignment::new(lv, rv, AssignType::Copy, span);
298                                    cur_bb.assignments.push(assign);
299                                }
300                            }
301                            Operand::Move(ref p) => {
302                                let rv_local = p.local.as_usize();
303                                if values[lv_local].may_drop && values[rv_local].may_drop {
304                                    let rv = *p;
305                                    let assign = Assignment::new(lv, rv, AssignType::Move, span);
306                                    cur_bb.assignments.push(assign);
307                                }
308                            }
309                            Operand::Constant(_) => {}
310                        },
311                        Rvalue::Aggregate(_, ref x) => {
312                            for each_x in x {
313                                match each_x {
314                                    Operand::Copy(ref p) | Operand::Move(ref p) => {
315                                        let rv_local = p.local.as_usize();
316                                        if values[lv_local].may_drop && values[rv_local].may_drop {
317                                            let rv = *p;
318                                            let assign =
319                                                Assignment::new(lv, rv, AssignType::Copy, span);
320                                            cur_bb.assignments.push(assign);
321                                        }
322                                    }
323                                    Operand::Constant(_) => {}
324                                }
325                            }
326                        }
327                        Rvalue::Discriminant(ref p) => {
328                            let rv = *p;
329                            let assign = Assignment::new(lv, rv, AssignType::Variant, span);
330                            cur_bb.assignments.push(assign);
331                            disc_map.insert(lv_local, p.local.as_usize());
332                        }
333                        _ => {}
334                    }
335                }
336            }
337
338            terms.push(terminator.kind.clone());
339            // handle terminator statements
340            match terminator.kind {
341                TerminatorKind::Goto { ref target } => {
342                    cur_bb.add_next(target.as_usize());
343                }
344                TerminatorKind::SwitchInt {
345                    discr: _,
346                    ref targets,
347                } => {
348                    cur_bb.switch_stmts.push(terminator.clone());
349                    for (_, ref target) in targets.iter() {
350                        cur_bb.add_next(target.as_usize());
351                    }
352                    cur_bb.add_next(targets.otherwise().as_usize());
353                }
354                TerminatorKind::UnwindResume
355                | TerminatorKind::Return
356                | TerminatorKind::UnwindTerminate(_)
357                | TerminatorKind::CoroutineDrop
358                | TerminatorKind::Unreachable => {}
359                TerminatorKind::Drop {
360                    place: _,
361                    ref target,
362                    ref unwind,
363                    replace: _,
364                    drop: _,
365                    async_fut: _,
366                } => {
367                    cur_bb.add_next(target.as_usize());
368                    if let UnwindAction::Cleanup(target) = unwind {
369                        cur_bb.add_next(target.as_usize());
370                    }
371                }
372                TerminatorKind::Call {
373                    func: _,
374                    args: _,
375                    destination: _,
376                    ref target,
377                    ref unwind,
378                    call_source: _,
379                    fn_span: _,
380                } => {
381                    if let Some(tt) = target {
382                        cur_bb.add_next(tt.as_usize());
383                    }
384                    if let UnwindAction::Cleanup(tt) = unwind {
385                        cur_bb.add_next(tt.as_usize());
386                    }
387                    cur_bb.calls.push(terminator.clone());
388                }
389                TerminatorKind::TailCall { .. } => todo!(),
390                TerminatorKind::Assert {
391                    cond: _,
392                    expected: _,
393                    msg: _,
394                    ref target,
395                    ref unwind,
396                } => {
397                    cur_bb.add_next(target.as_usize());
398                    if let UnwindAction::Cleanup(target) = unwind {
399                        cur_bb.add_next(target.as_usize());
400                    }
401                }
402                TerminatorKind::Yield {
403                    value: _,
404                    ref resume,
405                    resume_arg: _,
406                    ref drop,
407                } => {
408                    cur_bb.add_next(resume.as_usize());
409                    if let Some(target) = drop {
410                        cur_bb.add_next(target.as_usize());
411                    }
412                }
413                TerminatorKind::FalseEdge {
414                    ref real_target,
415                    imaginary_target: _,
416                } => {
417                    cur_bb.add_next(real_target.as_usize());
418                }
419                TerminatorKind::FalseUnwind {
420                    ref real_target,
421                    unwind: _,
422                } => {
423                    cur_bb.add_next(real_target.as_usize());
424                }
425                TerminatorKind::InlineAsm {
426                    template: _,
427                    operands: _,
428                    options: _,
429                    line_spans: _,
430                    ref unwind,
431                    targets,
432                    asm_macro: _,
433                } => {
434                    for target in targets {
435                        cur_bb.add_next(target.as_usize());
436                    }
437                    if let UnwindAction::Cleanup(target) = unwind {
438                        cur_bb.add_next(target.as_usize());
439                    }
440                }
441            }
442            blocks.push(cur_bb);
443        }
444
445        MopGraph {
446            def_id,
447            tcx,
448            span: body.span,
449            blocks,
450            values,
451            arg_size,
452            scc_indices,
453            alias_set: alias,
454            constant: FxHashMap::default(),
455            ret_alias: MopAAResult::new(arg_size),
456            visit_times: 0,
457            child_scc: FxHashMap::default(),
458            disc_map,
459            terms,
460        }
461    }
462
463    pub fn tarjan(
464        &mut self,
465        index: usize,
466        stack: &mut Vec<usize>,
467        instack: &mut FxHashSet<usize>,
468        dfn: &mut Vec<usize>,
469        low: &mut Vec<usize>,
470        time: &mut usize,
471    ) {
472        dfn[index] = *time;
473        low[index] = *time;
474        *time += 1;
475        instack.insert(index);
476        stack.push(index);
477        let out_set = self.blocks[index].next.clone();
478        for target in out_set {
479            if dfn[target] == 0 {
480                self.tarjan(target, stack, instack, dfn, low, time);
481                low[index] = min(low[index], low[target]);
482            } else if instack.contains(&target) {
483                low[index] = min(low[index], dfn[target]);
484            }
485        }
486        // generate SCC
487        if dfn[index] == low[index] {
488            let mut modified_set = FxHashSet::<usize>::default();
489            let mut switch_target = Vec::new();
490            let mut scc_block_set = FxHashSet::<usize>::default();
491            let init_block = self.blocks[index].clone();
492            loop {
493                let node = stack.pop().unwrap();
494                self.scc_indices[node] = index;
495                instack.remove(&node);
496                if index == node {
497                    // we have found all nodes of the current scc.
498                    break;
499                }
500                self.blocks[index].scc_sub_blocks.push(node);
501                scc_block_set.insert(node);
502
503                for value in &self.blocks[index].modified_value {
504                    modified_set.insert(*value);
505                }
506                if let Some(target) = self.switch_target(node) {
507                    if !self.blocks[index].switch_stmts.is_empty() {
508                        switch_target.push((target, self.blocks[index].switch_stmts[0].clone()));
509                    }
510                }
511
512                let nexts = self.blocks[node].next.clone();
513                for i in nexts {
514                    self.blocks[index].next.insert(i);
515                }
516            }
517            switch_target.retain(|v| !modified_set.contains(&(v.0)));
518
519            if !switch_target.is_empty() && switch_target.len() == 1 {
520                //let target_index = switch_target[0].0;
521                let target_terminator = switch_target[0].1.clone();
522
523                let TerminatorKind::SwitchInt { discr: _, targets } = target_terminator.kind else {
524                    unreachable!();
525                };
526
527                self.child_scc
528                    .insert(index, (init_block, targets, scc_block_set));
529            }
530            /* remove next nodes which are already in the current SCC */
531            let mut to_remove = Vec::new();
532            for i in self.blocks[index].next.iter() {
533                if self.scc_indices[*i] == index {
534                    to_remove.push(*i);
535                }
536            }
537            for i in to_remove {
538                self.blocks[index].next.remove(&i);
539            }
540            /* To ensure a resonable order of blocks within one SCC,
541             * so that the scc can be directly used for followup analysis without referencing the
542             * original graph.
543             * */
544            self.blocks[index].scc_sub_blocks.reverse();
545        }
546    }
547
548    // handle SCC
549    pub fn solve_scc(&mut self) {
550        let mut stack = Vec::<usize>::new();
551        let mut instack = FxHashSet::<usize>::default();
552        let mut dfn = vec![0_usize; self.blocks.len()];
553        let mut low = vec![0_usize; self.blocks.len()];
554        let mut time = 0;
555        self.tarjan(0, &mut stack, &mut instack, &mut dfn, &mut low, &mut time);
556    }
557
558    pub fn switch_target(&mut self, block_index: usize) -> Option<usize> {
559        let block = &self.blocks[block_index];
560        if block.switch_stmts.is_empty() {
561            return None;
562        }
563
564        let res = if let TerminatorKind::SwitchInt {
565            ref discr,
566            targets: _,
567        } = &block.switch_stmts[0].kind
568        {
569            match discr {
570                Operand::Copy(p) | Operand::Move(p) => {
571                    let place = self.projection(false, p.clone());
572                    Some(place)
573                }
574                _ => None,
575            }
576        } else {
577            None
578        };
579
580        res
581    }
582}
rapx/analysis/core/alias_analysis/mop/graph.rs

rapx/analysis/core/alias_analysis/mop/
graph.rs
