<commit_before>use crate::cell::UnsafeCell;
use crate::collections::VecDeque;
use crate::ffi::c_void;
use crate::hint;
use crate::ops::{Deref, DerefMut, Drop};
use crate::ptr;
use crate::sync::atomic::{AtomicUsize, Ordering};
use crate::sys::hermit::abi;

/// This type provides a lock based on busy waiting to realize mutual exclusion
///
/// # Description
///
/// This structure behaves a lot like a common mutex. There are some differences:
///
/// - By using busy waiting, it can be used outside the runtime.
/// - It is a so called ticket lock and is completely fair.
#[cfg_attr(target_arch = "x86_64", repr(align(128)))]
#[cfg_attr(not(target_arch = "x86_64"), repr(align(64)))]
struct Spinlock<T: ?Sized> {
    queue: AtomicUsize,
    dequeue: AtomicUsize,
    data: UnsafeCell<T>,
}

unsafe impl<T: ?Sized + Send> Sync for Spinlock<T> {}
unsafe impl<T: ?Sized + Send> Send for Spinlock<T> {}

/// A guard to which the protected data can be accessed
///
/// When the guard falls out of scope it will release the lock.
struct SpinlockGuard<'a, T: ?Sized + 'a> {
    dequeue: &'a AtomicUsize,
    data: &'a mut T,
}

impl<T> Spinlock<T> {
    pub const fn new(user_data: T) -> Spinlock<T> {
        Spinlock {
            queue: AtomicUsize::new(0),
            dequeue: AtomicUsize::new(1),
            data: UnsafeCell::new(user_data),
        }
    }

    #[inline]
    fn obtain_lock(&self) {
        let ticket = self.queue.fetch_add(1, Ordering::SeqCst) + 1;
        while self.dequeue.load(Ordering::SeqCst) != ticket {
            hint::spin_loop();
        }
    }

    #[inline]
    pub unsafe fn lock(&self) -> SpinlockGuard<'_, T> {
        self.obtain_lock();
        SpinlockGuard { dequeue: &self.dequeue, data: &mut *self.data.get() }
    }
}

impl<T: ?Sized + Default> Default for Spinlock<T> {
    fn default() -> Spinlock<T> {
        Spinlock::new(Default::default())
    }
}

impl<'a, T: ?Sized> Deref for SpinlockGuard<'a, T> {
    type Target = T;
    fn deref(&self) -> &T {
        &*self.data
    }
}

impl<'a, T: ?Sized> DerefMut for SpinlockGuard<'a, T> {
    fn deref_mut(&mut self) -> &mut T {
        &mut *self.data
    }
}

impl<'a, T: ?Sized> Drop for SpinlockGuard<'a, T> {
    /// The dropping of the SpinlockGuard will release the lock it was created from.
    fn drop(&mut self) {
        self.dequeue.fetch_add(1, Ordering::SeqCst);
    }
}

/// Realize a priority queue for tasks
struct PriorityQueue {
    queues: [Option<VecDeque<abi::Tid>>; abi::NO_PRIORITIES],
    prio_bitmap: u64,
}

impl PriorityQueue {
    pub const fn new() -> PriorityQueue {
        PriorityQueue {
            queues: [
                None, None, None, None, None, None, None, None, None, None, None, None, None, None,
                None, None, None, None, None, None, None, None, None, None, None, None, None, None,
                None, None, None,
            ],
            prio_bitmap: 0,
        }
    }

    /// Add a task id by its priority to the queue
    pub fn push(&mut self, prio: abi::Priority, id: abi::Tid) {
        let i: usize = prio.into().into();
        self.prio_bitmap |= (1 << i) as u64;
        if let Some(queue) = &mut self.queues[i] {
            queue.push_back(id);
        } else {
            let mut queue = VecDeque::new();
            queue.push_back(id);
            self.queues[i] = Some(queue);
        }
    }

    fn pop_from_queue(&mut self, queue_index: usize) -> Option<abi::Tid> {
        if let Some(queue) = &mut self.queues[queue_index] {
            let id = queue.pop_front();

            if queue.is_empty() {
                self.prio_bitmap &= !(1 << queue_index as u64);
            }

            id
        } else {
            None
        }
    }

    /// Pop the task handle with the highest priority from the queue
    pub fn pop(&mut self) -> Option<abi::Tid> {
        for i in 0..abi::NO_PRIORITIES {
            if self.prio_bitmap & (1 << i) != 0 {
                return self.pop_from_queue(i);
            }
        }

        None
    }
}

struct MutexInner {
    locked: bool,
    blocked_task: PriorityQueue,
}

impl MutexInner {
    pub const fn new() -> MutexInner {
        MutexInner { locked: false, blocked_task: PriorityQueue::new() }
    }
}

pub struct Mutex {
    inner: Spinlock<MutexInner>,
}

pub type MovableMutex = Mutex;

unsafe impl Send for Mutex {}
unsafe impl Sync for Mutex {}

impl Mutex {
    pub const fn new() -> Mutex {
        Mutex { inner: Spinlock::new(MutexInner::new()) }
    }

    #[inline]
    pub unsafe fn init(&mut self) {
        self.inner = Spinlock::new(MutexInner::new());
    }

    #[inline]
    pub unsafe fn lock(&self) {
        loop {
            let mut guard = self.inner.lock();
            if guard.locked == false {
                guard.locked = true;
                return;
            } else {
                let prio = abi::get_priority();
                let id = abi::getpid();

                guard.blocked_task.push(prio, id);
                abi::block_current_task();
                drop(guard);
                abi::yield_now();
            }
        }
    }

    #[inline]
    pub unsafe fn unlock(&self) {
        let mut guard = self.inner.lock();
        guard.locked = false;
        if let Some(tid) = guard.blocked_task.pop() {
            abi::wakeup_task(tid);
        }
    }

    #[inline]
    pub unsafe fn try_lock(&self) -> bool {
        let mut guard = self.inner.lock();
        if guard.locked == false {
            guard.locked = true;
        }
        guard.locked
    }

    #[inline]
    pub unsafe fn destroy(&self) {}
}

pub struct ReentrantMutex {
    inner: *const c_void,
}

impl ReentrantMutex {
    pub const unsafe fn uninitialized() -> ReentrantMutex {
        ReentrantMutex { inner: ptr::null() }
    }

    #[inline]
    pub unsafe fn init(&self) {
        let _ = abi::recmutex_init(&self.inner as *const *const c_void as *mut _);
    }

    #[inline]
    pub unsafe fn lock(&self) {
        let _ = abi::recmutex_lock(self.inner);
    }

    #[inline]
    pub unsafe fn try_lock(&self) -> bool {
        true
    }

    #[inline]
    pub unsafe fn unlock(&self) {
        let _ = abi::recmutex_unlock(self.inner);
    }

    #[inline]
    pub unsafe fn destroy(&self) {
        let _ = abi::recmutex_destroy(self.inner);
    }
}
<commit_msg>Rollup merge of #91176 - hermitcore:spin, r=kennytm<commit_after>use crate::cell::UnsafeCell;
use crate::collections::VecDeque;
use crate::ffi::c_void;
use crate::hint;
use crate::ops::{Deref, DerefMut, Drop};
use crate::ptr;
use crate::sync::atomic::{AtomicUsize, Ordering};
use crate::sys::hermit::abi;

