mirror of
https://github.com/n0-computer/noq.git
synced 2026-09-22 03:03:39 +00:00
Benjamin Saunders
288 lines
8.4 KiB
Rust
288 lines
8.4 KiB
Rust
use std::collections::BTreeMap;
|
|
use std::collections::Bound::{Included, Excluded};
|
|
use std::collections::btree_map;
|
|
use std::ops::Range;
|
|
use std::cmp;
|
|
|
|
/// A set of u64 values optimized for long runs and random insert/delete/contains
|
|
#[derive(Debug, Clone)]
|
|
pub struct RangeSet(BTreeMap<u64, u64>);
|
|
|
|
impl RangeSet {
|
|
pub fn new() -> Self { RangeSet(BTreeMap::new()) }
|
|
|
|
pub fn contains(&self, x: u64) -> bool {
|
|
self.pred(x).map_or(false, |(_, end)| end > x)
|
|
}
|
|
|
|
pub fn insert_one(&mut self, x: u64) -> bool {
|
|
if let Some((start, end)) = self.pred(x) {
|
|
if end > x {
|
|
// Wholly contained
|
|
return false;
|
|
} else if end == x {
|
|
// Extend existing
|
|
self.0.remove(&start);
|
|
let mut new_end = x+1;
|
|
if let Some((next_start, next_end)) = self.succ(x) {
|
|
if next_start == new_end {
|
|
self.0.remove(&next_start);
|
|
new_end = next_end;
|
|
}
|
|
}
|
|
self.0.insert(start, new_end);
|
|
return true;
|
|
}
|
|
}
|
|
let mut new_end = x+1;
|
|
if let Some((next_start, next_end)) = self.succ(x) {
|
|
if next_start == new_end {
|
|
self.0.remove(&next_start);
|
|
new_end = next_end;
|
|
}
|
|
}
|
|
self.0.insert(x, new_end);
|
|
true
|
|
}
|
|
|
|
pub fn insert(&mut self, mut x: Range<u64>) -> bool {
|
|
if let Some((start, end)) = self.pred(x.start) {
|
|
if end >= x.end {
|
|
// Wholly contained
|
|
return false;
|
|
} else if end >= x.start {
|
|
// Overlaps with pred
|
|
self.0.remove(&start);
|
|
while let Some((next_start, next_end)) = self.succ(x.start) {
|
|
if next_start > x.end { break; }
|
|
// ..and succ
|
|
self.0.remove(&next_start);
|
|
x.end = cmp::max(next_end, x.end);
|
|
}
|
|
self.0.insert(start, x.end);
|
|
return true;
|
|
}
|
|
}
|
|
while let Some((next_start, next_end)) = self.succ(x.start) {
|
|
if next_start > x.end { break; }
|
|
// Overlaps with succ
|
|
self.0.remove(&next_start);
|
|
x.end = cmp::max(next_end, x.end);
|
|
}
|
|
self.0.insert(x.start, x.end);
|
|
true
|
|
}
|
|
|
|
fn pred(&self, x: u64) -> Option<(u64, u64)> {
|
|
self.0.range((Included(0), Included(x))).rev().next().map(|(&x, &y)| (x, y))
|
|
}
|
|
|
|
fn succ(&self, x: u64) -> Option<(u64, u64)> {
|
|
self.0.range((Excluded(x), Included(u64::max_value()))).next().map(|(&x, &y)| (x, y))
|
|
}
|
|
|
|
pub fn remove(&mut self, x: Range<u64>) -> bool {
|
|
let before = if let Some((start, end)) = self.pred(x.start) {
|
|
if end > x.start {
|
|
self.0.remove(&start);
|
|
if start < x.start {
|
|
self.0.insert(start, x.start);
|
|
}
|
|
if end > x.end {
|
|
self.0.insert(x.end, end);
|
|
}
|
|
// Short-circuit if we cannot possibly overlap with another range
|
|
if end >= x.end { return true; }
|
|
true
|
|
} else { false }
|
|
} else { false };
|
|
let mut after = false;;
|
|
while let Some((start, end)) = self.succ(x.start) {
|
|
if start >= x.end { break; }
|
|
after = true;
|
|
self.0.remove(&start);
|
|
if end > x.end {
|
|
self.0.insert(x.end, end);
|
|
break;
|
|
}
|
|
};
|
|
before || after
|
|
}
|
|
|
|
pub fn add(&mut self, other: &RangeSet) {
|
|
for (&start, &end) in &other.0 {
|
|
self.insert(start..end);
|
|
}
|
|
}
|
|
|
|
pub fn subtract(&mut self, other: &RangeSet) {
|
|
for (&start, &end) in &other.0 {
|
|
self.remove(start..end);
|
|
}
|
|
}
|
|
|
|
pub fn is_empty(&self) -> bool { self.0.is_empty() }
|
|
|
|
pub fn min(&self) -> Option<u64> { self.iter().next().map(|x| x.start) }
|
|
pub fn max(&self) -> Option<u64> { self.iter().rev().next().map(|x| x.end-1) }
|
|
|
|
pub fn len(&self) -> usize { self.0.len() }
|
|
pub fn iter(&self) -> Iter { Iter(self.0.iter()) }
|
|
pub fn elts(&self) -> EltIter { EltIter { inner: self.0.iter(), next: 0, end: 0 } }
|
|
|
|
pub fn pop_min(&mut self) -> Option<Range<u64>> {
|
|
let (&start, &end) = self.0.iter().next()?;
|
|
self.0.remove(&start);
|
|
Some(start..end)
|
|
}
|
|
}
|
|
|
|
pub struct Iter<'a>(btree_map::Iter<'a, u64, u64>);
|
|
|
|
impl<'a> Iterator for Iter<'a> {
|
|
type Item = Range<u64>;
|
|
fn next(&mut self) -> Option<Range<u64>> {
|
|
let (&start, &end) = self.0.next()?;
|
|
Some(start..end)
|
|
}
|
|
}
|
|
|
|
impl<'a> DoubleEndedIterator for Iter<'a> {
|
|
fn next_back(&mut self) -> Option<Range<u64>> {
|
|
let (&start, &end) = self.0.next_back()?;
|
|
Some(start..end)
|
|
}
|
|
}
|
|
|
|
impl<'a> IntoIterator for &'a RangeSet {
|
|
type Item = Range<u64>;
|
|
type IntoIter = Iter<'a>;
|
|
fn into_iter(self) -> Iter<'a> { self.iter() }
|
|
}
|
|
|
|
pub struct EltIter<'a> {
|
|
inner: btree_map::Iter<'a, u64, u64>,
|
|
next: u64,
|
|
end: u64,
|
|
}
|
|
|
|
impl<'a> Iterator for EltIter<'a> {
|
|
type Item = u64;
|
|
fn next(&mut self) -> Option<u64> {
|
|
if self.next == self.end {
|
|
let (&start, &end) = self.inner.next()?;
|
|
self.next = start;
|
|
self.end = end;
|
|
}
|
|
let x = self.next;
|
|
self.next += 1;
|
|
Some(x)
|
|
}
|
|
}
|
|
|
|
impl<'a> DoubleEndedIterator for EltIter<'a> {
|
|
fn next_back(&mut self) -> Option<u64> {
|
|
if self.next == self.end {
|
|
let (&start, &end) = self.inner.next_back()?;
|
|
self.next = start;
|
|
self.end = end;
|
|
}
|
|
self.end -= 1;
|
|
Some(self.end)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn merge_and_split() {
|
|
let mut set = RangeSet::new();
|
|
assert!(set.insert(0..2));
|
|
assert!(set.insert(2..4));
|
|
assert!(!set.insert(1..3));
|
|
assert_eq!(set.len(), 1);
|
|
assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 2, 3]);
|
|
assert!(!set.contains(4));
|
|
assert!(set.remove(2..3));
|
|
assert_eq!(set.len(), 2);
|
|
assert!(!set.contains(2));
|
|
assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 3]);
|
|
}
|
|
|
|
#[test]
|
|
fn double_merge_exact() {
|
|
let mut set = RangeSet::new();
|
|
assert!(set.insert(0..2));
|
|
assert!(set.insert(4..6));
|
|
assert_eq!(set.len(), 2);
|
|
assert!(set.insert(2..4));
|
|
assert_eq!(set.len(), 1);
|
|
assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 2, 3, 4, 5]);
|
|
}
|
|
|
|
#[test]
|
|
fn single_merge_low() {
|
|
let mut set = RangeSet::new();
|
|
assert!(set.insert(0..2));
|
|
assert!(set.insert(4..6));
|
|
assert_eq!(set.len(), 2);
|
|
assert!(set.insert(2..3));
|
|
assert_eq!(set.len(), 2);
|
|
assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 2, 4, 5]);
|
|
}
|
|
|
|
#[test]
|
|
fn single_merge_high() {
|
|
let mut set = RangeSet::new();
|
|
assert!(set.insert(0..2));
|
|
assert!(set.insert(4..6));
|
|
assert_eq!(set.len(), 2);
|
|
assert!(set.insert(3..4));
|
|
assert_eq!(set.len(), 2);
|
|
assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 3, 4, 5]);
|
|
}
|
|
|
|
#[test]
|
|
fn double_merge_wide() {
|
|
let mut set = RangeSet::new();
|
|
assert!(set.insert(0..2));
|
|
assert!(set.insert(4..6));
|
|
assert_eq!(set.len(), 2);
|
|
assert!(set.insert(1..5));
|
|
assert_eq!(set.len(), 1);
|
|
assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 2, 3, 4, 5]);
|
|
}
|
|
|
|
#[test]
|
|
fn double_remove() {
|
|
let mut set = RangeSet::new();
|
|
assert!(set.insert(0..2));
|
|
assert!(set.insert(4..6));
|
|
assert!(set.remove(1..5));
|
|
assert_eq!(set.len(), 2);
|
|
assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 5]);
|
|
}
|
|
|
|
#[test]
|
|
fn insert_multiple() {
|
|
let mut set = RangeSet::new();
|
|
assert!(set.insert(0..1));
|
|
assert!(set.insert(2..3));
|
|
assert!(set.insert(4..5));
|
|
assert!(set.insert(0..5));
|
|
assert_eq!(set.len(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn remove_multiple() {
|
|
let mut set = RangeSet::new();
|
|
assert!(set.insert(0..1));
|
|
assert!(set.insert(2..3));
|
|
assert!(set.insert(4..5));
|
|
assert!(set.remove(0..5));
|
|
assert!(set.is_empty());
|
|
}
|
|
}
|