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

macro_rules! define_inet_type {
    ($($vis:ident)? struct $name:ident {
        $(
            $field:ident: $field_ty:ty
        ),*
        $(,)*
    }) => {
        #[allow(unused_imports)]
        #[cfg(any(test, feature = "generator"))]
        use bolero_generator::*;

        #[derive(Clone, Copy, Default, Eq, PartialEq, PartialOrd, Ord, zerocopy::FromBytes, zerocopy::AsBytes, zerocopy::Unaligned)]
        #[cfg_attr(any(test, feature = "generator"), derive(bolero_generator::TypeGenerator))]
        #[repr(C)]
        $($vis)? struct $name {
            $(
                pub(crate) $field: $field_ty,
            )*
        }

        // By letting the compiler derive PartialEq, we can do structural matching on these
        // structs. But we also want hashing to be on the byte level, since we know the struct
        // has no allocations and has a simple layout.
        //
        // See: https://godbolt.org/z/czohnrWxK
        #[allow(clippy::derived_hash_with_manual_eq)]
        impl core::hash::Hash for $name {
            #[inline]
            fn hash<H: core::hash::Hasher>(&self, hasher: &mut H) {
                self.as_bytes().hash(hasher);
            }
        }

        impl $name {
            #[allow(non_camel_case_types, clippy::too_many_arguments)]
            #[inline]
            pub fn new<$($field: Into<$field_ty>),*>($($field: $field),*) -> Self {
                Self {
                    $(
                        $field: $field.into()
                    ),*
                }
            }

            #[inline]
            pub fn as_bytes(&self) -> &[u8] {
                zerocopy::AsBytes::as_bytes(self)
            }

            #[inline]
            pub fn as_bytes_mut(&mut self) -> &mut [u8] {
                zerocopy::AsBytes::as_bytes_mut(self)
            }
        }

        s2n_codec::zerocopy_value_codec!($name);
    };
}

macro_rules! test_inet_snapshot {
    ($test:ident, $name:ident, $ty:ty) => {
        #[test]
        fn $name() {
            s2n_codec::assert_codec_round_trip_sample_file!(
                $ty,
                concat!("src/inet/test_samples/", stringify!($test), ".bin")
            );
        }
    };
}


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

#[macro_use]
mod macros;

pub mod checksum;
pub mod datagram;
pub mod ecn;
pub mod ethernet;
pub mod ip;
pub mod ipv4;
pub mod ipv6;
pub mod udp;
pub mod unspecified;

pub use datagram::{AncillaryData, DatagramInfo};
pub use ecn::*;
pub use ip::*;
pub use ipv4::{IpV4Address, SocketAddressV4};
pub use ipv6::{IpV6Address, SocketAddressV6};
pub use unspecified::*;


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

use core::fmt;
use s2n_codec::zerocopy::U16;

// NOTE: duvet doesn't know how to parse this RFC since it doesn't follow more modern formatting
//# https://www.rfc-editor.org/rfc/rfc768#format
//# Format
//# ------
//#                  0      7 8     15 16    23 24    31
//#                 +--------+--------+--------+--------+
//#                 |     Source      |   Destination   |
//#                 |      Port       |      Port       |
//#                 +--------+--------+--------+--------+
//#                 |                 |                 |
//#                 |     Length      |    Checksum     |
//#                 +--------+--------+--------+--------+
//#                 |
//#                 |          data octets ...
//#                 +---------------- ...

define_inet_type!(
    pub struct Header {
        source: U16,
        destination: U16,
        len: U16,
        checksum: U16,
    }
);

impl Header {
    /// Swaps the direction of the header
    #[inline]
    pub fn swap(&mut self) {
        core::mem::swap(&mut self.source, &mut self.destination)
    }

    #[inline]
    pub const fn source(&self) -> &U16 {
        &self.source
    }

    #[inline]
    pub fn source_mut(&mut self) -> &mut U16 {
        &mut self.source
    }

    #[inline]
    pub const fn destination(&self) -> &U16 {
        &self.destination
    }

    #[inline]
    pub fn destination_mut(&mut self) -> &mut U16 {
        &mut self.destination
    }

    #[inline]
    pub const fn len(&self) -> &U16 {
        &self.len
    }

    #[inline]
    pub fn len_mut(&mut self) -> &mut U16 {
        &mut self.len
    }

    #[inline]
    pub const fn checksum(&self) -> &U16 {
        &self.checksum
    }

    #[inline]
    pub fn checksum_mut(&mut self) -> &mut U16 {
        &mut self.checksum
    }
}

impl fmt::Debug for Header {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("udp::Header")
            .field("source", &self.source)
            .field("destination", &self.destination)
            .field("len", &self.len)
            .field("checksum", &format_args!("0x{:04x}", &self.checksum.get()))
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bolero::check;
    use s2n_codec::DecoderBuffer;

    #[test]
    #[cfg_attr(miri, ignore)]
    fn snapshot_test() {
        let mut buffer = vec![0u8; core::mem::size_of::<Header>()];
        for (idx, byte) in buffer.iter_mut().enumerate() {
            *byte = idx as u8;
        }
        let decoder = DecoderBuffer::new(&buffer);
        let (header, _) = decoder.decode::<&Header>().unwrap();
        insta::assert_debug_snapshot!("snapshot_test", header);

        for byte in &mut buffer {
            *byte = 255;
        }
        let decoder = DecoderBuffer::new(&buffer);
        let (header, _) = decoder.decode::<&Header>().unwrap();
        insta::assert_debug_snapshot!("snapshot_filled_test", header);
    }

    #[test]
    fn header_round_trip_test() {
        check!().for_each(|buffer| {
            s2n_codec::assert_codec_round_trip_bytes!(Header, buffer);
        });
    }
}


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

/// A trait to determine if the value is left unspecified,
/// usually containing the default value.
///
/// See: <https://en.wikipedia.org/wiki/IPv6_address#Unspecified_address>
pub trait Unspecified: Sized {
    /// Returns true if the value is unspecified
    fn is_unspecified(&self) -> bool;

    /// Coerce a potentially unspecified value into an Option<Self>
    fn filter_unspecified(self) -> Option<Self> {
        if self.is_unspecified() {
            None
        } else {
            Some(self)
        }
    }
}

macro_rules! unspecified_integer {
    ($name:ident) => {
        impl Unspecified for s2n_codec::zerocopy::$name {
            fn is_unspecified(&self) -> bool {
                Self::default().eq(self)
            }
        }
    };
}

unspecified_integer!(U16);
unspecified_integer!(U32);
unspecified_integer!(U64);
unspecified_integer!(U128);


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

#![cfg_attr(kani, allow(dead_code))]

use super::{Accumulator, LargeWriteFn, State, LARGE_WRITE_LEN};

#[cfg(target_arch = "x86")]
pub use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
pub use core::arch::x86_64::*;

/// Returns the optimized function for the given platform
///
/// If possible, this uses runtime feature detection so this should be cached.
#[inline]
pub fn probe() -> Option<LargeWriteFn> {
    probe_avx().or_else(probe_sse)
}

macro_rules! probe {
    ($feature:tt, $fun:ident) => {{
        #[cfg(all(feature = "std", not(any(miri, kani))))]
        {
            if std::is_x86_feature_detected!($feature) {
                return Some($fun);
            }
        }

        // no way to reliably detect features in no_std
        None
    }};
}

#[inline]
fn probe_avx() -> Option<LargeWriteFn> {
    probe!("avx", write_sized_avx)
}

#[inline]
fn probe_sse() -> Option<LargeWriteFn> {
    probe!("sse4.1", write_sized_sse)
}

/// Enable sse4.1 optimizations for the implementation
#[target_feature(enable = "sse4.1")]
pub unsafe fn write_sized_sse<'a>(state: &mut State, bytes: &'a [u8]) -> &'a [u8] {
    write_sized(state, bytes)
}

/// Enable avx optimizations for the implementation
#[target_feature(enable = "avx")]
pub unsafe fn write_sized_avx<'a>(state: &mut State, bytes: &'a [u8]) -> &'a [u8] {
    write_sized(state, bytes)
}

#[inline(always)]
unsafe fn write_sized<'a>(state: &mut State, mut bytes: &'a [u8]) -> &'a [u8] {
    assume!(
        bytes.len() >= LARGE_WRITE_LEN,
        "large write function should only be called with large byte buffers"
    );

    // Maintain two different sums to increase instruction parallelism
    let mut sum_a = Sum::new();
    let mut sum_b = Sum::new();

    // Unroll the loop to read 256 bytes at a time (16 16-byte blocks)
    while bytes.len() >= 256 {
        let (chunks, remaining) = bytes.split_at(256);
        bytes = remaining;

        let ptr = chunks.as_ptr() as *const __m128i;

        sum_a += ptr;
        sum_a += ptr.add(1);
        sum_a += ptr.add(2);
        sum_a += ptr.add(3);
        sum_a += ptr.add(4);
        sum_a += ptr.add(5);
        sum_a += ptr.add(6);
        sum_a += ptr.add(7);

        sum_b += ptr.add(8);
        sum_b += ptr.add(9);
        sum_b += ptr.add(10);
        sum_b += ptr.add(11);
        sum_b += ptr.add(12);
        sum_b += ptr.add(13);
        sum_b += ptr.add(14);
        sum_b += ptr.add(15);
    }

    // Unroll the loop to read 64 bytes at a time (4 16-byte blocks)
    while bytes.len() >= 64 {
        let (chunks, remaining) = bytes.split_at(64);
        bytes = remaining;

        let ptr = chunks.as_ptr() as *const __m128i;

        sum_a += ptr;
        sum_a += ptr.add(1);
        sum_b += ptr.add(2);
        sum_b += ptr.add(3);
    }

    // Finish reading the full 16-byte blocks
    while bytes.len() >= 16 {
        let (chunks, remaining) = bytes.split_at(16);
        bytes = remaining;

        let ptr = chunks.as_ptr() as *const __m128i;
        sum_a += ptr;
    }

    // Add up all of the sums and merge them into a single u32
    let sum = sum_a + sum_b;

    sum.extract(state);

    bytes
}

#[derive(Clone)]
struct Sum {
    a: __m128i,
    b: __m128i,
}

impl Sum {
    #[inline(always)]
    unsafe fn new() -> Self {
        Self {
            a: _mm_setzero_si128(),
            b: _mm_setzero_si128(),
        }
    }