/// This type provides a lock based on busy waiting to realize mutual exclusion
///
/// # Description
///
/// This structure behaves a lot like a common mutex. There are some differences:
///
/// - By using busy waiting, it can be used outside the runtime.
/// - It is a so called ticket lock and is completely fair.
#[cfg_attr(target_arch = "x86_64", repr(align(128)))]
#[cfg_attr(not(target_arch = "x86_64"), repr(align(64)))]
struct Spinlock<T: ?Sized> {
    queue: AtomicUsize,
    dequeue: AtomicUsize,
    data: UnsafeCell<T>,
}

unsafe impl<T: ?Sized + Send> Sync for Spinlock<T> {}
unsafe impl<T: ?Sized + Send> Send for Spinlock<T> {}

/// A guard to which the protected data can be accessed
///
/// When the guard falls out of scope it will release the lock.
struct SpinlockGuard<'a, T: ?Sized + 'a> {
    dequeue: &'a AtomicUsize,
    data: &'a mut T,
}

impl<T> Spinlock<T> {
    pub const fn new(user_data: T) -> Spinlock<T> {
        Spinlock {
            queue: AtomicUsize::new(0),
            dequeue: AtomicUsize::new(1),
            data: UnsafeCell::new(user_data),
        }
    }

    #[inline]
    fn obtain_lock(&self) {
        let ticket = self.queue.fetch_add(1, Ordering::SeqCst) + 1;
        let mut counter: u16 = 0;
        while self.dequeue.load(Ordering::SeqCst) != ticket {
            counter += 1;
            if counter < 100 {
                hint::spin_loop();
            } else {
                counter = 0;
                unsafe {
                    abi::yield_now();
                }
            }
        }
    }

    #[inline]
    pub unsafe fn lock(&self) -> SpinlockGuard<'_, T> {
        self.obtain_lock();
        SpinlockGuard { dequeue: &self.dequeue, data: &mut *self.data.get() }
    }
}

impl<T: ?Sized + Default> Default for Spinlock<T> {
    fn default() -> Spinlock<T> {
        Spinlock::new(Default::default())
    }
}

impl<'a, T: ?Sized> Deref for SpinlockGuard<'a, T> {
    type Target = T;
    fn deref(&self) -> &T {
        &*self.data
    }
}

impl<'a, T: ?Sized> DerefMut for SpinlockGuard<'a, T> {
    fn deref_mut(&mut self) -> &mut T {
        &mut *self.data
    }
}

impl<'a, T: ?Sized> Drop for SpinlockGuard<'a, T> {
    /// The dropping of the SpinlockGuard will release the lock it was created from.
    fn drop(&mut self) {
        self.dequeue.fetch_add(1, Ordering::SeqCst);
    }
}

/// Realize a priority queue for tasks
struct PriorityQueue {
    queues: [Option<VecDeque<abi::Tid>>; abi::NO_PRIORITIES],
    prio_bitmap: u64,
}

impl PriorityQueue {
    pub const fn new() -> PriorityQueue {
        PriorityQueue {
            queues: [
                None, None, None, None, None, None, None, None, None, None, None, None, None, None,
                None, None, None, None, None, None, None, None, None, None, None, None, None, None,
                None, None, None,
            ],
            prio_bitmap: 0,
        }
    }

    /// Add a task id by its priority to the queue
    pub fn push(&mut self, prio: abi::Priority, id: abi::Tid) {
        let i: usize = prio.into().into();
        self.prio_bitmap |= (1 << i) as u64;
        if let Some(queue) = &mut self.queues[i] {
            queue.push_back(id);
        } else {
            let mut queue = VecDeque::new();
            queue.push_back(id);
            self.queues[i] = Some(queue);
        }
    }

    fn pop_from_queue(&mut self, queue_index: usize) -> Option<abi::Tid> {
        if let Some(queue) = &mut self.queues[queue_index] {
            let id = queue.pop_front();

            if queue.is_empty() {
                self.prio_bitmap &= !(1 << queue_index as u64);
            }

            id
        } else {
            None
        }
    }

    /// Pop the task handle with the highest priority from the queue
    pub fn pop(&mut self) -> Option<abi::Tid> {
        for i in 0..abi::NO_PRIORITIES {
            if self.prio_bitmap & (1 << i) != 0 {
                return self.pop_from_queue(i);
            }
        }

        None
    }
}

struct MutexInner {
    locked: bool,
    blocked_task: PriorityQueue,
}

impl MutexInner {
    pub const fn new() -> MutexInner {
        MutexInner { locked: false, blocked_task: PriorityQueue::new() }
    }
}

pub struct Mutex {
    inner: Spinlock<MutexInner>,
}

pub type MovableMutex = Mutex;

unsafe impl Send for Mutex {}
unsafe impl Sync for Mutex {}

impl Mutex {
    pub const fn new() -> Mutex {
        Mutex { inner: Spinlock::new(MutexInner::new()) }
    }

    #[inline]
    pub unsafe fn init(&mut self) {
        self.inner = Spinlock::new(MutexInner::new());
    }

    #[inline]
    pub unsafe fn lock(&self) {
        loop {
            let mut guard = self.inner.lock();
            if guard.locked == false {
                guard.locked = true;
                return;
            } else {
                let prio = abi::get_priority();
                let id = abi::getpid();

                guard.blocked_task.push(prio, id);
                abi::block_current_task();
                drop(guard);
                abi::yield_now();
            }
        }
    }

    #[inline]
    pub unsafe fn unlock(&self) {
        let mut guard = self.inner.lock();
        guard.locked = false;
        if let Some(tid) = guard.blocked_task.pop() {
            abi::wakeup_task(tid);
        }
    }

    #[inline]
    pub unsafe fn try_lock(&self) -> bool {
        let mut guard = self.inner.lock();
        if guard.locked == false {
            guard.locked = true;
        }
        guard.locked
    }

    #[inline]
    pub unsafe fn destroy(&self) {}
}

pub struct ReentrantMutex {
    inner: *const c_void,
}

impl ReentrantMutex {
    pub const unsafe fn uninitialized() -> ReentrantMutex {
        ReentrantMutex { inner: ptr::null() }
    }

    #[inline]
    pub unsafe fn init(&self) {
        let _ = abi::recmutex_init(&self.inner as *const *const c_void as *mut _);
    }

    #[inline]
    pub unsafe fn lock(&self) {
        let _ = abi::recmutex_lock(self.inner);
    }

    #[inline]
    pub unsafe fn try_lock(&self) -> bool {
        true
    }

    #[inline]
    pub unsafe fn unlock(&self) {
        let _ = abi::recmutex_unlock(self.inner);
    }

    #[inline]
    pub unsafe fn destroy(&self) {
        let _ = abi::recmutex_destroy(self.inner);
    }
}
<|endoftext|>


<commit_before>//! Defines a generic 2D image type.

use ndarray;
use ndarray::prelude::*;
use num_traits::{Zero};

use std::iter::{IntoIterator};

use errors::*;
use traits::Pixel;

/// 2-dimensional image type.
pub struct Image2D

    where P: Pixel
{
    buffer: Array2

}

/// Type of immutable `Image2D` views.
pub type Image2DView<'a, P> = ndarray::ArrayView<'a, P, Ix2>;

