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

//! A model that ensures stream data is correctly sent and received between peers

#[cfg(feature = "generator")]
use bolero_generator::*;

#[cfg(test)]
use bolero::generator::*;

use bytes::Bytes;

static DATA: Bytes = {
    const INNER: [u8; DATA_LEN] = {
        let mut data = [0; DATA_LEN];
        let mut idx = 0;
        while idx < DATA_LEN {
            data[idx] = idx as u8;
            idx += 1;
        }
        data
    };

    Bytes::from_static(&INNER)
};

// when running with miri, set the values lower to make execution less expensive
const DATA_LEN: usize = (DEFAULT_STREAM_LEN as usize) * if cfg!(miri) { 1 } else { 128 };
const DEFAULT_STREAM_LEN: u64 = if cfg!(miri) { DATA_MOD as _ } else { 1024 };
const DATA_MOD: usize = 256; // Only the first 256 offsets of DATA are unique

#[cfg(any(feature = "generator", test))]
const GENERATOR: core::ops::RangeInclusive<u64> = 0..=DEFAULT_STREAM_LEN;

#[derive(Clone, Copy, Debug, PartialEq, PartialOrd, Eq, Ord)]
#[cfg_attr(any(feature = "generator", test), derive(TypeGenerator))]
pub struct Data {
    #[cfg_attr(any(feature = "generator", test), generator(GENERATOR))]
    len: u64,
    #[cfg_attr(any(feature = "generator", test), generator(constant(0)))]
    offset: u64,
}

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

impl Data {
    pub const fn new(len: u64) -> Self {
        Self { len, offset: 0 }
    }

    /// Notifies the data that a set of chunks were received
    pub fn receive<Chunk: AsRef<[u8]>>(&mut self, chunks: &[Chunk]) {
        Self::receive_check(&mut self.offset, self.len, chunks)
    }

    /// Notifies the data that a set of chunks were received
    pub fn receive_at<Chunk: AsRef<[u8]>>(&self, mut start: u64, chunks: &[Chunk]) {
        Self::receive_check(&mut start, self.len, chunks)
    }

    fn receive_check<Chunk: AsRef<[u8]>>(start: &mut u64, len: u64, chunks: &[Chunk]) {
        for mut chunk in chunks.iter().map(AsRef::as_ref) {
            while !chunk.is_empty() {
                let offset = ((*start) % DATA_MOD as u64) as usize;
                let len = chunk.len().min(DATA.len() - offset);
                assert_eq!(
                    &chunk[..len],
                    &DATA[offset..offset + len],
                    "receive stream data at offset {} has been corrupted",
                    *start,
                );
                *start += len as u64;
                chunk = &chunk[len..];
            }
        }

        assert!(
            len >= (*start),
            "receive stream data has exceeded beyond the expected amount by {}",
            (*start) - len
        );
    }

    /// Returns `true` if the stream is finished reading/writing
    pub fn is_finished(&self) -> bool {
        self.len <= self.offset
    }

    /// Asks the data what chunks should be sent next
    ///
    /// Returns the number of chunks that were filled
    pub fn send(&mut self, mut amount: usize, chunks: &mut [Bytes]) -> Option<usize> {
        if self.is_finished() {
            return None;
        }

        let mut count = 0;
        for chunk in chunks.iter_mut() {
            if let Some(data) = self.send_one(amount).filter(|data| !data.is_empty()) {
                amount -= data.len();
                count += 1;
                *chunk = data;
            } else {
                break;
            }
        }

        Some(count)
    }

    /// Sends a single chunk of data, up to the provided `amount`
    pub fn send_one(&mut self, amount: usize) -> Option<Bytes> {
        if self.is_finished() {
            return None;
        }

        let offset = (self.offset % DATA_MOD as u64) as usize;
        let to_send = ((self.len - self.offset) as usize)
            .min(amount)
            .min(DATA.len() - offset);

        if to_send == 0 {
            return Some(Bytes::new());
        }

        let chunk = DATA.slice(offset..offset + to_send);

        self.seek_forward(to_send);

        Some(chunk)
    }

    /// Returns the current offset being received or sent
    pub fn offset(&self) -> u64 {
        self.offset
    }

    /// Moves the current offset forward by the provided `len`
    pub fn seek_forward(&mut self, len: usize) {
        self.offset += len as u64;
    }
}

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

    #[test]
    fn receive_test() {
        let mut receiver = Data::new(10);
        assert!(!receiver.is_finished());

        receiver.receive(&[&[0, 1, 2]]);
        assert!(!receiver.is_finished());

        receiver.receive(&[&[3, 4, 5]]);
        assert!(!receiver.is_finished());

        receiver.receive(&[&[6, 7, 8, 9]]);
        assert!(receiver.is_finished());
    }

    #[test]
    #[should_panic]
    fn receive_corruption_test() {
        let mut receiver = Data::new(10);
        receiver.receive(&[&[1, 2, 3]]);
    }

    #[test]
    #[should_panic]
    fn receive_overflow_test() {
        let mut receiver = Data::new(2);
        receiver.receive(&[&[0, 1, 2]]);
    }

    #[test]
    fn send_test() {
        let mut sender = Data::new(10);
        assert!(!sender.is_finished());

        let mut output = vec![Bytes::new(); 10];

        assert_eq!(sender.send(3, &mut output), Some(1));
        assert_eq!(output[0].as_ref(), &[0, 1, 2][..]);
        assert!(!sender.is_finished());

        assert_eq!(sender.send(4, &mut output), Some(1));
        assert_eq!(output[0].as_ref(), &[3, 4, 5, 6][..]);
        assert!(!sender.is_finished());

        assert_eq!(sender.send(3, &mut output), Some(1));
        assert_eq!(output[0].as_ref(), &[7, 8, 9][..]);
        assert!(sender.is_finished());

        assert_eq!(sender.send(1, &mut output), None);
    }

    #[test]
    #[cfg_attr(miri, ignore)] // This test breaks in CI but can't be reproduced locally - https://github.com/aws/s2n-quic/issues/867
    fn send_receive() {
        let g = (
            (1..(DEFAULT_STREAM_LEN * 16)),
            (1..(DEFAULT_STREAM_LEN * 2)),
        );

        check!()
            .with_generator(g)
            .cloned()
            .for_each(|(stream_len, send_amount)| {
                let mut sender = Data::new(stream_len);
                let mut receiver = Data::new(stream_len);

                let mut buf = vec![Bytes::new(); 5];

                while let Some(count) = sender.send(send_amount as usize, &mut buf[..]) {
                    assert_ne!(count, 0, "sender should return None if no chunks were sent");
                    receiver.receive(&buf[..count]);
                }

                assert!(sender.is_finished());
                assert!(receiver.is_finished());
            })
    }
}


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

/// The Stream Type defines whether data can be transmitted in both directions
/// or only in a single direction.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum StreamType {
    /// Data can be transmitted on the Stream in both directions
    Bidirectional,
    /// Data can be transmitted on the Stream only in a single direction
    Unidirectional,
}

impl StreamType {
    pub fn is_bidirectional(self) -> bool {
        self == StreamType::Bidirectional
    }

    pub fn is_unidirectional(self) -> bool {
        self == StreamType::Unidirectional
    }
}


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

use crate::sync::primitive::{Arc, AtomicBool, Ordering};
use core::{
    pin::Pin,
    task::{Context, Poll, Waker},
};
use crossbeam_utils::CachePadded;

pub use crate::sync::primitive::AtomicWaker;

/// Creates a pair of attached atomic wakers
pub fn pair() -> (Handle, Handle) {
    let storage = Arc::pin(Storage::default());

    let a_ptr = &storage.a as *const _;
    let b_ptr = &storage.b as *const _;
    let is_open = &*storage.is_open as *const _;

    let a = Handle {
        local: a_ptr,
        remote: b_ptr,
        is_open,
        storage: storage.clone(),
    };

    let b = Handle {
        local: b_ptr,
        remote: a_ptr,
        is_open,
        storage: storage.clone(),
    };

    (a, b)
}

/// An attached atomic waker
#[derive(Debug)]
pub struct Handle {
    // store pointers so we don't have to go through the Arc pointer first
    local: *const AtomicWaker,
    remote: *const AtomicWaker,
    is_open: *const AtomicBool,
    #[allow(dead_code)]
    storage: Pin<Arc<Storage>>,
}

/// Safety: Pointers live as long as the storage
unsafe impl Send for Handle {}
/// Safety: Pointers live as long as the storage
unsafe impl Sync for Handle {}

impl Handle {
    /// Registers the local task for notifications from the other handle
    #[inline]
    pub fn register(&self, waker: &Waker) {
        unsafe { (*self.local).register(waker) }
    }

    /// Notifies the other handle that it should be woken up
    #[inline]
    pub fn wake(&self) {
        unsafe { (*self.remote).wake() }
    }

    /// Returns if the handle is open
    ///
    /// If `false` is returned, the other handle has been dropped and not interested in
    /// notifications anymore.
    #[inline]
    pub fn is_open(&self) -> bool {
        unsafe { (*self.is_open).load(Ordering::Acquire) }
    }

    /// Polls the handle until the peer handle has been closed
    #[inline]
    pub fn poll_close(&mut self, cx: &mut Context) -> Poll<()> {
        if !self.is_open() {
            return Poll::Ready(());
        }

        self.register(cx.waker());

        if !self.is_open() {
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }
}

#[derive(Debug)]
struct Storage {
    a: AtomicWaker,
    b: AtomicWaker,
    is_open: CachePadded<AtomicBool>,
}

impl Default for Storage {
    fn default() -> Self {
        Self {
            a: Default::default(),
            b: Default::default(),
            is_open: CachePadded::new(AtomicBool::new(true)),
        }
    }
}

impl Drop for Handle {
    #[inline]
    fn drop(&mut self) {
        // set that we've closed our handle and notify the peer
        unsafe {
            (*self.is_open).store(false, Ordering::Release);
        }
        self.wake();
    }
}


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

use core::{
    marker::PhantomData,
    num::Wrapping,
    ptr::NonNull,
    sync::atomic::{AtomicU32, Ordering},
};

pub struct Builder<T: Copy> {
    pub producer: NonNull<AtomicU32>,
    pub consumer: NonNull<AtomicU32>,
    pub data: NonNull<T>,
    pub size: u32,
}

impl<T: Copy> Builder<T> {
    /// Builds a cursor for a producer
    ///
    /// # Safety
    ///
    /// * This should only be called for the producer
    /// * The pointers should outlive the `Cursor`
    #[inline]
    pub unsafe fn build_producer(self) -> Cursor<T> {
        let mut cursor = self.build();
        cursor.init_producer();
        cursor
    }

