range_tree_rs/

lib.rs

#![cfg_attr(docsrs, feature(doc_cfg))]
#![cfg_attr(docsrs, allow(unused_attributes))]
//! This crate provides a range-tree implementation, automatically merge and split range section.
//!
//! The usage scenario for persistent slab allocators (Similar to the range-tree in ZFS),
//! while the underlayer is not avl, but [b+tree](https://docs.rs/embed-collections/latest/embed_collections/#btree),
//! which has more friendly for CPU cache and smaller in memory usage.

use core::{
    cmp::Ordering,
    fmt,
    ops::{AddAssign, Bound, RangeBounds, SubAssign},
};
use embed_btree::btree::{BTreeMap, Cursor, Entry, IntoIter as _IntoIter, Iter as _Iter};
use num_traits::*;

pub trait RangeTreeKey:
    Unsigned + AddAssign + SubAssign + Ord + Copy + fmt::Debug + fmt::Display + Default + 'static
{
}

impl<T> RangeTreeKey for T where
    T: Unsigned
        + AddAssign
        + SubAssign
        + Ord
        + Copy
        + fmt::Debug
        + fmt::Display
        + Default
        + 'static
{
}

/// A b+tree stores range segment of [start, start+size)
pub struct RangeTree<T: RangeTreeKey> {
    tree: BTreeMap<T, T>,
    space: T,
}

/// Trait for allocator
///
/// when range tree merge/split, need to mirror the adding and removal range from size_tree
pub trait RangeTreeOps<T: RangeTreeKey> {
    /// Callback for manage size tree
    fn op_add(&mut self, start: T, size: T);
    /// Callback for manage size tree
    fn op_remove(&mut self, start: T, size: T);
}

struct DummyOps();

impl<T: RangeTreeKey> RangeTreeOps<T> for DummyOps {
    #[inline]
    fn op_add(&mut self, _start: T, _size: T) {}

    #[inline]
    fn op_remove(&mut self, _start: T, _size: T) {}
}

impl<T: RangeTreeKey> RangeTree<T> {
    pub fn new() -> Self {
        Self { tree: BTreeMap::new(), space: T::zero() }
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.tree.is_empty()
    }

    #[inline(always)]
    pub fn get_space(&self) -> T {
        self.space
    }

    #[inline(always)]
    pub fn len(&self) -> usize {
        self.tree.len()
    }

    /// Load discrete range segment
    ///
    /// # Safety
    ///
    /// You should ganrantee the range never adjacent or overlaps
    pub unsafe fn load_unchecked(&mut self, start: T, size: T) -> Result<(), (T, T)> {
        match self.tree.entry(start) {
            Entry::Vacant(ent) => {
                ent.insert(size);
                self.space += size;
                Ok(())
            }
            Entry::Occupied(_ent) => Err((start, size)),
        }
    }

    /// Add range segment, merge with adjacent ranges, assuming no intersections.
    ///
    /// # Return value
    ///
    /// - `Ok(())` if there are no intersection;
    /// - otherwise returns the overlapping range as `Err((existing_start, existing_size))`.
    ///
    /// # Compatibility
    ///
    /// This equals to add + add_find_overlap in v0.1
    #[inline]
    pub fn add(&mut self, start: T, size: T) -> Result<(), (T, T)> {
        self.add_with(start, size, &mut DummyOps {})
    }