impl<'a, P> Image2D

    where P: Pixel
{
    /// Create a new image of specified dimensions from a `Vec`.
    ///
    /// **Error**: `InvalidDimensions` if the dimensions do not match the length of `v`.
    pub fn from_vec(w: u32, h: u32, v: Vec
) -> Result<Image2D
> {
        let buf = try!(Array2::from_shape_vec((w as usize, h as usize), v));
        Ok(Image2D{ buffer: buf })
    }

    /// Return the pixel at the specified coordinates.
    ///
    /// **Panics** if the index is out of bounds.
    pub fn get_pixel(&self, x: u32, y: u32) -> P {
        self.buffer[[x as usize, y as usize]].clone()
    }

    /// Set the pixel at the specified coordinates to the specified value.
    ///
    /// **Panics** if the index is out of bounds.
    pub fn put_pixel(&mut self, x: u32, y: u32, pixel: P) {
        self.buffer[[x as usize, y as usize]] = pixel;
    }

    /// Return the width of the image.
    pub fn width(&self) -> u32 { self.buffer.cols() as u32 }
    /// Return the height of the image.
    pub fn height(&self) -> u32 { self.buffer.rows() as u32 }
    /// Return the dimensions of the image as a `(width, height)` tuple.
    pub fn dimensions(&self) -> (u32, u32) { (self.width(), self.height()) }

    /// Return a view to a subset of the image of specified dimensions starting at the specified
    /// coordinates.
    ///
    /// **Panics** if the specified region crosses image boundaries.
    pub fn sub_image(&'a self, x: u32, y: u32, w: u32, h: u32) -> Image2DView<'a, P> {
        let x = x as isize;
        let y = y as isize;
        let w = w as isize;
        let h = h as isize;
        self.buffer.slice(s![x..x + w, y..y + h])
    }
}

impl<'a, P> IntoIterator for &'a Image2D

    where P: Pixel + 'a
{
    type Item = &'a P;
    type IntoIter = ndarray::iter::Iter<'a, P, Ix2>;

    fn into_iter(self) -> Self::IntoIter {
        self.buffer.into_iter()
    }
}

impl<'a, P> IntoIterator for &'a mut Image2D

    where P: Pixel + 'a
{
    type Item = &'a mut P;
    type IntoIter = ndarray::iter::IterMut<'a, P, Ix2>;

    fn into_iter(self) -> Self::IntoIter {
        self.buffer.iter_mut()
    }
}

impl
 Image2D

    where P: Pixel + Zero
{
    /// Create a new image of specified dimensions filled with zeros.
    pub fn zeros(width: u32, height: u32) -> Image2D
 {
        Image2D { buffer: Array2::zeros((width as usize, height as usize)) }
    }
}

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

    use std::iter::FromIterator;
    use std::fmt::Debug;

    #[test]
    fn test_from_vec() {
        let v1 = Vec::from_iter((1u8..10u8).map(|n| Luma::new([n])));
        let v2 = Vec::from_iter((1u8..7u8).map(|n| Luma::new([n])));

        assert!(Image2D::from_vec(3, 3, v1.clone()).is_ok());
        assert!(Image2D::from_vec(2, 3, v2.clone()).is_ok());
        assert!(Image2D::from_vec(3, 2, v2.clone()).is_ok());

        assert!(Image2D::from_vec(3, 3, v2.clone()).is_err());
        assert!(Image2D::from_vec(4, 2, v2.clone()).is_err());
    }

    #[test]
    fn test_zeros() {
        fn test_zeros_helper
(w: u32, h: u32)
            where P: Pixel + Zero + Debug
        {
            let img = Image2D::
::zeros(w, h);
            for pixel in &img {
                assert!(pixel.is_zero());
            }
        };
        test_zeros_helper::<Luma<u8>>(100, 100);
        test_zeros_helper::<Luma<f32>>(100, 100);
    }

    #[test]
    fn test_into_iter() {
        let v: Vec<Luma<u8>> = (1..10).map(|i| Luma::from(i)).collect();
        let img = Image2D::from_vec(3, 3, v.clone()).unwrap();

        for (p, i) in img.into_iter().zip(v.into_iter()) {
            assert!(&i == p);
        }
    }

    #[test]
    fn test_sub_image() {
        let v: Vec<Luma<usize>> = (1..10).map(|n| Luma::new([n])).collect();
        let img = Image2D::from_vec(3, 3, v).unwrap();
        let subimg1 = img.sub_image(1, 1, 2, 2);

        fn subimg_vec_eq<'a>(subimg: Image2DView<'a, Luma<usize>>, v: &Vec<usize>) -> bool {
            let img_iter = subimg.into_iter();
            let v_iter   = v.into_iter();
            match img_iter.len() == v_iter.len() {
                true => !img_iter.zip(v_iter).any(|(p, i)| p.data[0] != *i),
                false => false
            }
        }

        let subimg1_vec: Vec<usize> = vec![5, 6, 8, 9];

        assert!(subimg_vec_eq(subimg1, &subimg1_vec));
    }
}
<commit_msg>Add mutable sub-image view<commit_after>//! Defines a generic 2D image type.

use ndarray;
use ndarray::prelude::*;
use num_traits::{Zero};

use std::iter::{IntoIterator};

use errors::*;
use traits::Pixel;

/// 2-dimensional image type.
pub struct Image2D

    where P: Pixel
{
    buffer: Array2

}

/// Type of immutable `Image2D` views.
pub type Image2DView<'a, P> = ndarray::ArrayView<'a, P, Ix2>;
pub type Image2DViewMut<'a, P> = ndarray::ArrayViewMut<'a, P, Ix2>;