    /// Builds a cursor for a consumer
    ///
    /// # Safety
    ///
    /// * This should only be called for the consumer
    /// * The pointers should outlive the `Cursor`
    #[inline]
    pub unsafe fn build_consumer(self) -> Cursor<T> {
        self.build()
    }

    #[inline]
    const fn build(self) -> Cursor<T> {
        let Self {
            producer,
            consumer,
            data,
            size,
        } = self;

        debug_assert!(size.is_power_of_two());

        let mask = size - 1;

        Cursor {
            cached_consumer: Wrapping(0),
            cached_producer: Wrapping(0),
            cached_len: 0,
            size,
            mask,
            producer,
            consumer,
            data,
            entry: PhantomData,
        }
    }
}

/// A structure for tracking a ring shared between a producer and consumer
///
/// See [xsk.h](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/headers/xdp/xsk.h#L34-L42).
#[derive(Debug)]
pub struct Cursor<T: Copy> {
    /// A cached value for the producer cursor index
    ///
    /// This is stored locally to avoid atomic synchronization, if possible
    cached_producer: Wrapping<u32>,
    /// A cached value for the consumer cursor index
    ///
    /// This is stored locally to avoid atomic synchronization, if possible
    cached_consumer: Wrapping<u32>,
    /// A mask value to ensure validity of cursor indexes
    ///
    /// This value assumes that the size of the ring is a power of two
    mask: u32,
    /// The number of entries in the ring
    ///
    /// This value MUST be a power of two
    size: u32,
    /// Points to the producer cursor index
    producer: NonNull<AtomicU32>,
    /// Points to the consumer cursor index
    consumer: NonNull<AtomicU32>,
    /// Points to the values in the ring
    data: NonNull<T>,
    /// A cached value of the computed number of entries for the owner of the `Cursor`
    ///
    /// Since the `acquire` paths are critical to efficiency, we store a derived length to avoid
    /// performing the math over and over again. As such this value needs to be kept in sync with
    /// the `cached_consumer` and `cached_producer`.
    cached_len: u32,
    /// Holds the type of the entries in the ring
    entry: PhantomData<T>,
}

impl<T: Copy> Cursor<T> {
    /// Initializes a producer cursor
    ///
    /// # Safety
    ///
    /// This should only be called by a producer
    #[inline]
    unsafe fn init_producer(&mut self) {
        // increment the consumer cursor by the total size to avoid doing an addition inside
        // `cached_producer`
        //
        // See
        // https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/headers/xdp/xsk.h#L99-L104
        self.cached_consumer += self.size;
        self.cached_len = self.cached_producer_len();

        debug_assert!(self.cached_len <= self.size);
    }

    /// Returns a reference to the producer atomic cursor
    #[inline]
    pub fn producer(&self) -> &AtomicU32 {
        unsafe { &*self.producer.as_ptr() }
    }

    /// Returns a reference to the producer atomic cursor
    #[inline]
    pub fn consumer(&self) -> &AtomicU32 {
        unsafe { &*self.consumer.as_ptr() }
    }

    /// Returns the overall size of the ring
    pub const fn capacity(&self) -> u32 {
        self.size
    }

    /// Acquires a cursor index for a producer half
    ///
    /// The `watermark` can be provided to avoid synchronization by reusing the cached cursor
    /// value.
    ///
    /// See [xsk.h](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/headers/xdp/xsk.h#L92).
    #[inline]
    pub fn acquire_producer(&mut self, watermark: u32) -> u32 {
        // cap the watermark by the max size of the ring to prevent needless loads
        let watermark = watermark.min(self.size);
        let free = self.cached_len;

        // if we have enough space, then return the cached value
        if free >= watermark {
            return free;
        }

        let mut new_value = self.consumer().load(Ordering::Acquire);

        // Our cached copy has the size added so we also need to add the size here when comparing
        //
        // See `Self::init_producer` for more details
        new_value = new_value.wrapping_add(self.size);

        if self.cached_consumer.0 == new_value {
            return free;
        }

        self.cached_consumer.0 = new_value;

        self.cached_len = self.cached_producer_len();

        debug_assert!(self.cached_len <= self.size);

        self.cached_len
    }

    /// Returns the cached producer cursor which is also maxed by the cursor mask
    ///
    /// See [xsk.h](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/headers/xdp/xsk.h#L60).
    #[inline]
    pub fn cached_producer(&self) -> u32 {
        // Wrap the cursor around the size of the ring
        //
        // Masking with a `2^N - 1` value is the same as a mod operation, just more efficient
        self.cached_producer.0 & self.mask
    }

    /// Returns the cached number of available entries for the consumer
    ///
    /// See [xsk.h](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/headers/xdp/xsk.h#L94).
    #[inline]
    pub fn cached_producer_len(&self) -> u32 {
        (self.cached_consumer - self.cached_producer).0
    }

    /// Releases a `len` number of entries from the producer to the consumer.
    ///
    /// See [xsk.h](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/headers/xdp/xsk.h#L135).
    ///
    /// The provided `len` should not exceed the number from `acquire_producer`. With
    /// debug_assertions enabled, this will panic if it occurs.
    #[inline]
    pub fn release_producer(&mut self, len: u32) {
        if cfg!(debug_assertions) {
            let max_len = self.cached_producer_len();
            assert!(max_len >= len, "available: {}, requested: {}", max_len, len);
        }
        self.cached_producer += len;
        self.cached_len -= len;

        debug_assert!(self.cached_len <= self.size);

        self.producer().fetch_add(len, Ordering::Release);
    }

    /// Acquires a cursor index for a consumer half
    ///
    /// The `watermark` can be provided to avoid synchronization by reusing the cached cursor
    /// value.
    ///
    /// See [xsk.h](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/headers/xdp/xsk.h#L112).
    #[inline]
    pub fn acquire_consumer(&mut self, watermark: u32) -> u32 {
        // cap the watermark by the max size of the ring to prevent needless loads
        let watermark = watermark.min(self.size);
        let filled = self.cached_len;

        if filled >= watermark {
            return filled;
        }

        let new_value = self.producer().load(Ordering::Acquire);

        if self.cached_producer.0 == new_value {
            return filled;
        }

        self.cached_producer.0 = new_value;

        self.cached_len = self.cached_consumer_len();

        debug_assert!(self.cached_len <= self.size);

        self.cached_len
    }

    /// Returns the cached consumer cursor which is also maxed by the cursor mask
    ///
    /// See [xsk.h](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/headers/xdp/xsk.h#L68).
    #[inline]
    pub fn cached_consumer(&self) -> u32 {
        // Wrap the cursor around the size of the ring
        //
        // Masking with a `2^N - 1` value is the same as a mod operation, just more efficient
        self.cached_consumer.0 & self.mask
    }

    /// Returns the cached number of available entries for the consumer
    ///
    /// See [xsk.h](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/headers/xdp/xsk.h#L114).
    #[inline]
    pub fn cached_consumer_len(&self) -> u32 {
        (self.cached_producer - self.cached_consumer).0
    }

    /// Releases a `len` number of entries from the consumer to the producer.
    ///
    /// See [xsk.h](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/headers/xdp/xsk.h#L160).
    ///
    /// The provided `len` should not exceed the number from `acquire_consumer`. With
    /// debug_assertions enabled, this will panic if it occurs.
    #[inline]
    pub fn release_consumer(&mut self, len: u32) {
        if cfg!(debug_assertions) {
            let max_len = self.cached_consumer_len();
            assert!(max_len >= len, "available: {}, requested: {}", max_len, len);
        }
        self.cached_consumer += len;
        self.cached_len -= len;

        debug_assert!(self.cached_len <= self.size);

        self.consumer().fetch_add(len, Ordering::Release);
    }

    /// Returns the current consumer entries
    ///
    /// # Safety
    ///
    /// This function MUST only be used by the consumer side.
    #[inline]
    pub unsafe fn consumer_data(&mut self) -> (&mut [T], &mut [T]) {
        let idx = self.cached_consumer();
        let len = self.cached_len;

        debug_assert_eq!(len, self.cached_consumer_len());

        self.mut_slices(idx as _, len as _)
    }

    /// Returns the current producer entries
    ///
    /// # Safety
    ///
    /// This function MUST only be used by the producer side.
    #[inline]
    pub unsafe fn producer_data(&mut self) -> (&mut [T], &mut [T]) {
        let idx = self.cached_producer();
        let len = self.cached_len;

        debug_assert_eq!(len, self.cached_producer_len());

        self.mut_slices(idx as _, len as _)
    }

    #[inline]
    pub const fn data_ptr(&self) -> NonNull<T> {
        self.data
    }

    /// Creates a pair of slices for a given cursor index and len
    #[inline]
    fn mut_slices(&mut self, idx: u64, len: u64) -> (&mut [T], &mut [T]) {
        if len == 0 {
            return (&mut [][..], &mut [][..]);
        }

        let ptr = self.data.as_ptr();

        if let Some(tail_len) = (idx + len).checked_sub(self.size as _) {
            let head_len = self.size as u64 - idx;
            debug_assert_eq!(head_len + tail_len, len);
            let head = unsafe { core::slice::from_raw_parts_mut(ptr.add(idx as _), head_len as _) };
            let tail = unsafe { core::slice::from_raw_parts_mut(ptr, tail_len as _) };
            (head, tail)
        } else {
            let slice = unsafe { core::slice::from_raw_parts_mut(ptr.add(idx as _), len as _) };
            (slice, &mut [][..])
        }
    }
}

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

    #[derive(Clone, Copy, Debug, TypeGenerator)]
    enum Op {
        ConsumerAcquire(u16),
        ConsumerRelease(u16),
        ProducerAcquire(u16),
        ProducerRelease(u16),
    }

    /// Implements a FIFO queue with a monotonic value
    #[derive(Clone, Debug, Default)]
    struct Oracle {
        size: u32,
        producer: u32,
        producer_value: u32,
        consumer: u32,
        consumer_value: u32,
    }

    impl Oracle {
        fn acquire_consumer(&mut self, actual: u32) {
            self.consumer = actual;
            self.invariants();
        }

        fn release_consumer(&mut self, count: u16) -> u32 {
            let count = self.consumer.min(count as u32);

            self.consumer -= count;
            self.consumer_value += count;

            self.invariants();
            count
        }

