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

use crate::{connection, endpoint};
use core::{ops::RangeInclusive, time::Duration};

mod generated;
pub use generated::*;

/// All event types which can be emitted from this library.
pub trait Event: core::fmt::Debug {
    const NAME: &'static str;
}

pub trait IntoEvent<Target> {
    fn into_event(self) -> Target;
}

macro_rules! ident_into_event {
    ($($name:ty),* $(,)?) => {
        $(
            impl IntoEvent<$name> for $name {
                #[inline]
                fn into_event(self) -> Self {
                    self
                }
            }
        )*
    };
}

macro_rules! borrowed_into_event {
    ($($name:ty),* $(,)?) => {
        $(
            impl<'a> IntoEvent<&'a $name> for &'a $name {
                #[inline]
                fn into_event(self) -> Self {
                    self
                }
            }
        )*
    };
}

ident_into_event!(
    u8,
    i8,
    u16,
    i16,
    u32,
    i32,
    u64,
    i64,
    usize,
    isize,
    f32,
    Duration,
    bool,
    connection::Error,
    endpoint::Location,
);
borrowed_into_event!([u8; 4], [u8; 16], [u8], [u32], [&'a [u8]]);

impl<T: IntoEvent<U>, U> IntoEvent<Option<U>> for Option<T> {
    #[inline]
    fn into_event(self) -> Option<U> {
        self.map(IntoEvent::into_event)
    }
}

impl<'a> IntoEvent<&'a str> for &'a str {
    #[inline]
    fn into_event(self) -> Self {
        self
    }
}

impl<T> IntoEvent<RangeInclusive<T>> for RangeInclusive<T> {
    #[inline]
    fn into_event(self) -> RangeInclusive<T> {
        self
    }
}

#[derive(Clone, Debug, Copy)]
pub struct Timestamp(crate::time::Timestamp);

impl Timestamp {
    /// The duration since the start of the s2n-quic process.
    ///
    /// Record the start `SystemTime` at the start of the program
    /// to derive the absolute time at which an event is emitted.
    ///
    /// ```rust
    /// # use s2n_quic_core::{
    /// #    endpoint,
    /// #    event::{self, IntoEvent},
    /// #    time::{Duration, Timestamp},
    /// # };
    ///
    /// let start_time = std::time::SystemTime::now();
    /// // `meta` is included as part of each event
    /// # let meta: event::api::ConnectionMeta = event::builder::ConnectionMeta {
    /// #     endpoint_type: endpoint::Type::Server,
    /// #     id: 0,
    /// #     timestamp: unsafe { Timestamp::from_duration(Duration::from_secs(1) )},
    /// # }.into_event();
    /// let event_time = start_time + meta.timestamp.duration_since_start();
    /// ```
    pub fn duration_since_start(&self) -> Duration {
        // Safety: the duration is relative to start of program. This function along
        // with it's documentation captures this intent.
        unsafe { self.0.as_duration() }
    }

    /// Returns the `Duration` which elapsed since an earlier `Timestamp`.
    /// If `earlier` is more recent, the method returns a `Duration` of 0.
    #[inline]
    pub fn saturating_duration_since(self, earlier: Self) -> Duration {
        self.0.saturating_duration_since(earlier.0)
    }
}

impl IntoEvent<Timestamp> for crate::time::Timestamp {
    #[inline]
    fn into_event(self) -> Timestamp {
        Timestamp(self)
    }
}


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

use crate::{
    frame::{self, new_connection_id::STATELESS_RESET_TOKEN_LEN},
    stream::StreamType,
    varint,
};
use core::{convert::TryInto, ops::Range};
pub use s2n_codec::{Encoder, EncoderBuffer, EncoderValue};

pub trait Random {
    fn fill(&mut self, bytes: &mut [u8]);

    fn gen_range(&mut self, range: Range<u64>) -> u64;

    #[inline]
    fn shuffle(&mut self, bytes: &mut [u8]) {
        if bytes.is_empty() {
            return;
        }

        let len = bytes.len() as u64;

        let count = self.gen_range(0..len);
        for _ in 0..count {
            let from = self.gen_range(0..len);
            let to = self.gen_range(0..len);
            bytes.swap(from as usize, to as usize);
        }
    }

    #[inline]
    fn gen_slice<'a>(&mut self, bytes: &'a mut [u8]) -> &'a mut [u8] {
        if bytes.is_empty() {
            return bytes;
        }

        let len = bytes.len() as u64;
        let len = self.gen_range(0..len);
        let bytes = &mut bytes[..len as usize];
        self.fill(bytes);
        bytes
    }

    #[inline]
    fn gen_bool(&mut self) -> bool {
        self.gen_u8() & 0b1 == 0b1
    }

    #[inline]
    fn gen_u8(&mut self) -> u8 {
        let mut o = [0];
        self.fill(&mut o);
        o[0]
    }

    #[inline]
    fn gen_u32(&mut self) -> u32 {
        let mut o = [0; 4];
        self.fill(&mut o);
        u32::from_le_bytes(o)
    }

    #[inline]
    fn gen_u64(&mut self) -> u64 {
        let mut o = [0; 8];
        self.fill(&mut o);
        u64::from_le_bytes(o)
    }

    #[inline]
    fn gen_varint(&mut self) -> varint::VarInt {
        use varint::VarInt;
        let max = VarInt::MAX.as_u64();
        let v = self.gen_range(0..(max + 1));
        unsafe { VarInt::new_unchecked(v as _) }
    }
}

pub trait Strategy: Sized {
    /// Applies the havoc strategy to the supplied buffer
    fn havoc<R: Random>(&mut self, rand: &mut R, buffer: &mut EncoderBuffer);

    /// Applies the havoc strategy to the supplied buffer slice and returns the new buffer length
    #[inline]
    fn havoc_slice<R: Random>(&mut self, rand: &mut R, buffer: &mut [u8]) -> usize {
        let mut buffer = EncoderBuffer::new(buffer);
        self.havoc(rand, &mut buffer);
        buffer.len()
    }

    /// Alternate between two strategies with the supplied `period`
    fn alternate<B: Strategy>(self, b: B, period: Range<usize>) -> Alternate<Self, B> {
        Alternate::new(self, b, period)
    }

    /// Apply the strategy `count` times
    fn repeat(self, count: Range<usize>) -> Repeat<Self> {
        Repeat::new(self, count)
    }

    /// Randomly apply the strategy
    fn randomly(self) -> Randomly<Self> {
        Randomly { strategy: self }
    }

    /// Toggle the strategy on and off for the supplied `period`
    fn toggle(self, period: Range<usize>) -> Toggle<Self> {
        Toggle::new(self, period)
    }

    /// Applies two strategies in order
    fn and_then<B: Strategy>(self, b: B) -> AndThen<Self, B> {
        AndThen { a: self, b }
    }

    /// Repeatedly applies the strategy as long as the buffer has capacity
    fn while_has_capacity(self) -> WhileHasCapacity<Self> {
        WhileHasCapacity { strategy: self }
    }
}

#[derive(Clone, Copy, Debug, Default)]
pub struct Alternate<A: Strategy, B: Strategy> {
    a: A,
    b: B,
    min: u64,
    max: u64,
    is_a: bool,
    remaining: u64,
}

impl<A: Strategy, B: Strategy> Alternate<A, B> {
    pub fn new(a: A, b: B, range: Range<usize>) -> Self {
        debug_assert_ne!(range.start, 0);
        debug_assert!(range.start <= range.end);
        Self {
            a,
            b,
            min: range.start as _,
            max: range.end as _,
            is_a: false,
            remaining: 0,
        }
    }
}

impl<A: Strategy, B: Strategy> Strategy for Alternate<A, B> {
    #[inline]
    fn havoc<R: Random>(&mut self, rand: &mut R, buffer: &mut EncoderBuffer) {
        loop {
            if let Some(remaining) = self.remaining.checked_sub(1) {
                self.remaining = remaining;

                if self.is_a {
                    self.a.havoc(rand, buffer);
                } else {
                    self.b.havoc(rand, buffer);
                }

                break;
            }

            self.remaining = rand.gen_range(self.min..self.max);
            self.is_a = !self.is_a;
        }
    }
}

#[derive(Clone, Copy, Debug, Default)]
pub struct AndThen<A: Strategy, B: Strategy> {
    pub a: A,
    pub b: B,
}

impl<A: Strategy, B: Strategy> Strategy for AndThen<A, B> {
    #[inline]
    fn havoc<R: Random>(&mut self, rand: &mut R, buffer: &mut EncoderBuffer) {
        self.a.havoc(rand, buffer);
        self.b.havoc(rand, buffer);
    }
}

#[derive(Clone, Copy, Debug, Default)]
pub struct Repeat<S: Strategy> {
    strategy: S,
    min: u64,
    max: u64,
}

impl<S: Strategy> Repeat<S> {
    pub fn new(strategy: S, range: Range<usize>) -> Self {
        Self {
            strategy,
            min: range.start as _,
            max: range.end as _,
        }
    }
}

impl<S: Strategy> Strategy for Repeat<S> {
    #[inline]
    fn havoc<R: Random>(&mut self, rand: &mut R, buffer: &mut EncoderBuffer) {
        let count = rand.gen_range(self.min..self.max);
        for _ in 0..count {
            self.strategy.havoc(rand, buffer);
        }
    }
}

#[derive(Clone, Copy, Debug, Default)]
pub struct WhileHasCapacity<S: Strategy> {
    strategy: S,
}

impl<S: Strategy> WhileHasCapacity<S> {
    pub fn new(strategy: S) -> Self {
        Self { strategy }
    }
}

impl<S: Strategy> Strategy for WhileHasCapacity<S> {
    #[inline]
    fn havoc<R: Random>(&mut self, rand: &mut R, buffer: &mut EncoderBuffer) {
        // limit the number of iterations in case the strategy isn't filling the buffer
        for _ in 0..25 {
            if !buffer.remaining_capacity() == 0 {
                break;
            }

            self.strategy.havoc(rand, buffer);
        }
    }
}