impl<'a, P> Image2D

    where P: Pixel
{
    /// Create a new image of specified dimensions from a `Vec`.
    ///
    /// **Error**: `InvalidDimensions` if the dimensions do not match the length of `v`.
    pub fn from_vec(w: u32, h: u32, v: Vec
) -> Result<Image2D
> {
        let buf = try!(Array2::from_shape_vec((w as usize, h as usize), v));
        Ok(Image2D{ buffer: buf })
    }

    /// Return the pixel at the specified coordinates.
    ///
    /// **Panics** if the index is out of bounds.
    pub fn get_pixel(&self, x: u32, y: u32) -> P {
        self.buffer[[x as usize, y as usize]].clone()
    }

    /// Set the pixel at the specified coordinates to the specified value.
    ///
    /// **Panics** if the index is out of bounds.
    pub fn put_pixel(&mut self, x: u32, y: u32, pixel: P) {
        self.buffer[[x as usize, y as usize]] = pixel;
    }

    /// Return the width of the image.
    pub fn width(&self) -> u32 { self.buffer.cols() as u32 }
    /// Return the height of the image.
    pub fn height(&self) -> u32 { self.buffer.rows() as u32 }
    /// Return the dimensions of the image as a `(width, height)` tuple.
    pub fn dimensions(&self) -> (u32, u32) { (self.width(), self.height()) }

    /// Return a view to a subset of the image of specified dimensions starting at the specified
    /// coordinates.
    ///
    /// **Panics** if the specified region crosses image boundaries.
    pub fn sub_image(&'a self, x: u32, y: u32, w: u32, h: u32) -> Image2DView<'a, P> {
        let x = x as isize;
        let y = y as isize;
        let w = w as isize;
        let h = h as isize;
        self.buffer.slice(s![x..x + w, y..y + h])
    }

    /// Return a mutable view to a subset of the image of specified dimensions starting at the specified
    /// coordinates.
    ///
    /// **Panics** if the specified region crosses image boundaries.
    pub fn sub_image_mut(&'a mut self, x: u32, y: u32, w: u32, h: u32) -> Image2DViewMut<'a, P> {
        let x = x as isize;
        let y = y as isize;
        let w = w as isize;
        let h = h as isize;
        self.buffer.slice_mut(s![x..x + w, y..y + h])
    }
}

impl<'a, P> IntoIterator for &'a Image2D

    where P: Pixel + 'a
{
    type Item = &'a P;
    type IntoIter = ndarray::iter::Iter<'a, P, Ix2>;

    fn into_iter(self) -> Self::IntoIter {
        self.buffer.into_iter()
    }
}

impl<'a, P> IntoIterator for &'a mut Image2D

    where P: Pixel + 'a
{
    type Item = &'a mut P;
    type IntoIter = ndarray::iter::IterMut<'a, P, Ix2>;

    fn into_iter(self) -> Self::IntoIter {
        self.buffer.iter_mut()
    }
}

impl
 Image2D

    where P: Pixel + Zero
{
    /// Create a new image of specified dimensions filled with zeros.
    pub fn zeros(width: u32, height: u32) -> Image2D
 {
        Image2D { buffer: Array2::zeros((width as usize, height as usize)) }
    }
}

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

    use std::iter::FromIterator;
    use std::fmt::Debug;

    #[test]
    fn test_from_vec() {
        let v1 = Vec::from_iter((1u8..10u8).map(|n| Luma::new([n])));
        let v2 = Vec::from_iter((1u8..7u8).map(|n| Luma::new([n])));

        assert!(Image2D::from_vec(3, 3, v1.clone()).is_ok());
        assert!(Image2D::from_vec(2, 3, v2.clone()).is_ok());
        assert!(Image2D::from_vec(3, 2, v2.clone()).is_ok());

        assert!(Image2D::from_vec(3, 3, v2.clone()).is_err());
        assert!(Image2D::from_vec(4, 2, v2.clone()).is_err());
    }

    #[test]
    fn test_zeros() {
        fn test_zeros_helper
(w: u32, h: u32)
            where P: Pixel + Zero + Debug
        {
            let img = Image2D::
::zeros(w, h);
            for pixel in &img {
                assert!(pixel.is_zero());
            }
        };
        test_zeros_helper::<Luma<u8>>(100, 100);
        test_zeros_helper::<Luma<f32>>(100, 100);
    }

    #[test]
    fn test_into_iter() {
        let v: Vec<Luma<u8>> = (1..10).map(|i| Luma::from(i)).collect();
        let img = Image2D::from_vec(3, 3, v.clone()).unwrap();

        for (p, i) in img.into_iter().zip(v.into_iter()) {
            assert!(&i == p);
        }
    }

    #[test]
    fn test_sub_image() {
        let v: Vec<Luma<usize>> = (1..10).map(|n| Luma::new([n])).collect();
        let img = Image2D::from_vec(3, 3, v).unwrap();
        let subimg1 = img.sub_image(1, 1, 2, 2);

        fn subimg_vec_eq<'a>(subimg: Image2DView<'a, Luma<usize>>, v: &Vec<usize>) -> bool {
            let img_iter = subimg.into_iter();
            let v_iter   = v.into_iter();
            match img_iter.len() == v_iter.len() {
                true => !img_iter.zip(v_iter).any(|(p, i)| p.data[0] != *i),
                false => false
            }
        }

        let subimg1_vec: Vec<usize> = vec![5, 6, 8, 9];

        assert!(subimg_vec_eq(subimg1, &subimg1_vec));
    }
}
<|endoftext|>


<commit_before>use std::{fmt, num::ParseIntError, str::FromStr};

/// A type capable of representing either an `i64` or a `u64`.
#[derive(Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct Integer {
    value: i128,
}

impl Integer {
    pub fn as_signed(self) -> Option<i64> {
        if self.value >= i64::min_value().into() && self.value <= i64::max_value().into() {
            Some(self.value as i64)
        } else {
            None
        }
    }

    pub fn as_unsigned(self) -> Option<u64> {
        if self.value >= 0 && self.value <= u64::max_value().into() {
            Some(self.value as u64)
        } else {
            None
        }
    }
}

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

impl fmt::Display for Integer {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.value.fmt(f)
    }
}

impl FromStr for Integer {
    type Err = ParseIntError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("0x") {
            // NetBSD dialect adds the `0x` numeric objects,
            // which are always unsigned.
            // See the `PROP_NUMBER(3)` man page
            let s = s.trim_start_matches("0x");
            u64::from_str_radix(s, 16).map(Into::into)
        } else {
            // Match Apple's implementation in CFPropertyList.h - always try to parse as an i64 first.
            // TODO: Use IntErrorKind once stable and retry parsing on overflow only.
            Ok(match s.parse::<i64>() {
                Ok(v) => v.into(),
                Err(_) => s.parse::<u64>()?.into(),
            })
        }
    }
}

impl From<i64> for Integer {
    fn from(value: i64) -> Integer {
        Integer {
            value: value.into(),
        }
    }
}

impl From<i32> for Integer {
    fn from(value: i32) -> Integer {
        Integer {
            value: value.into(),
        }
    }
}

impl From<i16> for Integer {
    fn from(value: i16) -> Integer {
        Integer {
            value: value.into(),
        }
    }
}

impl From<i8> for Integer {
    fn from(value: i8) -> Integer {
        Integer {
            value: value.into(),
        }
    }
}

impl From<u64> for Integer {
    fn from(value: u64) -> Integer {
        Integer {
            value: value.into(),
        }
    }
}

impl From<u32> for Integer {
    fn from(value: u32) -> Integer {
        Integer {
            value: value.into(),
        }
    }
}

impl From<u16> for Integer {
    fn from(value: u16) -> Integer {
        Integer {
            value: value.into(),
        }
    }
}

impl From<u8> for Integer {
    fn from(value: u8) -> Integer {
        Integer {
            value: value.into(),
        }
    }
}

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

    #[test]
    fn from_str_limits() {
        assert_eq!("-1".parse::<Integer>(), Ok((-1).into()));
        assert_eq!("0".parse::<Integer>(), Ok(0.into()));
        assert_eq!("1".parse::<Integer>(), Ok(1.into()));
        assert_eq!(
            "-9223372036854775808".parse::<Integer>(),
            Ok((-9223372036854775808i64).into())
        );
        assert!("-9223372036854775809".parse::<Integer>().is_err());
        assert_eq!(
            "18446744073709551615".parse::<Integer>(),
            Ok(18446744073709551615u64.into())
        );
        assert!("18446744073709551616".parse::<Integer>().is_err());
    }
}
<commit_msg>Tweak `Integer` implementation and docs.<commit_after>use std::{fmt, num::ParseIntError, str::FromStr};

/// An integer that can be represented by either an `i64` or a `u64`.
#[derive(Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct Integer {
    value: i128,
}

impl Integer {
    /// Returns the value as an `i64` if it can be represented by that type.
    pub fn as_signed(self) -> Option<i64> {
        if self.value >= i128::from(i64::min_value()) && self.value <= i128::from(i64::max_value())
        {
            Some(self.value as i64)
        } else {
            None
        }
    }