        fn validate_consumer(&self, (a, b): (&mut [u32], &mut [u32])) {
            for (actual, expected) in a.iter().chain(b.iter()).zip(self.consumer_value..) {
                assert_eq!(
                    expected, *actual,
                    "entry values should match {a:?} {b:?} {self:?}"
                );
            }
        }

        fn acquire_producer(&mut self, actual: u32) {
            self.producer = actual;
            self.invariants();
        }

        fn release_producer(&mut self, count: u16) -> u32 {
            let count = self.producer.min(count as u32);

            self.producer -= count;
            self.producer_value += count;

            self.invariants();
            count
        }

        fn fill_producer(&self, (a, b): (&mut [u32], &mut [u32])) {
            for (entry, value) in a.iter_mut().chain(b).zip(self.producer_value..) {
                *entry = value;
            }
        }

        fn invariants(&self) {
            assert!(
                self.size >= self.producer + self.consumer,
                "The producer and consumer indexes should always be less than the size"
            );
        }
    }

    fn stack_cursors<T, F, R>(init_cursor: u32, desc: &mut [T], exec: F) -> R
    where
        T: Copy,
        F: FnOnce(&mut Cursor<T>, &mut Cursor<T>) -> R,
    {
        let size = desc.len() as u32;
        debug_assert!(size.is_power_of_two());
        let producer_v = UnsafeCell::new(AtomicU32::new(init_cursor));
        let consumer_v = UnsafeCell::new(AtomicU32::new(init_cursor));
        let desc = UnsafeCell::new(desc);

        let producer_v = producer_v.get();
        let consumer_v = consumer_v.get();
        let desc = unsafe { (*desc.get()).as_mut_ptr() as *mut _ };

        let cached_consumer = Wrapping(init_cursor);
        let cached_producer = Wrapping(init_cursor);

        let mut producer: Cursor<T> = unsafe {
            Builder {
                size,
                producer: NonNull::new(producer_v).unwrap(),
                consumer: NonNull::new(consumer_v).unwrap(),
                data: NonNull::new(desc).unwrap(),
            }
            .build_producer()
        };

        producer.cached_consumer = cached_consumer;
        // the producer increments the consumer by `size` to optimize the math so we need to do the
        // same here
        producer.cached_consumer += size;
        producer.cached_producer = cached_producer;
        producer.cached_len = size;

        assert_eq!(producer.acquire_producer(u32::MAX), size);
        assert_eq!(producer.cached_len, producer.cached_producer_len());

        let mut consumer: Cursor<T> = unsafe {
            Builder {
                size,
                producer: NonNull::new(producer_v).unwrap(),
                consumer: NonNull::new(consumer_v).unwrap(),
                data: NonNull::new(desc).unwrap(),
            }
            .build_consumer()
        };

        consumer.cached_consumer = cached_consumer;
        consumer.cached_producer = cached_producer;
        consumer.cached_len = 0;

        assert_eq!(consumer.acquire_consumer(u32::MAX), 0);
        assert_eq!(consumer.cached_len, consumer.cached_consumer_len());

        exec(&mut producer, &mut consumer)
    }

    fn model(power_of_two: u8, init_cursor: u32, ops: &[Op]) {
        let size = (1 << power_of_two) as u32;

        #[cfg(not(kani))]
        let mut desc = vec![u32::MAX; size as usize];

        #[cfg(kani)]
        let mut desc = &mut [u32::MAX; (1 << MAX_POWER_OF_TWO) as usize][..size as usize];

        stack_cursors(init_cursor, &mut desc, |producer, consumer| {
            let mut oracle = Oracle {
                size,
                producer: size,
                ..Default::default()
            };

            for op in ops.iter().copied() {
                oracle.fill_producer(unsafe { producer.producer_data() });

                match op {
                    Op::ConsumerAcquire(count) => {
                        let actual = consumer.acquire_consumer(count as _);
                        oracle.acquire_consumer(actual);
                    }
                    Op::ConsumerRelease(count) => {
                        let oracle_count = oracle.release_consumer(count);
                        consumer.release_consumer(oracle_count);
                    }
                    Op::ProducerAcquire(count) => {
                        let actual = producer.acquire_producer(count as _);
                        oracle.acquire_producer(actual);
                    }
                    Op::ProducerRelease(count) => {
                        let oracle_count = oracle.release_producer(count);
                        producer.release_producer(oracle_count);
                    }
                }

                oracle.validate_consumer(unsafe { consumer.consumer_data() });
            }

            // final assertions
            let actual = consumer.acquire_consumer(u32::MAX);
            oracle.acquire_consumer(actual);
            let data = unsafe { consumer.consumer_data() };
            oracle.validate_consumer(data);
        });
    }

    #[cfg(not(kani))]
    type Ops = Vec<Op>;
    #[cfg(kani)]
    type Ops = crate::testing::InlineVec<Op, 4>;

    const MAX_POWER_OF_TWO: u8 = if cfg!(kani) { 2 } else { 10 };

    #[test]
    #[cfg_attr(miri, ignore)] // this test is too expensive for miri to run
    #[cfg_attr(kani, kani::proof, kani::unwind(5), kani::solver(kissat))]
    fn oracle_test() {
        check!()
            .with_generator((1..=MAX_POWER_OF_TWO, gen(), gen::<Ops>()))
            .for_each(|(power_of_two, init_cursor, ops)| model(*power_of_two, *init_cursor, ops));
    }
}


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

#[cfg(all(loom, test))]
mod loom_primitive {
    use ::core::task::Waker;
    use ::loom::future::AtomicWaker as Inner;

    pub use ::loom::sync::{atomic::*, Arc};

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

    impl AtomicWaker {
        pub fn new() -> Self {
            Self(Inner::new())
        }

        pub fn wake(&self) {
            self.0.wake();
        }

        pub fn take(&self) -> Option<Waker> {
            self.0.take_waker()
        }

        pub fn register(&self, waker: &Waker) {
            self.0.register_by_ref(waker);
        }
    }
}

#[cfg(all(loom, test))]
pub use self::loom_primitive::*;

mod core_primitive {
    pub use ::core::sync::atomic::*;
    pub use alloc::sync::Arc;
    pub use atomic_waker::AtomicWaker;
}

#[cfg(not(all(loom, test)))]
pub use self::core_primitive::*;

/// Indicates if the type is a zero-sized type
///
/// This can be used to optimize the code to avoid needless calculations.
pub trait IsZst {
    const IS_ZST: bool;
}

impl<T> IsZst for T {
    const IS_ZST: bool = ::core::mem::size_of::<T>() == 0;
}


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

mod recv;
mod send;
mod slice;
mod state;

use slice::*;
use state::*;

pub use recv::{Receiver, RecvSlice};
pub use send::{SendSlice, Sender};

#[inline]
pub fn channel<T>(capacity: usize) -> (Sender<T>, Receiver<T>) {
    let state = State::new(capacity);
    let sender = Sender(state.clone());
    let receiver = Receiver(state);
    (sender, receiver)
}

#[cfg(test)]
mod tests;

type Result<T, E = ClosedError> = core::result::Result<T, E>;

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

#[derive(Clone, Copy, Debug)]
pub enum PushError<T> {
    Full(T),
    Closed,
}

impl<T> From<ClosedError> for PushError<T> {
    fn from(_error: ClosedError) -> Self {
        Self::Closed
    }
}


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

use crate::sync::primitive::{Arc, AtomicUsize, AtomicWaker, Ordering};
use core::{
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};
use crossbeam_utils::CachePadded;

/// Creates a worker channel with a Sender and Receiver
pub fn channel() -> (Sender, Receiver) {
    let state = Arc::new(State::default());
    let sender = Sender(state.clone());
    let receiver = Receiver { state, credits: 0 };
    (sender, receiver)
}

/// A handle to the receiver side of the worker channel
///
/// This handle is used by the worker to wake up when there is work to do.
pub struct Receiver {
    state: Arc<State>,
    credits: usize,
}

impl Receiver {
    /// Acquires work to be processed for the Receiver
    ///
    /// `None` is returned when there are no more active Senders.
    #[inline]
    pub async fn acquire(&mut self) -> Option<usize> {
        Acquire(self).await
    }

    /// Polls work to be processed for the receiver
    ///
    /// `None` is returned when there are no more active Senders.
    #[inline]
    pub fn poll_acquire(&mut self, cx: &mut Context) -> Poll<Option<usize>> {
        let state = &*self.state;

        macro_rules! acquire {
            () => {{
                // take the credits that we've been given by the senders
                self.credits += state.remaining.swap(0, Ordering::Acquire);

                // if we have any credits then return
                if self.credits > 0 {
                    return Poll::Ready(Some(self.credits));
                }
            }};
        }

        // first try to acquire credits
        acquire!();

        // if we didn't get any credits then register the waker
        state.receiver.register(cx.waker());

        // make one last effort to acquire credits in case a sender submitted some while we were
        // registering the waker
        acquire!();

        // If we're the only ones with a handle to the state then we're done
        if state.senders.load(Ordering::Acquire) == 0 {
            return Poll::Ready(None);
        }

        Poll::Pending
    }

    /// Marks `count` jobs as finished
    #[inline]
    pub fn finish(&mut self, count: usize) {
        debug_assert!(self.credits >= count);
        // decrement the number of credits we have
        self.credits -= count;
    }
}

/// A handle to submit work to be done to a worker receiver
///
/// Multiple Sender handles can be created with `.clone()`.
#[derive(Clone)]
pub struct Sender(Arc<State>);

impl Sender {
    /// Submits `count` jobs to be executed by the worker receiver
    #[inline]
    pub fn submit(&self, count: usize) {
        let state = &*self.0;

        // increment the work counter
        state.remaining.fetch_add(count, Ordering::Release);

        // wake up the receiver if possible
        state.receiver.wake();
    }
}

impl Drop for Sender {
    #[inline]
    fn drop(&mut self) {
        let state = &*self.0;

        state.senders.fetch_sub(1, Ordering::Release);

        // wake up the receiver to notify that one of the senders has dropped
        state.receiver.wake();
    }
}

struct State {
    remaining: CachePadded<AtomicUsize>,
    receiver: AtomicWaker,
    senders: CachePadded<AtomicUsize>,
}

impl Default for State {
    fn default() -> Self {
        Self {
            remaining: Default::default(),
            receiver: Default::default(),
            senders: AtomicUsize::new(1).into(),
        }
    }
}

struct Acquire<'a>(&'a mut Receiver);

impl<'a> Future for Acquire<'a> {
    type Output = Option<usize>;

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

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