    /// The same with [add()](Self::add), with additional mirror operation to `RangeTreeOps`
    #[inline]
    pub fn add_with<O>(&mut self, start: T, size: T, ops: &mut O) -> Result<(), (T, T)>
    where
        O: RangeTreeOps<T>,
    {
        assert!(size > T::zero(), "range tree add size={} error", size);
        let end = start + size;
        match self.tree.entry(start) {
            Entry::Vacant(ent) => {
                if let Some((_prev_start, _prev_size)) = ent.peek_backward() {
                    let prev_end = *_prev_start + *_prev_size;
                    match prev_end.cmp(&start) {
                        Ordering::Less => {}
                        Ordering::Equal => {
                            if let Some((_next_start, _next_size)) = ent.peek_forward() {
                                match end.cmp(_next_start) {
                                    Ordering::Less => {} // cannot merge
                                    Ordering::Equal => {
                                        // merge with prev and next
                                        ops.op_remove(*_prev_start, *_prev_size);
                                        ops.op_remove(*_next_start, *_next_size);
                                        let new_size = *_prev_size + size + *_next_size;
                                        ops.op_add(*_prev_start, new_size);
                                        let mut ent_prev = ent.move_backward().expect("merge prev");
                                        *ent_prev.get_mut() = new_size;
                                        self.space += size;
                                        let ent_next = ent_prev.move_forward().expect("merge next");
                                        ent_next.remove();
                                        return Ok(());
                                    }
                                    Ordering::Greater => return Err((*_next_start, *_next_size)),
                                }
                            }

                            // merge with previous
                            let prev_start = *_prev_start;
                            ops.op_remove(prev_start, *_prev_size);
                            let new_size = *_prev_size + size;
                            ops.op_add(*_prev_start, new_size);
                            let mut ent_prev = ent.move_backward().expect("merge prev");
                            *ent_prev.get_mut() = new_size;
                            self.space += size;
                            return Ok(());
                        }
                        Ordering::Greater => return Err((*_prev_start, *_prev_size)),
                    }
                }
                if let Some((_next_start, _next_size)) = ent.peek_forward() {
                    match end.cmp(_next_start) {
                        Ordering::Less => {}
                        Ordering::Equal => {
                            // merge with next
                            let mut ent_next = ent.move_forward().expect("merge next");
                            ops.op_remove(*_next_start, *_next_size);
                            let new_size = size + *_next_size;
                            *ent_next.get_mut() = new_size;
                            ent_next.alter_key(start).expect("merge next alter_key");
                            ops.op_add(start, new_size);
                            self.space += size;
                            return Ok(());
                        }
                        Ordering::Greater => return Err((*_next_start, *_next_size)),
                    }
                }
                ops.op_add(start, size);
                ent.insert(size);
                self.space += size;
                return Ok(());
            }
            Entry::Occupied(ent) => Err((*ent.key(), *ent.get())),
        }
    }

    /// Add to the tree with [start, end)
    ///
    /// # Return value
    ///
    /// - `Ok(())` if there are no intersection;
    /// - otherwise returns the overlapping range as `Err((existing_start, existing_size))`.
    #[inline(always)]
    pub fn add_abs(&mut self, start: T, end: T) -> Result<(), (T, T)> {
        assert!(start < end, "range tree add start={} end={}", start, end);
        self.add(start, end - start)
    }

    /// Add range which may have multiple intersections with existing range, ensuring union result
    /// stores in the tree
    #[inline]
    pub fn add_loosely(&mut self, start: T, size: T) {
        assert!(size > T::zero(), "range tree add size error");
        let new_end = start + size;
        let base_ent = match self.tree.entry(start) {
            Entry::Vacant(ve) => match ve.peek_backward() {
                None => Entry::Vacant(ve),
                Some((pre_start, pre_size)) => {
                    let cur_end = *pre_start + *pre_size;
                    if cur_end >= new_end {
                        return;
                    }
                    if cur_end >= start {
                        Entry::Occupied(ve.move_backward().expect("move back to merge"))
                    } else {
                        Entry::Vacant(ve)
                    }
                }
            },
            Entry::Occupied(oe) => {
                if start + *oe.get() >= new_end {
                    return;
                }
                Entry::Occupied(oe)
            }
        };

        macro_rules! remove_intersect {
            ($next_start: expr, $new_end: expr) => {
                if let Some((last_start, last_size)) = self.tree.remove_range_with(
                    $next_start..=$new_end,
                    |_removed_start, removed_size| {
                        self.space -= *removed_size;
                    },
                ) {
                    let last_end = last_start + last_size;
                    if last_end > new_end {
                        let _size = last_end - new_end;
                        // add back and join with previous range
                        self.add(new_end, _size)
                            .expect("add {new_end:?}:{_size:?} should not fail");
                    }
                }
            };
        }
        match base_ent {
            Entry::Occupied(mut oe) => {
                let base_start = *oe.key();
                let old_size = *oe.get();

                // extend the size to final size
                let final_size = new_end - base_start;
                self.space += final_size - old_size;
                *oe.get_mut() = final_size;

                if let Some((_next_start, _next_size)) = oe.peek_forward() {
                    let next_start = *_next_start;
                    let next_size = *_next_size;
                    if next_start < new_end {
                        _ = oe;
                        remove_intersect!(next_start, new_end);
                    } else if next_start == new_end {
                        // space is neutral (moving between segments)
                        *oe.get_mut() += next_size;
                        self.tree.remove(&next_start);
                    }
                }
            }
            Entry::Vacant(ve) => {
                let base_start = start;
                self.space += size;
                match ve.peek_forward() {
                    None => {
                        ve.insert(size);
                    }
                    Some((_next_start, _next_size)) => {
                        let next_start = *_next_start;
                        let next_size = *_next_size;
                        if next_start < new_end {
                            ve.insert(size);
                            remove_intersect!(next_start, new_end);
                        } else if next_start == new_end {
                            let final_size = new_end - base_start + next_size;
                            ve.insert(final_size);
                            self.tree.remove(&next_start);
                        } else {
                            ve.insert(size);
                        }
                    }
                }
            }
        }
    }

