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

use crate::{
    contexts::{OnTransmitError, WriteContext},
    transmission,
};
use bytes::Bytes;
use core::convert::TryInto;
use s2n_quic_core::{ack, interval_set::IntervalSet, packet::number::PacketNumber, varint::VarInt};

mod buffer;
mod traits;
mod transmissions;
pub mod writer;

pub use buffer::View;
use s2n_quic_core::stream::StreamError;
pub use traits::*;

/// Enumerates states of the [`DataSender`]
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum State {
    /// Outgoing data is accepted and transmitted
    Sending,
    /// The finish procedure has been initiated. New outgoing data is no longer
    /// accepted. The stream will continue to transmit data until all outgoing
    /// data has been transmitted and acknowledged successfully.
    /// In that case the `FinishAcknowledged` state will be entered.
    Finishing(FinState),
    /// All outgoing data including the FIN flag had been acknowledged.
    /// The Stream is thereby finalized.
    Finished,
    /// Sending data was cancelled due to a Stream RESET.
    Cancelled(StreamError),
}

impl State {
    fn fin_state_mut(&mut self) -> Option<&mut FinState> {
        if let Self::Finishing(state) = self {
            Some(state)
        } else {
            None
        }
    }

    fn can_transmit_fin(&self, constraint: transmission::Constraint, is_blocked: bool) -> bool {
        match self {
            // lost frames are not blocked by flow control since we've already acquired those
            // credits on the initial transmission
            Self::Finishing(FinState::Lost) => constraint.can_retransmit(),
            Self::Finishing(FinState::Pending) => !is_blocked && constraint.can_transmit(),
            _ => false,
        }
    }

    fn is_inflight(&self) -> bool {
        matches!(self, Self::Finishing(FinState::InFlight(_)))
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum FinState {
    Pending,
    InFlight(PacketNumber),
    Lost,
    Acknowledged,
}

impl FinState {
    pub fn is_acknowledged(self) -> bool {
        matches!(self, Self::Acknowledged)
    }

    /// This method gets called when a packet delivery got acknowledged
    fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A) {
        if let Self::InFlight(packet) = self {
            if ack_set.contains(*packet) {
                *self = Self::Acknowledged;
            }
        }
    }

    /// This method gets called when a packet loss is reported
    ///
    /// Returns `true` if the fin bit was lost
    fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) -> bool {
        if let Self::InFlight(packet) = self {
            if ack_set.contains(*packet) {
                *self = Self::Lost;
                return true;
            }
        }

        false
    }

    fn on_transmit(&mut self, packet: PacketNumber) {
        if matches!(self, Self::Pending | Self::Lost) {
            *self = Self::InFlight(packet);
        }
    }
}

/// Manages the transmission of all `Stream` and `Crypto` data frames towards
/// the peer as long as the `Stream` has not been reset or closed.
#[derive(Debug)]
pub struct DataSender<FlowController, ChunkToFrameWriter> {
    /// The data that needs to get transmitted
    buffer: buffer::Buffer,
    /// The current transmissions for the sender
    transmissions: transmissions::Transmissions<FlowController, ChunkToFrameWriter>,
    /// The offset of data waiting to be transmitted for the first time
    transmission_offset: VarInt,
    /// All of the intervals in the buffer that are waiting to be transmitted, ACKed or lost
    pending: IntervalSet<VarInt>,
    /// All of the intervals that have been declared lost
    lost: IntervalSet<VarInt>,
    /// The maximum amount of bytes that are buffered within the sending stream.
    /// This capacity will not be exceeded - even if the remote provides us a
    /// bigger flow control window.
    max_buffer_capacity: VarInt,
    /// Whether the size of the send stream is known and a FIN flag is already
    /// enqueued.
    state: State,
}

impl<FlowController: OutgoingDataFlowController, Writer: FrameWriter>
    DataSender<FlowController, Writer>
{
    /// Creates a new `DataSender` instance.
    ///
    /// `initial_window` denotes the amount of data we are allowed to send to the
    /// peer as known through transport parameters.
    /// `maximum_buffer_capacity` is the maximum amount of data the queue will
    /// hold. If users try to enqueue more data, it will be rejected in order to
    /// provide back-pressure on the `Stream`.
    pub fn new(flow_controller: FlowController, max_buffer_capacity: u32) -> Self {
        Self {
            buffer: Default::default(),
            transmissions: transmissions::Transmissions::new(flow_controller),
            transmission_offset: VarInt::from_u32(0),
            pending: IntervalSet::new(),
            lost: IntervalSet::new(),
            max_buffer_capacity: VarInt::from_u32(max_buffer_capacity),
            state: State::Sending,
        }
    }

    /// Declares all inflight packets as lost.
    pub fn on_all_lost(&mut self) {
        self.on_packet_loss(&self.transmissions.get_inflight_range());
    }

    /// Creates a new `DataSender` instance in its final
    /// [`DataSenderState::FinishAcknowledged`] state.
    pub fn new_finished(flow_controller: FlowController, max_buffer_capacity: u32) -> Self {
        let mut result = Self::new(flow_controller, max_buffer_capacity);
        result.state = State::Finished;
        result
    }

    /// Returns the flow controller for this `DataSender`
    pub fn flow_controller(&self) -> &FlowController {
        &self.transmissions.flow_controller
    }

    /// Returns the flow controller for this `DataSender`
    pub fn flow_controller_mut(&mut self) -> &mut FlowController {
        &mut self.transmissions.flow_controller
    }

    /// Stops sending out outgoing data.
    ///
    /// This is a one-way operation - sending can not be resumed.
    ///
    /// Calling the method removes all pending outgoing data as well as
    /// all tracking information from the buffer.
    pub fn stop_sending(&mut self, error: StreamError) {
        if self.state == State::Finished {
            return;
        }

        self.state = State::Cancelled(error);
        self.buffer.clear();
        self.pending.clear();
        self.lost.clear();
        self.transmissions.finish();
        self.transmission_offset = VarInt::from_u8(0);
        self.check_integrity();
    }

    /// Returns the amount of bytes that have ever been enqueued for writing on
    /// this Stream. This equals the offset of the highest enqueued byte + 1.
    pub fn total_enqueued_len(&self) -> VarInt {
        self.buffer.total_len()
    }

    /// Returns true if the data sender doesn't have any data enqueued for sending
    pub fn is_empty(&self) -> bool {
        self.buffer.is_empty()
    }

    /// Returns the state of the sender
    pub fn state(&self) -> State {
        self.state
    }

    /// Returns `true` if the delivery is currently in progress.
    pub fn is_inflight(&self) -> bool {
        !self.transmissions.is_empty() || self.state.is_inflight()
    }

    /// Overwrites the amount of total received and acknowledged bytes.
    ///
    /// This method is only used for testing purposes, in order to simulate a
    /// large number of already received bytes. The value is normally updated
    /// as an implementation detail!
    #[cfg(test)]
    pub fn set_total_acknowledged_len(&mut self, total_acknowledged: VarInt) {
        assert_eq!(
            VarInt::from_u8(0),
            self.total_enqueued_len(),
            "set_total_acknowledged_len can only be called on a new stream"
        );
        self.buffer.set_offset(total_acknowledged);
    }

    /// Returns the amount of data that can be additionally buffered for sending
    ///
    /// This depends on the configured maximum buffer size.
    /// We do not utilize the window size that the peer provides us in order to
    /// avoid excessive buffering in case the peer would provide a very big window.
    pub fn available_buffer_space(&self) -> usize {
        self.max_buffer_capacity
            .saturating_sub(self.buffer.enqueued_len())
            .try_into()
            .unwrap_or(usize::MAX)
    }

    /// Enqueues the data for transmission.
    ///
    /// It is only allowed to enqueue bytes if they do not overflow the maximum
    /// allowed Stream window (of the maximum VarInt size). This is not checked
    /// inside this method. The already enqueued bytes can be retrieved by
    /// calling [`total_enqueued_len()`].
    pub fn push(&mut self, data: Bytes) {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-4.5
        //# An endpoint MUST NOT send data on a stream at or beyond the final
        //# size.
        debug_assert!(
            self.state == State::Sending,
            "Data transmission is not allowed after finish() was called"
        );
        debug_assert!(
            data.len() <= core::u32::MAX as usize,
            "Maximum data size exceeded"
        );

        if data.is_empty() {
            return;
        }

        self.pending
            .insert(self.buffer.push(data))
            .expect("pending should not have a limit");

        self.check_integrity();
    }

    /// Starts the finalization process of a `Stream` by enqueuing a `FIN` frame.
    pub fn finish(&mut self) {
        if self.state != State::Sending {
            return;
        }

        if Writer::WRITES_FIN {
            self.state = State::Finishing(FinState::Pending);
        } else {
            self.state = State::Finishing(FinState::Acknowledged);
        }

        self.check_integrity();
    }

    /// This method gets called when a packet delivery got acknowledged
    pub fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A) {
        // If we do not get acknowledgements for any in flight data don't try
        // to release buffer chunks

        let pending = &mut self.pending;

        let any_acked = self.transmissions.on_ack_signal(ack_set, |range| {
            pending
                .remove(range)
                .expect("output should not have a limit");
        });

        if Writer::WRITES_FIN {
            if let Some(fin_state) = self.state.fin_state_mut() {
                fin_state.on_packet_ack(ack_set);
            }
        }

        if any_acked {
            // Remove any newly acked intervals from lost
            self.lost
                .intersection(pending)
                .expect("lost has no interval limit");
            if let Some(first) = self.pending.min_value() {
                self.buffer.release(first);
            } else {
                // the pending list was completely cleared
                self.buffer.release_all();
                // We don't need to track transmissions for already acked ranges
                self.transmissions.clear();
            }
        }

        // If the FIN was enqueued, and all outgoing data had been transmitted,
        // then we have finalized the stream.
        if matches!(self.state, State::Finishing(FinState::Acknowledged)) && self.is_idle() {
            self.state = State::Finished;
            self.flow_controller_mut().finish();
            self.buffer.release_all();
        }

        self.check_integrity();
    }

    fn is_idle(&self) -> bool {
        self.transmissions.is_empty() && self.pending.is_empty() && self.lost.is_empty()
    }

    /// This method gets called when a packet loss is reported
    pub fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) {
        let lost = &mut self.lost;

        let mut any_lost = self.transmissions.on_ack_signal(ack_set, |range| {
            lost.insert(range).expect("output should not have a limit");
        });

        if Writer::WRITES_FIN {
            if let Some(fin_state) = self.state.fin_state_mut() {
                any_lost |= fin_state.on_packet_loss(ack_set);
            }
        }

        // Since a packet needs to get retransmitted, we are
        // no longer blocked on waiting for flow control windows
        if any_lost {
            self.flow_controller_mut().clear_blocked();
            // Remove any lost intervals that had already been sent
            // and acked in a different packet
            self.lost
                .intersection(&self.pending)
                .expect("lost has no interval limit");
        }

        self.check_integrity();
    }

    /// Queries the component for any outgoing frames that need to get sent
    #[inline]
    pub fn on_transmit<W: WriteContext>(
        &mut self,
        writer_context: Writer::Context,
        context: &mut W,
    ) -> Result<(), OnTransmitError> {
        let initial_capacity = context.remaining_capacity();

        match self.on_transmit_impl(writer_context, context) {
            // only return an error if we didn't write anything
            Err(_) if context.remaining_capacity() < initial_capacity => Ok(()),
            other => other,
        }
    }

    fn on_transmit_impl<W: WriteContext>(
        &mut self,
        writer_context: Writer::Context,
        context: &mut W,
    ) -> Result<(), OnTransmitError> {
        let constraint = context.transmission_constraint();

        let mut transmitted_lost = false;
        // try to retransmit any lost ranges first
        if constraint.can_retransmit() {
            transmitted_lost = self.transmissions.transmit_set(
                &self.buffer,
                &mut self.lost,
                &mut self.state,
                writer_context,
                context,
            )?;
        }

        let is_blocked = self.flow_controller().is_blocked();

        // try to transmit the enqueued ranges
        let total_len = self.buffer.total_len();

        let starting_transmission_offset = self.transmission_offset;

        if !is_blocked && constraint.can_transmit() && self.transmission_offset < total_len {
            let mut viewer = self.buffer.viewer();
            self.transmission_offset = self
                .transmissions
                .transmit_interval(
                    &mut viewer,
                    (self.transmission_offset..total_len).into(),
                    &mut self.state,
                    writer_context,
                    context,
                )?
                .end_exclusive();
        }

        if Writer::WRITES_FIN && self.state.can_transmit_fin(constraint, is_blocked) {
            self.transmissions.transmit_fin(
                &self.buffer,
                &mut self.state,
                writer_context,
                context,
            )?;
        }

        // If the current transmission is a loss recovery probe, we can include some already
        // transmitted, unacknowledged data in the probe packet since there is a higher likelihood
        // this data has been lost. If lost data has already been written to the packet, we
        // skip this feature as an optimization to avoid having to filter out already written
        // lost data. Since it is unlikely there is lost data requiring retransmission at the
        // same time as a probe transmission is being sent, this optimization does not have
        // much impact on the effectiveness of this feature.
        let retransmit_unacked_data_in_probe = Writer::RETRANSMIT_IN_PROBE
            && context.transmission_mode().is_loss_recovery_probing()
            && !transmitted_lost;

        if retransmit_unacked_data_in_probe {
            let mut viewer = self.buffer.viewer();

            for interval in self.pending.intervals() {
                if interval.start_inclusive() >= starting_transmission_offset {
                    // Don't write data we've already written to this packet
                    break;
                }

                let interval_end = interval.end_exclusive().min(starting_transmission_offset);

                self.transmissions.transmit_interval(
                    &mut viewer,
                    (interval.start_inclusive()..interval_end).into(),
                    &mut self.state,
                    writer_context,
                    context,
                )?;
            }
        }

        self.check_integrity();

        Ok(())
    }

    #[inline]
    fn check_integrity(&self) {
        if cfg!(debug_assertions) {
            // TODO: assert!(self.lost.is_subset(&self.pending));
            if self.pending.is_empty() {
                assert!(self.lost.is_empty());
                assert!(self.transmissions.is_empty());
                assert_eq!(self.buffer.head(), self.buffer.total_len());
                assert_eq!(
                    self.transmission_offset,
                    self.buffer.total_len(),
                    "transmission offset should equal buffer length when pending is empty"
                );
            }

            if !self.pending.is_empty() {
                assert!(
                    !self.lost.is_empty() || self.transmission_offset < self.buffer.total_len() || !self.transmissions.is_empty(),
                    "pending: {:?}, lost: {:?}, enqueued: {:?}, total_len: {:?}, transmissions: {:?}",
                    self.pending,
                    self.lost,
                    self.transmission_offset,
                    self.buffer.total_len(),
                    self.transmissions.is_empty()
                );
            }

            if let Some(start) = self.pending.min_value() {
                assert_eq!(self.buffer.head(), start);
            }

            if self.flow_controller().is_blocked() {
                use transmission::interest::{Interest, Provider};
                assert_ne!(self.get_transmission_interest(), Interest::NewData);
            }
        }
    }
}

