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

1use crate::{
2    analysis::core::alias_analysis::AAResult, 
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,
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 father: usize,
97    pub field_id: usize, // the field id of its father node.
98    pub fields: FxHashMap<usize, usize>,
99}
100
101impl ValueNode {
102    pub fn new(index: usize, local: usize) -> Self {
103        ValueNode {
104            index,
105            local,
106            father: local,
107            field_id: usize::MAX,
108            fields: FxHashMap::default(),
109        }
110    }
111}
112
113pub struct MopGraph<'tcx> {
114    pub def_id: DefId,
115    pub tcx: TyCtxt<'tcx>,
116    pub span: Span,
117    // contains all varibles (including fields) as values.
118    pub values: Vec<ValueNode>,
119    // contains all blocks in the CFG
120    pub blocks: Vec<BlockNode<'tcx>>,
121    pub arg_size: usize,
122    // we shrink a SCC into a node and use a scc node to represent the SCC.
123    pub scc_indices: Vec<usize>,
124    // record the constant value during safedrop checking, i.e., which id has what value.
125    pub constant: FxHashMap<usize, usize>,
126    // contains the return results for inter-procedure analysis.
127    pub ret_alias: AAResult,
128    // a threhold to avoid path explosion.
129    pub visit_times: usize,
130    pub alias_set: Vec<usize>,
131    pub child_scc: FxHashMap<
132        usize,
133        (
134            BlockNode<'tcx>,
135            rustc_middle::mir::SwitchTargets,
136            FxHashSet<usize>,
137        ),
138    >,
139    pub disc_map: FxHashMap<usize, usize>,
140    pub terms: Vec<TerminatorKind<'tcx>>,
141}
142
143impl<'tcx> MopGraph<'tcx> {
144    pub fn new(tcx: TyCtxt<'tcx>, def_id: DefId) -> MopGraph<'tcx> {
145        let fn_name = get_fn_name(tcx, def_id);
146        rap_debug!("Building a new MoP graph for: {:?}", fn_name);
147        // handle variables
148        let body = tcx.optimized_mir(def_id);
149        let locals = &body.local_decls;
150        let arg_size = body.arg_count;
151        let mut values = Vec::<ValueNode>::new();
152        let mut alias = Vec::<usize>::new();
153        for (local, _local_decl) in locals.iter_enumerated() {
154            let node = ValueNode::new(local.as_usize(), local.as_usize());
155            alias.push(values.len());
156            values.push(node);
157        }
158
159        let basicblocks = &body.basic_blocks;
160        let mut blocks = Vec::<BlockNode<'tcx>>::new();
161        let mut scc_indices = Vec::<usize>::new();
162        let mut disc_map = FxHashMap::default();
163        let mut terms = Vec::new();
164
165        // handle each basicblock
166        for i in 0..basicblocks.len() {
167            scc_indices.push(i);
168            let iter = BasicBlock::from(i);
169            let terminator = basicblocks[iter].terminator.clone().unwrap();
170            let mut cur_bb = BlockNode::new(i);
171
172            // handle general statements
173            for stmt in &basicblocks[iter].statements {
174                /* Assign is a tuple defined as Assign(Box<(Place<'tcx>, Rvalue<'tcx>)>) */
175                let span = stmt.source_info.span;
176                if let StatementKind::Assign(ref assign) = stmt.kind {
177                    let lv_local = assign.0.local.as_usize(); // assign.0 is a Place
178                    let lv = assign.0;
179                    cur_bb.modified_value.insert(lv_local);
180                    match assign.1 {
181                        // assign.1 is a Rvalue
182                        Rvalue::Use(ref x) => {
183                            match x {
184                                Operand::Copy(ref p) => {
185                                    let rv = *p;
186                                    let assign = Assignment::new(lv, rv, AssignType::Copy, span);
187                                    cur_bb.assignments.push(assign);
188                                }
189                                Operand::Move(ref p) => {
190                                    let rv = *p;
191                                    let assign = Assignment::new(lv, rv, AssignType::Move, span);
192                                    cur_bb.assignments.push(assign);
193                                }
194                                Operand::Constant(ref constant) => {
195                                    /* We should check the correctness due to the update of rustc
196                                     * https://doc.rust-lang.org/beta/nightly-rustc/rustc_middle/mir/enum.Const.html
197                                     */
198                                    match constant.const_ {
199                                        Const::Ty(_ty, const_value) => {
200                                            if let Some(val) = const_value.try_to_target_usize(tcx)
201                                            {
202                                                cur_bb.const_value.push((lv_local, val as usize));
203                                            }
204                                        }
205                                        Const::Unevaluated(_const_value, _ty) => {}
206                                        Const::Val(const_value, _ty) => {
207                                            //if let Some(val) = const_value.try_to_target_usize(tcx) {
208                                            if let Some(scalar) = const_value.try_to_scalar_int() {
209                                                let val = scalar.to_uint(scalar.size());
210                                                cur_bb.const_value.push((lv_local, val as usize));
211                                            }
212                                        }
213                                    }
214                                }
215                            }
216                        }
217                        Rvalue::Ref(_, _, ref p)
218                        | Rvalue::RawPtr(_, ref p)
219                        | Rvalue::CopyForDeref(ref p) => {
220                            let rv = *p;
221                            let assign = Assignment::new(lv, rv, AssignType::Copy, span);
222                            cur_bb.assignments.push(assign);
223                        }
224                        Rvalue::ShallowInitBox(ref x, _) => {
225                            /*
226                             * Original ShllowInitBox is a two-level pointer: lvl0 -> lvl1 -> lvl2
227                             * Since our alias analysis does not consider multi-level pointer,
228                             * We simplify it as: lvl0
229                             */
230                            if !values[lv_local].fields.contains_key(&0) {
231                                let mut lvl0 = ValueNode::new(values.len(), lv_local);
232                                lvl0.field_id = 0;
233                                values[lv_local].fields.insert(0, lvl0.index);
234                                alias.push(values.len());
235                                values.push(lvl0);
236                            }
237                            match x {
238                                Operand::Copy(ref p) | Operand::Move(ref p) => {
239                                    let rv = *p;
240                                    let assign = Assignment::new(lv, rv, AssignType::InitBox, span);
241                                    cur_bb.assignments.push(assign);
242                                }
243                                Operand::Constant(_) => {}
244                            }
245                        }
246                        Rvalue::Cast(_, ref x, _) => match x {
247                            Operand::Copy(ref p) => {
248                                let rv = *p;
249                                let assign = Assignment::new(lv, rv, AssignType::Copy, span);
250                                cur_bb.assignments.push(assign);
251                            }
252                            Operand::Move(ref p) => {
253                                let rv = *p;
254                                let assign = Assignment::new(lv, rv, AssignType::Move, span);
255                                cur_bb.assignments.push(assign);
256                            }
257                            Operand::Constant(_) => {}
258                        },
259                        Rvalue::Aggregate(_, ref x) => {
260                            for each_x in x {
261                                match each_x {
262                                    Operand::Copy(ref p) | Operand::Move(ref p) => {
263                                        let rv = *p;
264                                        let assign =
265                                            Assignment::new(lv, rv, AssignType::Copy, span);
266                                        cur_bb.assignments.push(assign);
267                                    }
268                                    Operand::Constant(_) => {}
269                                }
270                            }
271                        }
272                        Rvalue::Discriminant(ref p) => {
273                            let rv = *p;
274                            let assign = Assignment::new(lv, rv, AssignType::Variant, span);
275                            cur_bb.assignments.push(assign);
276                            disc_map.insert(lv_local, p.local.as_usize());
277                        }
278                        _ => {}
279                    }
280                }
281            }
282
283            terms.push(terminator.kind.clone());
284            // handle terminator statements
285            match terminator.kind {
286                TerminatorKind::Goto { ref target } => {
287                    cur_bb.add_next(target.as_usize());
288                }
289                TerminatorKind::SwitchInt {
290                    discr: _,
291                    ref targets,
292                } => {
293                    cur_bb.switch_stmts.push(terminator.clone());
294                    for (_, ref target) in targets.iter() {
295                        cur_bb.add_next(target.as_usize());
296                    }
297                    cur_bb.add_next(targets.otherwise().as_usize());
298                }
299                TerminatorKind::UnwindResume
300                | TerminatorKind::Return
301                | TerminatorKind::UnwindTerminate(_)
302                | TerminatorKind::CoroutineDrop
303                | TerminatorKind::Unreachable => {}
304                TerminatorKind::Drop {
305                    place: _,
306                    ref target,
307                    ref unwind,
308                    replace: _,
309                    drop: _,
310                    async_fut: _,
311                } => {
312                    cur_bb.add_next(target.as_usize());
313                    if let UnwindAction::Cleanup(target) = unwind {
314                        cur_bb.add_next(target.as_usize());
315                    }
316                }
317                TerminatorKind::Call {
318                    func: _,
319                    args: _,
320                    destination: _,
321                    ref target,
322                    ref unwind,
323                    call_source: _,
324                    fn_span: _,
325                } => {
326                    if let Some(tt) = target {
327                        cur_bb.add_next(tt.as_usize());
328                    }
329                    if let UnwindAction::Cleanup(tt) = unwind {
330                        cur_bb.add_next(tt.as_usize());
331                    }
332                    cur_bb.calls.push(terminator.clone());
333                }
334                TerminatorKind::TailCall { .. } => todo!(),
335                TerminatorKind::Assert {
336                    cond: _,
337                    expected: _,
338                    msg: _,
339                    ref target,
340                    ref unwind,
341                } => {
342                    cur_bb.add_next(target.as_usize());
343                    if let UnwindAction::Cleanup(target) = unwind {
344                        cur_bb.add_next(target.as_usize());
345                    }
346                }
347                TerminatorKind::Yield {
348                    value: _,
349                    ref resume,
350                    resume_arg: _,
351                    ref drop,
352                } => {
353                    cur_bb.add_next(resume.as_usize());
354                    if let Some(target) = drop {
355                        cur_bb.add_next(target.as_usize());
356                    }
357                }
358                TerminatorKind::FalseEdge {
359                    ref real_target,
360                    imaginary_target: _,
361                } => {
362                    cur_bb.add_next(real_target.as_usize());
363                }
364                TerminatorKind::FalseUnwind {
365                    ref real_target,
366                    unwind: _,
367                } => {
368                    cur_bb.add_next(real_target.as_usize());
369                }
370                TerminatorKind::InlineAsm {
371                    template: _,
372                    operands: _,
373                    options: _,
374                    line_spans: _,
375                    ref unwind,
376                    targets,
377                    asm_macro: _,
378                } => {
379                    for target in targets {
380                        cur_bb.add_next(target.as_usize());
381                    }
382                    if let UnwindAction::Cleanup(target) = unwind {
383                        cur_bb.add_next(target.as_usize());
384                    }
385                }
386            }
387            blocks.push(cur_bb);
388        }
389
390        MopGraph {
391            def_id,
392            tcx,
393            span: body.span,
394            blocks,
395            values,
396            arg_size,
397            scc_indices,
398            alias_set: alias,
399            constant: FxHashMap::default(),
400            ret_alias: AAResult::new(arg_size),
401            visit_times: 0,
402            child_scc: FxHashMap::default(),
403            disc_map,
404            terms,
405        }
406    }
407
408    pub fn tarjan(
409        &mut self,
410        index: usize,
411        stack: &mut Vec<usize>,
412        instack: &mut FxHashSet<usize>,
413        dfn: &mut Vec<usize>,
414        low: &mut Vec<usize>,
415        time: &mut usize,
416    ) {
417        dfn[index] = *time;
418        low[index] = *time;
419        *time += 1;
420        instack.insert(index);
421        stack.push(index);
422        let out_set = self.blocks[index].next.clone();
423        for target in out_set {
424            if dfn[target] == 0 {
425                self.tarjan(target, stack, instack, dfn, low, time);
426                low[index] = min(low[index], low[target]);
427            } else if instack.contains(&target) {
428                low[index] = min(low[index], dfn[target]);
429            }
430        }
431        // generate SCC
432        if dfn[index] == low[index] {
433            let mut modified_set = FxHashSet::<usize>::default();
434            let mut switch_target = Vec::new();
435            let mut scc_block_set = FxHashSet::<usize>::default();
436            let init_block = self.blocks[index].clone();
437            loop {
438                let node = stack.pop().unwrap();
439                self.scc_indices[node] = index;
440                instack.remove(&node);
441                if index == node {
442                    // we have found all nodes of the current scc.
443                    break;
444                }
445                self.blocks[index].scc_sub_blocks.push(node);
446                scc_block_set.insert(node);
447
448                for value in &self.blocks[index].modified_value {
449                    modified_set.insert(*value);
450                }
451                if let Some(target) = self.switch_target(node) {
452                    if !self.blocks[index].switch_stmts.is_empty() {
453                        switch_target.push((target, self.blocks[index].switch_stmts[0].clone()));
454                    }
455                }
456
457                let nexts = self.blocks[node].next.clone();
458                for i in nexts {
459                    self.blocks[index].next.insert(i);
460                }
461            }
462            switch_target.retain(|v| !modified_set.contains(&(v.0)));
463
464            if !switch_target.is_empty() && switch_target.len() == 1 {
465                //let target_index = switch_target[0].0;
466                let target_terminator = switch_target[0].1.clone();
467
468                let TerminatorKind::SwitchInt { discr: _, targets } = target_terminator.kind else {
469                    unreachable!();
470                };
471
472                self.child_scc
473                    .insert(index, (init_block, targets, scc_block_set));
474            }
475            /* remove next nodes which are already in the current SCC */
476            let mut to_remove = Vec::new();
477            for i in self.blocks[index].next.iter() {
478                if self.scc_indices[*i] == index {
479                    to_remove.push(*i);
480                }
481            }
482            for i in to_remove {
483                self.blocks[index].next.remove(&i);
484            }
485            /* To ensure a resonable order of blocks within one SCC,
486             * so that the scc can be directly used for followup analysis without referencing the
487             * original graph.
488             * */
489            self.blocks[index].scc_sub_blocks.reverse();
490        }
491    }
492
493    // handle SCC
494    pub fn solve_scc(&mut self) {
495        let mut stack = Vec::<usize>::new();
496        let mut instack = FxHashSet::<usize>::default();
497        let mut dfn = vec![0_usize; self.blocks.len()];
498        let mut low = vec![0_usize; self.blocks.len()];
499        let mut time = 0;
500        self.tarjan(0, &mut stack, &mut instack, &mut dfn, &mut low, &mut time);
501    }
502
503    pub fn switch_target(&mut self, block_index: usize) -> Option<usize> {
504        let block = &self.blocks[block_index];
505        if block.switch_stmts.is_empty() {
506            return None;
507        }
508
509        let res = if let TerminatorKind::SwitchInt {
510            ref discr,
511            targets: _,
512        } = &block.switch_stmts[0].kind
513        {
514            match discr {
515                Operand::Copy(p) | Operand::Move(p) => {
516                    let place = self.projection(false, p.clone());
517                    Some(place)
518                }
519                _ => None,
520            }
521        } else {
522            None
523        };
524
525        res
526    }
527}
rapx/analysis/core/alias_analysis/default/graph.rs

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