#[derive(Clone, Copy, Debug, Default)]
pub struct Toggle<S: Strategy> {
    alt: Alternate<Disabled, S>,
}

impl<S: Strategy> Toggle<S> {
    pub fn new(strategy: S, range: Range<usize>) -> Self {
        Self {
            alt: Alternate::new(Disabled, strategy, range),
        }
    }
}

impl<S: Strategy> Strategy for Toggle<S> {
    #[inline]
    fn havoc<R: Random>(&mut self, rand: &mut R, buffer: &mut EncoderBuffer) {
        self.alt.havoc(rand, buffer);
    }
}

#[derive(Clone, Copy, Debug, Default)]
pub struct Randomly<S: Strategy> {
    pub strategy: S,
}

impl<S: Strategy> Strategy for Randomly<S> {
    #[inline]
    fn havoc<R: Random>(&mut self, rand: &mut R, buffer: &mut EncoderBuffer) {
        if rand.gen_bool() {
            self.strategy.havoc(rand, buffer);
        }
    }
}

#[derive(Clone, Copy, Debug, Default)]
pub struct Reset;

impl Strategy for Reset {
    #[inline]
    fn havoc<R: Random>(&mut self, _rand: &mut R, buffer: &mut EncoderBuffer) {
        buffer.set_position(0);
    }
}

#[derive(Clone, Copy, Debug, Default)]
pub struct Zero;

impl Strategy for Zero {
    #[inline]
    fn havoc<R: Random>(&mut self, _rand: &mut R, buffer: &mut EncoderBuffer) {
        for byte in buffer.as_mut_slice() {
            *byte = 0;
        }
    }
}

#[derive(Clone, Copy, Debug, Default)]
pub struct Shuffle;

impl Strategy for Shuffle {
    #[inline]
    fn havoc<R: Random>(&mut self, rand: &mut R, buffer: &mut EncoderBuffer) {
        if !buffer.is_empty() {
            rand.shuffle(buffer.as_mut_slice());
        }
    }
}

#[derive(Clone, Copy, Debug, Default)]
pub struct Swap;

impl Strategy for Swap {
    #[inline]
    fn havoc<R: Random>(&mut self, rand: &mut R, buffer: &mut EncoderBuffer) {
        let len = buffer.len() as u64;
        if len > 0 {
            let from = rand.gen_range(0..len) as usize;
            let to = rand.gen_range(0..len) as usize;
            buffer.as_mut_slice().swap(from, to);
        }
    }
}

#[derive(Clone, Copy, Debug, Default)]
pub struct Truncate;

impl Strategy for Truncate {
    #[inline]
    fn havoc<R: Random>(&mut self, rand: &mut R, buffer: &mut EncoderBuffer) {
        let len = buffer.capacity() as u64;
        if len > 0 {
            let new_len = rand.gen_range(0..len) as usize;
            buffer.set_position(new_len);
        }
    }
}

#[derive(Clone, Copy, Debug, Default)]
pub struct Mutate;

impl Strategy for Mutate {
    #[inline]
    fn havoc<R: Random>(&mut self, rand: &mut R, buffer: &mut EncoderBuffer) {
        let len = buffer.len() as u64;
        if len > 0 {
            let index = rand.gen_range(0..len) as usize;
            let value = rand.gen_u8();
            buffer.as_mut_slice()[index] = value;
        }
    }
}

#[derive(Clone, Copy, Debug, Default)]
pub struct Disabled;

impl Strategy for Disabled {
    #[inline]
    fn havoc<R: Random>(&mut self, _rand: &mut R, _buffer: &mut EncoderBuffer) {}
}

#[derive(Clone, Copy, Debug, Default)]
pub struct VarInt {
    min: u64,
    max: u64,
}

impl VarInt {
    pub fn new(range: Range<varint::VarInt>) -> Self {
        let min = range.start.as_u64();
        let max = range.end.as_u64();
        Self { min, max }
    }
}

impl Strategy for VarInt {
    #[inline]
    fn havoc<R: Random>(&mut self, rand: &mut R, buffer: &mut EncoderBuffer) {
        let value = rand.gen_range(self.min..self.max);
        let value: varint::VarInt = value.try_into().unwrap();
        if value.encoding_size() <= buffer.remaining_capacity() {
            buffer.encode(&value);
        }
    }
}

#[derive(Clone, Debug)]
pub struct Frame;

impl Strategy for Frame {
    #[inline]
    fn havoc<R: Random>(&mut self, rand: &mut R, buffer: &mut EncoderBuffer) {
        type GenFrame<R> = for<'a> fn(
            rand: &'a mut R,
            payload: &'a mut [u8],
            remaining_capacity: usize,
        ) -> frame::Frame<'a, AckRanges<'a, R>, &'a mut [u8]>;

        let frames: &[GenFrame<R>] = &[
            |rand, _data, cap| {
                frame::Padding {
                    length: rand.gen_range(1..cap as _) as _,
                }
                .into()
            },
            |_rand, _data, _cap| frame::Ping.into(),
            // TODO ACK
            |rand, _data, _cap| {
                frame::ResetStream {
                    stream_id: rand.gen_varint(),
                    application_error_code: rand.gen_varint(),
                    final_size: rand.gen_varint(),
                }
                .into()
            },
            |rand, _data, _cap| {
                frame::StopSending {
                    stream_id: rand.gen_varint(),
                    application_error_code: rand.gen_varint(),
                }
                .into()
            },
            |rand, data, cap| {
                let data = rand.gen_slice(&mut data[..cap]);
                frame::Crypto {
                    offset: rand.gen_varint(),
                    data,
                }
                .into()
            },
            |rand, data, cap| {
                let token = rand.gen_slice(&mut data[..cap]);
                frame::NewToken { token }.into()
            },
            |rand, data, cap| {
                let data = rand.gen_slice(&mut data[..cap]);
                frame::Stream {
                    stream_id: rand.gen_varint(),
                    offset: rand.gen_varint(),
                    data,
                    is_last_frame: rand.gen_bool(),
                    is_fin: rand.gen_bool(),
                }
                .into()
            },
            |rand, _data, _cap| {
                frame::MaxData {
                    maximum_data: rand.gen_varint(),
                }
                .into()
            },
            |rand, _data, _cap| {
                frame::MaxStreamData {
                    stream_id: rand.gen_varint(),
                    maximum_stream_data: rand.gen_varint(),
                }
                .into()
            },
            |rand, _data, _cap| {
                frame::MaxStreams {
                    stream_type: if rand.gen_bool() {
                        StreamType::Unidirectional
                    } else {
                        StreamType::Bidirectional
                    },
                    maximum_streams: rand.gen_varint(),
                }
                .into()
            },
            |rand, _data, _cap| {
                frame::DataBlocked {
                    data_limit: rand.gen_varint(),
                }
                .into()
            },
            |rand, _data, _cap| {
                frame::StreamDataBlocked {
                    stream_id: rand.gen_varint(),
                    stream_data_limit: rand.gen_varint(),
                }
                .into()
            },
            |rand, _data, _cap| {
                frame::StreamsBlocked {
                    stream_type: if rand.gen_bool() {
                        StreamType::Unidirectional
                    } else {
                        StreamType::Bidirectional
                    },
                    stream_limit: rand.gen_varint(),
                }
                .into()
            },
            |rand, data, cap| {
                let (stateless_reset_token, data) = data.split_at_mut(STATELESS_RESET_TOKEN_LEN);

                rand.fill(stateless_reset_token);
                let stateless_reset_token = (&*stateless_reset_token).try_into().unwrap();

                let connection_id = rand.gen_slice(&mut data[..cap]);

                frame::NewConnectionId {
                    sequence_number: rand.gen_varint(),
                    retire_prior_to: rand.gen_varint(),
                    connection_id,
                    stateless_reset_token,
                }
                .into()
            },
            |rand, _data, _cap| {
                frame::RetireConnectionId {
                    sequence_number: rand.gen_varint(),
                }
                .into()
            },
            |rand, data, _cap| {
                let data = &mut data[..frame::path_challenge::DATA_LEN];
                rand.fill(data);
                let data = (&*data).try_into().unwrap();
                frame::PathChallenge { data }.into()
            },
            |rand, data, _cap| {
                let data = &mut data[..frame::path_challenge::DATA_LEN];
                rand.fill(data);
                let data = (&*data).try_into().unwrap();
                frame::PathResponse { data }.into()
            },
            |rand, data, cap| {
                frame::ConnectionClose {
                    error_code: rand.gen_varint(),
                    frame_type: if rand.gen_bool() {
                        Some(rand.gen_varint())
                    } else {
                        None
                    },
                    reason: if rand.gen_bool() {
                        let reason = rand.gen_slice(&mut data[..cap]);
                        Some(reason)
                    } else {
                        None
                    },
                }
                .into()
            },
            |_rand, _data, _cap| frame::HandshakeDone.into(),
            |rand, data, cap| {
                let data = rand.gen_slice(&mut data[..cap]);
                frame::Datagram {
                    is_last_frame: rand.gen_bool(),
                    data,
                }
                .into()
            },
        ];

        let index = rand.gen_range(0..frames.len() as u64) as usize;
        let mut payload = [0u8; 1500];
        let frame = frames[index](rand, &mut payload, buffer.remaining_capacity());

        if frame.encoding_size() <= buffer.remaining_capacity() {
            buffer.encode(&frame);
        }
    }
}

// TODO implement this without allocating a bunch for the entries
struct AckRanges<'a, R> {
    #[allow(dead_code)]
    rand: &'a mut R,
}

impl<'a, R> frame::ack::AckRanges for AckRanges<'a, R> {
    type Iter = frame::ack::AckRangesIter<'a>;

    fn ack_ranges(&self) -> Self::Iter {
        todo!()
    }

    fn largest_acknowledged(&self) -> crate::varint::VarInt {
        todo!()
    }
}

#[cfg(any(test, feature = "testing"))]
pub mod testing {
    use super::*;