    fn loom_scenario(iterations: usize, send_batch_size: usize, recv_batch_size: usize) {
        assert_ne!(send_batch_size, 0);
        assert_ne!(recv_batch_size, 0);

        loom::model(move || {
            let (send, mut recv) = channel();

            let sender = loom::thread::spawn(move || {
                for _ in 0..iterations {
                    send.submit(send_batch_size);
                    loom::hint::spin_loop();
                }
            });

            let receiver = loom::thread::spawn(move || {
                loom::future::block_on(async move {
                    let mut total = 0;
                    while let Some(mut count) = recv.acquire().await {
                        assert_ne!(count, 0);

                        while count > 0 {
                            let to_finish = count.min(recv_batch_size);
                            recv.finish(to_finish);
                            total += to_finish;
                            count -= to_finish;
                        }
                    }

                    assert_eq!(total, iterations * send_batch_size);
                })
            });

            // loom tests will still run after returning so we don't need to join
            if cfg!(not(loom)) {
                sender.join().unwrap();
                receiver.join().unwrap();
            }
        });
    }

    /// Async loom tests seem to spin forever if the number of iterations is higher than 1.
    /// Ideally, this value would be a bit bigger to test more permutations of orderings.
    const ITERATIONS: usize = if cfg!(loom) { 1 } else { 100 };
    const SEND_BATCH_SIZE: usize = if cfg!(loom) { 2 } else { 8 };
    const RECV_BATCH_SIZE: usize = if cfg!(loom) { 2 } else { 8 };

    #[test]
    fn loom_no_items() {
        loom_scenario(0, 1, 1);
    }

    #[test]
    fn loom_single_item() {
        loom_scenario(ITERATIONS, 1, 1);
    }

    #[test]
    fn loom_send_batch() {
        loom_scenario(ITERATIONS, SEND_BATCH_SIZE, 1);
    }

    #[test]
    fn loom_recv_batch() {
        loom_scenario(ITERATIONS, 1, RECV_BATCH_SIZE);
    }

    #[test]
    fn loom_both_batch() {
        loom_scenario(ITERATIONS, SEND_BATCH_SIZE, RECV_BATCH_SIZE);
    }
}


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

use super::{state::Side, Cursor, Result, State};
use core::{
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};

pub struct Receiver<T>(pub(super) State<T>);

impl<T> Receiver<T> {
    #[inline]
    pub fn capacity(&self) -> usize {
        self.0.cursor.capacity()
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.0.cursor.recv_len()
    }

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

    #[inline]
    pub fn is_full(&self) -> bool {
        self.0.cursor.is_full()
    }

    /// Returns the currently acquired slice of entries for the receiver
    ///
    /// Callers should call [`Self::acquire`] or [`Self::poll_slice`] before calling this method.
    #[inline]
    pub fn slice(&mut self) -> RecvSlice<T> {
        let cursor = self.0.cursor;
        RecvSlice(&mut self.0, cursor)
    }

    /// Blocks until at least one entry is available for consumption
    #[inline]
    pub async fn acquire(&mut self) -> Result<()> {
        Acquire { receiver: self }.await
    }

    #[inline]
    pub fn poll_slice(&mut self, cx: &mut Context) -> Poll<Result<RecvSlice<T>>> {
        macro_rules! acquire_filled {
            () => {
                match self.0.acquire_filled() {
                    Ok(true) => {
                        let cursor = self.0.cursor;
                        return Ok(RecvSlice(&mut self.0, cursor)).into();
                    }
                    Ok(false) => {
                        // the queue is full
                    }
                    Err(err) => {
                        // the channel was closed
                        return Err(err).into();
                    }
                }
            };
        }

        // check capacity before registering a waker
        acquire_filled!();

        // register the waker
        self.0.receiver.register(cx.waker());

        // check once more to avoid a loss of notification
        acquire_filled!();

        Poll::Pending
    }

    #[inline]
    pub fn try_slice(&mut self) -> Result<Option<RecvSlice<T>>> {
        Ok(if self.0.acquire_filled()? {
            let cursor = self.0.cursor;
            Some(RecvSlice(&mut self.0, cursor))
        } else {
            None
        })
    }
}

impl<T> Drop for Receiver<T> {
    #[inline]
    fn drop(&mut self) {
        self.0.close(Side::Receiver);
    }
}

pub struct RecvSlice<'a, T>(&'a mut State<T>, Cursor);

impl<'a, T> RecvSlice<'a, T> {
    #[inline]
    pub fn peek(&mut self) -> (&mut [T], &mut [T]) {
        let _ = self.0.acquire_filled();
        let (slice, _) = self.0.as_pairs();
        unsafe {
            // Safety: the first pair of returned slices is the `initialized` half
            slice.assume_init().into_mut()
        }
    }

    #[inline]
    pub fn pop(&mut self) -> Option<T> {
        if self.0.cursor.is_empty() && !self.0.acquire_filled().unwrap_or(false) {
            return None;
        }

        let (pair, _) = self.0.as_pairs();
        let value = unsafe {
            // Safety: the state's cursor indicates that the first slot contains initialized data
            pair.take(0)
        };
        self.0.cursor.increment_head(1);
        Some(value)
    }

    #[inline]
    pub fn clear(&mut self) -> usize {
        // don't try to `acquire_filled` so the caller can observe any updates through peek/pop

        let (pair, _) = self.0.as_pairs();
        let len = pair.len();

        for entry in pair.iter() {
            unsafe {
                // Safety: the state's cursor indicates that each slot in the `iter` contains data
                let _ = entry.take();
            }
        }

        self.0.cursor.increment_head(len);

        len
    }

    /// Releases `len` entries back to the sender
    #[inline]
    pub fn release(&mut self, len: usize) {
        let (pair, _) = self.0.as_pairs();

        debug_assert!(pair.len() >= len, "cannot release more than was acquired");

        for entry in pair.iter().take(len) {
            unsafe {
                // Safety: the state's cursor indicates that each slot in the `iter` contains data
                let _ = entry.take();
            }
        }

        self.0.cursor.increment_head(len);
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.0.cursor.recv_len()
    }

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

    /// Synchronizes any updates from the sender
    ///
    /// This can be useful for when `slice` is called without polling for entries first.
    #[inline]
    pub fn sync(&mut self) -> Result<(), super::ClosedError> {
        self.0.acquire_filled()?;
        Ok(())
    }
}

impl<'a, T> Drop for RecvSlice<'a, T> {
    #[inline]
    fn drop(&mut self) {
        self.0.persist_head(self.1);
    }
}

struct Acquire<'a, T> {
    receiver: &'a mut Receiver<T>,
}

impl<'a, T> Future for Acquire<'a, T> {
    type Output = Result<()>;

    #[inline]
    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        match self.receiver.poll_slice(cx) {
            Poll::Ready(v) => Poll::Ready(v.map(|_| ())),
            Poll::Pending => Poll::Pending,
        }
    }
}


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

use super::{state::Side, Cursor, PushError, Result, State};
use core::{
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};

pub struct Sender<T>(pub(super) State<T>);

impl<T> Sender<T> {
    #[inline]
    pub fn capacity(&self) -> usize {
        self.0.cursor.capacity()
    }

    /// Returns the currently acquired slice of entries for the sender
    ///
    /// Callers should call [`Self::acquire`] or [`Self::poll_slice`] before calling this method.
    #[inline]
    pub fn slice(&mut self) -> SendSlice<T> {
        let cursor = self.0.cursor;
        SendSlice(&mut self.0, cursor)
    }

    /// Blocks until at least one entry is available for sending
    #[inline]
    pub async fn acquire(&mut self) -> Result<()> {
        Acquire { sender: self }.await
    }

    #[inline]
    pub fn poll_slice(&mut self, cx: &mut Context) -> Poll<Result<SendSlice<T>>> {
        macro_rules! acquire_capacity {
            () => {
                match self.0.acquire_capacity() {
                    Ok(true) => {
                        let cursor = self.0.cursor;
                        return Ok(SendSlice(&mut self.0, cursor)).into();
                    }
                    Ok(false) => {
                        // the queue is full
                    }
                    Err(err) => {
                        // the channel was closed
                        return Err(err).into();
                    }
                }
            };
        }

        // check capacity before registering a waker
        acquire_capacity!();

        // register the waker
        self.0.sender.register(cx.waker());

        // check once more to avoid a loss of notification
        acquire_capacity!();

        Poll::Pending
    }

    #[inline]
    pub fn try_slice(&mut self) -> Result<Option<SendSlice<T>>> {
        Ok(if self.0.acquire_capacity()? {
            let cursor = self.0.cursor;
            Some(SendSlice(&mut self.0, cursor))
        } else {
            None
        })
    }
}

impl<T> Drop for Sender<T> {
    #[inline]
    fn drop(&mut self) {
        self.0.close(Side::Sender);
    }
}

pub struct SendSlice<'a, T>(&'a mut State<T>, Cursor);

impl<'a, T> SendSlice<'a, T> {
    #[inline]
    pub fn push(&mut self, value: T) -> Result<(), PushError<T>> {
        if self.0.cursor.is_full() && !self.0.acquire_capacity()? {
            return Err(PushError::Full(value));
        }

        let (_, pair) = self.0.as_pairs();

        unsafe {
            // Safety: the second pair of slices contains uninitialized memory and the cursor
            // indicates we have capacity to write at least one value
            pair.write(0, value);
        }

        self.0.cursor.increment_tail(1);

        Ok(())
    }

    pub fn extend<I: Iterator<Item = T>>(&mut self, iter: &mut I) -> Result<()> {
        if self.0.acquire_capacity()? {
            let (_, pair) = self.0.as_pairs();

            let mut idx = 0;
            let capacity = self.capacity();

            while idx < capacity {
                if let Some(value) = iter.next() {
                    unsafe {
                        // Safety: the second pair of slices contains uninitialized memory
                        pair.write(idx, value);
                    }
                    idx += 1;
                } else {
                    break;
                }
            }

            self.0.cursor.increment_tail(idx);
        }

        Ok(())
    }

    #[inline]
    pub fn capacity(&self) -> usize {
        self.0.cursor.send_capacity()
    }

    /// Synchronizes any updates from the receiver
    ///
    /// This can be useful for when `slice` is called without polling for entries first.
    #[inline]
    pub fn sync(&mut self) -> Result<(), super::ClosedError> {
        self.0.acquire_capacity()?;
        Ok(())
    }
}

impl<'a, T> Drop for SendSlice<'a, T> {
    #[inline]
    fn drop(&mut self) {
        self.0.persist_tail(self.1);
    }
}

