rapx/analysis/safedrop/
graph.rs

1use super::bug_records::*;
2use crate::analysis::{
3    core::alias_analysis::default::types::*, core::ownedheap_analysis::OHAResultMap,
4    utils::intrinsic_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                            if id.index.as_usize() == DROP
398                                || id.index.as_usize() == DROP_IN_PLACE
399                                || id.index.as_usize() == MANUALLYDROP
400                                || id.index.as_usize() == DEALLOC
401                            {
402                                cur_bb.drops.push(terminator.clone());
403                            }
404                        }
405                    }
406
407                    if let Some(tt) = target {
408                        cur_bb.add_next(tt.as_usize());
409                    }
410                    if let UnwindAction::Cleanup(tt) = unwind {
411                        cur_bb.add_next(tt.as_usize());
412                    }
413                    cur_bb.calls.push(terminator.clone());
414                }
415                TerminatorKind::TailCall { .. } => todo!(),
416                TerminatorKind::Assert {
417                    cond: _,
418                    expected: _,
419                    msg: _,
420                    ref target,
421                    ref unwind,
422                } => {
423                    cur_bb.add_next(target.as_usize());
424                    if let UnwindAction::Cleanup(target) = unwind {
425                        cur_bb.add_next(target.as_usize());
426                    }
427                }
428                TerminatorKind::Yield {
429                    value: _,
430                    ref resume,
431                    resume_arg: _,
432                    ref drop,
433                } => {
434                    cur_bb.add_next(resume.as_usize());
435                    if let Some(target) = drop {
436                        cur_bb.add_next(target.as_usize());
437                    }
438                }
439                TerminatorKind::FalseEdge {
440                    ref real_target,
441                    imaginary_target: _,
442                } => {
443                    cur_bb.add_next(real_target.as_usize());
444                }
445                TerminatorKind::FalseUnwind {
446                    ref real_target,
447                    unwind: _,
448                } => {
449                    cur_bb.add_next(real_target.as_usize());
450                }
451                TerminatorKind::CoroutineDrop {} => {
452                    // todo
453                }
454                TerminatorKind::InlineAsm {
455                    template: _,
456                    operands: _,
457                    options: _,
458                    line_spans: _,
459                    ref unwind,
460                    targets,
461                    asm_macro: _,
462                } => {
463                    for target in targets {
464                        cur_bb.add_next(target.as_usize());
465                    }
466                    if let UnwindAction::Cleanup(target) = unwind {
467                        cur_bb.add_next(target.as_usize());
468                    }
469                }
470            }
471            blocks.push(cur_bb);
472        }
473
474        SafeDropGraph {
475            def_id,
476            tcx,
477            span: body.span,
478            blocks,
479            values,
480            arg_size,
481            scc_indices,
482            constant: FxHashMap::default(),
483            return_set: FxHashSet::default(),
484            bug_records: BugRecords::new(),
485            visit_times: 0,
486            alias_set: alias,
487            dead_record: dead,
488            adt_owner,
489            child_scc: FxHashMap::default(),
490            disc_map,
491            terms,
492        }
493    }
494
495    pub fn tarjan(
496        &mut self,
497        index: usize,
498        stack: &mut Vec<usize>,
499        instack: &mut FxHashSet<usize>,
500        dfn: &mut Vec<usize>,
501        low: &mut Vec<usize>,
502        time: &mut usize,
503    ) {
504        dfn[index] = *time;
505        low[index] = *time;
506        *time += 1;
507        instack.insert(index);
508        stack.push(index);
509        let out_set = self.blocks[index].next.clone();
510        for target in out_set {
511            if dfn[target] == 0 {
512                self.tarjan(target, stack, instack, dfn, low, time);
513                low[index] = min(low[index], low[target]);
514            } else if instack.contains(&target) {
515                low[index] = min(low[index], dfn[target]);
516            }
517        }
518
519        // generate SCC
520        if dfn[index] == low[index] {
521            let mut modified_set = FxHashSet::<usize>::default();
522            let mut switch_target = Vec::new();
523            let mut scc_block_set = FxHashSet::<usize>::default();
524            let init_block = self.blocks[index].clone();
525            loop {
526                let node = stack.pop().unwrap();
527                self.scc_indices[node] = index;
528                instack.remove(&node);
529                if index == node {
530                    // we have found all nodes of the current scc.
531                    break;
532                }
533                self.blocks[index].scc_sub_blocks.push(node);
534                scc_block_set.insert(node);
535
536                for value in &self.blocks[index].modified_value {
537                    modified_set.insert(*value);
538                }
539                if let Some(target) = self.switch_target(self.tcx, node) {
540                    if !self.blocks[index].switch_stmts.is_empty() {
541                        switch_target.push((target, self.blocks[index].switch_stmts[0].clone()));
542                    }
543                }
544                let nexts = self.blocks[node].next.clone();
545                for i in nexts {
546                    self.blocks[index].next.insert(i);
547                }
548            }
549            switch_target.retain(|v| !modified_set.contains(&(v.0)));
550
551            if !switch_target.is_empty() && switch_target.len() == 1 {
552                //let target_index = switch_target[0].0;
553                let target_terminator = switch_target[0].1.clone();
554
555                let TerminatorKind::SwitchInt { discr: _, targets } = target_terminator.kind else {
556                    unreachable!();
557                };
558
559                self.child_scc
560                    .insert(index, (init_block, targets, scc_block_set));
561            }
562
563            /* remove next nodes which are already in the current SCC */
564            let mut to_remove = Vec::new();
565            for i in self.blocks[index].next.iter() {
566                if self.scc_indices[*i] == index {
567                    to_remove.push(*i);
568                }
569            }
570            for i in to_remove {
571                self.blocks[index].next.remove(&i);
572            }
573            /* To ensure a resonable order of blocks within one SCC,
574             * so that the scc can be directly used for followup analysis without referencing the
575             * original graph.
576             * */
577            self.blocks[index].scc_sub_blocks.reverse();
578        }
579    }
580
581    // handle SCC
582    pub fn solve_scc(&mut self) {
583        let mut stack = Vec::<usize>::new();
584        let mut instack = FxHashSet::<usize>::default();
585        let mut dfn = vec![0; self.blocks.len()];
586        let mut low = vec![0; self.blocks.len()];
587        let mut time = 0;
588        self.tarjan(0, &mut stack, &mut instack, &mut dfn, &mut low, &mut time);
589    }
590
591    pub fn dfs_on_spanning_tree(
592        &self,
593        index: usize,
594        stack: &mut Vec<usize>,
595        paths: &mut Vec<Vec<usize>>,
596    ) {
597        let curr_scc_index = self.scc_indices[index];
598        if self.blocks[curr_scc_index].next.is_empty() {
599            paths.push(stack.to_vec());
600        } else {
601            for child in self.blocks[curr_scc_index].next.iter() {
602                stack.push(*child);
603                self.dfs_on_spanning_tree(*child, stack, paths);
604            }
605        }
606        stack.pop();
607    }
608
609    pub fn get_paths(&self) -> Vec<Vec<usize>> {
610        let mut paths: Vec<Vec<usize>> = Vec::new();
611        let mut stack: Vec<usize> = vec![0];
612        self.dfs_on_spanning_tree(0, &mut stack, &mut paths);
613        paths
614    }
615
616    pub fn switch_target(&mut self, tcx: TyCtxt<'tcx>, block_index: usize) -> Option<usize> {
617        let block = &self.blocks[block_index];
618        if block.switch_stmts.is_empty() {
619            return None;
620        }
621
622        if let TerminatorKind::SwitchInt { discr, .. } = &block.switch_stmts[0].kind {
623            match discr {
624                Operand::Copy(p) | Operand::Move(p) => {
625                    let place = self.projection(tcx, false, *p);
626                    Some(place)
627                }
628                _ => None,
629            }
630        } else {
631            None
632        }
633    }
634}
rapx/analysis/safedrop/graph.rs

rapx/analysis/safedrop/
graph.rs
