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