struct Acquire<'a, T> {
    sender: &'a mut Sender<T>,
}

impl<'a, T> Future for Acquire<'a, T> {
    type Output = Result<()>;

    #[inline]
    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        match self.sender.poll_slice(cx) {
            Poll::Ready(v) => Poll::Ready(v.map(|_| ())),
            Poll::Pending => Poll::Pending,
        }
    }
}


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

use core::{cell::UnsafeCell, mem::MaybeUninit, ops::Deref};

#[repr(transparent)]
pub struct Cell<T>(MaybeUninit<UnsafeCell<T>>);

impl<T> Cell<T> {
    #[inline]
    pub unsafe fn write(&self, value: T) {
        UnsafeCell::raw_get(self.0.as_ptr()).write(value);
    }

    #[inline]
    pub unsafe fn take(&self) -> T {
        self.0.assume_init_ref().get().read()
    }
}

#[derive(Debug)]
pub struct Slice<'a, T>(pub(super) &'a [T]);

impl<'a, T> Slice<'a, Cell<T>> {
    /// Assumes that the slice of [`Cell`]s is initialized and converts it to a slice of
    /// [`UnsafeCell`]s.
    ///
    /// See [`core::mem::MaybeUninit::assume_init`]
    #[inline]
    pub unsafe fn assume_init(self) -> Slice<'a, UnsafeCell<T>> {
        Slice(&*(self.0 as *const [Cell<T>] as *const [UnsafeCell<T>]))
    }

    /// Writes a value into a cell at the provided index
    ///
    /// # Safety
    ///
    /// The cell at `index` must be uninitialized and the caller must have synchronized access.
    #[inline]
    pub unsafe fn write(&self, index: usize, value: T) {
        self.0.get_unchecked(index).write(value)
    }

    /// Reads and takes the memory at a cell at the provided index
    ///
    /// # Safety
    ///
    /// The cell at `index` must be initialized and the caller must have synchronized access.
    #[inline]
    pub unsafe fn take(&self, index: usize) -> T {
        self.0.get_unchecked(index).take()
    }
}

impl<'a, T> Slice<'a, UnsafeCell<T>> {
    /// Converts the slice of [`UnsafeCell`]s into a mutable slice
    ///
    /// # Safety
    ///
    /// The slice must be exclusively owned, otherwise data races may occur.
    #[inline]
    pub unsafe fn into_mut(self) -> &'a mut [T] {
        let ptr = self.0.as_ptr() as *mut T;
        let len = self.0.len();
        core::slice::from_raw_parts_mut(ptr, len)
    }
}

impl<'a, T> Deref for Slice<'a, T> {
    type Target = [T];

    #[inline]
    fn deref(&self) -> &[T] {
        self.0
    }
}

impl<'a, T: PartialEq> PartialEq<[T]> for Slice<'a, UnsafeCell<T>> {
    #[inline]
    fn eq(&self, other: &[T]) -> bool {
        if self.len() != other.len() {
            return false;
        }

        for (a, b) in self.iter().zip(other) {
            if unsafe { &*a.get() } != b {
                return false;
            }
        }

        true
    }
}

impl<'a, T: PartialEq> PartialEq<Slice<'a, UnsafeCell<T>>> for [T] {
    #[inline]
    fn eq(&self, other: &Slice<'a, UnsafeCell<T>>) -> bool {
        other.eq(self)
    }
}

impl<'a, T: PartialEq> PartialEq<Slice<'a, UnsafeCell<T>>> for &[T] {
    #[inline]
    fn eq(&self, other: &Slice<'a, UnsafeCell<T>>) -> bool {
        other.eq(self)
    }
}

#[derive(Debug)]
pub struct Pair<S> {
    pub head: S,
    pub tail: S,
}

impl<'a, T> Pair<Slice<'a, Cell<T>>> {
    #[inline]
    pub unsafe fn assume_init(self) -> Pair<Slice<'a, UnsafeCell<T>>> {
        Pair {
            head: self.head.assume_init(),
            tail: self.tail.assume_init(),
        }
    }

    #[inline]
    pub unsafe fn write(&self, index: usize, value: T) {
        self.cell(index).write(value)
    }

    #[inline]
    pub unsafe fn take(&self, index: usize) -> T {
        self.cell(index).take()
    }

    unsafe fn cell(&self, index: usize) -> &Cell<T> {
        if let Some(cell) = self.head.0.get(index) {
            cell
        } else {
            assume!(
                index >= self.head.0.len(),
                "index must always be equal or greater than the `head` len"
            );
            let index = index - self.head.0.len();

            assume!(
                self.tail.get(index).is_some(),
                "index must be in-bounds for the `tail` slice: head={}, tail={}, index={}",
                self.head.0.len(),
                self.tail.0.len(),
                index
            );
            self.tail.get_unchecked(index)
        }
    }

    #[inline]
    pub fn iter(&self) -> impl Iterator<Item = &Cell<T>> {
        self.head.0.iter().chain(self.tail.0)
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.head.len() + self.tail.len()
    }
}

impl<'a, T> Pair<Slice<'a, UnsafeCell<T>>> {
    #[inline]
    pub unsafe fn into_mut(self) -> (&'a mut [T], &'a mut [T]) {
        let head = self.head.into_mut();
        let tail = self.tail.into_mut();
        (head, tail)
    }
}


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

use super::{Cell, ClosedError, Result, Slice};
use crate::sync::primitive::{AtomicBool, AtomicUsize, AtomicWaker, IsZst, Ordering};
use alloc::alloc::Layout;
use core::{
    fmt,
    marker::PhantomData,
    ops::Deref,
    panic::{RefUnwindSafe, UnwindSafe},
    ptr::NonNull,
};
use crossbeam_utils::CachePadded;

type Pair<'a, T> = super::Pair<Slice<'a, Cell<T>>>;

const MINIMUM_CAPACITY: usize = 2;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Side {
    Sender,
    Receiver,
}

#[derive(Clone, Copy)]
pub struct Cursor {
    head: usize,
    tail: usize,
    capacity: usize,
}

impl fmt::Debug for Cursor {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Cursor")
            .field("head", &self.head)
            .field("tail", &self.tail)
            .field("len", &self.len())
            .field("capacity", &self.capacity())
            .field("is_empty", &self.is_empty())
            .field("is_full", &self.is_full())
            .field("is_contiguous", &self.is_contiguous())
            .finish()
    }
}

impl Cursor {
    #[inline]
    fn new(capacity: usize) -> Self {
        Self {
            head: 0,
            tail: 0,
            capacity,
        }
    }

    #[inline]
    fn invariants(&self) {
        unsafe {
            assume!(
                self.capacity >= MINIMUM_CAPACITY,
                "the capacity must be at least the MINIMUM_CAPACITY value"
            );
            assume!(
                self.head < self.capacity,
                "the `head` pointer should be strictly less than the capacity"
            );
            assume!(
                self.tail < self.capacity,
                "the `tail` pointer should be strictly less than the capacity"
            );
            let len = count(self.head, self.tail, self.capacity);
            assume!(
                len < self.capacity,
                "the computed `len` should be strictly less than the capacity"
            );
        }
    }

    #[inline]
    pub fn capacity(&self) -> usize {
        self.invariants();
        // To make cursor management easier, we never allow the callers to hit the total capacity.
        // We also account for this when allocating the state by adding 1 to the request capacity.
        self.capacity - 1
    }

    #[inline]
    fn cap(&self) -> usize {
        self.invariants();
        self.capacity
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.invariants();
        count(self.head, self.tail, self.cap())
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.invariants();
        self.tail == self.head
    }

    #[inline]
    pub fn is_full(&self) -> bool {
        self.invariants();
        count(self.tail, self.head, self.cap()) == 1
    }

    #[inline]
    pub fn recv_len(&self) -> usize {
        self.invariants();
        self.len()
    }

    #[inline]
    pub fn send_capacity(&self) -> usize {
        self.invariants();
        self.capacity() - self.recv_len()
    }

    #[inline]
    pub fn increment_head(&mut self, n: usize) {
        self.invariants();
        unsafe {
            assume!(
                n <= self.capacity(),
                "n should never exceed the total capacity"
            );
        }
        self.head = self.wrap_add(self.head, n);
        self.invariants();
    }

    #[inline]
    pub fn increment_tail(&mut self, n: usize) {
        self.invariants();
        unsafe {
            assume!(
                n <= self.capacity(),
                "n should never exceed the total capacity"
            );
        }
        self.tail = self.wrap_add(self.tail, n);
        self.invariants();
    }

    #[inline]
    fn wrap_add(&self, idx: usize, addend: usize) -> usize {
        wrap_index(idx.wrapping_add(addend), self.cap())
    }

    #[inline]
    fn is_contiguous(&self) -> bool {
        self.tail >= self.head
    }
}

/// Returns the index in the underlying buffer for a given logical element index.
#[inline]
fn wrap_index(index: usize, size: usize) -> usize {
    // size is always a power of 2
    unsafe {
        assume!(
            size.is_power_of_two(),
            "The calculations in the lengths rely on the capacity being a power of 2"
        );
        assume!(
            size >= MINIMUM_CAPACITY,
            "The calculations in the lengths rely on the capacity being at least {}",
            MINIMUM_CAPACITY
        );
    }
    index & (size - 1)
}

/// Calculate the number of elements left to be read in the buffer
#[inline]
fn count(head: usize, tail: usize, size: usize) -> usize {
    // size is always a power of 2
    unsafe {
        assume!(
            size.is_power_of_two(),
            "The calculations in the lengths rely on the capacity being a power of 2"
        );
        assume!(
            size >= MINIMUM_CAPACITY,
            "The calculations in the lengths rely on the capacity being at least {}",
            MINIMUM_CAPACITY
        );
    }
    (tail.wrapping_sub(head)) & (size - 1)
}

/// The synchronized state between two peers
///
/// The internal design of the cursor management is based on [`alloc::collections::VecDeque`].
pub struct State<T> {
    header: NonNull<Header<T>>,
    pub cursor: Cursor,
}

impl<T> fmt::Debug for State<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("State")
            .field("header", self.deref())
            .field("cursor", &self.cursor)
            .finish()
    }
}

/// Safety: synchronization of state is managed through atomic values
unsafe impl<T: Send> Send for State<T> {}

/// Safety: synchronization of state is managed through atomic values
unsafe impl<T: Sync> Sync for State<T> {}

/// The data behind the header pointer itself is unwind safe
impl<T: RefUnwindSafe> UnwindSafe for State<T> {}