    #[inline(always)]
    unsafe fn add(&mut self, rhs: __m128i) {
        // Reads pairs of bytes into a 32-bit value
        //
        // Since we have 16 bytes as input, we need two outputs since we're doubling the bit-width
        let mask_a = _mm_setr_epi8(
            0x0, 0x1, -1, -1, 0x2, 0x3, -1, -1, 0x4, 0x5, -1, -1, 0x6, 0x7, -1, -1,
        );
        let mask_b = _mm_setr_epi8(
            0x8, 0x9, -1, -1, 0xa, 0xb, -1, -1, 0xc, 0xd, -1, -1, 0xe, 0xf, -1, -1,
        );

        // Add the shuffled counts to our current state
        self.a = _mm_add_epi32(self.a, _mm_shuffle_epi8(rhs, mask_a));
        self.b = _mm_add_epi32(self.b, _mm_shuffle_epi8(rhs, mask_b));
    }

    #[inline(always)]
    unsafe fn extract(self, state: &mut State) {
        let total = _mm_add_epi32(self.a, self.b);
        let total = _mm_hadd_epi32(total, _mm_setzero_si128());
        let total = _mm_hadd_epi32(total, _mm_setzero_si128());

        *state += _mm_extract_epi32(total, 0) as Accumulator;
    }
}

impl core::ops::AddAssign<__m128i> for Sum {
    #[inline(always)]
    fn add_assign(&mut self, rhs: __m128i) {
        unsafe { self.add(rhs) }
    }
}

impl core::ops::AddAssign<*const __m128i> for Sum {
    #[inline(always)]
    fn add_assign(&mut self, rhs: *const __m128i) {
        unsafe { self.add(_mm_loadu_si128(rhs)) }
    }
}

impl core::ops::Add for Sum {
    type Output = Self;

    #[inline(always)]
    fn add(self, rhs: Self) -> Self {
        unsafe {
            let a = _mm_add_epi32(self.a, self.b);
            let b = _mm_add_epi32(rhs.a, rhs.b);
            Self { a, b }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{super::write_sized_generic_u16, *};
    use bolero::check;
    use core::num::Wrapping;

    /// Compares the x86-optimized function to the generic function and ensures consistency
    #[test]
    fn differential() {
        if let Some(write_sized_opt) = probe() {
            check!().for_each(|bytes| {
                if bytes.len() < LARGE_WRITE_LEN {
                    return;
                }

                let generic = {
                    let mut state = Wrapping(0);
                    write_sized_generic_u16::<2>(&mut state, bytes);
                    state.0
                };

                let actual = {
                    let mut state = Wrapping(0);
                    let bytes = unsafe { write_sized_opt(&mut state, bytes) };
                    write_sized_generic_u16::<2>(&mut state, bytes);
                    state.0
                };

                assert_eq!(generic.to_ne_bytes(), actual.to_ne_bytes())
            });
        }
    }
}


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

use crate::interval_set::{Interval, IntervalBound, IntervalSetError};
use alloc::collections::VecDeque;
use core::{
    cmp::{max, min, Ordering},
    num::NonZeroUsize,
    ops::Range,
};

#[inline]
pub(crate) fn insert<T: IntervalBound + Ord>(
    ranges: &mut VecDeque<Interval<T>>,
    mut range: Interval<T>,
    start_index: usize,
    limit: Option<NonZeroUsize>,
) -> Result<usize, IntervalSetError> {
    // this range is intentionally invalid and will only be
    // valid if the `scan` method finds a match
    #[allow(clippy::reversed_empty_ranges)]
    let replace_range = core::usize::MAX..0;

    let mut insertion = Insertion { replace_range };

    let iter = ranges.iter().enumerate().skip(start_index);

    if let Some(index) = insertion.scan(iter, &mut range) {
        return Ok(index);
    }

    insertion.apply(ranges, range, limit)
}

/// A structure to keep temporary state for an insertion
#[derive(Debug)]
struct Insertion {
    replace_range: Range<usize>,
}

impl Insertion {
    /// Scans over the Intervals and updates the `Insertion` state with the
    /// applicable interval ranges
    ///
    /// Returns `Some(index)` if the `Interval` is already present in the iterator,
    /// otherwise `None` is returned.
    #[inline]
    fn scan<'a, T: 'a + Ord + IntervalBound, I: Iterator<Item = (usize, &'a Interval<T>)>>(
        &mut self,
        ranges: I,
        range_a: &mut Interval<T>,
    ) -> Option<usize> {
        use Ordering::*;

        for (slot_index, range_b) in ranges {
            match (
                range_a.start.cmp(&range_b.start),
                range_a.end.cmp(&range_b.end),
            ) {
                // the ranges are equal
                //
                // range A: |---------|
                // range B: |---------|
                //
                (Equal, Equal) |

                // range A is a subset of range B
                //
                // range A:     |-----|
                // range B: |---------|
                //
                // do nothing
                //
                (Greater, Equal) => return Some(slot_index + 1),

                // range A is a subset of range B
                //
                // range A: |-----|
                // range B: |---------|
                //
                // do nothing
                //
                (Equal, Less) |

                // range A is a subset of range B
                //
                // range A:    |----|
                // range B: |---------|
                //
                (Greater, Less) => return Some(slot_index),

                // range A is part of range B.
                //
                // Before:
                //
                // range A:      |--------|
                // range B: |--------|
                //
                // After:
                //
                // range A: |-------------|
                // range B: |
                //
                (Greater, Greater) if range_a.should_coalesce(range_b) => {
                    range_a.start = range_b.start;
                    self.set_start(slot_index);
                    let next_slot = slot_index + 1;
                    self.set_end(next_slot);
                    continue;
                }

                // range A comes later
                //
                // range A:          |-----|
                // range B: |----|
                //
                (Greater, Greater) => {
                    continue;
                }

                // range A contains range B, spilling over into the next slot
                //
                // range A: |---------|
                // range B: |---|
                //
                // mark B as obsolete
                //
                (Equal, Greater) |

                // range A contains range B, spilling over into the next slot
                //
                // range A: |---------|
                // range B:    |---|
                //
                // mark B as obsolete
                //
                (Less, Greater) => {
                    self.set_start(slot_index);
                    let next_slot = slot_index + 1;
                    self.set_end(next_slot);
                    continue;
                }

                // range A ends with range B
                //
                // range A: |---------|
                // range B:       |---|
                //
                // mark B as obsolete and return
                //
                (Less, Equal) => {
                    self.set_start(slot_index);
                    let next_slot = slot_index + 1;
                    self.set_end(next_slot);
                    break;
                }

                // range A overlaps part of range B
                //
                // Before:
                //
                // range A: |--------|
                // range B:      |--------|
                //
                // After:
                //
                // range A: |-------------|
                // range B:               |
                //
                (Less, Less) if range_b.should_coalesce(range_a) => {
                    range_a.end = range_b.end;
                    self.set_start(slot_index);
                    let next_slot = slot_index + 1;
                    self.set_end(next_slot);
                    break;
                }

                // range A comes before range B
                //
                // range A: |---|
                // range B:        |--------|
                //
                // returns; no more searching needed
                //
                (Less, Less) => {
                    self.set_start(slot_index);
                    self.set_end(slot_index);
                    break;
                }
            }
        }

        None
    }

    /// Applies the `Insertion` to the given set of `Interval`s.
    #[inline]
    fn apply<T>(
        self,
        ranges: &mut VecDeque<Interval<T>>,
        range: Interval<T>,
        limit: Option<NonZeroUsize>,
    ) -> Result<usize, IntervalSetError> {
        let replace_range = self.replace_range;
        let prev_len = ranges.len();

        let ensure_can_insert = || {
            let under_limit = if let Some(limit) = limit {
                limit.get() > prev_len
            } else {
                true
            };

            if under_limit {
                Ok(())
            } else {
                Err(IntervalSetError::LimitExceeded)
            }
        };

        let index = replace_range.start;
        let replace_count = if let Some(count) = replace_range.end.checked_sub(index) {
            count
        } else {
            ensure_can_insert()?;

            // add it to the end
            ranges.push_back(range);
            return Ok(prev_len);
        };

        match replace_count {
            0 => {
                ensure_can_insert()?;
                ranges.insert(index, range);
            }
            1 => {
                ranges[index] = range;
            }
            2 => {
                ranges[index] = range;
                ranges.remove(index + 1);
            }
            _ => {
                ranges[index] = range;
                ranges.drain(index + 1..replace_range.end);
            }
        };

        Ok(index)
    }

    #[inline]
    fn set_start(&mut self, start: usize) {
        self.replace_range.start = min(self.replace_range.start, start);
    }

    #[inline]
    fn set_end(&mut self, end: usize) {
        self.replace_range.end = max(self.replace_range.end, end);
    }
}


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

use crate::interval_set::{Interval, IntervalBound};
use alloc::collections::vec_deque::{self, VecDeque};
use core::cmp::Ordering;

/// An iterator of `Intervals` over the intersection of two sets,
/// i.e. the values in both `A` and `B` will be returned.
#[derive(Debug)]
pub struct Intersection<'a, T> {
    set_a: vec_deque::Iter<'a, Interval<T>>,
    set_b: vec_deque::Iter<'a, Interval<T>>,
    interval_a: Option<Interval<T>>,
    interval_b: Option<Interval<T>>,
}

impl<'a, T: Copy> Intersection<'a, T> {
    pub(crate) fn new(set_a: &'a VecDeque<Interval<T>>, set_b: &'a VecDeque<Interval<T>>) -> Self {
        let mut set_a = set_a.iter();
        let interval_a = set_a.next().cloned();

        let mut set_b = set_b.iter();
        let interval_b = set_b.next().cloned();

        Self {
            set_a,
            set_b,
            interval_a,
            interval_b,
        }
    }
}

impl<'a, T: 'a + IntervalBound + Ord> Iterator for Intersection<'a, T> {
    type Item = Interval<T>;

