1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
use rustc_index::bit_set::BitSet;
use rustc_index::interval::SparseIntervalMatrix;
use rustc_index::{Idx, IndexVec};
use rustc_middle::mir::{self, BasicBlock, Body, Location};

use crate::framework::{visit_results, ResultsVisitable, ResultsVisitor};

/// Maps between a `Location` and a `PointIndex` (and vice versa).
pub struct DenseLocationMap {
    /// For each basic block, how many points are contained within?
    statements_before_block: IndexVec<BasicBlock, usize>,

    /// Map backward from each point to the basic block that it
    /// belongs to.
    basic_blocks: IndexVec<PointIndex, BasicBlock>,

    num_points: usize,
}

impl DenseLocationMap {
    #[inline]
    pub fn new(body: &Body<'_>) -> Self {
        let mut num_points = 0;
        let statements_before_block: IndexVec<BasicBlock, usize> = body
            .basic_blocks
            .iter()
            .map(|block_data| {
                let v = num_points;
                num_points += block_data.statements.len() + 1;
                v
            })
            .collect();

        let mut basic_blocks = IndexVec::with_capacity(num_points);
        for (bb, bb_data) in body.basic_blocks.iter_enumerated() {
            basic_blocks.extend((0..=bb_data.statements.len()).map(|_| bb));
        }

        Self { statements_before_block, basic_blocks, num_points }
    }

    /// Total number of point indices
    #[inline]
    pub fn num_points(&self) -> usize {
        self.num_points
    }

    /// Converts a `Location` into a `PointIndex`. O(1).
    #[inline]
    pub fn point_from_location(&self, location: Location) -> PointIndex {
        let Location { block, statement_index } = location;
        let start_index = self.statements_before_block[block];
        PointIndex::new(start_index + statement_index)
    }

    /// Returns the `PointIndex` for the first statement in the given `BasicBlock`. O(1).
    #[inline]
    pub fn entry_point(&self, block: BasicBlock) -> PointIndex {
        let start_index = self.statements_before_block[block];
        PointIndex::new(start_index)
    }

    /// Return the PointIndex for the block start of this index.
    #[inline]
    pub fn to_block_start(&self, index: PointIndex) -> PointIndex {
        PointIndex::new(self.statements_before_block[self.basic_blocks[index]])
    }

    /// Converts a `PointIndex` back to a location. O(1).
    #[inline]
    pub fn to_location(&self, index: PointIndex) -> Location {
        assert!(index.index() < self.num_points);
        let block = self.basic_blocks[index];
        let start_index = self.statements_before_block[block];
        let statement_index = index.index() - start_index;
        Location { block, statement_index }
    }

    /// Sometimes we get point-indices back from bitsets that may be
    /// out of range (because they round up to the nearest 2^N number
    /// of bits). Use this function to filter such points out if you
    /// like.
    #[inline]
    pub fn point_in_range(&self, index: PointIndex) -> bool {
        index.index() < self.num_points
    }
}

rustc_index::newtype_index! {
    /// A single integer representing a `Location` in the MIR control-flow
    /// graph. Constructed efficiently from `DenseLocationMap`.
    #[orderable]
    #[debug_format = "PointIndex({})"]
    pub struct PointIndex {}
}

/// Add points depending on the result of the given dataflow analysis.
pub fn save_as_intervals<'tcx, N, R>(
    elements: &DenseLocationMap,
    body: &mir::Body<'tcx>,
    mut results: R,
) -> SparseIntervalMatrix<N, PointIndex>
where
    N: Idx,
    R: ResultsVisitable<'tcx, FlowState = BitSet<N>>,
{
    let values = SparseIntervalMatrix::new(elements.num_points());
    let mut visitor = Visitor { elements, values };
    visit_results(
        body,
        body.basic_blocks.reverse_postorder().iter().copied(),
        &mut results,
        &mut visitor,
    );
    visitor.values
}

struct Visitor<'a, N: Idx> {
    elements: &'a DenseLocationMap,
    values: SparseIntervalMatrix<N, PointIndex>,
}

impl<'mir, 'tcx, R, N> ResultsVisitor<'mir, 'tcx, R> for Visitor<'_, N>
where
    N: Idx,
{
    type FlowState = BitSet<N>;

    fn visit_statement_after_primary_effect(
        &mut self,
        _results: &mut R,
        state: &Self::FlowState,
        _statement: &'mir mir::Statement<'tcx>,
        location: Location,
    ) {
        let point = self.elements.point_from_location(location);
        // Use internal iterator manually as it is much more efficient.
        state.iter().for_each(|node| {
            self.values.insert(node, point);
        });
    }

    fn visit_terminator_after_primary_effect(
        &mut self,
        _results: &mut R,
        state: &Self::FlowState,
        _terminator: &'mir mir::Terminator<'tcx>,
        location: Location,
    ) {
        let point = self.elements.point_from_location(location);
        // Use internal iterator manually as it is much more efficient.
        state.iter().for_each(|node| {
            self.values.insert(node, point);
        });
    }
}