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

1use once_cell::sync::OnceCell;
2
3use rustc_middle::ty::TyCtxt;
4
5use crate::{
6    analysis::dataflow::*,
7    check::opt::OptCheck,
8    helpers::def_path::DefPath,
9    utils::log::{relative_pos_range, span_to_filename, span_to_line_number, span_to_source_code},
10};
11use annotate_snippets::{Level, Renderer, Snippet};
12use rustc_span::Span;
13
14static DEFPATHS: OnceCell<DefPaths> = OnceCell::new();
15
16struct DefPaths {
17    flat_map: DefPath,
18    flatten: DefPath,
19    collect: DefPath,
20}
21
22impl DefPaths {
23    fn new(tcx: &TyCtxt<'_>) -> Self {
24        Self {
25            flat_map: DefPath::new("std::iter::Iterator::flat_map", tcx),
26            flatten: DefPath::new("std::iter::Iterator::flatten", tcx),
27            collect: DefPath::new("std::iter::Iterator::collect", tcx),
28        }
29    }
30}
31
32pub struct FlattenCollectCheck {
33    record: Vec<Span>,
34}
35
36fn is_flatten_node(node: &GraphNode) -> bool {
37    let def_paths = &DEFPATHS.get().unwrap();
38    for op in node.ops.iter() {
39        if let NodeOp::Call(def_id) = op {
40            if *def_id == def_paths.flat_map.last_def_id()
41                || *def_id == def_paths.flatten.last_def_id()
42            {
43                return true;
44            }
45        }
46    }
47    false
48}
49
50fn is_collect_node(node: &GraphNode) -> bool {
51    let def_paths = &DEFPATHS.get().unwrap();
52    for op in node.ops.iter() {
53        if let NodeOp::Call(def_id) = op {
54            if *def_id == def_paths.collect.last_def_id() {
55                return true;
56            }
57        }
58    }
59    false
60}
61
62impl OptCheck for FlattenCollectCheck {
63    fn new() -> Self {
64        Self { record: Vec::new() }
65    }
66
67    fn check(&mut self, graph: &Graph, tcx: &TyCtxt) {
68        let _ = &DEFPATHS.get_or_init(|| DefPaths::new(tcx));
69        for node in graph.nodes.iter() {
70            if is_flatten_node(node) {
71                for edge_idx in node.out_edges.iter() {
72                    let dst_idx = graph.edges[*edge_idx].dst;
73                    let dst_node = &graph.nodes[dst_idx];
74                    if is_collect_node(dst_node) {
75                        self.record.push(dst_node.span);
76                    }
77                }
78            }
79        }
80    }
81
82    fn report(&self, graph: &Graph) {
83        for span in self.record.iter() {
84            report_flatten_collect(graph, *span);
85        }
86    }
87
88    fn cnt(&self) -> usize {
89        self.record.len()
90    }
91}
92
93fn report_flatten_collect(graph: &Graph, span: Span) {
94    let code_source = span_to_source_code(graph.span);
95    let filename = span_to_filename(span);
96    let snippet: Snippet<'_> = Snippet::source(&code_source)
97        .line_start(span_to_line_number(graph.span))
98        .origin(&filename)
99        .fold(true)
100        .annotation(
101            Level::Error
102                .span(relative_pos_range(graph.span, span))
103                .label("Flatten then collect."),
104        );
105
106    let message = Level::Error
107        .title("Data collection inefficiency detected")
108        .snippet(snippet)
109        .footer(Level::Help.title("Use extend manually."));
110    let renderer = Renderer::styled();
111    rap_warn!("{}", renderer.render(message));
112}
rapx/check/opt/data_collection/reallocation/flatten_collect.rs

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