    pub struct RandomSlice<'a>(core::iter::Cycle<core::slice::Iter<'a, u8>>);

    impl<'a> RandomSlice<'a> {
        pub fn new(slice: &'a [u8]) -> Self {
            Self(slice.iter().cycle())
        }
    }

    impl<'a> Random for RandomSlice<'a> {
        #[inline]
        fn fill(&mut self, bytes: &mut [u8]) {
            for byte in bytes.iter_mut() {
                *byte = *self.0.next().unwrap_or(&0);
            }
        }

        #[inline]
        fn gen_range(&mut self, range: Range<u64>) -> u64 {
            let start = range.start.min(range.end);
            let end = range.start.max(range.end);

            // this method is biased and should not be used outside of this test

            let variance = end - start;

            // check to see if they're the same number
            if variance == 0 {
                return start;
            }

            let value = self.gen_u64();
            start + value % variance
        }
    }
}

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

    macro_rules! test {
        ($name:ident, $strategy:expr) => {
            #[test]
            #[cfg_attr(miri, ignore)] // no need to test with miri as there isn't any unsafe
            fn $name() {
                check!().for_each(|bytes| {
                    let mut rand = RandomSlice::new(bytes);
                    let mut buffer = [0u8; 256];
                    let buffer = &mut buffer[..rand.gen_u8() as usize];
                    let mut buffer = EncoderBuffer::new(buffer);

                    let mut strategy = $strategy;

                    for _ in 0..rand.gen_range(1..10) {
                        strategy.havoc(&mut rand, &mut buffer);
                    }
                });
            }
        };
    }

    test!(disabled_test, Disabled);
    test!(truncate_test, Truncate);
    test!(reset_test, Reset);
    test!(zero_test, Zero);
    test!(swap_test, Swap);
    test!(shuffle_test, Shuffle);
    test!(mutate_test, Mutate);
    test!(alternate_test, Disabled.alternate(Disabled, 1..5));
    test!(and_then_test, Disabled.and_then(Disabled));
    test!(repeat_test, Disabled.repeat(0..5));
    test!(varint_test, VarInt::new(0u8.into()..42u8.into()));
    test!(frame_test, Frame);
    test!(while_has_capacity_test, Frame.while_has_capacity());
    test!(toggle_test, Disabled.toggle(1..5));
    test!(randomly_test, Disabled.randomly());
}


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

#![cfg_attr(not(any(test, feature = "std")), no_std)]

#[cfg(feature = "alloc")]
extern crate alloc;

/// Asserts that a boolean expression is true at runtime, only if debug_assertions are enabled.
///
/// Otherwise, the compiler is told to assume that the expression is always true and can perform
/// additional optimizations.
///
/// # Safety
///
/// The caller _must_ ensure this condition is never possible, otherwise the compiler
/// may optimize based on false assumptions and behave incorrectly.
#[macro_export]
macro_rules! assume {
    ($cond:expr) => {
        $crate::assume!($cond, "assumption failed: {}", stringify!($cond));
    };
    ($cond:expr $(, $fmtarg:expr)* $(,)?) => {
        let v = $cond;

        debug_assert!(v $(, $fmtarg)*);
        if cfg!(not(debug_assertions)) && !v {
            core::hint::unreachable_unchecked();
        }
    };
}

/// Implements a future that wraps `T::poll_ready` and yields after ready
macro_rules! impl_ready_future {
    ($name:ident, $fut:ident, $output:ty) => {
        pub struct $fut<'a, T>(&'a mut T);

        impl<'a, T: $name> core::future::Future for $fut<'a, T> {
            type Output = $output;

            #[inline]
            fn poll(
                mut self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context,
            ) -> core::task::Poll<Self::Output> {
                self.0.poll_ready(cx)
            }
        }
    };
}

pub mod ack;
pub mod application;
#[cfg(feature = "alloc")]
pub mod buffer;
pub mod connection;
pub mod counter;
pub mod crypto;
pub mod ct;
pub mod datagram;
pub mod endpoint;
pub mod event;
pub mod frame;
pub mod havoc;
pub mod inet;
#[cfg(feature = "alloc")]
pub mod interval_set;
pub mod io;
pub mod memo;
pub mod number;
pub mod packet;
pub mod path;
pub mod query;
pub mod random;
pub mod recovery;
pub mod slice;
pub mod stateless_reset;
pub mod stream;
pub mod sync;
pub mod time;
pub mod token;
pub mod transmission;
pub mod transport;
pub mod varint;
pub mod xdp;

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


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

use core::{cell::Cell, fmt};

/// A datastructure that [memoizes](https://wikipedia.org/wiki/Memoization) a query function
///
/// This can be used for when queries rarely change and can potentially be expensive or on hot
/// code paths. After the `input` is mutated, the query value should be `clear`ed to signal that
/// the function needs to be executed again.
///
/// In debug mode the `get` call will always run the query and assert that the values match.
#[derive(Clone)]
pub struct Memo<T: Copy, Input, Check = DefaultConsistencyCheck> {
    value: Cell<Option<T>>,
    query: fn(&Input) -> T,
    check: Check,
}

impl<T: Copy + fmt::Debug, Input, Check> fmt::Debug for Memo<T, Input, Check> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_tuple("Memo").field(&self.value.get()).finish()
    }
}

impl<T: Copy + PartialEq + fmt::Debug, Input, Check: ConsistencyCheck> Memo<T, Input, Check> {
    /// Creates a new `Memo` over a query function
    #[inline]
    pub fn new(query: fn(&Input) -> T) -> Self {
        Self {
            value: Cell::new(None),
            query,
            check: Check::default(),
        }
    }

    /// Returns the current value of the query function, which may be cached
    #[inline]
    #[track_caller]
    pub fn get(&self, input: &Input) -> T {
        if let Some(value) = self.value.get() {
            // make sure the values match
            self.check.check_consistency(value, input, self.query);
            return value;
        }

        let value = (self.query)(input);
        self.value.set(Some(value));
        value
    }

    /// Clears the cached value of the query function
    #[inline]
    pub fn clear(&self) {
        self.value.set(None);
    }

    /// Asserts that the cached value reflects the current query result in debug mode
    #[inline]
    #[track_caller]
    pub fn check_consistency(&self, input: &Input) {
        if cfg!(debug_assertions) {
            // `get` will assert the value matches the query internally
            let _ = self.get(input);
        }
    }
}

/// Trait to configure consistency checking behavior
pub trait ConsistencyCheck: Clone + Default {
    /// Called when the `Memo` struct has a cached value
    ///
    /// An implementation can assert that the `cache` value matches the current `query` result
    fn check_consistency<T: PartialEq + fmt::Debug, Input>(
        &self,
        cache: T,
        input: &Input,
        query: fn(&Input) -> T,
    );
}

#[derive(Copy, Clone, Default)]
pub struct ConsistencyCheckAlways;

impl ConsistencyCheck for ConsistencyCheckAlways {
    #[inline]
    fn check_consistency<T: PartialEq + fmt::Debug, Input>(
        &self,
        actual: T,
        input: &Input,
        query: fn(&Input) -> T,
    ) {
        let expected = query(input);
        assert_eq!(expected, actual);
    }
}

#[derive(Copy, Clone, Default)]
pub struct ConsistencyCheckNever;

impl ConsistencyCheck for ConsistencyCheckNever {
    #[inline]
    fn check_consistency<T: PartialEq + fmt::Debug, Input>(
        &self,
        _cache: T,
        _input: &Input,
        _query: fn(&Input) -> T,
    ) {
        // noop
    }
}

#[cfg(debug_assertions)]
pub type DefaultConsistencyCheck = ConsistencyCheckAlways;
#[cfg(not(debug_assertions))]
pub type DefaultConsistencyCheck = ConsistencyCheckNever;

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

    #[derive(Debug, Default)]
    struct Input<Value> {
        value: Value,
        should_query: bool,
    }

    #[test]
    fn memo_test() {
        let memo = Memo::<u64, Input<_>, ConsistencyCheckNever>::new(|input| {
            assert!(
                input.should_query,
                "query was called when it wasn't expected"
            );
            input.value
        });

        assert_eq!(
            memo.get(&Input {
                value: 1,
                should_query: true,
            }),
            1
        );

        assert_eq!(
            memo.get(&Input {
                value: 2,
                should_query: false,
            }),
            1
        );

        memo.clear();

        assert_eq!(
            memo.get(&Input {
                value: 3,
                should_query: true,
            }),
            3
        );

        assert_eq!(
            memo.get(&Input {
                value: 4,
                should_query: false,
            }),
            3
        );
    }
}


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

/// Generates a generic number operation
macro_rules! number_trait {
    ($name:ident, $method:ident, $output:ty, $test:stmt) => {
        number_trait!($name, $method, $output, $test, [
            u8, i8, u16, i16, u32, i32, u64, i64, usize, isize, u128, i128
        ]);
    };
    ($name:ident, $method:ident, $output:ty, $test:stmt, [$($ty:ident),*]) => {
        pub trait $name<Rhs = Self>: Sized {
            type Output;
            fn $method(self, rhs: Rhs) -> $output;
        }

        $(
           impl<Rhs> $name<Rhs> for $ty
                where $ty: UpcastFrom<Rhs>
           {
               type Output = Self;

               fn $method(self, rhs: Rhs) -> $output {
                   $ty::$method(self, rhs.upcast())
               }
           }
        )*

        #[test]
        fn $method() {
            $({
                type Type = $ty;
                $test
            })*
        }
    };
}

/// Generates a generic assign operation trait
macro_rules! assign_trait {
    ($name:ident, $method:ident, $delegate:expr, $output:ty) => {
        assign_trait!($name, $method, $delegate, $output, [
            u8, i8, u16, i16, u32, i32, u64, i64, usize, isize, u128, i128
        ]);
    };
    ($name:ident, $method:ident, $delegate:expr, $output:ty, [$($ty:ident),*]) => {
        pub trait $name<Rhs = Self> {
            fn $method(&mut self, rhs: Rhs) -> $output;
        }

        $(
           impl<Rhs> $name<Rhs> for $ty
             where $ty: UpcastFrom<Rhs>
           {
               fn $method(&mut self, rhs: Rhs) -> $output {
                   let f: fn(&mut $ty, Rhs) -> $output = $delegate;
                   f(self, rhs)
               }
           }
        )*
    };
}

