flux_driver/collector/
annot_stats.rs

1use std::{collections::HashMap, fs, io};
2
3use flux_config as config;
4use rustc_ast::{DelimArgs, tokenstream::TokenTree};
5use rustc_attr_ir::AttrArgs;
6use rustc_hir::def_id::LOCAL_CRATE;
7use rustc_middle::ty::TyCtxt;
8use rustc_span::{Span, source_map::SourceMap};
9use serde::Serialize;
10
11#[derive(Default, Serialize)]
12pub struct Stats {
13    /// The number of times an attribute appears, e.g., how many times `flux::trusted` appears.
14    attr_count: HashMap<String, usize>,
15    /// The number of lines taken by each type of attribute, e.g., the sum of all lines taken by
16    /// `flux::sig` annotations.
17    loc_per_attr: HashMap<String, usize>,
18    /// This is the sum over `loc_per_attr`
19    loc: usize,
20}
21
22impl Stats {
23    pub fn save(&self, tcx: TyCtxt) -> io::Result<()> {
24        fs::create_dir_all(config::log_dir())?;
25        let crate_name = tcx.crate_name(LOCAL_CRATE);
26        let path = config::log_dir().join(format!("{crate_name}-annots.json"));
27        let mut file = fs::File::create(path)?;
28        serde_json::to_writer(&mut file, self)?;
29        Ok(())
30    }
31
32    pub fn add(&mut self, tcx: TyCtxt, name: &str, args: &AttrArgs) {
33        let sm = tcx.sess.source_map();
34        self.increase_count(name);
35        match args {
36            AttrArgs::Empty => {
37                self.increase_loc(name, 1);
38            }
39            AttrArgs::Delimited(delim_args) => {
40                self.increase_loc(name, count_lines(sm, delim_args));
41            }
42            AttrArgs::Eq { .. } => {}
43        }
44    }
45
46    fn increase_count(&mut self, name: &str) {
47        if let Some(count) = self.attr_count.get_mut(name) {
48            *count += 1;
49        } else {
50            self.attr_count.insert(name.to_string(), 1);
51        }
52    }
53
54    fn increase_loc(&mut self, name: &str, loc: usize) {
55        if let Some(count) = self.loc_per_attr.get_mut(name) {
56            *count += loc;
57        } else {
58            self.loc_per_attr.insert(name.to_string(), loc);
59        }
60        self.loc += loc;
61    }
62}
63
64fn count_lines(sm: &SourceMap, delim_args: &DelimArgs) -> usize {
65    fn go<'a>(
66        sm: &SourceMap,
67        line_set: &mut IntervalSet,
68        tokens: impl Iterator<Item = &'a TokenTree>,
69    ) {
70        for t in tokens {
71            match t {
72                TokenTree::Token(token, _) => {
73                    let info = get_lines(sm, token.span);
74                    line_set.insert(info.start_line, info.end_line);
75                }
76                TokenTree::Delimited(delim_span, _, _, token_stream) => {
77                    let open_info = get_lines(sm, delim_span.open);
78                    let close_info = get_lines(sm, delim_span.close);
79                    line_set.insert(open_info.start_line, open_info.end_line);
80                    line_set.insert(close_info.start_line, close_info.end_line);
81                    go(sm, line_set, token_stream.iter());
82                }
83            }
84        }
85    }
86    let mut line_set = IntervalSet::new();
87    go(sm, &mut line_set, delim_args.tokens.iter());
88    let mut lines = 0;
89    for (start, end) in line_set.iter_intervals() {
90        lines += end - start + 1;
91    }
92    lines
93}
94
95#[expect(dead_code)]
96struct LineInfo {
97    start_line: usize,
98    start_col: usize,
99    end_line: usize,
100    end_col: usize,
101}
102
103fn get_lines(sm: &SourceMap, span: Span) -> LineInfo {
104    let lines = sm.span_to_location_info(span);
105    LineInfo { start_line: lines.1, start_col: lines.2, end_line: lines.3, end_col: lines.4 }
106}
107
108struct IntervalSet {
109    map: Vec<(usize, usize)>,
110}
111
112impl IntervalSet {
113    fn new() -> Self {
114        Self { map: vec![] }
115    }
116
117    fn insert(&mut self, start: usize, end: usize) {
118        if start > end {
119            return;
120        }
121
122        // This condition looks a bit weird, but actually makes sense.
123        //
124        // if r.0 == end + 1, then we're actually adjacent, so we want to
125        // continue to the next range. We're looking here for the first
126        // range which starts *non-adjacently* to our end.
127        let next = self.map.partition_point(|r| r.0 <= end + 1);
128
129        if let Some(right) = next.checked_sub(1) {
130            let (prev_start, prev_end) = self.map[right];
131            if prev_end + 1 >= start {
132                // If the start for the inserted range is adjacent to the
133                // end of the previous, we can extend the previous range.
134                if start < prev_start {
135                    // The first range which ends *non-adjacently* to our start.
136                    // And we can ensure that left <= right.
137
138                    let left = self.map.partition_point(|l| l.1 + 1 < start);
139                    let min = std::cmp::min(self.map[left].0, start);
140                    let max = std::cmp::max(prev_end, end);
141                    self.map[right] = (min, max);
142                    if left != right {
143                        self.map.drain(left..right);
144                    }
145                } else {
146                    // We overlap with the previous range, increase it to
147                    // include us.
148                    //
149                    // Make sure we're actually going to *increase* it though --
150                    // it may be that end is just inside the previously existing
151                    // set.
152                    if end > prev_end {
153                        self.map[right].1 = end;
154                    }
155                }
156            } else {
157                // Otherwise, we don't overlap, so just insert
158                self.map.insert(right + 1, (start, end));
159            }
160        } else if self.map.is_empty() {
161            // Quite common in practice, and expensive to call memcpy
162            // with length zero.
163            self.map.push((start, end));
164        } else {
165            self.map.insert(next, (start, end));
166        }
167    }
168
169    fn iter_intervals(&self) -> impl Iterator<Item = (usize, usize)> {
170        self.map.iter().copied()
171    }
172}
flux_driver/collector/annot_stats.rs

flux_driver/collector/
annot_stats.rs