    /// Valid and remove specify range start:size
    ///
    /// # Return value
    ///
    /// - Only return Ok(()) when there's existing range equal to or contain the removal range in the tree,
    /// - return Err(None) when not found,
    /// - return Err(Some(start, size)) when a range intersect with the removal range, or when the
    ///   removal range larger than existing range.
    #[inline]
    pub fn remove(&mut self, start: T, size: T) -> Result<(), Option<(T, T)>> {
        self.remove_with(start, size, &mut DummyOps {})
    }

    /// The same with [remove()](Self::remove), with additional mirror operation to `RangeTreeOps`
    #[inline]
    pub fn remove_with<O>(&mut self, start: T, size: T, ops: &mut O) -> Result<(), Option<(T, T)>>
    where
        O: RangeTreeOps<T>,
    {
        let end = start + size;
        let ent = self.tree.entry(start);
        match ent {
            Entry::Occupied(mut oent) => {
                let rs_size = *oent.get();
                ops.op_remove(start, rs_size);
                if rs_size == size {
                    // Exact match or subset removed
                    oent.remove();
                    self.space -= rs_size;
                    Ok(())
                } else if rs_size > size {
                    // Shrink from front
                    let new_start = start + size;
                    let new_size = rs_size - size;
                    oent.alter_key(new_start).expect("shrink alter_key");
                    *oent.get_mut() = new_size;
                    ops.op_add(new_start, new_size);
                    self.space -= size;
                    Ok(())
                } else {
                    // existing range smaller than what need to remove
                    Err(Some((start, rs_size)))
                }
            }
            Entry::Vacant(vent) => {
                if let Some((&rs_start, &rs_size)) = vent.peek_backward() {
                    let rs_end = rs_start + rs_size;
                    if rs_end > start {
                        ops.op_remove(rs_start, rs_size);
                        let mut oent = vent.move_backward().expect("move back to overlapping");
                        if rs_end > end {
                            let new_size = start - rs_start;
                            // punch a hold in the middle
                            *oent.get_mut() = new_size;
                            ops.op_add(rs_start, new_size);
                            let new_size2 = rs_end - end;
                            // TODO optimize add insert after entry for btree
                            self.tree.insert(end, new_size2);
                            ops.op_add(end, new_size2);
                            self.space -= size;
                            Ok(())
                        } else if rs_end == end {
                            // Shrink from back
                            let new_size = start - rs_start;
                            *oent.get_mut() = new_size;
                            ops.op_add(rs_start, new_size);
                            self.space -= rs_end - start;
                            Ok(())
                        } else {
                            Err(Some((rs_start, rs_size)))
                        }
                    } else {
                        Err(None)
                    }
                } else {
                    Err(None)
                }
            }
        }
    }

    /// Remove all the intersection ranges in the tree
    ///
    /// Unlike the strict behavior of [RangeTree::remove()],
    /// this function allows the removal range start:size to to be larger than the existing range.
    ///
    /// # Return value
    ///
    /// return true if overlapping range found and removed.
    /// return false if overlapping range not found.
    ///
    /// # Compatibility
    ///
    /// Equals to remove_and_split in v0.1
    #[inline]
    pub fn remove_loosely(&mut self, mut start: T, mut size: T) -> bool {
        let end = start + size;
        let mut ent = self.tree.entry(start);
        let mut removed = false;
        loop {
            match ent {
                Entry::Occupied(mut oent) => {
                    let rs_size = *oent.get();
                    if rs_size == size {
                        // Exact match or subset removed
                        oent.remove();
                        self.space -= rs_size;
                        return true;
                    } else if rs_size > size {
                        // Shrink from front
                        let new_start = start + size;
                        let new_size = rs_size - size;
                        oent.alter_key(new_start).expect("shrink alter_key");
                        *oent.get_mut() = new_size;
                        self.space -= size;
                        return true;
                    } else {
                        if let Some((_next_start, _next_size)) = oent.peek_forward()
                            && *_next_start < end
                        {
                            start = *_next_start;
                            size = end - start;
                            self.space -= *oent.get();
                            oent.remove();
                            ent = self.tree.entry(start);
                            removed = true;
                            continue;
                        }
                        self.space -= rs_size;
                        oent.remove();
                        return true;
                    }
                }
                Entry::Vacant(vent) => {
                    if let Some((&rs_start, &rs_size)) = vent.peek_backward() {
                        let rs_end = rs_start + rs_size;
                        if rs_end > start {
                            let mut oent = vent.move_backward().expect("move back to overlapping");
                            if rs_end > end {
                                let new_size = start - rs_start;
                                // punch a hold in the middle
                                *oent.get_mut() = new_size;
                                let new_size2 = rs_end - end;
                                // TODO optimize add insert after entry for btree
                                self.tree.insert(end, new_size2);
                                self.space -= size;
                                return true;
                            } else {
                                // Shrink from back
                                let new_size = start - rs_start;
                                *oent.get_mut() = new_size;
                                self.space -= rs_end - start;
                                if rs_end == end {
                                    return true;
                                }
                                if let Some((next_start, _)) = oent.peek_forward()
                                    && *next_start < end
                                {
                                    start = *next_start;
                                    size = end - *next_start;
                                    ent = Entry::Occupied(
                                        oent.move_forward().expect("move forward to overlapping"),
                                    );
                                    continue;
                                }
                                return true;
                            }
                        }
                    }
                    // Handle the case where range starts before the first overlapping segment
                    if let Some((next_start, _)) = vent.peek_forward()
                        && *next_start < end
                    {
                        start = *next_start;
                        size = end - *next_start;
                        ent = Entry::Occupied(
                            vent.move_forward().expect("move forward to overlapping"),
                        );
                        continue;
                    }
                    return removed;
                }
            }
        }
    }