number_trait!(
    CheckedAdd,
    checked_add,
    Option<Self::Output>,
    assert!(Type::MAX.checked_add(Type::MAX).is_none())
);

number_trait!(
    CheckedSub,
    checked_sub,
    Option<Self::Output>,
    assert!(Type::MIN.checked_sub(Type::MAX).is_none())
);

number_trait!(
    CheckedMul,
    checked_mul,
    Option<Self::Output>,
    assert!(Type::MAX.checked_mul(Type::MAX).is_none())
);

number_trait!(
    CheckedDiv,
    checked_div,
    Option<Self::Output>,
    assert!(Type::MAX.checked_div(0).is_none())
);

assign_trait!(
    CheckedAddAssign,
    checked_add_assign,
    |value, rhs| {
        *value = CheckedAdd::checked_add(*value, rhs)?;
        Some(())
    },
    Option<()>
);

assign_trait!(
    CheckedSubAssign,
    checked_sub_assign,
    |value, rhs| {
        *value = CheckedSub::checked_sub(*value, rhs)?;
        Some(())
    },
    Option<()>
);

assign_trait!(
    CheckedMulAssign,
    checked_mul_assign,
    |value, rhs| {
        *value = CheckedMul::checked_mul(*value, rhs)?;
        Some(())
    },
    Option<()>
);

number_trait!(
    SaturatingAdd,
    saturating_add,
    Self::Output,
    assert_eq!(Type::MAX.saturating_add(Type::MAX), Type::MAX)
);

number_trait!(
    SaturatingSub,
    saturating_sub,
    Self::Output,
    assert_eq!(Type::MIN.saturating_sub(Type::MAX), Type::MIN)
);

number_trait!(
    SaturatingMul,
    saturating_mul,
    Self::Output,
    assert_eq!(Type::MAX.saturating_mul(Type::MAX), Type::MAX)
);

assign_trait!(
    SaturatingAddAssign,
    saturating_add_assign,
    |value, rhs| {
        *value = SaturatingAdd::saturating_add(*value, rhs);
    },
    ()
);

assign_trait!(
    SaturatingSubAssign,
    saturating_sub_assign,
    |value, rhs| {
        *value = SaturatingSub::saturating_sub(*value, rhs);
    },
    ()
);

assign_trait!(
    SaturatingMulAssign,
    saturating_mul_assign,
    |value, rhs| {
        *value = SaturatingMul::saturating_mul(*value, rhs);
    },
    ()
);

/// Losslessly upcasts from one type to another
pub trait UpcastFrom<T> {
    fn upcast_from(value: T) -> Self;
}

/// Losslessly upcasts one type into another
pub trait Upcast<T> {
    fn upcast(self) -> T;
}

/// Implement Upcast automatically for all types that implement UpcastFrom
impl<T, U> Upcast<T> for U
where
    T: UpcastFrom<U>,
{
    fn upcast(self) -> T {
        UpcastFrom::upcast_from(self)
    }
}

macro_rules! upcast_impl {
    ($($ty:ident),*) => {
        upcast_impl!(@impl [], [$($ty),*]);
    };
    (@impl [$($prev:ident),*], []) => {
        // done!
    };
    (@impl [$($prev:ident),*], [$current:ident $(, $rest:ident)*]) => {
        impl UpcastFrom<$current> for $current {
            fn upcast_from(value: Self) -> Self {
                value
            }
        }

        impl UpcastFrom<&$current> for $current {
            fn upcast_from(value: &Self) -> Self {
                *value
            }
        }

        $(
            impl UpcastFrom<$prev> for $current {
                fn upcast_from(value: $prev) -> Self {
                    value as $current
                }
            }

            impl UpcastFrom<&$prev> for $current {
                fn upcast_from(value: &$prev) -> Self {
                    (*value) as $current
                }
            }
        )*

        upcast_impl!(@impl [$current $(, $prev)*], [$($rest),*]);
    };
}

upcast_impl!(u8, u16, u32, u64, u128);
upcast_impl!(i8, i16, i32, i64, i128);

macro_rules! upcast_usize {
    ($target_width:literal, $unsigned:ident, $signed:ident) => {
        #[cfg(target_pointer_width = $target_width)]
        impl<Rhs> UpcastFrom<Rhs> for usize
        where
            $unsigned: UpcastFrom<Rhs>,
        {
            fn upcast_from(value: Rhs) -> Self {
                $unsigned::upcast_from(value) as usize
            }
        }

        #[cfg(target_pointer_width = $target_width)]
        impl<Rhs> UpcastFrom<Rhs> for isize
        where
            $signed: UpcastFrom<Rhs>,
        {
            fn upcast_from(value: Rhs) -> Self {
                $signed::upcast_from(value) as isize
            }
        }
    };
}

upcast_usize!("8", u8, i8);
upcast_usize!("16", u16, u16);
upcast_usize!("32", u32, i32);
upcast_usize!("64", u64, i64);
upcast_usize!("128", u128, i128);


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

//! This module provides `Query` and `QueryMut` traits, which are used for querying
//! and executing functions on several different providers. This includes
//! [`Subscriber::ConnectionContext`](crate::event::Subscriber::ConnectionContext)
//! on a Subscriber and the [`Sender`](crate::datagram::Sender) and
//! [`Receiver`](crate::datagram::Receiver) types on a datagram [`Endpoint`](crate::datagram::Endpoint).

use core::marker::PhantomData;

pub trait Query {
    fn execute(&mut self, context: &dyn core::any::Any) -> ControlFlow;
}

pub trait QueryMut {
    fn execute_mut(&mut self, context: &mut dyn core::any::Any) -> ControlFlow;
}

/// Used to tell a query whether it should exit early or go on as usual.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ControlFlow {
    Continue,
    Break,
}

impl ControlFlow {
    #[inline]
    #[must_use]
    pub fn and_then(self, f: impl FnOnce() -> Self) -> Self {
        match self {
            Self::Continue => f(),
            Self::Break => Self::Break,
        }
    }
}

/// A type that implements Query and QueryMut traits and only executes once.
///
/// This will execute and short-circuit on the first match.
pub struct Once<F, Context, Outcome> {
    query: Option<F>,
    result: Option<Outcome>,
    context: PhantomData<Context>,
}

impl<F, Context, Outcome> From<Once<F, Context, Outcome>> for Result<Outcome, Error> {
    #[inline]
    fn from(query: Once<F, Context, Outcome>) -> Self {
        query.result.ok_or(Error::ContextTypeMismatch)
    }
}

impl<F, ConnectionContext, Outcome> Once<F, ConnectionContext, Outcome>
where
    F: FnOnce(&ConnectionContext) -> Outcome,
    ConnectionContext: 'static,
{
    #[inline]
    pub fn new(query: F) -> Self {
        Self {
            query: Some(query),
            result: None,
            context: PhantomData,
        }
    }
}

impl<F, ConnectionContext, Outcome> Once<F, ConnectionContext, Outcome>
where
    F: FnOnce(&mut ConnectionContext) -> Outcome,
    ConnectionContext: 'static,
{
    #[inline]
    pub fn new_mut(query: F) -> Self {
        Self {
            query: Some(query),
            result: None,
            context: PhantomData,
        }
    }
}

impl<F, Context, Outcome> Query for Once<F, Context, Outcome>
where
    F: FnOnce(&Context) -> Outcome,
    Context: 'static,
{
    fn execute(&mut self, context: &dyn core::any::Any) -> ControlFlow {
        match context.downcast_ref::<Context>() {
            Some(context) => {
                let query = self.query.take().expect("can only match once");
                self.result = Some(query(context));
                ControlFlow::Break
            }
            None => ControlFlow::Continue,
        }
    }
}

impl<F, Context, Outcome> QueryMut for Once<F, Context, Outcome>
where
    F: FnOnce(&mut Context) -> Outcome,
    Context: 'static,
{
    fn execute_mut(&mut self, context: &mut dyn core::any::Any) -> ControlFlow {
        match context.downcast_mut::<Context>() {
            Some(context) => {
                let query = self.query.take().expect("can only match once");
                self.result = Some(query(context));
                ControlFlow::Break
            }
            None => ControlFlow::Continue,
        }
    }
}

#[non_exhaustive]
#[derive(Debug, Clone)]
/// Reason for the failed query.
pub enum Error {
    /// The connection lock is poisoned and the connection unusable.
    ConnectionLockPoisoned,

    /// The expected query type failed to match any of the configured types.
    ContextTypeMismatch,
}

impl core::fmt::Display for Error {
    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(f, "{self:?}")
    }
}

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


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

use core::ops::RangeInclusive;

/// A generator of random data. The two methods provide the same functionality for
/// different use cases. One for "public" randomly generated data that may appear
/// in the clear, and one for "private" data that should remain secret. This approach
/// lessens the risk of potential predictability weaknesses in random number generation
/// algorithms from leaking information across contexts.
pub trait Generator: 'static + Send {
    /// Fills `dest` with unpredictable bits that may be
    /// sent over the wire and viewable in the clear.
    fn public_random_fill(&mut self, dest: &mut [u8]);

    /// Fills `dest` with unpredictable bits that will only be
    /// used internally within the endpoint, remaining secret.
    fn private_random_fill(&mut self, dest: &mut [u8]);
}

