use core::time::Duration;
#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash, Default)]
pub struct Timestamp(pub u64);
impl Timestamp {
pub const ZERO: Timestamp = Timestamp(0);
pub const fn from_nanos(nanos: u64) -> Self {
Timestamp(nanos)
}
pub const fn as_nanos(self) -> u64 {
self.0
}
pub const fn saturating_sub(self, other: Timestamp) -> Duration {
Duration::from_nanos(self.0.saturating_sub(other.0))
}
pub const fn checked_add_nanos(self, nanos: u64) -> Self {
Timestamp(self.0.saturating_add(nanos))
}
pub fn saturating_add(self, duration: Duration) -> Self {
let nanos = u64::try_from(duration.as_nanos()).unwrap_or(u64::MAX);
self.checked_add_nanos(nanos)
}
}
#[cfg(feature = "std")]
impl Timestamp {
pub fn from_instant(base: std::time::Instant, now: std::time::Instant) -> Self {
let elapsed = now.saturating_duration_since(base);
let nanos = u64::try_from(elapsed.as_nanos()).unwrap_or(u64::MAX);
Timestamp(nanos)
}
}
#[non_exhaustive]
#[derive(Copy, Clone, PartialEq, Eq, Debug, Default)]
pub struct Demand {
pub want_bytes: usize,
pub saturated: bool,
}
impl Demand {
pub const fn new(want_bytes: usize) -> Self {
Demand {
want_bytes,
saturated: false,
}
}
pub const fn saturated() -> Self {
Demand {
want_bytes: 0,
saturated: true,
}
}
}
pub trait Stage {
type In<'a>;
type Out;
type Error;
fn feed(&mut self, input: Self::In<'_>, now: Timestamp) -> Result<(), Self::Error>;
fn poll(&mut self) -> Option<Self::Out>;
fn finish(&mut self) -> Result<(), Self::Error>;
fn next_deadline(&self) -> Option<Timestamp>;
fn on_deadline(&mut self, now: Timestamp);
fn demand(&self) -> Demand;
}
#[cfg(test)]
mod tests {
use super::*;
use alloc::collections::VecDeque;
use alloc::vec::Vec;
use core::convert::Infallible;
struct ByteCounter {
total: u64,
threshold: u64,
emitted_up_to: u64,
pending: VecDeque<u64>,
deadline: Option<Timestamp>,
finished: bool,
}
impl ByteCounter {
fn new(threshold: u64, deadline: Option<Timestamp>) -> Self {
ByteCounter {
total: 0,
threshold,
emitted_up_to: 0,
pending: VecDeque::new(),
deadline,
finished: false,
}
}
fn maybe_emit(&mut self) {
if self.total.saturating_sub(self.emitted_up_to) >= self.threshold {
self.pending.push_back(self.total);
self.emitted_up_to = self.total;
}
}
}
impl Stage for ByteCounter {
type In<'a> = &'a [u8];
type Out = u64;
type Error = Infallible;
fn feed(&mut self, input: &[u8], _now: Timestamp) -> Result<(), Self::Error> {
self.total += input.len() as u64;
self.maybe_emit();
Ok(())
}
fn poll(&mut self) -> Option<Self::Out> {
self.pending.pop_front()
}
fn finish(&mut self) -> Result<(), Self::Error> {
if !self.finished && self.total > self.emitted_up_to {
self.pending.push_back(self.total);
self.emitted_up_to = self.total;
}
self.finished = true;
Ok(())
}
fn next_deadline(&self) -> Option<Timestamp> {
self.deadline
}
fn on_deadline(&mut self, now: Timestamp) {
if Some(now) >= self.deadline && self.total > self.emitted_up_to {
self.pending.push_back(self.total);
self.emitted_up_to = self.total;
self.deadline = None;
}
}
fn demand(&self) -> Demand {
if self.total.saturating_sub(self.emitted_up_to) >= self.threshold {
Demand::saturated()
} else {
Demand::new(self.threshold as usize)
}
}
}
#[test]
fn no_std_implementor_drives_via_feed_poll() {
let mut stage = ByteCounter::new(4, None);
let mut outs: Vec<u64> = Vec::new();
stage.feed(&[1, 2], Timestamp::from_nanos(0)).unwrap();
assert_eq!(stage.poll(), None);
stage
.feed(&[3, 4, 5], Timestamp::from_nanos(1_000))
.unwrap();
while let Some(out) = stage.poll() {
outs.push(out);
}
assert_eq!(outs, alloc::vec![5]);
stage.finish().unwrap();
assert_eq!(stage.poll(), None);
}
#[test]
fn no_std_implementor_finish_flushes_remainder() {
let mut stage = ByteCounter::new(100, None);
stage.feed(&[1, 2, 3], Timestamp::from_nanos(0)).unwrap();
assert_eq!(stage.poll(), None);
stage.finish().unwrap();
assert_eq!(stage.poll(), Some(3));
assert_eq!(stage.poll(), None);
}
#[test]
fn no_std_implementor_on_deadline_fires_without_new_input() {
let deadline = Timestamp::from_nanos(5_000);
let mut stage = ByteCounter::new(100, Some(deadline));
stage.feed(&[1, 2, 3], Timestamp::from_nanos(0)).unwrap();
assert_eq!(stage.next_deadline(), Some(deadline));
assert_eq!(stage.poll(), None);
stage.on_deadline(deadline);
assert_eq!(stage.poll(), Some(3));
assert_eq!(stage.next_deadline(), None);
}
#[test]
fn demand_default_and_constructors() {
let d = Demand::default();
assert_eq!(d.want_bytes, 0);
assert!(!d.saturated);
let want = Demand::new(1024);
assert_eq!(want.want_bytes, 1024);
assert!(!want.saturated);
let full = Demand::saturated();
assert_eq!(full.want_bytes, 0);
assert!(full.saturated);
}
#[test]
fn timestamp_arithmetic_saturates_instead_of_panicking() {
let zero = Timestamp::ZERO;
let small = Timestamp::from_nanos(5);
let big = Timestamp::from_nanos(u64::MAX);
assert_eq!(
small.saturating_sub(Timestamp::from_nanos(100)),
Duration::from_nanos(0)
);
assert_eq!(
Timestamp::from_nanos(100).saturating_sub(small),
Duration::from_nanos(95)
);
assert_eq!(big.checked_add_nanos(10), Timestamp::from_nanos(u64::MAX));
assert_eq!(zero.checked_add_nanos(10), Timestamp::from_nanos(10));
assert_eq!(
zero.saturating_add(Duration::from_nanos(10)),
Timestamp::from_nanos(10)
);
assert_eq!(
big.saturating_add(Duration::from_secs(1)),
Timestamp::from_nanos(u64::MAX)
);
let huge = Duration::from_secs(u64::MAX);
assert_eq!(zero.saturating_add(huge), Timestamp::from_nanos(u64::MAX));
}
#[test]
fn timestamp_ordering_and_default() {
assert!(Timestamp::from_nanos(1) < Timestamp::from_nanos(2));
assert_eq!(Timestamp::default(), Timestamp::ZERO);
}
#[cfg(feature = "std")]
#[test]
fn from_instant_std_convenience() {
let base = std::time::Instant::now();
let later = base + Duration::from_millis(5);
let ts = Timestamp::from_instant(base, later);
assert_eq!(ts, Timestamp::from_nanos(5_000_000));
let earlier = base.checked_sub(Duration::from_millis(1)).unwrap_or(base);
let ts2 = Timestamp::from_instant(base, earlier);
assert_eq!(ts2, Timestamp::ZERO);
}
}