1 //! Implements a map from integer indices to data.
2 //! Rather than storing data for every index, internally, this maps entire ranges to the data.
3 //! To this end, the APIs all work on ranges, not on individual integers. Ranges are split as
4 //! necessary (e.g., when [0,5) is first associated with X, and then [1,2) is mutated).
5 //! Users must not depend on whether a range is coalesced or not, even though this is observable
6 //! via the iteration APIs.
10 use rustc_target::abi::Size;
12 #[derive(Clone, Debug)]
14 /// The range covered by this element; never empty.
15 range: ops::Range<u64>,
16 /// The data stored for this element.
19 #[derive(Clone, Debug)]
20 pub struct RangeMap<T> {
25 /// Creates a new `RangeMap` for the given size, and with the given initial value used for
28 pub fn new(size: Size, init: T) -> RangeMap<T> {
29 let size = size.bytes();
30 let mut map = RangeMap { v: Vec::new() };
32 map.v.push(Elem { range: 0..size, data: init });
37 /// Finds the index containing the given offset.
38 fn find_offset(&self, offset: u64) -> usize {
39 // We do a binary search.
40 let mut left = 0usize; // inclusive
41 let mut right = self.v.len(); // exclusive
43 debug_assert!(left < right, "find_offset: offset {offset} is out-of-bounds");
44 let candidate = left.checked_add(right).unwrap() / 2;
45 let elem = &self.v[candidate];
46 if offset < elem.range.start {
47 // We are too far right (offset is further left).
48 debug_assert!(candidate < right); // we are making progress
50 } else if offset >= elem.range.end {
51 // We are too far left (offset is further right).
52 debug_assert!(candidate >= left); // we are making progress
61 /// Provides read-only iteration over everything in the given range. This does
62 /// *not* split items if they overlap with the edges. Do not use this to mutate
63 /// through interior mutability.
65 /// The iterator also provides the offset of the given element.
66 pub fn iter(&self, offset: Size, len: Size) -> impl Iterator<Item = (Size, &T)> {
67 let offset = offset.bytes();
68 let len = len.bytes();
69 // Compute a slice starting with the elements we care about.
70 let slice: &[Elem<T>] = if len == 0 {
71 // We just need any empty iterator. We don't even want to
72 // yield the element that surrounds this position.
75 let first_idx = self.find_offset(offset);
78 // The first offset that is not included any more.
79 let end = offset + len;
81 end <= self.v.last().unwrap().range.end,
82 "iterating beyond the bounds of this RangeMap"
86 .take_while(move |elem| elem.range.start < end)
87 .map(|elem| (Size::from_bytes(elem.range.start), &elem.data))
90 pub fn iter_mut_all(&mut self) -> impl Iterator<Item = &mut T> {
91 self.v.iter_mut().map(|elem| &mut elem.data)
94 pub fn iter_all(&self) -> impl Iterator<Item = (ops::Range<u64>, &T)> {
95 self.v.iter().map(|elem| (elem.range.clone(), &elem.data))
98 // Splits the element situated at the given `index`, such that the 2nd one starts at offset
99 // `split_offset`. Do nothing if the element already starts there.
100 // Returns whether a split was necessary.
101 fn split_index(&mut self, index: usize, split_offset: u64) -> bool
105 let elem = &mut self.v[index];
106 if split_offset == elem.range.start || split_offset == elem.range.end {
111 elem.range.contains(&split_offset),
112 "the `split_offset` is not in the element to be split"
115 // Now we really have to split. Reduce length of first element.
116 let second_range = split_offset..elem.range.end;
117 elem.range.end = split_offset;
118 // Copy the data, and insert second element.
119 let second = Elem { range: second_range, data: elem.data.clone() };
120 self.v.insert(index + 1, second);
124 /// Provides mutable iteration over everything in the given range. As a side-effect,
125 /// this will split entries in the map that are only partially hit by the given range,
126 /// to make sure that when they are mutated, the effect is constrained to the given range.
127 /// Moreover, this will opportunistically merge neighbouring equal blocks.
129 /// The iterator also provides the offset of the given element.
130 pub fn iter_mut(&mut self, offset: Size, len: Size) -> impl Iterator<Item = (Size, &mut T)>
132 T: Clone + PartialEq,
134 let offset = offset.bytes();
135 let len = len.bytes();
136 // Compute a slice containing exactly the elements we care about
137 let slice: &mut [Elem<T>] = if len == 0 {
138 // We just need any empty iterator. We don't even want to
139 // yield the element that surrounds this position, nor do
143 // Make sure we got a clear beginning
144 let mut first_idx = self.find_offset(offset);
145 if self.split_index(first_idx, offset) {
146 // The newly created 2nd element is ours
150 let first_idx = first_idx;
151 // Find our end. Linear scan, but that's ok because the iteration
152 // is doing the same linear scan anyway -- no increase in complexity.
153 // We combine this scan with a scan for duplicates that we can merge, to reduce
154 // the number of elements.
155 // We stop searching after the first "block" of size 1, to avoid spending excessive
156 // amounts of time on the merging.
157 let mut equal_since_idx = first_idx;
158 // Once we see too many non-mergeable blocks, we stop.
159 // The initial value is chosen via... magic. Benchmarking and magic.
160 let mut successful_merge_count = 3usize;
161 // When the loop is done, this is the first excluded element.
162 let mut end_idx = first_idx;
164 // Compute if `end` is the last element we need to look at.
165 let done = self.v[end_idx].range.end >= offset + len;
166 // We definitely need to include `end`, so move the index.
169 done || end_idx < self.v.len(),
170 "iter_mut: end-offset {} is out-of-bounds",
173 // see if we want to merge everything in `equal_since..end` (exclusive at the end!)
174 if successful_merge_count > 0 {
175 if done || self.v[end_idx].data != self.v[equal_since_idx].data {
176 // Everything in `equal_since..end` was equal. Make them just one element covering
178 let removed_elems = end_idx - equal_since_idx - 1; // number of elements that we would remove
179 if removed_elems > 0 {
180 // Adjust the range of the first element to cover all of them.
181 let equal_until = self.v[end_idx - 1].range.end; // end of range of last of the equal elements
182 self.v[equal_since_idx].range.end = equal_until;
183 // Delete the rest of them.
184 self.v.splice(equal_since_idx + 1..end_idx, std::iter::empty());
185 // Adjust `end_idx` because we made the list shorter.
186 end_idx -= removed_elems;
187 // Adjust the count for the cutoff.
188 successful_merge_count += removed_elems;
190 // Adjust the count for the cutoff.
191 successful_merge_count -= 1;
193 // Go on scanning for the next block starting here.
194 equal_since_idx = end_idx;
197 // Leave loop if this is the last element.
202 // Move to last included instead of first excluded index.
203 let end_idx = end_idx - 1;
204 // We need to split the end as well. Even if this performs a
205 // split, we don't have to adjust our index as we only care about
206 // the first part of the split.
207 self.split_index(end_idx, offset + len);
208 // Now we yield the slice. `end` is inclusive.
209 &mut self.v[first_idx..=end_idx]
211 slice.iter_mut().map(|elem| (Size::from_bytes(elem.range.start), &mut elem.data))
219 /// Query the map at every offset in the range and collect the results.
220 fn to_vec<T: Copy>(map: &RangeMap<T>, offset: u64, len: u64) -> Vec<T> {
221 (offset..offset + len)
223 map.iter(Size::from_bytes(i), Size::from_bytes(1)).next().map(|(_, &t)| t).unwrap()
230 let mut map = RangeMap::<i32>::new(Size::from_bytes(20), -1);
232 for (_, x) in map.iter_mut(Size::from_bytes(10), Size::from_bytes(1)) {
236 assert_eq!(to_vec(&map, 10, 1), vec![42]);
237 assert_eq!(map.v.len(), 3);
239 // Insert with size 0.
240 for (_, x) in map.iter_mut(Size::from_bytes(10), Size::from_bytes(0)) {
243 for (_, x) in map.iter_mut(Size::from_bytes(11), Size::from_bytes(0)) {
246 assert_eq!(to_vec(&map, 10, 2), vec![42, -1]);
247 assert_eq!(map.v.len(), 3);
252 let mut map = RangeMap::<i32>::new(Size::from_bytes(20), -1);
253 for (_, x) in map.iter_mut(Size::from_bytes(11), Size::from_bytes(1)) {
256 for (_, x) in map.iter_mut(Size::from_bytes(15), Size::from_bytes(1)) {
259 assert_eq!(map.v.len(), 5);
260 assert_eq!(to_vec(&map, 10, 10), vec![-1, 42, -1, -1, -1, 43, -1, -1, -1, -1]);
262 for (_, x) in map.iter_mut(Size::from_bytes(10), Size::from_bytes(10)) {
267 assert_eq!(map.v.len(), 6);
268 assert_eq!(to_vec(&map, 10, 10), vec![23, 42, 23, 23, 23, 43, 23, 23, 23, 23]);
269 assert_eq!(to_vec(&map, 13, 5), vec![23, 23, 43, 23, 23]);
271 for (_, x) in map.iter_mut(Size::from_bytes(15), Size::from_bytes(5)) {
274 assert_eq!(map.v.len(), 6);
275 assert_eq!(to_vec(&map, 10, 10), vec![23, 42, 23, 23, 23, 19, 19, 19, 19, 19]);
276 // Should be seeing two blocks with 19.
278 map.iter(Size::from_bytes(15), Size::from_bytes(2))
280 .collect::<Vec<_>>(),
284 // A NOP `iter_mut` should trigger merging.
285 for _ in map.iter_mut(Size::from_bytes(15), Size::from_bytes(5)) {}
286 assert_eq!(map.v.len(), 5);
287 assert_eq!(to_vec(&map, 10, 10), vec![23, 42, 23, 23, 23, 19, 19, 19, 19, 19]);
292 fn out_of_range_iter_mut() {
293 let mut map = RangeMap::<i32>::new(Size::from_bytes(20), -1);
294 let _ = map.iter_mut(Size::from_bytes(11), Size::from_bytes(11));
299 fn out_of_range_iter() {
300 let map = RangeMap::<i32>::new(Size::from_bytes(20), -1);
301 let _ = map.iter(Size::from_bytes(11), Size::from_bytes(11));