feat(rangedict): Introduce initial version

src/lib.rs

@@ -10,5 +10,7 @@ pub mod config;
 pub mod rundir;
 /// String utilities.
 pub mod str;
+/// Range dictionary data structure
+pub mod rangedict;
 #[cfg(feature = "thread")]
 pub mod thread;

src/rangedict.rs

+use std::collections::BTreeMap;
+
+use thiserror::Error;
+
+/// Error type for [`RangeDict`]
+#[derive(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(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
+    }
+}
+
+#[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));
+    }
+}