implement removal and add more tests

Author: DSharifi

quinn-proto/src/range_set/btree_range_set.rs

@@ -121,18 +121,24 @@ impl RangeSet {
             }
             Some(_) | None => false,
         };
+
+        let removed_end = self
+            .0
+            .extract_if((Excluded(x.start), Excluded(x.end)), |_, _| true)
+            .last()
+            .map(|(_start, end)| end);
+
         let mut after = false;
-        while let Some((start, end)) = self.succ(x.start) {
-            if start >= x.end {
-                break;
-            }
+
+        if let Some(removed_end) = removed_end {
             after = true;
-            self.0.remove(&start);
-            if end > x.end {
-                self.0.insert(x.end, end);
-                break;
+
+            let over_removed = removed_end > x.end;
+            if over_removed {
+                self.0.insert(x.end, removed_end);
             }
         }
+
         before || after
     }
 
@@ -382,7 +388,21 @@ mod tests {
     }
 
     #[test]
-    fn insert_merges_abutting_range() {
+    fn insert_extends_prev_joint_range() {
+        // given
+        let mut set = RangeSet::new();
+        set.insert(0..10);
+
+        // when
+        set.insert(10..20);
+
+        // then
+        let expected_inner = BTreeMap::from_iter([(0, 20)]);
+        assert_eq!(set.0, expected_inner, "ranges should have merged");
+    }
+
+    #[test]
+    fn insert_merges_touching_range() {
         // given
         let mut set = RangeSet::new();
         set.insert(10..20);
@@ -392,7 +412,7 @@ mod tests {
 
         // then
         let expected_inner = BTreeMap::from_iter([(5, 20)]);
-        assert_eq!(set.0, expected_inner, "Ranges should have merged.");
+        assert_eq!(set.0, expected_inner, "ranges should have merged.");
     }
 
     #[test]
@@ -406,7 +426,7 @@ mod tests {
 
         // then
         let expected_inner = BTreeMap::from_iter([(10, 25)]);
-        assert_eq!(set.0, expected_inner, "Ranges should have merged.");
+        assert_eq!(set.0, expected_inner, "ranges should have merged.");
     }
 
     #[test]
@@ -416,11 +436,15 @@ mod tests {
         set.insert(10..20);
 
         // when
-        set.insert(12..15);
+        let insert_result = set.insert(14..16);
 
         // then
         let expected_inner = BTreeMap::from_iter([(10, 20)]);
-        assert_eq!(set.0, expected_inner, "Ranges should have merged.");
+        assert_eq!(set.0, expected_inner, "insertion should be ignored");
+        assert!(
+            !insert_result,
+            "insert should return false when inserting an already contained range"
+        );
     }
 
     #[test]
@@ -435,6 +459,158 @@ mod tests {
 
         // then
         let expected_inner = BTreeMap::from_iter([(0, 30)]);
-        assert_eq!(set.0, expected_inner, "Ranges should have merged.");
+        assert_eq!(set.0, expected_inner, "existing ranges should have merged");
+    }
+
+    #[test]
+    fn remove_splits_contained_range() {
+        // given
+        let mut set = RangeSet::new();
+        set.insert(0..10);
+
+        // when
+        let changed = set.remove(2..8);
+
+        // then
+        assert!(changed);
+        let expected_inner = BTreeMap::from_iter([(0, 2), (8, 10)]);
+        assert_eq!(set.0, expected_inner, "range should have split into two");
+    }
+
+    #[test]
+    fn remove_trims_start_of_range() {
+        // given
+        let mut set = RangeSet::new();
+        set.insert(10..20);
+
+        // when
+        let changed = set.remove(5..15);
+
+        // then
+        assert!(changed);
+        let expected_inner = BTreeMap::from_iter([(15, 20)]);
+        assert_eq!(
+            set.0, expected_inner,
+            "start of range should have been trimmed"
+        );
+    }
+
+    #[test]
+    fn remove_trims_end_of_range() {
+        // given
+        let mut set = RangeSet::new();
+        set.insert(0..10);
+
+        // when
+        let changed = set.remove(5..15);
+
+        // then
+        assert!(changed);
+        let expected_inner = BTreeMap::from_iter([(0, 5)]);
+        assert_eq!(
+            set.0, expected_inner,
+            "end of range should have been trimmed"
+        );
+    }
+
+    #[test]
+    fn remove_swallows_multiple_ranges() {
+        // given
+        let mut set = RangeSet::new();
+        set.insert(0..5);
+        set.insert(10..15);
+        set.insert(20..25);
+
+        // when
+        let changed = set.remove(0..25);
+
+        // then
+        assert!(changed);
+        let expected_inner = BTreeMap::new();
+        assert_eq!(
+            set.0, expected_inner,
+            "all ranges should be gone as the removal swallowed them"
+        );
+    }
+
+    #[test]
+    fn remove_specific_range_in_set() {
+        // given
+        let mut set = RangeSet::new();
+        set.insert(10..20);
+
+        // when
+        let changed = set.remove(10..20);
+
+        // then
+        assert!(changed);
+        let expected_inner = BTreeMap::new();
+        assert_eq!(set.0, expected_inner, "set should be empty");
+    }
+
+    #[test]
+    fn removing_range_touching_existing_ranges_is_ignored() {
+        // given
+        let mut set = RangeSet::new();
+        set.insert(10..20);
+
+        // when
+        let changed_start = set.remove(0..10);
+        let changed_end = set.remove(20..30);
+
+        // then
+        assert!(!changed_start);
+        assert!(!changed_end);
+        let expected_inner = BTreeMap::from_iter([(10, 20)]);
+        assert_eq!(
+            set.0, expected_inner,
+            "touching ranges should not remove anything."
+        );
+    }
+
+    #[test]
+    fn removing_range_crossing_existing_ranges_that_should_be_trimmed() {
+        // given
+        let mut set = RangeSet::new();
+        set.insert(0..5);
+        set.insert(10..15);
+
+        // when
+        let changed = set.remove(4..11);
+
+        // then
+        assert!(changed);
+        let expected_inner = BTreeMap::from_iter([(0, 4), (11, 15)]);
+        assert_eq!(set.0, expected_inner, "existing ranges should be trimmed");
+    }
+
+    #[test]
+    fn remove_crossing_range_trims_start() {
+        // given
+        let mut set = RangeSet::new();
+        set.insert(10..20);
+
+        // when
+        let changed = set.remove(9..12);
+
+        // then
+        assert!(changed);
+        let expected_inner = BTreeMap::from_iter([(12, 20)]);
+        assert_eq!(set.0, expected_inner);
+    }
+
+    #[test]
+    fn remove_no_overlap_before_and_after() {
+        // given
+        let mut set = RangeSet::new();
+        set.insert(10..20);
+
+        // when
+        let changed = set.remove(0..5);
+
+        // then
+        assert!(!changed);
+        let expected_inner = BTreeMap::from_iter([(10, 20)]);
+        assert_eq!(set.0, expected_inner, "disjoint removals should be ignored");
     }
 }