src/range_map.rs

   1 #![allow(unused)]
   2
   3 //! Implements a map from integer indices to data.
   4 //! Rather than storing data for every index, internally, this maps entire ranges to the data.
   5 //! To this end, the APIs all work on ranges, not on individual integers. Ranges are split as
   6 //! necessary (e.g. when [0,5) is first associated with X, and then [1,2) is mutated).
   7 //! Users must not depend on whether a range is coalesced or not, even though this is observable
   8 //! via the iteration APIs.
   9 use std::collections::BTreeMap;
  10 use std::ops;
  11
  12 use rustc::ty::layout::Size;
  13
  14 #[derive(Clone, Debug, PartialEq, Eq)]
  15 pub struct RangeMap<T> {
  16     map: BTreeMap<Range, T>,
  17 }
  18
  19 // The derived `Ord` impl sorts first by the first field, then, if the fields are the same,
  20 // by the second field.
  21 // This is exactly what we need for our purposes, since a range query on a BTReeSet/BTreeMap will give us all
  22 // `MemoryRange`s whose `start` is <= than the one we're looking for, but not > the end of the range we're checking.
  23 // At the same time the `end` is irrelevant for the sorting and range searching, but used for the check.
  24 // This kind of search breaks, if `end < start`, so don't do that!
  25 #[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
  26 struct Range {
  27     start: u64,
  28     end: u64, // Invariant: end > start
  29 }
  30
  31 impl Range {
  32     /// Compute a range of ranges that contains all ranges overlaping with [offset, offset+len)
  33     fn range(offset: u64, len: u64) -> ops::Range<Range> {
  34         if len == 0 {
  35             // We can produce an empty range, nothing overlaps with this.
  36             let r = Range { start: 0, end: 1 };
  37             return r..r;
  38         }
  39         // We select all elements that are within
  40         // the range given by the offset into the allocation and the length.
  41         // This is sound if all ranges that intersect with the argument range, are in the
  42         // resulting range of ranges.
  43         let left = Range {
  44             // lowest range to include `offset`
  45             start: 0,
  46             end: offset + 1,
  47         };
  48         let right = Range {
  49             // lowest (valid) range not to include `offset+len`
  50             start: offset + len,
  51             end: offset + len + 1,
  52         };
  53         left..right
  54     }
  55
  56     /// Tests if any element of [offset, offset+len) is contained in this range.
  57     #[inline(always)]
  58     fn overlaps(&self, offset: u64, len: u64) -> bool {
  59         if len == 0 {
  60             // `offset` totally does not matter, we cannot overlap with an empty interval
  61             false
  62         } else {
  63             offset < self.end && offset.checked_add(len).unwrap() >= self.start
  64         }
  65     }
  66 }
  67
  68 impl<T> RangeMap<T> {
  69     /// Create a new RangeMap for the given size, and with the given initial value used for
  70     /// the entire range.
  71     #[inline(always)]
  72     pub fn new(size: Size, init: T) -> RangeMap<T> {
  73         let mut map = RangeMap { map: BTreeMap::new() };
  74         if size.bytes() > 0 {
  75             map.map.insert(Range { start: 0, end: size.bytes() }, init);
  76         }
  77         map
  78     }
  79
  80     fn iter_with_range<'a>(
  81         &'a self,
  82         offset: u64,
  83         len: u64,
  84     ) -> impl Iterator<Item = (&'a Range, &'a T)> + 'a {
  85         self.map.range(Range::range(offset, len)).filter_map(
  86             move |(range, data)| {
  87                 debug_assert!(len > 0);
  88                 if range.overlaps(offset, len) {
  89                     Some((range, data))
  90                 } else {
  91                     None
  92                 }
  93             },
  94         )
  95     }
  96
  97     /// Provide read-only iteration over everything in the given range.  This does
  98     /// *not* split items if they overlap with the edges.  Do not use this to mutate
  99     /// through interior mutability.
 100     pub fn iter<'a>(&'a self, offset: Size, len: Size) -> impl Iterator<Item = &'a T> + 'a {
 101         self.iter_with_range(offset.bytes(), len.bytes()).map(|(_, data)| data)
 102     }
 103
 104     pub fn iter_mut_all<'a>(&'a mut self) -> impl Iterator<Item = &'a mut T> + 'a {
 105         self.map.values_mut()
 106     }
 107
 108     fn split_entry_at(&mut self, offset: u64)
 109     where
 110         T: Clone,
 111     {
 112         let range = match self.iter_with_range(offset, 1).next() {
 113             Some((&range, _)) => range,
 114             None => return,
 115         };
 116         assert!(
 117             range.start <= offset && range.end > offset,
 118             "We got a range that doesn't even contain what we asked for."
 119         );
 120         // There is an entry overlapping this position, see if we have to split it
 121         if range.start < offset {
 122             let data = self.map.remove(&range).unwrap();
 123             let old = self.map.insert(
 124                 Range {
 125                     start: range.start,
 126                     end: offset,
 127                 },
 128                 data.clone(),
 129             );
 130             assert!(old.is_none());
 131             let old = self.map.insert(
 132                 Range {
 133                     start: offset,
 134                     end: range.end,
 135                 },
 136                 data,
 137             );
 138             assert!(old.is_none());
 139         }
 140     }
 141
 142     /// Provide mutable iteration over everything in the given range.  As a side-effect,
 143     /// this will split entries in the map that are only partially hit by the given range,
 144     /// to make sure that when they are mutated, the effect is constrained to the given range.
 145     pub fn iter_mut<'a>(
 146         &'a mut self,
 147         offset: Size,
 148         len: Size,
 149     ) -> impl Iterator<Item = &'a mut T> + 'a
 150     where
 151         T: Clone,
 152     {
 153         let offset = offset.bytes();
 154         let len = len.bytes();
 155
 156         if len > 0 {
 157             // Preparation: Split first and last entry as needed.
 158             self.split_entry_at(offset);
 159             self.split_entry_at(offset + len);
 160         }
 161         // Now we can provide a mutable iterator
 162         self.map.range_mut(Range::range(offset, len)).filter_map(
 163             move |(&range, data)| {
 164                 debug_assert!(len > 0);
 165                 if range.overlaps(offset, len) {
 166                     assert!(
 167                         offset <= range.start && offset + len >= range.end,
 168                         "The splitting went wrong"
 169                     );
 170                     Some(data)
 171                 } else {
 172                     // Skip this one
 173                     None
 174                 }
 175             },
 176         )
 177     }
 178 }
 179
 180 #[cfg(test)]
 181 mod tests {
 182     use super::*;
 183
 184     /// Query the map at every offset in the range and collect the results.
 185     fn to_vec<T: Copy>(map: &RangeMap<T>, offset: u64, len: u64) -> Vec<T> {
 186         (offset..offset + len)
 187             .into_iter()
 188             .map(|i| map
 189                 .iter(Size::from_bytes(i), Size::from_bytes(1))
 190                 .next()
 191                 .map(|&t| t)
 192                 .unwrap()
 193             )
 194             .collect()
 195     }
 196
 197     #[test]
 198     fn basic_insert() {
 199         let mut map = RangeMap::<i32>::new(Size::from_bytes(20), -1);
 200         // Insert
 201         for x in map.iter_mut(Size::from_bytes(10), Size::from_bytes(1)) {
 202             *x = 42;
 203         }
 204         // Check
 205         assert_eq!(to_vec(&map, 10, 1), vec![42]);
 206
 207         // Insert with size 0
 208         for x in map.iter_mut(Size::from_bytes(10), Size::from_bytes(0)) {
 209             *x = 19;
 210         }
 211         for x in map.iter_mut(Size::from_bytes(11), Size::from_bytes(0)) {
 212             *x = 19;
 213         }
 214         assert_eq!(to_vec(&map, 10, 2), vec![42, -1]);
 215     }
 216
 217     #[test]
 218     fn gaps() {
 219         let mut map = RangeMap::<i32>::new(Size::from_bytes(20), -1);
 220         for x in map.iter_mut(Size::from_bytes(11), Size::from_bytes(1)) {
 221             *x = 42;
 222         }
 223         for x in map.iter_mut(Size::from_bytes(15), Size::from_bytes(1)) {
 224             *x = 43;
 225         }
 226         assert_eq!(
 227             to_vec(&map, 10, 10),
 228             vec![-1, 42, -1, -1, -1, 43, -1, -1, -1, -1]
 229         );
 230
 231         for x in map.iter_mut(Size::from_bytes(10), Size::from_bytes(10)) {
 232             if *x < 42 {
 233                 *x = 23;
 234             }
 235         }
 236
 237         assert_eq!(
 238             to_vec(&map, 10, 10),
 239             vec![23, 42, 23, 23, 23, 43, 23, 23, 23, 23]
 240         );
 241         assert_eq!(to_vec(&map, 13, 5), vec![23, 23, 43, 23, 23]);
 242
 243         // Now request a range that goes beyond the initial size
 244         for x in map.iter_mut(Size::from_bytes(15), Size::from_bytes(10)) {
 245             *x = 19;
 246         }
 247         assert_eq!(map.iter(Size::from_bytes(19), Size::from_bytes(1))
 248             .map(|&t| t).collect::<Vec<_>>(), vec![19]);
 249         assert_eq!(map.iter(Size::from_bytes(20), Size::from_bytes(1))
 250             .map(|&t| t).collect::<Vec<_>>(), vec![]);
 251     }
 252 }