rapx/analysis/safedrop/
graph.rs

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

rapx/analysis/safedrop/
graph.rs