impl<F: OutgoingDataFlowController, W: FrameWriter> transmission::interest::Provider
    for DataSender<F, W>
{
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        if W::WRITES_FIN {
            match self.state {
                State::Finishing(FinState::Lost) => {
                    return query.on_lost_data();
                }
                State::Finishing(FinState::Pending) if !self.flow_controller().is_blocked() => {
                    query.on_new_data()?;
                }
                _ => {}
            }
        };

        if !self.lost.is_empty() {
            query.on_lost_data()?;
        } else if self.transmission_offset < self.buffer.total_len()
            && !self.flow_controller().is_blocked()
        {
            query.on_new_data()?;
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        contexts::testing::{MockWriteContext, OutgoingFrameBuffer},
        transmission::{self, interest::Provider as _},
    };
    use bolero::{check, generator::*};
    use s2n_quic_core::{endpoint, frame, stream::testing as stream, time::clock::testing as time};
    use std::collections::HashSet;

    #[derive(Clone, Copy, Debug, TypeGenerator)]
    enum Event {
        Push(#[generator(1..)] u16),
        Finish,
        Transmit(u16, transmission::Constraint),
        IncFlowControl(u16),
        Ack(usize),
        Loss(usize),
    }

    #[derive(Debug, Default)]
    struct TestFlowController {
        max_offset: VarInt,
        is_blocked: bool,
    }

    impl OutgoingDataFlowController for TestFlowController {
        fn acquire_flow_control_window(&mut self, end_offset: VarInt) -> VarInt {
            if end_offset > self.max_offset {
                self.is_blocked = true;
                self.max_offset
            } else {
                end_offset
            }
        }

        fn is_blocked(&self) -> bool {
            self.is_blocked
        }

        fn clear_blocked(&mut self) {
            self.is_blocked = false;
        }

        fn finish(&mut self) {}
    }

    fn check_model(events: &[Event], id: &VarInt) -> OutgoingFrameBuffer {
        let mut send_data = stream::Data::new(u64::MAX);
        let mut sender: DataSender<_, writer::Stream> =
            DataSender::new(TestFlowController::default(), u32::MAX);
        let mut total_len = 0;
        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut context = MockWriteContext {
            current_time: time::now(),
            frame_buffer: &mut frame_buffer,
            transmission_constraint: transmission::Constraint::None,
            transmission_mode: transmission::Mode::Normal,
            endpoint: endpoint::Type::Server,
        };
        let mut lost = HashSet::new();
        let mut pending = HashSet::new();
        let mut acked = HashSet::new();
        let mut is_finished = false;

        for event in events.iter().copied() {
            match event {
                Event::Push(len) if !is_finished => {
                    let mut chunks = [Bytes::new(), Bytes::new()];
                    let count = send_data
                        .send(len as usize, &mut chunks)
                        .expect("stream should not end early");
                    for chunk in chunks.iter_mut().take(count) {
                        total_len += chunk.len() as u64;
                        sender.push(core::mem::replace(chunk, Bytes::new()));
                    }
                }
                Event::Finish if total_len > 0 => {
                    sender.finish();
                    is_finished = true;
                }
                Event::Transmit(capacity, constraint) => {
                    let interest = sender.get_transmission_interest();

                    let prev_len = context.frame_buffer.len();
                    context.transmission_constraint = constraint;

                    context
                        .frame_buffer
                        .set_max_packet_size(Some(capacity as usize));

                    let _ = sender.on_transmit(*id, &mut context);
                    context.frame_buffer.flush();

                    if interest.is_none() {
                        assert_eq!(
                            context.frame_buffer.len(),
                            prev_len,
                            "frames should only transmit with interest"
                        );
                    }

                    let packets = context
                        .frame_buffer
                        .frames
                        .iter()
                        .skip(prev_len)
                        .map(|frame| frame.packet_nr);
                    pending.extend(packets);
                }
                Event::Ack(index) if !context.frame_buffer.is_empty() => {
                    let index = index % context.frame_buffer.len();
                    let packet = context.frame_buffer.frames[index].packet_nr;
                    sender.on_packet_ack(&packet);
                    pending.remove(&packet);
                    acked.insert(packet);
                }
                Event::Loss(index) if !context.frame_buffer.is_empty() => {
                    let index = index % context.frame_buffer.len();
                    let packet = context.frame_buffer.frames[index].packet_nr;
                    lost.insert(packet);
                    sender.on_packet_loss(&packet);
                }
                Event::IncFlowControl(amount) => {
                    let flow_controller = sender.flow_controller_mut();
                    flow_controller.max_offset = flow_controller
                        .max_offset
                        .saturating_add(VarInt::from_u16(amount));
                }
                _ => {}
            }
        }

        // make sure the stream was finished
        sender.finish();

        for packet in pending {
            sender.on_packet_ack(&packet);
            acked.insert(packet);
        }

        context.frame_buffer.set_max_packet_size(Some(usize::MAX));
        context.transmission_constraint = transmission::Constraint::None;
        sender.flow_controller_mut().clear_blocked();
        sender.flow_controller_mut().max_offset = VarInt::MAX;

        while sender.has_transmission_interest() {
            let prev_len = context.frame_buffer.len();
            let _ = sender.on_transmit(*id, &mut context);
            context.frame_buffer.flush();
            let packets = context
                .frame_buffer
                .frames
                .iter()
                .skip(prev_len)
                .map(|frame| frame.packet_nr);

            let mut did_transmit = false;

            for packet in packets {
                sender.on_packet_ack(&packet);
                acked.insert(packet);
                did_transmit = true;
            }

            assert!(
                did_transmit,
                "transmission_interest was expressed but sender did not transmit: {:#?}",
                sender
            );
        }

        assert!(
            !frame_buffer.is_empty(),
            "the test should transmit at least one frame: {:#?}",
            sender,
        );

        let receiver = stream::Data::new(total_len);
        let mut received_ranges = IntervalSet::new();
        let mut transmitted_fin = false;

        for frame in &mut frame_buffer.frames {
            if acked.contains(&frame.packet_nr) {
                if let frame::Frame::Stream(frame) = frame.as_frame() {
                    let offset = frame.offset.as_u64();
                    let len = frame.data.len() as u64;
                    if len > 0 {
                        receiver.receive_at(offset, &[frame.data.as_less_safe_slice()]);
                        received_ranges.insert(offset..offset + len).unwrap();
                    }
                    transmitted_fin |= frame.is_fin;
                } else {
                    panic!("invalid frame");
                }
            }
        }

        if total_len != 0 {
            assert_eq!(
                received_ranges.interval_len(),
                1,
                "not all data was transmitted",
            );

            assert_eq!(received_ranges.max_value(), Some(total_len - 1));
        } else {
            assert!(received_ranges.is_empty(), "{:#?}", sender);
        }

        assert!(transmitted_fin);

        frame_buffer
    }

    #[test]
    fn model() {
        check!()
            .with_generator((gen(), gen::<Vec<Event>>().with().len(1usize..128)))
            .for_each(|(id, events)| {
                check_model(events, id);
            });
    }
}


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

//! Sends a "flag" frame towards the peer
//!
//! This is intended to be used by frames, like PING and HANDSHAKE_DONE, that don't have any
//! content other than the frame tag itself. At the cost of a single byte per packet, it will passively
//! transmit the flag in any outgoing packets until the peer ACKs the frame. This is to increase
//! the likelihood the peer receives the flag, even in a high-loss environment.

use crate::{
    contexts::{OnTransmitError, WriteContext},
    transmission,
};
use s2n_quic_core::{ack, packet::number::PacketNumber};

#[derive(Debug, Default)]
pub struct Flag<W: Writer> {
    delivery: DeliveryState,
    writer: W,
}

pub trait Writer: Default {
    fn write_frame<W: WriteContext>(&mut self, context: &mut W) -> Option<PacketNumber>;
}

#[derive(Debug, PartialEq)]
enum DeliveryState {
    /// The flag has not been requested
    Idle,

    /// The flag needs to be transmitted
    RequiresTransmission,

    /// The flag was lost and needs to be retransmitted
    RequiresRetransmission,

    /// The flag has been transmitted and is pending acknowledgement.
    ///
    /// Note that in this state, flags are being passively transmitted to ensure
    /// the peer can make progress.
    InFlight {
        /// A stable flag transmission
        ///
        /// In this case, "stable" means the oldest transmission that
        /// hasn't been acked by the peer.
        ///
        /// This packet number is stored to ensure the transmission is either confirmed or declared
        /// lost. Without it, the latest packet number would be a moving target and never
        /// transition to the `Delivered` state
        stable: PacketNumber,

        /// The latest flag transmission
        latest: PacketNumber,
    },

    /// The flag has been delivered
    Delivered,
}

impl Default for DeliveryState {
    fn default() -> Self {
        Self::Idle
    }
}

impl<W: Writer> Flag<W> {
    /// Returns `true` if the flag hasn't been sent
    pub fn is_idle(&self) -> bool {
        matches!(self.delivery, DeliveryState::Idle)
    }

    /// Returns `true` if the flag has been delivered
    pub fn is_delivered(&self) -> bool {
        matches!(self.delivery, DeliveryState::Delivered)
    }

    /// Stars sending the flag to the peer
    pub fn send(&mut self) {
        if self.is_idle() || self.is_delivered() {
            self.delivery = DeliveryState::RequiresTransmission;
        }
    }

    /// Mark the flag as delivered
    pub fn finish(&mut self) {
        self.delivery = DeliveryState::Delivered
    }

    /// This method gets called when a packet delivery got acknowledged
    pub fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A) -> bool {
        if let DeliveryState::InFlight { stable, latest } = &self.delivery {
            if ack_set.contains(*stable) || ack_set.contains(*latest) {
                self.finish();
                return true;
            }
        }

        false
    }

    /// This method gets called when a packet loss is reported
    pub fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) -> bool {
        let mut lost = false;
        if let DeliveryState::InFlight { stable, latest } = &mut self.delivery {
            // If stable is lost, fall back on latest
            if ack_set.contains(*stable) {
                lost = true;
                *stable = *latest;
            }

            // Force retransmission
            if ack_set.contains(*latest) {
                lost = true;
                self.delivery = DeliveryState::RequiresRetransmission;
            }
        }
        lost
    }

    /// Queries the component for any outgoing frames that need to get sent
    pub fn on_transmit<C: WriteContext>(&mut self, context: &mut C) -> Result<(), OnTransmitError> {
        let constraint = context.transmission_constraint();
        match &mut self.delivery {
            DeliveryState::RequiresTransmission if constraint.can_transmit() => {
                if let Some(packet_number) = self.writer.write_frame(context) {
                    self.delivery = DeliveryState::InFlight {
                        stable: packet_number,
                        latest: packet_number,
                    }
                }
            }
            DeliveryState::RequiresRetransmission if constraint.can_retransmit() => {
                if let Some(packet_number) = self.writer.write_frame(context) {
                    self.delivery = DeliveryState::InFlight {
                        stable: packet_number,
                        latest: packet_number,
                    }
                }
            }
            DeliveryState::InFlight { latest, .. } if constraint.can_transmit() => {
                if let Some(packet_number) = self.writer.write_frame(context) {
                    *latest = packet_number;
                }
            }
            _ => {}
        }

        Ok(())
    }
}

