1use crate::{
22 interval::Interval,
23 iterators::{DifferenceIterator, IntervalIterator, VoidsIterator},
24 numerated::Numerated,
25};
26use alloc::{collections::BTreeMap, fmt, fmt::Debug, vec::Vec};
27use core::{fmt::Formatter, ops::RangeInclusive};
28use num_traits::{CheckedAdd, Zero};
29use scale_info::{
30 TypeInfo,
31 scale::{Decode, Encode},
32};
33
34#[derive(Clone, PartialEq, Eq, Hash, TypeInfo, Encode, Decode)]
99#[codec(crate = scale_info::scale)]
100pub struct IntervalsTree<T> {
101 inner: BTreeMap<T, T>,
102}
103
104impl<T: Copy> IntervalsTree<T> {
105 pub const fn new() -> Self {
107 Self {
108 inner: BTreeMap::new(),
109 }
110 }
111 pub fn intervals_amount(&self) -> usize {
115 self.inner.len()
116 }
117
118 pub fn is_empty(&self) -> bool {
121 self.inner.is_empty()
122 }
123}
124
125impl<T: Copy + Ord> IntervalsTree<T> {
126 pub fn end(&self) -> Option<T> {
128 self.inner.last_key_value().map(|(_, &e)| e)
129 }
130 pub fn start(&self) -> Option<T> {
132 self.inner.first_key_value().map(|(&s, _)| s)
133 }
134}
135
136impl<T: Copy> Default for IntervalsTree<T> {
137 fn default() -> Self {
138 Self::new()
139 }
140}
141
142impl<T: Debug + Numerated> Debug for IntervalsTree<T> {
143 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
144 f.debug_list().entries(self.iter()).finish()
145 }
146}
147
148impl<T: Numerated> IntervalsTree<T> {
149 fn into_start_end<I: Into<IntervalIterator<T>>>(interval: I) -> Option<(T, T)> {
150 Into::<IntervalIterator<T>>::into(interval)
151 .inner()
152 .map(|i| i.into_parts())
153 }
154
155 #[track_caller]
156 fn put(&mut self, start: T, end: T) {
157 debug_assert!(start <= end, "Must be guarantied");
158 self.inner.insert(start, end);
159 }
160
161 pub fn iter(&self) -> impl Iterator<Item = Interval<T>> + '_ {
163 self.inner.iter().map(|(&start, &end)| unsafe {
164 Interval::<T>::new_unchecked(start, end)
166 })
167 }
168
169 pub fn contains<I: Into<IntervalIterator<T>>>(&self, interval: I) -> bool {
171 let Some((start, end)) = Self::into_start_end(interval) else {
172 return true;
174 };
175 self.inner
176 .range(..=end)
177 .next_back()
178 .map(|(&s, &e)| s <= start && end <= e)
179 .unwrap_or(false)
180 }
181
182 pub fn insert<I: Into<IntervalIterator<T>>>(&mut self, interval: I) -> bool {
192 let Some((start, end)) = Self::into_start_end(interval) else {
193 return false;
195 };
196
197 let Some(last) = self.end() else {
198 self.put(start, end);
200 return true;
201 };
202
203 let iter_end = end.inc_if_lt(last).unwrap_or(end);
206 let mut iter = self.inner.range(..=iter_end).map(|(&s, &e)| (s, e));
207
208 let Some((right_start, right_end)) = iter.next_back() else {
213 self.put(start, end);
215 return true;
216 };
217
218 if let Some(right_end) = right_end.inc_if_lt(start) {
219 if right_end == start {
221 self.put(right_start, end);
223 } else {
224 self.put(start, end);
226 }
227 return true;
228 } else if right_start <= start {
229 if right_end < end {
230 self.put(right_start, end);
232 return true;
233 } else {
234 return false;
236 }
237 }
238
239 let mut left_interval = None;
242 let mut intervals_to_remove = Vec::new();
243 while let Some((s, e)) = iter.next_back() {
244 if s <= start {
245 left_interval = Some((s, e));
246 break;
247 }
248 intervals_to_remove.push(s);
249 }
250
251 for start in intervals_to_remove {
254 self.inner.remove(&start);
255 }
256
257 self.inner.remove(&right_start);
259
260 let end = right_end.max(end);
261
262 let Some((left_start, left_end)) = left_interval else {
263 self.put(start, end);
265 return true;
266 };
267
268 debug_assert!(left_end < right_start && left_start <= start);
269 let Some(left_end) = left_end.inc_if_lt(right_start) else {
270 debug_assert!(false, "`T: Numerated` impl error");
272 return false;
273 };
274
275 if left_end >= start {
276 self.put(left_start, end);
278 } else {
279 self.put(start, end);
281 }
282
283 true
285 }
286
287 pub fn remove<I: Into<IntervalIterator<T>>>(&mut self, interval: I) -> bool {
297 let Some((start, end)) = Self::into_start_end(interval) else {
298 return false;
300 };
301
302 let mut iter = self.inner.range(..=end);
304
305 let Some((&right_start, &right_end)) = iter.next_back() else {
307 return false;
308 };
309
310 if right_end < start {
311 return false;
313 }
314
315 let mut left_interval = None;
318 let mut intervals_to_remove = Vec::new();
319 while let Some((&s, &e)) = iter.next_back() {
320 if s < start {
321 if e >= start {
322 left_interval = Some(s);
323 }
324 break;
325 }
326
327 intervals_to_remove.push(s)
328 }
329
330 for start in intervals_to_remove {
333 self.inner.remove(&start);
334 }
335
336 if let Some(start) = start.dec_if_gt(right_start) {
337 self.put(right_start, start);
340 } else {
341 debug_assert!(
344 right_start <= end,
345 "Must be, because of method it was found"
346 );
347 self.inner.remove(&right_start);
348 }
349
350 if let Some(end) = end.inc_if_lt(right_end) {
351 self.put(end, right_end);
354 } else {
355 debug_assert!(start <= right_end);
356 }
357
358 if let Some(left_start) = left_interval {
359 if let Some(start) = start.dec_if_gt(left_start) {
362 self.put(left_start, start);
363 } else {
364 debug_assert!(false, "`T: Numerated` impl error");
365 }
366 }
367
368 true
371 }
372
373 pub fn voids<I: Into<IntervalIterator<T>>>(
381 &self,
382 interval: I,
383 ) -> VoidsIterator<T, impl Iterator<Item = Interval<T>> + '_> {
384 let Some((mut start, end)) = Self::into_start_end(interval) else {
385 return VoidsIterator { inner: None };
387 };
388
389 if let Some((_, &e)) = self.inner.range(..=start).next_back() {
390 if let Some(e) = e.inc_if_lt(end) {
391 if e > start {
392 start = e;
393 }
394 } else {
395 return VoidsIterator { inner: None };
397 }
398 }
399
400 let iter = self.inner.range(start..=end).map(|(&start, &end)| {
401 unsafe { Interval::new_unchecked(start, end) }
403 });
404
405 let interval = unsafe { Interval::new_unchecked(start, end) };
407
408 VoidsIterator {
409 inner: Some((iter, interval)),
410 }
411 }
412
413 pub fn difference<'a>(&'a self, other: &'a Self) -> impl Iterator<Item = Interval<T>> + 'a {
420 DifferenceIterator {
421 iter1: self.iter(),
422 iter2: other.iter(),
423 interval1: None,
424 interval2: None,
425 }
426 }
427
428 pub fn points_amount(&self) -> Option<T::Distance> {
432 let mut res = T::Distance::zero();
433 for interval in self.iter() {
434 res = res.checked_add(&interval.raw_len()?)?;
435 }
436 Some(res)
437 }
438
439 pub fn points_iter(&self) -> impl Iterator<Item = T> + '_ {
441 self.inner.iter().flat_map(|(&s, &e)| unsafe {
442 Interval::new_unchecked(s, e).iter()
444 })
445 }
446
447 pub fn to_vec(&self) -> Vec<RangeInclusive<T>> {
449 self.iter().map(Into::into).collect()
450 }
451}
452
453impl<T: Numerated, D: Into<IntervalIterator<T>>> FromIterator<D> for IntervalsTree<T> {
454 fn from_iter<I: IntoIterator<Item = D>>(iter: I) -> Self {
455 let mut tree = Self::new();
456 for interval in iter {
457 tree.insert(interval);
458 }
459 tree
460 }
461}
462
463#[allow(clippy::reversed_empty_ranges)]
464#[cfg(test)]
465mod tests {
466 use super::*;
467 use alloc::vec;
468
469 #[test]
470 fn insert() {
471 let mut tree = IntervalsTree::new();
472 assert!(tree.insert(Interval::try_from(1..=2).unwrap()));
473 assert_eq!(tree.to_vec(), vec![1..=2]);
474
475 let mut tree = IntervalsTree::new();
476 assert!(tree.insert(Interval::try_from(-1..=2).unwrap()));
477 assert!(tree.insert(Interval::try_from(4..=5).unwrap()));
478 assert_eq!(tree.to_vec(), vec![-1..=2, 4..=5]);
479
480 let mut tree = IntervalsTree::new();
481 assert!(tree.insert(Interval::try_from(-1..=2).unwrap()));
482 assert!(tree.insert(Interval::try_from(3..=4).unwrap()));
483 assert_eq!(tree.to_vec(), vec![-1..=4]);
484
485 let mut tree = IntervalsTree::new();
486 assert!(tree.insert(1));
487 assert!(tree.insert(2));
488 assert_eq!(tree.to_vec(), vec![1..=2]);
489
490 let mut tree = IntervalsTree::new();
491 assert!(tree.insert(Interval::try_from(-1..=3).unwrap()));
492 assert!(tree.insert(Interval::try_from(5..=7).unwrap()));
493 assert!(tree.insert(Interval::try_from(2..=6).unwrap()));
494 assert!(
495 !tree.insert(Interval::try_from(7..=7).unwrap()),
496 "Expected false, because point 7 already in tree"
497 );
498 assert!(tree.insert(Interval::try_from(19..=25).unwrap()));
499 assert_eq!(tree.to_vec(), vec![-1..=7, 19..=25]);
500
501 let mut tree = IntervalsTree::new();
502 assert!(tree.insert(Interval::try_from(-1..=3).unwrap()));
503 assert!(tree.insert(Interval::try_from(10..=14).unwrap()));
504 assert!(tree.insert(Interval::try_from(4..=9).unwrap()));
505 assert_eq!(tree.to_vec(), vec![-1..=14]);
506
507 let mut tree = IntervalsTree::new();
508 assert!(tree.insert(Interval::try_from(-111..=3).unwrap()));
509 assert!(tree.insert(Interval::try_from(10..=14).unwrap()));
510 assert!(tree.insert(Interval::try_from(3..=10).unwrap()));
511 assert_eq!(tree.to_vec(), vec![-111..=14]);
512
513 let mut tree = IntervalsTree::new();
514 assert!(tree.insert(..=10));
515 assert!(
516 !tree.insert(Interval::try_from(3..=4).unwrap()),
517 "Expected false, because no unique points to insert in 3..=4"
518 );
519 assert_eq!(tree.to_vec(), vec![i32::MIN..=10]);
520
521 let mut tree = IntervalsTree::new();
522 assert!(tree.insert(Interval::try_from(1..=10).unwrap()));
523 assert!(
524 !tree.insert(Interval::try_from(3..=4).unwrap()),
525 "Expected false, because non-empty interval has no unique points to insert in 3..=4"
526 );
527 assert!(
528 !tree.insert(Interval::try_from(5..=6).unwrap()),
529 "Expected false, because non-empty interval has no unique points to insert in 5..=6"
530 );
531 assert_eq!(tree.to_vec(), vec![1..=10]);
532
533 let mut tree = IntervalsTree::new();
534 assert_eq!(tree.to_vec(), vec![]);
535 assert!(tree.insert(0..));
536 assert_eq!(tree.to_vec(), vec![0..=u32::MAX]);
537
538 let mut tree = IntervalsTree::<i32>::new();
539 assert!(
540 !tree.insert(IntervalIterator::empty()),
541 "Expected false, because empty interval don't change self"
542 );
543 }
544
545 #[test]
546 fn remove() {
547 let mut tree: IntervalsTree<i32> = [1].into_iter().collect();
548 assert!(tree.remove(1));
549 assert_eq!(tree.to_vec(), vec![]);
550
551 let mut tree: IntervalsTree<i32> = [1, 2].into_iter().collect();
552 assert!(tree.remove(Interval::try_from(1..=2).unwrap()));
553 assert_eq!(tree.to_vec(), vec![]);
554
555 let mut tree: IntervalsTree<i32> = [-1, 0, 1, 2, 4, 5].into_iter().collect();
556 assert!(tree.remove(Interval::try_from(-1..=2).unwrap()));
557 assert_eq!(tree.to_vec(), vec![4..=5]);
558
559 let mut tree: IntervalsTree<i32> = [-1, 0, 1, 2, 4, 5].into_iter().collect();
560 tree.remove(Interval::try_from(4..=5).unwrap());
561 assert_eq!(tree.to_vec(), vec![-1..=2]);
562
563 let mut tree: IntervalsTree<i32> = [1, 2, 4, 5].into_iter().collect();
564 assert!(tree.remove(Interval::try_from(2..=4).unwrap()));
565 assert_eq!(tree.to_vec(), vec![1..=1, 5..=5]);
566
567 let mut tree: IntervalsTree<i32> = [-1, 0, 1, 2, 4, 5].into_iter().collect();
568 assert!(tree.remove(Interval::try_from(3..=4).unwrap()));
569 assert_eq!(tree.to_vec(), vec![-1..=2, 5..=5]);
570
571 let mut tree: IntervalsTree<i32> = [-1, 0, 1, 2, 4, 5].into_iter().collect();
572 assert!(tree.remove(Interval::try_from(-1..=5).unwrap()));
573 assert_eq!(tree.to_vec(), vec![]);
574
575 let mut tree: IntervalsTree<i32> = [1, 2, 4, 5].into_iter().collect();
576 assert!(tree.remove(Interval::try_from(2..=5).unwrap()));
577 assert_eq!(tree.to_vec(), vec![1..=1]);
578
579 let mut tree: IntervalsTree<i32> = [1, 2, 4, 5].into_iter().collect();
580 assert!(tree.remove(Interval::try_from(1..=4).unwrap()));
581 assert_eq!(tree.to_vec(), vec![5..=5]);
582
583 let mut tree: IntervalsTree<i32> = [1, 2, 4, 5].into_iter().collect();
584 assert!(
585 !tree.remove(Interval::try_from(3..=3).unwrap()),
586 "Expected false, because there is no point 3 in tree"
587 );
588 assert!(tree.remove(Interval::try_from(1..=3).unwrap()));
589 assert_eq!(tree.to_vec(), vec![4..=5]);
590 assert_eq!(tree.to_vec(), vec![4..=5]);
591
592 let mut tree: IntervalsTree<u32> = [1, 2, 5, 6, 7, 9, 10, 11].into_iter().collect();
593 assert_eq!(tree.to_vec(), vec![1..=2, 5..=7, 9..=11]);
594 assert!(tree.remove(Interval::try_from(1..2).unwrap()));
595 assert_eq!(tree.to_vec(), vec![2..=2, 5..=7, 9..=11]);
596 assert!(tree.remove(..7));
597 assert_eq!(tree.to_vec(), vec![7..=7, 9..=11]);
598 assert!(tree.remove(..));
599 assert_eq!(tree.to_vec(), vec![]);
600
601 let mut tree = IntervalsTree::<i32>::new();
602 assert!(
603 !tree.remove(IntervalIterator::empty()),
604 "Expected false, because empty interval don't change self"
605 );
606 }
607
608 #[test]
609 fn voids() {
610 let tree: IntervalsTree<u32> = [1..=7, 19..=25]
611 .into_iter()
612 .map(|i| Interval::try_from(i).unwrap())
613 .collect();
614
615 assert_eq!(
616 tree.voids(Interval::try_from(0..100).unwrap())
617 .map(RangeInclusive::from)
618 .collect::<Vec<_>>(),
619 vec![0..=0, 8..=18, 26..=99],
620 );
621 assert_eq!(
622 tree.voids(..).map(RangeInclusive::from).collect::<Vec<_>>(),
623 vec![0..=0, 8..=18, 26..=u32::MAX],
624 );
625 assert_eq!(
626 tree.voids(IntervalIterator::empty()).collect::<Vec<_>>(),
627 Vec::<RangeInclusive<_>>::new()
628 );
629 assert_eq!(
630 tree.voids(0).map(RangeInclusive::from).collect::<Vec<_>>(),
631 vec![0..=0],
632 );
633 }
634
635 #[test]
636 fn contains() {
637 let tree: IntervalsTree<u64> = [0, 100, 101, 102, 45678, 45679, 1, 2, 3]
638 .into_iter()
639 .collect();
640 assert_eq!(tree.to_vec(), vec![0..=3, 100..=102, 45678..=45679]);
641 assert!(tree.contains(0));
642 assert!(!tree.contains(4));
643 assert!(tree.contains(100));
644 assert!(!tree.contains(103));
645 assert!(tree.contains(45678));
646 assert!(!tree.contains(45680));
647 assert!(tree.contains(Interval::try_from(0..=3).unwrap()));
648 assert!(!tree.contains(Interval::try_from(0..5).unwrap()));
649 assert!(tree.contains(IntervalIterator::empty()));
650 assert!(!tree.contains(..));
651 assert!(tree.contains(..1));
652 }
653
654 #[test]
655 fn amount() {
656 let tree: IntervalsTree<i32> = [-100, -99, 100, 101, 102, 1000].into_iter().collect();
657 assert_eq!(tree.intervals_amount(), 3);
658 assert_eq!(tree.points_amount(), Some(6));
659
660 let tree: IntervalsTree<i32> = [..].into_iter().collect();
661 assert_eq!(tree.intervals_amount(), 1);
662 assert_eq!(tree.points_amount(), None);
663
664 let tree: IntervalsTree<i32> = Default::default();
665 assert_eq!(tree.intervals_amount(), 0);
666 assert_eq!(tree.points_amount(), Some(0));
667 }
668
669 #[test]
670 fn start_end() {
671 let tree: IntervalsTree<u64> = [0u64, 100, 101, 102, 45678, 45679, 1, 2, 3]
672 .into_iter()
673 .collect();
674 assert_eq!(tree.to_vec(), vec![0..=3, 100..=102, 45678..=45679]);
675 assert_eq!(tree.start(), Some(0));
676 assert_eq!(tree.end(), Some(45679));
677 }
678
679 #[test]
680 fn difference() {
681 let tree: IntervalsTree<u64> = [0, 1, 2, 3, 4, 8, 9, 100, 101, 102].into_iter().collect();
682 let tree1: IntervalsTree<u64> = [3, 4, 7, 8, 9, 10, 45, 46, 100, 102].into_iter().collect();
683 let v: Vec<RangeInclusive<u64>> = tree.difference(&tree1).map(Into::into).collect();
684 assert_eq!(v, vec![0..=2, 101..=101]);
685
686 let tree1: IntervalsTree<u64> = [..].into_iter().collect();
687 let v: Vec<RangeInclusive<u64>> = tree.difference(&tree1).map(Into::into).collect();
688 assert_eq!(v, vec![]);
689
690 let tree1: IntervalsTree<u64> = [..=100].into_iter().collect();
691 let v: Vec<RangeInclusive<u64>> = tree.difference(&tree1).map(Into::into).collect();
692 assert_eq!(v, vec![101..=102]);
693
694 let tree1: IntervalsTree<u64> = [101..].into_iter().collect();
695 let v: Vec<RangeInclusive<u64>> = tree.difference(&tree1).map(Into::into).collect();
696 assert_eq!(v, vec![0..=4, 8..=9, 100..=100]);
697
698 let tree1: IntervalsTree<u64> = [6, 10, 110].into_iter().collect();
699 let v: Vec<RangeInclusive<u64>> = tree.difference(&tree1).map(Into::into).collect();
700 assert_eq!(v, vec![0..=4, 8..=9, 100..=102]);
701 }
702}