src/tools/miri/src/range_map.rs

   1 //! Implements a map from integer indices to data.
   2 //! Rather than storing data for every index, internally, this maps entire ranges to the data.
   3 //! To this end, the APIs all work on ranges, not on individual integers. Ranges are split as
   4 //! necessary (e.g., when [0,5) is first associated with X, and then [1,2) is mutated).
   5 //! Users must not depend on whether a range is coalesced or not, even though this is observable
   6 //! via the iteration APIs.
   7
   8 use std::ops;
   9
  10 use rustc_target::abi::Size;
  11
  12 #[derive(Clone, Debug)]
  13 struct Elem<T> {
  14     /// The range covered by this element; never empty.
  15     range: ops::Range<u64>,
  16     /// The data stored for this element.
  17     data: T,
  18 }
  19 #[derive(Clone, Debug)]
  20 pub struct RangeMap<T> {
  21     v: Vec<Elem<T>>,
  22 }
  23
  24 impl<T> RangeMap<T> {
  25     /// Creates a new `RangeMap` for the given size, and with the given initial value used for
  26     /// the entire range.
  27     #[inline(always)]
  28     pub fn new(size: Size, init: T) -> RangeMap<T> {
  29         let size = size.bytes();
  30         let mut map = RangeMap { v: Vec::new() };
  31         if size > 0 {
  32             map.v.push(Elem { range: 0..size, data: init });
  33         }
  34         map
  35     }
  36
  37     /// Finds the index containing the given offset.
  38     fn find_offset(&self, offset: u64) -> usize {
  39         // We do a binary search.
  40         let mut left = 0usize; // inclusive
  41         let mut right = self.v.len(); // exclusive
  42         loop {
  43             debug_assert!(left < right, "find_offset: offset {offset} is out-of-bounds");
  44             let candidate = left.checked_add(right).unwrap() / 2;
  45             let elem = &self.v[candidate];
  46             if offset < elem.range.start {
  47                 // We are too far right (offset is further left).
  48                 debug_assert!(candidate < right); // we are making progress
  49                 right = candidate;
  50             } else if offset >= elem.range.end {
  51                 // We are too far left (offset is further right).
  52                 debug_assert!(candidate >= left); // we are making progress
  53                 left = candidate + 1;
  54             } else {
  55                 // This is it!
  56                 return candidate;
  57             }
  58         }
  59     }
  60
  61     /// Provides read-only iteration over everything in the given range. This does
  62     /// *not* split items if they overlap with the edges. Do not use this to mutate
  63     /// through interior mutability.
  64     ///
  65     /// The iterator also provides the offset of the given element.
  66     pub fn iter(&self, offset: Size, len: Size) -> impl Iterator<Item = (Size, &T)> {
  67         let offset = offset.bytes();
  68         let len = len.bytes();
  69         // Compute a slice starting with the elements we care about.
  70         let slice: &[Elem<T>] = if len == 0 {
  71             // We just need any empty iterator. We don't even want to
  72             // yield the element that surrounds this position.
  73             &[]
  74         } else {
  75             let first_idx = self.find_offset(offset);
  76             &self.v[first_idx..]
  77         };
  78         // The first offset that is not included any more.
  79         let end = offset + len;
  80         assert!(
  81             end <= self.v.last().unwrap().range.end,
  82             "iterating beyond the bounds of this RangeMap"
  83         );
  84         slice
  85             .iter()
  86             .take_while(move |elem| elem.range.start < end)
  87             .map(|elem| (Size::from_bytes(elem.range.start), &elem.data))
  88     }
  89
  90     pub fn iter_mut_all(&mut self) -> impl Iterator<Item = &mut T> {
  91         self.v.iter_mut().map(|elem| &mut elem.data)
  92     }
  93
  94     pub fn iter_all(&self) -> impl Iterator<Item = (ops::Range<u64>, &T)> {
  95         self.v.iter().map(|elem| (elem.range.clone(), &elem.data))
  96     }
  97
  98     // Splits the element situated at the given `index`, such that the 2nd one starts at offset
  99     // `split_offset`. Do nothing if the element already starts there.
 100     // Returns whether a split was necessary.
 101     fn split_index(&mut self, index: usize, split_offset: u64) -> bool
 102     where
 103         T: Clone,
 104     {
 105         let elem = &mut self.v[index];
 106         if split_offset == elem.range.start || split_offset == elem.range.end {
 107             // Nothing to do.
 108             return false;
 109         }
 110         debug_assert!(
 111             elem.range.contains(&split_offset),
 112             "the `split_offset` is not in the element to be split"
 113         );
 114
 115         // Now we really have to split. Reduce length of first element.
 116         let second_range = split_offset..elem.range.end;
 117         elem.range.end = split_offset;
 118         // Copy the data, and insert second element.
 119         let second = Elem { range: second_range, data: elem.data.clone() };
 120         self.v.insert(index + 1, second);
 121         true
 122     }
 123
 124     /// Provides mutable iteration over everything in the given range. As a side-effect,
 125     /// this will split entries in the map that are only partially hit by the given range,
 126     /// to make sure that when they are mutated, the effect is constrained to the given range.
 127     /// Moreover, this will opportunistically merge neighbouring equal blocks.
 128     ///
 129     /// The iterator also provides the offset of the given element.
 130     pub fn iter_mut(&mut self, offset: Size, len: Size) -> impl Iterator<Item = (Size, &mut T)>
 131     where
 132         T: Clone + PartialEq,
 133     {
 134         let offset = offset.bytes();
 135         let len = len.bytes();
 136         // Compute a slice containing exactly the elements we care about
 137         let slice: &mut [Elem<T>] = if len == 0 {
 138             // We just need any empty iterator. We don't even want to
 139             // yield the element that surrounds this position, nor do
 140             // any splitting.
 141             &mut []
 142         } else {
 143             // Make sure we got a clear beginning
 144             let mut first_idx = self.find_offset(offset);
 145             if self.split_index(first_idx, offset) {
 146                 // The newly created 2nd element is ours
 147                 first_idx += 1;
 148             }
 149             // No more mutation.
 150             let first_idx = first_idx;
 151             // Find our end. Linear scan, but that's ok because the iteration
 152             // is doing the same linear scan anyway -- no increase in complexity.
 153             // We combine this scan with a scan for duplicates that we can merge, to reduce
 154             // the number of elements.
 155             // We stop searching after the first "block" of size 1, to avoid spending excessive
 156             // amounts of time on the merging.
 157             let mut equal_since_idx = first_idx;
 158             // Once we see too many non-mergeable blocks, we stop.
 159             // The initial value is chosen via... magic. Benchmarking and magic.
 160             let mut successful_merge_count = 3usize;
 161             // When the loop is done, this is the first excluded element.
 162             let mut end_idx = first_idx;
 163             loop {
 164                 // Compute if `end` is the last element we need to look at.
 165                 let done = self.v[end_idx].range.end >= offset + len;
 166                 // We definitely need to include `end`, so move the index.
 167                 end_idx += 1;
 168                 debug_assert!(
 169                     done || end_idx < self.v.len(),
 170                     "iter_mut: end-offset {} is out-of-bounds",
 171                     offset + len
 172                 );
 173                 // see if we want to merge everything in `equal_since..end` (exclusive at the end!)
 174                 if successful_merge_count > 0 {
 175                     if done || self.v[end_idx].data != self.v[equal_since_idx].data {
 176                         // Everything in `equal_since..end` was equal. Make them just one element covering
 177                         // the entire range.
 178                         let removed_elems = end_idx - equal_since_idx - 1; // number of elements that we would remove
 179                         if removed_elems > 0 {
 180                             // Adjust the range of the first element to cover all of them.
 181                             let equal_until = self.v[end_idx - 1].range.end; // end of range of last of the equal elements
 182                             self.v[equal_since_idx].range.end = equal_until;
 183                             // Delete the rest of them.
 184                             self.v.splice(equal_since_idx + 1..end_idx, std::iter::empty());
 185                             // Adjust `end_idx` because we made the list shorter.
 186                             end_idx -= removed_elems;
 187                             // Adjust the count for the cutoff.
 188                             successful_merge_count += removed_elems;
 189                         } else {
 190                             // Adjust the count for the cutoff.
 191                             successful_merge_count -= 1;
 192                         }
 193                         // Go on scanning for the next block starting here.
 194                         equal_since_idx = end_idx;
 195                     }
 196                 }
 197                 // Leave loop if this is the last element.
 198                 if done {
 199                     break;
 200                 }
 201             }
 202             // Move to last included instead of first excluded index.
 203             let end_idx = end_idx - 1;
 204             // We need to split the end as well. Even if this performs a
 205             // split, we don't have to adjust our index as we only care about
 206             // the first part of the split.
 207             self.split_index(end_idx, offset + len);
 208             // Now we yield the slice. `end` is inclusive.
 209             &mut self.v[first_idx..=end_idx]
 210         };
 211         slice.iter_mut().map(|elem| (Size::from_bytes(elem.range.start), &mut elem.data))
 212     }
 213 }
 214
 215 #[cfg(test)]
 216 mod tests {
 217     use super::*;
 218
 219     /// Query the map at every offset in the range and collect the results.
 220     fn to_vec<T: Copy>(map: &RangeMap<T>, offset: u64, len: u64) -> Vec<T> {
 221         (offset..offset + len)
 222             .into_iter()
 223             .map(|i| {
 224                 map.iter(Size::from_bytes(i), Size::from_bytes(1)).next().map(|(_, &t)| t).unwrap()
 225             })
 226             .collect()
 227     }
 228
 229     #[test]
 230     fn basic_insert() {
 231         let mut map = RangeMap::<i32>::new(Size::from_bytes(20), -1);
 232         // Insert.
 233         for (_, x) in map.iter_mut(Size::from_bytes(10), Size::from_bytes(1)) {
 234             *x = 42;
 235         }
 236         // Check.
 237         assert_eq!(to_vec(&map, 10, 1), vec![42]);
 238         assert_eq!(map.v.len(), 3);
 239
 240         // Insert with size 0.
 241         for (_, x) in map.iter_mut(Size::from_bytes(10), Size::from_bytes(0)) {
 242             *x = 19;
 243         }
 244         for (_, x) in map.iter_mut(Size::from_bytes(11), Size::from_bytes(0)) {
 245             *x = 19;
 246         }
 247         assert_eq!(to_vec(&map, 10, 2), vec![42, -1]);
 248         assert_eq!(map.v.len(), 3);
 249     }
 250
 251     #[test]
 252     fn gaps() {
 253         let mut map = RangeMap::<i32>::new(Size::from_bytes(20), -1);
 254         for (_, x) in map.iter_mut(Size::from_bytes(11), Size::from_bytes(1)) {
 255             *x = 42;
 256         }
 257         for (_, x) in map.iter_mut(Size::from_bytes(15), Size::from_bytes(1)) {
 258             *x = 43;
 259         }
 260         assert_eq!(map.v.len(), 5);
 261         assert_eq!(to_vec(&map, 10, 10), vec![-1, 42, -1, -1, -1, 43, -1, -1, -1, -1]);
 262
 263         for (_, x) in map.iter_mut(Size::from_bytes(10), Size::from_bytes(10)) {
 264             if *x < 42 {
 265                 *x = 23;
 266             }
 267         }
 268         assert_eq!(map.v.len(), 6);
 269         assert_eq!(to_vec(&map, 10, 10), vec![23, 42, 23, 23, 23, 43, 23, 23, 23, 23]);
 270         assert_eq!(to_vec(&map, 13, 5), vec![23, 23, 43, 23, 23]);
 271
 272         for (_, x) in map.iter_mut(Size::from_bytes(15), Size::from_bytes(5)) {
 273             *x = 19;
 274         }
 275         assert_eq!(map.v.len(), 6);
 276         assert_eq!(to_vec(&map, 10, 10), vec![23, 42, 23, 23, 23, 19, 19, 19, 19, 19]);
 277         // Should be seeing two blocks with 19.
 278         assert_eq!(
 279             map.iter(Size::from_bytes(15), Size::from_bytes(2))
 280                 .map(|(_, &t)| t)
 281                 .collect::<Vec<_>>(),
 282             vec![19, 19]
 283         );
 284
 285         // A NOP `iter_mut` should trigger merging.
 286         for _ in map.iter_mut(Size::from_bytes(15), Size::from_bytes(5)) {}
 287         assert_eq!(map.v.len(), 5);
 288         assert_eq!(to_vec(&map, 10, 10), vec![23, 42, 23, 23, 23, 19, 19, 19, 19, 19]);
 289     }
 290
 291     #[test]
 292     #[should_panic]
 293     fn out_of_range_iter_mut() {
 294         let mut map = RangeMap::<i32>::new(Size::from_bytes(20), -1);
 295         let _ = map.iter_mut(Size::from_bytes(11), Size::from_bytes(11));
 296     }
 297
 298     #[test]
 299     #[should_panic]
 300     fn out_of_range_iter() {
 301         let map = RangeMap::<i32>::new(Size::from_bytes(20), -1);
 302         let _ = map.iter(Size::from_bytes(11), Size::from_bytes(11));
 303     }
 304 }