impl<W: Writer> transmission::interest::Provider for Flag<W> {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        match &self.delivery {
            DeliveryState::RequiresTransmission => query.on_new_data(),
            DeliveryState::RequiresRetransmission => query.on_lost_data(),
            _ => Ok(()),
        }
    }
}

pub type Ping = Flag<PingWriter>;

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

impl Writer for PingWriter {
    fn write_frame<W: WriteContext>(&mut self, context: &mut W) -> Option<PacketNumber> {
        if context.ack_elicitation().is_ack_eliciting() {
            // we don't need to write a PING frame but we'll store the PacketNumber since it'll be
            // ACKed as if we did
            Some(context.packet_number())
        } else {
            context.write_frame(&s2n_quic_core::frame::Ping)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{contexts::testing::*, transmission::interest::Provider};
    use s2n_quic_core::{endpoint, time::clock::testing as time};

    #[test]
    fn ping_test() {
        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut context = MockWriteContext::new(
            time::now(),
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::Normal,
            endpoint::Type::Server,
        );

        let mut pinger = Ping::default();

        assert!(pinger.is_idle());
        assert!(
            !pinger.has_transmission_interest(),
            "status should not express interest in default state"
        );

        pinger.on_transmit(&mut context).unwrap();

        assert!(
            context.frame_buffer.is_empty(),
            "status should not transmit in default state"
        );

        pinger.send();
        assert!(!pinger.is_idle());

        assert_eq!(
            pinger.get_transmission_interest(),
            transmission::Interest::NewData,
            "status should express interest in deliver after calling send"
        );

        pinger.send();
        assert_eq!(
            pinger.get_transmission_interest(),
            transmission::Interest::NewData,
            "status should accept duplicate calls to send"
        );

        context.transmission_constraint = transmission::Constraint::CongestionLimited;
        pinger.on_transmit(&mut context).unwrap();
        assert!(!pinger.is_idle());

        assert_eq!(
            pinger.delivery,
            DeliveryState::RequiresTransmission,
            "status should not transmit when congestion limited"
        );
        assert!(
            context.frame_buffer.is_empty(),
            "status should not transmit in congestion limited state"
        );

        context.transmission_constraint = transmission::Constraint::None;

        pinger.on_transmit(&mut context).unwrap();
        assert!(!pinger.is_idle());

        let stable_packet_number = context
            .frame_buffer
            .pop_front()
            .expect("status should write PING frames")
            .packet_nr;
        context.frame_buffer.clear();

        assert_eq!(
            pinger.delivery,
            DeliveryState::InFlight {
                stable: stable_packet_number,
                latest: stable_packet_number
            }
        );

        context.transmission_constraint = transmission::Constraint::RetransmissionOnly;

        pinger.on_transmit(&mut context).unwrap();
        assert!(!pinger.is_idle());

        assert!(
            context.frame_buffer.is_empty(),
            "status should not passively write frames when transmission constrained"
        );

        context.transmission_constraint = transmission::Constraint::None;

        pinger.on_transmit(&mut context).unwrap();
        assert!(!pinger.is_idle());

        let latest_packet_number = context
            .frame_buffer
            .pop_front()
            .expect("status should passively write PING frames")
            .packet_nr;
        context.frame_buffer.clear();

        assert_eq!(
            pinger.delivery,
            DeliveryState::InFlight {
                stable: stable_packet_number,
                latest: latest_packet_number,
            }
        );

        pinger.on_packet_loss(&stable_packet_number);

        assert_eq!(
            pinger.delivery,
            DeliveryState::InFlight {
                stable: latest_packet_number,
                latest: latest_packet_number,
            },
            "status should transition to latest on stable packet loss"
        );

        pinger.on_packet_loss(&latest_packet_number);

        assert_eq!(
            pinger.get_transmission_interest(),
            transmission::Interest::LostData,
            "transmission should be active on latest packet loss"
        );
        assert_eq!(
            pinger.delivery,
            DeliveryState::RequiresRetransmission,
            "status should force retransmission on loss"
        );

        context.transmission_constraint = transmission::Constraint::CongestionLimited;
        pinger.on_transmit(&mut context).unwrap();
        assert!(!pinger.is_idle());

        assert_eq!(
            pinger.delivery,
            DeliveryState::RequiresRetransmission,
            "status should not transmit when congestion limited"
        );

        context.transmission_constraint = transmission::Constraint::RetransmissionOnly;

        pinger.on_transmit(&mut context).unwrap();

        let latest_packet_number = context
            .frame_buffer
            .pop_front()
            .expect("status should passively write PING frames")
            .packet_nr;
        context.frame_buffer.clear();

        pinger.on_packet_ack(&latest_packet_number);

        assert_eq!(pinger.delivery, DeliveryState::Delivered);

        assert!(
            !pinger.has_transmission_interest(),
            "status should not express interest after complete",
        );

        pinger.on_transmit(&mut context).unwrap();
        assert!(
            context.frame_buffer.is_empty(),
            "status should not transmit after complete"
        );
    }
}


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

//! Synchronizes a strictly increasing value of type `T` towards the remote peer.

use crate::{
    contexts::{OnTransmitError, WriteContext},
    sync::{DeliveryState, InFlightDelivery, InflightPacketInfo, ValueToFrameWriter},
    transmission,
};
use s2n_quic_core::{ack, stream::StreamId};

/// Synchronizes a strictly increasing value of type `T` towards the remote peer.
///
/// IncrementalValueSync will only send an update if it is significant enough (above a
/// certain `threshold`), or if the last update had been lost.
/// `S` is of type [`ValueToFrameWriter`] and used to serialize the latest value
/// into an outgoing frame.
#[derive(Debug)]
pub struct IncrementalValueSync<T, S> {
    latest_value: T,
    value_ackd_up_to: T,
    threshold: T,
    delivery: DeliveryState<T>,
    writer: S,
}

impl<
        T: Copy + Clone + core::fmt::Debug + Eq + PartialEq + PartialOrd + core::ops::Sub<Output = T>,
        S: ValueToFrameWriter<T>,
    > IncrementalValueSync<T, S>
{
    pub fn new(latest_value: T, value_ackd_up_to: T, threshold: T) -> Self {
        debug_assert!(
            latest_value >= value_ackd_up_to,
            "value to sync must be bigger or equal than the last acknowledged value"
        );
        let mut sync = IncrementalValueSync {
            latest_value,
            value_ackd_up_to,
            delivery: DeliveryState::NotRequested,
            threshold,
            writer: S::default(),
        };

        sync.request_delivery_if_necessary();

        sync
    }

    /// Returns the latest value that needs to get synchronized
    pub fn latest_value(&self) -> T {
        self.latest_value
    }

    /// Returns `true` if the synchronization has been cancelled
    pub fn is_cancelled(&self) -> bool {
        self.delivery.is_cancelled()
    }

    /// Returns `true` if the delivery is current in progress.
    /// A packet has been sent, but no acknowledgement has been retrieved so far.
    pub fn is_inflight(&self) -> bool {
        self.delivery.is_inflight()
    }

    /// Sets the new value that needs to get synchronized to the peer.
    /// Returns true if new value requires `on_transmit` to be called as soon as
    /// possible.
    pub fn update_latest_value(&mut self, value: T) {
        debug_assert!(value >= self.latest_value);
        self.latest_value = value;
        self.request_delivery_if_necessary();
    }

    /// Stop to synchronize the value to the peer
    pub fn stop_sync(&mut self) {
        self.delivery.cancel();
    }

    /// If the latest value is high enough to require sending an update, this
    /// sets the delivery state to `Requested`.
    fn request_delivery_if_necessary(&mut self) {
        if self.should_send_update() {
            self.delivery = DeliveryState::Requested(self.latest_value);
        }
    }

    /// Returns whether an update about the state of the stored value should be
    /// sent to the peer.
    ///
    /// Updates are only sent if the value that needs to get synchronized to the
    /// peer exceeds the value acknowledged by the peer by the configured `threshold`.
    fn should_send_update(&self) -> bool {
        if self.delivery.is_cancelled() {
            return false;
        }

        // Check if the window has already been fully transmitted and ackd
        if self.latest_value != self.value_ackd_up_to {
            // Check if a frame is already in transmission
            if let DeliveryState::InFlight(in_flight) = self.delivery {
                // Check if the new update is significant enough to supersede the old one
                // or whether we think the old update had been lost.
                if self.latest_value - in_flight.value >= self.threshold {
                    // TODO: Using the 10% threshold means we will overwrite the
                    // tracking information about older in-flight packets all the
                    // time, and only adjust `value_ackd_up_to` if the latest
                    // packet gets acknowledged. That might lead us in some cases
                    // to send a re-transmit even when not strictly necessary, since
                    // we don't observe a previous transmit being acknowledged.
                    // Either tracking more pending transmissions, or increasing
                    // the threshold could improve on that.
                    return true;
                }
            } else {
                // Check if the new update is significant enough
                if self.latest_value - self.value_ackd_up_to >= self.threshold {
                    return true;
                }
            }
        }

        false
    }

    /// This method gets called when a packet delivery got acknowledged
    pub fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A) {
        // If the frame gets acknowledged, remove the in_flight information
        if let DeliveryState::InFlight(in_flight) = self.delivery {
            if ack_set.contains(in_flight.packet.packet_nr) {
                self.value_ackd_up_to = in_flight.value;
                self.delivery = DeliveryState::NotRequested;
                // There is no need to call request_delivery_if_necessary() here:
                // If the value would have been updated significantly enough the
                // update would already have been superseded.
            }
        }
    }

    /// This method gets called when a packet loss is reported
    pub fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) {
        // If the frame was lost, remove the in_flight information.
        // This will trigger resending it.
        if let DeliveryState::InFlight(in_flight) = self.delivery {
            if ack_set.contains(in_flight.packet.packet_nr) {
                self.delivery = DeliveryState::Lost(self.latest_value);
            }
        }
    }

