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

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

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