impl<T> Clone for State<T> {
    #[inline]
    fn clone(&self) -> Self {
        Self {
            header: self.header,
            cursor: self.cursor,
        }
    }
}

impl<T> Deref for State<T> {
    type Target = Header<T>;

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

impl<T> State<T> {
    #[inline]
    pub fn new(capacity: usize) -> Self {
        // If we're sending a zero-sized type, set the capacity to the maximum value, since we're
        // not sending any data and just coordinating cursors at this point
        let capacity = if T::IS_ZST {
            // The total capacity must be a power of two
            usize::MAX / 2 + 1
        } else {
            // Add 1 to the requested capacity so it's easier to manage cursor wrapping
            core::cmp::max(capacity + 1, MINIMUM_CAPACITY).next_power_of_two()
        };
        let header = Header::alloc(capacity).expect("could not allocate channel");
        let cursor = Cursor::new(capacity);
        Self { header, cursor }
    }

    /// Tries to acquire more unfilled slots on the channel
    ///
    /// If the channel is closed, an error is returned. If the channel has at least one slot of
    /// capacity, `true` is returned. Otherwise `false` is returned.
    #[inline]
    pub fn acquire_capacity(&mut self) -> Result<bool> {
        if !self.open.load(Ordering::Acquire) {
            return Err(ClosedError);
        }

        // update the cached version
        self.cursor.head = self.head.load(Ordering::Acquire);

        let is_full = self.cursor.is_full();

        Ok(!is_full)
    }

    /// Tries to acquire more filled slots on the channel
    ///
    /// If the channel is closed, an error is returned. If the channel has at least one slot of
    /// capacity, `true` is returned. Otherwise `false` is returned.
    #[inline]
    pub fn acquire_filled(&mut self) -> Result<bool> {
        self.cursor.tail = self.tail.load(Ordering::Acquire);

        if !self.cursor.is_empty() {
            return Ok(true);
        }

        if !self.open.load(Ordering::Acquire) {
            // make one more effort to load the remaining items
            self.cursor.tail = self.tail.load(Ordering::Acquire);

            if !self.cursor.is_empty() {
                return Ok(true);
            }

            return Err(ClosedError);
        }

        Ok(false)
    }

    /// Notifies the peer of `head` updates for the given cursor
    #[inline]
    pub fn persist_head(&self, prev: Cursor) {
        // nothing changed
        if prev.head == self.cursor.head {
            return;
        }

        self.head.store(self.cursor.head, Ordering::Release);

        self.sender.wake();
    }

    /// Notifies the peer of `tail` updates for the given cursor
    #[inline]
    pub fn persist_tail(&self, prev: Cursor) {
        // nothing changed
        if prev.tail == self.cursor.tail {
            return;
        }

        self.tail.store(self.cursor.tail, Ordering::Release);

        self.receiver.wake();
    }

    #[inline]
    fn data(&self) -> &[Cell<T>] {
        unsafe {
            // Safety: the state must still be allocated and the cursor inbounds
            let ptr = self.data_ptr();
            let capacity = self.cursor.capacity;
            core::slice::from_raw_parts(ptr, capacity)
        }
    }

    #[inline]
    fn data_ptr(&self) -> *const Cell<T> {
        unsafe {
            // If the type is zero-sized, no need to calculate offsets
            if T::IS_ZST {
                return NonNull::<Cell<T>>::dangling().as_ptr();
            }

            // Safety: the state must still be allocated and the cursor inbounds
            let capacity = self.cursor.capacity;
            let (_, offset) = Header::<T>::layout_unchecked(capacity);

            let ptr = self.header.as_ptr() as *const u8;
            let ptr = ptr.add(offset);
            ptr as *const Cell<T>
        }
    }

    /// Closes one side of the channel and notifies the peer of the event
    #[inline]
    pub fn close(&mut self, side: Side) {
        // notify the other side that we've closed the channel
        match side {
            Side::Sender => self.receiver.wake(),
            Side::Receiver => self.sender.wake(),
        }

        let was_open = self.open.swap(false, Ordering::SeqCst);

        // make sure the peer is notified before fully dropping the contents
        match side {
            Side::Sender => self.receiver.wake(),
            Side::Receiver => self.sender.wake(),
        }

        if !was_open {
            unsafe {
                // Safety: we synchronization closing between the two peers through atomic
                // variables. At this point both sides have agreed on its final state.
                self.drop_contents();
            }
        }
    }

    /// Returns the channel slots as two pairs of filled and unfilled slices
    #[inline]
    pub fn as_pairs(&self) -> (Pair<T>, Pair<T>) {
        let data = self.data();
        self.data_to_pairs(data)
    }

    #[inline]
    fn data_to_pairs<'a>(&self, data: &'a [Cell<T>]) -> (Pair<'a, T>, Pair<'a, T>) {
        self.cursor.invariants();

        let head = self.cursor.head;
        let tail = self.cursor.tail;

        let (filled, unfilled) = if self.cursor.is_contiguous() {
            unsafe {
                assume!(data.len() >= tail, "data must span the tail length");
            }
            let (data, unfilled_head) = data.split_at(tail);

            unsafe {
                assume!(data.len() >= head, "data must span the head length");
            }
            let (unfilled_tail, filled_head) = data.split_at(head);

            let filled = Pair {
                head: Slice(filled_head),
                tail: Slice(&[]),
            };
            let unfilled = Pair {
                head: Slice(unfilled_head),
                tail: Slice(unfilled_tail),
            };
            (filled, unfilled)
        } else {
            unsafe {
                assume!(data.len() >= head, "data must span the head length");
            }
            let (data, filled_head) = data.split_at(head);

            unsafe {
                assume!(data.len() >= tail, "data must span the tail length");
            }
            let (filled_tail, unfilled_head) = data.split_at(tail);

            let filled = Pair {
                head: Slice(filled_head),
                tail: Slice(filled_tail),
            };
            let unfilled = Pair {
                head: Slice(unfilled_head),
                tail: Slice(&[]),
            };
            (filled, unfilled)
        };

        unsafe {
            assume!(
                filled.len() == self.cursor.recv_len(),
                "filled len should agree with the cursor len {} == {}\n{:?}",
                filled.len(),
                self.cursor.recv_len(),
                self.cursor
            );
        }

        (filled, unfilled)
    }

    /// Frees the contents of the channel
    ///
    /// # Safety
    ///
    /// Each side must have synchronized and agreed on the final state before calling this
    #[inline]
    unsafe fn drop_contents(&mut self) {
        // refresh the cursor from the shared state
        self.cursor.head = self.head.load(Ordering::Acquire);
        self.cursor.tail = self.tail.load(Ordering::Acquire);

        // release all of the filled data
        let (filled, _unfilled) = self.as_pairs();
        if !T::IS_ZST {
            for cell in filled.iter() {
                drop(cell.take());
            }
        }

        // make sure we free any stored wakers
        let header = self.header.as_mut();
        drop(header.receiver.take());
        drop(header.sender.take());

        // free the header
        let ptr = self.header.as_ptr() as *mut u8;
        let capacity = self.cursor.capacity;
        let (layout, _offset) = Header::<T>::layout_unchecked(capacity);
        alloc::alloc::dealloc(ptr, layout)
    }
}

pub struct Header<T> {
    head: CachePadded<AtomicUsize>,
    tail: CachePadded<AtomicUsize>,
    open: CachePadded<AtomicBool>,
    pub receiver: AtomicWaker,
    pub sender: AtomicWaker,
    data: PhantomData<T>,
}

impl<T> fmt::Debug for Header<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Header")
            .field("head", &self.head.load(Ordering::Relaxed))
            .field("tail", &self.tail.load(Ordering::Relaxed))
            .field("open", &self.open.load(Ordering::Relaxed))
            .field("receiver", &self.receiver)
            .field("sender", &self.sender)
            .finish()
    }
}

impl<T> Header<T> {
    /// Allocates a header and data slice for the given capacity
    fn alloc(capacity: usize) -> Option<NonNull<Self>> {
        unsafe {
            // Safety: we assume that `alloc` gives us a valid pointer to write to
            let (layout, _offset) = Self::layout(capacity).ok()?;
            let state = alloc::alloc::alloc(layout);
            let state = state as *mut Self;
            let state = NonNull::new(state)?;

            state.as_ptr().write(Self::new());

            Some(state)
        }
    }

    #[inline]
    fn new() -> Self {
        Self {
            head: CachePadded::new(AtomicUsize::new(0)),
            tail: CachePadded::new(AtomicUsize::new(0)),
            sender: AtomicWaker::new(),
            receiver: AtomicWaker::new(),
            open: CachePadded::new(AtomicBool::new(true)),
            data: PhantomData,
        }
    }

    /// Computes the checked layout for the header
    #[inline]
    fn layout(capacity: usize) -> Result<(Layout, usize), alloc::alloc::LayoutError> {
        let header_layout = Layout::new::<Self>();
        // A slice of cells is allocated in the same region as the header
        let data_layout = Layout::array::<Cell<T>>(capacity)?;
        let (layout, offset) = header_layout.extend(data_layout)?;
        Ok((layout, offset))
    }

    /// Computes the memory layout of the header without checking of its validatity
    ///
    /// # Safety
    ///
    /// The layout must have been previously checked before calling this.
    #[inline]
    unsafe fn layout_unchecked(capacity: usize) -> (Layout, usize) {
        if let Ok(v) = Self::layout(capacity) {
            v
        } else {
            core::hint::unreachable_unchecked()
        }
    }
}


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

use super::*;
use crate::testing::loom;
use bolero::{check, generator::*};
use core::task::{Context, Poll, Waker};
use futures_test::task::{new_count_waker, AwokenCount};
use std::collections::VecDeque;

#[derive(Clone, Copy, Debug, TypeGenerator)]
enum Operation {
    Push(u16),
    AsyncPush(u16),
    Pop(u16),
    AsyncPop(u16),
    Clear,
    DropSend,
    DropRecv,
}

struct WakeState {
    waker: Waker,
    wake_count: AwokenCount,
    snapshot: Option<usize>,
}

impl Default for WakeState {
    fn default() -> Self {
        let (waker, wake_count) = new_count_waker();
        Self {
            waker,
            wake_count,
            snapshot: None,
        }
    }
}

impl WakeState {
    fn context(&self) -> Context {
        Context::from_waker(&self.waker)
    }

    fn snapshot(&mut self) {
        self.snapshot = Some(self.wake_count.get());
    }

    fn assert_wake(&mut self) {
        if let Some(prev) = self.snapshot.take() {
            assert_eq!(self.wake_count.get(), prev + 1);
        }
    }
}