    /// Queries the component for any outgoing frames that need to get sent
    #[inline]
    pub fn on_transmit<W: WriteContext>(
        &mut self,
        stream_id: StreamId,
        context: &mut W,
    ) -> Result<(), OnTransmitError> {
        if self
            .delivery
            .try_transmit(context.transmission_constraint())
            .is_some()
        {
            // We grab the latest value here, even if an older one had been
            // requested for delivery. That makes sure we always transmit the
            // highest available value to the peer.
            let value = self.latest_value;

            let packet_nr = self
                .writer
                .write_value_as_frame(value, stream_id, context)
                .ok_or(OnTransmitError::CouldNotWriteFrame)?;

            // Overwrite the information about the in_flight transmission of the
            // latest value.
            self.delivery = DeliveryState::InFlight(InFlightDelivery {
                value,
                packet: InflightPacketInfo {
                    packet_nr,
                    timestamp: context.current_time(),
                },
            });
        }

        Ok(())
    }
}

impl<T, S> transmission::interest::Provider for IncrementalValueSync<T, S> {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        self.delivery.transmission_interest(query)
    }
}


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

//! Tools for synchronizing data between peers

use crate::{contexts::WriteContext, transmission};
use s2n_quic_core::{packet::number::PacketNumber, stream::StreamId, time::Timestamp};

pub mod data_sender;
pub mod flag;
mod incremental_value_sync;
mod once_sync;
mod periodic_sync;

pub use data_sender::DataSender;
pub use flag::Flag;
pub use incremental_value_sync::IncrementalValueSync;
pub use once_sync::OnceSync;
pub use periodic_sync::{PeriodicSync, DEFAULT_SYNC_PERIOD};

/// Carries information about the packet in which a frame is transmitted
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct InflightPacketInfo {
    /// The number of the packet that was used to send the message.
    packet_nr: PacketNumber,
    /// The timestamp when the message had been sent.
    timestamp: Timestamp,
}

/// A value delivery which is currently in progress
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct InFlightDelivery<T> {
    /// The value which is sent
    value: T,
    /// The packet in which the value was sent
    packet: InflightPacketInfo,
}

/// Tracks the delivery of a particular piece of information to the peer.
#[derive(PartialEq, Eq, Debug)]
pub enum DeliveryState<T> {
    /// The delivery of the information has not yet been requested
    NotRequested,
    /// The delivery had been requested, but not yet started
    Requested(T),
    /// The original delivery was lost and needs to be retransmitted
    Lost(T),
    /// The delivery of the information has been requested and is in progress
    InFlight(InFlightDelivery<T>),
    /// The delivery of the information has succeeded
    Delivered(T),
    /// The delivery was cancelled. If a delivery of a value was previously
    /// requested it will be stored in the `Option`. Otherwise `None` will be
    /// stored.
    Cancelled(Option<T>),
}

impl<T> DeliveryState<T> {
    /// Moves the DeliverState into the `Cancelled` state.
    /// If a delivery was previously requested, the `Option` stored in
    /// `DeliverState::Cancelled` will contain the last value which was scheduled
    /// for delivery.
    #[inline]
    pub fn cancel(&mut self) {
        let old_state = core::mem::replace(self, DeliveryState::Cancelled(None));
        *self = match old_state {
            DeliveryState::NotRequested => DeliveryState::Cancelled(None),
            DeliveryState::Requested(value)
            | DeliveryState::Lost(value)
            | DeliveryState::Delivered(value)
            | DeliveryState::InFlight(InFlightDelivery { value, .. }) => {
                DeliveryState::Cancelled(Some(value))
            }
            state @ DeliveryState::Cancelled(_) => state,
        };
    }

    /// Returns `true` if the delivery of the value had been cancelled
    #[inline]
    pub fn is_cancelled(&self) -> bool {
        matches!(self, Self::Cancelled(_))
    }

    /// Returns `true` if the delivery is current in progress.
    /// A packet has been sent, but no acknowledgement has been retrieved so far.
    #[inline]
    pub fn is_inflight(&self) -> bool {
        matches!(self, Self::InFlight(_))
    }

    /// Tries to transmit the delivery with the given transmission constraint
    #[inline]
    pub fn try_transmit(&self, constraint: transmission::Constraint) -> Option<&T> {
        match self {
            DeliveryState::Requested(value) if constraint.can_transmit() => Some(value),
            DeliveryState::Lost(value) if constraint.can_retransmit() => Some(value),
            _ => None,
        }
    }
}

impl<T> transmission::interest::Provider for DeliveryState<T> {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        match self {
            Self::Requested(_) => query.on_new_data(),
            Self::Lost(_) => query.on_lost_data(),
            _ => Ok(()),
        }
    }
}

/// Writes values of type `T` into frames.
pub trait ValueToFrameWriter<T>: Default {
    /// Creates a QUIC frame out of the given value, and writes it using the
    /// provided [`WriteContext`].
    /// The method returns the `PacketNumber` of the packet containing the value
    /// if the write was successful, and `None` otherwise.
    fn write_value_as_frame<W: WriteContext>(
        &self,
        value: T,
        stream_id: StreamId,
        context: &mut W,
    ) -> Option<PacketNumber>;
}


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

//! Synchronizes a value of type `T` exactly once towards the remote peer.

use crate::{
    contexts::{OnTransmitError, WriteContext},
    sync::{DeliveryState, InFlightDelivery, InflightPacketInfo, ValueToFrameWriter},
    transmission,
};
use core::task::Poll;
use s2n_quic_core::{ack, stream::StreamId};

/// Synchronizes a value of type `T` exactly once towards the remote peer.
///
/// Retransmissions of the value will be performed if it got lost. However it is
/// not possible to replace the value with a newer value.
///
/// `S` is of type [`ValueToFrameWriter`] and used to serialize the value
/// into an outgoing frame.
#[derive(Debug)]
pub struct OnceSync<T, S> {
    delivery: DeliveryState<T>,
    writer: S,
}

