// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
use super::*;
use crate::{
    event,
    packet::number::PacketNumberSpace,
    path,
    recovery::congestion_controller::PathPublisher,
    time::{Clock, NoopClock},
};
use core::time::Duration;

#[macro_export]
macro_rules! assert_delta {
    ($x:expr, $y:expr, $d:expr) => {
        assert!(
            ($x - $y).abs() < $d,
            "assertion failed: `({:?} - {:?}).abs() < {:?})`",
            $x,
            $y,
            $d
        );
    };
}

fn bytes_to_packets(bytes: f32, max_datagram_size: u16) -> f32 {
    bytes / max_datagram_size as f32
}

#[test]
//= https://www.rfc-editor.org/rfc/rfc8312#section-4.1
//= type=test
fn w_cubic() {
    let max_datagram_size = 1200;
    let mut cubic = Cubic::new(max_datagram_size);

    // 2_764_800 is used because it can be divided by 1200 and then have a cubic
    // root result in an integer value.
    cubic.multiplicative_decrease(2_764_800.0);
    assert_delta!(
        cubic.w_max,
        bytes_to_packets(2_764_800.0, max_datagram_size),
        0.001
    );

    let mut t = Duration::from_secs(0);

    // W_cubic(0)=W_max*beta_cubic
    assert_delta!(cubic.w_max * BETA_CUBIC, cubic.w_cubic(t), 0.001);

    // K = cubic_root(W_max*(1-beta_cubic)/C)
    // K = cubic_root(2304 * 0.75) = 12
    assert_eq!(cubic.k, Duration::from_secs(12));

    //= https://www.rfc-editor.org/rfc/rfc8312#section-5.1
    //= type=test
    //# Therefore, C SHOULD be set to 0.4.

    // W_cubic(t) = C*(t-K)^3 + W_max
    // W_cubic(t) = .4*(t-12)^3 + 2304
    // W_cubic(15) = .4*27 + 2304 = 2314.8
    t = Duration::from_secs(15);
    assert_delta!(cubic.w_cubic(t), 2314.8, 0.001);

    // W_cubic(10) = .4*-8 + 2304 = 2300.8
    t = Duration::from_secs(10);
    assert_delta!(cubic.w_cubic(t), 2300.8, 0.001);
}

#[test]
//= https://www.rfc-editor.org/rfc/rfc8312#section-4.6
//= type=test
fn w_est() {
    let max_datagram_size = 1200;
    let mut cubic = Cubic::new(max_datagram_size);
    cubic.w_max = 100.0;
    let t = Duration::from_secs(6);
    let rtt = Duration::from_millis(300);

    // W_est(t) = W_max*beta_cubic + [3*(1-beta_cubic)/(1+beta_cubic)] * (t/RTT)
    // W_est(6) = 100*.7 + [3*(1-.7)/(1+.7)] * (6/.3)
    // W_est(6) = 70 + 0.5294117647 * 20 = 80.588235294

    assert_delta!(cubic.w_est(t, rtt), 80.5882, 0.001);
}

#[allow(clippy::float_cmp)]
#[test]
//= https://www.rfc-editor.org/rfc/rfc8312#section-4.5
//= type=test
fn multiplicative_decrease() {
    let max_datagram_size = 1200.0;
    let mut cubic = Cubic::new(max_datagram_size as u16);
    cubic.w_max = bytes_to_packets(10000.0, max_datagram_size as u16);

    assert_eq!(
        cubic.multiplicative_decrease(100_000.0),
        (100_000.0 * BETA_CUBIC)
    );
    // Window max was not less than the last max, so not fast convergence
    assert_delta!(cubic.w_last_max, cubic.w_max, 0.001);
    assert_delta!(cubic.w_max, 100_000.0 / max_datagram_size, 0.001);

    assert_eq!(
        cubic.multiplicative_decrease(80000.0),
        (80000.0 * BETA_CUBIC)
    );
    //= https://www.rfc-editor.org/rfc/rfc8312#section-4.6
    //= type=test
    //# To speed up this bandwidth release by
    //# existing flows, the following mechanism called "fast convergence"
    //# SHOULD be implemented.
    // Window max was less than the last max, so fast convergence applies
    assert_delta!(cubic.w_last_max, 80000.0 / max_datagram_size, 0.001);
    // W_max = W_max*(1.0+beta_cubic)/2.0 = W_max * .85
    assert_delta!(cubic.w_max, 80000.0 * 0.85 / max_datagram_size, 0.001);

    //= https://www.rfc-editor.org/rfc/rfc8312#section-4.5
    //= type=test
    //# Parameter beta_cubic SHOULD be set to 0.7.
    assert_eq!(0.7, BETA_CUBIC);
}

#[test]
//= https://www.rfc-editor.org/rfc/rfc9002#section-7.8
//= type=test
fn is_congestion_limited() {
    let max_datagram_size = 1000;
    let mut cc = CubicCongestionController::new(max_datagram_size);
    cc.congestion_window = 1000.0;
    cc.bytes_in_flight = BytesInFlight::new(100);

    assert!(cc.is_congestion_limited());

    cc.congestion_window = 1100.0;

    assert!(!cc.is_congestion_limited());

    cc.bytes_in_flight = BytesInFlight::new(2000);

    assert!(cc.is_congestion_limited());
}

#[test]
fn is_congestion_window_under_utilized() {
    let max_datagram_size = 1200;
    let mut cc = CubicCongestionController::new(max_datagram_size);
    cc.congestion_window = 12000.0;

    // In Slow Start, the window is under utilized if it is less than half full
    cc.bytes_in_flight = BytesInFlight::new(5999);
    cc.state = SlowStart;
    assert!(cc.is_congestion_window_under_utilized());

    cc.bytes_in_flight = BytesInFlight::new(6000);
    assert!(!cc.is_congestion_window_under_utilized());

    cc.state = State::congestion_avoidance(NoopClock.get_time());
    assert!(cc.is_congestion_window_under_utilized());

    // In Congestion Avoidance, the window is under utilized if there are more than
    // 3 * MTU bytes available in the congestion window (12000 - 3 * 1200 = 8400)
    cc.bytes_in_flight = BytesInFlight::new(8399);
    assert!(cc.is_congestion_window_under_utilized());

    cc.bytes_in_flight = BytesInFlight::new(8400);
    assert!(!cc.is_congestion_window_under_utilized());
}

//= https://www.rfc-editor.org/rfc/rfc9002#section-7.2
//= type=test
//# Endpoints SHOULD use an initial congestion
//# window of ten times the maximum datagram size (max_datagram_size),
//# while limiting the window to the larger of 14,720 bytes or twice the
//# maximum datagram size.
#[test]
fn initial_window() {
    let mut max_datagram_size = 1200;
    assert_eq!(
        (max_datagram_size * 10) as u32,
        CubicCongestionController::initial_window(max_datagram_size)
    );

    max_datagram_size = 2000;
    assert_eq!(
        14720,
        CubicCongestionController::initial_window(max_datagram_size)
    );

    max_datagram_size = 8000;
    assert_eq!(
        (max_datagram_size * 2) as u32,
        CubicCongestionController::initial_window(max_datagram_size)
    );
}

//= https://www.rfc-editor.org/rfc/rfc9002#section-7.2
//= type=test
//# The RECOMMENDED
//# value is 2 * max_datagram_size.
#[test]
fn minimum_window_equals_two_times_max_datagram_size() {
    let max_datagram_size = 1200;
    let cc = CubicCongestionController::new(max_datagram_size);

    assert_delta!(
        (2 * max_datagram_size) as f32,
        cc.cubic.minimum_window(),
        0.001
    );
}

