miri/data_structures/
range_object_map.rs

1//! Implements a map from allocation ranges to data. This is somewhat similar to RangeMap, but the
2//! ranges and data are discrete and non-splittable -- they represent distinct "objects". An
3//! allocation in the map will always have the same range until explicitly removed
4
5use std::ops::Range;
6
7use rustc_abi::Size;
8use rustc_const_eval::interpret::AllocRange;
9
10#[derive(Clone, Debug)]
11struct Elem<T> {
12    /// The range covered by this element; never empty.
13    range: AllocRange,
14    /// The data stored for this element.
15    data: T,
16}
17
18/// Index of an allocation within the map
19type Position = usize;
20
21#[derive(Clone, Debug)]
22pub struct RangeObjectMap<T> {
23    v: Vec<Elem<T>>,
24}
25
26#[derive(Clone, Debug, PartialEq)]
27pub enum AccessType {
28    /// The access perfectly overlaps (same offset and range) with an existing object.
29    PerfectlyOverlapping(Position),
30    /// The access does not touch any existing object.
31    Empty(Position),
32    /// The access overlaps with one or more existing objects.
33    ImperfectlyOverlapping(Range<Position>),
34}
35
36impl<T> RangeObjectMap<T> {
37    pub fn new() -> Self {
38        Self { v: Vec::new() }
39    }
40
41    /// Finds the position of the allocation containing the given offset. If the offset is not
42    /// in an existing allocation, then returns Err containing the position
43    /// where such allocation should be inserted
44    fn find_offset(&self, offset: Size) -> Result<Position, Position> {
45        self.v.binary_search_by(|elem| -> std::cmp::Ordering {
46            if offset < elem.range.start {
47                // We are too far right (offset is further left).
48                // (`Greater` means that `elem` is greater than the desired target.)
49                std::cmp::Ordering::Greater
50            } else if offset >= elem.range.end() {
51                // We are too far left (offset is further right).
52                std::cmp::Ordering::Less
53            } else {
54                // This is it!
55                std::cmp::Ordering::Equal
56            }
57        })
58    }
59
60    /// Determines whether a given access on `range` overlaps with
61    /// an existing allocation
62    pub fn access_type(&self, range: AllocRange) -> AccessType {
63        match self.find_offset(range.start) {
64            Ok(pos) => {
65                // Start of the range belongs to an existing object, now let's check the overlapping situation
66                let elem = &self.v[pos];
67                // FIXME: derive Eq for AllocRange in rustc
68                if elem.range.start == range.start && elem.range.size == range.size {
69                    // Happy case: perfectly overlapping access
70                    AccessType::PerfectlyOverlapping(pos)
71                } else {
72                    // FIXME: add a last() method to AllocRange that returns the last inclusive offset (end() is exclusive)
73                    let end_pos = match self.find_offset(range.end() - Size::from_bytes(1)) {
74                        // If the end lands in an existing object, add one to get the exclusive position
75                        Ok(inclusive_pos) => inclusive_pos + 1,
76                        Err(exclusive_pos) => exclusive_pos,
77                    };
78
79                    AccessType::ImperfectlyOverlapping(pos..end_pos)
80                }
81            }
82            Err(pos) => {
83                // Start of the range doesn't belong to an existing object
84                match self.find_offset(range.end() - Size::from_bytes(1)) {
85                    // Neither does the end
86                    Err(end_pos) =>
87                        if pos == end_pos {
88                            // There's nothing between the start and the end, so the range thing is empty
89                            AccessType::Empty(pos)
90                        } else {
91                            // Otherwise we have entirely covered an existing object
92                            AccessType::ImperfectlyOverlapping(pos..end_pos)
93                        },
94                    // Otherwise at least part of it overlaps with something else
95                    Ok(end_pos) => AccessType::ImperfectlyOverlapping(pos..end_pos + 1),
96                }
97            }
98        }
99    }
100
101    /// Inserts an object and its occupied range at given position
102    // The Position can be calculated from AllocRange, but the only user of AllocationMap
103    // always calls access_type before calling insert/index/index_mut, and we don't
104    // want to repeat the binary search on each time, so we ask the caller to supply Position
105    pub fn insert_at_pos(&mut self, pos: Position, range: AllocRange, data: T) {
106        assert!(range.size > Size::ZERO);
107
108        self.v.insert(pos, Elem { range, data });
109
110        // Validate position.
111        // If we aren't the first element, then our start must be greater than the previous element's end
112        if pos > 0 {
113            assert!(self.v[pos - 1].range.end() <= range.start);
114        }
115        // If we aren't the last element, then our end must be smaller than next element's start
116        if pos < self.v.len() - 1 {
117            assert!(range.end() <= self.v[pos + 1].range.start);
118        }
119    }
120
121    pub fn remove_pos_range(&mut self, pos_range: Range<Position>) {
122        self.v.drain(pos_range);
123    }
124
125    pub fn remove_from_pos(&mut self, pos: Position) -> T {
126        self.v.remove(pos).data
127    }
128
129    pub fn iter(&self) -> impl Iterator<Item = &T> {
130        self.v.iter().map(|e| &e.data)
131    }
132
133    pub fn range(&self, range: AllocRange) -> impl Iterator<Item = &T> {
134        let begin = self.find_offset(range.start).unwrap_or_else(|i| i);
135        self.v[begin..]
136            .iter()
137            .take_while(move |elem| {
138                debug_assert!(elem.range.end() > range.start);
139                // `elem` overlaps `range` if it starts before the end of `range`.
140                elem.range.start < range.end()
141            })
142            .map(|e| &e.data)
143    }
144
145    pub fn remove_range_if(&mut self, range: AllocRange, pred: impl Fn(&T) -> bool) {
146        let mut idx = self.find_offset(range.start).unwrap_or_else(|i| i);
147        loop {
148            let Some(elem) = self.v.get(idx) else { break };
149            // If we no longer overlap, we can end the loop.
150            if elem.range.start >= range.end() {
151                break;
152            }
153            // See if we should remove this element.
154            if pred(&elem.data) {
155                self.v.remove(idx);
156                // `idx` now refers to the next element.
157            } else {
158                // If not, go look at next element.
159                idx += 1;
160            }
161        }
162    }
163
164    /// Get the element that starts at the given offset.
165    pub fn get_start(&self, offset: Size) -> Option<&T> {
166        let idx = self.find_offset(offset).ok()?;
167        let elem = &self.v[idx];
168        (elem.range.start == offset).then_some(&elem.data)
169    }
170
171    /// Get the element that starts at the given offset.
172    pub fn get_start_mut(&mut self, offset: Size) -> Option<&mut T> {
173        let idx = self.find_offset(offset).ok()?;
174        let elem = &mut self.v[idx];
175        (elem.range.start == offset).then_some(&mut elem.data)
176    }
177
178    pub fn get_pos(&self, pos: Position) -> Option<&T> {
179        self.v.get(pos).map(|e| &e.data)
180    }
181
182    pub fn get_mut_pos(&mut self, pos: Position) -> Option<&mut T> {
183        self.v.get_mut(pos).map(|e| &mut e.data)
184    }
185
186    pub fn is_empty(&self) -> bool {
187        self.v.is_empty()
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use rustc_const_eval::interpret::alloc_range;
194
195    use super::*;
196
197    #[test]
198    fn empty_map() {
199        // FIXME: make Size::from_bytes const
200        let four = Size::from_bytes(4);
201        let map = RangeObjectMap::<()>::new();
202
203        // Correctly tells where we should insert the first element (at position 0)
204        assert_eq!(map.find_offset(Size::from_bytes(3)), Err(0));
205
206        // Correctly tells the access type along with the supposed position
207        assert_eq!(map.access_type(alloc_range(Size::ZERO, four)), AccessType::Empty(0));
208    }
209
210    #[test]
211    #[should_panic]
212    fn no_overlapping_inserts() {
213        let four = Size::from_bytes(4);
214
215        let mut map = RangeObjectMap::<&str>::new();
216
217        // |_|_|_|_|#|#|#|#|_|_|_|_|...
218        //  0 1 2 3 4 5 6 7 8 9 a b c d
219        map.insert_at_pos(0, alloc_range(four, four), "#");
220        // |_|_|_|_|#|#|#|#|_|_|_|_|...
221        //  0 ^ ^ ^ ^ 5 6 7 8 9 a b c d
222        map.insert_at_pos(0, alloc_range(Size::from_bytes(1), four), "@");
223    }
224
225    #[test]
226    fn boundaries() {
227        let four = Size::from_bytes(4);
228
229        let mut map = RangeObjectMap::<&str>::new();
230
231        // |#|#|#|#|_|_|...
232        //  0 1 2 3 4 5
233        map.insert_at_pos(0, alloc_range(Size::ZERO, four), "#");
234        // |#|#|#|#|_|_|...
235        //  0 1 2 3 ^ 5
236        assert_eq!(map.find_offset(four), Err(1));
237        // |#|#|#|#|_|_|_|_|_|...
238        //  0 1 2 3 ^ ^ ^ ^ 8
239        assert_eq!(map.access_type(alloc_range(four, four)), AccessType::Empty(1));
240
241        let eight = Size::from_bytes(8);
242        // |#|#|#|#|_|_|_|_|@|@|@|@|_|_|...
243        //  0 1 2 3 4 5 6 7 8 9 a b c d
244        map.insert_at_pos(1, alloc_range(eight, four), "@");
245        // |#|#|#|#|_|_|_|_|@|@|@|@|_|_|...
246        //  0 1 2 3 4 5 6 ^ 8 9 a b c d
247        assert_eq!(map.find_offset(Size::from_bytes(7)), Err(1));
248        // |#|#|#|#|_|_|_|_|@|@|@|@|_|_|...
249        //  0 1 2 3 ^ ^ ^ ^ 8 9 a b c d
250        assert_eq!(map.access_type(alloc_range(four, four)), AccessType::Empty(1));
251    }
252
253    #[test]
254    fn perfectly_overlapping() {
255        let four = Size::from_bytes(4);
256
257        let mut map = RangeObjectMap::<&str>::new();
258
259        // |#|#|#|#|_|_|...
260        //  0 1 2 3 4 5
261        map.insert_at_pos(0, alloc_range(Size::ZERO, four), "#");
262        // |#|#|#|#|_|_|...
263        //  ^ ^ ^ ^ 4 5
264        assert_eq!(map.find_offset(Size::ZERO), Ok(0));
265        assert_eq!(
266            map.access_type(alloc_range(Size::ZERO, four)),
267            AccessType::PerfectlyOverlapping(0)
268        );
269
270        // |#|#|#|#|@|@|@|@|_|...
271        //  0 1 2 3 4 5 6 7 8
272        map.insert_at_pos(1, alloc_range(four, four), "@");
273        // |#|#|#|#|@|@|@|@|_|...
274        //  0 1 2 3 ^ ^ ^ ^ 8
275        assert_eq!(map.find_offset(four), Ok(1));
276        assert_eq!(map.access_type(alloc_range(four, four)), AccessType::PerfectlyOverlapping(1));
277    }
278
279    #[test]
280    fn straddling() {
281        let four = Size::from_bytes(4);
282
283        let mut map = RangeObjectMap::<&str>::new();
284
285        // |_|_|_|_|#|#|#|#|_|_|_|_|...
286        //  0 1 2 3 4 5 6 7 8 9 a b c d
287        map.insert_at_pos(0, alloc_range(four, four), "#");
288        // |_|_|_|_|#|#|#|#|_|_|_|_|...
289        //  0 1 ^ ^ ^ ^ 6 7 8 9 a b c d
290        assert_eq!(
291            map.access_type(alloc_range(Size::from_bytes(2), four)),
292            AccessType::ImperfectlyOverlapping(0..1)
293        );
294        // |_|_|_|_|#|#|#|#|_|_|_|_|...
295        //  0 1 2 3 4 5 ^ ^ ^ ^ a b c d
296        assert_eq!(
297            map.access_type(alloc_range(Size::from_bytes(6), four)),
298            AccessType::ImperfectlyOverlapping(0..1)
299        );
300        // |_|_|_|_|#|#|#|#|_|_|_|_|...
301        //  0 1 ^ ^ ^ ^ ^ ^ ^ ^ a b c d
302        assert_eq!(
303            map.access_type(alloc_range(Size::from_bytes(2), Size::from_bytes(8))),
304            AccessType::ImperfectlyOverlapping(0..1)
305        );
306
307        // |_|_|_|_|#|#|#|#|_|_|@|@|_|_|...
308        //  0 1 2 3 4 5 6 7 8 9 a b c d
309        map.insert_at_pos(1, alloc_range(Size::from_bytes(10), Size::from_bytes(2)), "@");
310        // |_|_|_|_|#|#|#|#|_|_|@|@|_|_|...
311        //  0 1 2 3 4 5 ^ ^ ^ ^ ^ ^ ^ ^
312        assert_eq!(
313            map.access_type(alloc_range(Size::from_bytes(6), Size::from_bytes(8))),
314            AccessType::ImperfectlyOverlapping(0..2)
315        );
316    }
317}
miri/data_structures/range_object_map.rs

miri/data_structures/
range_object_map.rs