    /// Returns the value as a `u64` if it can be represented by that type.
    pub fn as_unsigned(self) -> Option<u64> {
        if self.value >= 0 && self.value <= i128::from(u64::max_value()) {
            Some(self.value as u64)
        } else {
            None
        }
    }
}

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

impl fmt::Display for Integer {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.value.fmt(f)
    }
}

impl FromStr for Integer {
    type Err = ParseIntError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("0x") {
            // NetBSD dialect adds the `0x` numeric objects,
            // which are always unsigned.
            // See the `PROP_NUMBER(3)` man page
            let s = s.trim_start_matches("0x");
            u64::from_str_radix(s, 16).map(Into::into)
        } else {
            // Match Apple's implementation in CFPropertyList.h - always try to parse as an i64 first.
            // TODO: Use IntErrorKind once stable and retry parsing on overflow only.
            Ok(match s.parse::<i64>() {
                Ok(v) => v.into(),
                Err(_) => s.parse::<u64>()?.into(),
            })
        }
    }
}

impl From<i64> for Integer {
    fn from(value: i64) -> Integer {
        Integer {
            value: value.into(),
        }
    }
}

impl From<i32> for Integer {
    fn from(value: i32) -> Integer {
        Integer {
            value: value.into(),
        }
    }
}

impl From<i16> for Integer {
    fn from(value: i16) -> Integer {
        Integer {
            value: value.into(),
        }
    }
}

impl From<i8> for Integer {
    fn from(value: i8) -> Integer {
        Integer {
            value: value.into(),
        }
    }
}

impl From<u64> for Integer {
    fn from(value: u64) -> Integer {
        Integer {
            value: value.into(),
        }
    }
}

impl From<u32> for Integer {
    fn from(value: u32) -> Integer {
        Integer {
            value: value.into(),
        }
    }
}

impl From<u16> for Integer {
    fn from(value: u16) -> Integer {
        Integer {
            value: value.into(),
        }
    }
}

impl From<u8> for Integer {
    fn from(value: u8) -> Integer {
        Integer {
            value: value.into(),
        }
    }
}

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

    #[test]
    fn from_str_limits() {
        assert_eq!("-1".parse::<Integer>(), Ok((-1).into()));
        assert_eq!("0".parse::<Integer>(), Ok(0.into()));
        assert_eq!("1".parse::<Integer>(), Ok(1.into()));
        assert_eq!(
            "-9223372036854775808".parse::<Integer>(),
            Ok((-9223372036854775808i64).into())
        );
        assert!("-9223372036854775809".parse::<Integer>().is_err());
        assert_eq!(
            "18446744073709551615".parse::<Integer>(),
            Ok(18446744073709551615u64.into())
        );
        assert!("18446744073709551616".parse::<Integer>().is_err());
    }
}
<|endoftext|>


<commit_before>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError, Offset};
use log::{error, trace};
use mio::{Evented, Events, Poll, PollOpt, Ready, Registration, SetReadiness, Token};
use parking_lot::Mutex;
use std::{
    io::{Read, Write},
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    thread::{self, Builder as ThreadBuilder, JoinHandle},
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const DATA: Token = Token(2);
const CONTINUE: Token = Token(3);

const FRAMES_STORAGE: usize = 32;

type ThreadHandle = JoinHandle<Result<()>>;

#[derive(Clone, Debug)]
pub(crate) struct IoLoopHandle {
    handle: Arc<Mutex<Option<ThreadHandle>>>,
}

impl Default for IoLoopHandle {
    fn default() -> Self {
        Self {
            handle: Arc::new(Mutex::new(None)),
        }
    }
}

impl IoLoopHandle {
    pub(crate) fn register(&self, handle: JoinHandle<Result<()>>) {
        *self.handle.lock() = Some(handle);
    }

    pub(crate) fn wait(&self) -> Result<()> {
        if let Some(handle) = self.handle.lock().take() {
            handle.join().expect("io loop")?
        }
        Ok(())
    }
}

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    registration: Registration,
    set_readiness: SetReadiness,
    hb_handle: Option<JoinHandle<()>>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    has_data: bool,
    send_heartbeat: Arc<AtomicBool>,
    poll_timeout: Option<Duration>,
}

