use crate::{
interval::Interval,
iterators::{DifferenceIterator, IntervalIterator, VoidsIterator},
numerated::Numerated,
};
use alloc::{collections::BTreeMap, fmt, fmt::Debug, vec::Vec};
use core::{fmt::Formatter, ops::RangeInclusive};
use num_traits::{CheckedAdd, Zero};
use scale_info::{
TypeInfo,
scale::{Decode, Encode},
};
#[derive(Clone, PartialEq, Eq, Hash, TypeInfo, Encode, Decode)]
#[codec(crate = scale_info::scale)]
pub struct IntervalsTree<T> {
inner: BTreeMap<T, T>,
}
impl<T: Copy> IntervalsTree<T> {
pub const fn new() -> Self {
Self {
inner: BTreeMap::new(),
}
}
pub fn intervals_amount(&self) -> usize {
self.inner.len()
}
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
}
impl<T: Copy + Ord> IntervalsTree<T> {
pub fn end(&self) -> Option<T> {
self.inner.last_key_value().map(|(_, &e)| e)
}
pub fn start(&self) -> Option<T> {
self.inner.first_key_value().map(|(&s, _)| s)
}
}
impl<T: Copy> Default for IntervalsTree<T> {
fn default() -> Self {
Self::new()
}
}
impl<T: Debug + Numerated> Debug for IntervalsTree<T> {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
f.debug_list().entries(self.iter()).finish()
}
}
impl<T: Numerated> IntervalsTree<T> {
fn into_start_end<I: Into<IntervalIterator<T>>>(interval: I) -> Option<(T, T)> {
Into::<IntervalIterator<T>>::into(interval)
.inner()
.map(|i| i.into_parts())
}
#[track_caller]
fn put(&mut self, start: T, end: T) {
debug_assert!(start <= end, "Must be guarantied");
self.inner.insert(start, end);
}
pub fn iter(&self) -> impl Iterator<Item = Interval<T>> + '_ {
self.inner.iter().map(|(&start, &end)| unsafe {
Interval::<T>::new_unchecked(start, end)
})
}
pub fn contains<I: Into<IntervalIterator<T>>>(&self, interval: I) -> bool {
let Some((start, end)) = Self::into_start_end(interval) else {
return true;
};
self.inner
.range(..=end)
.next_back()
.map(|(&s, &e)| s <= start && end <= e)
.unwrap_or(false)
}
pub fn insert<I: Into<IntervalIterator<T>>>(&mut self, interval: I) -> bool {
let Some((start, end)) = Self::into_start_end(interval) else {
return false;
};
let Some(last) = self.end() else {
self.put(start, end);
return true;
};
let iter_end = end.inc_if_lt(last).unwrap_or(end);
let mut iter = self.inner.range(..=iter_end).map(|(&s, &e)| (s, e));
let Some((right_start, right_end)) = iter.next_back() else {
self.put(start, end);
return true;
};
if let Some(right_end) = right_end.inc_if_lt(start) {
if right_end == start {
self.put(right_start, end);
} else {
self.put(start, end);
}
return true;
} else if right_start <= start {
if right_end < end {
self.put(right_start, end);
return true;
} else {
return false;
}
}
let mut left_interval = None;
let mut intervals_to_remove = Vec::new();
while let Some((s, e)) = iter.next_back() {
if s <= start {
left_interval = Some((s, e));
break;
}
intervals_to_remove.push(s);
}
for start in intervals_to_remove {
self.inner.remove(&start);
}
self.inner.remove(&right_start);
let end = right_end.max(end);
let Some((left_start, left_end)) = left_interval else {
self.put(start, end);
return true;
};
debug_assert!(left_end < right_start && left_start <= start);
let Some(left_end) = left_end.inc_if_lt(right_start) else {
debug_assert!(false, "`T: Numerated` impl error");
return false;
};
if left_end >= start {
self.put(left_start, end);
} else {
self.put(start, end);
}
true
}
pub fn remove<I: Into<IntervalIterator<T>>>(&mut self, interval: I) -> bool {
let Some((start, end)) = Self::into_start_end(interval) else {
return false;
};
let mut iter = self.inner.range(..=end);
let Some((&right_start, &right_end)) = iter.next_back() else {
return false;
};
if right_end < start {
return false;
}
let mut left_interval = None;
let mut intervals_to_remove = Vec::new();
while let Some((&s, &e)) = iter.next_back() {
if s < start {
if e >= start {
left_interval = Some(s);
}
break;
}
intervals_to_remove.push(s)
}
for start in intervals_to_remove {
self.inner.remove(&start);
}
if let Some(start) = start.dec_if_gt(right_start) {
self.put(right_start, start);
} else {
debug_assert!(
right_start <= end,
"Must be, because of method it was found"
);
self.inner.remove(&right_start);
}
if let Some(end) = end.inc_if_lt(right_end) {
self.put(end, right_end);
} else {
debug_assert!(start <= right_end);
}
if let Some(left_start) = left_interval {
if let Some(start) = start.dec_if_gt(left_start) {
self.put(left_start, start);
} else {
debug_assert!(false, "`T: Numerated` impl error");
}
}
true
}
pub fn voids<I: Into<IntervalIterator<T>>>(
&self,
interval: I,
) -> VoidsIterator<T, impl Iterator<Item = Interval<T>> + '_> {
let Some((mut start, end)) = Self::into_start_end(interval) else {
return VoidsIterator { inner: None };
};
if let Some((_, &e)) = self.inner.range(..=start).next_back() {
if let Some(e) = e.inc_if_lt(end) {
if e > start {
start = e;
}
} else {
return VoidsIterator { inner: None };
}
}
let iter = self.inner.range(start..=end).map(|(&start, &end)| {
unsafe { Interval::new_unchecked(start, end) }
});
let interval = unsafe { Interval::new_unchecked(start, end) };
VoidsIterator {
inner: Some((iter, interval)),
}
}
pub fn difference<'a>(&'a self, other: &'a Self) -> impl Iterator<Item = Interval<T>> + 'a {
DifferenceIterator {
iter1: self.iter(),
iter2: other.iter(),
interval1: None,
interval2: None,
}
}
pub fn points_amount(&self) -> Option<T::Distance> {
let mut res = T::Distance::zero();
for interval in self.iter() {
res = res.checked_add(&interval.raw_len()?)?;
}
Some(res)
}
pub fn points_iter(&self) -> impl Iterator<Item = T> + '_ {
self.inner.iter().flat_map(|(&s, &e)| unsafe {
Interval::new_unchecked(s, e).iter()
})
}
pub fn to_vec(&self) -> Vec<RangeInclusive<T>> {
self.iter().map(Into::into).collect()
}
}
impl<T: Numerated, D: Into<IntervalIterator<T>>> FromIterator<D> for IntervalsTree<T> {
fn from_iter<I: IntoIterator<Item = D>>(iter: I) -> Self {
let mut tree = Self::new();
for interval in iter {
tree.insert(interval);
}
tree
}
}
#[allow(clippy::reversed_empty_ranges)]
#[cfg(test)]
mod tests {
use super::*;
use alloc::vec;
#[test]
fn insert() {
let mut tree = IntervalsTree::new();
assert!(tree.insert(Interval::try_from(1..=2).unwrap()));
assert_eq!(tree.to_vec(), vec![1..=2]);
let mut tree = IntervalsTree::new();
assert!(tree.insert(Interval::try_from(-1..=2).unwrap()));
assert!(tree.insert(Interval::try_from(4..=5).unwrap()));
assert_eq!(tree.to_vec(), vec![-1..=2, 4..=5]);
let mut tree = IntervalsTree::new();
assert!(tree.insert(Interval::try_from(-1..=2).unwrap()));
assert!(tree.insert(Interval::try_from(3..=4).unwrap()));
assert_eq!(tree.to_vec(), vec![-1..=4]);
let mut tree = IntervalsTree::new();
assert!(tree.insert(1));
assert!(tree.insert(2));
assert_eq!(tree.to_vec(), vec![1..=2]);
let mut tree = IntervalsTree::new();
assert!(tree.insert(Interval::try_from(-1..=3).unwrap()));
assert!(tree.insert(Interval::try_from(5..=7).unwrap()));
assert!(tree.insert(Interval::try_from(2..=6).unwrap()));
assert!(
!tree.insert(Interval::try_from(7..=7).unwrap()),
"Expected false, because point 7 already in tree"
);
assert!(tree.insert(Interval::try_from(19..=25).unwrap()));
assert_eq!(tree.to_vec(), vec![-1..=7, 19..=25]);
let mut tree = IntervalsTree::new();
assert!(tree.insert(Interval::try_from(-1..=3).unwrap()));
assert!(tree.insert(Interval::try_from(10..=14).unwrap()));
assert!(tree.insert(Interval::try_from(4..=9).unwrap()));
assert_eq!(tree.to_vec(), vec![-1..=14]);
let mut tree = IntervalsTree::new();
assert!(tree.insert(Interval::try_from(-111..=3).unwrap()));
assert!(tree.insert(Interval::try_from(10..=14).unwrap()));
assert!(tree.insert(Interval::try_from(3..=10).unwrap()));
assert_eq!(tree.to_vec(), vec![-111..=14]);
let mut tree = IntervalsTree::new();
assert!(tree.insert(..=10));
assert!(
!tree.insert(Interval::try_from(3..=4).unwrap()),
"Expected false, because no unique points to insert in 3..=4"
);
assert_eq!(tree.to_vec(), vec![i32::MIN..=10]);
let mut tree = IntervalsTree::new();
assert!(tree.insert(Interval::try_from(1..=10).unwrap()));
assert!(
!tree.insert(Interval::try_from(3..=4).unwrap()),
"Expected false, because non-empty interval has no unique points to insert in 3..=4"
);
assert!(
!tree.insert(Interval::try_from(5..=6).unwrap()),
"Expected false, because non-empty interval has no unique points to insert in 5..=6"
);
assert_eq!(tree.to_vec(), vec![1..=10]);
let mut tree = IntervalsTree::new();
assert_eq!(tree.to_vec(), vec![]);
assert!(tree.insert(0..));
assert_eq!(tree.to_vec(), vec![0..=u32::MAX]);
let mut tree = IntervalsTree::<i32>::new();
assert!(
!tree.insert(IntervalIterator::empty()),
"Expected false, because empty interval don't change self"
);
}
#[test]
fn remove() {
let mut tree: IntervalsTree<i32> = [1].into_iter().collect();
assert!(tree.remove(1));
assert_eq!(tree.to_vec(), vec![]);
let mut tree: IntervalsTree<i32> = [1, 2].into_iter().collect();
assert!(tree.remove(Interval::try_from(1..=2).unwrap()));
assert_eq!(tree.to_vec(), vec![]);
let mut tree: IntervalsTree<i32> = [-1, 0, 1, 2, 4, 5].into_iter().collect();
assert!(tree.remove(Interval::try_from(-1..=2).unwrap()));
assert_eq!(tree.to_vec(), vec![4..=5]);
let mut tree: IntervalsTree<i32> = [-1, 0, 1, 2, 4, 5].into_iter().collect();
tree.remove(Interval::try_from(4..=5).unwrap());
assert_eq!(tree.to_vec(), vec![-1..=2]);
let mut tree: IntervalsTree<i32> = [1, 2, 4, 5].into_iter().collect();
assert!(tree.remove(Interval::try_from(2..=4).unwrap()));
assert_eq!(tree.to_vec(), vec![1..=1, 5..=5]);
let mut tree: IntervalsTree<i32> = [-1, 0, 1, 2, 4, 5].into_iter().collect();
assert!(tree.remove(Interval::try_from(3..=4).unwrap()));
assert_eq!(tree.to_vec(), vec![-1..=2, 5..=5]);
let mut tree: IntervalsTree<i32> = [-1, 0, 1, 2, 4, 5].into_iter().collect();
assert!(tree.remove(Interval::try_from(-1..=5).unwrap()));
assert_eq!(tree.to_vec(), vec![]);
let mut tree: IntervalsTree<i32> = [1, 2, 4, 5].into_iter().collect();
assert!(tree.remove(Interval::try_from(2..=5).unwrap()));
assert_eq!(tree.to_vec(), vec![1..=1]);
let mut tree: IntervalsTree<i32> = [1, 2, 4, 5].into_iter().collect();
assert!(tree.remove(Interval::try_from(1..=4).unwrap()));
assert_eq!(tree.to_vec(), vec![5..=5]);
let mut tree: IntervalsTree<i32> = [1, 2, 4, 5].into_iter().collect();
assert!(
!tree.remove(Interval::try_from(3..=3).unwrap()),
"Expected false, because there is no point 3 in tree"
);
assert!(tree.remove(Interval::try_from(1..=3).unwrap()));
assert_eq!(tree.to_vec(), vec![4..=5]);
assert_eq!(tree.to_vec(), vec![4..=5]);
let mut tree: IntervalsTree<u32> = [1, 2, 5, 6, 7, 9, 10, 11].into_iter().collect();
assert_eq!(tree.to_vec(), vec![1..=2, 5..=7, 9..=11]);
assert!(tree.remove(Interval::try_from(1..2).unwrap()));
assert_eq!(tree.to_vec(), vec![2..=2, 5..=7, 9..=11]);
assert!(tree.remove(..7));
assert_eq!(tree.to_vec(), vec![7..=7, 9..=11]);
assert!(tree.remove(..));
assert_eq!(tree.to_vec(), vec![]);
let mut tree = IntervalsTree::<i32>::new();
assert!(
!tree.remove(IntervalIterator::empty()),
"Expected false, because empty interval don't change self"
);
}
#[test]
fn voids() {
let tree: IntervalsTree<u32> = [1..=7, 19..=25]
.into_iter()
.map(|i| Interval::try_from(i).unwrap())
.collect();
assert_eq!(
tree.voids(Interval::try_from(0..100).unwrap())
.map(RangeInclusive::from)
.collect::<Vec<_>>(),
vec![0..=0, 8..=18, 26..=99],
);
assert_eq!(
tree.voids(..).map(RangeInclusive::from).collect::<Vec<_>>(),
vec![0..=0, 8..=18, 26..=u32::MAX],
);
assert_eq!(
tree.voids(IntervalIterator::empty()).collect::<Vec<_>>(),
Vec::<RangeInclusive<_>>::new()
);
assert_eq!(
tree.voids(0).map(RangeInclusive::from).collect::<Vec<_>>(),
vec![0..=0],
);
}
#[test]
fn contains() {
let tree: IntervalsTree<u64> = [0, 100, 101, 102, 45678, 45679, 1, 2, 3]
.into_iter()
.collect();
assert_eq!(tree.to_vec(), vec![0..=3, 100..=102, 45678..=45679]);
assert!(tree.contains(0));
assert!(!tree.contains(4));
assert!(tree.contains(100));
assert!(!tree.contains(103));
assert!(tree.contains(45678));
assert!(!tree.contains(45680));
assert!(tree.contains(Interval::try_from(0..=3).unwrap()));
assert!(!tree.contains(Interval::try_from(0..5).unwrap()));
assert!(tree.contains(IntervalIterator::empty()));
assert!(!tree.contains(..));
assert!(tree.contains(..1));
}
#[test]
fn amount() {
let tree: IntervalsTree<i32> = [-100, -99, 100, 101, 102, 1000].into_iter().collect();
assert_eq!(tree.intervals_amount(), 3);
assert_eq!(tree.points_amount(), Some(6));
let tree: IntervalsTree<i32> = [..].into_iter().collect();
assert_eq!(tree.intervals_amount(), 1);
assert_eq!(tree.points_amount(), None);
let tree: IntervalsTree<i32> = Default::default();
assert_eq!(tree.intervals_amount(), 0);
assert_eq!(tree.points_amount(), Some(0));
}
#[test]
fn start_end() {
let tree: IntervalsTree<u64> = [0u64, 100, 101, 102, 45678, 45679, 1, 2, 3]
.into_iter()
.collect();
assert_eq!(tree.to_vec(), vec![0..=3, 100..=102, 45678..=45679]);
assert_eq!(tree.start(), Some(0));
assert_eq!(tree.end(), Some(45679));
}
#[test]
fn difference() {
let tree: IntervalsTree<u64> = [0, 1, 2, 3, 4, 8, 9, 100, 101, 102].into_iter().collect();
let tree1: IntervalsTree<u64> = [3, 4, 7, 8, 9, 10, 45, 46, 100, 102].into_iter().collect();
let v: Vec<RangeInclusive<u64>> = tree.difference(&tree1).map(Into::into).collect();
assert_eq!(v, vec![0..=2, 101..=101]);
let tree1: IntervalsTree<u64> = [..].into_iter().collect();
let v: Vec<RangeInclusive<u64>> = tree.difference(&tree1).map(Into::into).collect();
assert_eq!(v, vec![]);
let tree1: IntervalsTree<u64> = [..=100].into_iter().collect();
let v: Vec<RangeInclusive<u64>> = tree.difference(&tree1).map(Into::into).collect();
assert_eq!(v, vec![101..=102]);
let tree1: IntervalsTree<u64> = [101..].into_iter().collect();
let v: Vec<RangeInclusive<u64>> = tree.difference(&tree1).map(Into::into).collect();
assert_eq!(v, vec![0..=4, 8..=9, 100..=100]);
let tree1: IntervalsTree<u64> = [6, 10, 110].into_iter().collect();
let v: Vec<RangeInclusive<u64>> = tree.difference(&tree1).map(Into::into).collect();
assert_eq!(v, vec![0..=4, 8..=9, 100..=102]);
}
}