#[test]
fn on_packet_sent() {
    let mut cc = CubicCongestionController::new(1000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let mut rtt_estimator = RttEstimator::default();
    let now = NoopClock.get_time();

    cc.congestion_window = 100_000.0;

    // Last sent packet time updated to t10
    cc.on_packet_sent(
        now + Duration::from_secs(10),
        1,
        None,
        &rtt_estimator,
        &mut publisher,
    );

    assert_eq!(cc.bytes_in_flight, 1);

    // Latest RTT is 100ms
    rtt_estimator.update_rtt(
        Duration::from_millis(0),
        Duration::from_millis(100),
        now,
        true,
        PacketNumberSpace::ApplicationData,
    );

    //= https://www.rfc-editor.org/rfc/rfc8312#section-4.8
    //= type=test
    //# CUBIC MUST employ a slow-start algorithm, when the cwnd is no more
    //# than ssthresh.  Among the slow-start algorithms, CUBIC MAY choose the
    //# standard TCP slow start [RFC5681] in general networks, or the limited
    //# slow start [RFC3742] or hybrid slow start [HR08] for fast and long-
    //# distance networks.

    // Round one of hybrid slow start
    cc.on_rtt_update(now, now, &rtt_estimator, &mut publisher);

    assert!(cc.state.is_slow_start());

    // Latest RTT is 200ms
    rtt_estimator.update_rtt(
        Duration::from_millis(0),
        Duration::from_millis(200),
        now,
        true,
        PacketNumberSpace::ApplicationData,
    );

    // Last sent packet time updated to t20
    cc.on_packet_sent(
        now + Duration::from_secs(20),
        1,
        None,
        &rtt_estimator,
        &mut publisher,
    );

    assert_eq!(cc.bytes_in_flight, 2);
    assert!(cc.state.is_slow_start());

    // Round two of hybrid slow start
    for _i in 1..=8 {
        cc.on_rtt_update(
            now + Duration::from_secs(10),
            now + Duration::from_secs(10),
            &rtt_estimator,
            &mut publisher,
        );
    }

    assert!(!cc.state.is_slow_start());
    assert_delta!(cc.slow_start.threshold, 100_000.0, 0.001);
}

#[test]
fn on_packet_sent_application_limited() {
    let mut cc = CubicCongestionController::new(1000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let rtt_estimator = RttEstimator::default();
    let now = NoopClock.get_time();

    cc.congestion_window = 100_000.0;
    cc.bytes_in_flight = BytesInFlight::new(92_500);
    cc.state = SlowStart;

    // t0: Send a packet in Slow Start
    cc.on_packet_sent(now, 1000, Some(true), &rtt_estimator, &mut publisher);

    assert_eq!(cc.bytes_in_flight, 93_500);
    assert_eq!(cc.time_of_last_sent_packet, Some(now));

    // t10: Enter Congestion Avoidance
    cc.state = State::congestion_avoidance(now + Duration::from_secs(10));

    assert!(!cc.under_utilized);

    // t15: Send a packet in Congestion Avoidance
    cc.on_packet_sent(
        now + Duration::from_secs(15),
        1000,
        Some(true),
        &rtt_estimator,
        &mut publisher,
    );

    assert_eq!(cc.bytes_in_flight, 94_500);
    assert_eq!(
        cc.time_of_last_sent_packet,
        Some(now + Duration::from_secs(15))
    );
    assert!(cc.under_utilized);

    // t20: Send packets to fully utilize the congestion window
    while cc.bytes_in_flight < cc.congestion_window() {
        cc.on_packet_sent(
            now + Duration::from_secs(20),
            1000,
            Some(true),
            &rtt_estimator,
            &mut publisher,
        );
    }

    assert!(!cc.under_utilized);
}

// Confirm `under_utilized` is set properly even when `app_limited` is `None`
#[test]
fn on_packet_sent_none_application_limited() {
    let mut cc = CubicCongestionController::new(1000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let rtt_estimator = RttEstimator::default();
    let now = NoopClock.get_time();

    cc.congestion_window = 100_000.0;
    cc.bytes_in_flight = BytesInFlight::new(92_500);
    cc.state = SlowStart;

    // t0: Send a packet in Slow Start
    cc.on_packet_sent(now, 1000, None, &rtt_estimator, &mut publisher);

    assert_eq!(cc.bytes_in_flight, 93_500);
    assert_eq!(cc.time_of_last_sent_packet, Some(now));

    // t10: Enter Congestion Avoidance
    cc.state = State::congestion_avoidance(now + Duration::from_secs(10));

    assert!(!cc.under_utilized);

    // t15: Send a packet in Congestion Avoidance
    cc.on_packet_sent(
        now + Duration::from_secs(15),
        1000,
        None,
        &rtt_estimator,
        &mut publisher,
    );

    assert_eq!(cc.bytes_in_flight, 94_500);
    assert_eq!(
        cc.time_of_last_sent_packet,
        Some(now + Duration::from_secs(15))
    );
    assert!(cc.under_utilized);

    // t20: Send packets to fully utilize the congestion window
    while cc.bytes_in_flight < cc.congestion_window() {
        cc.on_packet_sent(
            now + Duration::from_secs(20),
            1000,
            None,
            &rtt_estimator,
            &mut publisher,
        );
    }

    assert!(!cc.under_utilized);
}

#[test]
fn on_packet_sent_fast_retransmission() {
    let mut cc = CubicCongestionController::new(1000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let rtt_estimator = RttEstimator::default();
    let now = NoopClock.get_time();

    cc.congestion_window = 100_000.0;
    cc.bytes_in_flight = BytesInFlight::new(99900);
    cc.state = Recovery(now, RequiresTransmission);

    cc.on_packet_sent(
        now + Duration::from_secs(10),
        100,
        None,
        &rtt_estimator,
        &mut publisher,
    );

    assert_eq!(cc.state, Recovery(now, Idle));
}

//= https://www.rfc-editor.org/rfc/rfc8312#section-5.8
//= type=test
//# In case of long periods when cwnd has not been updated due
//# to the application rate limit, such as idle periods, t in Eq. 1 MUST
//# NOT include these periods; otherwise, W_cubic(t) might be very high
//# after restarting from these periods.
#[test]
fn congestion_avoidance_after_idle_period() {
    let mut cc = CubicCongestionController::new(1000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let now = NoopClock.get_time();
    let rtt_estimator = &RttEstimator::default();
    let random = &mut random::testing::Generator::default();

    cc.congestion_window = 6000.0;
    cc.bytes_in_flight = BytesInFlight::new(0);
    cc.state = SlowStart;

    // t0: Send a packet in Slow Start
    cc.on_packet_sent(now, 1000, Some(true), rtt_estimator, &mut publisher);

    assert_eq!(cc.bytes_in_flight, 1000);

    // t10: Enter Congestion Avoidance
    cc.cubic.w_max = 6.0;
    cc.state = State::congestion_avoidance(now + Duration::from_secs(10));

    // t15: Send a packet in Congestion Avoidance while under utilized
    cc.on_packet_sent(
        now + Duration::from_secs(15),
        1000,
        Some(true),
        rtt_estimator,
        &mut publisher,
    );
    assert!(cc.is_congestion_window_under_utilized());

    assert_eq!(cc.bytes_in_flight, 2000);

    // t16: Ack a packet in Congestion Avoidance
    cc.on_ack(
        now,
        1000,
        (),
        rtt_estimator,
        random,
        now + Duration::from_secs(16),
        &mut publisher,
    );
    // Verify the app limited time is set
    assert_eq!(
        cc.state,
        CongestionAvoidance(CongestionAvoidanceTiming {
            start_time: now + Duration::from_secs(10),
            window_increase_time: now + Duration::from_secs(10),
            app_limited_time: Some(now + Duration::from_secs(16)),
        })
    );

    assert_eq!(cc.bytes_in_flight, 1000);

    // t20: Send packets to fully utilize the congestion window
    while cc.bytes_in_flight < cc.congestion_window() {
        cc.on_packet_sent(
            now + Duration::from_secs(20),
            1000,
            Some(false),
            rtt_estimator,
            &mut publisher,
        );
    }

    assert!(!cc.is_congestion_window_under_utilized());

    // t25: Ack a packet in Congestion Avoidance
    cc.on_ack(
        now,
        1000,
        (),
        rtt_estimator,
        random,
        now + Duration::from_secs(25),
        &mut publisher,
    );

    // Verify congestion avoidance start time was moved from t10 to t16 to account
    // for the 6 seconds of under utilized time and the app_limited_time was reset
    assert_eq!(
        cc.state,
        CongestionAvoidance(CongestionAvoidanceTiming {
            start_time: now + Duration::from_secs(16),
            window_increase_time: now + Duration::from_secs(25),
            app_limited_time: None,
        })
    );
    // Verify t does not include the app limited time
    if let CongestionAvoidance(timing) = cc.state {
        assert_eq!(
            Duration::from_secs(9),
            timing.t(now + Duration::from_secs(25))
        );
    }
}

#[test]
fn congestion_avoidance_after_fast_convergence() {
    let max_datagram_size = 1200;
    let mut cc = CubicCongestionController::new(max_datagram_size);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let now = NoopClock.get_time();
    let random = &mut random::testing::Generator::default();
    cc.bytes_in_flight = BytesInFlight::new(100);
    cc.congestion_window = 80_000.0;
    cc.cubic.w_last_max = bytes_to_packets(100_000.0, max_datagram_size);

    cc.on_packet_lost(100, (), false, false, random, now, &mut publisher);
    assert_delta!(cc.congestion_window, 80_000.0 * BETA_CUBIC, 0.001);

    // Window max was less than the last max, so fast convergence applies
    assert_delta!(
        cc.cubic.w_last_max,
        80000.0 / max_datagram_size as f32,
        0.001
    );
    // W_max = W_max*(1.0+beta_cubic)/2.0 = W_max * .85
    assert_delta!(
        cc.cubic.w_max,
        80000.0 * 0.85 / max_datagram_size as f32,
        0.001
    );

    let prev_cwnd = cc.congestion_window;

    // Enter congestion avoidance
    cc.congestion_avoidance(
        Duration::from_millis(10),
        Duration::from_millis(100),
        100,
        f32::MAX,
    );

    // Verify congestion window has increased
    assert!(cc.congestion_window > prev_cwnd);
}

#[test]
fn congestion_avoidance_after_rtt_improvement() {
    let max_datagram_size = 1200;
    let mut cc = CubicCongestionController::new(max_datagram_size);
    cc.bytes_in_flight = BytesInFlight::new(100);
    cc.congestion_window = 80_000.0;
    cc.cubic.w_max = cc.congestion_window / 1200.0;

    // Enter congestion avoidance with a long rtt
    cc.congestion_avoidance(
        Duration::from_millis(10),
        Duration::from_millis(750),
        100,
        f32::MAX,
    );

    // At this point the target is less than the congestion window
    let prev_cwnd = cc.congestion_window;
    assert!(
        cc.cubic.w_cubic(Duration::from_secs(0)) < bytes_to_packets(prev_cwnd, max_datagram_size)
    );

    // Receive another ack, now with a short rtt
    cc.congestion_avoidance(
        Duration::from_millis(20),
        Duration::from_millis(10),
        100,
        f32::MAX,
    );

    // Verify congestion window did not change
    assert_delta!(cc.congestion_window, prev_cwnd, 0.001);
}

#[test]
fn congestion_avoidance_with_small_min_rtt() {
    let max_datagram_size = 1200;
    let mut cc = CubicCongestionController::new(max_datagram_size);
    cc.bytes_in_flight = BytesInFlight::new(100);
    cc.congestion_window = 80_000.0;
    cc.cubic.w_max = cc.congestion_window / 1200.0;

    cc.congestion_avoidance(
        Duration::from_millis(100),
        Duration::from_millis(1),
        100,
        f32::MAX,
    );

    // Verify the window grew by half the sent bytes
    assert_delta!(cc.congestion_window, 80_050.0, 0.001);
}

#[test]
fn congestion_avoidance_max_cwnd() {
    let max_datagram_size = 1200;
    let mut cc = CubicCongestionController::new(max_datagram_size);
    cc.bytes_in_flight = BytesInFlight::new(100);
    cc.congestion_window = 80_000.0;
    cc.cubic.w_max = bytes_to_packets(100_000.0, max_datagram_size);

    cc.congestion_avoidance(
        Duration::from_millis(300),
        Duration::from_millis(100),
        1200,
        80_100.0,
    );

    // Verify the window did not exceed the max cwnd
    assert_delta!(cc.congestion_window, 80_100.0, 0.001);
}

#[test]
fn on_packet_lost() {
    let mut cc = CubicCongestionController::new(1000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let now = NoopClock.get_time();
    let random = &mut random::testing::Generator::default();
    cc.congestion_window = 100_000.0;
    cc.bytes_in_flight = BytesInFlight::new(100_000);
    cc.state = SlowStart;

    cc.on_packet_lost(
        100,
        (),
        false,
        false,
        random,
        now + Duration::from_secs(10),
        &mut publisher,
    );

    assert_eq!(cc.bytes_in_flight, 100_000u32 - 100);
    //= https://www.rfc-editor.org/rfc/rfc9002#section-7.3.1
    //= type=test
    //# The sender MUST exit slow start and enter a recovery period when a
    //# packet is lost or when the ECN-CE count reported by its peer
    //# increases.
    assert_eq!(
        cc.state,
        Recovery(now + Duration::from_secs(10), RequiresTransmission)
    );
    //= https://www.rfc-editor.org/rfc/rfc9002#section-7.3.2
    //= type=test
    //# Implementations MAY reduce the congestion window immediately upon
    //# entering a recovery period or use other mechanisms, such as
    //# Proportional Rate Reduction [PRR], to reduce the congestion window
    //# more gradually.
    assert_delta!(cc.congestion_window, 100_000.0 * BETA_CUBIC, 0.001);
    assert_delta!(cc.slow_start.threshold, 100_000.0 * BETA_CUBIC, 0.001);
}

#[test]
fn on_packet_lost_below_minimum_window() {
    let mut cc = CubicCongestionController::new(1000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let now = NoopClock.get_time();
    let random = &mut random::testing::Generator::default();
    cc.congestion_window = cc.cubic.minimum_window();
    cc.bytes_in_flight = BytesInFlight::new(cc.congestion_window());
    cc.state = State::congestion_avoidance(now);

    cc.on_packet_lost(
        100,
        (),
        false,
        false,
        random,
        now + Duration::from_secs(10),
        &mut publisher,
    );

    assert_delta!(cc.congestion_window, cc.cubic.minimum_window(), 0.001);
}

#[test]
fn on_packet_lost_already_in_recovery() {
    let mut cc = CubicCongestionController::new(1000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let now = NoopClock.get_time();
    let random = &mut random::testing::Generator::default();
    cc.congestion_window = 10000.0;
    cc.bytes_in_flight = BytesInFlight::new(1000);
    cc.state = Recovery(now, Idle);

    // break up on_packet_loss into two call to confirm double call
    // behavior is valid (50 + 50 = 100 lost bytes)
    cc.on_packet_lost(50, (), false, false, random, now, &mut publisher);
    cc.on_packet_lost(50, (), false, false, random, now, &mut publisher);

    // No change to the congestion window
    assert_delta!(cc.congestion_window, 10000.0, 0.001);
}

//= https://www.rfc-editor.org/rfc/rfc9002#section-7.6.2
//= type=test
//# When persistent congestion is declared, the sender's congestion
//# window MUST be reduced to the minimum congestion window
//# (kMinimumWindow), similar to a TCP sender's response on an RTO
//# [RFC5681].
#[test]
fn on_packet_lost_persistent_congestion() {
    let mut cc = CubicCongestionController::new(1000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let now = NoopClock.get_time();
    let random = &mut random::testing::Generator::default();
    cc.congestion_window = 10000.0;
    cc.bytes_in_flight = BytesInFlight::new(1000);
    cc.state = Recovery(now, Idle);

    cc.on_packet_lost(100, (), true, false, random, now, &mut publisher);

    assert!(cc.state.is_slow_start());
    assert_eq!(cc.state, SlowStart);
    assert_delta!(cc.congestion_window, cc.cubic.minimum_window(), 0.001);
    assert_delta!(cc.cubic.w_max, 0.0, 0.001);
    assert_delta!(cc.cubic.w_last_max, 0.0, 0.001);
    assert_eq!(cc.cubic.k, Duration::from_millis(0));
}

//= https://www.rfc-editor.org/rfc/rfc9002#section-7.2
//= type=test
//# If the maximum datagram size changes during the connection, the
//# initial congestion window SHOULD be recalculated with the new size.

//= https://www.rfc-editor.org/rfc/rfc8899#section-3
//= type=test
//# A PL that maintains the congestion window in terms of a limit to
//# the number of outstanding fixed-size packets SHOULD adapt this
//# limit to compensate for the size of the actual packets.
#[test]
fn on_mtu_update_increase() {
    let mut mtu = 5000;
    let cwnd_in_packets = 100_000f32;
    let cwnd_in_bytes = cwnd_in_packets / mtu as f32;
    let mut cc = CubicCongestionController::new(mtu);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    cc.congestion_window = cwnd_in_packets;

    mtu = 10000;
    cc.on_mtu_update(mtu, &mut publisher);
    assert_eq!(cc.max_datagram_size, mtu);
    assert_eq!(cc.cubic.max_datagram_size, mtu);

    assert_delta!(cc.congestion_window, 200_000.0, 0.001);

    //= https://www.rfc-editor.org/rfc/rfc8899#section-3
    //= type=test
    //# An update to the PLPMTU (or MPS) MUST NOT increase the congestion
    //# window measured in bytes [RFC4821].
    assert_delta!(cc.congestion_window / mtu as f32, cwnd_in_bytes, 0.001);
}

//= https://www.rfc-editor.org/rfc/rfc9002#section-6.4
//= type=test
//# The sender MUST discard all recovery state associated with
//# those packets and MUST remove them from the count of bytes in flight.
#[test]
fn on_packet_discarded() {
    let mut cc = CubicCongestionController::new(5000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    cc.bytes_in_flight = BytesInFlight::new(10000);

    cc.on_packet_discarded(1000, &mut publisher);

    assert_eq!(cc.bytes_in_flight, 10000 - 1000);

    let now = NoopClock.get_time();
    cc.state = Recovery(now, FastRetransmission::RequiresTransmission);

    cc.on_packet_discarded(1000, &mut publisher);

    assert_eq!(Recovery(now, FastRetransmission::Idle), cc.state);
}

//= https://www.rfc-editor.org/rfc/rfc9002#section-7.8
//= type=test
//# When bytes in flight is smaller than the congestion window and
//# sending is not pacing limited, the congestion window is
//# underutilized.  This can happen due to insufficient application data
//# or flow control limits.  When this occurs, the congestion window
//# SHOULD NOT be increased in either slow start or congestion avoidance.
#[test]
fn on_packet_ack_limited() {
    let mut cc = CubicCongestionController::new(5000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let now = NoopClock.get_time();
    let random = &mut random::testing::Generator::default();
    cc.congestion_window = 100_000.0;
    cc.bytes_in_flight = BytesInFlight::new(10000);
    cc.under_utilized = true;
    cc.state = SlowStart;

    cc.on_ack(
        now,
        1,
        (),
        &RttEstimator::default(),
        random,
        now,
        &mut publisher,
    );

    assert_delta!(cc.congestion_window, 100_000.0, 0.001);

    cc.state = State::congestion_avoidance(now);

    cc.on_ack(
        now,
        1,
        (),
        &RttEstimator::default(),
        random,
        now,
        &mut publisher,
    );

    assert_delta!(cc.congestion_window, 100_000.0, 0.001);
}

#[test]
#[should_panic]
fn on_packet_ack_timestamp_regression() {
    let mut cc = CubicCongestionController::new(5000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let now = NoopClock.get_time() + Duration::from_secs(1);
    let rtt_estimator = RttEstimator::default();
    let random = &mut random::testing::Generator::default();
    cc.congestion_window = 100_000.0;
    cc.bytes_in_flight = BytesInFlight::new(10000);
    cc.under_utilized = true;
    cc.state = State::congestion_avoidance(now);

    cc.on_ack(now, 1, (), &rtt_estimator, random, now, &mut publisher);

    assert_eq!(
        State::CongestionAvoidance(CongestionAvoidanceTiming {
            start_time: now,
            window_increase_time: now,
            app_limited_time: Some(now),
        }),
        cc.state
    );

    cc.on_ack(
        now,
        1,
        (),
        &rtt_estimator,
        random,
        now - Duration::from_secs(1),
        &mut publisher,
    );
}

#[test]
fn on_packet_ack_utilized_then_under_utilized() {
    let mut cc = CubicCongestionController::new(5000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let now = NoopClock.get_time();
    let mut rtt_estimator = RttEstimator::default();
    let random = &mut random::testing::Generator::default();
    rtt_estimator.update_rtt(
        Duration::from_secs(0),
        Duration::from_millis(200),
        now,
        true,
        PacketNumberSpace::ApplicationData,
    );
    cc.congestion_window = 100_000.0;
    cc.state = SlowStart;

    cc.on_packet_sent(now, 60_000, Some(true), &rtt_estimator, &mut publisher);
    cc.on_ack(now, 10_000, (), &rtt_estimator, random, now, &mut publisher);
    let cwnd = cc.congestion_window();

    assert!(!cc.under_utilized);
    assert!(cwnd > 100_000);

    // Now the window is under utilized, but we still grow the window until more packets are sent
    assert!(cc.is_congestion_window_under_utilized());
    cc.on_ack(
        now,
        1200,
        (),
        &rtt_estimator,
        random,
        now + Duration::from_millis(100),
        &mut publisher,
    );
    assert!(cc.congestion_window() > cwnd);

    // A large amount is acked, but the window should only grow to the max_cwnd
    cc.on_ack(
        now,
        40_000,
        (),
        &rtt_estimator,
        random,
        now + Duration::from_millis(100),
        &mut publisher,
    );
    // 60_000 is the highest bytes in flight * 2 for the slow start max_cwnd multiplier
    assert_eq!(60_000 * 2, cc.congestion_window());

    let cwnd = cc.congestion_window();

    // Now the application has had a chance to send more data, but it didn't send enough to
    // utilize the congestion window, so the window does not grow.
    cc.on_packet_sent(now, 1200, Some(true), &rtt_estimator, &mut publisher);
    assert!(cc.under_utilized);
    cc.on_ack(
        now,
        1200,
        (),
        &rtt_estimator,
        random,
        now + Duration::from_millis(201),
        &mut publisher,
    );
    assert_eq!(cc.congestion_window(), cwnd);
}

#[test]
fn on_packet_ack_congestion_avoidance_max_cwnd() {
    let mut cc = CubicCongestionController::new(5000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let now = NoopClock.get_time();
    let mut rtt_estimator = RttEstimator::default();
    let random = &mut random::testing::Generator::default();
    rtt_estimator.update_rtt(
        Duration::from_secs(0),
        Duration::from_millis(200),
        now,
        true,
        PacketNumberSpace::ApplicationData,
    );
    cc.congestion_window = 89_000.0;
    cc.state = State::congestion_avoidance(now);
    cc.cubic.w_max = 100_000.0;

    cc.on_packet_sent(now, 60_000, Some(false), &rtt_estimator, &mut publisher);
    cc.on_ack(now, 60_000, (), &rtt_estimator, random, now, &mut publisher);

    // 60_000 is the highest bytes in flight * 1.5 for the congestion avoidance max_cwnd multiplier
    assert_eq!(90_000, cc.congestion_window());
}

#[test]
//= https://www.rfc-editor.org/rfc/rfc9002#section-7.3.2
//= type=test
fn on_packet_ack_recovery_to_congestion_avoidance() {
    let mut cc = CubicCongestionController::new(5000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let now = NoopClock.get_time();
    let random = &mut random::testing::Generator::default();

    cc.cubic.w_max = bytes_to_packets(25000.0, 5000);
    cc.state = Recovery(now, Idle);
    cc.bytes_in_flight = BytesInFlight::new(25000);
    cc.under_utilized = false;

    cc.on_ack(
        now + Duration::from_millis(1),
        1,
        (),
        &RttEstimator::default(),
        random,
        now + Duration::from_millis(2),
        &mut publisher,
    );

    assert_eq!(
        cc.state,
        State::congestion_avoidance(now + Duration::from_millis(2))
    );
}

#[test]
//= https://www.rfc-editor.org/rfc/rfc9002#section-7.3.2
//= type=test
fn on_packet_ack_slow_start_to_congestion_avoidance() {
    let mut cc = CubicCongestionController::new(5000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let now = NoopClock.get_time();
    let random = &mut random::testing::Generator::default();

    cc.state = SlowStart;
    cc.congestion_window = 10000.0;
    cc.bytes_in_flight = BytesInFlight::new(10000);
    cc.slow_start.threshold = 10050.0;
    cc.under_utilized = false;

    cc.on_ack(
        now,
        100,
        (),
        &RttEstimator::default(),
        random,
        now + Duration::from_millis(2),
        &mut publisher,
    );

    assert_delta!(cc.congestion_window, 10100.0, 0.001);
    assert_delta!(
        cc.packets_to_bytes(cc.cubic.w_max),
        cc.congestion_window,
        0.001
    );
    assert_eq!(cc.cubic.k, Duration::from_secs(0));
    assert_eq!(
        cc.state,
        State::congestion_avoidance(now + Duration::from_millis(2))
    );
}

#[test]
fn on_packet_ack_recovery() {
    let mut cc = CubicCongestionController::new(5000);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let now = NoopClock.get_time();
    let random = &mut random::testing::Generator::default();

    cc.state = Recovery(now, Idle);
    cc.congestion_window = 10000.0;
    cc.bytes_in_flight = BytesInFlight::new(10000);

    cc.on_ack(
        now,
        100,
        (),
        &RttEstimator::default(),
        random,
        now + Duration::from_millis(2),
        &mut publisher,
    );

    // Congestion window stays the same in recovery
    assert_delta!(cc.congestion_window, 10000.0, 0.001);
    assert_eq!(cc.state, Recovery(now, Idle));
}

#[test]
fn on_packet_ack_congestion_avoidance() {
    let max_datagram_size = 5000;
    let mut cc = CubicCongestionController::new(max_datagram_size);
    let mut cc2 = CubicCongestionController::new(max_datagram_size);
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let now = NoopClock.get_time();
    let random = &mut random::testing::Generator::default();

    cc.state = State::congestion_avoidance(now + Duration::from_millis(3300));
    cc.congestion_window = 10000.0;
    cc.bytes_in_flight = BytesInFlight::new(10000);
    cc.cubic.w_max = bytes_to_packets(10000.0, max_datagram_size);
    cc.under_utilized = false;

    cc2.congestion_window = 10000.0;
    cc2.bytes_in_flight = BytesInFlight::new(10000);
    cc2.cubic.w_max = bytes_to_packets(10000.0, max_datagram_size);

    let mut rtt_estimator = RttEstimator::default();
    rtt_estimator.update_rtt(
        Duration::from_secs(0),
        Duration::from_millis(275),
        now,
        true,
        PacketNumberSpace::ApplicationData,
    );

    cc.on_ack(
        now,
        1000,
        (),
        &rtt_estimator,
        random,
        now + Duration::from_millis(4750),
        &mut publisher,
    );

    let t = Duration::from_millis(4750) - Duration::from_millis(3300);
    let rtt = rtt_estimator.min_rtt();

    cc2.congestion_avoidance(t, rtt, 1000, f32::MAX);

    assert_delta!(cc.congestion_window, cc2.congestion_window, 0.001);
}

//= https://www.rfc-editor.org/rfc/rfc8312#section-4.2
//= type=test
//# If so, CUBIC is in the TCP-friendly region and cwnd SHOULD
//# be set to W_est(t) at each reception of an ACK.
#[test]
fn on_packet_ack_congestion_avoidance_tcp_friendly_region() {
    let mut cc = CubicCongestionController::new(5000);

    cc.congestion_window = 10000.0;
    cc.cubic.w_max = 2.5;
    cc.cubic.k = Duration::from_secs_f32(2.823);

    let t = Duration::from_millis(300);
    let rtt = Duration::from_millis(250);

    cc.congestion_avoidance(t, rtt, 5000, f32::MAX);

    assert!(cc.cubic.w_cubic(t) < cc.cubic.w_est(t, rtt));
    assert_delta!(cc.congestion_window, cc.cubic.w_est(t, rtt) * 5000.0, 0.001);
}

//= https://www.rfc-editor.org/rfc/rfc8312#section-4.3
//= type=test
//# In this region, cwnd MUST be incremented by
//# (W_cubic(t+RTT) - cwnd)/cwnd for each received ACK, where
//# W_cubic(t+RTT) is calculated using Eq. 1.
#[test]
fn on_packet_ack_congestion_avoidance_concave_region() {
    let max_datagram_size = 1200;
    let mut cc = CubicCongestionController::new(max_datagram_size as u16);

    cc.congestion_window = 2_400_000.0;
    cc.cubic.w_max = 2304.0;
    cc.cubic.k = Duration::from_secs(12);

    let t = Duration::from_millis(9800);
    let rtt = Duration::from_millis(200);

    cc.congestion_avoidance(t, rtt, 1000, f32::MAX);

    assert!(cc.cubic.w_cubic(t) > cc.cubic.w_est(t, rtt));

    // W_cubic(t+RTT) = C*(t-K)^3 + W_max
    // W_cubic(10) = .4*(-2)^3 + 2304
    // W_cubic(10) = 2300.8

    // cwnd = (W_cubic(t+RTT) - cwnd)/cwnd + cwnd
    // cwnd = ((2300.8 - 2000)/2000 + 2000) * max_datagram_size
    // cwnd = 2400180.48

    assert_delta!(cc.congestion_window, 2_400_180.5, 0.001);
}

//= https://www.rfc-editor.org/rfc/rfc8312#section-4.4
//= type=test
//# In this region, cwnd MUST be incremented by
//# (W_cubic(t+RTT) - cwnd)/cwnd for each received ACK, where
//# W_cubic(t+RTT) is calculated using Eq. 1.
#[test]
fn on_packet_ack_congestion_avoidance_convex_region() {
    let max_datagram_size = 1200;
    let mut cc = CubicCongestionController::new(max_datagram_size);

    cc.congestion_window = 3_600_000.0;
    cc.cubic.w_max = 2304.0;
    cc.cubic.k = Duration::from_secs(12);

    let t = Duration::from_millis(25800);
    let rtt = Duration::from_millis(200);

    cc.congestion_avoidance(t, rtt, 1000, f32::MAX);

    assert!(cc.cubic.w_cubic(t) > cc.cubic.w_est(t, rtt));

    // W_cubic(t+RTT) = C*(t-K)^3 + W_max
    // W_cubic(26) = .4*(14)^3 + 2304
    // W_cubic(26) = 3401.6

    // cwnd = (W_cubic(t+RTT) - cwnd)/cwnd + cwnd
    // cwnd = ((3401.6 - 3000)/3000 + 3000) * max_datagram_size
    // cwnd = 3600160.64

    assert_eq!(cc.congestion_window(), 3_600_160);
}

#[test]
fn on_packet_ack_congestion_avoidance_too_large_increase() {
    let max_datagram_size = 1200;
    let mut cc = CubicCongestionController::new(max_datagram_size);

    cc.congestion_window = 3_600_000.0;
    cc.cubic.w_max = bytes_to_packets(2_764_800.0, max_datagram_size);

    let t = Duration::from_millis(125_800);
    let rtt = Duration::from_millis(200);

    cc.congestion_avoidance(t, rtt, 1000, f32::MAX);

    assert!(cc.cubic.w_cubic(t) > cc.cubic.w_est(t, rtt));
    assert_delta!(cc.congestion_window, 3_600_000.0 + 1000.0 / 2.0, 0.001);
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    event,
    packet::number::PacketNumberSpace,
    path,
    path::MINIMUM_MTU,
    recovery::{
        congestion_controller::PathPublisher,
        pacing::{Pacer, INITIAL_INTERVAL, N, SLOW_START_N},
        RttEstimator, MAX_BURST_PACKETS,
    },
    time::{Clock, NoopClock, Timestamp},
};
use bolero::{check, generator::*};
use core::time::Duration;
use num_traits::ToPrimitive;

#[test]
fn earliest_departure_time() {
    let mut pacer = Pacer::default();
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    assert_eq!(None, pacer.next_packet_departure_time);

    let now = NoopClock.get_time();
    let rtt = RttEstimator::default();
    let cwnd = 12000;

    pacer.on_packet_sent(
        now,
        MINIMUM_MTU as usize,
        &rtt,
        cwnd,
        MINIMUM_MTU,
        false,
        &mut publisher,
    );

    // The initial interval is added for the second packet
    assert_eq!(
        Some(now + INITIAL_INTERVAL),
        pacer.earliest_departure_time()
    );
}

#[test]
fn on_packet_sent_large_bytes_sent() {
    let mut pacer = Pacer::default();
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());

    let now = NoopClock.get_time();
    let rtt = RttEstimator::default();
    let cwnd = 12000;

    pacer.on_packet_sent(
        now,
        usize::MAX,
        &rtt,
        cwnd,
        MINIMUM_MTU,
        false,
        &mut publisher,
    );

    assert_eq!(
        Some(now + INITIAL_INTERVAL),
        pacer.earliest_departure_time()
    );
}

#[test]
fn slow_start() {
    test_one_rtt(true);
}

#[test]
fn post_slow_start() {
    test_one_rtt(false);
}

fn test_one_rtt(slow_start: bool) {
    let mut pacer = Pacer::default();
    let now = NoopClock.get_time();
    let rtt = RttEstimator::default();
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());

    let cwnd = MINIMUM_MTU as u32 * 100;
    let n = if slow_start {
        SLOW_START_N.to_f32().unwrap()
    } else {
        N.to_f32().unwrap()
    };
    // If in slow start we should be sending 2X the cwnd in one RTT,
    // otherwise we should be sending 1.25X the cwnd.
    let bytes_to_send = ((cwnd as f32) * n) as u32;

    // Send one packet to move beyond the initial interval
    pacer.on_packet_sent(
        now,
        MINIMUM_MTU as usize,
        &rtt,
        cwnd,
        MINIMUM_MTU,
        slow_start,
        &mut publisher,
    );
    assert_eq!(
        Some(now + INITIAL_INTERVAL),
        pacer.earliest_departure_time()
    );

    let mut sent_bytes = 0;
    let now = now + INITIAL_INTERVAL;
    while sent_bytes < bytes_to_send {
        // Confirm the current departure time is less than 1 rtt
        assert!(pacer
            .earliest_departure_time()
            .map_or(true, |departure_time| departure_time
                < now + rtt.smoothed_rtt()));
        pacer.on_packet_sent(
            now,
            MINIMUM_MTU as usize,
            &rtt,
            cwnd,
            MINIMUM_MTU,
            slow_start,
            &mut publisher,
        );
        sent_bytes += MINIMUM_MTU as u32;
    }
    assert!(pacer
        .earliest_departure_time()
        .unwrap()
        .has_elapsed(now + rtt.smoothed_rtt()));
}

#[test]
fn earliest_departure_time_before_now() {
    let mut pacer = Pacer::default();
    let now = NoopClock.get_time();
    let rtt = RttEstimator::default();
    let mut publisher = event::testing::Publisher::snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());

    let cwnd = MINIMUM_MTU as u32 * 100;
    loop {
        pacer.on_packet_sent(
            now,
            MINIMUM_MTU as usize,
            &rtt,
            cwnd,
            MINIMUM_MTU,
            false,
            &mut publisher,
        );
        if pacer.capacity == 0 {
            break;
        }
    }
    assert_eq!(
        Some(now + INITIAL_INTERVAL),
        pacer.earliest_departure_time()
    );

    // The first timeslot has reached capacity, so calling `on_packet_sent` normally would move to the
    // next interval. Since the timestamp we pass in is after the next interval, the earliest departure time
    // becomes the new "now".
    let now = now + Duration::from_secs(1);
    pacer.on_packet_sent(
        now,
        MINIMUM_MTU as usize,
        &rtt,
        cwnd,
        MINIMUM_MTU,
        false,
        &mut publisher,
    );
    assert_eq!(Some(now), pacer.earliest_departure_time());
}

#[test]
fn interval_change() {
    let mut pacer = Pacer::default();
    let now = NoopClock.get_time();
    let mut rtt = RttEstimator::default();

    let mut cwnd = MINIMUM_MTU as u32 * 100;

    if INITIAL_INTERVAL > Duration::ZERO {
        let interval = get_interval(now, &mut pacer, &rtt, cwnd, MINIMUM_MTU, false);
        assert_eq!(INITIAL_INTERVAL, interval);
    }

    let interval = get_interval(now, &mut pacer, &rtt, cwnd, MINIMUM_MTU, false);

    cwnd += MINIMUM_MTU as u32;
    let new_interval = get_interval(now, &mut pacer, &rtt, cwnd, MINIMUM_MTU, false);

    // Interval decreases after the congestion window increases, as more bursts need to be
    // distributed evenly across the same time period (1 rtt)
    assert!(new_interval < interval);

    let interval = new_interval;
    rtt.update_rtt(
        Duration::default(),
        Duration::from_millis(750),
        now,
        true,
        PacketNumberSpace::ApplicationData,
    );
    let new_interval = get_interval(now, &mut pacer, &rtt, cwnd, MINIMUM_MTU, false);

    // Interval increases after the RTT increases, as the same amount of data is distributed over
    // a longer time period
    assert!(new_interval > interval);

    let interval = new_interval;
    let max_datagram_size = MINIMUM_MTU + 300;
    let new_interval = get_interval(now, &mut pacer, &rtt, cwnd, max_datagram_size, false);

    // Interval increases after the MTU increases, as each burst contains more data, so less
    // bursts need to be distributed across the same time period
    assert!(new_interval > interval);

    let interval = new_interval;
    let new_interval = get_interval(now, &mut pacer, &rtt, cwnd, max_datagram_size, true);

    // Interval decreases in slow start, as the rapidly increasing congestion window in slow start
    // means we want to allow for packets to be sent faster to avoid under utilizing the congestion window.
    assert!(new_interval < interval);
}

/// This test aims to compare the rate based implementation of pacing with the inter-packet interval
/// based implementation of pacing described in RFC 9002. Due to rounding issues while multiplying
/// and dividing, the two implementations do not match exactly, so this test asserts that the
/// two implementations differ by less than 1.1 ms.
#[test]
#[cfg_attr(miri, ignore)] // This test is too expensive for miri to complete in a reasonable amount of time
fn interval_differential_test() {
    check!()
        .with_generator((
            1_000_000..u32::MAX, // RTT ranges from 1ms to ~4sec
            2400..u32::MAX, // congestion window ranges from the minimum window (2 * MINIMUM_MTU) to u32::MAX
            MINIMUM_MTU..=9000, // max_datagram_size ranges from MINIMUM_MTU to 9000
            gen(),
        ))
        .cloned()
        .for_each(|(rtt, congestion_window, max_datagram_size, slow_start)| {
            let mut publisher = event::testing::Publisher::no_snapshot();
            let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
            let rtt = Duration::from_nanos(rtt as _);
            let actual = Pacer::interval(
                rtt,
                congestion_window,
                max_datagram_size,
                slow_start,
                &mut publisher,
            );

            let expected = rfc_interval(rtt, congestion_window, max_datagram_size, slow_start);

            assert!(
                abs_difference(actual, expected) < Duration::from_nanos(1_100_000),
                "expected: {:?}; actual: {:?}",
                expected,
                actual
            );
        });
}

//= https://www.rfc-editor.org/rfc/rfc9002#section-7.7
//= type=test
//# A perfectly paced sender spreads packets exactly evenly over time.
//# For a window-based congestion controller, such as the one in this
//# document, that rate can be computed by averaging the congestion
//# window over the RTT. Expressed as a rate in units of bytes per time,
//# where congestion_window is in bytes:
//#
//# rate = N * congestion_window / smoothed_rtt
//#
//# Or expressed as an inter-packet interval in units of time:
//#
//# interval = ( smoothed_rtt * packet_size / congestion_window ) / N
//(rtt_estimator.smoothed_rtt() * packet_size / congestion_window) / n
fn rfc_interval(
    rtt: Duration,
    congestion_window: u32,
    max_datagram_size: u16,
    slow_start: bool,
) -> Duration {
    let packet_size = MAX_BURST_PACKETS * max_datagram_size as u32;
    let result = rtt * packet_size / congestion_window;
    let n = if slow_start { SLOW_START_N } else { N };

    // Divide by n by multiplying by the inverse
    result * *n.denom() as u32 / *n.numer() as u32
}

fn abs_difference(a: Duration, b: Duration) -> Duration {
    if a > b {
        a - b
    } else {
        b - a
    }
}

// Calls `on_packet_sent` until the earliest departure time has increased, and returns the interval
// between the new earliest departure time and the original earliest departure time
fn get_interval(
    now: Timestamp,
    pacer: &mut Pacer,
    rtt_estimator: &RttEstimator,
    congestion_window: u32,
    max_datagram_size: u16,
    slow_start: bool,
) -> Duration {
    let mut publisher = event::testing::Publisher::no_snapshot();
    let mut publisher = PathPublisher::new(&mut publisher, path::Id::test_id());
    let starting_departure_time = pacer.earliest_departure_time().unwrap_or(now);

    loop {
        pacer.on_packet_sent(
            now,
            max_datagram_size as usize,
            rtt_estimator,
            congestion_window,
            max_datagram_size,
            slow_start,
            &mut publisher,
        );
        if let Some(departure_time) = pacer.earliest_departure_time() {
            if departure_time > starting_departure_time {
                return departure_time - starting_departure_time;
            }
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

//! Defines the Stateless Reset token

use core::convert::{TryFrom, TryInto};
use s2n_codec::{decoder_value, Encoder, EncoderValue};
use subtle::ConstantTimeEq;

//= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
//# Stateless Reset {
//#   Fixed Bits (2) = 1,
//#   Unpredictable Bits (38..),
//#   Stateless Reset Token (128),
//# }
pub const LEN: usize = 128 / 8;

// The implemented PartialEq will have the same results as
// a derived version, except it is constant-time. Therefore
// Hash can still be derived.
#[allow(clippy::derived_hash_with_manual_eq)]
#[derive(Copy, Clone, Debug, Eq, Hash)]
pub struct Token([u8; LEN]);

impl Token {
    /// A zeroed out stateless reset token
    pub const ZEROED: Self = Self([0; LEN]);

    /// Unwraps this token, returning the underlying array
    pub fn into_inner(self) -> [u8; LEN] {
        self.0
    }
}

impl From<[u8; LEN]> for Token {
    fn from(bytes: [u8; LEN]) -> Self {
        Self(bytes)
    }
}

impl TryFrom<&[u8]> for Token {
    type Error = core::array::TryFromSliceError;

    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
        let bytes = bytes.try_into()?;
        Ok(Self(bytes))
    }
}

impl AsRef<[u8]> for Token {
    fn as_ref(&self) -> &[u8] {
        &self.0[..]
    }
}

impl AsMut<[u8]> for Token {
    fn as_mut(&mut self) -> &mut [u8] {
        self.0.as_mut()
    }
}

impl PartialEq for Token {
    //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3.1
    //# When comparing a datagram to stateless reset token values, endpoints
    //# MUST perform the comparison without leaking information about the
    //# value of the token.
    fn eq(&self, other: &Self) -> bool {
        self.0.ct_eq(&other.0).into()
    }
}

decoder_value!(
    impl<'a> Token {
        fn decode(buffer: Buffer) -> Result<Self> {
            let (value, buffer) = buffer.decode_slice(LEN)?;
            let value: &[u8] = value.into_less_safe_slice();
            let token = Token::try_from(value).expect("slice len already verified");

            Ok((token, buffer))
        }
    }
);

impl EncoderValue for Token {
    fn encoding_size(&self) -> usize {
        LEN
    }

    fn encode<E: Encoder>(&self, encoder: &mut E) {
        self.as_ref().encode(encoder)
    }
}

/// A generator for a stateless reset token
pub trait Generator: 'static + Send {
    /// If enabled, a stateless reset packet containing the token generated
    /// by this Generator will be sent when a packet is received that cannot
    /// be matched to an existing connection. Otherwise, the packet will be
    /// dropped with no further action.
    const ENABLED: bool = true;

    /// Generates a stateless reset token.
    ///
    /// The stateless reset token MUST be difficult to guess.
    ///
    /// To enable stateless reset functionality, the stateless reset token must
    /// be generated the same for a given `local_connection_id` before and after loss of state.
    fn generate(&mut self, local_connection_id: &[u8]) -> Token;
}

#[cfg(any(test, feature = "testing"))]
pub mod testing {
    use crate::{
        stateless_reset,
        stateless_reset::token::{Token, LEN},
    };

    pub const TEST_TOKEN_1: Token = Token(11111111123456578987654321u128.to_be_bytes());
    pub const TEST_TOKEN_2: Token = Token(222222222123456578987654321u128.to_be_bytes());
    pub const TEST_TOKEN_3: Token = Token(333333333123456578987654321u128.to_be_bytes());
    pub const TEST_TOKEN_4: Token = Token(444444444123456578987654321u128.to_be_bytes());

    const KEY: u8 = 123;

    #[derive(Debug, Default)]
    pub struct Generator();

    impl stateless_reset::token::Generator for Generator {
        fn generate(&mut self, connection_id: &[u8]) -> Token {
            let mut token = [0; LEN];

            for (index, byte) in connection_id.as_ref().iter().enumerate() {
                token[index] = byte ^ KEY;
            }

            token.into()
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::stateless_reset::token::{testing::TEST_TOKEN_1, LEN};

    //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3.1
    //= type=test
    //# When comparing a datagram to stateless reset token values, endpoints
    //# MUST perform the comparison without leaking information about the
    //# value of the token.
    #[test]
    fn equality_test() {
        let token_1 = TEST_TOKEN_1;
        let token_2 = TEST_TOKEN_1;

        assert_eq!(token_1, token_2);

        for i in 0..LEN {
            let mut token = TEST_TOKEN_1;
            token.0[i] = !TEST_TOKEN_1.0[i];

            assert_ne!(TEST_TOKEN_1, token);
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{application, connection, frame::ConnectionClose, transport};
use core::{fmt, panic};

/// Errors that a stream can encounter.
#[derive(PartialEq, Eq, Debug, Copy, Clone)]
#[non_exhaustive]
pub enum StreamError {
    /// The Stream ID which was referenced is invalid
    ///
    /// This could mean the ID is no longer tracked by the Connection.
    #[non_exhaustive]
    InvalidStream {
        source: &'static panic::Location<'static>,
    },
    /// The Stream had been reset by the peer via a `RESET_STREAM` frame.
    ///
    /// Inside this frame the peer will deliver an error code, which will be
    /// provided by the parameter.
    #[non_exhaustive]
    StreamReset {
        error: application::Error,
        source: &'static panic::Location<'static>,
    },
    /// A send attempt had been performed on a Stream after it was closed
    #[non_exhaustive]
    SendAfterFinish {
        source: &'static panic::Location<'static>,
    },
    /// Attempting to write data would exceed the stream limit
    ///
    /// This is caused because the maximum possible amount
    /// of data (2^62-1 bytes) had already been written to the
    /// Stream.
    #[non_exhaustive]
    MaxStreamDataSizeExceeded {
        source: &'static panic::Location<'static>,
    },
    /// The Stream was reset due to a Connection Error
    #[non_exhaustive]
    ConnectionError { error: connection::Error },
    /// The stream is not readable
    #[non_exhaustive]
    NonReadable {
        source: &'static panic::Location<'static>,
    },
    /// The stream is not writable
    #[non_exhaustive]
    NonWritable {
        source: &'static panic::Location<'static>,
    },
    /// The stream is blocked on writing data
    ///
    /// This is caused by trying to send data before polling readiness
    #[non_exhaustive]
    SendingBlocked {
        source: &'static panic::Location<'static>,
    },
    /// The stream was provided a non-empty placeholder buffer for receiving data.
    ///
    /// The application should ensure only empty buffers are provided to receive calls,
    /// otherwise it can lead to data loss on the stream.
    #[non_exhaustive]
    NonEmptyOutput {
        source: &'static panic::Location<'static>,
    },
}

#[cfg(feature = "std")]
impl std::error::Error for StreamError {}

impl fmt::Display for StreamError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::InvalidStream { .. } => {
                write!(f, "The Stream ID which was referenced is invalid")
            }
            Self::StreamReset { error, .. } => write!(
                f,
                "The Stream had been reset with the error {:?} by {}",
                error,
                crate::endpoint::Location::Remote,
            ),
            Self::SendAfterFinish { .. } => write!(
                f,
                "A send attempt had been performed on a Stream after it was closed"
            ),
            Self::MaxStreamDataSizeExceeded { .. } => {
                write!(f, "Attempting to write data would exceed the stream limit")
            }
            Self::ConnectionError { error, .. } => error.fmt(f),
            Self::NonReadable { .. } => write!(f, "The stream is not readable"),
            Self::NonWritable { .. } => write!(f, "The stream is not writable"),
            Self::SendingBlocked { .. } => write!(f, "The stream is blocked on writing data"),
            Self::NonEmptyOutput { .. } => write!(
                f,
                "The stream was provided a non-empty placeholder buffer for receiving data."
            ),
        }
    }
}

impl StreamError {
    /// Returns the [`panic::Location`] for the error
    pub fn source(&self) -> &'static panic::Location<'static> {
        match self {
            StreamError::InvalidStream { source } => source,
            StreamError::StreamReset { source, .. } => source,
            StreamError::SendAfterFinish { source } => source,
            StreamError::MaxStreamDataSizeExceeded { source } => source,
            StreamError::ConnectionError { error } => error.source(),
            StreamError::NonReadable { source } => source,
            StreamError::NonWritable { source } => source,
            StreamError::SendingBlocked { source } => source,
            StreamError::NonEmptyOutput { source } => source,
        }
    }

    #[track_caller]
    #[inline]
    #[doc(hidden)]
    pub fn invalid_stream() -> StreamError {
        let source = panic::Location::caller();
        StreamError::InvalidStream { source }
    }

    #[track_caller]
    #[inline]
    #[doc(hidden)]
    pub fn stream_reset(error: application::Error) -> StreamError {
        let source = panic::Location::caller();
        StreamError::StreamReset { source, error }
    }

    #[track_caller]
    #[inline]
    #[doc(hidden)]
    pub fn send_after_finish() -> StreamError {
        let source = panic::Location::caller();
        StreamError::SendAfterFinish { source }
    }

    #[track_caller]
    #[inline]
    #[doc(hidden)]
    pub fn max_stream_data_size_exceeded() -> StreamError {
        let source = panic::Location::caller();
        StreamError::MaxStreamDataSizeExceeded { source }
    }

    #[track_caller]
    #[inline]
    #[doc(hidden)]
    pub fn non_readable() -> StreamError {
        let source = panic::Location::caller();
        StreamError::NonReadable { source }
    }

    #[track_caller]
    #[inline]
    #[doc(hidden)]
    pub fn non_writable() -> StreamError {
        let source = panic::Location::caller();
        StreamError::NonWritable { source }
    }

    #[track_caller]
    #[inline]
    #[doc(hidden)]
    pub fn sending_blocked() -> StreamError {
        let source = panic::Location::caller();
        StreamError::SendingBlocked { source }
    }

    #[track_caller]
    #[inline]
    #[doc(hidden)]
    pub fn non_empty_output() -> StreamError {
        let source = panic::Location::caller();
        StreamError::NonEmptyOutput { source }
    }
}

impl application::error::TryInto for StreamError {
    fn application_error(&self) -> Option<application::Error> {
        if let StreamError::ConnectionError { error, .. } = self {
            error.application_error()
        } else {
            None
        }
    }
}

impl From<connection::Error> for StreamError {
    fn from(error: connection::Error) -> Self {
        Self::ConnectionError { error }
    }
}

impl From<transport::Error> for StreamError {
    #[track_caller]
    fn from(error: transport::Error) -> Self {
        let error: connection::Error = error.into();
        error.into()
    }
}

impl<'a> From<ConnectionClose<'a>> for StreamError {
    #[track_caller]
    fn from(error: ConnectionClose) -> Self {
        let error: connection::Error = error.into();
        error.into()
    }
}

#[cfg(feature = "std")]
impl From<StreamError> for std::io::Error {
    fn from(error: StreamError) -> Self {
        let kind = error.into();
        std::io::Error::new(kind, error)
    }
}

#[cfg(feature = "std")]
impl From<StreamError> for std::io::ErrorKind {
    fn from(error: StreamError) -> Self {
        use std::io::ErrorKind;
        match error {
            StreamError::InvalidStream { .. } => ErrorKind::NotFound,
            StreamError::StreamReset { .. } => ErrorKind::ConnectionReset,
            StreamError::SendAfterFinish { .. } => ErrorKind::BrokenPipe,
            StreamError::MaxStreamDataSizeExceeded { .. } => ErrorKind::Other,
            StreamError::ConnectionError { error, .. } => error.into(),
            StreamError::NonReadable { .. } => ErrorKind::Other,
            StreamError::NonWritable { .. } => ErrorKind::Other,
            StreamError::SendingBlocked { .. } => ErrorKind::WouldBlock,
            StreamError::NonEmptyOutput { .. } => ErrorKind::InvalidInput,
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

//! Types and utilities around the QUIC Stream identifier

use crate::{endpoint, stream::StreamType, varint::VarInt};
#[cfg(any(test, feature = "generator"))]
use bolero_generator::*;

/// The ID of a stream.
///
/// A stream ID is a 62-bit integer (0 to 2^62-1) that is unique for all streams
/// on a connection.
#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Copy, Clone, Hash)]
#[cfg_attr(any(feature = "generator", test), derive(TypeGenerator))]
pub struct StreamId(VarInt);

// Stream IDs can be converted into `VarInt` and `u64`

impl From<StreamId> for VarInt {
    fn from(id: StreamId) -> Self {
        id.0
    }
}

impl From<StreamId> for u64 {
    fn from(id: StreamId) -> Self {
        id.0.as_u64()
    }
}

impl StreamId {
    /// Creates a Stream ID from a [`VarInt`].
    ///
    /// This is always a safe conversion, since Stream IDs and [`VarInt`]s
    /// share the same range.
    #[inline]
    pub const fn from_varint(id: VarInt) -> StreamId {
        StreamId(id)
    }

    /// Converts the stream id into a [`VarInt`]
    #[inline]
    pub const fn as_varint(self) -> VarInt {
        self.0
    }

    /// Returns the initial Stream ID for a given stream type.
    ///
    /// E.g. the initial Stream ID for a server initiated unidirectional Stream
    /// is Stream ID `3`.
    ///
    /// Example:
    ///
    /// ```
    /// # use s2n_quic_core::{endpoint, stream::{StreamId, StreamType}};
    /// let stream_id = StreamId::initial(endpoint::Type::Server, StreamType::Unidirectional);
    /// // Initial server initiated unidirectional Stream ID is 3
    /// assert_eq!(3u64, stream_id.as_varint().as_u64());
    /// ```
    #[inline]
    pub fn initial(initiator: endpoint::Type, stream_type: StreamType) -> StreamId {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-2.1
        //# The two least significant bits from a stream ID therefore identify a
        //# stream as one of four types, as summarized in Table 1.
        //#
        //#        +======+==================================+
        //#        | Bits | Stream Type                      |
        //#        +======+==================================+
        //#        | 0x00 | Client-Initiated, Bidirectional  |
        //#        +------+----------------------------------+
        //#        | 0x01 | Server-Initiated, Bidirectional  |
        //#        +------+----------------------------------+
        //#        | 0x02 | Client-Initiated, Unidirectional |
        //#        +------+----------------------------------+
        //#        | 0x03 | Server-Initiated, Unidirectional |
        //#        +------+----------------------------------+

        match (
            stream_type == StreamType::Bidirectional,
            initiator == endpoint::Type::Client,
        ) {
            (true, true) => StreamId(VarInt::from_u32(0)),
            (true, false) => StreamId(VarInt::from_u32(1)),
            (false, true) => StreamId(VarInt::from_u32(2)),
            (false, false) => StreamId(VarInt::from_u32(3)),
        }
    }

    /// Returns the n-th `StreamId` for a certain type of `Stream`.
    ///
    /// The 0th `StreamId` thereby represents the `StreamId` which is returned
    /// by the [`Self::initial`] method. All further `StreamId`s of a certain type
    /// will be spaced apart by 4.
    ///
    /// nth() will return `None` if the resulting `StreamId` would not be valid.
    #[inline]
    pub fn nth(initiator: endpoint::Type, stream_type: StreamType, n: u64) -> Option<StreamId> {
        let initial = Self::initial(initiator, stream_type);
        // We calculate as much as possible with u64, to reduce the number of
        // overflow checks for the maximum Stream ID to the last operation
        let id = VarInt::new(n.checked_mul(4)?.checked_add(initial.into())?).ok()?;
        Some(StreamId(id))
    }

    /// Returns the next [`StreamId`] which is of the same type the one referred
    /// to. E.g. if the method is called on a Stream ID for an unidirectional
    /// client initiated stream, the Stream ID of the next unidirectional client
    /// initiated stream will be returned.
    ///
    /// Returns `None` if the next Stream ID would not be valid, due to being out
    /// of bounds.
    ///
    /// Example:
    ///
    /// ```
    /// # use s2n_quic_core::{endpoint, stream::{StreamId, StreamType}};
    /// let stream_id = StreamId::initial(endpoint::Type::Client, StreamType::Unidirectional);
    /// // Initial client initiated unidirectional Stream ID is 2
    /// assert_eq!(2u64, stream_id.as_varint().as_u64());
    /// // Get the next client initiated Stream ID
    /// let next_stream_id = stream_id.next_of_type();
    /// assert_eq!(6u64, next_stream_id.expect("Next Stream ID is valid").as_varint().as_u64());
    /// ```
    #[inline]
    pub fn next_of_type(self) -> Option<StreamId> {
        // Stream IDs increase in steps of 4, since the 2 least significant bytes
        // are used to indicate the stream type
        self.0
            .checked_add(VarInt::from_u32(4))
            .map(StreamId::from_varint)
    }

    /// Returns whether the client or server initiated the Stream
    #[inline]
    pub fn initiator(self) -> endpoint::Type {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-2.1
        //# The least significant bit (0x1) of the stream ID identifies the
        //# initiator of the stream.  Client-initiated streams have even-numbered
        //# stream IDs (with the bit set to 0)
        if Into::<u64>::into(self.0) & 0x01u64 == 0 {
            endpoint::Type::Client
        } else {
            endpoint::Type::Server
        }
    }

    /// Returns whether the Stream is unidirectional or bidirectional.
    #[inline]
    pub fn stream_type(self) -> StreamType {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-2.1
        //# The second least significant bit (0x2) of the stream ID distinguishes
        //# between bidirectional streams (with the bit set to 0) and
        //# unidirectional streams (with the bit set to 1).
        if Into::<u64>::into(self.0) & 0x02 == 0 {
            StreamType::Bidirectional
        } else {
            StreamType::Unidirectional
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::varint::MAX_VARINT_VALUE;

    #[test]
    fn initial_stream_ids() {
        for stream_type in &[StreamType::Bidirectional, StreamType::Unidirectional] {
            for initiator in &[endpoint::Type::Client, endpoint::Type::Server] {
                let id = StreamId::initial(*initiator, *stream_type);
                assert_eq!(*stream_type, id.stream_type());
                assert_eq!(*initiator, id.initiator());
            }
        }
    }

    #[test]
    fn stream_id_overflow() {
        // Check that the highest possible Stream ID works
        let max_stream_id_varint = VarInt::new((1 << 62) - 1).unwrap();
        let _max_stream_id = StreamId::from_varint(max_stream_id_varint);

        let max_increaseable_stream_id_varint = max_stream_id_varint - 4;
        let max_inreasable_stream_id = StreamId::from_varint(max_increaseable_stream_id_varint);
        assert!(max_inreasable_stream_id.next_of_type().is_some());

        // Check all the variants where the base ID is still valid but the
        // increment is no longer.
        for increment in 1..5 {
            let id_varint = max_increaseable_stream_id_varint + increment;
            let stream_id = StreamId::from_varint(id_varint);
            assert!(stream_id.next_of_type().is_none());
        }
    }

    #[test]
    fn nth_stream_id() {
        for stream_type in &[StreamType::Bidirectional, StreamType::Unidirectional] {
            for initiator in &[endpoint::Type::Client, endpoint::Type::Server] {
                // The first StreamId is the initial one
                let first = StreamId::nth(*initiator, *stream_type, 0).unwrap();
                assert_eq!(StreamId::initial(*initiator, *stream_type), first);

                for n in 1..10 {
                    let nth = StreamId::nth(*initiator, *stream_type, n).unwrap();
                    assert_eq!(VarInt::from_u32(n as u32 * 4), nth.0 - first.0);
                }
            }
        }
    }

    #[test]
    fn invalid_nth_stream_id() {
        for stream_type in &[StreamType::Bidirectional, StreamType::Unidirectional] {
            for initiator in &[endpoint::Type::Client, endpoint::Type::Server] {
                assert_eq!(
                    None,
                    StreamId::nth(
                        *initiator,
                        *stream_type,
                        Into::<u64>::into(MAX_VARINT_VALUE / 2)
                    )
                );
            }
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::StreamId;
#[cfg(any(test, feature = "generator"))]
use bolero_generator::*;

/// An Iterator over Stream Ids of a particular type.
#[derive(Debug, Clone, Copy)]
#[cfg_attr(any(feature = "generator", test), derive(TypeGenerator))]
pub struct StreamIter {
    start_stream_id: StreamId,
    max_stream_id: StreamId,
    finished: bool,
}

impl StreamIter {
    #[inline]
    pub fn new(start_stream_id: StreamId, max_stream_id: StreamId) -> Self {
        debug_assert_eq!(start_stream_id.stream_type(), max_stream_id.stream_type());
        debug_assert_eq!(start_stream_id.initiator(), max_stream_id.initiator());
        debug_assert!(start_stream_id <= max_stream_id);

        Self {
            start_stream_id,
            max_stream_id,
            finished: false,
        }
    }

    #[inline]
    pub fn max_stream_id(self) -> StreamId {
        self.max_stream_id
    }
}

impl Iterator for StreamIter {
    type Item = StreamId;

    fn next(&mut self) -> Option<Self::Item> {
        // short circuit when finished
        if self.finished {
            return None;
        }

        match self.start_stream_id.cmp(&self.max_stream_id) {
            core::cmp::Ordering::Less => {
                let ret = self.start_stream_id;
                // The Stream ID can be expected to be valid, since `max_stream_id`
                // is a valid `StreamId` and all IDs we iterate over are lower.
                self.start_stream_id = self
                    .start_stream_id
                    .next_of_type()
                    .expect("Expect a valid Stream ID");
                Some(ret)
            }
            core::cmp::Ordering::Equal => {
                // Avoid incrementing beyond `max_stream_id` and mark finished to
                // to avoid returning max value again
                self.finished = true;
                Some(self.start_stream_id)
            }
            core::cmp::Ordering::Greater => {
                debug_assert!(false, "The `new` method should verify valid ranges");

                // finished
                self.finished = true;
                None
            }
        }
    }
}

#[cfg(test)]
mod fuzz_target {
    use super::*;

    #[test]
    #[cfg_attr(miri, ignore)] // This test is too expensive for miri to complete in a reasonable amount of time
    #[cfg_attr(kani, kani::proof, kani::unwind(1), kani::solver(kissat))]
    fn fuzz_builder() {
        bolero::check!()
            .with_type::<(StreamId, StreamId)>()
            .cloned()
            .for_each(|(min, max)| {
                // enforce min <= max
                if min > max {
                    return;
                }
                // enforce same initiator type
                if min.initiator() != max.initiator() {
                    return;
                }
                // enforce same stream type
                if min.stream_type() != max.stream_type() {
                    return;
                }

                // All other combinations of min/max StreamId should be valid
                StreamIter::new(min, max);
            });
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    transport::parameters::{InitialMaxStreamsBidi, InitialMaxStreamsUni, ValidationError},
    varint::VarInt,
};

// TODO investigate a good default
/// The default send buffer size for Streams
const DEFAULT_STREAM_MAX_SEND_BUFFER_SIZE: u32 = 512 * 1024;

pub trait LocalLimits {
    fn as_varint(&self) -> VarInt;
}

/// Per-stream limits
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct Limits {
    /// The maximum send buffer size for a Stream
    pub max_send_buffer_size: MaxSendBufferSize,
    /// The maximum number of unidirectional streams that may
    /// be opened concurrently by the local endpoint. This value
    /// is not communicated to the peer, it is only used for limiting
    /// concurrent streams opened locally by the application.
    pub max_open_local_unidirectional_streams: LocalUnidirectional,
    /// The maximum number of bidirectional streams that may
    /// be opened concurrently by the local endpoint. This value
    /// is not communicated to the peer, it is only used for limiting
    /// concurrent streams opened locally by the application.
    pub max_open_local_bidirectional_streams: LocalBidirectional,
}

impl Default for Limits {
    fn default() -> Self {
        Self::RECOMMENDED
    }
}

impl Limits {
    pub const RECOMMENDED: Self = Self {
        max_send_buffer_size: MaxSendBufferSize::RECOMMENDED,
        max_open_local_unidirectional_streams: LocalUnidirectional::RECOMMENDED,
        max_open_local_bidirectional_streams: LocalBidirectional::RECOMMENDED,
    };
}

macro_rules! local_limits {
    ($name:ident($encodable_type:ty)) => {
        #[derive(Debug, Copy, Clone, PartialEq, Eq)]
        pub struct $name($encodable_type);
    };
}

// VarInt specific functionality
macro_rules! varint_local_limits {
    ($name:ident($encodable_type:ty)) => {
        local_limits!($name($encodable_type));

        impl LocalLimits for $name {
            #[inline]
            fn as_varint(&self) -> VarInt {
                self.0
            }
        }

        impl TryFrom<u64> for $name {
            type Error = ValidationError;

            #[inline]
            fn try_from(value: u64) -> Result<Self, Self::Error> {
                let value = VarInt::new(value)?;
                Ok(Self(value))
            }
        }
    };
}

local_limits!(MaxSendBufferSize(u32));

impl MaxSendBufferSize {
    pub const RECOMMENDED: Self = Self(DEFAULT_STREAM_MAX_SEND_BUFFER_SIZE);

    #[inline]
    pub fn as_u32(self) -> u32 {
        self.0
    }
}

impl TryFrom<u32> for MaxSendBufferSize {
    type Error = ValidationError;

    #[inline]
    fn try_from(value: u32) -> Result<Self, Self::Error> {
        Ok(Self(value))
    }
}

varint_local_limits!(LocalUnidirectional(VarInt));

impl LocalUnidirectional {
    pub const RECOMMENDED: Self = Self(InitialMaxStreamsUni::RECOMMENDED.as_varint());
}

varint_local_limits!(LocalBidirectional(VarInt));

impl LocalBidirectional {
    pub const RECOMMENDED: Self = Self(InitialMaxStreamsBidi::RECOMMENDED.as_varint());
}

impl From<InitialMaxStreamsBidi> for LocalBidirectional {
    fn from(value: InitialMaxStreamsBidi) -> Self {
        Self(value.as_varint())
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

mod error;
mod id;
pub mod iter;
pub mod limits;
#[cfg(feature = "alloc")]
pub mod ops;
mod type_;

pub use error::*;
pub use id::*;
pub use limits::Limits;
pub use type_::*;

#[cfg(any(test, feature = "testing"))]
pub mod testing;


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

//! Bulk operations performed on streams
//!
//! By representing stream operations as structs, callers can request multiple tasks to be
//! performed in a single call, which reduces context switching.
//!
//! Consider the following scenario where we send 3 chunks of data and finish the stream:
//!
//! ```rust,ignore
//! stream.send(a).await?;
//! stream.send(b).await?;
//! stream.send(c).await?;
//! stream.finish().await?;
//! ```
//!
//! This will result in at least 4 context switches (and potentially even more if the stream
//! is currently blocked on sending).
//!
//! Using the bulk operation API greatly reduces this amount:
//!
//! ```rust,ignore
//! stream
//!     .request()
//!     .send(&mut [a, b, c])
//!     .finish()
//!     .await?;
//! ```

use crate::{application, stream};
use core::task::Poll;

/// A request made on a stream
#[derive(Default, Debug)]
pub struct Request<'a> {
    /// The `tx` options of the request
    pub tx: Option<tx::Request<'a>>,

    /// The `rx` options of the request
    pub rx: Option<rx::Request<'a>>,
}

impl<'a> Request<'a> {
    /// Requests a slice of chunks to be sent on the tx stream
    pub fn send(&mut self, chunks: &'a mut [bytes::Bytes]) -> &mut Self {
        self.tx_mut().chunks = Some(chunks);
        self
    }

    /// Resets the tx stream with an error code
    pub fn reset(&mut self, error: application::Error) -> &mut Self {
        self.tx_mut().reset = Some(error);
        self
    }

    /// Flushes any pending tx data to be ACKed before unblocking
    pub fn flush(&mut self) -> &mut Self {
        self.tx_mut().flush = true;
        self
    }

    /// Marks the tx stream as finished (e.g. no more data will be sent)
    pub fn finish(&mut self) -> &mut Self {
        self.tx_mut().finish = true;
        self
    }

    /// Requests data on the rx stream to be received into the provided slice of chunks
    pub fn receive(&mut self, chunks: &'a mut [bytes::Bytes]) -> &mut Self {
        self.rx_mut().chunks = Some(chunks);
        self
    }

    /// Requests the peer to stop sending data on the rx stream
    pub fn stop_sending(&mut self, error: application::Error) -> &mut Self {
        self.rx_mut().stop_sending = Some(error);
        self
    }

    /// Sets the watermarks for the rx stream
    pub fn with_watermark(&mut self, low: usize, high: usize) -> &mut Self {
        let rx = self.rx_mut();
        rx.low_watermark = low.min(high);
        rx.high_watermark = high.max(low);
        self
    }

    /// Sets the low watermark for the rx stream
    ///
    /// If the watermark is set to `0`, the caller will be notified as soon as there is data
    /// available on the stream.
    ///
    /// If the watermark is greater than `0`, the caller will be notified as soon as there is at
    /// least `low` bytes available to be read. Note that the stream may be woken earlier.
    pub fn with_low_watermark(&mut self, low: usize) -> &mut Self {
        let rx = self.rx_mut();
        rx.low_watermark = low;
        // raise the high watermark to be at least the lower
        rx.high_watermark = rx.high_watermark.max(low);
        self
    }

    /// Sets the high watermark for the rx stream
    ///
    /// The stream ensures that all the received data will not exceed the watermark amount. This
    /// can be useful for receiving at most `n` bytes.
    pub fn with_high_watermark(&mut self, high: usize) -> &mut Self {
        let rx = self.rx_mut();
        rx.high_watermark = high;
        // lower the low watermark to be less than the higher
        rx.low_watermark = rx.low_watermark.min(high);
        self
    }

    pub fn detach_tx(&mut self) -> &mut Self {
        let tx = self.tx_mut();
        tx.detached = true;
        self
    }

    pub fn detach_rx(&mut self) -> &mut Self {
        let rx = self.rx_mut();
        rx.detached = true;
        self
    }

    /// Lazily creates and returns the `tx` request
    fn tx_mut(&mut self) -> &mut tx::Request<'a> {
        if self.tx.is_none() {
            self.tx = Some(Default::default());
        }
        self.tx.as_mut().expect("tx should always be initialized")
    }

    /// Lazily creates and returns the `rx` request
    fn rx_mut(&mut self) -> &mut rx::Request<'a> {
        if self.rx.is_none() {
            self.rx = Some(Default::default());
        }
        self.rx.as_mut().expect("rx should always be initialized")
    }
}

/// A response received after executing a request
#[derive(Debug, Default, PartialEq, Eq)]
pub struct Response {
    /// The `tx` information of the response
    pub tx: Option<tx::Response>,

    /// The `rx` information of the response
    pub rx: Option<rx::Response>,
}

impl Response {
    /// Returns `true` if either the `rx` or `tx` requests will wake the provided waker at a later
    /// point in time.
    pub fn is_pending(&self) -> bool {
        self.tx.iter().any(|tx| tx.is_pending()) || self.rx.iter().any(|rx| rx.is_pending())
    }

    /// Returns the `tx` response
    pub fn tx(&self) -> Option<&tx::Response> {
        self.tx.as_ref()
    }

    /// Returns the `rx` response
    pub fn rx(&self) -> Option<&rx::Response> {
        self.rx.as_ref()
    }
}

/// Request and response related to transmitting on a stream
pub mod tx {
    use super::*;

    /// A request on a `tx` stream
    #[derive(Default, Debug)]
    pub struct Request<'a> {
        /// Optionally transmit chunks onto the stream
        ///
        /// The chunks will be replaced with empty buffers as they are stored in the transmission
        /// buffer. The response will indicate how many chunks and bytes were consumed from
        /// this slice.
        pub chunks: Option<&'a mut [bytes::Bytes]>,

        /// Optionally reset the stream with an error
        pub reset: Option<application::Error>,

        /// Waits for an ACK on resets and finishes
        pub flush: bool,

        /// Marks the tx stream as finished (e.g. no more data will be sent)
        pub finish: bool,

        /// Marks the tx stream as detached, which makes the stream make progress, regardless of
        /// application observations.
        pub detached: bool,
    }

    /// The result of a tx request
    #[derive(Debug, PartialEq, Eq)]
    pub struct Response {
        /// Information about the bytes that were sent
        pub bytes: Bytes,

        /// Information about the chunks that were sent
        pub chunks: Chunks,

        /// Indicates if the operation resulted in storing the provided waker to notify when the
        /// request may be polled again.
        pub will_wake: bool,

        /// The current status of the stream
        pub status: Status,
    }

    impl Default for Response {
        fn default() -> Self {
            Self {
                bytes: Bytes::default(),
                chunks: Chunks::default(),
                will_wake: false,
                status: Status::Open,
            }
        }
    }

    impl Response {
        /// Returns true if provided waker will be woken
        pub fn is_pending(&self) -> bool {
            self.will_wake
        }

        /// Returns the `tx` response
        pub fn tx(&self) -> Option<&Self> {
            Some(self)
        }
    }
}

/// Request and response related to receiving on a stream
pub mod rx {
    use super::*;

    /// A request on a `rx` stream
    #[derive(Debug)]
    pub struct Request<'a> {
        /// Optionally receive chunks from the stream
        ///
        /// At least one of the provided chunks should be empty, as it will be replaced by the
        /// received data from the stream. The response will indicate how many chunks and
        /// bytes were consumed from the stream into the provided slice.
        pub chunks: Option<&'a mut [bytes::Bytes]>,

        /// Sets the low watermark for the rx stream
        ///
        /// If the watermark is set to `0`, the caller will be notified as soon as there is data
        /// available on the stream.
        ///
        /// If the watermark is greater than `0`, the caller will be notified as soon as there is at
        /// least `low` bytes available to be read. Note that the stream may be woken earlier.
        pub low_watermark: usize,

        /// Sets the high watermark for the rx stream
        ///
        /// The stream ensures that all the received data will not exceed the watermark amount. This
        /// can be useful for receiving at most `n` bytes.
        pub high_watermark: usize,

        /// Optionally requests the peer to stop sending data with an error
        pub stop_sending: Option<application::Error>,

        /// Marks the rx stream as detached, which makes the stream make progress, regardless of
        /// application observations.
        pub detached: bool,
    }

    impl<'a> Default for Request<'a> {
        fn default() -> Self {
            Self {
                chunks: None,
                low_watermark: 0,
                high_watermark: core::usize::MAX,
                stop_sending: None,
                detached: false,
            }
        }
    }

    /// The result of a pop operation
    #[derive(Debug, PartialEq, Eq)]
    pub struct Response {
        /// Information about the bytes that were received
        pub bytes: Bytes,

        /// Information about the chunks that were received
        pub chunks: Chunks,

        /// Indicates if the operation resulted in storing the provided waker to notify when the
        /// request may be polled again.
        pub will_wake: bool,

        /// The current status of the stream
        pub status: Status,
    }

    impl Default for Response {
        fn default() -> Self {
            Self {
                bytes: Bytes::default(),
                chunks: Chunks::default(),
                will_wake: false,
                status: Status::Open,
            }
        }
    }

    impl Response {
        /// Returns true if provided waker will be woken
        pub fn is_pending(&self) -> bool {
            self.will_wake
        }

        /// Returns the `rx` response
        pub fn rx(&self) -> Option<&Self> {
            Some(self)
        }
    }
}

#[derive(Debug, Default, PartialEq, Eq)]
pub struct Bytes {
    /// The number of bytes that were consumed by the operation.
    ///
    /// In the case of `tx` operations, this is the number of bytes that were sent on the
    /// stream.
    ///
    /// In the case of `rx` operations, this is the number of bytes that were received from the
    /// stream.
    pub consumed: usize,

    /// The number of bytes that are available on the stream.
    ///
    /// In the case of `tx` operations, this is the number of additional bytes that can be sent
    /// in the stream. Note that this is not a hard limit on accepting a chunk of data.
    ///
    /// In the case of `rx` operations, this is the number of additional bytes that can be received
    /// from the stream.
    pub available: usize,
}

#[derive(Debug, Default, PartialEq, Eq)]
pub struct Chunks {
    /// The number of chunks that were consumed by the operation.
    ///
    /// In the case of `tx` operations, this is the number of chunks that were sent on the
    /// stream.
    ///
    /// In the case of `rx` operations, this is the number of chunks that were received from the
    /// stream.
    pub consumed: usize,

    /// The number of chunks that are available on the stream.
    ///
    /// In the case of `tx` operations, this is the number of additional chunks that can be sent
    /// in the stream. This value will be based on the assumption of 1 byte chunks and will
    /// contain the same value of `bytes.available`.
    ///
    /// In the case of `rx` operations, this is the number of additional chunks that can be received
    /// from the stream.
    pub available: usize,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum Status {
    /// The stream is open and writable
    Open,

    /// The stream is finishing but still has data to be flushed
    Finishing,

    /// The stream is finished and completely flushed
    Finished,

    /// The stream has been reset locally but has not been acknowledged by the peer
    Resetting,

    /// The stream was reset either by the peer or locally
    Reset(stream::StreamError),
}

macro_rules! impl_status {
    (| $self:ident | $value:expr) => {
        /// Returns `true` if the status is `Open`
        pub fn is_open(&self) -> bool {
            matches!(self.status(), Status::Open)
        }

        /// Returns `true` if the status is `Finishing`
        pub fn is_finishing(&self) -> bool {
            matches!(self.status(), Status::Finishing)
        }

        /// Returns `true` if the status is `Finished`
        pub fn is_finished(&self) -> bool {
            matches!(self.status(), Status::Finished)
        }

        /// Returns `true` if the status is `Resetting`
        pub fn is_resetting(&self) -> bool {
            matches!(self.status(), Status::Resetting)
        }

        /// Returns `true` if the status is `Reset`
        pub fn is_reset(&self) -> bool {
            matches!(self.status(), Status::Reset(_))
        }

        /// Returns `true` if the status is `Finishing` or `Resetting`
        pub fn is_closing(&self) -> bool {
            self.is_finishing() || self.is_resetting()
        }

        /// Returns `true` if the status is `Finished` or `Reset`
        pub fn is_closed(&self) -> bool {
            self.is_finished() || self.is_reset()
        }

        const fn status(&$self) -> Status {
            $value
        }
    };
}

impl Status {
    impl_status!(|self| *self);
}

impl rx::Response {
    impl_status!(|self| self.status);
}

impl tx::Response {
    impl_status!(|self| self.status);
}

macro_rules! conversions {
    ($name:path) => {
        impl $name {
            /// Converts the response into a `Poll<Self>`
            pub fn into_poll(self) -> Poll<Self> {
                if self.is_pending() {
                    Poll::Pending
                } else {
                    Poll::Ready(self)
                }
            }
        }

        impl From<$name> for () {
            fn from(_: $name) {}
        }

        impl<T, E> From<$name> for Poll<Result<T, E>>
        where
            $name: Into<T>,
        {
            fn from(v: $name) -> Poll<Result<T, E>> {
                if v.is_pending() {
                    Poll::Pending
                } else {
                    Poll::Ready(Ok(v.into()))
                }
            }
        }
    };
}

conversions!(Response);
conversions!(tx::Response);
conversions!(rx::Response);

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn request_builder_test() {
        let mut request = Request::default();
        let mut send_chunks = [bytes::Bytes::from_static(&[1])];
        let mut receive_chunks = [
            bytes::Bytes::from_static(&[2]),
            bytes::Bytes::from_static(&[3]),
        ];

        request
            .send(&mut send_chunks)
            .finish()
            .flush()
            .reset(application::Error::new(1).unwrap())
            .receive(&mut receive_chunks)
            .with_watermark(5, 10)
            .stop_sending(application::Error::new(2).unwrap());

        assert!(matches!(
            request,
            Request {
                tx: Some(tx::Request {
                    chunks: Some(tx_chunks),
                    finish: true,
                    flush: true,
                    reset: Some(reset),
                    detached: false,
                }),
                rx: Some(rx::Request {
                    chunks: Some(rx_chunks),
                    low_watermark: 5,
                    high_watermark: 10,
                    stop_sending: Some(stop_sending),
                    detached: false,
                })
            } if reset == application::Error::new(1).unwrap()
              && stop_sending == application::Error::new(2).unwrap()
              && tx_chunks.len() == 1
              && rx_chunks.len() == 2
        ));
    }

    #[test]
    fn response_pending_test() {
        for rx_pending in [false, true] {
            for tx_pending in [false, true] {
                let response = Response {
                    tx: Some(tx::Response {
                        will_wake: tx_pending,
                        ..Default::default()
                    }),
                    rx: Some(rx::Response {
                        will_wake: rx_pending,
                        ..Default::default()
                    }),
                };

                assert_eq!(response.is_pending(), rx_pending || tx_pending);

                if rx_pending || tx_pending {
                    assert_eq!(response.into_poll(), Poll::Pending);
                } else {
                    assert_eq!(
                        response.into_poll(),
                        Poll::Ready(Response {
                            tx: Some(tx::Response::default()),
                            rx: Some(rx::Response::default()),
                        })
                    );
                }
            }
        }
    }
}