/// Generates a random usize within the given inclusive range
///
/// NOTE: This will have slight bias towards the lower end of the range. Usages that
/// require uniform sampling should implement rejection sampling or other methodologies
/// and not copy this implementation.
pub(crate) fn gen_range_biased<R: Generator + ?Sized>(
    random_generator: &mut R,
    range: RangeInclusive<usize>,
) -> usize {
    if range.start() == range.end() {
        return *range.start();
    }

    let mut dest = [0; core::mem::size_of::<usize>()];
    random_generator.public_random_fill(&mut dest);
    let result = usize::from_le_bytes(dest);

    let max_variance = (range.end() - range.start()).saturating_add(1);
    range.start() + result % max_variance
}

#[cfg(any(test, feature = "testing"))]
pub mod testing {
    use crate::random;

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

    impl random::Generator for Generator {
        fn public_random_fill(&mut self, dest: &mut [u8]) {
            let seed = self.0;

            for (i, elem) in dest.iter_mut().enumerate() {
                *elem = seed ^ i as u8;
            }

            self.0 = self.0.wrapping_add(1)
        }

        fn private_random_fill(&mut self, dest: &mut [u8]) {
            let seed = u8::max_value() - self.0;

            for (i, elem) in dest.iter_mut().enumerate() {
                *elem = seed ^ i as u8;
            }

            self.0 = self.0.wrapping_add(1)
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::random;

    #[test]
    #[cfg_attr(miri, ignore)] // This test is too expensive for miri to complete in a reasonable amount of time
    #[cfg_attr(kani, kani::proof, kani::unwind(10), kani::solver(kissat))]
    fn gen_range_biased_test() {
        bolero::check!()
            .with_type()
            .cloned()
            .for_each(|(seed, mut min, mut max)| {
                if min > max {
                    core::mem::swap(&mut min, &mut max);
                }
                let mut generator = random::testing::Generator(seed);
                let result = random::gen_range_biased(&mut generator, min..=max);
                assert!(result >= min);
                assert!(result <= max);
            });
    }
}


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

use core::ops::{Deref, DerefMut};

/// Copies vectored slices from one slice into another
///
/// The number of copied items is limited by the minimum of the lengths of each of the slices.
///
/// Returns the number of entries that were copied
#[inline]
pub fn vectored_copy<A, B, T>(from: &[A], to: &mut [B]) -> usize
where
    A: Deref<Target = [T]>,
    B: Deref<Target = [T]> + DerefMut,
    T: Copy,
{
    zip_chunks(from, to, |a, b| {
        b.copy_from_slice(a);
    })
}

/// Zips entries from one slice to another
///
/// The number of copied items is limited by the minimum of the lengths of each of the slices.
///
/// Returns the number of entries that were processed
#[inline]
pub fn zip<A, At, B, Bt, F>(from: &[A], to: &mut [B], mut on_item: F) -> usize
where
    A: Deref<Target = [At]>,
    B: Deref<Target = [Bt]> + DerefMut,
    F: FnMut(&At, &mut Bt),
{
    zip_chunks(from, to, |a, b| {
        for (a, b) in a.iter().zip(b) {
            on_item(a, b);
        }
    })
}

/// Zips overlapping chunks from one slice to another
///
/// The number of copied items is limited by the minimum of the lengths of each of the slices.
///
/// Returns the number of entries that were processed
#[inline]
pub fn zip_chunks<A, At, B, Bt, F>(from: &[A], to: &mut [B], mut on_slice: F) -> usize
where
    A: Deref<Target = [At]>,
    B: Deref<Target = [Bt]> + DerefMut,
    F: FnMut(&[At], &mut [Bt]),
{
    let mut count = 0;

    let mut from_index = 0;
    let mut from_offset = 0;

    let mut to_index = 0;
    let mut to_offset = 0;

    // The compiler isn't smart enough to remove all of the bounds checks so we resort to
    // `get_unchecked`.
    //
    // https://godbolt.org/z/45cG1v

    // iterate until we reach one of the ends
    while from_index < from.len() && to_index < to.len() {
        let from = unsafe {
            // Safety: this length is already checked in the while condition
            debug_assert!(from.len() > from_index);
            from.get_unchecked(from_index)
        };

        let to = unsafe {
            // Safety: this length is already checked in the while condition
            debug_assert!(to.len() > to_index);
            to.get_unchecked_mut(to_index)
        };

        {
            // calculate the current views
            let from = unsafe {
                // Safety: the slice offsets are checked at the end of the while loop
                debug_assert!(from.len() >= from_offset);
                from.get_unchecked(from_offset..)
            };

            let to = unsafe {
                // Safety: the slice offsets are checked at the end of the while loop
                debug_assert!(to.len() >= to_offset);
                to.get_unchecked_mut(to_offset..)
            };

            let len = from.len().min(to.len());

            unsafe {
                // Safety: by using the min of the two lengths we will never exceed
                //         either slice's buffer
                debug_assert!(from.len() >= len);
                debug_assert!(to.len() >= len);

                let at = from.get_unchecked(..len);
                let bt = to.get_unchecked_mut(..len);

                on_slice(at, bt);
            }

            // increment the offsets
            from_offset += len;
            to_offset += len;
            count += len;
        }

        // check if the `from` is done
        if from.len() == from_offset {
            from_index += 1;
            from_offset = 0;
        }

        // check if the `to` is done
        if to.len() == to_offset {
            to_index += 1;
            to_offset = 0;
        }
    }

    count
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::testing::InlineVec;
    use bolero::check;

    fn assert_eq_slices<A, B, T>(a: &[A], b: &[B])
    where
        A: Deref<Target = [T]>,
        B: Deref<Target = [T]>,
        T: PartialEq + core::fmt::Debug,
    {
        let a = a.iter().flat_map(|a| a.iter());
        let b = b.iter().flat_map(|b| b.iter());

        // make sure all of the values match
        //
        // Note: this doesn't use Iterator::eq, as the slice lengths may be different
        for (a, b) in a.zip(b) {
            assert_eq!(a, b);
        }
    }

    #[test]
    fn vectored_copy_test() {
        let from = [
            &[0][..],
            &[1, 2, 3][..],
            &[4, 5, 6, 7][..],
            &[][..],
            &[8, 9, 10, 11][..],
        ];

        for len in 0..6 {
            let mut to = vec![vec![0; 2]; len];
            let copied_len = vectored_copy(&from, &mut to);
            assert_eq!(copied_len, len * 2);
            assert_eq_slices(&from, &to);
        }
    }

    const LEN: usize = if cfg!(kani) { 2 } else { 32 };

    #[test]
    #[cfg_attr(kani, kani::proof, kani::unwind(5), kani::solver(cadical))]
    #[cfg_attr(miri, ignore)] // This test is too expensive for miri to complete in a reasonable amount of time
    fn vectored_copy_fuzz_test() {
        check!()
            .with_type::<(
                InlineVec<InlineVec<u8, LEN>, LEN>,
                InlineVec<InlineVec<u8, LEN>, LEN>,
            )>()
            .cloned()
            .for_each(|(from, mut to)| {
                vectored_copy(&from, &mut to);
                assert_eq_slices(&from, &to);
            })
    }
}


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

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

pub use token::Token;

pub mod token;


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

#[cfg(feature = "alloc")]
mod primitive;

#[cfg(feature = "crossbeam-utils")]
pub use crossbeam_utils::CachePadded;
#[cfg(feature = "atomic-waker")]
pub mod atomic_waker;
#[cfg(target_has_atomic = "32")]
pub mod cursor;
#[cfg(feature = "alloc")]
pub mod spsc;
#[cfg(feature = "alloc")]
pub mod worker;


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

use bolero_generator::TypeGenerator;

#[derive(Clone, Copy, Debug, TypeGenerator)]
pub struct InlineVec<T, const LEN: usize> {
    values: [T; LEN],

    #[generator(_code = "0..LEN")]
    len: usize,
}

impl<T, const LEN: usize> core::ops::Deref for InlineVec<T, LEN> {
    type Target = [T];

    fn deref(&self) -> &Self::Target {
        &self.values[..self.len]
    }
}

impl<T, const LEN: usize> core::ops::DerefMut for InlineVec<T, LEN> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.values[..self.len]
    }
}

#[cfg(all(test, not(loom)))]
pub mod loom {
    pub use std::*;

    pub mod future {
        pub use futures::executor::block_on;
    }

    pub fn model<F: 'static + FnOnce() -> R, R>(f: F) -> R {
        f()
    }
}

#[cfg(all(test, loom))]
pub use loom;


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

use crate::{
    connection,
    event::{api::SocketAddress, IntoEvent},
    inet, random,
};

#[non_exhaustive]
pub struct Context<'a> {
    pub remote_address: SocketAddress<'a>,
    pub peer_connection_id: &'a [u8],
    pub random: &'a mut dyn random::Generator,
}

impl<'a> Context<'a> {
    #[inline]
    #[doc(hidden)]
    pub fn new(
        remote_address: &'a inet::SocketAddress,
        peer_connection_id: &'a connection::PeerId,
        random: &'a mut dyn random::Generator,
    ) -> Self {
        Self {
            remote_address: remote_address.into_event(),
            peer_connection_id: peer_connection_id.as_bytes(),
            random,
        }
    }
}

pub trait Format: 'static + Send {
    const TOKEN_LEN: usize;

    /// Generate a signed token to be delivered in a NEW_TOKEN frame.
    /// This function will only be called if the provider support NEW_TOKEN frames.
    fn generate_new_token(
        &mut self,
        context: &mut Context<'_>,
        source_connection_id: &connection::LocalId,
        output_buffer: &mut [u8],
    ) -> Option<()>;

    /// Generate a signed token to be delivered in a Retry Packet
    fn generate_retry_token(
        &mut self,
        context: &mut Context<'_>,
        original_destination_connection_id: &connection::InitialId,
        output_buffer: &mut [u8],
    ) -> Option<()>;

    /// Return the original destination connection id of a valid token.
    /// If the token is invalid, return None.
    /// Callers should detect duplicate tokens and treat them as invalid.
    fn validate_token(
        &mut self,
        context: &mut Context<'_>,
        token: &[u8],
    ) -> Option<connection::InitialId>;
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Source {
    RetryPacket,
    NewTokenFrame,
}