    fn next(&mut self) -> Option<Self::Item> {
        use Ordering::*;

        let interval_a = self.interval_a.as_mut()?;
        let interval_b = self.interval_b.as_mut()?;

        macro_rules! advance_set_a {
            () => {
                self.interval_a = self.set_a.next().cloned();
            };
        }

        macro_rules! advance_set_b {
            () => {
                self.interval_b = self.set_b.next().cloned();
            };
        }

        loop {
            match (
                interval_a.start.cmp(&interval_b.start),
                interval_a.end.cmp(&interval_b.end),
            ) {
                // interval A is a subset of interval B
                //
                // interval A: |-----|
                // interval B: |---------|
                //
                // Returns:
                //             |-----|
                //
                (Equal, Less) |

                // interval A is a subset of interval B
                //
                // interval A:    |----|
                // interval B: |---------|
                //
                // Returns:
                //                |----|
                //
                (Greater, Less) => {
                    let item = *interval_a;
                    interval_b.start = interval_a.end_exclusive();
                    advance_set_a!();
                    return Some(item);
                }

                // interval A contains interval B, spilling over into the next slot
                //
                // interval A: |---------|
                // interval B: |---|
                //
                // Returns:
                //             |---|
                //
                (Equal, Greater) |

                // interval A contains interval B, spilling over into the next slot
                //
                // interval A: |---------|
                // interval B:    |---|
                //
                // Returns:
                //                |---|
                //
                (Less, Greater) => {
                    let item = *interval_b;
                    interval_a.start = interval_b.end_exclusive();
                    advance_set_b!();
                    return Some(item);
                }

                // interval A is a subset of interval B
                //
                // interval A:     |-----|
                // interval B: |---------|
                //
                // Returns:
                //                 |-----|
                //
                (Greater, Equal) |

                // the intervals are equal
                //
                // interval A: |---------|
                // interval B: |---------|
                //
                // Returns:
                //             |---------|
                //
                (Equal, Equal) => {
                    let item = *interval_a;
                    advance_set_a!();
                    advance_set_b!();
                    return Some(item);
                }

                // interval A ends with interval B
                //
                // interval A: |---------|
                // interval B:       |---|
                //
                // Returns:
                //                   |---|
                //
                (Less, Equal) => {
                    let item = *interval_b;
                    advance_set_a!();
                    advance_set_b!();
                    return Some(item);
                }

                // interval A is part of interval B.
                //
                // interval A:      |--------|
                // interval B: |--------|
                //
                // Returns:
                //                  |---|
                //
                (Greater, Greater) if interval_a.start <= interval_b.end => {
                    let mut item = *interval_a;
                    item.end = interval_b.end;
                    interval_a.start = item.end_exclusive();
                    advance_set_b!();
                    return Some(item);
                }

                // interval A overlaps part of interval B
                //
                // interval A: |--------|
                // interval B:      |--------|
                //
                // Returns:
                //                  |---|
                //
                (Less, Less) if interval_a.end >= interval_b.start => {
                    let mut item = *interval_b;
                    item.end = interval_a.end;
                    interval_b.start = item.end_exclusive();
                    advance_set_a!();
                    return Some(item);
                }

                // interval A comes later
                //
                // interval A:          |-----|
                // interval B: |----|
                //
                // continue to next B
                //
                (Greater, Greater) => {
                    *interval_b = self.set_b.next().cloned()?;
                    continue;
                }

                // interval A comes before interval B
                //
                // interval A: |---|
                // interval B:        |--------|
                //
                // continue to next A
                //
                (Less, Less) => {
                    *interval_a = self.set_a.next().cloned()?;
                    continue;
                }
            }
        }
    }
}

/// Apply the intersection of `set_a` with `set_b` to `set_a`
#[inline]
pub(super) fn apply<T: IntervalBound>(
    set_a: &mut VecDeque<Interval<T>>,
    set_b: &VecDeque<Interval<T>>,
) {
    use Ordering::*;

    if set_a.is_empty() {
        return;
    }

    if set_b.is_empty() {
        set_a.clear();
        return;
    }

    let mut set_b = set_b.iter();
    let mut interval_b = set_b.next().expect("set_b is not empty");
    let mut a_index = 0;

    macro_rules! advance_set_a {
        () => {
            a_index += 1;
        };
    }

    macro_rules! advance_set_b {
        () => {
            if let Some(next_interval_b) = set_b.next() {
                interval_b = next_interval_b;
            } else {
                // the remainder of set_a is not in the intersection, since we've
                // reached the end of set_b
                set_a.truncate(a_index);
                return;
            }
        };
    }

    // Inserts a new interval into the next index of set_a that starts where the current interval b
    // ends (with a gap of 1) and ends where the current interval a ends
    macro_rules! split_off_a {
        ($interval_a:ident) => {
            // step_up_saturating is used because the next intersecting interval should not be
            // contiguous with the current interval, so we can expect a gap of at least 1.
            let new_interval_a: Interval<T> =
                (interval_b.end_exclusive().step_up_saturating()..=$interval_a.end).into();

            if new_interval_a.is_valid() {
                // if interval_a had only overlapped interval_b by 1, then the new interval
                // will not be valid and we can skip inserting it
                // for example: interval_a = [0..=48], interval_b = [0..=47]
                // new_interval_a = [49..=48] since we know 48 is not in interval_b
                set_a.insert(a_index + 1, new_interval_a);
            }
        };
    }

    while let Some(interval_a) = set_a.get(a_index) {
        match (interval_a.start.cmp(&interval_b.start),
               interval_a.end.cmp(&interval_b.end),
        ) {
            // interval A is a subset of interval B
            //
            // interval A: |-----|
            // interval B: |---------|
            //
            // End state:
            //             |-----|
            // a_index:           ^
            (Equal, Less) |

            // interval A is a subset of interval B
            //
            // interval A:    |----|
            // interval B: |---------|
            //
            // End state:
            //                |----|
            // a_index:             ^
            (Greater, Less) => {
                advance_set_a!();
            }

            // interval A contains interval B, spilling over into the next slot
            //
            // interval A: |---------|
            // interval B: |---|
            //
            // End state:
            //             |---| |---|
            // a_index:          ^
            (Equal, Greater) |

            // interval A contains interval B, spilling over into the next slot
            //
            // interval A: |---------|
            // interval B:   |---|
            //
            // End state:
            //               |---| |-|
            // a_index:            ^
            (Less, Greater) => {
                // split off the part of interval A that is spilling into the next slot
                split_off_a!(interval_a);
                set_a[a_index] = *interval_b;
                advance_set_a!();
                advance_set_b!();
            }

            // interval A is a subset of interval B
            //
            // interval A:     |-----|
            // interval B: |---------|
            //
            // End state:
            //                 |-----|
            // a_index:               ^
            (Greater, Equal) |

            // the intervals are equal
            //
            // interval A: |---------|
            // interval B: |---------|
            //
            // End state:
            //             |---------|
            // a_index:               ^
            (Equal, Equal) => {
                advance_set_a!();
                advance_set_b!();
            }

            // interval A ends with interval B
            //
            // interval A: |---------|
            // interval B:       |---|
            //
            // End state:
            //                   |---|
            // a_index:               ^
            (Less, Equal) => {
                set_a[a_index].start = interval_b.start;
                advance_set_a!();
                advance_set_b!();
            }

            // interval A overlaps part of interval B.
            //
            // interval A:      |--------|
            // interval B: |--------|
            //
            // End state:
            //                  |---| |--|
            // a_index:               ^
            (Greater, Greater) if interval_a.start <= interval_b.end => {
                // split off the part of interval A that is spilling into the next slot
                split_off_a!(interval_a);
                set_a[a_index].end = interval_b.end;
                advance_set_a!();
                advance_set_b!();
            }

            // interval A overlaps part of interval B
            //
            // interval A: |--------|
            // interval B:      |--------|
            //
            // End state:
            //                  |---|
            // a_index:              ^
            (Less, Less) if interval_a.end >= interval_b.start => {
                set_a[a_index].start = interval_b.start;
                advance_set_a!();
            }

            // interval A comes later
            //
            // interval A:          |-----|
            // interval B: |----|
            //
            // continue to next B
            //
            (Greater, Greater) => {
                advance_set_b!();
            }

            // interval A comes before interval B
            //
            // interval A: |---|
            // interval B:        |--------|
            //
            // remove A and continue to next A
            //
            (Less, Less) => {
                set_a.remove(a_index);
            }
        }
    }
}


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

use super::IntervalSetError;
use core::{
    cmp::Ordering,
    fmt,
    ops::{Bound, Range, RangeBounds, RangeInclusive},
};

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Interval<T> {
    pub(super) start: T,
    pub(super) end: T,
}

impl<T: IntervalBound> Interval<T> {
    pub fn from_range_bounds<B: RangeBounds<T>>(bounds: B) -> Result<Self, IntervalSetError> {
        let start = match bounds.start_bound() {
            Bound::Included(start) => *start,
            Bound::Excluded(start) => start.step_down().ok_or(IntervalSetError::InvalidInterval)?,
            _ => return Err(IntervalSetError::InvalidInterval),
        };

        let end = match bounds.end_bound() {
            Bound::Included(end) => *end,
            Bound::Excluded(end) => end.step_down().ok_or(IntervalSetError::InvalidInterval)?,
            _ => return Err(IntervalSetError::InvalidInterval),
        };

        let interval = Self { start, end };

        if interval.is_valid() {
            Ok(interval)
        } else {
            Err(IntervalSetError::InvalidInterval)
        }
    }

    #[inline]
    pub fn start_inclusive(&self) -> T {
        self.start
    }

    #[inline]
    pub fn start_exclusive(&self) -> T {
        self.start.step_down_saturating()
    }

    #[inline]
    pub fn end_inclusive(&self) -> T {
        self.end
    }

    #[inline]
    pub fn end_exclusive(&self) -> T {
        self.end.step_up_saturating()
    }

    #[inline]
    pub fn len(&self) -> usize {
        // Interval always has at least 1
        let interval_base_len = 1;
        interval_base_len + self.start.steps_between(&self.end)
    }

    #[inline]
    pub fn set_len(&mut self, len: usize) {
        debug_assert_ne!(len, 0, "Intervals cannot be 0 length");
        self.end = self.start.steps_between_len(len);
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        // Interval always has at least 1
        false
    }

    #[inline]
    pub fn is_valid(&self) -> bool {
        self.end >= self.start
    }

    #[inline]
    pub(crate) fn should_coalesce(&self, other: &Self) -> bool {
        self.start <= other.end_exclusive()
    }
}

impl<T: IntervalBound> IntoIterator for &Interval<T> {
    type IntoIter = Interval<T>;
    type Item = T;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        *self
    }
}

