From 2c13f3de57df0796d77c75faec1aad0d2e808d88 Mon Sep 17 00:00:00 2001
From: Xiretza <xiretza@xiretza.xyz>
Date: Sun, 4 Dec 2022 21:30:45 +0100
Subject: [PATCH] common/rust: implement union and difference for SectionRange

---
 common/rust/src/section_range.rs | 154 ++++++++++++++++++++++++++++++-
 1 file changed, 153 insertions(+), 1 deletion(-)

diff --git a/common/rust/src/section_range.rs b/common/rust/src/section_range.rs
index 3d8778a..83f3993 100644
--- a/common/rust/src/section_range.rs
+++ b/common/rust/src/section_range.rs
@@ -2,10 +2,12 @@ use std::{
     error::Error,
     fmt,
     num::NonZeroU64,
-    ops::{BitAnd, RangeInclusive},
+    ops::{BitAnd, BitOr, RangeInclusive, Sub},
     str::FromStr,
 };
 
+use crate::UpToTwo;
+
 /// Error returned when an attempt is made to construct an empty `SectionRange`.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub struct EmptyRange;
@@ -131,6 +133,65 @@ impl BitAnd for SectionRange {
     }
 }
 
+impl BitOr for SectionRange {
+    type Output = UpToTwo<Self>;
+
+    fn bitor(self, other: Self) -> Self::Output {
+        let first_start = u64::min(self.start, other.start);
+        let first_end = u64::min(self.end, other.end);
+
+        let second_start = u64::max(self.start, other.start);
+        let second_end = u64::max(self.end, other.end);
+
+        if first_end < second_start {
+            let first = Self {
+                start: first_start,
+                end: first_end,
+            };
+            let second = Self {
+                start: second_start,
+                end: second_end,
+            };
+            UpToTwo::Two(first, second)
+        } else {
+            UpToTwo::One(Self {
+                start: first_start,
+                end: second_end,
+            })
+        }
+    }
+}
+
+impl Sub for SectionRange {
+    type Output = UpToTwo<Self>;
+
+    fn sub(self, other: Self) -> Self::Output {
+        if other.encompasses(&self) {
+            return UpToTwo::Zero;
+        }
+
+        // Closures to prevent integer overflow
+        let first = || Self {
+            start: self.start,
+            end: u64::min(self.end, other.start - 1),
+        };
+        let second = || Self {
+            start: u64::max(self.start, other.end + 1),
+            end: self.end,
+        };
+
+        // The other range does not encompass this one entirely - find out if the remaining bits
+        // are at the start, at the end or both
+        if other.end >= self.end {
+            UpToTwo::One(first())
+        } else if other.start <= self.start {
+            UpToTwo::One(second())
+        } else {
+            UpToTwo::Two(first(), second())
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -168,10 +229,12 @@ mod tests {
         check_intersection(range1, range3, None);
         check_intersection(range2, range3, Some(range3));
 
+        // self-intersection is always encompassing
         assert!(range1.encompasses(&range1));
         assert!(range2.encompasses(&range2));
         assert!(range3.encompasses(&range3));
 
+        // only 2 includes all of 3
         assert!(!range1.encompasses(&range2));
         assert!(!range1.encompasses(&range3));
 
@@ -181,4 +244,93 @@ mod tests {
         assert!(!range3.encompasses(&range1));
         assert!(!range3.encompasses(&range2));
     }
+
+    #[test]
+    fn section_range_operations() {
+        fn elements(r: SectionRange) -> Vec<u64> {
+            RangeInclusive::<_>::from(r).collect()
+        }
+
+        fn elements_multi(r: UpToTwo<SectionRange>) -> Vec<u64> {
+            let is_empty = r.is_empty();
+            let sections: Vec<_> = r.into_iter().flat_map(elements).collect();
+
+            if !is_empty {
+                let mut it = sections.iter().copied();
+                let mut prev_section = it.next().unwrap();
+
+                for section in it {
+                    assert!(section > prev_section);
+                    prev_section = section;
+                }
+            }
+
+            sections
+        }
+
+        let ranges = (0..=5).flat_map(|start| (start..=5).map(move |end| start..=end));
+
+        for range_left in ranges.clone() {
+            for range_right in ranges.clone() {
+                // SETUP
+                println!("testing {:?} against {:?}", range_left, range_right);
+
+                let left = SectionRange::try_from(range_left.clone()).unwrap();
+                let right = SectionRange::try_from(range_right.clone()).unwrap();
+
+                let e_left: Vec<_> = elements(left);
+                let e_right: Vec<_> = elements(right);
+
+                // TESTS
+                assert_eq!(RangeInclusive::<_>::from(left), range_left);
+                assert_eq!(RangeInclusive::<_>::from(right), range_right);
+
+                // len()
+                assert_eq!(
+                    left.len().get(),
+                    *range_left.end() - *range_left.start() + 1
+                );
+
+                // contains()
+                for i in 0..=10 {
+                    assert_eq!(
+                        left.contains(i),
+                        i >= *range_left.start() && i <= *range_left.end()
+                    );
+                }
+
+                // encompasses()
+                assert!(left.encompasses(&left));
+                assert_eq!(
+                    left.encompasses(&right),
+                    e_right.iter().all(|e| e_left.contains(e))
+                );
+
+                // intersection
+                let e_intersection: Vec<_> = e_right
+                    .iter()
+                    .copied()
+                    .filter(|e| e_left.contains(e))
+                    .collect();
+                match left & right {
+                    None => assert_eq!(e_intersection, []),
+                    Some(intersection) => assert_eq!(elements(intersection), e_intersection),
+                }
+
+                // union
+                let mut e_union: Vec<_> = e_left.iter().chain(e_right.iter()).copied().collect();
+                e_union.sort_unstable();
+                e_union.dedup();
+                assert_eq!(elements_multi(left | right), e_union);
+
+                // difference
+                let e_difference: Vec<_> = e_left
+                    .iter()
+                    .copied()
+                    .filter(|e| !e_right.contains(e))
+                    .collect();
+                assert_eq!(elements_multi(left - right), e_difference);
+            }
+        }
+    }
 }