#[cfg(any(test, feature = "testing"))]
pub mod testing {
    use super::*;
    use crate::crypto::retry;

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

    impl super::Format for Format {
        const TOKEN_LEN: usize = retry::example::TOKEN_LEN;

        fn generate_new_token(
            &mut self,
            _context: &mut Context<'_>,
            _source_connection_id: &connection::LocalId,
            _output_buffer: &mut [u8],
        ) -> Option<()> {
            // TODO implement one for testing
            None
        }

        fn generate_retry_token(
            &mut self,
            _context: &mut Context<'_>,
            _original_destination_connection_id: &connection::InitialId,
            output_buffer: &mut [u8],
        ) -> Option<()> {
            output_buffer.copy_from_slice(&retry::example::TOKEN);
            Some(())
        }

        fn validate_token(
            &mut self,
            _context: &mut Context<'_>,
            token: &[u8],
        ) -> Option<connection::InitialId> {
            if token == retry::example::TOKEN {
                return Some(connection::InitialId::TEST_ID);
            }

            None
        }
    }

    impl Format {
        pub fn new() -> Self {
            Self(0)
        }
    }
}


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

pub mod bpf;
pub mod decoder;
pub mod encoder;
pub mod path;


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

//! A set of acknowledgments

use crate::packet::number::{PacketNumber, PacketNumberRange};
use core::ops::RangeInclusive;

/// A set of Acknowledgments
///
/// The implementation of the set is allowed to store packet numbers in
/// an arbitrary form.
pub trait Set {
    /// Returns whether the [`Set`] contains a given packet number
    fn contains(&self, packet_number: PacketNumber) -> bool;

    /// Smallest packet number in the set
    fn smallest(&self) -> PacketNumber;

    /// Largest packet number in the set
    fn largest(&self) -> PacketNumber;

    /// Returns the range of packet numbers contained in the set
    fn as_range(&self) -> PacketNumberRange {
        PacketNumberRange::new(self.smallest(), self.largest())
    }
}

// A single packet number is also a set

impl Set for PacketNumber {
    fn contains(&self, packet_number: PacketNumber) -> bool {
        *self == packet_number
    }

    fn smallest(&self) -> PacketNumber {
        *self
    }

    fn largest(&self) -> PacketNumber {
        *self
    }
}

impl Set for RangeInclusive<PacketNumber> {
    fn contains(&self, packet_number: PacketNumber) -> bool {
        RangeInclusive::contains(self, &packet_number)
    }

    fn smallest(&self) -> PacketNumber {
        *self.start()
    }

    fn largest(&self) -> PacketNumber {
        *self.end()
    }
}

impl Set for PacketNumberRange {
    fn contains(&self, packet_number: PacketNumber) -> bool {
        PacketNumberRange::contains(self, packet_number)
    }

    fn smallest(&self) -> PacketNumber {
        self.start()
    }

    fn largest(&self) -> PacketNumber {
        self.end()
    }

    fn as_range(&self) -> PacketNumberRange {
        *self
    }
}


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

use crate::{
    transport::parameters::{AckDelayExponent, MaxAckDelay},
    varint::VarInt,
};
use core::{convert::TryInto, time::Duration};

// After running simulations, this seemed to be a good baseline
// TODO experiment more with this
/// The recommended value for the ack_elicitation_interval setting
const RECOMMENDED_ELICITATION_INTERVAL: u8 = 4;

// TODO experiment more with this
/// The recommended number of packet number ranges that an endpoint should store
const RECOMMENDED_RANGES_LIMIT: u8 = 10;

/// Settings for ACK frames
#[derive(Clone, Copy, Debug)]
pub struct Settings {
    /// The maximum ACK delay indicates the maximum amount of time by which the
    /// endpoint will delay sending acknowledgments.
    pub max_ack_delay: Duration,
    /// The ACK delay exponent is an integer value indicating an exponent used
    /// to decode the ACK Delay field in the ACK frame
    pub ack_delay_exponent: u8,

    //= https://www.rfc-editor.org/rfc/rfc9000#section-13.2.4
    //# A receiver that sends only non-ack-eliciting packets, such as ACK
    //# frames, might not receive an acknowledgement for a long period of
    //# time.  This could cause the receiver to maintain state for a large
    //# number of ACK frames for a long period of time, and ACK frames it
    //# sends could be unnecessarily large.  In such a case, a receiver could
    //# send a PING or other small ack-eliciting frame occasionally, such as
    //# once per round trip, to elicit an ACK from the peer.
    /// The number of packets received before sending an ACK-eliciting packet
    pub ack_elicitation_interval: u8,

    /// The number of packet number intervals an endpoint is willing to store
    pub ack_ranges_limit: u8,
}

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

//= https://www.rfc-editor.org/rfc/rfc9000#section-19.3
//# ACK Delay:  A variable-length integer encoding the acknowledgement
//#    delay in microseconds; see Section 13.2.5.  It is decoded by
//#    multiplying the value in the field by 2 to the power of the
//#    ack_delay_exponent transport parameter sent by the sender of the
//#    ACK frame; see Section 18.2.  Compared to simply expressing the
//#    delay as an integer, this encoding allows for a larger range of
//#    values within the same number of bytes, at the cost of lower
//#    resolution.

impl Settings {
    //= https://www.rfc-editor.org/rfc/rfc9000#section-13.2.1
    //# An endpoint MUST acknowledge all ack-eliciting Initial and Handshake
    //# packets immediately
    pub const EARLY: Self = Self {
        max_ack_delay: Duration::from_secs(0),
        ack_delay_exponent: 0,
        ..Self::RECOMMENDED
    };

    pub const RECOMMENDED: Self = Self {
        max_ack_delay: MaxAckDelay::RECOMMENDED.as_duration(),
        ack_delay_exponent: AckDelayExponent::RECOMMENDED.as_u8(),
        ack_elicitation_interval: RECOMMENDED_ELICITATION_INTERVAL,
        ack_ranges_limit: RECOMMENDED_RANGES_LIMIT,
    };

    /// Decodes the peer's `Ack Delay` field
    pub fn decode_ack_delay(&self, delay: VarInt) -> Duration {
        Duration::from_micros(*delay) * self.scale()
    }

    /// Encodes the local `Ack Delay` field
    pub fn encode_ack_delay(&self, delay: Duration) -> VarInt {
        let micros = delay.as_micros();
        let scale = self.scale() as u128;
        (micros / scale).try_into().unwrap_or(VarInt::MAX)
    }

    /// Computes the scale from the exponent
    fn scale(&self) -> u32 {
        2u32.pow(self.ack_delay_exponent as u32)
    }
}

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

    #[test]
    #[cfg_attr(miri, ignore)] // this test is too expensive for miri
    fn ack_settings_test() {
        for ack_delay_exponent in 0..=20 {
            let settings = Settings {
                ack_delay_exponent,
                ..Default::default()
            };
            // use an epsilon instead of comparing the values directly,
            // as there will be some precision loss
            let epsilon = settings.scale() as u128;

            for delay in (0..1000).map(|v| v * 100).map(Duration::from_micros) {
                let delay_varint = settings.encode_ack_delay(delay);
                let expected_us = delay.as_micros();
                let actual_us = settings.decode_ack_delay(delay_varint).as_micros();
                let actual_difference = expected_us - actual_us;
                assert!(actual_difference < epsilon);
            }

            // ensure MAX values are handled correctly and don't overflow
            let delay = settings.decode_ack_delay(VarInt::MAX);
            let delay_varint = settings.encode_ack_delay(delay);
            assert_eq!(VarInt::MAX, delay_varint);
        }
    }
}


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

//! Defines QUIC Application Error Codes

use crate::{
    frame::ConnectionClose,
    varint::{VarInt, VarIntError},
};
use core::{convert::TryFrom, fmt, ops};

//= https://www.rfc-editor.org/rfc/rfc9000#section-20.2
//# The management of application error codes is left to application
//# protocols.

/// Application Error Codes are 62-bit unsigned integer values which
/// may be used by applications to exchange errors.
#[derive(Copy, Clone, Eq, PartialEq)]
pub struct Error(VarInt);

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

impl fmt::Debug for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "application::Error({})", self.0.as_u64())
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "QUIC application error code: {}", self.0.as_u64())
    }
}

impl Error {
    /// An error code that can be used when the application cannot provide
    /// a more meaningful code.
    pub const UNKNOWN: Self = Self(VarInt::from_u8(0));

    /// Creates an `ApplicationErrorCode` from an unsigned integer.
    ///
    /// This will return the error code if the given value is inside the valid
    /// range for error codes and return `Err` otherwise.
    pub fn new(value: u64) -> Result<Self, VarIntError> {
        Ok(Self(VarInt::new(value)?))
    }
}

impl ops::Deref for Error {
    type Target = u64;

    fn deref(&self) -> &Self::Target {
        self.0.deref()
    }
}

impl TryInto for Error {
    fn application_error(&self) -> Option<Error> {
        Some(*self)
    }
}

impl From<VarInt> for Error {
    fn from(value: VarInt) -> Self {
        Self(value)
    }
}

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

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

impl<'a> From<Error> for ConnectionClose<'a> {
    fn from(error: Error) -> Self {
        ConnectionClose {
            error_code: error.0,
            frame_type: None,
            reason: None,
        }
    }
}

macro_rules! convert {
    ($ty:ident) => {
        impl From<$ty> for Error {
            fn from(value: $ty) -> Self {
                Self(VarInt::from(value))
            }
        }
    };
}

convert!(u8);
convert!(u16);
convert!(u32);

impl TryFrom<u64> for Error {
    type Error = VarIntError;

    fn try_from(value: u64) -> Result<Self, Self::Error> {
        Ok(VarInt::try_from(value)?.into())
    }
}

/// Conversion trait for errors that have an associated [`Error`]
pub trait TryInto {
    /// Returns the associated [`Error`], if any
    fn application_error(&self) -> Option<Error>;
}


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

pub mod error;
#[cfg(feature = "alloc")]
pub(crate) mod server_name;

