Skip to content

perf(proto): use extract_if in RangeSet for better performance #2474

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
DSharifi wants to merge 5 commits into
base: main
Choose a base branch
from
+308 −1
Draft

perf(proto): use extract_if in RangeSet for better performance #2474

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
d828e05
chore: use `extract_if` in `RangeSet::insert`
DSharifi Dec 12, 2025
e3f363f
implement removal and add more tests
DSharifi Dec 12, 2025
7ff0f56
Merge remote-tracking branch 'origin' into dsharifi/use-extract-if
DSharifi Dec 29, 2025
0f4bf15
use rustversion to feature gate extract_if usage
DSharifi Dec 29, 2025
70a87d8
fix cargo check
DSharifi Dec 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ rand = "0.9"
rcgen = "0.14"
ring = "0.17"
rustc-hash = "2"
rustversion = "1"
rustls = { version = "0.23.5", default-features = false, features = ["std"] }
rustls-platform-verifier = "0.6"
rustls-pki-types = "1.7"
Expand Down
1 change: 1 addition & 0 deletions quinn-proto/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ slab = { workspace = true }
thiserror = { workspace = true }
tinyvec = { workspace = true, features = ["alloc"] }
tracing = { workspace = true }
rustversion = { workspace = true }

# Feature flags & dependencies for wasm
# wasm-bindgen is assumed for a wasm*-*-unknown target
Expand Down
306 changes: 305 additions & 1 deletion quinn-proto/src/range_set/btree_range_set.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#[cfg(test)]
use std::cmp::Ordering;

use std::{
cmp,
collections::{BTreeMap, btree_map},
Expand Down Expand Up @@ -58,6 +59,37 @@ impl RangeSet {
true
}

#[rustversion::since(1.91)]
pub(crate) fn insert(&mut self, mut x: Range<u64>) -> bool {
if x.is_empty() {
return false;
}
if let Some((start, end)) = self.pred(x.start) {
if end >= x.end {
// Wholly contained
return false;
} else if end >= x.start {
// Extend overlapping predecessor
self.0.remove(&start);
x.start = start;
}
}

let existing_end = self
.0
.extract_if((Excluded(x.start), Included(x.end)), |_, _| true)
.last()
.map(|(_start, end)| end);

if let Some(existing_end) = existing_end {
x.end = cmp::max(x.end, existing_end);
}

self.0.insert(x.start, x.end);
true
}

#[rustversion::before(1.91)]
pub(crate) fn insert(&mut self, mut x: Range<u64>) -> bool {
if x.is_empty() {
return false;
Expand All @@ -83,7 +115,6 @@ impl RangeSet {
self.0.insert(x.start, x.end);
true
}

/// Find closest range to `x` that begins at or before it
fn pred(&self, x: u64) -> Option<(u64, u64)> {
self.0
Expand All @@ -101,6 +132,52 @@ impl RangeSet {
}

#[cfg(test)]
#[rustversion::since(1.91)]
pub(super) fn remove(&mut self, x: Range<u64>) -> bool {
if x.is_empty() {
return false;
}

let before = match self.pred(x.start) {
Some((start, end)) if end > x.start => {
self.0.remove(&start);
if start < x.start {
self.0.insert(start, x.start);
}
if end > x.end {
self.0.insert(x.end, end);
}
// Short-circuit if we cannot possibly overlap with another range
if end >= x.end {
return true;
}
true
}
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;

if let Some(removed_end) = removed_end {
after = true;

let over_removed = removed_end > x.end;
if over_removed {
self.0.insert(x.end, removed_end);
}
}

before || after
}

#[cfg(test)]
#[rustversion::before(1.91)]
pub(super) fn remove(&mut self, x: Range<u64>) -> bool {
if x.is_empty() {
return false;
Expand Down Expand Up @@ -392,4 +469,231 @@ mod tests {
assert_eq!(set.len(), 1);
assert_eq!(set.peek_min().unwrap(), 0..4);
}

#[test]
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);

// when
set.insert(5..10);

// then
let expected_inner = BTreeMap::from_iter([(5, 20)]);
assert_eq!(set.0, expected_inner, "ranges should have merged.");
}

#[test]
fn insert_merges_existing_crossing_range() {
// given
let mut set = RangeSet::new();
set.insert(10..20);

// when
set.insert(12..25);

// then
let expected_inner = BTreeMap::from_iter([(10, 25)]);
assert_eq!(set.0, expected_inner, "ranges should have merged.");
}

#[test]
fn insert_ignored_if_wholly_contained() {
// given
let mut set = RangeSet::new();
set.insert(10..20);

// when
let insert_result = set.insert(14..16);

// then
let expected_inner = BTreeMap::from_iter([(10, 20)]);
assert_eq!(set.0, expected_inner, "insertion should be ignored");
assert!(
!insert_result,
"insert should return false when inserting an already contained range"
);
}

#[test]
fn insert_bridges_disjoint_ranges() {
// given
let mut set = RangeSet::new();
set.insert(0..10);
set.insert(20..30);

// when
set.insert(10..20);

// then
let expected_inner = BTreeMap::from_iter([(0, 30)]);
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");
}
}