rapx/check/opt/data_collection/reallocation/
unreserved_vec.rs

1use std::collections::HashSet;
2
3use crate::{
4    analysis::dataflow::*,
5    check::opt::OptCheck,
6    helpers::def_path::DefPath,
7    utils::log::{relative_pos_range, span_to_filename, span_to_line_number, span_to_source_code},
8};
9use once_cell::sync::OnceCell;
10use rustc_middle::{mir::Local, ty::TyCtxt};
11
12use annotate_snippets::{Level, Renderer, Snippet};
13use rustc_span::Span;
14
15static DEFPATHS: OnceCell<DefPaths> = OnceCell::new();
16
17use super::super::super::LEVEL;
18use rustc_hir::{Expr, ExprKind, intravisit};
19use rustc_middle::ty::TypeckResults;
20
21struct DefPaths {
22    vec_new: DefPath,
23    vec_push: DefPath,
24    vec_with_capacity: DefPath,
25    vec_reserve: DefPath,
26}
27
28impl DefPaths {
29    pub fn new(tcx: &TyCtxt<'_>) -> Self {
30        Self {
31            vec_new: DefPath::new("std::vec::Vec::new", tcx),
32            vec_push: DefPath::new("std::vec::Vec::push", tcx),
33            vec_with_capacity: DefPath::new("std::vec::Vec::with_capacity", tcx),
34            vec_reserve: DefPath::new("std::vec::Vec::reserve", tcx),
35        }
36    }
37}
38
39pub struct LoopFinder<'tcx> {
40    pub typeck_results: &'tcx TypeckResults<'tcx>,
41    pub record: Vec<(Span, Vec<Span>)>,
42}
43
44pub struct PushFinder<'tcx> {
45    typeck_results: &'tcx TypeckResults<'tcx>,
46    record: Vec<Span>,
47}
48
49impl<'tcx> intravisit::Visitor<'tcx> for PushFinder<'tcx> {
50    fn visit_expr(&mut self, ex: &'tcx Expr<'tcx>) {
51        if let ExprKind::MethodCall(.., span) = ex.kind {
52            let def_id = self
53                .typeck_results
54                .type_dependent_def_id(ex.hir_id)
55                .unwrap();
56            let target_def_id = (&DEFPATHS.get().unwrap()).vec_push.last_def_id();
57            if def_id == target_def_id {
58                self.record.push(span);
59            }
60        }
61        intravisit::walk_expr(self, ex);
62    }
63}
64
65impl<'tcx> intravisit::Visitor<'tcx> for LoopFinder<'tcx> {
66    fn visit_expr(&mut self, ex: &'tcx Expr<'tcx>) {
67        if let ExprKind::Loop(block, ..) = ex.kind {
68            let mut push_finder = PushFinder {
69                typeck_results: self.typeck_results,
70                record: Vec::new(),
71            };
72            intravisit::walk_block(&mut push_finder, block);
73            // if !push_finder.record.is_empty() {
74            //     self.record.push((ex.span, push_finder.record));
75            // }
76            if push_finder.record.len() == 1 {
77                // we only use simple cases
78                self.record.push((ex.span, push_finder.record));
79            }
80        }
81        intravisit::walk_expr(self, ex);
82    }
83}
84
85pub struct UnreservedVecCheck {
86    record: Vec<Span>,
87}
88
89fn is_vec_new_node(node: &GraphNode) -> bool {
90    for op in node.ops.iter() {
91        if let NodeOp::Call(def_id) = op {
92            let def_paths = &DEFPATHS.get().unwrap();
93            if *def_id == def_paths.vec_new.last_def_id() {
94                return true;
95            }
96        }
97    }
98    false
99}
100
101fn is_vec_push_node(node: &GraphNode) -> bool {
102    for op in node.ops.iter() {
103        if let NodeOp::Call(def_id) = op {
104            let def_paths = &DEFPATHS.get().unwrap();
105            if *def_id == def_paths.vec_push.last_def_id() {
106                return true;
107            }
108        }
109    }
110    false
111}
112
113fn find_upside_reservation(graph: &Graph, node_idx: Local) -> Option<Local> {
114    let mut reservation_node_idx = None;
115    let def_paths = &DEFPATHS.get().unwrap();
116    let mut node_operator = |graph: &Graph, idx: Local| -> DFSStatus {
117        let node = &graph.nodes[idx];
118        for op in node.ops.iter() {
119            if let NodeOp::Call(def_id) = op {
120                if *def_id == def_paths.vec_with_capacity.last_def_id()
121                    || *def_id == def_paths.vec_reserve.last_def_id()
122                {
123                    reservation_node_idx = Some(idx);
124                    return DFSStatus::Stop;
125                }
126            }
127        }
128        DFSStatus::Continue
129    };
130    let mut seen = HashSet::new();
131    graph.dfs(
132        node_idx,
133        Direction::Upside,
134        &mut node_operator,
135        &mut Graph::equivalent_edge_validator,
136        false,
137        &mut seen,
138    );
139    reservation_node_idx
140}
141
142impl OptCheck for UnreservedVecCheck {
143    fn new() -> Self {
144        Self { record: Vec::new() }
145    }
146
147    fn check(&mut self, graph: &Graph, tcx: &TyCtxt) {
148        let _ = &DEFPATHS.get_or_init(|| DefPaths::new(tcx));
149        let level = LEVEL.lock().unwrap();
150        if *level == 2 {
151            for (node_idx, node) in graph.nodes.iter_enumerated() {
152                if is_vec_new_node(node) {
153                    self.record.push(node.span);
154                }
155                if is_vec_push_node(node) {
156                    if let None = find_upside_reservation(graph, node_idx) {
157                        self.record.push(node.span);
158                    }
159                }
160            }
161        }
162
163        let def_id = graph.def_id;
164        let body = tcx.hir_body_owned_by(def_id.as_local().unwrap());
165        let typeck_results = tcx.typeck(def_id.as_local().unwrap());
166        let mut loop_finder = LoopFinder {
167            typeck_results,
168            record: Vec::new(),
169        };
170        intravisit::walk_body(&mut loop_finder, body);
171        for (_, push_record) in loop_finder.record {
172            for push_span in push_record {
173                if let Some((node_idx, _)) = graph.query_node_by_span(push_span, false) {
174                    if let None = find_upside_reservation(graph, node_idx) {
175                        self.record.push(push_span);
176                    }
177                }
178            }
179        }
180    }
181
182    fn report(&self, graph: &Graph) {
183        for span in self.record.iter() {
184            report_unreserved_vec_bug(graph, *span);
185        }
186    }
187
188    fn cnt(&self) -> usize {
189        self.record.len()
190    }
191}
192
193fn report_unreserved_vec_bug(graph: &Graph, span: Span) {
194    let code_source = span_to_source_code(graph.span);
195    let filename = span_to_filename(span);
196    let snippet: Snippet<'_> = Snippet::source(&code_source)
197        .line_start(span_to_line_number(graph.span))
198        .origin(&filename)
199        .fold(true)
200        .annotation(
201            Level::Error
202                .span(relative_pos_range(graph.span, span))
203                .label("Space unreserved."),
204        );
205    let message = Level::Warning
206        .title("Improper data collection detected")
207        .snippet(snippet)
208        .footer(Level::Help.title("Reserve enough space."));
209    let renderer = Renderer::styled();
210    rap_warn!("{}", renderer.render(message));
211}