use std::{
ops::{Bound, RangeBounds, RangeFull, RangeTo, RangeToInclusive},
task::Poll,
time::Duration,
};
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct Subscription {
pub priority: u8,
pub max_age: Duration,
pub start: Option<Position>,
pub end: Option<Position>,
}
impl Default for Subscription {
fn default() -> Self {
Self {
priority: 0,
max_age: Duration::ZERO,
start: None,
end: None,
}
}
}
impl Subscription {
pub fn with_priority(mut self, priority: u8) -> Self {
self.priority = priority;
self
}
pub fn with_max_age(mut self, max_age: Duration) -> Self {
self.max_age = max_age;
self
}
pub fn with_start(mut self, start: impl Into<Option<Position>>) -> Self {
self.start = start.into();
self
}
pub fn with_end(mut self, end: impl Into<Option<Position>>) -> Self {
self.end = end.into();
self
}
pub fn with_groups(mut self, groups: impl RangeBounds<u64>) -> Self {
let unbounded_start = matches!(groups.start_bound(), Bound::Unbounded);
let (start, end) = sequence_bounds(groups);
self.start = (!unbounded_start).then(|| Position::group(start));
self.end = end.map(Position::group);
self
}
pub(super) fn poll_combined(&self, combined: &Option<Subscription>) -> Poll<Subscription> {
let Some(combined) = combined else {
return Poll::Ready(self.clone());
};
let merged = Subscription {
priority: self.priority.max(combined.priority),
max_age: self.max_age.max(combined.max_age),
start: min_floored(self.start, combined.start),
end: max_unbounded(self.end, combined.end),
};
if &merged != combined {
return Poll::Ready(merged);
}
Poll::Pending
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
pub struct Position {
pub group: u64,
pub frame: u64,
}
impl Position {
pub fn group(group: u64) -> Self {
Self { group, frame: 0 }
}
pub fn after(group: u64, frame: u64) -> Option<Self> {
match frame.checked_add(1) {
Some(frame) => Some(Self { group, frame }),
None => Self::after_group(group),
}
}
pub fn after_group(group: u64) -> Option<Self> {
Some(Self::group(group.checked_add(1)?))
}
pub fn group_end(self) -> Bound<u64> {
if self.frame == 0 {
Bound::Excluded(self.group)
} else {
Bound::Included(self.group)
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct Cap(Option<u64>);
impl Cap {
pub(crate) fn exclusive(self) -> Option<u64> {
self.0
}
}
impl From<Bound<u64>> for Cap {
fn from(bound: Bound<u64>) -> Self {
Self(match bound {
Bound::Included(index) => index.checked_add(1),
Bound::Excluded(index) => Some(index),
Bound::Unbounded => None,
})
}
}
impl From<RangeTo<u64>> for Cap {
fn from(range: RangeTo<u64>) -> Self {
Self(Some(range.end))
}
}
impl From<RangeToInclusive<u64>> for Cap {
fn from(range: RangeToInclusive<u64>) -> Self {
Bound::Included(range.end).into()
}
}
impl From<RangeFull> for Cap {
fn from(_: RangeFull) -> Self {
Self(None)
}
}
pub(super) fn sequence_bounds(range: impl RangeBounds<u64>) -> (u64, Option<u64>) {
let start = match range.start_bound() {
Bound::Included(&start) => start,
Bound::Excluded(&start) => match start.checked_add(1) {
Some(start) => start,
None => return (u64::MAX, Some(u64::MAX)),
},
Bound::Unbounded => 0,
};
(start, Cap::from(range.end_bound().cloned()).exclusive())
}
pub(super) fn before_end(sequence: u64, end: Option<u64>) -> bool {
end.is_none_or(|end| sequence < end)
}
pub(super) fn min_some<T: Ord>(a: Option<T>, b: Option<T>) -> Option<T> {
match (a, b) {
(Some(a), Some(b)) => Some(a.min(b)),
(Some(a), None) | (None, Some(a)) => Some(a),
(None, None) => None,
}
}
pub(super) fn max_some<T: Ord>(a: Option<T>, b: Option<T>) -> Option<T> {
match (a, b) {
(Some(a), Some(b)) => Some(a.max(b)),
(Some(a), None) | (None, Some(a)) => Some(a),
(None, None) => None,
}
}
pub(super) fn min_floored<T: Ord>(a: Option<T>, b: Option<T>) -> Option<T> {
match (a, b) {
(Some(a), Some(b)) => Some(a.min(b)),
(None, _) | (_, None) => None,
}
}
pub(super) fn max_unbounded<T: Ord>(a: Option<T>, b: Option<T>) -> Option<T> {
match (a, b) {
(Some(a), Some(b)) => Some(a.max(b)),
(None, _) | (_, None) => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
fn combine(subscriptions: &[Subscription]) -> Option<Subscription> {
let mut combined = None;
for sub in subscriptions {
if let Poll::Ready(merged) = sub.poll_combined(&combined) {
combined = Some(merged);
}
}
combined
}
#[test]
fn group_ranges_build_whole_group_positions() {
let sub = Subscription::default().with_groups(2..=5);
assert_eq!(sub.start, Some(Position::group(2)));
assert_eq!(sub.end, Some(Position::group(6)));
let sub = Subscription::default().with_groups(2..6);
assert_eq!(sub.end, Some(Position::group(6)));
let sub = Subscription::default().with_groups(..6);
assert_eq!(sub.start, None);
assert_eq!(sub.end, Some(Position::group(6)));
let sub = Subscription::default().with_groups(2..);
assert_eq!(sub.start, Some(Position::group(2)));
assert_eq!(sub.end, None);
let sub = Subscription::default().with_groups(..=u64::MAX);
assert_eq!(sub.end, None);
let sub = Subscription::default().with_groups(2..=5).with_groups(..);
assert_eq!((sub.start, sub.end), (None, None));
}
#[test]
fn positions_are_total_at_the_extremes() {
assert_eq!(Position::after(5, u64::MAX), Some(Position::group(6)));
assert_eq!(Position::after_group(u64::MAX), None);
assert_eq!(Position::after(u64::MAX, u64::MAX), None);
assert_eq!(
Subscription::default().with_end(Position::after_group(u64::MAX)).end,
None
);
assert_eq!(Position::group(0).group_end(), Bound::Excluded(0));
assert_eq!(Position::after_group(5).unwrap().group_end(), Bound::Excluded(6));
assert_eq!(Position::after(5, 2).unwrap().group_end(), Bound::Included(5));
assert_eq!(Cap::from(..0).exclusive(), Some(0));
assert_eq!(Cap::from(..=5).exclusive(), Some(6));
assert_eq!(Cap::from(..=u64::MAX).exclusive(), None);
assert_eq!(Cap::from(..).exclusive(), None);
assert_eq!(Cap::from(Bound::Included(5)), Cap::from(..6));
assert_eq!(Cap::from(Bound::Unbounded), Cap::from(..));
}
#[test]
fn combined_group_start_keeps_the_loosest_floor() {
let catchup = Subscription::default().with_start(Position::group(10));
let older_catchup = Subscription::default().with_start(Position::group(5));
let combined = combine(&[catchup.clone(), older_catchup]).unwrap();
assert_eq!(combined.start, Some(Position::group(5)));
let unfloored = Subscription::default();
let combined = combine(&[catchup, unfloored]).unwrap();
assert_eq!(combined.start, None);
}
#[test]
fn combined_group_end_keeps_live_subscription_unbounded() {
let live = Subscription::default();
let bounded = Subscription::default().with_end(Position::after_group(10));
let combined = combine(&[live, bounded]).unwrap();
assert_eq!(combined.end, None);
}
#[test]
fn combined_start_folds_the_whole_position() {
let early_frame = Subscription::default().with_start(Position { group: 5, frame: 2 });
let late_frame = Subscription::default().with_start(Position { group: 5, frame: 9 });
let combined = combine(&[late_frame.clone(), early_frame.clone()]).unwrap();
assert_eq!(combined.start, Some(Position { group: 5, frame: 2 }));
let earlier_group = Subscription::default().with_start(Position { group: 4, frame: 7 });
let combined = combine(&[early_frame, earlier_group]).unwrap();
assert_eq!(combined.start, Some(Position { group: 4, frame: 7 }));
}
#[test]
fn combined_end_folds_the_whole_position() {
let short = Subscription::default().with_end(Position::after(5, 2));
let long = Subscription::default().with_end(Position::after(5, 9));
let combined = combine(&[short.clone(), long.clone()]).unwrap();
assert_eq!(combined.end, Some(Position { group: 5, frame: 10 }));
let whole = Subscription::default().with_end(Position::after_group(5));
let combined = combine(&[long, whole]).unwrap();
assert_eq!(combined.end, Some(Position::group(6)));
let later_group = Subscription::default().with_end(Position::after(6, 1));
let combined = combine(&[short, later_group]).unwrap();
assert_eq!(combined.end, Some(Position { group: 6, frame: 2 }));
}
#[test]
fn combined_group_end_uses_latest_bounded_end() {
let early = Subscription::default().with_end(Position::after_group(10));
let late = Subscription::default().with_end(Position::after_group(20));
let combined = combine(&[early, late]).unwrap();
assert_eq!(combined.end, Some(Position::group(21)));
}
}