impl<T, S: Default> Default for OnceSync<T, S> {
    fn default() -> Self {
        Self {
            delivery: DeliveryState::NotRequested,
            writer: S::default(),
        }
    }
}

impl<T: Copy + Clone + Eq + PartialEq, S: ValueToFrameWriter<T>> OnceSync<T, S> {
    /// Creates a new OnceSync
    pub fn new() -> Self {
        Self::default()
    }

    /// Returns `true` if the delivery is current in progress.
    /// A packet has been sent, but no acknowledgement has been retrieved so far.
    #[inline]
    pub fn is_inflight(&self) -> bool {
        self.delivery.is_inflight()
    }

    /// Returns `true` if the synchronization has been cancelled
    #[inline]
    pub fn is_cancelled(&self) -> bool {
        self.delivery.is_cancelled()
    }

    /// Requested delivery of the given value.
    pub fn request_delivery(&mut self, value: T) {
        if let DeliveryState::NotRequested = self.delivery {
            self.delivery = DeliveryState::Requested(value);
        }
    }

    /// Forces transmission of the given value.
    pub fn force_delivery(&mut self, value: T) {
        self.delivery = DeliveryState::Requested(value);
    }

    /// Stop to synchronize the value to the peer
    #[inline]
    pub fn stop_sync(&mut self) {
        self.delivery.cancel();
    }

    /// This method gets called when a packet delivery got acknowledged
    pub fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A) -> Poll<()> {
        // If the packet containing the frame gets acknowledged, mark the delivery as
        // succeeded.
        if let DeliveryState::InFlight(in_flight) = self.delivery {
            if ack_set.contains(in_flight.packet.packet_nr) {
                self.delivery = DeliveryState::Delivered(in_flight.value);
                return Poll::Ready(());
            }
        }

        Poll::Pending
    }

    /// This method gets called when a packet loss is reported
    pub fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) {
        // If the packet containing the frame was lost, remove the in_flight information.
        // This will trigger resending it.
        if let DeliveryState::InFlight(in_flight) = self.delivery {
            if ack_set.contains(in_flight.packet.packet_nr) {
                self.delivery = DeliveryState::Lost(in_flight.value);
            }
        }
    }

    /// Queries the component for any outgoing frames that need to get sent
    #[inline]
    pub fn on_transmit<W: WriteContext>(
        &mut self,
        stream_id: StreamId,
        context: &mut W,
    ) -> Result<(), OnTransmitError> {
        if let Some(value) = self
            .delivery
            .try_transmit(context.transmission_constraint())
            .cloned()
        {
            let packet_nr = self
                .writer
                .write_value_as_frame(value, stream_id, context)
                .ok_or(OnTransmitError::CouldNotWriteFrame)?;

            // Overwrite the information about the pending delivery
            self.delivery = DeliveryState::InFlight(InFlightDelivery {
                packet: InflightPacketInfo {
                    packet_nr,
                    timestamp: context.current_time(),
                },
                value,
            });
        }

        Ok(())
    }
}

impl<T, S> transmission::interest::Provider for OnceSync<T, S> {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        self.delivery.transmission_interest(query)
    }
}


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

//! Synchronizes a value of type `T` periodically towards the remote peer.

use crate::{
    contexts::{OnTransmitError, WriteContext},
    sync::{DeliveryState, InFlightDelivery, InflightPacketInfo, ValueToFrameWriter},
    transmission,
};
use core::time::Duration;
use s2n_quic_core::{
    ack,
    counter::{self, Counter},
    stream::StreamId,
    time::{timer, Timer, Timestamp},
};

/// The default period for synchronizing the value. This value is only used prior to a more
/// precise value calculated based on idle timeout and current RTT estimates and provided
/// in the `update_sync_period` method. The value is based on 3 * DEFAULT_INITIAL_RTT
pub const DEFAULT_SYNC_PERIOD: Duration = Duration::from_millis(999);

/// The initial backoff value for transmission
const INITIAL_BACKOFF: u16 = 1;

/// Synchronizes a monotonically increasing value of type `T` periodically towards the remote peer.
///
/// Retransmissions of the value will be performed if it got lost.
///
/// `S` is of type [`ValueToFrameWriter`] and used to serialize the value
/// into an outgoing frame.
#[derive(Debug)]
pub struct PeriodicSync<T, S> {
    latest_value: T,
    sync_period: Duration,
    delivery_timer: Timer,
    delivery: DeliveryState<T>,
    writer: S,
    delivered: bool,
    transmission_backoff: Counter<u16, counter::Saturating>,
}

impl<T: Copy + Clone + Default + Eq + PartialEq + PartialOrd, S: ValueToFrameWriter<T>>
    PeriodicSync<T, S>
{
    /// Creates a new PeriodicSync. The value will transmitted when `request_delivery` is called
    /// and every subsequent `sync_period` until `stop_sync` is called.
    pub fn new() -> Self {
        Self {
            latest_value: T::default(),
            sync_period: DEFAULT_SYNC_PERIOD,
            delivery_timer: Timer::default(),
            delivery: DeliveryState::NotRequested,
            writer: S::default(),
            delivered: false,
            transmission_backoff: Counter::new(INITIAL_BACKOFF),
        }
    }

    /// Requested delivery of the given value. If delivery has already been requested, the
    /// original value will be overwritten. The new value must be greater than or equal
    /// to the original value.
    pub fn request_delivery(&mut self, value: T) {
        debug_assert!(value >= self.latest_value);

        self.latest_value = value;

        if matches!(
            self.delivery,
            DeliveryState::NotRequested | DeliveryState::Cancelled(_)
        ) {
            self.delivery = DeliveryState::Requested(value);
        }
    }

    /// Skip delivery for this sync period. A delivery will be scheduled for the next sync period.
    pub fn skip_delivery(&mut self, now: Timestamp) {
        match self.delivery {
            DeliveryState::Requested(_) | DeliveryState::Lost(_) => {
                self.delivery = DeliveryState::NotRequested;
                self.transmission_backoff *= 2u16;
                self.update_timer(now);
            }
            DeliveryState::Delivered(_) => {
                self.transmission_backoff *= 2u16;
                self.update_timer(now);
            }
            _ => {}
        }
    }

    /// Called when the connection timer expires
    pub fn on_timeout(&mut self, now: Timestamp) {
        if self.delivery_timer.poll_expiration(now).is_ready() {
            self.delivery = DeliveryState::Requested(self.latest_value);
        }
    }

    /// Stop synchronizing the value to the peer
    pub fn stop_sync(&mut self) {
        self.delivery_timer.cancel();
        self.delivery.cancel();
        self.delivered = false;
        self.transmission_backoff.set(INITIAL_BACKOFF);
    }

    /// This method gets called when a packet delivery got acknowledged
    pub fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A) {
        // If the packet containing the frame gets acknowledged, schedule a delivery for the
        // next delivery period
        if let DeliveryState::InFlight(in_flight) = self.delivery {
            if ack_set.contains(in_flight.packet.packet_nr) {
                self.delivered = true;
                self.update_timer(in_flight.packet.timestamp);
                self.delivery = DeliveryState::Delivered(in_flight.value);
            }
        }
    }

    /// Sets the sync period. The given `sync_period` will be used the next time
    /// the delivery timer is armed; the existing timer will be unaffected.
    pub fn update_sync_period(&mut self, sync_period: Duration) {
        self.sync_period = sync_period;
    }

    /// This method gets called when a packet loss is reported
    pub fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) {
        // If the packet containing the frame was lost, remove the in_flight information.
        // This will trigger resending it.
        if let DeliveryState::InFlight(in_flight) = self.delivery {
            if ack_set.contains(in_flight.packet.packet_nr) {
                self.delivery = DeliveryState::Lost(in_flight.value);
            }
        }
    }

    /// Queries the component for any outgoing frames that need to get sent
    #[inline]
    pub fn on_transmit<W: WriteContext>(
        &mut self,
        stream_id: StreamId,
        context: &mut W,
    ) -> Result<(), OnTransmitError> {
        if self
            .delivery
            .try_transmit(context.transmission_constraint())
            .is_some()
        {
            let value = self.latest_value;

            let packet_nr = self
                .writer
                .write_value_as_frame(value, stream_id, context)
                .ok_or(OnTransmitError::CouldNotWriteFrame)?;

            // Overwrite the information about the pending delivery
            self.delivery = DeliveryState::InFlight(InFlightDelivery {
                packet: InflightPacketInfo {
                    packet_nr,
                    timestamp: context.current_time(),
                },
                value,
            });

            self.transmission_backoff *= 2u16;
        }

        Ok(())
    }

    /// Returns whether the value has been delivered at least once since delivery was first
    /// requested, or requested again after `stop_sync` was called.
    #[inline]
    pub fn has_delivered(&self) -> bool {
        self.delivered
    }

    fn update_timer(&mut self, now: Timestamp) {
        self.delivery_timer.set(now + self.sync_period());
    }

    fn sync_period(&self) -> Duration {
        // exponentially back off on the sync period so we don't overwhelm the receiver
        self.sync_period * *self.transmission_backoff as _
    }
}

impl<T, S> timer::Provider for PeriodicSync<T, S> {
    #[inline]
    fn timers<Q: timer::Query>(&self, query: &mut Q) -> timer::Result {
        self.delivery_timer.timers(query)?;
        Ok(())
    }
}

impl<T, S> transmission::interest::Provider for PeriodicSync<T, S> {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        self.delivery.transmission_interest(query)
    }
}

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

    #[test]
    fn default_period() {
        assert_eq!(DEFAULT_INITIAL_RTT * 3, DEFAULT_SYNC_PERIOD);
    }
}


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

use alloc::collections::VecDeque;
use bytes::{Buf, Bytes};
use core::{
    convert::{TryFrom, TryInto},
    fmt,
};
use s2n_codec::{Encoder, EncoderValue};
use s2n_quic_core::{frame::FitError, interval_set::Interval, varint::VarInt};

#[derive(Debug, Default)]
pub struct Buffer {
    chunks: VecDeque<Chunk>,
    head: VarInt,
    pending_len: VarInt,
}

impl Buffer {
    /// Pushes a chunk of data into to the buffer for transmission
    pub fn push(&mut self, data: Bytes) -> Interval<VarInt> {
        debug_assert!(
            self.capacity().as_u64() >= data.len() as u64,
            "capacity should be checked before pushing"
        );
        let start = self.total_len();
        let len = VarInt::try_from(data.len()).expect("cannot send more than VarInt::MAX");
        self.pending_len += len;
        self.chunks.push_back(Chunk { data });
        self.check_integrity();

        // sub 1 so we don't overflow
        let end = start + (len - 1);

        (start..=end).into()
    }

