diff --git a/ostd/src/util.rs b/ostd/src/util.rs index f11ff514..649a8153 100644 --- a/ostd/src/util.rs +++ b/ostd/src/util.rs @@ -1,5 +1,7 @@ // SPDX-License-Identifier: MPL-2.0 +use core::ops::Range; + /// Asserts that a boolean expression is `true` at compile-time. /// /// Rust provides [`const` blocks], which can be used flexibly within methods, but cannot be used @@ -21,3 +23,82 @@ macro_rules! const_assert { /// /// Types that implement `SameSizeAs` must have the same size as `T`. pub unsafe trait SameSizeAs {} + +/// Calculates the [difference] of two [`Range`]s, i.e., `a - b`. +/// +/// This method will return 0, 1, or 2 ranges. All returned ranges are +/// guaranteed to be non-empty and non-overlapping. The returned ranges +/// will be sorted in ascending order. +/// +/// [difference]: https://en.wikipedia.org/wiki/Set_(mathematics)#Set_difference +pub fn range_difference( + a: &Range, + b: &Range, +) -> impl Iterator> { + use core::cmp::{max, min}; + + let r = if b.is_empty() { + [a.clone(), b.clone()] + } else { + [a.start..min(a.end, b.start), max(a.start, b.end)..a.end] + }; + + r.into_iter().filter(|v| !v.is_empty()) +} + +#[cfg(ktest)] +#[expect(clippy::single_range_in_vec_init)] +mod test { + use super::*; + use crate::prelude::ktest; + + #[track_caller] + fn assert_range_difference( + a: Range, + b: Range, + expected: [Range; N], + ) { + let mut res = range_difference(&a, &b); + expected + .into_iter() + .for_each(|val| assert_eq!(res.next(), Some(val))); + assert!(res.next().is_none()); + } + + #[ktest] + fn range_difference_contained() { + assert_range_difference(0..10, 3..7, [0..3, 7..10]); + } + #[ktest] + fn range_difference_all_same() { + assert_range_difference(0..10, 0..10, []); + } + #[ktest] + fn range_difference_left_same() { + assert_range_difference(0..10, 0..5, [5..10]); + } + #[ktest] + fn range_difference_right_same() { + assert_range_difference(0..10, 5..10, [0..5]); + } + #[ktest] + fn range_difference_b_empty() { + assert_range_difference(0..10, 0..0, [0..10]); + } + #[ktest] + fn range_difference_a_empty() { + assert_range_difference(0..0, 0..10, []); + } + #[ktest] + fn range_difference_all_empty() { + assert_range_difference(0..0, 0..0, []); + } + #[ktest] + fn range_difference_left_intersected() { + assert_range_difference(5..10, 0..6, [6..10]); + } + #[ktest] + fn range_difference_right_intersected() { + assert_range_difference(5..10, 6..12, [5..6]); + } +}