impl<T: Evented + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll.map(|t| Ok((t.0, true))).unwrap_or_else(|| {
            Poll::new()
                .map(|poll| (poll, false))
                .map_err(Arc::new)
                .map_err(Error::IOError)
        })?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let (registration, set_readiness) = Registration::new2();
        let inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            registration,
            set_readiness,
            hb_handle: None,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            has_data: false,
            send_heartbeat: Arc::new(AtomicBool::new(false)),
            poll_timeout: None,
        };
        if registered {
            inner
                .poll
                .reregister(&inner.socket, SOCKET, Ready::all(), PollOpt::edge())
                .map_err(Arc::new)
                .map_err(Error::IOError)?;
        } else {
            inner
                .poll
                .register(&inner.socket, SOCKET, Ready::all(), PollOpt::edge())
                .map_err(Arc::new)
                .map_err(Error::IOError)?;
        }
        inner
            .poll
            .register(&inner.connection, DATA, Ready::readable(), PollOpt::edge())
            .map_err(Arc::new)
            .map_err(Error::IOError)?;
        inner
            .poll
            .register(
                &inner.registration,
                CONTINUE,
                Ready::readable(),
                PollOpt::edge(),
            )
            .map_err(Arc::new)
            .map_err(Error::IOError)?;
        Ok(inner)
    }

    fn start_heartbeat(&mut self, interval: Duration) -> Result<()> {
        let connection = self.connection.clone();
        let send_hartbeat = self.send_heartbeat.clone();
        let hb_handle = ThreadBuilder::new()
            .name("heartbeat".to_owned())
            .spawn(move || {
                while connection.status().connected() {
                    let start = Instant::now();
                    let mut remaining = interval;

                    loop {
                        thread::park_timeout(remaining);
                        let elapsed = start.elapsed();
                        if elapsed >= remaining {
                            break;
                        }
                        remaining -= interval - elapsed;
                    }

                    send_hartbeat.store(true, Ordering::Relaxed);
                }
            })
            .map_err(Arc::new)
            .map_err(Error::IOError)?;
        self.hb_handle = Some(hb_handle);
        Ok(())
    }

    fn heartbeat(&mut self) -> Result<()> {
        if self.send_heartbeat.load(Ordering::Relaxed) {
            trace!("send heartbeat");
            self.connection.send_heartbeat()?;
            self.send_heartbeat.store(false, Ordering::Relaxed);
        }
        Ok(())
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection.status().connected() {
            let frame_max = self.connection.configuration().frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.connection.configuration().heartbeat();
            if heartbeat != 0 {
                trace!("io_loop: start heartbeat");
                let heartbeat = Duration::from_secs(u64::from(heartbeat));
                self.start_heartbeat(heartbeat)?;
                self.poll_timeout = Some(heartbeat);
                trace!("io_loop: heartbeat started");
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data && !self.connection.status().blocked()
    }

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

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        let connection_status = self.connection.status();
        (self.status == Status::Initial
            || connection_status.connected()
            || connection_status.closing())
            && self.status != Status::Stop
            && !connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        self.run(&mut events)?;
                    }
                    if let Some(hb_handle) = self.hb_handle.take() {
                        hb_handle.thread().unpark();
                        hb_handle.join().expect("heartbeat loop failed");
                    }
                    Ok(())
                })
                .map_err(Arc::new)
                .map_err(Error::IOError)?,
        );
        Ok(())
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll
            .poll(events, self.poll_timeout)
            .map_err(Arc::new)
            .map_err(Error::IOError)?;
        trace!("io_loop poll done");
        for event in events.iter() {
            match event.token() {
                SOCKET => {
                    if event.readiness().is_readable() {
                        self.can_read = true;
                    }
                    if event.readiness().is_writable() {
                        self.can_write = true;
                    }
                }
                DATA => self.has_data = true,
                _ => {}
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data
        );
        loop {
            self.heartbeat()?;
            self.write()?;
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data,
            self.status
        );
        Ok(())
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection.status().errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    if let ConnectionState::SentProtocolHeader(pinky, ..) =
                        self.connection.status().state()
                    {
                        pinky.swear(Err(Error::ConnectionRefused));
                        self.status = Status::Stop;
                    }
                    self.connection.set_error()?;
                    return Err(e);
                }
            }
            self.send_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.connection.set_error()?;
                    return Err(e);
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.set_readiness
            .set_readiness(Ready::readable())
            .map_err(Arc::new)
            .map_err(Error::IOError)
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self
            .socket
            .write(&self.send_buffer.data())
            .map_err(Arc::new)
            .map_err(Error::IOError)?;

        trace!("wrote {} bytes", sz);
        self.send_buffer.consume(sz);
        if self.send_buffer.available_data() > 0 {
            // We didn't write all the data yet
            self.connection.set_readable()?;
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection.status().state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => self
                .socket
                .read(&mut self.receive_buffer.space())
                .map(|sz| {
                    trace!("read {} bytes", sz);
                    self.receive_buffer.fill(sz);
                })
                .map_err(Arc::new)
                .map_err(Error::IOError),
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((send_id, next_msg)) = self.connection.next_frame() {
            trace!("will write to buffer: {:?}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(_) => {
                    self.connection.mark_sent(send_id);
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame(send_id, next_msg)?;
                            self.send_buffer.shift();
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.connection.set_error()?;
                            Err(Error::SerialisationError(Arc::new(e)))
                        }
                    }
                }
            }
        } else {
            self.has_data = false;
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.connection.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.data().offset(i);
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    self.connection.set_error()?;
                    Err(Error::ParsingError(format!("{:?}", e)))
                }
            }
        }
    }
}
<commit_msg>io_loop: add an extra check<commit_after>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError, Offset};
use log::{error, trace};
use mio::{Evented, Events, Poll, PollOpt, Ready, Registration, SetReadiness, Token};
use parking_lot::Mutex;
use std::{
    io::{Read, Write},
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    thread::{self, Builder as ThreadBuilder, JoinHandle},
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const DATA: Token = Token(2);
const CONTINUE: Token = Token(3);

const FRAMES_STORAGE: usize = 32;

type ThreadHandle = JoinHandle<Result<()>>;

#[derive(Clone, Debug)]
pub(crate) struct IoLoopHandle {
    handle: Arc<Mutex<Option<ThreadHandle>>>,
}

impl Default for IoLoopHandle {
    fn default() -> Self {
        Self {
            handle: Arc::new(Mutex::new(None)),
        }
    }
}

impl IoLoopHandle {
    pub(crate) fn register(&self, handle: JoinHandle<Result<()>>) {
        *self.handle.lock() = Some(handle);
    }

    pub(crate) fn wait(&self) -> Result<()> {
        if let Some(handle) = self.handle.lock().take() {
            handle.join().expect("io loop")?
        }
        Ok(())
    }
}

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    registration: Registration,
    set_readiness: SetReadiness,
    hb_handle: Option<JoinHandle<()>>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    has_data: bool,
    send_heartbeat: Arc<AtomicBool>,
    poll_timeout: Option<Duration>,
}

impl<T: Evented + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll.map(|t| Ok((t.0, true))).unwrap_or_else(|| {
            Poll::new()
                .map(|poll| (poll, false))
                .map_err(Arc::new)
                .map_err(Error::IOError)
        })?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let (registration, set_readiness) = Registration::new2();
        let inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            registration,
            set_readiness,
            hb_handle: None,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            has_data: false,
            send_heartbeat: Arc::new(AtomicBool::new(false)),
            poll_timeout: None,
        };
        if registered {
            inner
                .poll
                .reregister(&inner.socket, SOCKET, Ready::all(), PollOpt::edge())
                .map_err(Arc::new)
                .map_err(Error::IOError)?;
        } else {
            inner
                .poll
                .register(&inner.socket, SOCKET, Ready::all(), PollOpt::edge())
                .map_err(Arc::new)
                .map_err(Error::IOError)?;
        }
        inner
            .poll
            .register(&inner.connection, DATA, Ready::readable(), PollOpt::edge())
            .map_err(Arc::new)
            .map_err(Error::IOError)?;
        inner
            .poll
            .register(
                &inner.registration,
                CONTINUE,
                Ready::readable(),
                PollOpt::edge(),
            )
            .map_err(Arc::new)
            .map_err(Error::IOError)?;
        Ok(inner)
    }

    fn start_heartbeat(&mut self, interval: Duration) -> Result<()> {
        let connection = self.connection.clone();
        let send_hartbeat = self.send_heartbeat.clone();
        let hb_handle = ThreadBuilder::new()
            .name("heartbeat".to_owned())
            .spawn(move || {
                while connection.status().connected() {
                    let start = Instant::now();
                    let mut remaining = interval;

                    loop {
                        thread::park_timeout(remaining);
                        let elapsed = start.elapsed();
                        if elapsed >= remaining {
                            break;
                        }
                        remaining -= interval - elapsed;
                    }

                    send_hartbeat.store(true, Ordering::Relaxed);
                }
            })
            .map_err(Arc::new)
            .map_err(Error::IOError)?;
        self.hb_handle = Some(hb_handle);
        Ok(())
    }

    fn heartbeat(&mut self) -> Result<()> {
        if self.send_heartbeat.load(Ordering::Relaxed) {
            trace!("send heartbeat");
            self.connection.send_heartbeat()?;
            self.send_heartbeat.store(false, Ordering::Relaxed);
        }
        Ok(())
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection.status().connected() {
            let frame_max = self.connection.configuration().frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.connection.configuration().heartbeat();
            if heartbeat != 0 {
                trace!("io_loop: start heartbeat");
                let heartbeat = Duration::from_secs(u64::from(heartbeat));
                self.start_heartbeat(heartbeat)?;
                self.poll_timeout = Some(heartbeat);
                trace!("io_loop: heartbeat started");
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data && !self.connection.status().blocked()
    }

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

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        let connection_status = self.connection.status();
        (self.status == Status::Initial
            || connection_status.connected()
            || connection_status.closing())
            && self.status != Status::Stop
            && !connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        self.run(&mut events)?;
                    }
                    if let Some(hb_handle) = self.hb_handle.take() {
                        hb_handle.thread().unpark();
                        hb_handle.join().expect("heartbeat loop failed");
                    }
                    Ok(())
                })
                .map_err(Arc::new)
                .map_err(Error::IOError)?,
        );
        Ok(())
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll
            .poll(events, self.poll_timeout)
            .map_err(Arc::new)
            .map_err(Error::IOError)?;
        trace!("io_loop poll done");
        for event in events.iter() {
            match event.token() {
                SOCKET => {
                    if event.readiness().is_readable() {
                        self.can_read = true;
                    }
                    if event.readiness().is_writable() {
                        self.can_write = true;
                    }
                }
                DATA => self.has_data = true,
                _ => {}
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data
        );
        loop {
            self.heartbeat()?;
            self.write()?;
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data,
            self.status
        );
        Ok(())
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection.status().errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    if let ConnectionState::SentProtocolHeader(pinky, ..) =
                        self.connection.status().state()
                    {
                        pinky.swear(Err(Error::ConnectionRefused));
                        self.status = Status::Stop;
                    }
                    self.connection.set_error()?;
                    return Err(e);
                }
            }
            self.send_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.connection.set_error()?;
                    return Err(e);
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.set_readiness
            .set_readiness(Ready::readable())
            .map_err(Arc::new)
            .map_err(Error::IOError)
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self
            .socket
            .write(&self.send_buffer.data())
            .map_err(Arc::new)
            .map_err(Error::IOError)?;

        trace!("wrote {} bytes", sz);
        self.send_buffer.consume(sz);
        if sz > 0 && self.send_buffer.available_data() > 0 {
            // We didn't write all the data yet
            self.connection.set_readable()?;
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection.status().state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => self
                .socket
                .read(&mut self.receive_buffer.space())
                .map(|sz| {
                    trace!("read {} bytes", sz);
                    self.receive_buffer.fill(sz);
                })
                .map_err(Arc::new)
                .map_err(Error::IOError),
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((send_id, next_msg)) = self.connection.next_frame() {
            trace!("will write to buffer: {:?}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(_) => {
                    self.connection.mark_sent(send_id);
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame(send_id, next_msg)?;
                            self.send_buffer.shift();
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.connection.set_error()?;
                            Err(Error::SerialisationError(Arc::new(e)))
                        }
                    }
                }
            }
        } else {
            self.has_data = false;
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.connection.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.data().offset(i);
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    self.connection.set_error()?;
                    Err(Error::ParsingError(format!("{:?}", e)))
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error,
    PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker};
