range_tree_rs/
lib.rs

1#![cfg_attr(docsrs, feature(doc_cfg))]
2#![cfg_attr(docsrs, allow(unused_attributes))]
3//! This crate provides a range-tree implementation, automatically merge and split range section.
4//!
5//! The usage scenario for persistent slab allocators (Similar to the range-tree in ZFS),
6//! while the underlayer is not avl, but [b+tree](https://docs.rs/embed-collections/latest/embed_collections/#btree),
7//! which has more friendly for CPU cache and smaller in memory usage.
8
9use core::{
10    cmp::Ordering,
11    fmt,
12    ops::{AddAssign, Bound, RangeBounds, SubAssign},
13};
14use embed_btree::btree::{BTreeMap, Cursor, Entry, IntoIter as _IntoIter, Iter as _Iter};
15use num_traits::*;
16
17pub trait RangeTreeKey:
18    Unsigned + AddAssign + SubAssign + Ord + Copy + fmt::Debug + fmt::Display + Default + 'static
19{
20}
21
22impl<T> RangeTreeKey for T where
23    T: Unsigned
24        + AddAssign
25        + SubAssign
26        + Ord
27        + Copy
28        + fmt::Debug
29        + fmt::Display
30        + Default
31        + 'static
32{
33}
34
35/// A b+tree stores range segment of [start, start+size)
36pub struct RangeTree<T: RangeTreeKey> {
37    tree: BTreeMap<T, T>,
38    space: T,
39}
40
41/// Trait for allocator
42///
43/// when range tree merge/split, need to mirror the adding and removal range from size_tree
44pub trait RangeTreeOps<T: RangeTreeKey> {
45    /// Callback for manage size tree
46    fn op_add(&mut self, start: T, size: T);
47    /// Callback for manage size tree
48    fn op_remove(&mut self, start: T, size: T);
49}
50
51struct DummyOps();
52
53impl<T: RangeTreeKey> RangeTreeOps<T> for DummyOps {
54    #[inline]
55    fn op_add(&mut self, _start: T, _size: T) {}
56
57    #[inline]
58    fn op_remove(&mut self, _start: T, _size: T) {}
59}
60
61impl<T: RangeTreeKey> RangeTree<T> {
62    pub fn new() -> Self {
63        Self { tree: BTreeMap::new(), space: T::zero() }
64    }
65
66    #[inline]
67    pub fn is_empty(&self) -> bool {
68        self.tree.is_empty()
69    }
70
71    #[inline(always)]
72    pub fn get_space(&self) -> T {
73        self.space
74    }
75
76    #[inline(always)]
77    pub fn len(&self) -> usize {
78        self.tree.len()
79    }
80
81    /// Load discrete range segment
82    ///
83    /// # Safety
84    ///
85    /// You should ganrantee the range never adjacent or overlaps
86    pub unsafe fn load_unchecked(&mut self, start: T, size: T) -> Result<(), (T, T)> {
87        match self.tree.entry(start) {
88            Entry::Vacant(ent) => {
89                ent.insert(size);
90                self.space += size;
91                Ok(())
92            }
93            Entry::Occupied(_ent) => Err((start, size)),
94        }
95    }
96
97    /// Add range segment, merge with adjacent ranges, assuming no intersections.
98    ///
99    /// # Return value
100    ///
101    /// - `Ok(())` if there are no intersection;
102    /// - otherwise returns the overlapping range as `Err((existing_start, existing_size))`.
103    ///
104    /// # Compatibility
105    ///
106    /// This equals to add + add_find_overlap in v0.1
107    #[inline]
108    pub fn add(&mut self, start: T, size: T) -> Result<(), (T, T)> {
109        self.add_with(start, size, &mut DummyOps {})
110    }
111
112    /// The same with [add()](Self::add), with additional mirror operation to `RangeTreeOps`
113    #[inline]
114    pub fn add_with<O>(&mut self, start: T, size: T, ops: &mut O) -> Result<(), (T, T)>
115    where
116        O: RangeTreeOps<T>,
117    {
118        assert!(size > T::zero(), "range tree add size={} error", size);
119        let end = start + size;
120        match self.tree.entry(start) {
121            Entry::Vacant(ent) => {
122                if let Some((prev_start, prev_size)) = ent.peek_backward().map(|(k, v)| (*k, *v)) {
123                    let prev_end = prev_start + prev_size;
124                    match prev_end.cmp(&start) {
125                        Ordering::Less => {}
126                        Ordering::Equal => {
127                            // merge with previous
128                            ops.op_remove(prev_start, prev_size);
129                            if let Some((_next_start, _next_size)) = ent.peek_forward() {
130                                match end.cmp(_next_start) {
131                                    Ordering::Less => {} // cannot merge
132                                    Ordering::Equal => {
133                                        // merge with prev and next
134                                        ops.op_remove(*_next_start, *_next_size);
135                                        let new_size = prev_size + size + *_next_size;
136                                        ops.op_add(prev_start, new_size);
137                                        let mut ent_prev = ent.move_backward().expect("merge prev");
138                                        *ent_prev.get_mut() = new_size;
139                                        self.space += size;
140                                        let ent_next = ent_prev.move_forward().expect("merge next");
141                                        ent_next.remove();
142                                        return Ok(());
143                                    }
144                                    Ordering::Greater => return Err((*_next_start, *_next_size)),
145                                }
146                            }
147                            let new_size = prev_size + size;
148                            ops.op_add(prev_start, new_size);
149                            let mut ent_prev = ent.move_backward().expect("merge prev");
150                            *ent_prev.get_mut() = new_size;
151                            self.space += size;
152                            return Ok(());
153                        }
154                        Ordering::Greater => return Err((prev_start, prev_size)),
155                    }
156                }
157                if let Some((next_start, next_size)) = ent.peek_forward().map(|(k, v)| (*k, *v)) {
158                    match end.cmp(&next_start) {
159                        Ordering::Less => {}
160                        Ordering::Equal => {
161                            // merge with next
162                            let mut ent_next = ent.move_forward().expect("merge next");
163                            ops.op_remove(next_start, next_size);
164                            let new_size = size + next_size;
165                            *ent_next.get_mut() = new_size;
166                            ent_next.alter_key(start).expect("merge next alter_key");
167                            ops.op_add(start, new_size);
168                            self.space += size;
169                            return Ok(());
170                        }
171                        Ordering::Greater => return Err((next_start, next_size)),
172                    }
173                }
174                ops.op_add(start, size);
175                ent.insert(size);
176                self.space += size;
177                Ok(())
178            }
179            Entry::Occupied(ent) => Err((*ent.key(), *ent.get())),
180        }
181    }
182
183    /// Add to the tree with [start, end)
184    ///
185    /// # Return value
186    ///
187    /// - `Ok(())` if there are no intersection;
188    /// - otherwise returns the overlapping range as `Err((existing_start, existing_size))`.
189    #[inline(always)]
190    pub fn add_abs(&mut self, start: T, end: T) -> Result<(), (T, T)> {
191        assert!(start < end, "range tree add start={} end={}", start, end);
192        self.add(start, end - start)
193    }
194
195    /// Add range which may have multiple intersections with existing range, ensuring union result
196    /// stores in the tree
197    #[inline]
198    pub fn add_loosely(&mut self, start: T, size: T) {
199        assert!(size > T::zero(), "range tree add size error");
200        let new_end = start + size;
201        let base_ent = match self.tree.entry(start) {
202            Entry::Vacant(ve) => match ve.peek_backward() {
203                None => Entry::Vacant(ve),
204                Some((pre_start, pre_size)) => {
205                    let cur_end = *pre_start + *pre_size;
206                    if cur_end >= new_end {
207                        return;
208                    }
209                    if cur_end >= start {
210                        Entry::Occupied(ve.move_backward().expect("move back to merge"))
211                    } else {
212                        Entry::Vacant(ve)
213                    }
214                }
215            },
216            Entry::Occupied(oe) => {
217                if start + *oe.get() >= new_end {
218                    return;
219                }
220                Entry::Occupied(oe)
221            }
222        };
223
224        macro_rules! remove_intersect {
225            ($next_start: expr, $new_end: expr) => {
226                if let Some((last_start, last_size)) = self.tree.remove_range_with(
227                    $next_start..=$new_end,
228                    |_removed_start, removed_size| {
229                        self.space -= *removed_size;
230                    },
231                ) {
232                    let last_end = last_start + last_size;
233                    if last_end > new_end {
234                        let _size = last_end - new_end;
235                        // add back and join with previous range
236                        self.add(new_end, _size)
237                            .expect("add {new_end:?}:{_size:?} should not fail");
238                    }
239                }
240            };
241        }
242        match base_ent {
243            Entry::Occupied(mut oe) => {
244                let base_start = *oe.key();
245                let old_size = *oe.get();
246
247                // extend the size to final size
248                let final_size = new_end - base_start;
249                self.space += final_size - old_size;
250                *oe.get_mut() = final_size;
251
252                if let Some((_next_start, _next_size)) = oe.peek_forward() {
253                    let next_start = *_next_start;
254                    let next_size = *_next_size;
255                    if next_start < new_end {
256                        _ = oe;
257                        remove_intersect!(next_start, new_end);
258                    } else if next_start == new_end {
259                        // space is neutral (moving between segments)
260                        *oe.get_mut() += next_size;
261                        self.tree.remove(&next_start);
262                    }
263                }
264            }
265            Entry::Vacant(ve) => {
266                let base_start = start;
267                self.space += size;
268                match ve.peek_forward() {
269                    None => {
270                        ve.insert(size);
271                    }
272                    Some((_next_start, _next_size)) => {
273                        let next_start = *_next_start;
274                        let next_size = *_next_size;
275                        if next_start < new_end {
276                            ve.insert(size);
277                            remove_intersect!(next_start, new_end);
278                        } else if next_start == new_end {
279                            let final_size = new_end - base_start + next_size;
280                            ve.insert(final_size);
281                            self.tree.remove(&next_start);
282                        } else {
283                            ve.insert(size);
284                        }
285                    }
286                }
287            }
288        }
289    }
290
291    /// Valid and remove specify range start:size
292    ///
293    /// # Return value
294    ///
295    /// - Only return Ok(()) when there's existing range equal to or contain the removal range in the tree,
296    /// - return Err(None) when not found,
297    /// - return Err(Some(start, size)) when a range intersect with the removal range, or when the
298    ///   removal range larger than existing range.
299    #[inline]
300    pub fn remove(&mut self, start: T, size: T) -> Result<(), Option<(T, T)>> {
301        self.remove_with(start, size, &mut DummyOps {})
302    }
303
304    /// The same with [remove()](Self::remove), with additional mirror operation to `RangeTreeOps`
305    #[inline]
306    pub fn remove_with<O>(&mut self, start: T, size: T, ops: &mut O) -> Result<(), Option<(T, T)>>
307    where
308        O: RangeTreeOps<T>,
309    {
310        let end = start + size;
311        let ent = self.tree.entry(start);
312        match ent {
313            Entry::Occupied(mut oent) => {
314                let rs_size = *oent.get();
315                ops.op_remove(start, rs_size);
316                if rs_size == size {
317                    // Exact match or subset removed
318                    oent.remove();
319                    self.space -= rs_size;
320                    Ok(())
321                } else if rs_size > size {
322                    // Shrink from front
323                    let new_start = start + size;
324                    let new_size = rs_size - size;
325                    oent.alter_key(new_start).expect("shrink alter_key");
326                    *oent.get_mut() = new_size;
327                    ops.op_add(new_start, new_size);
328                    self.space -= size;
329                    Ok(())
330                } else {
331                    // existing range smaller than what need to remove
332                    Err(Some((start, rs_size)))
333                }
334            }
335            Entry::Vacant(vent) => {
336                if let Some((&rs_start, &rs_size)) = vent.peek_backward() {
337                    let rs_end = rs_start + rs_size;
338                    if rs_end > start {
339                        ops.op_remove(rs_start, rs_size);
340                        let mut oent = vent.move_backward().expect("move back to overlapping");
341                        if rs_end > end {
342                            let new_size = start - rs_start;
343                            // punch a hold in the middle
344                            *oent.get_mut() = new_size;
345                            ops.op_add(rs_start, new_size);
346                            let new_size2 = rs_end - end;
347                            // TODO optimize add insert after entry for btree
348                            self.tree.insert(end, new_size2);
349                            ops.op_add(end, new_size2);
350                            self.space -= size;
351                            Ok(())
352                        } else if rs_end == end {
353                            // Shrink from back
354                            let new_size = start - rs_start;
355                            *oent.get_mut() = new_size;
356                            ops.op_add(rs_start, new_size);
357                            self.space -= rs_end - start;
358                            Ok(())
359                        } else {
360                            Err(Some((rs_start, rs_size)))
361                        }
362                    } else {
363                        Err(None)
364                    }
365                } else {
366                    Err(None)
367                }
368            }
369        }
370    }
371
372    /// Remove all the intersection ranges in the tree
373    ///
374    /// Unlike the strict behavior of [RangeTree::remove()],
375    /// this function allows the removal range start:size to to be larger than the existing range.
376    ///
377    /// # Return value
378    ///
379    /// return true if overlapping range found and removed.
380    /// return false if overlapping range not found.
381    ///
382    /// # Compatibility
383    ///
384    /// Equals to remove_and_split in v0.1
385    #[inline]
386    pub fn remove_loosely(&mut self, mut start: T, mut size: T) -> bool {
387        let end = start + size;
388        let mut ent = self.tree.entry(start);
389        let mut removed = false;
390        loop {
391            match ent {
392                Entry::Occupied(mut oent) => {
393                    let rs_size = *oent.get();
394                    if rs_size == size {
395                        // Exact match or subset removed
396                        oent.remove();
397                        self.space -= rs_size;
398                        return true;
399                    } else if rs_size > size {
400                        // Shrink from front
401                        let new_start = start + size;
402                        let new_size = rs_size - size;
403                        oent.alter_key(new_start).expect("shrink alter_key");
404                        *oent.get_mut() = new_size;
405                        self.space -= size;
406                        return true;
407                    } else {
408                        if let Some((_next_start, _next_size)) = oent.peek_forward()
409                            && *_next_start < end
410                        {
411                            start = *_next_start;
412                            size = end - start;
413                            self.space -= *oent.get();
414                            oent.remove();
415                            ent = self.tree.entry(start);
416                            removed = true;
417                            continue;
418                        }
419                        self.space -= rs_size;
420                        oent.remove();
421                        return true;
422                    }
423                }
424                Entry::Vacant(vent) => {
425                    if let Some((&rs_start, &rs_size)) = vent.peek_backward() {
426                        let rs_end = rs_start + rs_size;
427                        if rs_end > start {
428                            let mut oent = vent.move_backward().expect("move back to overlapping");
429                            if rs_end > end {
430                                let new_size = start - rs_start;
431                                // punch a hold in the middle
432                                *oent.get_mut() = new_size;
433                                let new_size2 = rs_end - end;
434                                // TODO optimize add insert after entry for btree
435                                self.tree.insert(end, new_size2);
436                                self.space -= size;
437                                return true;
438                            } else {
439                                // Shrink from back
440                                let new_size = start - rs_start;
441                                *oent.get_mut() = new_size;
442                                self.space -= rs_end - start;
443                                if rs_end == end {
444                                    return true;
445                                }
446                                if let Some((next_start, _)) = oent.peek_forward()
447                                    && *next_start < end
448                                {
449                                    start = *next_start;
450                                    size = end - *next_start;
451                                    ent = Entry::Occupied(
452                                        oent.move_forward().expect("move forward to overlapping"),
453                                    );
454                                    continue;
455                                }
456                                return true;
457                            }
458                        }
459                    }
460                    // Handle the case where range starts before the first overlapping segment
461                    if let Some((next_start, _)) = vent.peek_forward()
462                        && *next_start < end
463                    {
464                        start = *next_start;
465                        size = end - *next_start;
466                        ent = Entry::Occupied(
467                            vent.move_forward().expect("move forward to overlapping"),
468                        );
469                        continue;
470                    }
471                    return removed;
472                }
473            }
474        }
475    }
476
477    /// return only when segment overlaps with start..end
478    ///
479    /// # NOTE
480    ///
481    /// Be aware the first range may have `_start < start`, and the last range may have `_start + _size > end`
482    #[inline]
483    pub fn range<'a, R: RangeBounds<T>>(&'a self, r: R) -> RangeIter<'a, T> {
484        let start = match r.start_bound() {
485            Bound::Included(start) => Some(*start),
486            Bound::Excluded(start) => Some(*start),
487            _ => None,
488        };
489        let cursor = if let Some(_start) = start {
490            let mut _cursor = self.tree.cursor(&_start);
491            if let Some((pre_start, pre_size)) = _cursor.peek_backward() {
492                let pre_end = *pre_start + *pre_size;
493                if pre_end > _start {
494                    _cursor.previous();
495                }
496                // TODO what if we find pre_start < start but pre_start + size >= start
497            }
498            _cursor
499        } else {
500            self.tree.first_cursor()
501        };
502        RangeIter { cursor, end: r.end_bound().cloned(), not_empty: true }
503    }
504
505    pub fn collect(&self) -> Vec<(T, T)> {
506        let mut v = Vec::with_capacity(self.len());
507        for (start, size) in &self.tree {
508            v.push((*start, *size))
509        }
510        v
511    }
512
513    #[inline]
514    pub fn iter(&self) -> Iter<'_, T> {
515        Iter(self.tree.iter())
516    }
517
518    pub fn validate(&self) {
519        self.tree.validate();
520    }
521
522    #[inline]
523    pub fn memory_used(&self) -> usize {
524        self.tree.memory_used()
525    }
526}
527
528impl<'a, T: RangeTreeKey> IntoIterator for &'a RangeTree<T> {
529    type Item = (T, T);
530    type IntoIter = Iter<'a, T>;
531
532    #[inline]
533    fn into_iter(self) -> Self::IntoIter {
534        self.iter()
535    }
536}
537
538impl<T: RangeTreeKey> IntoIterator for RangeTree<T> {
539    type Item = (T, T);
540    type IntoIter = IntoIter<T>;
541
542    #[inline]
543    fn into_iter(self) -> Self::IntoIter {
544        IntoIter(self.tree.into_iter())
545    }
546}
547
548/// An iterator without bound, acquire from [RangeTree::iter()]
549pub struct Iter<'a, T: RangeTreeKey>(_Iter<'a, T, T>);
550
551impl<'a, T: RangeTreeKey> Iterator for Iter<'a, T> {
552    type Item = (T, T);
553
554    #[inline]
555    fn next(&mut self) -> Option<Self::Item> {
556        self.0.next().map(|(start, size)| (*start, *size))
557    }
558}
559
560/// An consuming iterator for RangeTree
561pub struct IntoIter<T: RangeTreeKey>(_IntoIter<T, T>);
562
563impl<T: RangeTreeKey> Iterator for IntoIter<T> {
564    type Item = (T, T);
565
566    #[inline]
567    fn next(&mut self) -> Option<Self::Item> {
568        self.0.next()
569    }
570}
571
572/// An iterator for RangeTree with bound, acquire from [RangeTree::range()]
573pub struct RangeIter<'a, T: RangeTreeKey> {
574    cursor: Cursor<'a, T, T>,
575    end: Bound<T>,
576    not_empty: bool,
577}
578
579impl<'a, T: RangeTreeKey> Iterator for RangeIter<'a, T> {
580    type Item = (T, T);
581
582    #[inline]
583    fn next(&mut self) -> Option<Self::Item> {
584        if self.not_empty {
585            if let Some((start, size)) = self.cursor.next() {
586                match self.end {
587                    Bound::Unbounded => return Some((*start, *size)),
588                    Bound::Excluded(end) => {
589                        if *start < end {
590                            return Some((*start, *size));
591                        }
592                        self.not_empty = false;
593                        return None;
594                    }
595                    Bound::Included(end) => {
596                        if *start <= end {
597                            return Some((*start, *size));
598                        }
599                        self.not_empty = false;
600                        return None;
601                    }
602                }
603            }
604            self.not_empty = false;
605        }
606        None
607    }
608}
range_tree_rs/lib.rs

range_tree_rs/
lib.rs