struct Model<T> {
    oracle: VecDeque<T>,
    send: Option<Sender<T>>,
    send_waker: WakeState,
    recv: Option<Receiver<T>>,
    recv_waker: WakeState,
    capacity: usize,
}

impl<T: Clone + PartialEq + core::fmt::Debug> Model<T> {
    fn new(capacity: usize) -> Self {
        let (send, recv) = channel(capacity);
        let capacity = send.capacity();
        Self {
            oracle: Default::default(),
            send: Some(send),
            send_waker: Default::default(),
            recv: Some(recv),
            recv_waker: Default::default(),
            capacity,
        }
    }

    fn apply_all(&mut self, operations: &[Operation], mut generator: impl FnMut() -> T) {
        for op in operations {
            self.apply(*op, &mut generator);
        }
    }

    fn apply(&mut self, operation: Operation, mut generator: impl FnMut() -> T) {
        match operation {
            Operation::Push(count) => {
                let mut did_push = false;
                if let Some(send) = self.send.as_mut() {
                    match send.try_slice() {
                        Ok(Some(mut slice)) => {
                            for _ in 0..count {
                                let value = generator();
                                if slice.push(value.clone()).is_ok() {
                                    self.oracle.push_back(value);
                                    did_push = true;
                                }
                            }
                        }
                        Ok(None) => {
                            assert_eq!(
                                self.oracle.len(),
                                self.capacity,
                                "slice should return None when at capacity"
                            );
                        }
                        Err(_) => {
                            assert!(self.recv.is_none());
                        }
                    }
                }

                if did_push {
                    self.recv_waker.assert_wake();
                }
            }
            Operation::AsyncPush(count) => {
                let mut did_push = false;
                if let Some(send) = self.send.as_mut() {
                    match send.poll_slice(&mut self.send_waker.context()) {
                        Poll::Ready(Ok(mut slice)) => {
                            for _ in 0..count {
                                let value = generator();
                                if slice.push(value.clone()).is_ok() {
                                    self.oracle.push_back(value);

                                    did_push = true;
                                }
                            }
                        }
                        Poll::Ready(Err(_)) => {
                            assert!(self.recv.is_none());
                        }
                        Poll::Pending => {
                            assert_eq!(
                                self.oracle.len(),
                                self.capacity,
                                "slice should return Pending when at capacity"
                            );
                            self.send_waker.snapshot();
                        }
                    }
                }

                if did_push {
                    self.recv_waker.assert_wake();
                }
            }
            Operation::Pop(count) => {
                let mut did_pop = false;
                if let Some(recv) = self.recv.as_mut() {
                    match recv.try_slice() {
                        Ok(Some(mut slice)) => {
                            for _ in 0..count {
                                let value = slice.pop();
                                assert_eq!(value, self.oracle.pop_front());
                                did_pop |= value.is_some();
                            }
                        }
                        Ok(None) => {
                            assert!(self.oracle.is_empty());
                        }
                        Err(_) => {
                            assert!(self.send.is_none());
                            assert!(self.oracle.is_empty());
                        }
                    }
                }

                if did_pop {
                    self.send_waker.assert_wake();
                }
            }
            Operation::AsyncPop(count) => {
                let mut did_pop = false;
                if let Some(recv) = self.recv.as_mut() {
                    match recv.poll_slice(&mut self.recv_waker.context()) {
                        Poll::Ready(Ok(mut slice)) => {
                            for _ in 0..count {
                                let value = slice.pop();
                                assert_eq!(value, self.oracle.pop_front());
                                did_pop |= value.is_some();
                            }
                        }
                        Poll::Ready(Err(_)) => {
                            assert!(self.send.is_none());
                            assert!(self.oracle.is_empty());
                        }
                        Poll::Pending => {
                            assert!(self.oracle.is_empty());
                            self.recv_waker.snapshot();
                        }
                    }
                }

                if did_pop {
                    self.send_waker.assert_wake();
                }
            }
            Operation::Clear => {
                let mut did_pop = false;
                if let Some(recv) = self.recv.as_mut() {
                    match recv.try_slice() {
                        Ok(Some(mut slice)) => {
                            // we need to pull in the latest values to clear everything
                            while slice.clear() > 0 {
                                did_pop = true;
                                let _ = slice.peek();
                            }
                            self.oracle.clear();
                        }
                        Ok(None) => {
                            assert!(self.oracle.is_empty());
                        }
                        Err(_) => {
                            assert!(self.send.is_none());
                            assert!(self.oracle.is_empty());
                        }
                    }
                }

                if did_pop {
                    self.send_waker.assert_wake();
                }
            }
            Operation::DropSend => {
                self.send = None;
            }
            Operation::DropRecv => {
                self.recv = None;
            }
        }
    }

    fn finish(mut self) {
        loop {
            self.apply(Operation::Pop(u16::MAX), || unimplemented!());
            if self.oracle.is_empty() || self.recv.is_none() {
                return;
            }
        }
    }
}

#[cfg(any(kani, miri))]
type Operations = crate::testing::InlineVec<Operation, 2>;
#[cfg(not(any(kani, miri)))]
type Operations = Vec<Operation>;

#[test]
fn model() {
    let max_capacity = if cfg!(any(kani, miri)) { 2 } else { 64 };

    let generator = (1usize..max_capacity, gen::<Operations>());

    check!()
        .with_generator(generator)
        .for_each(|(capacity, ops)| {
            let mut model = Model::new(*capacity);
            let mut cursor = 0;
            let generator = || {
                let v = cursor;
                cursor += 1;
                Box::new(v)
            };

            model.apply_all(ops, generator);
            model.finish();
        })
}

#[test]
fn model_zst() {
    let max_capacity = if cfg!(any(kani, miri)) { 2 } else { 64 };

    let generator = (1usize..max_capacity, gen::<Operations>());

    check!()
        .with_generator(generator)
        .for_each(|(capacity, ops)| {
            let mut model = Model::new(*capacity);
            let generator = || ();

            model.apply_all(ops, generator);
            model.finish();
        })
}

#[test]
#[cfg_attr(kani, kani::proof, kani::unwind(3), kani::solver(kissat))]
fn alloc_test() {
    let capacity = if cfg!(any(kani, miri)) {
        1usize..3
    } else {
        1usize..4096
    };

    check!()
        .with_generator((capacity, gen::<u8>()))
        .cloned()
        .for_each(|(capacity, push_value)| {
            let (mut send, mut recv) = channel(capacity);

            send.try_slice().unwrap().unwrap().push(push_value).unwrap();

            let pop_value = recv.try_slice().unwrap().unwrap().pop().unwrap();
            assert_eq!(pop_value, push_value);
        })
}

const CAPACITY: usize = if cfg!(loom) { 2 } else { 10 };
const BATCH_COUNT: usize = if cfg!(loom) { 2 } else { 100 };
const BATCH_SIZE: usize = if cfg!(loom) { 3 } else { 20 };
const EXPECTED_COUNT: usize = BATCH_COUNT * BATCH_SIZE;

// TODO The async rx loom tests seem to take an unbounded amount if time if the batch count/size is
// anything bigger than 1. Ideally, these sizes would be bigger to test more permutations of
// orderings so we should investigate what's causing loom to endless spin.
const ASYNC_RX_BATCH_COUNT: usize = if cfg!(loom) { 1 } else { BATCH_COUNT };
const ASYNC_RX_BATCH_SIZE: usize = if cfg!(loom) { 1 } else { BATCH_SIZE };
const ASYNC_RX_EXPECTED_COUNT: usize = ASYNC_RX_BATCH_COUNT * ASYNC_RX_BATCH_SIZE;

#[test]
#[cfg_attr(miri, ignore)] // TODO https://github.com/aws/s2n-quic/issues/1635
fn loom_spin_tx_spin_rx_test() {
    loom_scenario(
        CAPACITY,
        |send| loom_spin_tx(send, BATCH_COUNT, BATCH_SIZE),
        |recv| loom_spin_rx(recv, EXPECTED_COUNT),
    )
}

#[test]
#[cfg_attr(miri, ignore)] // TODO https://github.com/aws/s2n-quic/issues/1635
fn loom_spin_tx_async_rx_test() {
    loom_scenario(
        CAPACITY,
        |send| loom_spin_tx(send, ASYNC_RX_BATCH_COUNT, ASYNC_RX_BATCH_SIZE),
        |recv| loom_async_rx(recv, ASYNC_RX_EXPECTED_COUNT),
    )
}

#[test]
#[cfg_attr(miri, ignore)] // TODO https://github.com/aws/s2n-quic/issues/1635
fn loom_async_tx_spin_rx_test() {
    loom_scenario(
        CAPACITY,
        |send| loom_async_tx(send, BATCH_COUNT, BATCH_SIZE),
        |recv| loom_spin_rx(recv, EXPECTED_COUNT),
    )
}

#[test]
#[cfg_attr(miri, ignore)] // TODO https://github.com/aws/s2n-quic/issues/1635
fn loom_async_tx_async_rx_test() {
    loom_scenario(
        CAPACITY,
        |send| loom_async_tx(send, ASYNC_RX_BATCH_COUNT, ASYNC_RX_BATCH_SIZE),
        |recv| loom_async_rx(recv, ASYNC_RX_EXPECTED_COUNT),
    )
}

fn loom_scenario(capacity: usize, sender: fn(Sender<u32>), receiver: fn(Receiver<u32>)) {
    loom::model(move || {
        let (send, recv) = channel(capacity);

        let a = loom::thread::spawn(move || sender(send));

        let b = loom::thread::spawn(move || receiver(recv));

        // loom tests will still run after returning so we don't need to join
        if cfg!(not(loom)) {
            a.join().unwrap();
            b.join().unwrap();
        }
    });
}

fn loom_spin_rx(mut recv: Receiver<u32>, expected: usize) {
    use loom::hint;

    let mut value = 0u32;
    loop {
        match recv.try_slice() {
            Ok(Some(mut slice)) => {
                while let Some(actual) = slice.pop() {
                    assert_eq!(actual, value);
                    value += 1;
                }
            }
            Ok(None) => hint::spin_loop(),
            Err(_) => {
                assert_eq!(value as usize, expected);
                return;
            }
        }
    }
}

fn loom_async_rx(mut recv: Receiver<u32>, expected: usize) {
    use futures::{future::poll_fn, ready};

    loom::future::block_on(async move {
        let mut value = 0u32;
        poll_fn(|cx| loop {
            match ready!(recv.poll_slice(cx)) {
                Ok(mut slice) => {
                    while let Some(actual) = slice.pop() {
                        assert_eq!(actual, value);
                        value += 1;
                    }
                }
                Err(_err) => return Poll::Ready(()),
            }
        })
        .await;

        assert_eq!(value as usize, expected);
    });
}