pub use error::Error;
#[cfg(feature = "alloc")]
pub use server_name::ServerName;


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

use bytes::Bytes;

/// ServerName holds a negotiated
/// [Server Name Indication](https://en.wikipedia.org/wiki/Server_Name_Indication)
/// value, encoded as UTF-8.
///
/// ServerName should be a valid UTF-8 string, therefore this struct can only be
/// constructed from a `&str` or `String`.
///
/// ```rust
/// # use s2n_quic_core::application::ServerName;
/// let string: String = String::from("a valid utf-8 string");
/// let name: ServerName = string.into();
///
/// let str_arr: &str = &"a valid utf-8 str array";
/// let name: ServerName = str_arr.into();
/// ```
///
/// `ServerName` serves a dual purpose:
/// - It can be converted into [`Bytes`] which supports zero-copy slicing and
/// reference counting.
/// - It can be accessed as `&str` so that applications can reason about the string value.
#[derive(Clone, PartialEq)]
pub struct ServerName(Bytes);

/// A static value for localhost
#[allow(dead_code)] // this is used by conditional modules so don't warn
pub(crate) static LOCALHOST: ServerName = ServerName(Bytes::from_static(b"localhost"));

impl ServerName {
    #[inline]
    pub fn into_bytes(self) -> Bytes {
        self.0
    }

    #[inline]
    fn as_str(&self) -> &str {
        // Safety: the byte array is validated as a valid UTF-8 string
        // before creating an instance of Sni.
        unsafe { core::str::from_utf8_unchecked(&self.0) }
    }
}

impl From<&str> for ServerName {
    #[inline]
    fn from(data: &str) -> Self {
        Self(Bytes::copy_from_slice(data.as_bytes()))
    }
}

#[cfg(feature = "alloc")]
impl From<alloc::string::String> for ServerName {
    #[inline]
    fn from(data: alloc::string::String) -> Self {
        Self(data.into_bytes().into())
    }
}

impl core::fmt::Debug for ServerName {
    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
        self.as_str().fmt(f)
    }
}

impl core::ops::Deref for ServerName {
    type Target = str;

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


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

mod receive_buffer;

pub use receive_buffer::*;


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

//! This module contains data structures for buffering incoming and outgoing data
//! in Quic streams.

use crate::varint::VarInt;
use alloc::collections::{vec_deque, VecDeque};
use bytes::BytesMut;

mod request;
mod slot;

#[cfg(test)]
mod tests;

use request::Request;
use slot::Slot;

/// Enumerates error that can occur while inserting data into the Receive Buffer
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum ReceiveBufferError {
    /// An invalid data range was provided
    OutOfRange,
}

/// The default buffer size for slots that the [`ReceiveBuffer`] uses.
///
/// This value was picked as it is typically used for the default memory page size.
const MIN_BUFFER_ALLOCATION_SIZE: usize = 4096;

//= https://www.rfc-editor.org/rfc/rfc9000#section-2.2
//# Endpoints MUST be able to deliver stream data to an application as an
//# ordered byte-stream.

/// `ReceiveBuffer` is a buffer structure for combining chunks of bytes in an
/// ordered stream, which might arrive out of order.
///
/// `ReceiveBuffer` will accumulate the bytes, and provide them to its users
/// once a contiguous range of bytes at the current position of the stream has
/// been accumulated.
///
/// `ReceiveBuffer` is optimized for minimizing memory allocations and for
/// offering it's users chunks of sizes that minimize call overhead.
///
/// If data is received in smaller chunks, only the first chunk will trigger a
/// memory allocation. All other chunks can be copied into the already allocated
/// region.
///
/// When users want to consume data from the buffer, the consumable part of the
/// internal receive buffer is split off and passed back to the caller. Due to
/// this chunk being a view onto a reference-counted internal buffer of type
/// [`BytesMut`] this is also efficient and does not require additional memory
/// allocation or copy.
///
/// ## Usage
///
/// ```rust,ignore
/// use s2n_quic_transport::buffer::ReceiveBuffer;
///
/// let mut buffer = ReceiveBuffer::new();
///
/// // write a chunk of bytes at offset 4, which can not be consumed yet
/// assert!(buffer.write_at(4u32.into(), &[4, 5, 6, 7]).is_ok());
/// assert_eq!(0, buffer.len());
/// assert_eq!(None, buffer.pop());
///
/// // write a chunk of bytes at offset 0, which allows for consumption
/// assert!(buffer.write_at(0u32.into(), &[0, 1, 2, 3]).is_ok());
/// assert_eq!(8, buffer.len());
///
/// // Pop chunks. Since they all fitted into a single internal buffer,
/// // they will be returned in combined fashion.
/// assert_eq!(&[0u8, 1, 2, 3, 4, 5, 6, 7], &buffer.pop().unwrap()[..]);
/// ```
#[derive(Debug)]
pub struct ReceiveBuffer {
    slots: VecDeque<Slot>,
    start_offset: u64,
}

impl Default for ReceiveBuffer {
    fn default() -> Self {
        Self::new()
    }
}

impl ReceiveBuffer {
    /// Creates a new `ReceiveBuffer`
    pub fn new() -> ReceiveBuffer {
        ReceiveBuffer {
            slots: VecDeque::new(),
            start_offset: 0,
        }
    }

    /// Returns the amount of bytes available for reading.
    /// This equals the amount of data that is stored in contiguous fashion at
    /// the start of the buffer.
    #[inline]
    pub fn len(&self) -> usize {
        self.report().0
    }

    /// Returns true if no bytes are available for reading
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Returns the number of bytes and chunks available for consumption
    #[inline]
    pub fn report(&self) -> (usize, usize) {
        let mut bytes = 0;
        let mut chunks = 0;
        for chunk in self.iter() {
            bytes += chunk.len();
            chunks += 1;
        }
        (bytes, chunks)
    }

    /// Pushes a slice at a certain offset
    #[inline]
    pub fn write_at(&mut self, offset: VarInt, data: &[u8]) -> Result<(), ReceiveBufferError> {
        // create a request
        let request = Request::new(offset, data)?;

        // trim off any data that we've already read
        let (_, mut request) = request.split(self.start_offset);

        // if the request is empty we're done
        if request.is_empty() {
            return Ok(());
        }

        // start from the back with the assumption that most data arrives in order
        for mut idx in (0..self.slots.len()).rev() {
            let slot = &mut self.slots[idx];

            let slot::Outcome { lower, mid, upper } = slot.try_write(request);

            // if this slot was completed, we should try and unsplit with the next slot
            if slot.is_full() {
                let current_block =
                    Self::align_offset(slot.start(), Self::allocation_size(slot.start()));
                let end = slot.end();

                if let Some(next) = self.slots.get(idx + 1) {
                    let next_block =
                        Self::align_offset(next.start(), Self::allocation_size(next.start()));

                    if next.start() == end && current_block == next_block {
                        if let Some(next) = self.slots.remove(idx + 1) {
                            self.slots[idx].unsplit(next);
                        } else {
                            debug_assert!(false, "slot should be available");
                            unsafe {
                                // Safety: we've already checked that `idx + 1` exists
                                core::hint::unreachable_unchecked();
                            }
                        }
                    }
                }
            }

            idx += 1;
            self.allocate_request(idx, upper);

            if let Some(mid) = mid {
                self.insert(idx, mid);
            }

            request = lower;

            if request.is_empty() {
                break;
            }
        }

        self.allocate_request(0, request);

        self.check_consistency();

        Ok(())
    }

    /// Iterates over all of the chunks waiting to be received
    #[inline]
    pub fn iter(&self) -> impl Iterator<Item = &[u8]> {
        Iter::new(self)
    }

    /// Pops a buffer from the front of the receive queue if available
    #[inline]
    pub fn pop(&mut self) -> Option<BytesMut> {
        self.pop_transform(|buffer| buffer.split())
    }

    /// Pops a buffer from the front of the receive queue, who's length is always guaranteed to be
    /// less than the provided `watermark`.
    #[inline]
    pub fn pop_watermarked(&mut self, watermark: usize) -> Option<BytesMut> {
        self.pop_transform(|buffer| {
            // make sure the buffer doesn't exceed the watermark
            let watermark = watermark.min(buffer.len());

            // if the watermark is 0 then don't needlessly increment refcounts
            if watermark == 0 {
                return BytesMut::new();
            }

            buffer.split_to(watermark)
        })
    }

    /// Pops a buffer from the front of the receive queue as long as the `transform` function returns a
    /// non-empty buffer.
    #[inline]
    fn pop_transform<F: Fn(&mut BytesMut) -> BytesMut>(
        &mut self,
        transform: F,
    ) -> Option<BytesMut> {
        let slot = self.slots.front_mut()?;

        if !slot.is_occupied(self.start_offset) {
            return None;
        }

        let buffer = slot.data_mut();

        let out = transform(buffer);

        // filter out empty buffers
        if out.is_empty() {
            return None;
        }

        slot.add_start(out.len());

        if slot.start() == slot.end_allocated() {
            // remove empty buffers
            self.slots.pop_front();
        }

        self.start_offset += out.len() as u64;

        Some(out)
    }

    /// Returns the amount of data that had already been consumed from the
    /// receive buffer.
    #[inline]
    pub fn consumed_len(&self) -> u64 {
        self.start_offset
    }

    /// Returns the total amount of contiguous received data.
    ///
    /// This includes the already consumed data as well as the data that is still
    /// buffered and available for consumption.
    #[inline]
    pub fn total_received_len(&self) -> u64 {
        self.consumed_len() + self.len() as u64
    }

    /// Resets the receive buffer.
    ///
    /// This will drop all previously received data.
    #[inline]
    pub fn reset(&mut self) {
        self.slots.clear();
        self.start_offset = Default::default();
    }

    #[inline(always)]
    fn insert(&mut self, idx: usize, slot: Slot) {
        if self.slots.len() < idx {
            debug_assert_eq!(self.slots.len() + 1, idx);
            self.slots.push_back(slot);
        } else {
            self.slots.insert(idx, slot);
        }
    }