impl<T> RangeBounds<T> for Interval<T> {
    #[inline]
    fn start_bound(&self) -> Bound<&T> {
        Bound::Included(&self.start)
    }

    #[inline]
    fn end_bound(&self) -> Bound<&T> {
        Bound::Included(&self.end)
    }
}

impl<T> RangeBounds<T> for &Interval<T> {
    #[inline]
    fn start_bound(&self) -> Bound<&T> {
        Bound::Included(&self.start)
    }

    #[inline]
    fn end_bound(&self) -> Bound<&T> {
        Bound::Included(&self.end)
    }
}

impl<T: Ord> PartialEq<T> for Interval<T> {
    #[inline]
    fn eq(&self, value: &T) -> bool {
        self.partial_cmp(value) == Some(Ordering::Equal)
    }
}

impl<T: Ord> PartialOrd<T> for Interval<T> {
    #[inline]
    fn partial_cmp(&self, value: &T) -> Option<Ordering> {
        use Ordering::*;
        Some(match (self.start.cmp(value), self.end.cmp(value)) {
            (Equal, _) => Equal,
            (_, Equal) => Equal,
            (Greater, Less) => Equal,
            (Less, Greater) => Equal,
            (Less, Less) => Less,
            (Greater, Greater) => Greater,
        })
    }
}

impl<T: IntervalBound> From<&Interval<T>> for Interval<T> {
    #[inline]
    fn from(interval: &Interval<T>) -> Self {
        *interval
    }
}

macro_rules! range_impls {
    ($range_ty:ident, $into:expr, $from:expr) => {
        impl<T: IntervalBound> From<Interval<T>> for $range_ty<T> {
            #[inline]
            fn from(range: Interval<T>) -> Self {
                #[allow(clippy::redundant_closure_call)]
                ($into)(range)
            }
        }

        impl<T: IntervalBound> From<&Interval<T>> for $range_ty<T> {
            #[inline]
            fn from(range: &Interval<T>) -> Self {
                #[allow(clippy::redundant_closure_call)]
                ($into)(*range)
            }
        }

        impl<T: IntervalBound> From<$range_ty<T>> for Interval<T> {
            #[inline]
            fn from(range: $range_ty<T>) -> Self {
                #[allow(clippy::redundant_closure_call)]
                ($from)(range)
            }
        }

        impl<'a, T: IntervalBound> From<&'a $range_ty<T>> for Interval<T> {
            #[inline]
            fn from(range: &'a $range_ty<T>) -> Self {
                #[allow(clippy::redundant_closure_call)]
                ($from)(range.clone())
            }
        }

        impl<T: IntervalBound> PartialEq<$range_ty<T>> for Interval<T> {
            #[inline]
            fn eq(&self, other: &$range_ty<T>) -> bool {
                self.partial_cmp(other) == Some(Ordering::Equal)
            }
        }

        impl<T: IntervalBound> PartialEq<$range_ty<T>> for &Interval<T> {
            #[inline]
            fn eq(&self, other: &$range_ty<T>) -> bool {
                self.partial_cmp(other) == Some(Ordering::Equal)
            }
        }

        impl<T: IntervalBound> PartialOrd<$range_ty<T>> for Interval<T> {
            #[inline]
            fn partial_cmp(&self, other: &$range_ty<T>) -> Option<Ordering> {
                let other: Self = other.into();
                self.partial_cmp(&other)
            }
        }

        impl<T: IntervalBound> PartialOrd<$range_ty<T>> for &Interval<T> {
            #[inline]
            fn partial_cmp(&self, other: &$range_ty<T>) -> Option<Ordering> {
                let other: Interval<_> = other.into();
                (*self).partial_cmp(&other)
            }
        }
    };
}

range_impls!(
    Range,
    |interval: Interval<_>| interval.start..interval.end_exclusive(),
    |range: Range<_>| Self {
        start: range.start,
        end: range.end.step_down_saturating(),
    }
);

range_impls!(
    RangeInclusive,
    |interval: Interval<_>| interval.start..=interval.end,
    |range: RangeInclusive<_>| {
        let (start, end) = range.into_inner();
        Self { start, end }
    }
);

impl<T: IntervalBound> Iterator for Interval<T> {
    type Item = T;

    #[inline]
    fn next(&mut self) -> Option<T> {
        let current = self.start;
        if current > self.end {
            return None;
        }
        if let Some(next) = current.step_up() {
            self.start = next;
        } else {
            self.end = current.step_down()?;
        }
        Some(current)
    }
}

impl<T: IntervalBound> DoubleEndedIterator for Interval<T> {
    #[inline]
    fn next_back(&mut self) -> Option<T> {
        let current = self.end;
        if current < self.start {
            return None;
        }
        if let Some(next) = current.step_down() {
            self.end = next;
        } else {
            self.start = current.step_up()?;
        }
        Some(current)
    }
}

impl<T: fmt::Debug> fmt::Debug for Interval<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        ((&self.start)..=(&self.end)).fmt(f)
    }
}

pub trait IntervalBound: Copy + Ord + Sized {
    fn step_up(self) -> Option<Self>;
    fn step_down(self) -> Option<Self>;
    fn steps_between(&self, upper: &Self) -> usize;
    fn steps_between_len(&self, len: usize) -> Self;

    fn step_up_saturating(self) -> Self {
        self.step_up().unwrap_or(self)
    }

    fn step_down_saturating(self) -> Self {
        self.step_down().unwrap_or(self)
    }
}

macro_rules! integer_bounds {
    ($type:ident) => {
        impl IntervalBound for $type {
            #[inline]
            fn step_up(self) -> Option<Self> {
                self.checked_add(1)
            }

            #[inline]
            fn step_down(self) -> Option<Self> {
                self.checked_sub(1)
            }

            #[inline]
            fn steps_between(&self, upper: &Self) -> usize {
                use core::convert::TryInto;
                (upper - self).try_into().unwrap_or(core::usize::MAX)
            }

            #[inline]
            fn steps_between_len(&self, len: usize) -> Self {
                use core::convert::TryInto;
                let len: Self = (len - 1).try_into().unwrap();
                *self + len
            }
        }
    };
}

integer_bounds!(u8);
integer_bounds!(i8);
integer_bounds!(u16);
integer_bounds!(i16);
integer_bounds!(u32);
integer_bounds!(i32);
integer_bounds!(u64);
integer_bounds!(i64);
integer_bounds!(u128);
integer_bounds!(i128);
integer_bounds!(usize);
integer_bounds!(isize);

use crate::{packet::number::PacketNumber, varint::VarInt};

impl IntervalBound for VarInt {
    #[inline]
    fn step_up(self) -> Option<Self> {
        self.checked_add(Self::from_u8(1))
    }

    #[inline]
    fn step_down(self) -> Option<Self> {
        self.checked_sub(Self::from_u8(1))
    }

    #[inline]
    fn steps_between(&self, upper: &Self) -> usize {
        <u64 as IntervalBound>::steps_between(self, upper)
    }

    #[inline]
    fn steps_between_len(&self, len: usize) -> Self {
        *self + (len - 1)
    }
}

impl IntervalBound for PacketNumber {
    #[inline]
    fn step_up(self) -> Option<Self> {
        self.next()
    }

    #[inline]
    fn step_down(self) -> Option<Self> {
        let space = self.space();
        let value = PacketNumber::as_varint(self);
        Some(space.new_packet_number(value.step_down()?))
    }

    #[inline]
    fn steps_between(&self, upper: &Self) -> usize {
        let lower = PacketNumber::as_varint(*self);
        let upper = PacketNumber::as_varint(*upper);
        lower.steps_between(&upper)
    }

    #[inline]
    fn steps_between_len(&self, len: usize) -> Self {
        let start = PacketNumber::as_varint(*self);
        let end = start.steps_between_len(len);
        self.space().new_packet_number(end)
    }
}


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

#![forbid(unsafe_code)]

mod insert;
mod intersection;
pub mod interval;
mod remove;

#[cfg(test)]
mod tests;

use alloc::collections::vec_deque::{self, VecDeque};
use core::{
    fmt,
    iter::FromIterator,
    num::NonZeroUsize,
    ops::{Bound, Range, RangeBounds, RangeInclusive},
};
use insert::insert;
pub use intersection::Intersection;
pub use interval::*;
use remove::remove;

#[derive(Clone, Copy, Debug, PartialEq, PartialOrd, Eq, Ord, Hash)]
pub enum IntervalSetError {
    LimitExceeded,
    InvalidInterval,
}

/// `IntervalSet` is an efficient structure for storing sets of consecutive numbers. Instead
/// of storing an individual entry per value, only the lower and upper bounds (`Interval`) are stored.
///
/// ## Usage
///
/// ```rust,ignore
/// use s2n_quic_transport::interval_set::IntervalSet;
///
/// let mut set = IntervalSet::new();
///
/// set.insert_value(0);
/// set.insert_value(1);
/// set.insert_value(2);
/// set.insert_value(3);
///
/// // because 0 to 3 are consecutive, only a single interval is stored
/// assert_eq!(set.interval_len(), 1);
///
/// set.insert_value(5);
/// set.insert_value(6);
///
/// // 5 and 6 are not consecutive with 0 to 3 so a new entry is created
/// assert_eq!(set.interval_len(), 2);
///
/// set.insert_value(4);
///
/// // inserting a 4 merges all of the intervals into a single entry
/// assert_eq!(set.interval_len(), 1);
///
/// // ranges of numbers can be inserted at the same time
/// set.insert(12..15);
/// set.insert(18..=21);
///
/// assert_eq!(set.interval_len(), 3);
///
/// // ranges can also be removed
/// set.remove(0..22);
///
/// assert_eq!(set.interval_len(), 0);
/// ```
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct IntervalSet<T> {
    limit: Option<NonZeroUsize>,
    intervals: VecDeque<Interval<T>>,
}

impl<T> Default for IntervalSet<T> {
    fn default() -> Self {
        Self {
            limit: None,
            intervals: VecDeque::new(),
        }
    }
}

impl<T> IntervalSet<T> {
    /// Creates an empty `IntervalSet`
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// assert!(set.insert(0..4).is_ok());
    /// ```
    #[inline]
    pub fn new() -> IntervalSet<T> {
        Self::default()
    }