use std::{
    collections::VecDeque,
    io::{Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    waker: Arc<Waker>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
    sent_initial_heartbeat: bool,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let waker = Arc::new(Waker::new(poll.registry(), WAKER)?);
        let mut inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            waker,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
            sent_initial_heartbeat: false,
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection.status().connected() {
            let frame_max = self.connection.configuration().frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.connection.configuration().heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat * 500)); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn has_data(&self) -> bool {
        self.connection.has_pending_frames()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection.status().blocked()
    }

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

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        let connection_status = self.connection.status();
        (self.status == Status::Initial
            || connection_status.connected()
            || connection_status.closing())
            && self.status != Status::Stop
            && !connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        self.run(&mut events)?;
                    }
                    Ok(())
                })?,
            waker,
        )
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                if event.is_readable() {
                    self.can_read = true;
                }
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.connection.poll_internal_promises()?;
            if self.should_heartbeat() {
                self.sent_initial_heartbeat = true;
                self.connection.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection.status().errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let ConnectionState::SentProtocolHeader(resolver, ..) = self.connection.status().state()
        {
            resolver.swear(Err(error.clone()));
            self.status = Status::Stop;
        }
        self.connection.set_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()?;
        Ok(())
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            // If we didn't send the first heartbeat, don't update last_write to trick us into sending it
            if self.sent_initial_heartbeat {
                self.last_write = Instant::now();
            }

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz, true);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                self.send_continue()?;
            }
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection.status().state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer
                    .read_from(&mut self.socket, false)
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz, false);
                    })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.connection.next_frame() {
            trace!("will write to buffer: {:?}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame((next_msg, resolver))?;
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.connection.set_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.connection.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed, false);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.connection.set_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<commit_msg>IoLoop: properly forward errors to serilized but not written frames<commit_after>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error,
    PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker};