fn loom_spin_tx(mut send: Sender<u32>, batch_count: usize, batch_size: usize) {
    use loom::hint;

    let max_value = (batch_count * batch_size) as u32;
    let mut value = 0u32;

    'done: while max_value > value {
        let mut remaining = batch_size;
        while remaining > 0 {
            match send.try_slice() {
                Ok(Some(mut slice)) => {
                    let num_items = remaining;
                    for _ in 0..num_items {
                        if slice.push(value).is_err() {
                            hint::spin_loop();
                            continue;
                        }
                        value += 1;
                        remaining -= 1;
                    }
                }
                Ok(None) => {
                    // we don't have capacity to send so yield the thread
                    hint::spin_loop();
                }
                Err(_) => {
                    // The peer dropped the channel so bail
                    break 'done;
                }
            }
        }
    }

    assert_eq!(value, max_value);
}

fn loom_async_tx(mut send: Sender<u32>, batch_count: usize, batch_size: usize) {
    use futures::{future::poll_fn, ready};

    loom::future::block_on(async move {
        let max_value = (batch_count * batch_size) as u32;
        let mut value = 0u32;

        while max_value > value {
            let mut remaining = batch_size;
            let result = poll_fn(|cx| loop {
                return match ready!(send.poll_slice(cx)) {
                    Ok(mut slice) => {
                        let num_items = remaining;
                        for _ in 0..num_items {
                            if slice.push(value).is_err() {
                                // try polling the slice capacity again
                                break;
                            }
                            value += 1;
                            remaining -= 1;
                        }

                        if remaining > 0 {
                            continue;
                        }

                        Ok(())
                    }
                    Err(err) => Err(err),
                }
                .into();
            })
            .await;

            if result.is_err() {
                return;
            }
        }

        assert_eq!(value, max_value);
    });
}


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

use crate::time::timestamp::Timestamp;
use core::{
    task::{Context, Poll},
    time::Duration,
};

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

#[cfg(any(test, feature = "std"))]
pub use self::std::*;

/// A `Clock` is a source of [`Timestamp`]s.
pub trait Clock {
    /// Returns the current [`Timestamp`]
    fn get_time(&self) -> Timestamp;
}

pub trait ClockWithTimer: Clock {
    type Timer: Timer;

    fn timer(&self) -> Self::Timer;
}

pub trait Timer: Sized {
    #[inline]
    fn ready(&mut self) -> TimerReady<Self> {
        TimerReady(self)
    }

    fn poll_ready(&mut self, cx: &mut Context) -> Poll<()>;
    fn update(&mut self, timestamp: Timestamp);
}

impl_ready_future!(Timer, TimerReady, ());

/// A clock which always returns a Timestamp of value 1us
#[derive(Clone, Copy, Debug)]
pub struct NoopClock;

impl Clock for NoopClock {
    fn get_time(&self) -> Timestamp {
        unsafe { Timestamp::from_duration(Duration::from_micros(1)) }
    }
}


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

pub mod clock;
pub mod timer;
mod timestamp;

pub use clock::*;
pub use core::time::Duration;
pub use timer::Timer;
pub use timestamp::*;


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

use crate::time::timestamp::Timestamp;
use core::task::Poll;

/// A timer that does not trigger an update in a timer
/// list. These are usually owned by individual components
/// and needs to be explicitly polled.
///
/// Note: The timer doesn't implement Copy to ensure it isn't accidentally moved
///       and have the expiration discarded.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Timer {
    expiration: Option<Timestamp>,
}

impl Timer {
    /// Sets the timer to expire at the given timestamp
    #[inline]
    pub fn set(&mut self, time: Timestamp) {
        self.expiration = Some(time);
    }

    /// Cancels the timer.
    /// After cancellation, a timer will no longer report as expired.
    #[inline]
    pub fn cancel(&mut self) {
        self.expiration = None;
    }

    /// Returns true if the timer has expired
    #[inline]
    pub fn is_expired(&self, current_time: Timestamp) -> bool {
        match self.expiration {
            Some(timeout) => timeout.has_elapsed(current_time),
            _ => false,
        }
    }

    /// Returns true if the timer is armed
    #[inline]
    pub fn is_armed(&self) -> bool {
        self.expiration.is_some()
    }

    /// Notifies the timer of the current time.
    /// If the timer's expiration occurs before the current time, it will be cancelled.
    /// The method returns whether the timer was expired and had been
    /// cancelled.
    #[inline]
    pub fn poll_expiration(&mut self, current_time: Timestamp) -> Poll<()> {
        if self.is_expired(current_time) {
            self.cancel();
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }
}

impl From<Option<Timestamp>> for Timer {
    #[inline]
    fn from(expiration: Option<Timestamp>) -> Self {
        Self { expiration }
    }
}

/// Returned when a `Query` wants to end a timer query
#[derive(Clone, Copy, Debug, Default)]
pub struct QueryBreak;

/// The return type of a `timers` call
pub type Result<T = (), E = QueryBreak> = core::result::Result<T, E>;

/// A trait for a components that owns at least one timer
pub trait Provider {
    /// Notifies the query of any timers owned by the provider
    ///
    /// The provider should also delegate to subcomponents that own timers as well.
    fn timers<Q: Query>(&self, query: &mut Q) -> Result;

    /// Returns the next `Timestamp` at which the earliest timer is armed in the provider
    #[inline]
    fn next_expiration(&self) -> Option<Timestamp> {
        let mut timeout: Option<Timestamp> = None;
        let _ = self.timers(&mut timeout);
        timeout
    }

    /// Returns `true` if there are any timers armed
    #[inline]
    fn is_armed(&self) -> bool {
        let mut is_armed = IsArmed::default();
        let _ = self.timers(&mut is_armed);
        is_armed.0
    }

    /// Counts the number of armed timers in the provider
    #[inline]
    fn armed_timer_count(&self) -> usize {
        let mut count = ArmedCount::default();
        let _ = self.timers(&mut count);
        count.0
    }

    /// Iterates over each timer in the provider and calls the provided function
    #[inline]
    fn for_each_timer<F: FnMut(&Timer) -> Result>(&self, f: F) {
        let mut for_each = ForEach(f);
        let _ = self.timers(&mut for_each);
    }
}

impl Provider for Timer {
    #[inline]
    fn timers<Q: Query>(&self, query: &mut Q) -> Result {
        query.on_timer(self)
    }
}

impl<T: Provider> Provider for &T {
    #[inline]
    fn timers<Q: Query>(&self, query: &mut Q) -> Result {
        (**self).timers(query)
    }
}

impl<T: Provider> Provider for &mut T {
    #[inline]
    fn timers<Q: Query>(&self, query: &mut Q) -> Result {
        (**self).timers(query)
    }
}

/// Implement Provider for a 2-element tuple to make it easy to do joins
impl<A: Provider, B: Provider> Provider for (A, B) {
    #[inline]
    fn timers<Q: Query>(&self, query: &mut Q) -> Result {
        self.0.timers(query)?;
        self.1.timers(query)?;
        Ok(())
    }
}

impl<T: Provider> Provider for Option<T> {
    #[inline]
    fn timers<Q: Query>(&self, query: &mut Q) -> Result {
        if let Some(t) = self.as_ref() {
            t.timers(query)?;
        }
        Ok(())
    }
}

/// A query to be executed against a provider
pub trait Query {
    /// Called for each timer owned by the provider
    fn on_timer(&mut self, timer: &Timer) -> Result;
}

/// Implement Query for `Option<Timestamp>` to make it easy to get the earliest armed timestamp
impl Query for Option<Timestamp> {
    #[inline]
    fn on_timer(&mut self, timer: &Timer) -> Result {
        match (self, timer.expiration) {
            // Take the minimum of the two timers
            (Some(a), Some(b)) => *a = (*a).min(b),
            // We don't have a time yet so just assign the expiration of the other
            (a @ None, b) => *a = b,
            // do nothing for everything else
            _ => {}
        }
        Ok(())
    }
}

/// Counts all of the armed timers
#[derive(Debug, Default)]
pub struct ArmedCount(pub usize);

impl Query for ArmedCount {
    #[inline]
    fn on_timer(&mut self, timer: &Timer) -> Result {
        if timer.is_armed() {
            self.0 += 1;
        }
        Ok(())
    }
}

/// Returns `true` if any of the timers are armed
#[derive(Debug, Default)]
pub struct IsArmed(pub bool);

impl Query for IsArmed {
    #[inline]
    fn on_timer(&mut self, timer: &Timer) -> Result {
        if timer.is_armed() {
            self.0 = true;
            return Err(QueryBreak);
        }
        Ok(())
    }
}

/// Iterates over each timer in the provider and calls a function
#[derive(Debug, Default)]
pub struct ForEach<F: FnMut(&Timer) -> Result>(F);

impl<F: FnMut(&Timer) -> Result> Query for ForEach<F> {
    #[inline]
    fn on_timer(&mut self, timer: &Timer) -> Result {
        (self.0)(timer)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::time::clock::{Clock, NoopClock};
    use core::time::Duration;

    #[test]
    fn is_armed_test() {
        let now = NoopClock.get_time();
        let mut timer = Timer::default();

        assert!(!timer.is_armed());

        timer.set(now);
        assert!(timer.is_armed());

        timer.cancel();
        assert!(!timer.is_armed());
    }

    #[test]
    fn is_expired_test() {
        let mut now = NoopClock.get_time();
        let mut timer = Timer::default();

        assert!(!timer.is_expired(now));

        timer.set(now + Duration::from_millis(100));

        now += Duration::from_millis(99);
        assert!(!timer.is_expired(now));

        assert!(
            timer.is_expired(now + Duration::from_micros(1)),
            "if a timer is less than 1ms in the future is should expire"
        );

        now += Duration::from_millis(1);
        assert!(timer.is_expired(now));

        timer.cancel();
        assert!(!timer.is_expired(now));
    }

    #[test]
    fn poll_expiration_test() {
        let mut now = NoopClock.get_time();
        let mut timer = Timer::default();

        timer.set(now + Duration::from_millis(100));

        assert!(!timer.poll_expiration(now).is_ready());
        assert!(timer.is_armed());

        now += Duration::from_millis(100);

        assert!(timer.poll_expiration(now).is_ready());
        assert!(!timer.is_armed());
    }
}