diff --git a/src/lib.rs b/src/lib.rs index 9b5b4f13c37d638081ca4588dc03a0fc068df168..bddf793d3f54ff3aae5678528f4acb220f02e832 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,8 +6,8 @@ pub mod cli; /// Configuration retrieval and parsing. pub mod config; -/// Range dictionary data structure -pub mod rangedict; +/// Range dictionary/map data structure +pub mod rangemap; /// Range set data structure pub mod rangeset; /// Ephemeral working directories. diff --git a/src/rangedict.rs b/src/rangedict.rs deleted file mode 100644 index 5e83df327498a8637b4ac155759f91a32b867212..0000000000000000000000000000000000000000 --- a/src/rangedict.rs +++ /dev/null @@ -1,173 +0,0 @@ -use std::collections::BTreeMap; - -use thiserror::Error; - -/// Error type for [`RangeDict`] -#[derive(Clone, Error, Debug, PartialEq)] -pub enum RangeDictError { - /// Used when an invalid range is provided (e.g. End is before start). - #[error("Invalid range provided")] - InvalidRange, - /// Used when an insertion would result in an overlap. - #[error("Overlapping ranges")] - RangeOverlap, -} - -/// A dictionary keyed by ranges -/// -/// Useful for sharding implementations -/// -/// Powereded by a BTreeMap under the hood. -/// -/// Note: Ranges are inclusive. -#[derive(Clone, Debug)] -pub struct RangeDict<RK: Ord + Copy, V> { - ranges: BTreeMap<RK, (RK, V)>, // Key: start of range, Value: (end, associated value) -} - -impl<RK: Ord + Copy, V> RangeDict<RK, V> { - /// Makes a new, empty [`RangeDict`]. - pub fn new() -> Self { - Self { - ranges: BTreeMap::new(), - } - } - - /// Inserts a value in the map at the given range - /// - /// A [`RangeDictError::InvalidRange`] error will be returned if the range - /// is invalid. - /// A [`RangeDictError::RangeOverlap`] error will be returned if the - /// insertion would result in an overlap. - pub fn insert(&mut self, start: RK, end: RK, value: V) -> Result<(), RangeDictError> { - if end < start { - return Err(RangeDictError::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(RangeDictError::RangeOverlap); - } - } - - self.ranges.insert(start, (end, value)); - - Ok(()) - } - - /// Looks up a value at the specified key - /// - /// If the key is within an existing range, the matching value is - /// returned. Otherwise, [`None`] is returned. - 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 - } - - /// Looks up a value at the specified key and removes it from the - /// dictionary. - /// - /// Lookup works just like [`lookup`]. - 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 - } -} - -impl<RK: Ord + Copy, V> Default for RangeDict<RK, V> { - fn default() -> Self { - RangeDict::new() - } -} - -#[cfg(test)] -mod tests { - use crate::rangedict::RangeDictError; - - use super::RangeDict; - - #[test] - pub fn test_insertions() { - let mut rd: RangeDict<i64, &str> = RangeDict::new(); - - assert_eq!(rd.lookup(-1), None); - assert_eq!(rd.lookup(0), None); - assert_eq!(rd.lookup(6), None); - assert_eq!(rd.lookup(10), None); - assert_eq!(rd.lookup(11), None); - assert_eq!(rd.lookup(45), None); - - assert_eq!(rd.insert(0, 10, "A"), Ok(())); - - assert_eq!(rd.lookup(-1), None); - assert_eq!(rd.lookup(0), Some(&"A")); - assert_eq!(rd.lookup(6), Some(&"A")); - assert_eq!(rd.lookup(10), Some(&"A")); - assert_eq!(rd.lookup(11), None); - assert_eq!(rd.lookup(45), None); - - assert_eq!(rd.insert(11, 30, "B"), Ok(())); - - assert_eq!(rd.lookup(-1), None); - assert_eq!(rd.lookup(0), Some(&"A")); - assert_eq!(rd.lookup(6), Some(&"A")); - assert_eq!(rd.lookup(10), Some(&"A")); - assert_eq!(rd.lookup(11), Some(&"B")); - assert_eq!(rd.lookup(15), Some(&"B")); - assert_eq!(rd.lookup(30), Some(&"B")); - assert_eq!(rd.lookup(31), None); - assert_eq!(rd.lookup(45), None); - - assert_eq!(rd.insert(500, 1000, "C"), Ok(())); - - assert_eq!(rd.lookup(-1), None); - assert_eq!(rd.lookup(0), Some(&"A")); - assert_eq!(rd.lookup(6), Some(&"A")); - assert_eq!(rd.lookup(10), Some(&"A")); - assert_eq!(rd.lookup(11), Some(&"B")); - assert_eq!(rd.lookup(15), Some(&"B")); - assert_eq!(rd.lookup(30), Some(&"B")); - assert_eq!(rd.lookup(31), None); - assert_eq!(rd.lookup(45), None); - assert_eq!(rd.lookup(499), None); - assert_eq!(rd.lookup(500), Some(&"C")); - assert_eq!(rd.lookup(550), Some(&"C")); - assert_eq!(rd.lookup(1000), Some(&"C")); - assert_eq!(rd.lookup(1001), None); - - assert_eq!(rd.insert(2000, 2000, "D"), Ok(())); - - assert_eq!(rd.lookup(1999), None); - assert_eq!(rd.lookup(2000), Some(&"D")); - assert_eq!(rd.lookup(2001), None); - } - - #[test] - pub fn test_insertions_with_overlaps() { - let mut rd: RangeDict<i64, &str> = RangeDict::new(); - - assert_eq!(rd.insert(0, 10, "A"), Ok(())); - assert_eq!(rd.insert(5, 7, "B"), Err(RangeDictError::RangeOverlap)); - assert_eq!(rd.insert(400, 500, "B"), Ok(())); - assert_eq!(rd.insert(40, 450, "C"), Err(RangeDictError::RangeOverlap)); - } - - #[test] - pub fn test_insertions_with_invalid() { - let mut rd: RangeDict<i64, &str> = RangeDict::new(); - - assert_eq!(rd.insert(41, 20, "A"), Err(RangeDictError::InvalidRange)); - } -} diff --git a/src/rangemap.rs b/src/rangemap.rs new file mode 100644 index 0000000000000000000000000000000000000000..91eef9a8774e8ed5e2b244ead2e6959bf63552b9 --- /dev/null +++ b/src/rangemap.rs @@ -0,0 +1,173 @@ +use std::collections::BTreeMap; + +use thiserror::Error; + +/// Error type for [`RangeMap`] +#[derive(Clone, Error, Debug, PartialEq)] +pub enum RangeMapError { + /// Used when an invalid range is provided (e.g. End is before start). + #[error("Invalid range provided")] + InvalidRange, + /// Used when an insertion would result in an overlap. + #[error("Overlapping ranges")] + RangeOverlap, +} + +/// A dictionary keyed by ranges +/// +/// Useful for sharding implementations +/// +/// Powereded by a BTreeMap under the hood. +/// +/// Note: Ranges are inclusive. +#[derive(Clone, Debug)] +pub struct RangeMap<RK: Ord + Copy, V> { + ranges: BTreeMap<RK, (RK, V)>, // Key: start of range, Value: (end, associated value) +} + +impl<RK: Ord + Copy, V> RangeMap<RK, V> { + /// Makes a new, empty [`RangeMap`]. + pub fn new() -> Self { + Self { + ranges: BTreeMap::new(), + } + } + + /// Inserts a value in the map at the given range + /// + /// A [`RangeMapError::InvalidRange`] error will be returned if the range + /// is invalid. + /// A [`RangeMapError::RangeOverlap`] error will be returned if the + /// insertion would result in an overlap. + 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(()) + } + + /// Looks up a value at the specified key + /// + /// If the key is within an existing range, the matching value is + /// returned. Otherwise, [`None`] is returned. + 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 + } + + /// Looks up a value at the specified key and removes it from the + /// dictionary. + /// + /// Lookup works just like [`lookup`]. + 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 + } +} + +impl<RK: Ord + Copy, V> Default for RangeMap<RK, V> { + fn default() -> Self { + RangeMap::new() + } +} + +#[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)); + } +}