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::bptree::{BTree, Cursor, Entry, IntoIter as _IntoIter, Iter as _Iter, Stats};
pub use embed_btree::{compact::Compact, large::TreeInfo};
use num_traits::*;

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

pub type RangeTreeCompact<T> = RangeTree<T, Compact>;
pub type RangeTreeLarge<T> = RangeTree<T, TreeInfo>;

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, S: Stats = Compact> {
    tree: BTree<T, T, S>,
    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, S: Stats> RangeTree<T, S> {
    pub fn new() -> Self {
        Self { tree: BTree::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().map(|(k, v)| (*k, *v)) {
                    let prev_end = prev_start + prev_size;
                    match prev_end.cmp(&start) {
                        Ordering::Less => {}
                        Ordering::Equal => {
                            // merge with previous
                            ops.op_remove(prev_start, prev_size);
                            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(*_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)),
                                }
                            }
                            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().map(|(k, v)| (*k, *v)) {
                    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;
                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 next_start: T;
        {
            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)
                }
            };
            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() {
                        next_start = *_next_start;
                        let next_size = *_next_size;
                        if next_start < new_end {
                            // fallback to remove range to escape lifetime
                        } else if next_start == new_end {
                            // space is neutral (moving between segments)
                            *oe.get_mut() += next_size;
                            // fallback to remove(next_start); to escape liftime
                        } else {
                            return;
                        }
                    } else {
                        return;
                    }
                }
                Entry::Vacant(ve) => {
                    let base_start = start;
                    self.space += size;
                    match ve.peek_forward() {
                        None => {
                            ve.insert(size);
                            return;
                        }
                        Some((_next_start, _next_size)) => {
                            next_start = *_next_start;
                            let next_size = *_next_size;
                            if next_start < new_end {
                                ve.insert(size);
                                // fallback to remove range
                            } else if next_start == new_end {
                                let final_size = new_end - base_start + next_size;
                                ve.insert(final_size);
                                // fallback to remove(next_start);
                            } else {
                                ve.insert(size);
                                return;
                            }
                        }
                    }
                }
            }
        }
        if next_start < new_end {
            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");
                }
            }
        } else {
            assert_eq!(next_start, new_end);
            self.tree.remove(&next_start);
        }
    }

    /// 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 new_size2: T;
        {
            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;
                        return 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;
                        return Ok(());
                    } else {
                        // existing range smaller than what need to remove
                        return 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);
                                new_size2 = rs_end - end;
                                // fallback to escape the lifetime borrowing
                                // TODO optimize add insert after entry for btree
                            } 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;
                                return Ok(());
                            } else {
                                return Err(Some((rs_start, rs_size)));
                            }
                        } else {
                            return Err(None);
                        }
                    } else {
                        return Err(None);
                    }
                }
            }
        }

        self.tree.insert(end, new_size2);
        ops.op_add(end, new_size2);
        self.space -= size;
        Ok(())
    }

    /// 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 removed = false;
        let size_after_end;
        'MAIN: loop {
            let mut oent = match self.tree.entry(start) {
                Entry::Occupied(oent) => oent,
                Entry::Vacant(vent) => {
                    if let Some((rs_start, rs_size)) = vent.peek_backward()
                        && *rs_start + *rs_size > start
                    {
                        let rs_end = *rs_start + *rs_size;
                        let new_size = start - *rs_start;
                        self.space -= *rs_size - new_size;
                        let mut oent = vent.move_backward().expect("move back to overlapping");
                        *oent.get_mut() = new_size;
                        if rs_end == end {
                            // rs_start < start <= rs_end = end
                            return true;
                        } else if rs_end < end {
                            // rs_start < start <= rs_end <  end
                            if let Some((next_start, _)) = oent.peek_forward()
                                && *next_start < end
                            {
                                start = *next_start;
                                size = end - *next_start;
                                oent.move_forward().expect("move forward to overlapping")
                            } else {
                                return true;
                            }
                        } else {
                            debug_assert!(rs_end > end);
                            size_after_end = rs_end - end;
                            // TODO optimize add insert after entry for btree
                            // break to escape lifetime borrowing and re-insert the rest
                            break 'MAIN;
                        }
                    } else {
                        // 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;
                            // fallthough
                            vent.move_forward().expect("move forward to overlapping")
                        } else {
                            return removed;
                        }
                    }
                }
            };
            let rs_size = *oent.get();
            if rs_size <= size {
                let mut next = None;
                if rs_size < size
                    && let Some((_next_start, _next_size)) = oent.peek_forward()
                    && *_next_start < end
                {
                    next.replace((*_next_start, end - *_next_start));
                }
                removed = true;
                self.space -= rs_size;
                oent.remove();
                if let Some((next_start, next_size)) = next {
                    start = next_start;
                    size = next_size;
                    // continue
                } else {
                    return true;
                }
            } else {
                debug_assert!(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;
            }
        }
        self.tree.insert(end, size_after_end);
        self.space += size_after_end;
        true
    }

    /// 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();
    }
}

impl<T: RangeTreeKey> RangeTree<T, TreeInfo> {
    #[inline]
    pub fn memory_used(&self) -> usize {
        self.tree.memory_used()
    }
}

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

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

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

    #[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, S: Stats>(_IntoIter<T, T, S>);

impl<T: RangeTreeKey, S: Stats> Iterator for IntoIter<T, S> {
    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
    }
}