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 impl<T> Default for RangeMap<T> {
  20     #[inline(always)]
  21     fn default() -> Self {
  22         RangeMap::new()
  23     }
  24 }
  25
  26 // The derived `Ord` impl sorts first by the first field, then, if the fields are the same,
  27 // by the second field.
  28 // This is exactly what we need for our purposes, since a range query on a BTReeSet/BTreeMap will give us all
  29 // `MemoryRange`s whose `start` is <= than the one we're looking for, but not > the end of the range we're checking.
  30 // At the same time the `end` is irrelevant for the sorting and range searching, but used for the check.
  31 // This kind of search breaks, if `end < start`, so don't do that!
  32 #[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
  33 struct Range {
  34     start: u64,
  35     end: u64, // Invariant: end > start
  36 }
  37
  38 impl Range {
  39     /// Compute a range of ranges that contains all ranges overlaping with [offset, offset+len)
  40     fn range(offset: u64, len: u64) -> ops::Range<Range> {
  41         if len == 0 {
  42             // We can produce an empty range, nothing overlaps with this.
  43             let r = Range { start: 0, end: 1 };
  44             return r..r;
  45         }
  46         // We select all elements that are within
  47         // the range given by the offset into the allocation and the length.
  48         // This is sound if all ranges that intersect with the argument range, are in the
  49         // resulting range of ranges.
  50         let left = Range {
  51             // lowest range to include `offset`
  52             start: 0,
  53             end: offset + 1,
  54         };
  55         let right = Range {
  56             // lowest (valid) range not to include `offset+len`
  57             start: offset + len,
  58             end: offset + len + 1,
  59         };
  60         left..right
  61     }
  62
  63     /// Tests if any element of [offset, offset+len) is contained in this range.
  64     #[inline(always)]
  65     fn overlaps(&self, offset: u64, len: u64) -> bool {
  66         if len == 0 {
  67             // `offset` totally does not matter, we cannot overlap with an empty interval
  68             false
  69         } else {
  70             offset < self.end && offset.checked_add(len).unwrap() >= self.start
  71         }
  72     }
  73 }
  74
  75 impl<T> RangeMap<T> {
  76     #[inline(always)]
  77     pub fn new() -> RangeMap<T> {
  78         RangeMap { map: BTreeMap::new() }
  79     }
  80
  81     fn iter_with_range<'a>(
  82         &'a self,
  83         offset: u64,
  84         len: u64,
  85     ) -> impl Iterator<Item = (&'a Range, &'a T)> + 'a {
  86         self.map.range(Range::range(offset, len)).filter_map(
  87             move |(range, data)| {
  88                 debug_assert!(len > 0);
  89                 if range.overlaps(offset, len) {
  90                     Some((range, data))
  91                 } else {
  92                     None
  93                 }
  94             },
  95         )
  96     }
  97
  98     pub fn iter<'a>(&'a self, offset: Size, len: Size) -> impl Iterator<Item = &'a T> + 'a {
  99         self.iter_with_range(offset.bytes(), len.bytes()).map(|(_, data)| data)
 100     }
 101
 102     pub fn iter_mut_all<'a>(&'a mut self) -> impl Iterator<Item = &'a mut T> + 'a {
 103         self.map.values_mut()
 104     }
 105
 106     fn split_entry_at(&mut self, offset: u64)
 107     where
 108         T: Clone,
 109     {
 110         let range = match self.iter_with_range(offset, 1).next() {
 111             Some((&range, _)) => range,
 112             None => return,
 113         };
 114         assert!(
 115             range.start <= offset && range.end > offset,
 116             "We got a range that doesn't even contain what we asked for."
 117         );
 118         // There is an entry overlapping this position, see if we have to split it
 119         if range.start < offset {
 120             let data = self.map.remove(&range).unwrap();
 121             let old = self.map.insert(
 122                 Range {
 123                     start: range.start,
 124                     end: offset,
 125                 },
 126                 data.clone(),
 127             );
 128             assert!(old.is_none());
 129             let old = self.map.insert(
 130                 Range {
 131                     start: offset,
 132                     end: range.end,
 133                 },
 134                 data,
 135             );
 136             assert!(old.is_none());
 137         }
 138     }
 139
 140     /// Provide mutable iteration over everything in the given range.  As a side-effect,
 141     /// this will split entries in the map that are only partially hit by the given range,
 142     /// to make sure that when they are mutated, the effect is constrained to the given range.
 143     /// If there are gaps, leave them be.
 144     pub fn iter_mut_with_gaps<'a>(
 145         &'a mut self,
 146         offset: Size,
 147         len: Size,
 148     ) -> impl Iterator<Item = &'a mut T> + 'a
 149     where
 150         T: Clone,
 151     {
 152         let offset = offset.bytes();
 153         let len = len.bytes();
 154
 155         if len > 0 {
 156             // Preparation: Split first and last entry as needed.
 157             self.split_entry_at(offset);
 158             self.split_entry_at(offset + len);
 159         }
 160         // Now we can provide a mutable iterator
 161         self.map.range_mut(Range::range(offset, len)).filter_map(
 162             move |(&range, data)| {
 163                 debug_assert!(len > 0);
 164                 if range.overlaps(offset, len) {
 165                     assert!(
 166                         offset <= range.start && offset + len >= range.end,
 167                         "The splitting went wrong"
 168                     );
 169                     Some(data)
 170                 } else {
 171                     // Skip this one
 172                     None
 173                 }
 174             },
 175         )
 176     }
 177
 178     /// Provide a mutable iterator over everything in the given range, with the same side-effects as
 179     /// iter_mut_with_gaps.  Furthermore, if there are gaps between ranges, fill them with the given default
 180     /// before yielding them in the iterator.
 181     /// This is also how you insert.
 182     pub fn iter_mut<'a>(&'a mut self, offset: Size, len: Size) -> impl Iterator<Item = &'a mut T> + 'a
 183     where
 184         T: Clone + Default,
 185     {
 186         if len.bytes() > 0 {
 187             let offset = offset.bytes();
 188             let len = len.bytes();
 189
 190             // Do a first iteration to collect the gaps
 191             let mut gaps = Vec::new();
 192             let mut last_end = offset;
 193             for (range, _) in self.iter_with_range(offset, len) {
 194                 if last_end < range.start {
 195                     gaps.push(Range {
 196                         start: last_end,
 197                         end: range.start,
 198                     });
 199                 }
 200                 last_end = range.end;
 201             }
 202             if last_end < offset + len {
 203                 gaps.push(Range {
 204                     start: last_end,
 205                     end: offset + len,
 206                 });
 207             }
 208
 209             // Add default for all gaps
 210             for gap in gaps {
 211                 let old = self.map.insert(gap, Default::default());
 212                 assert!(old.is_none());
 213             }
 214         }
 215
 216         // Now provide mutable iteration
 217         self.iter_mut_with_gaps(offset, len)
 218     }
 219
 220     pub fn retain<F>(&mut self, mut f: F)
 221     where
 222         F: FnMut(&T) -> bool,
 223     {
 224         let mut remove = Vec::new();
 225         for (range, data) in &self.map {
 226             if !f(data) {
 227                 remove.push(*range);
 228             }
 229         }
 230
 231         for range in remove {
 232             self.map.remove(&range);
 233         }
 234     }
 235 }
 236
 237 #[cfg(test)]
 238 mod tests {
 239     use super::*;
 240
 241     /// Query the map at every offset in the range and collect the results.
 242     fn to_vec<T: Copy>(map: &RangeMap<T>, offset: u64, len: u64, default: Option<T>) -> Vec<T> {
 243         (offset..offset + len)
 244             .into_iter()
 245             .map(|i| map
 246                 .iter(Size::from_bytes(i), Size::from_bytes(1))
 247                 .next()
 248                 .map(|&t| t)
 249                 .or(default)
 250                 .unwrap()
 251             )
 252             .collect()
 253     }
 254
 255     #[test]
 256     fn basic_insert() {
 257         let mut map = RangeMap::<i32>::new();
 258         // Insert
 259         for x in map.iter_mut(Size::from_bytes(10), Size::from_bytes(1)) {
 260             *x = 42;
 261         }
 262         // Check
 263         assert_eq!(to_vec(&map, 10, 1, None), vec![42]);
 264
 265         // Insert with size 0
 266         for x in map.iter_mut(Size::from_bytes(10), Size::from_bytes(0)) {
 267             *x = 19;
 268         }
 269         for x in map.iter_mut(Size::from_bytes(11), Size::from_bytes(0)) {
 270             *x = 19;
 271         }
 272         assert_eq!(to_vec(&map, 10, 2, Some(-1)), vec![42, -1]);
 273     }
 274
 275     #[test]
 276     fn gaps() {
 277         let mut map = RangeMap::<i32>::new();
 278         for x in map.iter_mut(Size::from_bytes(11), Size::from_bytes(1)) {
 279             *x = 42;
 280         }
 281         for x in map.iter_mut(Size::from_bytes(15), Size::from_bytes(1)) {
 282             *x = 43;
 283         }
 284         assert_eq!(
 285             to_vec(&map, 10, 10, Some(-1)),
 286             vec![-1, 42, -1, -1, -1, 43, -1, -1, -1, -1]
 287         );
 288
 289         // Now request a range that needs three gaps filled
 290         for x in map.iter_mut(Size::from_bytes(10), Size::from_bytes(10)) {
 291             if *x < 42 {
 292                 *x = 23;
 293             }
 294         }
 295
 296         assert_eq!(
 297             to_vec(&map, 10, 10, None),
 298             vec![23, 42, 23, 23, 23, 43, 23, 23, 23, 23]
 299         );
 300         assert_eq!(to_vec(&map, 13, 5, None), vec![23, 23, 43, 23, 23]);
 301     }
 302 }