    /// Creates an empty `IntervalSet` with a specific capacity.
    /// This preallocates enough memory for `capacity` elements,
    /// so that the `IntervalSet` does not have to be reallocated
    /// until it contains at least that many values.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::with_capacity(10);
    /// assert!(set.insert(0..4).is_ok());
    /// ```
    #[inline]
    pub fn with_capacity(capacity: usize) -> IntervalSet<T> {
        let intervals = VecDeque::with_capacity(capacity);

        IntervalSet {
            limit: None,
            intervals,
        }
    }

    /// Creates an empty `IntervalSet` with a specific limit.
    /// The number of elements in the set cannot exceed this
    /// amount, otherwise `insert` calls will be rejected.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// use core::num::NonZeroUsize;
    /// let mut set = IntervalSet::with_limit(NonZeroUsize::new(1).unwrap());
    /// assert!(set.insert(0..4).is_ok());
    /// assert!(set.insert(12..16).is_err());
    /// assert!(set.insert(4..12).is_ok());
    /// assert!(set.insert(12..16).is_ok());
    /// ```
    #[inline]
    pub fn with_limit(limit: NonZeroUsize) -> IntervalSet<T> {
        let mut set = Self::default();
        set.set_limit(limit);
        set
    }

    /// Sets an element limit for the given `IntervalSet`.
    /// The number of elements in the set cannot exceed this
    /// amount, otherwise `insert` calls will be rejected
    ///
    /// Note: calling this will not drop existing intervals
    /// that exceed the new limit and will only be
    /// applied to later calls to `insert`.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// use core::num::NonZeroUsize;
    /// let mut set = IntervalSet::new();
    /// assert!(set.insert(0..4).is_ok());
    /// set.set_limit(NonZeroUsize::new(1).unwrap());
    /// assert!(set.insert(4..8).is_ok());
    /// assert!(set.insert(12..16).is_err());
    /// ```
    #[inline]
    pub fn set_limit(&mut self, limit: NonZeroUsize) {
        self.limit = Some(limit);
    }

    /// Removes the element limit for the given `IntervalSet`.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// use core::num::NonZeroUsize;
    /// let mut set = IntervalSet::with_limit(NonZeroUsize::new(1).unwrap());
    /// assert!(set.insert(0..4).is_ok());
    /// assert!(set.insert(4..8).is_ok());
    /// assert!(set.insert(12..16).is_err());
    /// set.remove_limit();
    /// assert!(set.insert(12..16).is_ok());
    /// ```
    #[inline]
    pub fn remove_limit(&mut self) {
        self.limit = None;
    }

    /// Returns the number of intervals in `IntervalSet`.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// assert_eq!(set.interval_len(), 0);
    /// assert!(set.insert(0..4).is_ok());
    /// assert_eq!(set.interval_len(), 1);
    /// ```
    #[inline]
    pub fn interval_len(&self) -> usize {
        self.intervals.len()
    }

    /// Returns the allocated capacity of the `IntervalSet`.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::with_capacity(1);
    /// assert_eq!(set.capacity(), 1);
    /// assert!(set.insert(0..4).is_ok());
    /// assert!(set.insert(6..10).is_ok());
    /// assert!(set.capacity() > 1);
    /// ```
    #[inline]
    pub fn capacity(&self) -> usize {
        self.intervals.capacity()
    }

    /// Clears all elements contained in the `IntervalSet`.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// assert!(set.insert(0..4).is_ok());
    /// set.clear();
    /// assert!(set.is_empty());
    /// ```
    #[inline]
    pub fn clear(&mut self) {
        self.intervals.clear()
    }

    /// Removes the lowest `Interval` in the set, if any
    ///
    /// # Examples
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// assert_eq!(set.pop_min(), None);
    /// assert!(set.insert(0..4).is_ok());
    /// assert_eq!(set.pop_min(), Some((0..4).into()));
    /// ```
    #[inline]
    pub fn pop_min(&mut self) -> Option<Interval<T>> {
        self.intervals.pop_front()
    }
}

impl<T: IntervalBound> IntervalSet<T> {
    /// Returns the number of elements in `IntervalSet`.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// assert_eq!(set.count(), 0);
    /// assert!(set.insert(0..4).is_ok());
    /// assert_eq!(set.count(), 4);
    /// ```
    #[inline]
    pub fn count(&self) -> usize {
        self.intervals.iter().map(|interval| interval.len()).sum()
    }

    /// Returns `true` if the `IntervalSet` has no intervals.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// assert!(set.is_empty());
    /// assert!(set.insert(0..4).is_ok());
    /// assert!(!set.is_empty());
    /// ```
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.intervals.is_empty()
    }

    /// Inserts the supplied `interval` into the `IntervalSet`
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// assert!(set.insert(0..4).is_ok());
    /// assert!(set.contains(&3));
    /// assert!(!set.contains(&5));
    /// ```
    #[inline]
    pub fn insert<R: RangeBounds<T>>(&mut self, interval: R) -> Result<(), IntervalSetError> {
        let interval = Interval::from_range_bounds(interval)?;

        if self.intervals.is_empty() {
            self.intervals.push_front(interval);
            return Ok(());
        }

        let index = self.index_for(&interval);
        insert(&mut self.intervals, interval, index, self.limit)?;

        self.check_integrity();

        Ok(())
    }

    /// Inserts the supplied `interval` at the beginning of the `IntervalSet`.
    /// This method can be used to optimize insertion when the `interval` is less
    /// than all of the current intervals.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// assert!(set.insert_front(0..4).is_ok());
    /// assert!(set.contains(&3));
    /// assert!(!set.contains(&5));
    /// ```
    #[inline]
    pub fn insert_front<R: RangeBounds<T>>(&mut self, interval: R) -> Result<(), IntervalSetError> {
        let interval = Interval::from_range_bounds(interval)?;

        if self.intervals.is_empty() {
            self.intervals.push_front(interval);
            return Ok(());
        }

        insert(&mut self.intervals, interval, 0, self.limit)?;

        self.check_integrity();

        Ok(())
    }

    /// Inserts a single `value` into the `IntervalSet`
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// assert!(set.insert_value(1).is_ok());
    /// assert!(set.contains(&1));
    /// assert!(!set.contains(&0));
    /// assert!(!set.contains(&2));
    /// ```
    #[inline]
    pub fn insert_value(&mut self, value: T) -> Result<(), IntervalSetError> {
        self.insert((Bound::Included(value), Bound::Included(value)))
    }

    /// Performs a union, i.e., all the values in `self` or `other` will
    /// be present in `self`
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut a = IntervalSet::new();
    /// assert!(a.insert(0..4).is_ok());
    /// let mut b = IntervalSet::new();
    /// assert!(b.insert(4..8).is_ok());
    /// a.union(&b);
    /// assert_eq!(a.iter().collect::<Vec<_>>(), (0..8).collect::<Vec<_>>());
    /// ```
    #[inline]
    pub fn union(&mut self, other: &Self) -> Result<(), IntervalSetError> {
        if self.intervals.is_empty() {
            self.intervals = other.intervals.clone();
            return Ok(());
        }

        self.set_operation(other, insert)
    }

    /// Removes the supplied `interval` from the `IntervalSet`
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// assert!(set.insert(1..3).is_ok());
    /// assert!(set.remove(0..4).is_ok());
    /// assert!(set.is_empty());
    /// ```
    #[inline]
    pub fn remove<R: RangeBounds<T>>(&mut self, interval: R) -> Result<(), IntervalSetError> {
        let interval = Interval::from_range_bounds(interval)?;

        if self.intervals.is_empty() {
            return Ok(());
        }

        let index = self.index_for(&interval);
        remove(&mut self.intervals, interval, index, self.limit)?;

        self.check_integrity();

        Ok(())
    }

    /// Removes a single `value` from the `IntervalSet`
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// assert!(set.insert_value(1).is_ok());
    /// assert!(set.remove_value(1).is_ok());
    /// assert!(!set.contains(&1));
    /// ```
    #[inline]
    pub fn remove_value(&mut self, value: T) -> Result<(), IntervalSetError> {
        self.remove((Bound::Included(value), Bound::Included(value)))
    }

    /// Performs a difference, i.e., all the values that are in `self` but not
    /// in `other` will be present in `self`.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set_a = IntervalSet::new();
    /// assert!(set_a.insert(0..=10).is_ok());
    /// let mut set_b = IntervalSet::new();
    /// assert!(set_b.insert(4..=8).is_ok());
    /// assert!(set_a.difference(&set_b).is_ok());
    /// assert_eq!(set_a.iter().collect::<Vec<_>>(), vec![0, 1, 2, 3, 9, 10]);
    /// ```
    #[inline]
    pub fn difference(&mut self, other: &Self) -> Result<(), IntervalSetError> {
        if self.intervals.is_empty() {
            return Ok(());
        }

        self.set_operation(other, remove)
    }

    /// Performs an intersection, i.e., all the values in both `self` and `other` will
    /// be present in `self`.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set_a = IntervalSet::new();
    /// assert!(set_a.insert(0..=10).is_ok());
    /// let mut set_b = IntervalSet::new();
    /// assert!(set_b.insert(4..=8).is_ok());
    /// assert!(set_a.intersection(&set_b).is_ok());
    /// assert_eq!(set_a.iter().collect::<Vec<_>>(), vec![4, 5, 6, 7, 8]);
    /// ```
    #[inline]
    pub fn intersection(&mut self, other: &Self) -> Result<(), IntervalSetError> {
        intersection::apply(&mut self.intervals, &other.intervals);

        Ok(())
    }

    /// Returns an iterator of `Intervals` over the intersection, i.e., all
    /// the values in both `self` and `other` will be returned.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set_a = IntervalSet::new();
    /// assert!(set_a.insert(0..=10).is_ok());
    /// let mut set_b = IntervalSet::new();
    /// assert!(set_b.insert(4..=8).is_ok());
    /// let intersection = set_a.intersection_iter(&set_b).flatten();
    /// assert_eq!(intersection.collect::<Vec<_>>(), vec![4, 5, 6, 7, 8]);
    /// ```
    #[inline]
    pub fn intersection_iter<'a>(&'a self, other: &'a Self) -> Intersection<'a, T> {
        Intersection::new(&self.intervals, &other.intervals)
    }

    /// Returns an iterator over all of the values contained in the given `IntervalSet`.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// assert!(set.insert(0..5).is_ok());
    /// assert!(set.insert(10..15).is_ok());
    /// let items: Vec<_> = set.iter().collect();
    /// assert_eq!(vec![0, 1, 2, 3, 4, 10, 11, 12, 13, 14], items);
    /// ```
    #[inline]
    pub fn iter(&self) -> Iter<T> {
        Iter {
            iter: self.intervals.iter(),
            head: None,
            tail: None,
        }
    }

    /// Returns the smallest value in the given `IntervalSet`. If no items
    /// are present in the set, `None` is returned.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// assert_eq!(set.min_value(), None);
    /// assert!(set.insert(0..5).is_ok());
    /// assert_eq!(set.min_value(), Some(0));
    /// ```
    #[inline]
    pub fn min_value(&self) -> Option<T> {
        let interval = self.intervals.front()?;
        Some(interval.start)
    }

    /// Returns the largest value in the given `IntervalSet`. If no items
    /// are present in the set, `None` is returned.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// assert_eq!(set.max_value(), None);
    /// assert!(set.insert(0..5).is_ok());
    /// assert_eq!(set.max_value(), Some(4));
    /// ```
    #[inline]
    pub fn max_value(&self) -> Option<T> {
        let interval = self.intervals.back()?;
        Some(interval.end)
    }

    /// Returns `true` if the set contains a value
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// assert_eq!(set.contains(&1), false);
    /// assert!(set.insert(0..5).is_ok());
    /// assert_eq!(set.contains(&1), true);
    /// ```
    #[inline]
    pub fn contains(&self, value: &T) -> bool {
        self.binary_search_with(value, |_index| true, |_index| false, |_index| false)
    }

    /// Internal function for applying set operations
    #[inline]
    fn set_operation<
        F: Fn(
            &mut VecDeque<Interval<T>>,
            Interval<T>,
            usize,
            Option<NonZeroUsize>,
        ) -> Result<usize, IntervalSetError>,
    >(
        &mut self,
        other: &Self,
        apply: F,
    ) -> Result<(), IntervalSetError> {
        let mut iter = other.intervals.iter();
        let limit = self.limit;

        // get the first interval in `other` and find the applicable
        // index in `self`
        let interval = if let Some(interval) = iter.next() {
            interval
        } else {
            return Ok(());
        };

        let mut index = self.index_for(interval);

        // apply the set operation for the interval above
        index = apply(&mut self.intervals, *interval, index, limit)?;

        // apply the set operation for the rest of the intervals
        for interval in iter {
            index = apply(&mut self.intervals, *interval, index, limit)?;
        }

        self.check_integrity();

        Ok(())
    }

    /// Internal function for locating the optimal starting index for
    /// interval comparison
    #[inline]
    fn index_for(&self, interval: &Interval<T>) -> usize {
        // it's faster just to iterate through the set for smaller lengths
        if self.interval_len() < 16 {
            return 0;
        }

        self.binary_search_with(&interval.start, |index| index, |index| index, |index| index)
    }

    /// Internal function for searching for a value in the contained intervals
    #[inline]
    fn binary_search_with<
        V,
        EqualFn: Fn(usize) -> V,
        GreaterFn: Fn(usize) -> V,
        LessFn: Fn(usize) -> V,
    >(
        &self,
        value: &T,
        on_equal: EqualFn,
        on_greater: GreaterFn,
        on_less: LessFn,
    ) -> V {
        use core::cmp::Ordering::*;

        let intervals = &self.intervals;

        let mut size = intervals.len();
        if size == 0 {
            return on_greater(0);
        }

        let mut base = 0usize;
        while size > 1 {
            let half = size / 2;
            let mid = base + half;
            let subject = &intervals[mid];
            match subject.partial_cmp(value) {
                Some(Equal) => return on_equal(mid),
                Some(Greater) => {}
                Some(Less) => base = mid,
                None => return on_greater(0),
            };
            size -= half;
        }

        let subject = &intervals[base];
        match subject.partial_cmp(value) {
            Some(Equal) => on_equal(base),
            Some(Greater) => on_greater(base),
            Some(Less) => on_less(base),
            None => on_greater(0),
        }
    }

    /// Internal check for integrity - only used when debug_assertions are enabled
    #[inline]
    fn check_integrity(&self) {
        if cfg!(debug_assertions) {
            let mut prev_end: Option<T> = None;

            for interval in self.intervals.iter() {
                // make sure that a few items exist
                for value in (*interval).take(3) {
                    assert!(self.contains(&value), "set should contain value");
                }
                for value in (*interval).rev().take(3) {
                    assert!(self.contains(&value), "set should contain value");
                }

                if let Some(prev_end) = prev_end.as_ref() {
                    assert!(
                        *prev_end < interval.start_exclusive(),
                        "the previous end should be less than the next start",
                    );
                }

                assert!(interval.is_valid(), "interval should be valid");

                prev_end = Some(interval.end);
            }
        }
    }
}