    /// Returns the maximum capacity the buffer could ever hold
    #[inline]
    pub fn capacity(&self) -> VarInt {
        VarInt::MAX - self.total_len()
    }

    /// Clears and resets the buffer
    pub fn clear(&mut self) {
        self.chunks.clear();
        self.head = VarInt::from_u8(0);
        self.pending_len = VarInt::from_u8(0);
        self.check_integrity();
    }

    /// Returns the total number of bytes the buffer has and is currently holding
    #[inline]
    pub fn total_len(&self) -> VarInt {
        self.head + self.pending_len
    }

    /// Returns the head or offset at which the first chunk in the buffer starts
    #[inline]
    pub fn head(&self) -> VarInt {
        self.head
    }

    /// Returns the number of bytes enqueue for transmission/retransmission
    #[inline]
    pub fn enqueued_len(&self) -> VarInt {
        self.pending_len
    }

    /// Returns true if the buffer is currently empty
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.pending_len == VarInt::from_u8(0)
    }

    /// Sets the current offset of the buffer.
    ///
    /// This should only be used in testing.
    #[cfg(test)]
    pub fn set_offset(&mut self, head: VarInt) {
        self.head = head;
    }

    /// Releases all of the chunks up to the provided offset in the buffer
    ///
    /// This method should be called after a chunk of data has been transmitted
    /// and acknowledged, as there is no longer a need for it to be buffered.
    pub fn release(&mut self, up_to: VarInt) {
        // we've already released up to this offset
        if up_to <= self.head {
            return;
        }

        debug_assert!(
            self.total_len() >= up_to,
            "cannot release more than the total len"
        );

        while let Some(mut chunk) = self.chunks.pop_front() {
            let len = VarInt::try_from(chunk.len()).unwrap();
            let start = self.head;
            let end = start + len;

            // if the end of this chunk is less than the up_to, drop it entirely
            if end <= up_to {
                self.pending_len -= len;
                self.head = end;
                continue;
            }

            // only part of the chunk has been released
            self.head = up_to;

            // compute the consumed amount for the chunk
            let consumed = self.head - start;
            self.pending_len -= consumed;
            chunk.data.advance(consumed.try_into().unwrap());

            // push the chunk back for later
            self.chunks.push_front(chunk);

            break;
        }

        self.check_integrity();
    }

    /// Releases all of the currently enqueued chunks
    pub fn release_all(&mut self) {
        self.chunks.clear();
        self.head = self.total_len();
        self.pending_len = VarInt::from_u8(0);

        self.check_integrity();
    }

    /// Returns a Viewer for the buffer
    #[inline]
    pub fn viewer(&self) -> Viewer {
        Viewer {
            buffer: self,
            offset: *self.head,
            chunk_index: 0,
        }
    }

    #[inline]
    fn check_integrity(&self) {
        if cfg!(debug_assertions) {
            let actual: VarInt = self
                .chunks
                .iter()
                .map(|chunk| chunk.len())
                .sum::<usize>()
                .try_into()
                .unwrap();
            assert_eq!(
                actual, self.pending_len,
                "actual buffer lengths should equal `pending_len`"
            );
        }
    }
}

#[derive(Default)]
struct Chunk {
    data: Bytes,
}

impl fmt::Debug for Chunk {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Chunk")
            .field("len", &self.data.len())
            .finish()
    }
}

impl core::ops::Deref for Chunk {
    type Target = Bytes;

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.data
    }
}

#[derive(Clone, Copy, Debug)]
pub struct Viewer<'a> {
    buffer: &'a Buffer,
    offset: u64,
    chunk_index: usize,
}

impl<'a> Viewer<'a> {
    /// Returns the next view in the buffer for a given range
    #[inline]
    pub fn next_view(&mut self, range: Interval<VarInt>, has_fin: bool) -> View<'a> {
        View::new(
            self.buffer,
            range,
            has_fin,
            &mut self.offset,
            &mut self.chunk_index,
        )
    }
}

#[derive(Clone, Copy, Debug)]
pub struct View<'a> {
    buffer: &'a Buffer,
    chunk_index: usize,
    offset: usize,
    len: usize,
    is_fin: bool,
}

impl<'a> View<'a> {
    #[inline]
    fn new(
        buffer: &'a Buffer,
        range: Interval<VarInt>,
        has_fin: bool,
        stream_offset: &mut u64,
        chunk_index: &mut usize,
    ) -> Self {
        debug_assert!(
            buffer.head <= range.start_inclusive(),
            "range ({:?}) is referring to a chunk that has already been released: {:?}",
            range,
            buffer.head..buffer.total_len()
        );

        debug_assert!(
            *stream_offset <= range.start_inclusive().as_u64(),
            "viewer is trying to go backwards from offset {:?} to {:?}",
            stream_offset,
            range.start_inclusive()
        );
        debug_assert!(range.end_exclusive() <= buffer.total_len());

        let mut offset = 0;

        // find the chunk and offset where the range starts
        for chunk in buffer.chunks.iter().skip(*chunk_index) {
            let len = chunk.len() as u64;
            let start = *stream_offset;
            let end = start + len;

            if (start..end).contains(&range.start_inclusive()) {
                offset = (range.start_inclusive().as_u64() - start) as usize;
                break;
            }

            *stream_offset += len;
            *chunk_index += 1;
        }

        debug_assert!(*chunk_index < buffer.chunks.len());

        Self {
            buffer,
            chunk_index: *chunk_index,
            offset,
            len: range.len(),
            is_fin: has_fin && range.end_inclusive() == (buffer.total_len() - 1),
        }
    }

    /// Trims off an `amount` number of bytes from the end of the view
    ///
    /// If `amount` exceeds the view `len`, Err will be returned
    #[inline]
    pub fn trim_off(&mut self, amount: usize) -> Result<(), FitError> {
        self.len = self.len.checked_sub(amount).ok_or(FitError)?;

        // trimming data off the end invalidates this
        self.is_fin &= amount == 0;

        Ok(())
    }

    /// Returns the number of bytes in the current view
    #[inline]
    pub fn len(&self) -> VarInt {
        VarInt::try_from(self.len).expect("len should always fit in a VarInt")
    }

    /// Returns `true` if the view includes the last byte in the stream
    #[inline]
    pub fn is_fin(&self) -> bool {
        self.is_fin
    }

    #[inline]
    pub fn iter<'iter, S: Slice<'iter>>(&'iter self) -> ViewIter<'iter, S> {
        ViewIter {
            view: *self,
            slice: Default::default(),
        }
    }
}

pub struct ViewIter<'a, S: Slice<'a>> {
    view: View<'a>,
    slice: core::marker::PhantomData<S>,
}

impl<'a, S: Slice<'a>> Iterator for ViewIter<'a, S> {
    type Item = S;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        let view = &mut self.view;

        if view.len == 0 {
            return None;
        }

        let chunk = &view.buffer.chunks[view.chunk_index];

        let start = view.offset;
        // reset the offset to the beginning of the next chunk
        view.offset = 0;
        // compute the chunk len with the given offset
        let len = chunk.len() - start;
        // make sure we don't exceed that max len
        let len = view.len.min(len);

        // compute the end of the chunk slice
        let end = start + len;
        // decrement the remaining len
        view.len -= len;
        // move to the next chunk
        view.chunk_index += 1;

        debug_assert_eq!(chunk[start..end].len(), len);
        Some(S::from_chunk(chunk, start, end))
    }
}

/// Converts a Bytes into the implemented type
pub trait Slice<'a> {
    fn from_chunk(chunk: &'a Bytes, start: usize, end: usize) -> Self;
}

impl<'a> Slice<'a> for &'a [u8] {
    #[inline]
    fn from_chunk(chunk: &'a Bytes, start: usize, end: usize) -> Self {
        unsafe {
            debug_assert!(chunk.len() >= end);
            chunk.get_unchecked(start..end)
        }
    }
}

impl<'a> Slice<'a> for Bytes {
    #[inline]
    fn from_chunk(chunk: &'a Bytes, start: usize, end: usize) -> Self {
        chunk.slice(start..end)
    }
}

impl<'a> EncoderValue for &mut View<'a> {
    #[inline]
    fn encode<E: Encoder>(&self, encoder: &mut E) {
        // Specialize on writing byte chunks directly instead of copying the slices
        if E::SPECIALIZES_BYTES {
            for chunk in self.iter::<Bytes>() {
                encoder.write_bytes(chunk);
            }
            return;
        }

        encoder.write_sized(self.len, |slice| {
            let mut offset = 0;
            for chunk in self.iter::<&[u8]>() {
                let len = chunk.len();
                let end = offset + len;
                unsafe {
                    // Safety: we've already checked that the slice has enough
                    // capacity with `write_sized`
                    debug_assert!(slice.len() >= end);

                    // These copies are critical to performance so use use copy_nonoverlapping
                    // directly, rather than rely on compiler optimizations to ensure we
                    // don't pay any additional costs
                    core::ptr::copy_nonoverlapping(
                        chunk.as_ptr(),
                        slice.get_unchecked_mut(offset),
                        len,
                    );
                }
                offset += len;
            }
        });
    }