    #[inline]
    fn allocate_request(&mut self, mut idx: usize, mut request: Request) {
        while !request.is_empty() {
            let size = Self::allocation_size(request.start());
            let offset = Self::align_offset(request.start(), size);
            let buffer = BytesMut::with_capacity(size);
            let end = offset + size as u64;
            let mut slot = Slot::new(offset, end, buffer);

            let slot::Outcome { lower, mid, upper } = slot.try_write(request);

            debug_assert!(lower.is_empty(), "lower requests should always be empty");

            // first insert the newly-created Slot
            self.insert(idx, slot);
            idx += 1;

            // check if we have a mid-slot and insert that as well
            if let Some(mid) = mid {
                self.insert(idx, mid);
                idx += 1;
            }

            // set the current request to the upper slot and loop
            request = upper;
        }
    }

    /// Aligns an offset to a certain alignment size
    #[inline]
    fn align_offset(offset: u64, alignment: usize) -> u64 {
        (offset / (alignment as u64)) * (alignment as u64)
    }

    /// Returns the desired allocation size for the given offset
    ///
    /// The allocation size gradually increases as the offset increases. This is under
    /// the assumption that streams that receive a lot of data will continue to receive
    /// a lot of data.
    ///
    /// The current table is as follows:
    ///
    /// | offset         | allocation size |
    /// |----------------|-----------------|
    /// | 0              | 4096            |
    /// | 65536          | 16384           |
    /// | 262144         | 32768           |
    /// | >=1048575      | 65536           |
    #[inline]
    fn allocation_size(offset: u64) -> usize {
        for pow in (2..=4).rev() {
            let mult = 1 << pow;
            let square = mult * mult;
            let min_offset = (MIN_BUFFER_ALLOCATION_SIZE * square) as u64;
            let allocation_size = MIN_BUFFER_ALLOCATION_SIZE * mult;

            if offset >= min_offset {
                return allocation_size;
            }
        }

        MIN_BUFFER_ALLOCATION_SIZE
    }

    #[inline]
    fn check_consistency(&self) {
        if cfg!(debug_assertions) {
            let mut prev_end = self.start_offset;

            for slot in &self.slots {
                assert!(slot.start() >= prev_end, "{self:#?}");
                prev_end = slot.end_allocated();
            }
        }
    }
}

pub struct Iter<'a> {
    prev_end: u64,
    inner: vec_deque::Iter<'a, Slot>,
}

impl<'a> Iter<'a> {
    #[inline]
    fn new(buffer: &'a ReceiveBuffer) -> Self {
        Self {
            prev_end: buffer.start_offset,
            inner: buffer.slots.iter(),
        }
    }
}

impl<'a> Iterator for Iter<'a> {
    type Item = &'a [u8];

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

        if !slot.is_occupied(self.prev_end) {
            return None;
        }

        self.prev_end = slot.end();
        Some(slot.as_slice())
    }
}


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

use super::ReceiveBufferError;
use crate::varint::VarInt;
use bytes::{BufMut, BytesMut};
use core::fmt;

#[derive(PartialEq, Eq)]
pub struct Request<'a> {
    offset: u64,
    data: &'a [u8],
}

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

impl<'a> Request<'a> {
    #[inline]
    pub fn new(offset: VarInt, data: &'a [u8]) -> Result<Self, ReceiveBufferError> {
        offset
            .checked_add_usize(data.len())
            .ok_or(ReceiveBufferError::OutOfRange)?;
        Ok(Self {
            offset: offset.as_u64(),
            data,
        })
    }

    #[inline]
    pub fn split(self, offset: u64) -> (Self, Self) {
        let mid = offset.saturating_sub(self.offset);
        let mid = self.data.len().min(mid as _);
        let (a, b) = self.data.split_at(mid);

        let a_offset = self.offset.min(offset);
        let b_offset = self.offset.max(offset);

        let a = Self {
            offset: a_offset,
            data: a,
        };
        let b = Self {
            offset: b_offset,
            data: b,
        };
        (a, b)
    }

    #[inline]
    pub fn write(self, buffer: &mut BytesMut) {
        let chunk = buffer.chunk_mut();
        unsafe {
            let len = self.data.len();
            debug_assert!(len <= chunk.len(), "{:?} <= {:?}", len, chunk.len());

            // Safety: `chunk` is always going to be uninitialized memory which is allocated through `BytesMut`.
            //         Since the receive buffer owns this allocation, it's impossible for the request to overlap
            //         with this `chunk`.
            core::ptr::copy_nonoverlapping(self.data.as_ptr(), chunk.as_mut_ptr(), len);
            buffer.advance_mut(len);
        }
    }

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

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

    #[inline]
    pub fn into_option(self) -> Option<Self> {
        if self.data.is_empty() {
            None
        } else {
            Some(self)
        }
    }
}


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

use super::Request;
use bytes::BytesMut;
use core::fmt;

/// Possible states for slots in the [`ReceiveBuffer`]s queue
#[derive(PartialEq, Eq)]
pub struct Slot {
    start: u64,
    end: u64,
    data: BytesMut,
}

impl fmt::Debug for Slot {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Slot")
            .field("start", &self.start)
            .field("end", &self.end)
            .field("len", &self.data.len())
            .field("capacity", &self.data.capacity())
            .finish()
    }
}

#[derive(Debug, PartialEq, Eq)]
pub struct Outcome<'a> {
    pub lower: Request<'a>,
    pub mid: Option<Slot>,
    pub upper: Request<'a>,
}

impl Slot {
    pub fn new(start: u64, end: u64, data: BytesMut) -> Self {
        Self { start, end, data }
    }

    #[inline(always)]
    pub fn try_write<'a>(&mut self, request: Request<'a>) -> Outcome<'a> {
        // trim off chunks lower than the start
        let (lower, request) = request.split(self.start());

        // trim off chunks we've already copied
        let (_, request) = request.split(self.end());

        // trim off chunks higher than the allocated end
        let (to_write, upper) = request.split(self.end_allocated());

        let mut mid = None;

        if let Some(to_write) = to_write.into_option() {
            if to_write.start() > self.end() {
                // find the split point between the buffers
                let len = to_write.start() - self.start();

                // create a new mid slot
                let start = to_write.start();
                let mut data = self.data.split_off(len as usize);

                // copy the data to the buffer
                to_write.write(&mut data);

                mid = Some(Self {
                    start,
                    end: self.end,
                    data,
                });
                self.end = start;
            } else {
                // copy the data to the buffer
                to_write.write(&mut self.data);
            }
        }

        Outcome { lower, mid, upper }
    }

    #[inline]
    pub fn unsplit(&mut self, next: Self) {
        debug_assert_eq!(self.end(), self.end_allocated());
        debug_assert_eq!(self.end(), next.start());
        self.data.unsplit(next.data);
        self.end = next.end;
    }

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

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

    #[inline]
    pub fn is_occupied(&self, prev_offset: u64) -> bool {
        !self.is_empty() && self.start() == prev_offset
    }

    #[inline]
    pub fn as_slice(&self) -> &[u8] {
        &self.data
    }

    #[inline]
    pub fn data_mut(&mut self) -> &mut BytesMut {
        &mut self.data
    }

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

    #[inline]
    pub fn add_start(&mut self, len: usize) {
        self.start += len as u64;
    }

    #[inline]
    pub fn end(&self) -> u64 {
        self.start() + self.data.len() as u64
    }

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

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

    #[test]
    fn bytes_assumption() {
        let mut data = BytesMut::with_capacity(100);
        let half = data.split_off(50);

        // after splitting a BytesMut the capacity should also be split
        assert_eq!(data.capacity(), 50);
        assert_eq!(half.capacity(), 50);
    }

    fn slot(range: core::ops::Range<u64>, data: &[u8]) -> Slot {
        let mut buffer = BytesMut::with_capacity((range.end - range.start) as usize);
        buffer.extend_from_slice(data);
        Slot::new(range.start, range.end, buffer)
    }

    fn req(offset: u64, data: &[u8]) -> Request {
        Request::new(VarInt::new(offset).unwrap(), data).unwrap()
    }

    macro_rules! assert_write {
        ($slot:expr, $request:expr, $expected_slot:expr, $result:expr) => {{
            let mut slot = $slot;
            let req = $request;

            let result = slot.try_write(req);
            assert_eq!(slot, $expected_slot);
            assert_eq!(result, $result);
        }};
    }

    #[test]
    fn overlap() {
        assert_write!(
            slot(4..8, &[1]),
            req(0, &[42; 12]),
            slot(4..8, &[1, 42, 42, 42]),
            Outcome {
                lower: req(0, &[42; 4]),
                mid: None,
                upper: req(8, &[42; 4]),
            }
        );
    }

    #[test]
    fn upper() {
        assert_write!(
            slot(4..8, &[1]),
            req(8, &[42; 4]),
            slot(4..8, &[1]),
            Outcome {
                lower: req(4, &[]),
                mid: None,
                upper: req(8, &[42; 4]),
            }
        );
    }

    #[test]
    fn lower() {
        assert_write!(
            slot(4..8, &[1]),
            req(0, &[42; 4]),
            slot(4..8, &[1]),
            Outcome {
                lower: req(0, &[42; 4]),
                mid: None,
                upper: req(8, &[]),
            }
        );
    }

    #[test]
    fn mid() {
        assert_write!(
            slot(4..8, &[1]),
            req(6, &[42; 1]),
            slot(4..6, &[1]),
            Outcome {
                lower: req(4, &[]),
                mid: Some(slot(6..8, &[42])),
                upper: req(8, &[]),
            }
        );
    }

    #[test]
    fn mid_upper() {
        assert_write!(
            slot(4..8, &[1]),
            req(6, &[42; 4]),
            slot(4..6, &[1]),
            Outcome {
                lower: req(4, &[]),
                mid: Some(slot(6..8, &[42; 2])),
                upper: req(8, &[42; 2]),
            }
        );
    }
}