    /// return only when segment overlaps with start..end
    ///
    /// # NOTE
    ///
    /// Be aware the first range may have `_start < start`, and the last range may have `_start + _size > end`
    #[inline]
    pub fn range<'a, R: RangeBounds<T>>(&'a self, r: R) -> RangeIter<'a, T> {
        let start = match r.start_bound() {
            Bound::Included(start) => Some(*start),
            Bound::Excluded(start) => Some(*start),
            _ => None,
        };
        let cursor = if let Some(_start) = start {
            let mut _cursor = self.tree.cursor(&_start);
            if let Some((pre_start, pre_size)) = _cursor.peek_backward() {
                let pre_end = *pre_start + *pre_size;
                if pre_end > _start {
                    _cursor.previous();
                }
                // TODO what if we find pre_start < start but pre_start + size >= start
            }
            _cursor
        } else {
            self.tree.first_cursor()
        };
        RangeIter { cursor, end: r.end_bound().cloned(), not_empty: true }
    }

    pub fn collect(&self) -> Vec<(T, T)> {
        let mut v = Vec::with_capacity(self.len());
        for (start, size) in &self.tree {
            v.push((*start, *size))
        }
        v
    }

    #[inline]
    pub fn iter(&self) -> Iter<'_, T> {
        Iter(self.tree.iter())
    }

    pub fn validate(&self) {
        self.tree.validate();
    }

    #[inline]
    pub fn memory_used(&self) -> usize {
        self.tree.memory_used()
    }
}

impl<'a, T: RangeTreeKey> IntoIterator for &'a RangeTree<T> {
    type Item = (T, T);
    type IntoIter = Iter<'a, T>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<T: RangeTreeKey> IntoIterator for RangeTree<T> {
    type Item = (T, T);
    type IntoIter = IntoIter<T>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        IntoIter(self.tree.into_iter())
    }
}

/// An iterator without bound, acquire from [RangeTree::iter()]
pub struct Iter<'a, T: RangeTreeKey>(_Iter<'a, T, T>);

impl<'a, T: RangeTreeKey> Iterator for Iter<'a, T> {
    type Item = (T, T);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.0.next().map(|(start, size)| (*start, *size))
    }
}

/// An consuming iterator for RangeTree
pub struct IntoIter<T: RangeTreeKey>(_IntoIter<T, T>);

impl<T: RangeTreeKey> Iterator for IntoIter<T> {
    type Item = (T, T);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.0.next()
    }
}

/// An iterator for RangeTree with bound, acquire from [RangeTree::range()]
pub struct RangeIter<'a, T: RangeTreeKey> {
    cursor: Cursor<'a, T, T>,
    end: Bound<T>,
    not_empty: bool,
}

impl<'a, T: RangeTreeKey> Iterator for RangeIter<'a, T> {
    type Item = (T, T);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.not_empty {
            if let Some((start, size)) = self.cursor.next() {
                match self.end {
                    Bound::Unbounded => return Some((*start, *size)),
                    Bound::Excluded(end) => {
                        if *start < end {
                            return Some((*start, *size));
                        }
                        self.not_empty = false;
                        return None;
                    }
                    Bound::Included(end) => {
                        if *start <= end {
                            return Some((*start, *size));
                        }
                        self.not_empty = false;
                        return None;
                    }
                }
            }
            self.not_empty = false;
        }
        None
    }
}