    #[inline]
    fn encoding_size_for_encoder<E: Encoder>(&self, _encoder: &E) -> usize {
        self.len
    }
}

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

    fn almost_full_buffer() -> Buffer {
        Buffer {
            head: VarInt::MAX - 1,
            ..Default::default()
        }
    }

    #[test]
    fn partial_release_test() {
        let mut buffer = Buffer::default();

        buffer.push(Bytes::from_static(&[0, 1, 2]));
        assert_eq!(buffer.total_len(), VarInt::from_u8(3));
        assert_eq!(buffer.enqueued_len(), VarInt::from_u8(3));

        // trim off the first byte
        buffer.release(VarInt::from_u8(1));
        assert_eq!(buffer.total_len(), VarInt::from_u8(3));
        assert_eq!(buffer.enqueued_len(), VarInt::from_u8(2));
        assert_eq!(buffer.chunks.len(), 1);
        assert_eq!(buffer.chunks[0][..], [1, 2]);

        // duplicate releases should be ok
        buffer.release(VarInt::from_u8(1));
        assert_eq!(buffer.total_len(), VarInt::from_u8(3));
        assert_eq!(buffer.enqueued_len(), VarInt::from_u8(2));
        assert_eq!(buffer.chunks.len(), 1);
        assert_eq!(buffer.chunks[0][..], [1, 2]);
    }

    #[test]
    fn full_release_test() {
        let mut buffer = Buffer::default();

        buffer.push(Bytes::from_static(&[0, 1, 2]));
        assert_eq!(buffer.total_len(), VarInt::from_u8(3));
        assert_eq!(buffer.enqueued_len(), VarInt::from_u8(3));

        // trim off all bytes
        buffer.release(VarInt::from_u8(3));
        assert_eq!(buffer.total_len(), VarInt::from_u8(3));
        assert_eq!(buffer.enqueued_len(), VarInt::from_u8(0));
        assert!(buffer.chunks.is_empty());

        // duplicate releases should be ok
        buffer.release(VarInt::from_u8(3));
        assert_eq!(buffer.total_len(), VarInt::from_u8(3));
        assert_eq!(buffer.enqueued_len(), VarInt::from_u8(0));
        assert!(buffer.chunks.is_empty());
    }

    #[test]
    fn varint_max_test() {
        let mut buffer = almost_full_buffer();

        buffer.push(Bytes::from_static(&[0]));

        buffer.release(VarInt::MAX);

        assert_eq!(buffer.total_len(), VarInt::MAX);
        assert_eq!(buffer.enqueued_len(), VarInt::from_u8(0));
        assert!(buffer.chunks.is_empty());
    }

    #[test]
    #[should_panic]
    fn varint_overflow_test() {
        let mut buffer = almost_full_buffer();

        // pushing 2 bytes will exceed the capacity and panic
        buffer.push(Bytes::from_static(&[0, 1]));
    }

    fn check_view(buffer: &Buffer, interval: Interval<u64>, expected: &[u8]) {
        let interval = (VarInt::new(interval.start_inclusive()).unwrap()
            ..=VarInt::new(interval.end_inclusive()).unwrap())
            .into();
        let actual: Vec<u8> = View::new(buffer, interval, false, &mut buffer.head.as_u64(), &mut 0)
            .iter::<&[u8]>()
            .flatten()
            .copied()
            .collect();
        assert_eq!(actual, expected);
    }

    fn check_viewer(viewer: &mut Viewer, interval: Interval<u64>, expected: &[u8]) {
        let interval = (VarInt::new(interval.start_inclusive()).unwrap()
            ..=VarInt::new(interval.end_inclusive()).unwrap())
            .into();
        let actual: Vec<u8> = viewer
            .next_view(interval, false)
            .iter::<&[u8]>()
            .flatten()
            .copied()
            .collect();
        assert_eq!(actual, expected);
    }

    #[test]
    fn view_test() {
        let mut buffer = Buffer::default();

        buffer.push(Bytes::from_static(&[0, 1, 2]));
        buffer.push(Bytes::from_static(&[3, 4, 5]));

        check_view(&buffer, (0..1).into(), &[0]);
        check_view(&buffer, (0..2).into(), &[0, 1]);
        check_view(&buffer, (0..3).into(), &[0, 1, 2]);
        check_view(&buffer, (0..4).into(), &[0, 1, 2, 3]);
        check_view(&buffer, (0..5).into(), &[0, 1, 2, 3, 4]);
        check_view(&buffer, (0..6).into(), &[0, 1, 2, 3, 4, 5]);
        check_view(&buffer, (1..6).into(), &[1, 2, 3, 4, 5]);
        check_view(&buffer, (2..6).into(), &[2, 3, 4, 5]);
        check_view(&buffer, (3..6).into(), &[3, 4, 5]);
        check_view(&buffer, (4..6).into(), &[4, 5]);
        check_view(&buffer, (5..6).into(), &[5]);
    }

    #[test]
    fn viewer_test() {
        let mut buffer = Buffer::default();

        buffer.push(Bytes::from_static(&[0, 1, 2]));
        buffer.push(Bytes::from_static(&[3, 4, 5]));

        let mut viewer = buffer.viewer();

        check_viewer(&mut viewer, (0..1).into(), &[0]);
        check_viewer(&mut viewer, (2..4).into(), &[2, 3]);
        check_viewer(&mut viewer, (5..6).into(), &[5]);
    }
}


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

use super::View;
use crate::contexts::WriteContext;
use s2n_quic_core::{frame::FitError, varint::VarInt};

/// Manages the outgoing flow control window for a sending data on a particular
/// data stream.
pub trait OutgoingDataFlowController {
    /// Tries to acquire a flow control window for the described chunk of data.
    /// The implementation must return the **maximum** (exclusive) offset up to
    /// which the data sender is allowed to send.
    fn acquire_flow_control_window(&mut self, end_offset: VarInt) -> VarInt;

    /// Returns `true` if sending data on the `Stream` was blocked because the
    /// the call to `acquire_flow_control_window` did not return any available
    /// window. This means not even the request for the minimum window size could
    /// be fulfilled.
    fn is_blocked(&self) -> bool;

    /// Clears the `is_blocked` flag which is stored inside the `FlowController`.
    /// The next call to `is_blocked` will return `None`, until another call to
    /// `acquire_flow_control_window` will move it back into the blocked state.
    fn clear_blocked(&mut self);

    /// Signals the flow controller that no further data will be submitted on
    /// the stream and therefore no further flow control window will be requested.
    fn finish(&mut self);
}

/// Writes chunks of data into frames.
pub trait FrameWriter: Default {
    // A value to be passed to the frame writer
    type Context: Copy;

    /// Indicates that this writer will potentially write a FIN frame, which
    /// marks the end of a stream.
    ///
    /// If set to `false`, FIN tracking will not be performed by the data sender
    const WRITES_FIN: bool = true;

    /// The minimum payload size we want to be able to write in a single frame,
    /// in case the frame would get fragmented due to this.
    /// We want to avoid writing too small chunks, since every chunk requires us
    /// to allocate an associated tracking state on sender and receiver side.
    const MIN_WRITE_SIZE: usize = 32;

    /// Indicates that this writer will retransmit unacknowledged data in probe
    /// packets.
    ///
    /// If set to `true`, the data sender will retransmit pending data if the
    /// transmission::Constraint is `PROBING`
    const RETRANSMIT_IN_PROBE: bool = false;

    /// Asks the writer to write a frame for the given chunk of data at the offset
    /// provided. The implementation should ensure the view fits by calling `trim_off`.
    ///
    /// If the frame will not fit with the current buffer capacity, a `FitError` should
    /// be returned.
    fn write_chunk<W: WriteContext>(
        &self,
        offset: VarInt,
        payload: &mut View,
        writer_context: Self::Context,
        context: &mut W,
    ) -> Result<(), FitError>;

    /// Asks the writer to write a FIN frame for the given offset.
    ///
    /// If the frame will not fit with the current buffer capacity, a `FitError` should
    /// be returned.
    ///
    /// This method will only be called if `WRITES_FIN` is set to `true`.
    fn write_fin<W: WriteContext>(
        &self,
        offset: VarInt,
        writer_context: Self::Context,
        context: &mut W,
    ) -> Result<(), FitError>;
}


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

use super::{
    buffer::{Buffer, Viewer},
    FinState, FrameWriter, OutgoingDataFlowController, State,
};
use crate::contexts::{OnTransmitError, WriteContext};
use core::num::NonZeroU16;
use s2n_quic_core::{
    ack,
    interval_set::{Interval, IntervalSet},
    packet::number::{Map as PacketNumberMap, PacketNumber, PacketNumberRange},
    varint::VarInt,
};

#[derive(Debug)]
pub struct Transmissions<FlowController, Writer> {
    /// Tracking information for all data in transmission
    in_flight: Set,
    //= https://www.rfc-editor.org/rfc/rfc9000#section-4.4
    //# Both endpoints MUST maintain flow control state
    //# for the stream in the unterminated direction until that direction
    //# enters a terminal state.
    /// The flow controller which is used to determine whether data chunks can
    /// be sent.
    pub flow_controller: FlowController,
    /// Serializes chunks into frames and writes the frames
    writer: Writer,
}