impl<T: Copy> IntervalSet<T> {
    /// Returns an iterator of `Interval`s contained in the `IntervalSet`
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// set.insert(0..=10);
    /// assert_eq!(set.intervals().collect::<Vec<_>>(), vec![0..=10]);
    /// ```
    #[inline]
    pub fn intervals(&self) -> IntervalIter<T> {
        IntervalIter {
            iter: self.intervals.iter(),
        }
    }

    /// Returns an iterator of `Range`s contained in the `IntervalSet`
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// set.insert(0..=10);
    /// assert_eq!(set.ranges().collect::<Vec<_>>(), vec![0..11]);
    /// ```
    #[inline]
    pub fn ranges(&self) -> RangeIter<T> {
        RangeIter {
            iter: self.intervals.iter(),
        }
    }

    /// Returns an iterator of `RangeInclusive`s contained in the `IntervalSet`
    ///
    /// # Examples
    ///
    /// ```ignore
    /// # use s2n_quic_transport::interval_set::IntervalSet;
    /// let mut set = IntervalSet::new();
    /// set.insert(0..=10);
    /// assert_eq!(set.inclusive_ranges().collect::<Vec<_>>(), vec![0..=10]);
    /// ```
    #[inline]
    pub fn inclusive_ranges(&self) -> RangeInclusiveIter<T> {
        RangeInclusiveIter {
            iter: self.intervals.iter(),
        }
    }
}

impl<T: Copy + fmt::Debug> fmt::Debug for IntervalSet<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_set().entries(self.intervals.iter()).finish()
    }
}

/// Iterator over all of the values contained in an `IntervalSet`
pub struct Iter<'a, T> {
    iter: vec_deque::Iter<'a, Interval<T>>,
    head: Option<Interval<T>>,
    tail: Option<Interval<T>>,
}

impl<'a, T: IntervalBound> Iterator for Iter<'a, T> {
    type Item = T;

    #[inline]
    fn next(&mut self) -> Option<T> {
        loop {
            if let Some(item) = self.head.as_mut().and_then(Iterator::next) {
                return Some(item);
            }

            let item = self.iter.next().cloned().or_else(|| self.tail.take())?;
            self.head = Some(item);
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let (lower, _) = self.iter.size_hint();
        // Computing the upper length would require iterating through all of the ranges
        let upper = None;
        (lower, upper)
    }
}

impl<'a, T: IntervalBound> DoubleEndedIterator for Iter<'a, T> {
    #[inline]
    fn next_back(&mut self) -> Option<T> {
        loop {
            if let Some(item) = self.tail.as_mut().and_then(DoubleEndedIterator::next_back) {
                return Some(item);
            }

            let item = self
                .iter
                .next_back()
                .cloned()
                .or_else(|| self.head.take())?;
            self.tail = Some(item);
        }
    }
}

macro_rules! impl_iterator_conversions {
    ($item:ident, $iter:ident) => {
        #[derive(Clone, Debug)]
        pub struct $iter<'a, T> {
            iter: vec_deque::Iter<'a, Interval<T>>,
        }

        impl<'a, T: IntervalBound> Iterator for $iter<'a, T> {
            type Item = $item<T>;

            #[inline]
            fn next(&mut self) -> Option<Self::Item> {
                self.iter.next().map(|interval| interval.into())
            }

            #[inline]
            fn size_hint(&self) -> (usize, Option<usize>) {
                self.iter.size_hint()
            }
        }

        impl<'a, T: IntervalBound> DoubleEndedIterator for $iter<'a, T> {
            #[inline]
            fn next_back(&mut self) -> Option<Self::Item> {
                self.iter.next_back().map(|interval| interval.into())
            }
        }

        impl<'a, T: IntervalBound> ExactSizeIterator for $iter<'a, T> where
            vec_deque::Iter<'a, Interval<T>>: ExactSizeIterator
        {
        }

        impl<T: IntervalBound> FromIterator<$item<T>> for IntervalSet<T> {
            #[inline]
            fn from_iter<I: IntoIterator<Item = $item<T>>>(intervals: I) -> Self {
                let intervals = intervals.into_iter();
                let mut set = Self::with_capacity(intervals.size_hint().0);
                set.extend(intervals);
                set
            }
        }

        impl<'a, T: 'a + IntervalBound> FromIterator<&'a $item<T>> for IntervalSet<T> {
            #[inline]
            fn from_iter<I: IntoIterator<Item = &'a $item<T>>>(intervals: I) -> Self {
                let intervals = intervals.into_iter();
                let mut set = Self::with_capacity(intervals.size_hint().0);
                set.extend(intervals);
                set
            }
        }

        impl<T: IntervalBound> Extend<$item<T>> for IntervalSet<T> {
            #[inline]
            fn extend<I: IntoIterator<Item = $item<T>>>(&mut self, intervals: I) {
                for interval in intervals {
                    if self.insert(interval).is_err() {
                        return;
                    }
                }
            }
        }

        impl<'a, T: 'a + IntervalBound> Extend<&'a $item<T>> for IntervalSet<T> {
            #[inline]
            fn extend<I: IntoIterator<Item = &'a $item<T>>>(&mut self, intervals: I) {
                for interval in intervals {
                    let interval: Interval<T> = interval.into();
                    if self.insert(interval).is_err() {
                        return;
                    }
                }
            }
        }

        impl<T: IntervalBound> From<$item<T>> for IntervalSet<T> {
            #[inline]
            fn from(interval: $item<T>) -> Self {
                let mut set = Self::with_capacity(1);
                let _ = set.insert(interval);
                set
            }
        }
    };
}

impl_iterator_conversions!(Interval, IntervalIter);
impl_iterator_conversions!(Range, RangeIter);
impl_iterator_conversions!(RangeInclusive, RangeInclusiveIter);

