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;
12 #[derive(Clone, Debug, PartialEq, Eq)]
13 pub struct RangeMap<T> {
14 map: BTreeMap<Range, T>,
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)]
26 end: u64, // Invariant: end > start
30 fn range(offset: u64, len: u64) -> ops::Range<Range> {
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.
37 // lowest range to include `offset`
42 // lowest (valid) range not to include `offset+len`
44 end: offset + len + 1,
49 /// Tests if all of [offset, offset+len) are contained in this range.
50 fn overlaps(&self, offset: u64, len: u64) -> bool {
52 offset < self.end && offset + len >= self.start
57 pub fn new() -> RangeMap<T> {
58 RangeMap { map: BTreeMap::new() }
61 fn iter_with_range<'a>(
65 ) -> impl Iterator<Item = (&'a Range, &'a T)> + 'a {
67 self.map.range(Range::range(offset, len)).filter_map(
70 if range.overlaps(offset, len) {
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)
83 fn split_entry_at(&mut self, offset: u64)
87 let range = match self.iter_with_range(offset, 1).next() {
88 Some((&range, _)) => range,
92 range.start <= offset && range.end > offset,
93 "We got a range that doesn't even contain what we asked for."
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(
105 assert!(old.is_none());
106 let old = self.map.insert(
113 assert!(old.is_none());
117 pub fn iter_mut_all<'a>(&'a mut self) -> impl Iterator<Item = &'a mut T> + 'a {
118 self.map.values_mut()
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>(
128 ) -> impl Iterator<Item = &'a mut T> + 'a
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) {
141 offset <= range.start && offset + len >= range.end,
142 "The splitting went wrong"
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
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 {
170 last_end = range.end;
172 if last_end < offset + len {
179 // Add default for all gaps
181 let old = self.map.insert(gap, Default::default());
182 assert!(old.is_none());
185 // Now provide mutable iteration
186 self.iter_mut_with_gaps(offset, len)
189 pub fn retain<F>(&mut self, mut f: F)
191 F: FnMut(&T) -> bool,
193 let mut remove = Vec::new();
194 for (range, data) in &self.map {
200 for range in remove {
201 self.map.remove(&range);
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)
214 .map(|i| *map.iter(i, 1).next().unwrap())
220 let mut map = RangeMap::<i32>::new();
222 for x in map.iter_mut(10, 1) {
226 assert_eq!(to_vec(&map, 10, 1), vec![42]);
231 let mut map = RangeMap::<i32>::new();
232 for x in map.iter_mut(11, 1) {
235 for x in map.iter_mut(15, 1) {
239 // Now request a range that needs three gaps filled
240 for x in map.iter_mut(10, 10) {
247 to_vec(&map, 10, 10),
248 vec![23, 42, 23, 23, 23, 42, 23, 23, 23, 23]
250 assert_eq!(to_vec(&map, 13, 5), vec![23, 23, 42, 23, 23]);