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