impl<T: IntervalBound> FromIterator<T> for IntervalSet<T> {
    #[inline]
    fn from_iter<I: IntoIterator<Item = T>>(values: I) -> Self {
        let values = values.into_iter();
        let mut set = Self::with_capacity(values.size_hint().0);
        for value in values {
            if set.insert_value(value).is_err() {
                break;
            }
        }
        set
    }
}

impl<'a, T: 'a + IntervalBound> FromIterator<&'a T> for IntervalSet<T> {
    #[inline]
    fn from_iter<I: IntoIterator<Item = &'a T>>(values: I) -> Self {
        values.into_iter().cloned().collect()
    }
}


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

use crate::interval_set::{Interval, IntervalBound, IntervalSetError};
use alloc::collections::VecDeque;
use core::{
    cmp::{max, min, Ordering},
    num::NonZeroUsize,
    ops::Range,
};

#[inline]
pub(crate) fn remove<T: IntervalBound + Ord>(
    ranges: &mut VecDeque<Interval<T>>,
    range: Interval<T>,
    start_index: usize,
    limit: Option<NonZeroUsize>,
) -> Result<usize, IntervalSetError> {
    // this range is intentionally invalid and will only be
    // valid if the `scan` method finds a match
    #[allow(clippy::reversed_empty_ranges)]
    let replace_range = core::usize::MAX..0;

    let can_push_range = limit.map(|l| l.get() > ranges.len() + 1).unwrap_or(true);

    let mut removal = Removal {
        replace_range,
        push_range: None,
        can_push_range,
    };

    let iter = ranges.iter_mut().enumerate().skip(start_index);

    if let Some(index) = removal.scan(iter, &range) {
        return Ok(index);
    }

    removal.apply(ranges)
}

/// A structure to keep temporary state for a removal
#[derive(Debug)]
struct Removal<T> {
    replace_range: Range<usize>,
    push_range: Option<Interval<T>>,
    can_push_range: bool,
}

impl<'a, T: 'a + IntervalBound + Ord> Removal<T> {
    /// Scans over the Intervals and updates the `Removal` state with the
    /// applicable interval ranges
    ///
    /// Returns `Some(index)` if the removal can be applied to a single interval
    #[inline]
    fn scan<I: Iterator<Item = (usize, &'a mut Interval<T>)>>(
        &mut self,
        ranges: I,
        range_a: &Interval<T>,
    ) -> Option<usize> {
        use Ordering::*;

        for (slot_index, range_b) in ranges {
            match (
                range_a.start.cmp(&range_b.start),
                range_a.end.cmp(&range_b.end),
            ) {
                // range A is a subset of range B
                //
                // Before:
                //
                // range A: |-----|
                // range B: |---------|
                //
                // After:
                //
                // range A: |-----|
                // range B:       |---|
                //
                (Equal, Less) => {
                    range_b.start = range_a.end_exclusive();
                    return Some(slot_index);
                }

                // range A is a subset of range B
                //
                // Before:
                //
                // range A:     |-----|
                // range B: |---------|
                //
                // After:
                //
                // range A:     |-----|
                // range B: |---|
                //
                (Greater, Equal) => {
                    range_b.end = range_a.start_exclusive();
                    return Some(slot_index + 1);
                }

                // range A is a subset of range B
                //
                // Before:
                //
                // range A:    |----|
                // range B: |---------|
                //
                // After:
                //
                // range A:    |----|
                // range B: |--|    |-|
                //
                (Greater, Less) => {
                    self.push_range(range_a.end_exclusive(), range_b.end)?;
                    range_b.end = range_a.start_exclusive();
                    let next_slot = slot_index + 1;
                    self.set_start(next_slot);
                    self.set_end(next_slot);
                    break;
                }

                // range A is part of range B.
                //
                // Before:
                //
                // range A:      |--------|
                // range B: |--------|
                //
                // After:
                //
                // range A:      |--------|
                // range B: |----|
                //
                (Greater, Greater) if range_a.should_coalesce(range_b) => {
                    range_b.end = range_a.start_exclusive();
                    let next_slot = slot_index + 1;
                    self.set_start(next_slot);
                    continue;
                }

                // range A comes later
                //
                // range A:          |-----|
                // range B: |----|
                //
                (Greater, Greater) => {
                    continue;
                }

                // range A contains range B, spilling over into the next slot
                //
                // Before:
                //
                // range A: |---------|
                // range B: |---|
                //
                // After:
                //
                // range A: |---------|
                // range B: |
                //
                (Equal, Greater) |

                // range A contains range B, spilling over into the next slot
                //
                // range A: |---------|
                // range B:    |---|
                //
                // mark B as obsolete
                //
                (Less, Greater) => {
                    self.set_start(slot_index);
                    let next_slot = slot_index + 1;
                    self.set_end(next_slot);
                    continue;
                }

                // the ranges are equal
                //
                // range A: |---------|
                // range B: |---------|
                //
                (Equal, Equal) |

                // range A ends with range B
                //
                // range A: |---------|
                // range B:       |---|
                //
                // mark B as obsolete
                //
                (Less, Equal) => {
                    self.set_start(slot_index);
                    let next_slot = slot_index + 1;
                    self.set_end(next_slot);
                    break;
                }

                // range A overlaps part of range B
                //
                // Before:
                //
                // range A: |--------|
                // range B:      |--------|
                //
                // After:
                //
                // range A: |--------|
                // range B:          |----|
                //
                (Less, Less) if range_b.should_coalesce(range_a)  => {
                    range_b.start = range_a.end_exclusive();
                    self.set_end(slot_index);
                    break;
                }

                // range A comes before range B
                //
                // range A: |---|
                // range B:        |--------|
                //
                // returns; no more searching needed
                //
                (Less, Less) => {
                    break;
                }
            }
        }

        None
    }

    /// Applies the `Removal` to the given set of `Interval`s.
    #[inline]
    fn apply(self, ranges: &mut VecDeque<Interval<T>>) -> Result<usize, IntervalSetError> {
        let replace_range = self.replace_range;

        let index = replace_range.start;

        if let Some(interval) = self.push_range {
            if self.can_push_range {
                ranges.insert(index, interval);
                return Ok(index);
            } else {
                return Err(IntervalSetError::LimitExceeded);
            }
        }

        match replace_range.end.checked_sub(index) {
            None => Ok(0),
            Some(0) => Ok(index),
            Some(1) => {
                ranges.remove(index);
                Ok(index)
            }
            Some(_) => {
                ranges.drain(replace_range);
                Ok(index)
            }
        }
    }

    #[inline]
    fn set_start(&mut self, start: usize) {
        self.replace_range.start = min(self.replace_range.start, start);
    }

    #[inline]
    fn set_end(&mut self, end: usize) {
        self.replace_range.end = max(self.replace_range.end, end);
    }

    #[inline]
    fn push_range(&mut self, start: T, end: T) -> Option<()> {
        debug_assert!(self.push_range.is_none());
        self.push_range = Some(Interval { start, end });

        if self.can_push_range {
            Some(())
        } else {
            None
        }
    }
}


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

use crate::interval_set::IntervalSet;
use alloc::collections::BTreeSet;
use bolero::{check, generator::*};
use core::{
    num::NonZeroUsize,
    ops::{Range, RangeInclusive},
};

type RangeBound = u8;

#[derive(Clone, Debug, TypeGenerator)]
pub enum RangeValue {
    Range(Range<RangeBound>),
    RangeInclusive(RangeInclusive<RangeBound>),
}

#[derive(Clone, Debug, TypeGenerator)]
pub enum Operation {
    Insert { range: RangeValue },
    Remove { range: RangeValue },
}

#[derive(Clone, Debug, TypeGenerator)]
pub struct OperationTest {
    operations: Vec<Operation>,
    limit: Option<usize>,
}

macro_rules! assert_set_eq {
    ($expected:ident, $actual:ident) => {{
        // put all of the contained values into a vec for easy comparison
        let actual_items = $actual.iter().collect::<Vec<_>>();

        // make sure the contained values are exactly the same
        assert_eq!(
            $expected.iter().collect::<Vec<_>>(),
            actual_items,
            "values don't match\n expected:\t{:?}\n actual:\t{:?}",
            $expected,
            $actual
        );

        // test that the reverse iterator works as well
        let mut reversed_items = $actual.iter().rev().collect::<Vec<_>>();
        reversed_items.reverse();
        assert_eq!(
            actual_items, reversed_items,
            "reverse iterator isn't valid {:?}",
            reversed_items
        );

        // make sure the item counts are correct
        assert_eq!(
            $expected.count(),
            $actual.count(),
            "element lengths don't match\n expected:\t{:?}\n actual:\t{:?}",
            $expected,
            $actual
        );

        // make sure the number of intervals is the same
        assert_eq!(
            $expected.interval_len(),
            $actual.interval_len(),
            "interval lengths don't match\n expected:\t{:?}\n actual:\t{:?}",
            $expected,
            $actual
        );
    }};
}

#[test]
#[cfg_attr(miri, ignore)] // This test is too expensive for miri to complete in a reasonable amount of time
fn interval_set_test() {
    check!().with_type().for_each(
        |(initial_ops, union_ops, difference_ops, intersection_ops)| {
            // creates an initial set given a Vec of insertions and deletions
            let (oracle, subject) = process_operation(initial_ops);

            /// performs a specific set operation to both implementations
            /// and asserts equivalence
            macro_rules! set_op_test {
                ($name:ident, $ops:ident) => {{
                    let (mut expected, mut actual) = process_operation($ops);
                    expected.$name(&oracle);
                    actual.remove_limit();
                    actual.$name(&subject).expect("invalid op");
                    assert_set_eq!(expected, actual);
                }};
            }

            set_op_test!(union, union_ops);
            set_op_test!(difference, difference_ops);

            // Assert that both `intersection` and `intersection_iter`
            // produce the same results as the oracle
            let (mut expected, mut actual) = process_operation(intersection_ops);
            expected.intersection(&oracle);
            actual.remove_limit();
            let actual_iter = IntervalSet::from_iter(actual.intersection_iter(&subject));
            actual.intersection(&subject).expect("invalid op");
            assert_set_eq!(expected, actual);
            assert_set_eq!(expected, actual_iter);
        },
    );
}

