use std::{
collections::BTreeMap,
ops::{Index, IndexMut},
};
use thiserror::Error;
#[derive(Clone, Error, Debug, PartialEq)]
pub enum RangeMapError {
#[error("Invalid range provided")]
InvalidRange,
#[error("Overlapping ranges")]
RangeOverlap,
}
#[derive(Clone, Debug)]
pub struct RangeMap<RK: Ord + Copy, V> {
ranges: BTreeMap<RK, (RK, V)>, }
impl<RK: Ord + Copy, V> RangeMap<RK, V> {
pub fn new() -> Self {
Self {
ranges: BTreeMap::new(),
}
}
pub fn insert(&mut self, start: RK, end: RK, value: V) -> Result<(), RangeMapError> {
if end < start {
return Err(RangeMapError::InvalidRange);
}
if let Some((&existing_start, &(existing_end, _))) = self.ranges.range(..=end).next_back() {
if (existing_start <= start && start <= existing_end)
|| (existing_start <= end && end <= existing_end)
{
return Err(RangeMapError::RangeOverlap);
}
}
self.ranges.insert(start, (end, value));
Ok(())
}
pub fn get(&self, key: (RK, RK)) -> Option<&V> {
if let Some(entry) = self.ranges.get(&key.0) {
if entry.0 == key.1 {
return Some(&entry.1);
}
}
None
}
pub fn get_mut(&mut self, key: (RK, RK)) -> Option<&mut V> {
if let Some(entry) = self.ranges.get_mut(&key.0) {
if entry.0 == key.1 {
return Some(&mut entry.1);
}
}
None
}
pub fn lookup(&self, key: RK) -> Option<&V> {
if let Some((&_start, &(end, ref value))) = self.ranges.range(..=key).next_back() {
if key <= end {
return Some(value);
}
}
None
}
pub fn overlaps(&self, start: RK, end: RK) -> bool {
if let Some((&_existing_start, &(existing_end, _))) =
self.ranges.range(..=start).next_back()
{
if existing_end >= start {
return true;
}
}
if let Some((&existing_start, &(_, _))) = self.ranges.range(start..).next() {
if existing_start <= end {
return true;
}
}
false
}
pub fn remove(&mut self, key: RK) -> Option<V> {
if let Some((&start, &(end, ref _value))) = self.ranges.range(..=key).next_back() {
if key <= end {
return self.ranges.remove(&start).map(|entry| entry.1);
}
}
None
}
pub fn iter(&self) -> impl Iterator<Item = ((RK, RK), &V)> + '_ {
self.ranges
.iter()
.map(|(&s, value)| ((s, value.0), &value.1))
}
pub fn len(&self) -> usize {
self.ranges.len()
}
pub fn is_empty(&self) -> bool {
self.ranges.is_empty()
}
}
impl<RK: Ord + Copy, V> Default for RangeMap<RK, V> {
fn default() -> Self {
RangeMap::new()
}
}
impl<RK: Ord + Copy, V> Index<(RK, RK)> for RangeMap<RK, V> {
type Output = V;
fn index(&self, index: (RK, RK)) -> &Self::Output {
let entry = self.ranges.index(&index.0);
if entry.0 != index.1 {
panic!("no entry found for key");
}
&entry.1
}
}
impl<RK: Ord + Copy, V> IndexMut<(RK, RK)> for RangeMap<RK, V> {
fn index_mut(&mut self, index: (RK, RK)) -> &mut Self::Output {
let entry = self.ranges.get_mut(&index.0);
if let Some(entry) = entry {
if entry.0 != index.1 {
panic!("no entry found for key");
}
&mut entry.1
} else {
panic!("no entry found for key")
}
}
}
#[cfg(test)]
mod tests {
use crate::rangemap::RangeMapError;
use super::RangeMap;
#[test]
pub fn test_insertions() {
let mut rm: RangeMap<i64, &str> = RangeMap::new();
assert_eq!(rm.lookup(-1), None);
assert_eq!(rm.lookup(0), None);
assert_eq!(rm.lookup(6), None);
assert_eq!(rm.lookup(10), None);
assert_eq!(rm.lookup(11), None);
assert_eq!(rm.lookup(45), None);
assert_eq!(rm.insert(0, 10, "A"), Ok(()));
assert_eq!(rm.lookup(-1), None);
assert_eq!(rm.lookup(0), Some(&"A"));
assert_eq!(rm.lookup(6), Some(&"A"));
assert_eq!(rm.lookup(10), Some(&"A"));
assert_eq!(rm.lookup(11), None);
assert_eq!(rm.lookup(45), None);
assert_eq!(rm.insert(11, 30, "B"), Ok(()));
assert_eq!(rm.lookup(-1), None);
assert_eq!(rm.lookup(0), Some(&"A"));
assert_eq!(rm.lookup(6), Some(&"A"));
assert_eq!(rm.lookup(10), Some(&"A"));
assert_eq!(rm.lookup(11), Some(&"B"));
assert_eq!(rm.lookup(15), Some(&"B"));
assert_eq!(rm.lookup(30), Some(&"B"));
assert_eq!(rm.lookup(31), None);
assert_eq!(rm.lookup(45), None);
assert_eq!(rm.insert(500, 1000, "C"), Ok(()));
assert_eq!(rm.lookup(-1), None);
assert_eq!(rm.lookup(0), Some(&"A"));
assert_eq!(rm.lookup(6), Some(&"A"));
assert_eq!(rm.lookup(10), Some(&"A"));
assert_eq!(rm.lookup(11), Some(&"B"));
assert_eq!(rm.lookup(15), Some(&"B"));
assert_eq!(rm.lookup(30), Some(&"B"));
assert_eq!(rm.lookup(31), None);
assert_eq!(rm.lookup(45), None);
assert_eq!(rm.lookup(499), None);
assert_eq!(rm.lookup(500), Some(&"C"));
assert_eq!(rm.lookup(550), Some(&"C"));
assert_eq!(rm.lookup(1000), Some(&"C"));
assert_eq!(rm.lookup(1001), None);
assert_eq!(rm.insert(2000, 2000, "D"), Ok(()));
assert_eq!(rm.lookup(1999), None);
assert_eq!(rm.lookup(2000), Some(&"D"));
assert_eq!(rm.lookup(2001), None);
}
#[test]
pub fn test_insertions_with_overlaps() {
let mut rm: RangeMap<i64, &str> = RangeMap::new();
assert_eq!(rm.insert(0, 10, "A"), Ok(()));
assert_eq!(rm.insert(5, 7, "B"), Err(RangeMapError::RangeOverlap));
assert_eq!(rm.insert(400, 500, "B"), Ok(()));
assert_eq!(rm.insert(40, 450, "C"), Err(RangeMapError::RangeOverlap));
}
#[test]
pub fn test_insertions_with_invalid() {
let mut rm: RangeMap<i64, &str> = RangeMap::new();
assert_eq!(rm.insert(41, 20, "A"), Err(RangeMapError::InvalidRange));
}
#[test]
pub fn test_index_mut() -> anyhow::Result<()> {
let mut rm: RangeMap<i64, &str> = RangeMap::new();
rm.insert(1, 200, "hello")?;
assert_eq!(rm[(1, 200)], "hello");
rm[(1, 200)] = "world";
assert_eq!(rm[(1, 200)], "world");
Ok(())
}
#[test]
pub fn test_overlaps() -> anyhow::Result<()> {
let mut rm: RangeMap<i64, &str> = RangeMap::new();
rm.insert(0, 10, "A")?;
assert!(rm.overlaps(10, 13));
assert!(rm.overlaps(5, 7));
assert!(rm.overlaps(-20, 0));
assert!(rm.overlaps(-20, 20));
assert!(!rm.overlaps(50, 70));
assert!(!rm.overlaps(-50, -10));
rm.insert(400, 500, "B")?;
assert!(rm.overlaps(40, 450));
Ok(())
}
}