rapx/analysis/safedrop/
graph.rs

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

rapx/analysis/safedrop/
graph.rs