use std::{
    collections::VecDeque,
    io::{Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    waker: Arc<Waker>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
    sent_initial_heartbeat: bool,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let waker = Arc::new(Waker::new(poll.registry(), WAKER)?);
        let mut inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            waker,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
            sent_initial_heartbeat: false,
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection.status().connected() {
            let frame_max = self.connection.configuration().frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.connection.configuration().heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat * 500)); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn has_data(&self) -> bool {
        self.connection.has_pending_frames()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection.status().blocked()
    }

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

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        let connection_status = self.connection.status();
        (self.status == Status::Initial
            || connection_status.connected()
            || connection_status.closing())
            && self.status != Status::Stop
            && !connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        if let Err(err) = self.run(&mut events) {
                            self.cancel_serilized_frames(err)?;
                        }
                    }
                    Ok(())
                })?,
            waker,
        )
    }

    fn cancel_serilized_frames(&self, error: Error) -> Result<()> {
        for (_, resolver) in self.serialized_frames.iter() {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                if event.is_readable() {
                    self.can_read = true;
                }
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.connection.poll_internal_promises()?;
            if self.should_heartbeat() {
                self.sent_initial_heartbeat = true;
                self.connection.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection.status().errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let ConnectionState::SentProtocolHeader(resolver, ..) = self.connection.status().state()
        {
            resolver.swear(Err(error.clone()));
            self.status = Status::Stop;
        }
        self.connection.set_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()?;
        Ok(())
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            // If we didn't send the first heartbeat, don't update last_write to trick us into sending it
            if self.sent_initial_heartbeat {
                self.last_write = Instant::now();
            }

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz, true);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                self.send_continue()?;
            }
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection.status().state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer
                    .read_from(&mut self.socket, false)
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz, false);
                    })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.connection.next_frame() {
            trace!("will write to buffer: {:?}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame((next_msg, resolver))?;
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.connection.set_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.connection.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed, false);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.connection.set_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>const N: i32 = 10;

pub struct Territory {
    territory: [[u32; N as usize]; N as usize],
    islands_count: i32,
    id: u32
}

const OFFSETS: [(i32, i32); 4] = [
    (0, 1),
    (0, -1),
    (1, 0),
    (-1, 0)
];

const SEA: u32 = 0;

impl Territory {

    pub fn new() -> Territory {
        return Territory {
            islands_count: 0,
            id: 0,
            territory: [[0; N as usize]; N as usize]
        };
    }

    pub fn add_land(& mut self, r: i32, c: i32) -> i32 {
        let mut different_neighbours: i32 = 0;

        for offset in OFFSETS.iter() {
            // relies on undeflow
            let new_coordinates = ((r + offset.0) as usize, (c + offset.1) as usize);
            let upper_bound = (N - 1) as usize;
            let are_coordinates_valid = new_coordinates.0 < upper_bound && new_coordinates.1 < upper_bound;
            if are_coordinates_valid && self.territory[new_coordinates.0][new_coordinates.1] != SEA {
                different_neighbours += 1;
            }
        }
        let islands_difference = -1 * (different_neighbours - 1);
        self.islands_count += islands_difference;

        // always give the new indexed ground a new id
        self.id += 1;
        self.territory[r as usize][c as usize] = self.id;
        self.reindex_islands(r, c);

        return self.islands_count;
    }

    fn reindex_islands(&mut self, r: i32, c: i32) {
        let current_el = self.territory[r as usize][c as usize];
        if current_el == SEA || current_el == self.id {
            return;
        }

        self.territory[r as usize][c as usize] = self.id;

        // TODO: make some sort of iterator to avoid code duplication
        for offset in OFFSETS.iter() {
            let new_coordinates = ((r + offset.0) as usize, (c + offset.1) as usize);
            let upper_bound = (N - 1) as usize;
            let are_coordinates_valid = new_coordinates.0 < upper_bound && new_coordinates.1 < upper_bound;
            if are_coordinates_valid {
                self.reindex_islands(new_coordinates.0 as i32, new_coordinates.1 as i32);
            }
        }
    }
}
<commit_msg>Fix recursion, it would always return<commit_after>const N: i32 = 10;

pub struct Territory {
    territory: [[u32; N as usize]; N as usize],
    islands_count: i32,
    id: u32
}

const OFFSETS: [(i32, i32); 4] = [
    (0, 1),
    (0, -1),
    (1, 0),
    (-1, 0)
];

const SEA: u32 = 0;

impl Territory {

    pub fn new() -> Territory {
        return Territory {
            islands_count: 0,
            id: 0,
            territory: [[0; N as usize]; N as usize]
        };
    }

    pub fn add_land(& mut self, r: i32, c: i32) -> i32 {
        let mut different_neighbours: i32 = 0;

        for offset in OFFSETS.iter() {
            // relies on undeflow
            let new_coordinates = ((r + offset.0) as usize, (c + offset.1) as usize);
            let upper_bound = (N - 1) as usize;
            let are_coordinates_valid = new_coordinates.0 < upper_bound && new_coordinates.1 < upper_bound;
            if are_coordinates_valid && self.territory[new_coordinates.0][new_coordinates.1] != SEA {
                different_neighbours += 1;
            }
        }
        let islands_difference = -1 * (different_neighbours - 1);
        self.islands_count += islands_difference;

        // always give the new indexed ground a new id
        self.id += 1;
        self.reindex_islands(r, c, true);

        return self.islands_count;
    }

    fn reindex_islands(&mut self, r: i32, c: i32, initial: bool) {
        let current_el = self.territory[r as usize][c as usize];
        if (current_el == SEA || current_el == self.id) && !initial {
            return;
        }

        self.territory[r as usize][c as usize] = self.id;

        // TODO: make some sort of iterator to avoid code duplication
        for offset in OFFSETS.iter() {
            let new_coordinates = ((r + offset.0) as usize, (c + offset.1) as usize);
            let upper_bound = (N - 1) as usize;
            let are_coordinates_valid = new_coordinates.0 < upper_bound && new_coordinates.1 < upper_bound;
            if are_coordinates_valid {
                self.reindex_islands(new_coordinates.0 as i32, new_coordinates.1 as i32, false);
            }
        }
    }
}
<|endoftext|>


<commit_before>mod compiler;
mod rts;

pub use self::compiler::*;

use dynasmrt;

pub struct Program {
    code: dynasmrt::ExecutableBuffer,
    start: dynasmrt::AssemblyOffset,
}

#[cfg(test)]
mod tests {
    use test_helpers::*;
    use common::{BfResult, Error};

    #[test]
    fn move_right_once() {
        assert_parse_interpret(b">", "", Ok(""));
    }

    #[test]
    fn move_left_once() {
        assert_parse_interpret(b"<", "", Err(Error::PointerUnderflow));
    }

    #[test]
    fn move_right_forever() {
        assert_parse_interpret(b"+[>+]", "", Err(Error::PointerOverflow));
    }

    #[test]
    fn echo_one_byte() {
        assert_parse_interpret(b",.", "A", Ok("A"));
    }

    #[test]
    fn inc_echo_one_byte() {
        assert_parse_interpret(b",+.", "A", Ok("B"));
    }

    #[test]
    fn hello_world() {
        assert_parse_interpret(HELLO_WORLD_SRC, "", Ok("Hello, World!"));
    }

    #[test]
    fn factoring() {
        assert_parse_interpret(FACTOR_SRC, "2\n", Ok("2: 2\n"));
        assert_parse_interpret(FACTOR_SRC, "3\n", Ok("3: 3\n"));
        assert_parse_interpret(FACTOR_SRC, "6\n", Ok("6: 2 3\n"));
        assert_parse_interpret(FACTOR_SRC, "100\n", Ok("100: 2 2 5 5\n"));
    }

    fn assert_parse_interpret(program: &[u8], input: &str, output: BfResult<&str>) {
        let program = ::ast::parse_program(program).unwrap();
        let program = ::rle::compile(&program);
        let program = ::peephole::compile(&program);
        let program = ::jit::compile(&program, true);
        assert_interpret_result(&program, input.as_bytes(), output.map(|s| s.as_bytes()));
    }
}

<commit_msg>Much more documentation.<commit_after>//! Just-in-time compiles Brainfuck AST to x64 machine code (`--features jit`, nightly only)
//!
//! This uses the [`dynasm`](https://crates.io/search?q=dynasm) crate to generate x86-64
//! machine code from peephole-optimized AST. This is currently the fastest implementation,
//! but it is available only on nightly Rust because `dynasm` uses a plugin.
//!
//! To go even faster, pass the `--unchecked` flag to the `bfi` interpreter to disable
//! memory bounds checking in the generated code. Note that this runs Brainfuck in
//! unsafe mode, which means that programs that move the pointer outside the allocated
//! memory will access and possibly overwrite arbitrary memory locations.

mod compiler;
mod rts;

pub use self::compiler::*;

use dynasmrt;

pub struct Program {
    code: dynasmrt::ExecutableBuffer,
    start: dynasmrt::AssemblyOffset,
}

#[cfg(test)]
mod tests {
    use test_helpers::*;
    use common::{BfResult, Error};

    #[test]
    fn move_right_once() {
        assert_parse_interpret(b">", "", Ok(""));
    }

    #[test]
    fn move_left_once() {
        assert_parse_interpret(b"<", "", Err(Error::PointerUnderflow));
    }

    #[test]
    fn move_right_forever() {
        assert_parse_interpret(b"+[>+]", "", Err(Error::PointerOverflow));
    }

    #[test]
    fn echo_one_byte() {
        assert_parse_interpret(b",.", "A", Ok("A"));
    }

    #[test]
    fn inc_echo_one_byte() {
        assert_parse_interpret(b",+.", "A", Ok("B"));
    }

    #[test]
    fn hello_world() {
        assert_parse_interpret(HELLO_WORLD_SRC, "", Ok("Hello, World!"));
    }

    #[test]
    fn factoring() {
        assert_parse_interpret(FACTOR_SRC, "2\n", Ok("2: 2\n"));
        assert_parse_interpret(FACTOR_SRC, "3\n", Ok("3: 3\n"));
        assert_parse_interpret(FACTOR_SRC, "6\n", Ok("6: 2 3\n"));
        assert_parse_interpret(FACTOR_SRC, "100\n", Ok("100: 2 2 5 5\n"));
    }

    fn assert_parse_interpret(program: &[u8], input: &str, output: BfResult<&str>) {
        let program = ::ast::parse_program(program).unwrap();
        let program = ::rle::compile(&program);
        let program = ::peephole::compile(&program);
        let program = ::jit::compile(&program, true);
        assert_interpret_result(&program, input.as_bytes(), output.map(|s| s.as_bytes()));
    }
}

<|endoftext|>