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