#[test]
#[cfg_attr(kani, kani::proof, kani::unwind(2), kani::solver(kissat))]
#[cfg_attr(miri, ignore)] // This test is too expensive for miri to complete in a reasonable amount of time
fn interval_set_inset_range_test() {
    // Generate valid ranges (lb <= ub)
    let gen = gen::<(i32, i32, i32)>().filter_gen(|(a, b, _c)| a <= b);

    check!().with_generator(gen).for_each(|(lb, ub, elem)| {
        let mut set: IntervalSet<i32> = IntervalSet::new();
        let range = lb..=ub;
        assert!(set.insert(range.clone()).is_ok());
        assert_eq!(range.contains(&elem), set.contains(elem));
    });
}

fn process_operation(
    OperationTest { limit, operations }: &OperationTest,
) -> (Oracle, IntervalSet<RangeBound>) {
    let limit = limit.and_then(NonZeroUsize::new);
    let mut oracle = Oracle::new();
    let mut subject = if let Some(limit) = limit {
        IntervalSet::with_limit(limit)
    } else {
        IntervalSet::new()
    };

    macro_rules! perform_operation {
        ($operation:ident, $range:expr) => {{
            match $range {
                RangeValue::Range(range) => {
                    // if the operation didn't exceed the limit then apply it
                    // to the oracle
                    if subject.$operation(range.clone()).is_ok() {
                        oracle.$operation(range.clone());
                        subject.$operation(range.clone()).unwrap();
                    }
                }
                RangeValue::RangeInclusive(range) => {
                    // if the operation didn't exceed the limit then apply it
                    // to the oracle
                    if subject.$operation(range.clone()).is_ok() {
                        oracle.$operation(range.clone());
                        subject.$operation(range.clone()).unwrap();
                    }
                }
            }
        }};
    }

    for operation in operations {
        match operation {
            Operation::Insert { range } => {
                perform_operation!(insert, range);
            }
            Operation::Remove { range } => {
                perform_operation!(remove, range);
            }
        }

        if let Some(limit) = limit {
            assert!(
                subject.interval_len() <= limit.get(),
                "lhs: {}, rhs: {}",
                subject.interval_len(),
                limit
            );
        }
    }

    assert_set_eq!(oracle, subject);

    (oracle, subject)
}

/// `Oracle` is modeled as a `BTreeSet`. Instead of storing `Interval`s,
/// each number contained inside the `Interval` is inserted into
/// the set. This is semantically equivalent, but much less efficient.
/// We can assert that given the same operations, the two implementations
/// should result in the same values.
#[derive(Debug)]
struct Oracle {
    data: BTreeSet<RangeBound>,
}

impl Oracle {
    fn new() -> Self {
        Oracle {
            data: Default::default(),
        }
    }

    fn insert<V: Iterator<Item = RangeBound>>(&mut self, values: V) {
        for value in values {
            self.data.insert(value);
        }
    }

    fn remove<V: Iterator<Item = RangeBound>>(&mut self, values: V) {
        for value in values {
            self.data.remove(&value);
        }
    }

    fn iter(&self) -> impl Iterator<Item = RangeBound> + '_ {
        self.data.iter().cloned()
    }

    fn difference(&mut self, other: &Self) {
        self.data = self.data.difference(&other.data).cloned().collect();
    }

    fn union(&mut self, other: &Self) {
        self.data = self.data.union(&other.data).cloned().collect();
    }

    fn intersection(&mut self, other: &Self) {
        self.data = self.data.intersection(&other.data).cloned().collect();
    }

    fn count(&self) -> usize {
        self.data.len()
    }

    fn interval_len(&self) -> usize {
        if self.data.is_empty() {
            return 0;
        }

        let mut iter = self.iter();
        let mut prev = iter.next().unwrap();
        let mut count = 1;

        for item in iter {
            if item != prev + 1 {
                count += 1;
            }
            prev = item;
        }

        count
    }
}


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

use crate::{
    endpoint::Endpoint,
    event::{self, EndpointPublisher},
    io::{rx::Rx, tx::Tx},
    time::clock::{ClockWithTimer, Timer},
};
use core::pin::Pin;

pub mod select;
use select::Select;

pub struct EventLoop<E, C, R, T> {
    pub endpoint: E,
    pub clock: C,
    pub rx: R,
    pub tx: T,
}

impl<E, C, R, T> EventLoop<E, C, R, T>
where
    E: Endpoint,
    C: ClockWithTimer,
    R: Rx<PathHandle = E::PathHandle>,
    T: Tx<PathHandle = E::PathHandle>,
{
    /// Starts running the endpoint event loop in an async task
    pub async fn start(self) {
        let Self {
            mut endpoint,
            clock,
            mut rx,
            mut tx,
        } = self;

        /// Creates a event publisher with the endpoint's subscriber
        macro_rules! publisher {
            ($timestamp:expr) => {{
                let timestamp = $timestamp;
                let subscriber = endpoint.subscriber();
                event::EndpointPublisherSubscriber::new(
                    event::builder::EndpointMeta {
                        endpoint_type: E::ENDPOINT_TYPE,
                        timestamp,
                    },
                    None,
                    subscriber,
                )
            }};
        }

        let mut timer = clock.timer();

        loop {
            // Poll for RX readiness
            let rx_ready = rx.ready();

            // Poll for TX readiness
            let tx_ready = tx.ready();

            // Poll for any application-driven updates
            let mut wakeups = endpoint.wakeups(&clock);

            // TODO use the [pin macro](https://doc.rust-lang.org/std/pin/macro.pin.html) once
            // available in MSRV
            //
            // See https://github.com/aws/s2n-quic/issues/1751
            let wakeups = unsafe {
                // Safety: the wakeups future is on the stack and won't move
                Pin::new_unchecked(&mut wakeups)
            };

            // Poll for timer expiration
            let timer_ready = timer.ready();

            // Concurrently poll all of the futures and wake up on the first one that's ready
            let select = Select::new(rx_ready, tx_ready, wakeups, timer_ready);

            let select::Outcome {
                rx_result,
                tx_result,
                timeout_expired,
                application_wakeup,
            } = if let Ok(outcome) = select.await {
                outcome
            } else {
                // The endpoint has shut down; stop the event loop
                return;
            };

            // notify the application that we woke up and why
            let wakeup_timestamp = clock.get_time();
            publisher!(wakeup_timestamp).on_platform_event_loop_wakeup(
                event::builder::PlatformEventLoopWakeup {
                    timeout_expired,
                    rx_ready: rx_result.is_some(),
                    tx_ready: tx_result.is_some(),
                    application_wakeup,
                },
            );

            match rx_result {
                Some(Ok(())) => {
                    // we received some packets. give them to the endpoint.
                    rx.queue(|queue| {
                        endpoint.receive(queue, &clock);
                    });
                }
                Some(Err(error)) => {
                    // The RX provider has encountered an error. shut down the event loop
                    let mut publisher = publisher!(clock.get_time());
                    rx.handle_error(error, &mut publisher);
                    return;
                }
                None => {
                    // We didn't receive any packets; nothing to do
                }
            }

            match tx_result {
                Some(Ok(())) => {
                    // The TX queue was full and now has capacity. The endpoint can now continue to
                    // transmit
                }
                Some(Err(error)) => {
                    // The RX provider has encountered an error. shut down the event loop
                    let mut publisher = publisher!(clock.get_time());
                    tx.handle_error(error, &mut publisher);
                    return;
                }
                None => {
                    // The TX queue is either waiting to be flushed or has capacity. Either way, we
                    // call `endpoint.transmit` to at least update the clock and poll any timer
                    // expirations.
                }
            }

            // Let the endpoint transmit, if possible
            tx.queue(|queue| {
                endpoint.transmit(queue, &clock);
            });

            // Get the next expiration from the endpoint and update the timer
            let timeout = endpoint.timeout();
            if let Some(timeout) = timeout {
                timer.update(timeout);
            }

            let sleep_timestamp = clock.get_time();
            // compute the relative timeout to the current time
            let timeout = timeout.map(|t| t.saturating_duration_since(sleep_timestamp));
            // compute how long it took to process the current iteration
            let processing_duration = sleep_timestamp.saturating_duration_since(wakeup_timestamp);

            // publish the event to the application
            publisher!(sleep_timestamp).on_platform_event_loop_sleep(
                event::builder::PlatformEventLoopSleep {
                    timeout,
                    processing_duration,
                },
            );
        }
    }
}


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

pub mod event_loop;
pub mod rx;
pub mod tx;

#[cfg(test)]
mod testing;


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

use crate::{event, inet::datagram, path};
use core::task::{Context, Poll};

pub mod pair;

/// Handle to a receive IO provider
pub trait Rx: Sized {
    type PathHandle;
    // TODO make this generic over lifetime
    // See https://github.com/aws/s2n-quic/issues/1742
    type Queue: Queue<Handle = Self::PathHandle>;
    type Error;

    /// Returns a future that yields after a packet is ready to be received
    #[inline]
    fn ready(&mut self) -> RxReady<Self> {
        RxReady(self)
    }

    /// Polls the IO provider for a packet that is ready to be received
    fn poll_ready(&mut self, cx: &mut Context) -> Poll<Result<(), Self::Error>>;

    /// Calls the provided callback with the IO provider queue
    fn queue<F: FnOnce(&mut Self::Queue)>(&mut self, f: F);

    /// Handles the queue error and potentially publishes an event
    fn handle_error<E: event::EndpointPublisher>(self, error: Self::Error, event: &mut E);
}

impl_ready_future!(Rx, RxReady, Result<(), T::Error>);

/// Extension traits for Rx channels
pub trait RxExt: Rx {
    /// Creates single channel from a pair of channels
    #[inline]
    fn with_pair<Other>(self, other: Other) -> pair::Channel<Self, Other>
    where
        Other: Rx<PathHandle = Self::PathHandle>,
    {
        pair::Channel { a: self, b: other }
    }
}

/// Implement the extension traits for all Rx queues
impl<R: Rx> RxExt for R {}

/// A structure capable of queueing and receiving messages
pub trait Queue {
    type Handle: path::Handle;

    /// Iterates over all of the packets in the receive queue and processes them
    fn for_each<F: FnMut(datagram::Header<Self::Handle>, &mut [u8])>(&mut self, on_packet: F);

    /// Returns if there are items in the queue or not
    fn is_empty(&self) -> bool;
}