bl-range-set 0.1.0

use num_traits::PrimInt;

type Error = Box<dyn std::error::Error>;
type Result<T> = std::result::Result<T, Error>;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct DiscreteRange<T: PrimInt>(T, T);

impl<T: PrimInt> DiscreteRange<T> {
    pub fn new(start: T, end: T) -> Result<Self> {
        if start > end {
            return Err("Invalid range (negative size)".into());
        }

        Ok(DiscreteRange(start, end))
    }

    pub fn len(&self) -> T {
        self.1 - self.0 + T::one()
    }

    pub fn contains(&self, value: T) -> bool {
        value >= self.0 && value <= self.1
    }

    pub fn try_merge(&self, other: &Self) -> Result<Self> {
        if self.1 < other.0 - T::one() || other.1 < self.0 - T::one() {
            return Err("Disjoint ranges cannot be merged".into());
        }

        Ok(DiscreteRange(self.0.min(other.0), self.1.max(other.1)))
    }
}

impl<T: PrimInt> From<(T, T)> for DiscreteRange<T> {
    fn from(tuple: (T, T)) -> Self {
        DiscreteRange::new(tuple.0, tuple.1).unwrap()
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DiscreteRangeSet<T: PrimInt> {
    ranges: Vec<DiscreteRange<T>>,
}

impl<T: PrimInt> DiscreteRangeSet<T> {
    pub fn new() -> Self {
        DiscreteRangeSet { ranges: Vec::new() }
    }

    fn binary_search_by_first(&self, value: T) -> std::result::Result<usize, usize> {
        self.ranges
            .binary_search_by(|r| r.0.partial_cmp(&value).unwrap())
    }

    fn binary_search_contained_range(&self, value: T) -> std::result::Result<usize, usize> {
        self.ranges.binary_search_by(|r| {
            if r.contains(value) {
                std::cmp::Ordering::Equal
            } else if value < r.0 {
                std::cmp::Ordering::Greater
            } else {
                std::cmp::Ordering::Less
            }
        })
    }

    // Add range to the set in sorted order, merging overlapping ranges.
    pub fn add_range<R: Into<DiscreteRange<T>>>(&mut self, range: R) {
        let range = range.into();
        if range.len() == T::zero() {
            return;
        }

        if self.ranges.is_empty() {
            self.ranges.push(range);
            return;
        }

        let start_pos = match self.binary_search_contained_range(range.0) {
            Ok(pos) => pos,
            Err(0) => 0,
            Err(pos) => {
                if range.0 == self.ranges[pos - 1].1 + T::one() {
                    pos - 1
                } else {
                    pos
                }
            }
        };

        let end_pos = match self.binary_search_contained_range(range.1) {
            Ok(pos) => pos + 1,
            Err(0) => 0,
            Err(pos) => {
                if pos != self.ranges.len() && range.1 + T::one() == self.ranges[pos].0 {
                    pos + 1
                } else {
                    pos
                }
            }
        };

        if start_pos == end_pos {
            self.ranges.insert(start_pos, range);
        } else {
            let new_start = self.ranges[start_pos].0.min(range.0);
            let new_end = self.ranges[end_pos - 1].1.max(range.1);
            self.ranges[start_pos].0 = new_start;
            self.ranges[start_pos].1 = new_end;
            self.ranges.drain(start_pos + 1..end_pos);
        }
    }

    pub fn contains(&self, value: T) -> bool {
        match self.binary_search_by_first(value) {
            Ok(_) => true,
            Err(0) => false,
            Err(pos) => self.ranges[pos - 1].contains(value),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_inclusive_range_new() {
        assert!(DiscreteRange::new(1u32, 2).is_ok());
        assert!(DiscreteRange::new(2u32, 1).is_err());
    }

    #[test]
    fn test_inclusive_range_set_add_and_contains() {
        let mut range_set = DiscreteRangeSet::new();
        range_set.add_range(DiscreteRange::new(1u32, 3).unwrap());
        range_set.add_range(DiscreteRange::new(5u32, 7).unwrap());

        assert_eq!(range_set.ranges.len(), 2);

        let in_set = [1, 2, 3, 5, 6, 7];

        for &value in &in_set {
            assert!(range_set.contains(value));
        }

        let not_in_set = [0, 4, 8];
        for &value in &not_in_set {
            assert!(!range_set.contains(value));
        }

        range_set.add_range(DiscreteRange::new(4, 4).unwrap());
        assert_eq!(range_set.ranges.len(), 1);

        assert!(range_set.contains(1));
        assert!(range_set.contains(2));
        assert!(range_set.contains(3));
        assert!(range_set.contains(4));
        assert!(range_set.contains(5));
        assert!(range_set.contains(6));
        assert!(range_set.contains(7));

        assert!(!range_set.contains(0));
        assert!(!range_set.contains(8));
    }

    #[test]
    fn test_inclusive_range_max_value() {
        let max = u32::MAX;
        let mut range_set = DiscreteRangeSet::new();

        assert!(!range_set.contains(max));

        range_set.add_range(DiscreteRange::new(max, max).unwrap());

        assert!(range_set.contains(max));
        assert!(!range_set.contains(max - 1));
        assert!(!range_set.contains(u32::MIN));
    }
}