impl<FlowController: OutgoingDataFlowController, Writer: FrameWriter>
    Transmissions<FlowController, Writer>
{
    #[inline]
    pub fn new(flow_controller: FlowController) -> Self {
        Self {
            in_flight: Default::default(),
            flow_controller,
            writer: Default::default(),
        }
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.in_flight.is_empty()
    }

    #[inline]
    pub fn on_ack_signal<Set: ack::Set, F: FnMut(Interval<VarInt>)>(
        &mut self,
        ack_set: &Set,
        mut on_range: F,
    ) -> bool {
        let mut changed = false;

        let range = ack_set.as_range();

        for range in self.in_flight.remove_range(range) {
            on_range(range);
            changed = true;
        }

        changed
    }

    #[inline]
    pub fn transmit_set<W: WriteContext>(
        &mut self,
        buffer: &Buffer,
        set: &mut IntervalSet<VarInt>,
        state: &mut State,
        writer_context: Writer::Context,
        context: &mut W,
    ) -> Result<bool, OnTransmitError> {
        // make sure we've got something to transmit
        if set.is_empty() {
            return Ok(false);
        }

        let mut viewer = buffer.viewer();

        let mut has_transmitted = false;
        while let Some(mut interval) = set.pop_min() {
            match self.transmit_interval(&mut viewer, interval, state, writer_context, context) {
                Ok(transmitted) => {
                    has_transmitted = true;
                    let len = transmitted.len();
                    if len != interval.len() {
                        // only a part of the range was written so push back what wasn't
                        interval =
                            (interval.start_inclusive() + len..=interval.end_inclusive()).into();
                        debug_assert!(interval.is_valid());
                        set.insert_front(interval).unwrap();
                        return Ok(has_transmitted);
                    }
                }
                Err(err) => {
                    // if the interval failed to transmit it, put it back
                    set.insert_front(interval).unwrap();
                    return Err(err);
                }
            }
        }

        Ok(has_transmitted)
    }

    #[inline]
    pub fn transmit_interval<W: WriteContext>(
        &mut self,
        viewer: &mut Viewer,
        mut interval: Interval<VarInt>,
        state: &mut State,
        writer_context: Writer::Context,
        context: &mut W,
    ) -> Result<Interval<VarInt>, OnTransmitError> {
        // First trim the range to the current packet buffer
        let capacity = context.remaining_capacity();
        // Bound the capacity to u16::MAX - UDP payloads can't be larger anyway
        let capacity = capacity.min(u16::MAX as _);
        let mut interval_len = interval.len();

        // ensure the current packet buffer meets our minimum requirements
        if capacity == 0
            || interval_len >= Writer::MIN_WRITE_SIZE && capacity < Writer::MIN_WRITE_SIZE
            || !self.in_flight.has_capacity()
        {
            return Err(OnTransmitError::CouldNotAcquireEnoughSpace);
        }

        if capacity < interval_len {
            interval.set_len(capacity);
            interval_len = capacity;
        }

        let window_len = {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-2.2
            //# An endpoint MUST NOT send data on any stream without ensuring that it
            //# is within the flow control limits set by its peer.

            //= https://www.rfc-editor.org/rfc/rfc9000#section-4.1
            //# Senders MUST NOT send data in excess of either limit.
            self.flow_controller
                .acquire_flow_control_window(interval.end_exclusive())
                .checked_sub(interval.start_inclusive())
                .ok_or(OnTransmitError::CouldNotAcquireEnoughSpace)?
                .as_u64()
        };

        // ensure the window is non-zero
        if window_len == 0 {
            return Err(OnTransmitError::CouldNotAcquireEnoughSpace);
        }

        if window_len < interval_len as u64 {
            interval.set_len(window_len as usize);
        }

        let packet_number = context.packet_number();
        let mut view = viewer.next_view(interval, matches!(state, State::Finishing(_)));

        self.writer
            .write_chunk(
                interval.start_inclusive(),
                &mut view,
                writer_context,
                context,
            )
            .map_err(|_| OnTransmitError::CouldNotAcquireEnoughSpace)?;

        let len = view.len();
        debug_assert_ne!(len, 0u64, "cannot transmit an empty payload");

        interval.set_len(len.as_u64() as usize);

        debug_assert!(interval.is_valid());

        self.in_flight
            .insert(packet_number, interval.start_inclusive(), len);

        // Piggyback a fin transmission if we can
        if Writer::WRITES_FIN && view.is_fin() {
            if let Some(state) = state.fin_state_mut() {
                state.on_transmit(packet_number);
            }
        }

        Ok(interval)
    }

    #[inline]
    pub fn transmit_fin<W: WriteContext>(
        &mut self,
        buffer: &Buffer,
        state: &mut State,
        writer_context: Writer::Context,
        context: &mut W,
    ) -> Result<(), OnTransmitError> {
        // make sure we're not blocked before transmitting the fin bit
        if self.flow_controller.is_blocked() {
            return Err(OnTransmitError::CouldNotWriteFrame);
        }

        if let Some(state) = state.fin_state_mut() {
            if matches!(state, FinState::Pending | FinState::Lost) {
                let packet_number = context.packet_number();

                self.writer
                    .write_fin(buffer.total_len(), writer_context, context)
                    .map_err(|_| OnTransmitError::CouldNotAcquireEnoughSpace)?;

                state.on_transmit(packet_number);
            }
        }

        Ok(())
    }

    /// Remove all inflight transmissions
    #[inline]
    pub fn clear(&mut self) {
        self.in_flight.clear();
    }

    /// Remove all inflight transmissions and finish the flow controller
    #[inline]
    pub fn finish(&mut self) {
        self.clear();
        self.flow_controller.finish();
    }

    /// Get the inflight inclusive PacketNumberRange
    #[inline]
    pub fn get_inflight_range(&self) -> PacketNumberRange {
        self.in_flight.packets.get_range()
    }
}

/// Describes a chunk of bytes which has to be transmitted to the peer
#[derive(Copy, Clone, Debug)]
struct Transmission {
    /// The range of data that was sent in this transmission
    offset: VarInt,
    /// The length of data that was sent in the transmission
    len: u16,
    /// An optional next transmission in the same packet
    next: Option<TransmissionId>,
}

impl Transmission {
    #[inline]
    pub fn range(&self) -> Interval<VarInt> {
        (self.offset..self.offset + VarInt::from_u16(self.len)).into()
    }
}

#[derive(Debug, Default)]
struct Set {
    /// The packets that are currently in flight
    packets: PacketNumberMap<Transmission>,
    /// A slab of transmission ranges
    ///
    /// Because a packet number can have more than one transmission range,
    /// we need to store any extras outside of the map
    overflow: TransmissionSlab,
}

impl Set {
    #[inline]
    pub fn insert(&mut self, packet_number: PacketNumber, offset: VarInt, len: VarInt) {
        debug_assert!(len <= u16::MAX as u64);

        let transmission = Transmission {
            offset,
            len: len.as_u64() as _,
            next: None,
        };

        let transmissions = &mut self.overflow;

        self.packets
            .insert_or_update(packet_number, transmission, |prev| {
                // if we already have a entry for this packet number then chain the transmissions
                // together
                let idx = transmissions.insert(transmission);

                if let Some(prev) = prev.next {
                    transmissions.chain(prev, idx);
                } else {
                    prev.next = Some(idx);
                }
            });
    }

    #[inline]
    pub fn remove_range(&mut self, range: PacketNumberRange) -> SetRemoveIter {
        SetRemoveIter {
            inner: self.packets.remove_range(range),
            next: None,
            transmissions: &mut self.overflow,
        }
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.packets.is_empty()
    }

    #[inline]
    pub fn has_capacity(&self) -> bool {
        self.overflow.has_capacity()
    }

    #[inline]
    pub fn clear(&mut self) {
        self.packets.clear();
        self.overflow.clear();
    }
}

struct SetRemoveIter<'a> {
    inner: s2n_quic_core::packet::number::map::RemoveIter<'a, Transmission>,
    next: Option<TransmissionId>,
    transmissions: &'a mut TransmissionSlab,
}

impl<'a> Iterator for SetRemoveIter<'a> {
    type Item = Interval<VarInt>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if let Some(idx) = self.next.take() {
            let transmission = self.transmissions.remove(idx);
            self.next = transmission.next;
            return Some(transmission.range());
        }

        let (_, transmission) = self.inner.next()?;
        self.next = transmission.next;
        Some(transmission.range())
    }
}

#[derive(Debug, Default)]
struct TransmissionSlab {
    entries: Vec<TransmissionSlabEntry>,
    len: u16,
    next_free: u16,
}

/// An index into the transmission slab
///
/// We use a NonZeroU16 so it's the same size as `Option<TransmissionId>`
#[derive(Clone, Copy, Debug)]
struct TransmissionId(NonZeroU16);

#[derive(Debug)]
struct TransmissionSlabEntry {
    transmission: Transmission,
    next_free: u16,

    #[cfg(debug_assertions)]
    /// Tracks if the entry is occupied or not
    ///
    /// This is only needed with debug_assertions since the next_free logic actually
    /// takes care of this for us. This is all assuming the caller doesn't ever double-free
    /// any entries.
    ///
    /// This field exists to make sure these invariants stay true while testing.
    occupied: bool,
}

impl TransmissionSlab {
    #[inline]
    fn insert(&mut self, transmission: Transmission) -> TransmissionId {
        debug_assert!(self.has_capacity());
        let id = self.next_free;
        let index = id as usize;
        self.len += 1;

        let new_entry = TransmissionSlabEntry {
            transmission,
            next_free: 0,

            #[cfg(debug_assertions)]
            occupied: true,
        };

        if let Some(entry) = self.entries.get_mut(index) {
            #[cfg(debug_assertions)]
            assert!(!entry.occupied);

            self.next_free = entry.next_free;
            *entry = new_entry;
        } else {
            self.entries.push(new_entry);
            self.next_free += 1;
        }

        TransmissionId(NonZeroU16::new(1 + id).unwrap())
    }

    #[inline]
    fn remove(&mut self, index: TransmissionId) -> Transmission {
        let index = index.0.get() - 1;
        let entry = &mut self.entries[index as usize];

        #[cfg(debug_assertions)]
        {
            entry.occupied = false;
        }

        entry.next_free = self.next_free;
        self.next_free = index;
        self.len -= 1;
        entry.transmission
    }

    #[inline]
    fn chain(&mut self, prev: TransmissionId, next: TransmissionId) {
        let prev_entry = self.get_mut(prev);

        #[cfg(debug_assertions)]
        assert!(prev_entry.occupied);

        let next_entry = core::mem::replace(&mut prev_entry.transmission.next, Some(next));

        #[cfg(debug_assertions)]
        assert!(self.get_mut(next).occupied);

        self.get_mut(next).transmission.next = next_entry;
    }

    #[inline]
    fn get_mut(&mut self, idx: TransmissionId) -> &mut TransmissionSlabEntry {
        let index = (idx.0.get() - 1) as usize;

        #[cfg(debug_assertions)]
        assert!(self.entries[index].occupied);

        &mut self.entries[index]
    }

    #[inline]
    fn has_capacity(&self) -> bool {
        // we need to be able to store the correct index so it's 1 less than the max
        self.len < u16::MAX - 1
    }

    #[inline]
    fn clear(&mut self) {
        self.entries.clear();
        self.len = 0;
        self.next_free = 0;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bolero::{check, generator::*};

    #[derive(Clone, Copy, Debug, TypeGenerator)]
    enum Operation {
        Insert,
        Remove(usize),
        Clear,
    }

    #[test]
    fn transmission_slab_test() {
        check!().with_type::<Vec<Operation>>().for_each(|ops| {
            let mut subject = TransmissionSlab::default();
            let mut oracle = Vec::new();
            let mut offset = VarInt::from_u8(0);
            // keeps track of the peak len of the oracle so we can make sure it matches the slab
            // entries.len
            let mut peak_len = 0;

            for op in ops {
                match op {
                    Operation::Insert => {
                        let id = subject.insert(Transmission {
                            offset,
                            len: 1,
                            next: None,
                        });
                        oracle.push((offset, id));
                        offset += 1;
                        peak_len = peak_len.max(oracle.len());

                        assert_eq!(subject.entries.len(), peak_len);
                        assert_eq!(subject.len as usize, oracle.len());
                    }
                    Operation::Remove(index) => {
                        if oracle.is_empty() {
                            continue;
                        }

                        let index = index % oracle.len();
                        let (offset, id) = oracle.remove(index);
                        let transmission = subject.remove(id);
                        assert_eq!(offset, transmission.offset);

                        assert_eq!(subject.entries.len(), peak_len);
                        assert_eq!(subject.len as usize, oracle.len());
                    }
                    Operation::Clear => {
                        subject.clear();
                        oracle.clear();
                        peak_len = 0;
                    }
                }
            }
        });
    }

    #[test]
    #[cfg(target_pointer_width = "64")]
    fn size_test() {
        use core::mem::size_of;
        use insta::assert_debug_snapshot;

        assert_debug_snapshot!(
            "transmission_entry_size",
            size_of::<TransmissionSlabEntry>()
        );
        assert_eq!(
            size_of::<TransmissionId>(),
            size_of::<Option<TransmissionId>>()
        );
    }
}