<commit_before>//! The implementation is based on Dmitry Vyukov's bounded MPMC queue.
//!
//! Source:
//!   - http://www.1024cores.net/home/lock-free-algorithms/queues/bounded-mpmc-queue
//!
//! Copyright & License:
//!   - Copyright (c) 2010-2011 Dmitry Vyukov
//!   - Simplified BSD License and Apache License, Version 2.0
//!   - http://www.1024cores.net/home/code-license

use std::cell::UnsafeCell;
use std::fmt;
use std::marker::PhantomData;
use std::mem;
use std::ptr;
use std::sync::atomic::{self, AtomicUsize, Ordering};

use crossbeam_utils::{Backoff, CachePadded};

use err::{PopError, PushError};

/// A slot in a queue.
struct Slot<T> {
    /// The current stamp.
    ///
    /// If the stamp equals the tail, this node will be next written to. If it equals the head,
    /// this node will be next read from.
    stamp: AtomicUsize,

    /// The value in this slot.
    value: UnsafeCell<T>,
}

/// A bounded multi-producer multi-consumer queue.
///
/// This queue allocates a fixed-capacity buffer on construction, which is used to store pushed
/// elements. The queue cannot hold more elements that the buffer allows. Attempting to push an
/// element into a full queue will fail. Having a buffer allocated upfront makes this queue a bit
/// faster than [`SegQueue`].
///
/// [`SegQueue`]: struct.SegQueue.html
///
/// # Examples
///
/// ```
/// use crossbeam_queue::{ArrayQueue, PushError};
///
/// let q = ArrayQueue::new(2);
///
/// assert_eq!(q.push('a'), Ok(()));
/// assert_eq!(q.push('b'), Ok(()));
/// assert_eq!(q.push('c'), Err(PushError('c')));
/// assert_eq!(q.pop(), Ok('a'));
/// ```
pub struct ArrayQueue<T> {
    /// The head of the queue.
    ///
    /// This value is a "stamp" consisting of an index into the buffer and a lap, but packed into a
    /// single `usize`. The lower bits represent the index, while the upper bits represent the lap.
    ///
    /// Elements are popped from the head of the queue.
    head: CachePadded<AtomicUsize>,

    /// The tail of the queue.
    ///
    /// This value is a "stamp" consisting of an index into the buffer and a lap, but packed into a
    /// single `usize`. The lower bits represent the index, while the upper bits represent the lap.
    ///
    /// Elements are pushed into the tail of the queue.
    tail: CachePadded<AtomicUsize>,

    /// The buffer holding slots.
    buffer: *mut Slot<T>,

    /// The queue capacity.
    cap: usize,

    /// A stamp with the value of `{ lap: 1, index: 0 }`.
    one_lap: usize,

    /// Indicates that dropping an `ArrayQueue<T>` may drop elements of type `T`.
    _marker: PhantomData<T>,
}

unsafe impl<T: Send> Sync for ArrayQueue<T> {}
unsafe impl<T: Send> Send for ArrayQueue<T> {}

impl<T> ArrayQueue<T> {
    /// Creates a new bounded queue with the given capacity.
    ///
    /// # Panics
    ///
    /// Panics if the capacity is zero.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::ArrayQueue;
    ///
    /// let q = ArrayQueue::<i32>::new(100);
    /// ```
    pub fn new(cap: usize) -> ArrayQueue<T> {
        assert!(cap > 0, "capacity must be non-zero");

        // Head is initialized to `{ lap: 0, index: 0 }`.
        // Tail is initialized to `{ lap: 0, index: 0 }`.
        let head = 0;
        let tail = 0;

        // Allocate a buffer of `cap` slots.
        let buffer = {
            let mut v = Vec::<Slot<T>>::with_capacity(cap);
            let ptr = v.as_mut_ptr();
            mem::forget(v);
            ptr
        };

        // Initialize stamps in the slots.
        for i in 0..cap {
            unsafe {
                // Set the stamp to `{ lap: 0, index: i }`.
                let slot = buffer.add(i);
                ptr::write(&mut (*slot).stamp, AtomicUsize::new(i));
            }
        }

        // One lap is the smallest power of two greater than `cap`.
        let one_lap = (cap + 1).next_power_of_two();

        ArrayQueue {
            buffer,
            cap,
            one_lap,
            head: CachePadded::new(AtomicUsize::new(head)),
            tail: CachePadded::new(AtomicUsize::new(tail)),
            _marker: PhantomData,
        }
    }

    /// Attempts to push an element into the queue.
    ///
    /// If the queue is full, the element is returned back as an error.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PushError};
    ///
    /// let q = ArrayQueue::new(1);
    ///
    /// assert_eq!(q.push(10), Ok(()));
    /// assert_eq!(q.push(20), Err(PushError(20)));
    /// ```
    pub fn push(&self, value: T) -> Result<(), PushError<T>> {
        let backoff = Backoff::new();
        let mut tail = self.tail.load(Ordering::Relaxed);

        loop {
            // Deconstruct the tail.
            let index = tail & (self.one_lap - 1);
            let lap = tail & !(self.one_lap - 1);

            // Inspect the corresponding slot.
            let slot = unsafe { &*self.buffer.add(index) };
            let stamp = slot.stamp.load(Ordering::Acquire);

            // If the tail and the stamp match, we may attempt to push.
            if tail == stamp {
                let new_tail = if index + 1 < self.cap {
                    // Same lap, incremented index.
                    // Set to `{ lap: lap, index: index + 1 }`.
                    tail + 1
                } else {
                    // One lap forward, index wraps around to zero.
                    // Set to `{ lap: lap.wrapping_add(1), index: 0 }`.
                    lap.wrapping_add(self.one_lap)
                };

                // Try moving the tail.
                match self
                    .tail
                    .compare_exchange_weak(tail, new_tail, Ordering::SeqCst, Ordering::Relaxed)
                {
                    Ok(_) => {
                        // Write the value into the slot and update the stamp.
                        unsafe { slot.value.get().write(value); }
                        slot.stamp.store(tail + 1, Ordering::Release);
                        return Ok(());
                    }
                    Err(t) => {
                        tail = t;
                        backoff.spin();
                    }
                }
            } else if stamp.wrapping_add(self.one_lap) == tail + 1 {
                atomic::fence(Ordering::SeqCst);
                let head = self.head.load(Ordering::Relaxed);

                // If the head lags one lap behind the tail as well...
                if head.wrapping_add(self.one_lap) == tail {
                    // ...then the queue is full.
                    return Err(PushError(value));
                }

                backoff.spin();
                tail = self.tail.load(Ordering::Relaxed);
            } else {
                // Snooze because we need to wait for the stamp to get updated.
                backoff.snooze();
                tail = self.tail.load(Ordering::Relaxed);
            }
        }
    }

    /// Attempts to pop an element from the queue.
    ///
    /// If the queue is empty, an error is returned.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::new(1);
    /// assert_eq!(q.push(10), Ok(()));
    ///
    /// assert_eq!(q.pop(), Ok(10));
    /// assert_eq!(q.pop(), Err(PopError));
    /// ```
    pub fn pop(&self) -> Result<T, PopError> {
        let backoff = Backoff::new();
        let mut head = self.head.load(Ordering::Relaxed);

        loop {
            // Deconstruct the head.
            let index = head & (self.one_lap - 1);
            let lap = head & !(self.one_lap - 1);

            // Inspect the corresponding slot.
            let slot = unsafe { &*self.buffer.add(index) };
            let stamp = slot.stamp.load(Ordering::Acquire);

            // If the the stamp is ahead of the head by 1, we may attempt to pop.
            if head + 1 == stamp {
                let new = if index + 1 < self.cap {
                    // Same lap, incremented index.
                    // Set to `{ lap: lap, index: index + 1 }`.
                    head + 1
                } else {
                    // One lap forward, index wraps around to zero.
                    // Set to `{ lap: lap.wrapping_add(1), index: 0 }`.
                    lap.wrapping_add(self.one_lap)
                };

                // Try moving the head.
                match self
                    .head
                    .compare_exchange_weak(head, new, Ordering::SeqCst, Ordering::Relaxed)
                {
                    Ok(_) => {
                        // Read the value from the slot and update the stamp.
                        let msg = unsafe { slot.value.get().read() };
                        slot.stamp.store(head.wrapping_add(self.one_lap), Ordering::Release);
                        return Ok(msg);
                    }
                    Err(h) => {
                        head = h;
                        backoff.spin();
                    }
                }
            } else if stamp == head {
                atomic::fence(Ordering::SeqCst);
                let tail = self.tail.load(Ordering::Relaxed);

                // If the tail equals the head, that means the channel is empty.
                if tail == head {
                    return Err(PopError);
                }

                backoff.spin();
                head = self.head.load(Ordering::Relaxed);
            } else {
                // Snooze because we need to wait for the stamp to get updated.
                backoff.snooze();
                head = self.head.load(Ordering::Relaxed);
            }
        }
    }

    /// Returns the capacity of the queue.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::<i32>::new(100);
    ///
    /// assert_eq!(q.capacity(), 100);
    /// ```
    pub fn capacity(&self) -> usize {
        self.cap
    }

    /// Returns `true` if the queue is empty.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::new(100);
    ///
    /// assert!(q.is_empty());
    /// q.push(1).unwrap();
    /// assert!(!q.is_empty());
    /// ```
    pub fn is_empty(&self) -> bool {
        let head = self.head.load(Ordering::SeqCst);
        let tail = self.tail.load(Ordering::SeqCst);

        // Is the tail lagging one lap behind head?
        // Is the tail equal to the head?
        //
        // Note: If the head changes just before we load the tail, that means there was a moment
        // when the channel was not empty, so it is safe to just return `false`.
        tail == head
    }

    /// Returns `true` if the queue is full.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::new(1);
    ///
    /// assert!(!q.is_full());
    /// q.push(1).unwrap();
    /// assert!(q.is_full());
    /// ```
    pub fn is_full(&self) -> bool {
        let tail = self.tail.load(Ordering::SeqCst);
        let head = self.head.load(Ordering::SeqCst);

        // Is the head lagging one lap behind tail?
        //
        // Note: If the tail changes just before we load the head, that means there was a moment
        // when the queue was not full, so it is safe to just return `false`.
        head.wrapping_add(self.one_lap) == tail
    }

    /// Returns the number of elements in the queue.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::new(100);
    /// assert_eq!(q.len(), 0);
    ///
    /// q.push(10).unwrap();
    /// assert_eq!(q.len(), 1);
    ///
    /// q.push(20).unwrap();
    /// assert_eq!(q.len(), 2);
    /// ```
    pub fn len(&self) -> usize {
        loop {
            // Load the tail, then load the head.
            let tail = self.tail.load(Ordering::SeqCst);
            let head = self.head.load(Ordering::SeqCst);

            // If the tail didn't change, we've got consistent values to work with.
            if self.tail.load(Ordering::SeqCst) == tail {
                let hix = head & (self.one_lap - 1);
                let tix = tail & (self.one_lap - 1);

                return if hix < tix {
                    tix - hix
                } else if hix > tix {
                    self.cap - hix + tix
                } else if tail == head {
                    0
                } else {
                    self.cap
                };
            }
        }
    }
}

impl<T> Drop for ArrayQueue<T> {
    fn drop(&mut self) {
        // Get the index of the head.
        let hix = self.head.load(Ordering::Relaxed) & (self.one_lap - 1);

        // Loop over all slots that hold a message and drop them.
        for i in 0..self.len() {
            // Compute the index of the next slot holding a message.
            let index = if hix + i < self.cap {
                hix + i
            } else {
                hix + i - self.cap
            };

            unsafe {
                self.buffer.add(index).drop_in_place();
            }
        }

        // Finally, deallocate the buffer, but don't run any destructors.
        unsafe {
            Vec::from_raw_parts(self.buffer, 0, self.cap);
        }
    }
}

impl<T> fmt::Debug for ArrayQueue<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.pad("ArrayQueue { .. }")
    }
}
<commit_msg>Fix a typo<commit_after>//! The implementation is based on Dmitry Vyukov's bounded MPMC queue.
//!
//! Source:
//!   - http://www.1024cores.net/home/lock-free-algorithms/queues/bounded-mpmc-queue
//!
//! Copyright & License:
//!   - Copyright (c) 2010-2011 Dmitry Vyukov
//!   - Simplified BSD License and Apache License, Version 2.0
//!   - http://www.1024cores.net/home/code-license

use std::cell::UnsafeCell;
use std::fmt;
use std::marker::PhantomData;
use std::mem;
use std::ptr;
use std::sync::atomic::{self, AtomicUsize, Ordering};

use crossbeam_utils::{Backoff, CachePadded};

use err::{PopError, PushError};

/// A slot in a queue.
struct Slot<T> {
    /// The current stamp.
    ///
    /// If the stamp equals the tail, this node will be next written to. If it equals the head,
    /// this node will be next read from.
    stamp: AtomicUsize,

    /// The value in this slot.
    value: UnsafeCell<T>,
}

/// A bounded multi-producer multi-consumer queue.
///
/// This queue allocates a fixed-capacity buffer on construction, which is used to store pushed
/// elements. The queue cannot hold more elements than the buffer allows. Attempting to push an
/// element into a full queue will fail. Having a buffer allocated upfront makes this queue a bit
/// faster than [`SegQueue`].
///
/// [`SegQueue`]: struct.SegQueue.html
///
/// # Examples
///
/// ```
/// use crossbeam_queue::{ArrayQueue, PushError};
///
/// let q = ArrayQueue::new(2);
///
/// assert_eq!(q.push('a'), Ok(()));
/// assert_eq!(q.push('b'), Ok(()));
/// assert_eq!(q.push('c'), Err(PushError('c')));
/// assert_eq!(q.pop(), Ok('a'));
/// ```
pub struct ArrayQueue<T> {
    /// The head of the queue.
    ///
    /// This value is a "stamp" consisting of an index into the buffer and a lap, but packed into a
    /// single `usize`. The lower bits represent the index, while the upper bits represent the lap.
    ///
    /// Elements are popped from the head of the queue.
    head: CachePadded<AtomicUsize>,

    /// The tail of the queue.
    ///
    /// This value is a "stamp" consisting of an index into the buffer and a lap, but packed into a
    /// single `usize`. The lower bits represent the index, while the upper bits represent the lap.
    ///
    /// Elements are pushed into the tail of the queue.
    tail: CachePadded<AtomicUsize>,

    /// The buffer holding slots.
    buffer: *mut Slot<T>,

    /// The queue capacity.
    cap: usize,

    /// A stamp with the value of `{ lap: 1, index: 0 }`.
    one_lap: usize,

    /// Indicates that dropping an `ArrayQueue<T>` may drop elements of type `T`.
    _marker: PhantomData<T>,
}

unsafe impl<T: Send> Sync for ArrayQueue<T> {}
unsafe impl<T: Send> Send for ArrayQueue<T> {}

impl<T> ArrayQueue<T> {
    /// Creates a new bounded queue with the given capacity.
    ///
    /// # Panics
    ///
    /// Panics if the capacity is zero.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::ArrayQueue;
    ///
    /// let q = ArrayQueue::<i32>::new(100);
    /// ```
    pub fn new(cap: usize) -> ArrayQueue<T> {
        assert!(cap > 0, "capacity must be non-zero");

        // Head is initialized to `{ lap: 0, index: 0 }`.
        // Tail is initialized to `{ lap: 0, index: 0 }`.
        let head = 0;
        let tail = 0;

        // Allocate a buffer of `cap` slots.
        let buffer = {
            let mut v = Vec::<Slot<T>>::with_capacity(cap);
            let ptr = v.as_mut_ptr();
            mem::forget(v);
            ptr
        };

        // Initialize stamps in the slots.
        for i in 0..cap {
            unsafe {
                // Set the stamp to `{ lap: 0, index: i }`.
                let slot = buffer.add(i);
                ptr::write(&mut (*slot).stamp, AtomicUsize::new(i));
            }
        }

        // One lap is the smallest power of two greater than `cap`.
        let one_lap = (cap + 1).next_power_of_two();

        ArrayQueue {
            buffer,
            cap,
            one_lap,
            head: CachePadded::new(AtomicUsize::new(head)),
            tail: CachePadded::new(AtomicUsize::new(tail)),
            _marker: PhantomData,
        }
    }

    /// Attempts to push an element into the queue.
    ///
    /// If the queue is full, the element is returned back as an error.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PushError};
    ///
    /// let q = ArrayQueue::new(1);
    ///
    /// assert_eq!(q.push(10), Ok(()));
    /// assert_eq!(q.push(20), Err(PushError(20)));
    /// ```
    pub fn push(&self, value: T) -> Result<(), PushError<T>> {
        let backoff = Backoff::new();
        let mut tail = self.tail.load(Ordering::Relaxed);

        loop {
            // Deconstruct the tail.
            let index = tail & (self.one_lap - 1);
            let lap = tail & !(self.one_lap - 1);

            // Inspect the corresponding slot.
            let slot = unsafe { &*self.buffer.add(index) };
            let stamp = slot.stamp.load(Ordering::Acquire);

            // If the tail and the stamp match, we may attempt to push.
            if tail == stamp {
                let new_tail = if index + 1 < self.cap {
                    // Same lap, incremented index.
                    // Set to `{ lap: lap, index: index + 1 }`.
                    tail + 1
                } else {
                    // One lap forward, index wraps around to zero.
                    // Set to `{ lap: lap.wrapping_add(1), index: 0 }`.
                    lap.wrapping_add(self.one_lap)
                };

                // Try moving the tail.
                match self
                    .tail
                    .compare_exchange_weak(tail, new_tail, Ordering::SeqCst, Ordering::Relaxed)
                {
                    Ok(_) => {
                        // Write the value into the slot and update the stamp.
                        unsafe { slot.value.get().write(value); }
                        slot.stamp.store(tail + 1, Ordering::Release);
                        return Ok(());
                    }
                    Err(t) => {
                        tail = t;
                        backoff.spin();
                    }
                }
            } else if stamp.wrapping_add(self.one_lap) == tail + 1 {
                atomic::fence(Ordering::SeqCst);
                let head = self.head.load(Ordering::Relaxed);

                // If the head lags one lap behind the tail as well...
                if head.wrapping_add(self.one_lap) == tail {
                    // ...then the queue is full.
                    return Err(PushError(value));
                }

                backoff.spin();
                tail = self.tail.load(Ordering::Relaxed);
            } else {
                // Snooze because we need to wait for the stamp to get updated.
                backoff.snooze();
                tail = self.tail.load(Ordering::Relaxed);
            }
        }
    }

    /// Attempts to pop an element from the queue.
    ///
    /// If the queue is empty, an error is returned.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::new(1);
    /// assert_eq!(q.push(10), Ok(()));
    ///
    /// assert_eq!(q.pop(), Ok(10));
    /// assert_eq!(q.pop(), Err(PopError));
    /// ```
    pub fn pop(&self) -> Result<T, PopError> {
        let backoff = Backoff::new();
        let mut head = self.head.load(Ordering::Relaxed);

        loop {
            // Deconstruct the head.
            let index = head & (self.one_lap - 1);
            let lap = head & !(self.one_lap - 1);

            // Inspect the corresponding slot.
            let slot = unsafe { &*self.buffer.add(index) };
            let stamp = slot.stamp.load(Ordering::Acquire);

            // If the the stamp is ahead of the head by 1, we may attempt to pop.
            if head + 1 == stamp {
                let new = if index + 1 < self.cap {
                    // Same lap, incremented index.
                    // Set to `{ lap: lap, index: index + 1 }`.
                    head + 1
                } else {
                    // One lap forward, index wraps around to zero.
                    // Set to `{ lap: lap.wrapping_add(1), index: 0 }`.
                    lap.wrapping_add(self.one_lap)
                };

                // Try moving the head.
                match self
                    .head
                    .compare_exchange_weak(head, new, Ordering::SeqCst, Ordering::Relaxed)
                {
                    Ok(_) => {
                        // Read the value from the slot and update the stamp.
                        let msg = unsafe { slot.value.get().read() };
                        slot.stamp.store(head.wrapping_add(self.one_lap), Ordering::Release);
                        return Ok(msg);
                    }
                    Err(h) => {
                        head = h;
                        backoff.spin();
                    }
                }
            } else if stamp == head {
                atomic::fence(Ordering::SeqCst);
                let tail = self.tail.load(Ordering::Relaxed);

                // If the tail equals the head, that means the channel is empty.
                if tail == head {
                    return Err(PopError);
                }

                backoff.spin();
                head = self.head.load(Ordering::Relaxed);
            } else {
                // Snooze because we need to wait for the stamp to get updated.
                backoff.snooze();
                head = self.head.load(Ordering::Relaxed);
            }
        }
    }

    /// Returns the capacity of the queue.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::<i32>::new(100);
    ///
    /// assert_eq!(q.capacity(), 100);
    /// ```
    pub fn capacity(&self) -> usize {
        self.cap
    }

    /// Returns `true` if the queue is empty.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::new(100);
    ///
    /// assert!(q.is_empty());
    /// q.push(1).unwrap();
    /// assert!(!q.is_empty());
    /// ```
    pub fn is_empty(&self) -> bool {
        let head = self.head.load(Ordering::SeqCst);
        let tail = self.tail.load(Ordering::SeqCst);

        // Is the tail lagging one lap behind head?
        // Is the tail equal to the head?
        //
        // Note: If the head changes just before we load the tail, that means there was a moment
        // when the channel was not empty, so it is safe to just return `false`.
        tail == head
    }

    /// Returns `true` if the queue is full.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::new(1);
    ///
    /// assert!(!q.is_full());
    /// q.push(1).unwrap();
    /// assert!(q.is_full());
    /// ```
    pub fn is_full(&self) -> bool {
        let tail = self.tail.load(Ordering::SeqCst);
        let head = self.head.load(Ordering::SeqCst);

        // Is the head lagging one lap behind tail?
        //
        // Note: If the tail changes just before we load the head, that means there was a moment
        // when the queue was not full, so it is safe to just return `false`.
        head.wrapping_add(self.one_lap) == tail
    }

    /// Returns the number of elements in the queue.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::new(100);
    /// assert_eq!(q.len(), 0);
    ///
    /// q.push(10).unwrap();
    /// assert_eq!(q.len(), 1);
    ///
    /// q.push(20).unwrap();
    /// assert_eq!(q.len(), 2);
    /// ```
    pub fn len(&self) -> usize {
        loop {
            // Load the tail, then load the head.
            let tail = self.tail.load(Ordering::SeqCst);
            let head = self.head.load(Ordering::SeqCst);

            // If the tail didn't change, we've got consistent values to work with.
            if self.tail.load(Ordering::SeqCst) == tail {
                let hix = head & (self.one_lap - 1);
                let tix = tail & (self.one_lap - 1);

                return if hix < tix {
                    tix - hix
                } else if hix > tix {
                    self.cap - hix + tix
                } else if tail == head {
                    0
                } else {
                    self.cap
                };
            }
        }
    }
}

impl<T> Drop for ArrayQueue<T> {
    fn drop(&mut self) {
        // Get the index of the head.
        let hix = self.head.load(Ordering::Relaxed) & (self.one_lap - 1);

        // Loop over all slots that hold a message and drop them.
        for i in 0..self.len() {
            // Compute the index of the next slot holding a message.
            let index = if hix + i < self.cap {
                hix + i
            } else {
                hix + i - self.cap
            };

            unsafe {
                self.buffer.add(index).drop_in_place();
            }
        }

        // Finally, deallocate the buffer, but don't run any destructors.
        unsafe {
            Vec::from_raw_parts(self.buffer, 0, self.cap);
        }
    }
}

impl<T> fmt::Debug for ArrayQueue<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.pad("ArrayQueue { .. }")
    }
}
<|endoftext|>


<commit_before>use std::marker::PhantomData;
use std::mem;
use util::locksteparray;
use super::jpm::jpm_root::Jpm;
use ::Key;
use ::rudymap::results::InsertResult;
use std::iter;
use super::rootptr::RootPtr;

pub trait RootLeaf<K: Key, V> {
    fn get(&self, key: K) -> Option<&V>;
    fn get_mut(&mut self, key: K) -> Option<&mut V>;
    fn insert(&mut self, key: K, value: V) -> InsertResult<V>;
    fn expand(self: Box<Self>, key: K, value: V) -> RootPtr<K, V>;
    fn len(&self) -> usize;
}

pub struct Empty<K: Key, V>(PhantomData<(K, V)>);

impl<K: Key, V> Empty<K, V> {
    pub fn new() -> Empty<K, V> {
        Empty(PhantomData)
    }
}

impl<K: Key, V> RootLeaf<K, V> for Empty<K, V> {
    fn get(&self, key: K) -> Option<&V> {
        None
    }

    fn get_mut(&mut self, key: K) -> Option<&mut V> {
        None
    }

    fn insert(&mut self, key: K, value: V) -> InsertResult<V> {
        InsertResult::Resize(value)
    }

    fn expand(self: Box<Self>, key: K, value: V) -> RootPtr<K, V> {
        Box::new(Leaf1::new(key, value)).into()
    }

    fn len(&self) -> usize {
        0
    }
}

impl<'a, K: Key + 'a, V: 'a> IntoIterator for &'a Empty<K, V> {
    type Item = (K, &'a V);
    type IntoIter = iter::Empty<Self::Item>;
    fn into_iter(self) -> Self::IntoIter {
        iter::empty()
    }
}

pub struct Leaf1<K: Key, V> {
    key: K,
    value: V
}

impl<K: Key, V> Leaf1<K, V> {
    pub fn new(key: K, value: V) -> Leaf1<K, V> {
        Leaf1 { key, value }
    }
}

impl<'a, K: Key + 'a, V: 'a> IntoIterator for &'a Leaf1<K, V> {
    type Item = (K, &'a V);
    type IntoIter = iter::Once<Self::Item>;
    fn into_iter(self) -> Self::IntoIter {
        iter::once((self.key, &self.value))
    }
}

/// A leaf root with one item.
impl<K: Key, V> RootLeaf<K, V> for Leaf1<K, V> {
    fn get(&self, key: K) -> Option<&V> {
        if self.key == key {
            Some(&self.value)
        } else {
            None
        }
    }

    fn get_mut(&mut self, key: K) -> Option<&mut V> {
        if self.key == key {
            Some(&mut self.value)
        } else {
            None
        }
    }

    fn insert(&mut self, key: K, value: V) -> InsertResult<V> {
        if self.key == key {
            InsertResult::replace(&mut self.value, value)
        } else {
            InsertResult::Resize(value)
        }
    }

    fn expand(self: Box<Self>, key: K, value: V) -> RootPtr<K, V> {
        Box::new(Leaf2::new(self.key, self.value, key, value)).into()
    }

    fn len(&self) -> usize {
        1
    }
}

pub struct Leaf2<K: Key, V> {
    keys: [K; 2],
    values: [V; 2]
}

impl<K: Key, V> Leaf2<K, V> {
    pub fn new(key1: K, value1: V, key2: K, value2: V) -> Leaf2<K, V> {
        if key1 < key2 {
            Leaf2 {
                keys: [key1, key2],
                values: [value1, value2]
            }
        } else {
            Leaf2 {
                keys: [key2, key1],
                values: [value2, value1]
            }
        }
    }
}

impl<K: Key, V> RootLeaf<K, V> for Leaf2<K, V> {
    fn get(&self, key: K) -> Option<&V> {
        self.keys.iter()
            .zip(self.values.iter())
            .find(|&(&leaf_key, _)| leaf_key == key)
            .map(|(key, value)| value)
    }

    fn get_mut(&mut self, key: K) -> Option<&mut V> {
        self.keys.iter()
            .zip(self.values.iter_mut())
            .find(|&(&leaf_key, _)| leaf_key == key)
            .map(|(key, value)| value)
    }

    /// Attempt to insert, fail if we didn't find a key to replace
    fn insert(&mut self, key: K, value: V) -> InsertResult<V> {
        for (i, leaf_key) in self.keys.iter().enumerate() {
            if key == *leaf_key {
                return InsertResult::replace(&mut self.values[i], value);
            }
        }
        InsertResult::Resize(value)
    }

    fn expand(self: Box<Self>, key: K, value: V) -> RootPtr<K, V> {
        let Leaf2 { keys, values } = *self;
        let mut leaf = Box::new(VecLeaf::from_arrays(keys, values));
        leaf.insert(key, value).success();
        leaf.into()
    }

    fn len(&self) -> usize {
        2
    }
}

pub struct VecLeaf<K: Key, V> {
    array: locksteparray::LockstepArray<[K; 31], [V; 31]>
}

impl<K: Key, V> VecLeaf<K, V> {
    fn new() -> VecLeaf<K, V> {
        // TODO Copy memory from values
        VecLeaf {
            array: locksteparray::LockstepArray::new()
        }
    }

    fn from_arrays(keys: [K; 2], values: [V; 2]) -> VecLeaf<K, V> {
        VecLeaf {
            array: locksteparray::LockstepArray::from_arrays(keys, values)
        }
    }
}

impl<K: Key, V> IntoIterator for VecLeaf<K, V> {
    type Item = (K, V);
    type IntoIter = locksteparray::IntoIter<[K; 31], [V; 31]>;

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

impl<K: Key, V> RootLeaf<K, V> for VecLeaf<K, V> {
    fn get(&self, key: K) -> Option<&V> {
        self.array.array1()
            .binary_search(&key)
            .ok()
            .map(|index| &self.array.array2()[index])
    }

    fn get_mut(&mut self, key: K) -> Option<&mut V> {
        self.array.array1_mut()
            .binary_search(&key)
            .ok()
            .map(move |index| &mut self.array.array2_mut()[index])
    }

    fn insert(&mut self, key: K, value: V) -> InsertResult<V> {
        match self.array.array1().binary_search(&key) {
            Ok(replace) => {
                InsertResult::replace(&mut self.array.array2_mut()[replace],
                                      value)
            },
            Err(insert) => match self.array.insert(insert, key, value) {
                Ok(()) => InsertResult::Success(None),
                Err(locksteparray::InsertError::Overflow(key, value)) => {
                    InsertResult::Resize(value)
                },
                Err(locksteparray::InsertError::OutOfBounds(..)) => {
                    unreachable!()
                }
            }
        }
    }

    fn expand(self: Box<Self>, key: K, value: V) -> RootPtr<K, V> {
        let jpm: Jpm<K, V> = self.into_iter().collect();
        Box::new(jpm).into()
    }

    fn len(&self) -> usize {
        self.array.len()
    }
}
<commit_msg>Fix missed insertion<commit_after>use std::marker::PhantomData;
use std::mem;
use util::locksteparray;
use super::jpm::jpm_root::Jpm;
use ::Key;
use ::rudymap::results::InsertResult;
use std::iter;
use super::rootptr::RootPtr;

pub trait RootLeaf<K: Key, V> {
    fn get(&self, key: K) -> Option<&V>;
    fn get_mut(&mut self, key: K) -> Option<&mut V>;
    fn insert(&mut self, key: K, value: V) -> InsertResult<V>;
    fn expand(self: Box<Self>, key: K, value: V) -> RootPtr<K, V>;
    fn len(&self) -> usize;
}

pub struct Empty<K: Key, V>(PhantomData<(K, V)>);

impl<K: Key, V> Empty<K, V> {
    pub fn new() -> Empty<K, V> {
        Empty(PhantomData)
    }
}

impl<K: Key, V> RootLeaf<K, V> for Empty<K, V> {
    fn get(&self, key: K) -> Option<&V> {
        None
    }

    fn get_mut(&mut self, key: K) -> Option<&mut V> {
        None
    }

    fn insert(&mut self, key: K, value: V) -> InsertResult<V> {
        InsertResult::Resize(value)
    }

    fn expand(self: Box<Self>, key: K, value: V) -> RootPtr<K, V> {
        Box::new(Leaf1::new(key, value)).into()
    }

    fn len(&self) -> usize {
        0
    }
}

impl<'a, K: Key + 'a, V: 'a> IntoIterator for &'a Empty<K, V> {
    type Item = (K, &'a V);
    type IntoIter = iter::Empty<Self::Item>;
    fn into_iter(self) -> Self::IntoIter {
        iter::empty()
    }
}

pub struct Leaf1<K: Key, V> {
    key: K,
    value: V
}

impl<K: Key, V> Leaf1<K, V> {
    pub fn new(key: K, value: V) -> Leaf1<K, V> {
        Leaf1 { key, value }
    }
}

impl<'a, K: Key + 'a, V: 'a> IntoIterator for &'a Leaf1<K, V> {
    type Item = (K, &'a V);
    type IntoIter = iter::Once<Self::Item>;
    fn into_iter(self) -> Self::IntoIter {
        iter::once((self.key, &self.value))
    }
}

/// A leaf root with one item.
impl<K: Key, V> RootLeaf<K, V> for Leaf1<K, V> {
    fn get(&self, key: K) -> Option<&V> {
        if self.key == key {
            Some(&self.value)
        } else {
            None
        }
    }

    fn get_mut(&mut self, key: K) -> Option<&mut V> {
        if self.key == key {
            Some(&mut self.value)
        } else {
            None
        }
    }

    fn insert(&mut self, key: K, value: V) -> InsertResult<V> {
        if self.key == key {
            InsertResult::replace(&mut self.value, value)
        } else {
            InsertResult::Resize(value)
        }
    }

    fn expand(self: Box<Self>, key: K, value: V) -> RootPtr<K, V> {
        Box::new(Leaf2::new(self.key, self.value, key, value)).into()
    }

    fn len(&self) -> usize {
        1
    }
}

pub struct Leaf2<K: Key, V> {
    keys: [K; 2],
    values: [V; 2]
}

impl<K: Key, V> Leaf2<K, V> {
    pub fn new(key1: K, value1: V, key2: K, value2: V) -> Leaf2<K, V> {
        if key1 < key2 {
            Leaf2 {
                keys: [key1, key2],
                values: [value1, value2]
            }
        } else {
            Leaf2 {
                keys: [key2, key1],
                values: [value2, value1]
            }
        }
    }
}

impl<K: Key, V> RootLeaf<K, V> for Leaf2<K, V> {
    fn get(&self, key: K) -> Option<&V> {
        self.keys.iter()
            .zip(self.values.iter())
            .find(|&(&leaf_key, _)| leaf_key == key)
            .map(|(key, value)| value)
    }

    fn get_mut(&mut self, key: K) -> Option<&mut V> {
        self.keys.iter()
            .zip(self.values.iter_mut())
            .find(|&(&leaf_key, _)| leaf_key == key)
            .map(|(key, value)| value)
    }

    /// Attempt to insert, fail if we didn't find a key to replace
    fn insert(&mut self, key: K, value: V) -> InsertResult<V> {
        for (i, leaf_key) in self.keys.iter().enumerate() {
            if key == *leaf_key {
                return InsertResult::replace(&mut self.values[i], value);
            }
        }
        InsertResult::Resize(value)
    }

    fn expand(self: Box<Self>, key: K, value: V) -> RootPtr<K, V> {
        let Leaf2 { keys, values } = *self;
        let mut leaf = Box::new(VecLeaf::from_arrays(keys, values));
        leaf.insert(key, value).success();
        leaf.into()
    }

    fn len(&self) -> usize {
        2
    }
}

pub struct VecLeaf<K: Key, V> {
    array: locksteparray::LockstepArray<[K; 31], [V; 31]>
}

impl<K: Key, V> VecLeaf<K, V> {
    fn new() -> VecLeaf<K, V> {
        // TODO Copy memory from values
        VecLeaf {
            array: locksteparray::LockstepArray::new()
        }
    }

    fn from_arrays(keys: [K; 2], values: [V; 2]) -> VecLeaf<K, V> {
        VecLeaf {
            array: locksteparray::LockstepArray::from_arrays(keys, values)
        }
    }
}

impl<K: Key, V> IntoIterator for VecLeaf<K, V> {
    type Item = (K, V);
    type IntoIter = locksteparray::IntoIter<[K; 31], [V; 31]>;

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

impl<K: Key, V> RootLeaf<K, V> for VecLeaf<K, V> {
    fn get(&self, key: K) -> Option<&V> {
        self.array.array1()
            .binary_search(&key)
            .ok()
            .map(|index| &self.array.array2()[index])
    }

    fn get_mut(&mut self, key: K) -> Option<&mut V> {
        self.array.array1_mut()
            .binary_search(&key)
            .ok()
            .map(move |index| &mut self.array.array2_mut()[index])
    }

    fn insert(&mut self, key: K, value: V) -> InsertResult<V> {
        match self.array.array1().binary_search(&key) {
            Ok(replace) => {
                InsertResult::replace(&mut self.array.array2_mut()[replace],
                                      value)
            },
            Err(insert) => match self.array.insert(insert, key, value) {
                Ok(()) => InsertResult::Success(None),
                Err(locksteparray::InsertError::Overflow(key, value)) => {
                    InsertResult::Resize(value)
                },
                Err(locksteparray::InsertError::OutOfBounds(..)) => {
                    unreachable!()
                }
            }
        }
    }

    fn expand(self: Box<Self>, key: K, value: V) -> RootPtr<K, V> {
        let mut jpm: Jpm<K, V> = self.into_iter().collect();
        jpm.insert(key, value).success();
        Box::new(jpm).into()
    }

    fn len(&self) -> usize {
        self.array.len()
    }
}
<|endoftext|>


<commit_before>use byteread::*;

pub trait BitReadEndian {
    fn peek_val(&mut self, n:usize) -> u64;
    fn merge_val(msp:u64, lsp:u64, msb:usize, lsb:usize) -> u64;
    fn skip_rem(&mut self, n:usize) -> ();
}

pub trait BitReadFill {
    fn can_refill(&self) -> bool;
    fn fill32(&self) -> u64;
    fn fill64(&self) -> u64;
}

pub trait BitReadInternal : BitReadEndian + BitReadFill {
    #[inline]
    fn left(&self) -> usize;
    fn refill32(&mut self) -> ();
    fn refill64(&mut self) -> ();

    fn get_val(&mut self, n:usize) -> u64 {
        let ret = self.peek_val(n);

        self.skip_rem(n);

        return ret;
    }

}

pub trait BitRead<'a>: BitReadInternal+Copy {
    fn new(&'a[u8]) -> Self;
    fn consumed(&self) -> usize;
    fn available(&self) -> usize;

    fn skip_bits(&mut self, size : usize) -> ();

    #[inline]
    fn get_bit(&mut self) -> bool {
        if self.left() <= 0 {
            self.refill64();
        }

        self.get_val(1) != 0
    }

    #[inline]
    fn get_bits_64(&mut self, mut n:usize) -> u64 {
        let mut left = 0;
        let mut ret = 0;

        if n == 0 {
            return 0;
        }

        if self.left() < n {
            n   -= self.left();
            left = self.left();
            ret  = self.get_val(left);
            self.refill64();
        }

        Self::merge_val(self.get_val(n), ret, left, n)
    }

    #[inline]
    fn get_bits_32(&mut self, n:usize) -> u32 {
        if n == 0 {
            return 0;
        }

        if self.left() <= n {
            self.refill32();
        }

        return self.get_val(n) as u32;
    }

    #[inline]
    fn peek_bit(&mut self) -> bool {
        let mut tmp = *self;

        tmp.get_bit()
    }

    #[inline]
    fn peek_bits_32(&mut self, n:usize) -> u32 {
        let mut tmp = *self;

        tmp.get_bits_32(n)
    }

    #[inline]
    fn peek_bits_64(&self, n:usize) -> u64 {
        let mut tmp = *self;

        tmp.get_bits_64(n)
    }

    #[inline]
    fn align_bits(&mut self) -> () {
        let left = self.left() & 7;

        self.skip_bits(left);
    }
}

#[macro_export]
macro_rules! endian_reader {
    {$name: ident} => {
        #[derive(Debug, Clone, Copy)]
        pub struct $name<'a> {
            buffer : &'a[u8], /// read buffer, 8-bytes padded
            index : usize,
            cache : u64,
            left : usize,
        }
        impl <'a> BitReadInternal for $name<'a> {
            #[inline]
            fn left(&self) -> usize {
                self.left
            }
            #[inline]
            fn refill32(&mut self) -> () {
                if !self.can_refill() {
                    return;
                }
                let val = self.fill32();

                self.cache  = Self::merge_val(val, self.cache,
                                              self.left, 32 - self.left);
                self.index += 4;
                self.left  += 32;
            }
            #[inline]
            fn refill64(&mut self) -> () {
                if !self.can_refill() {
                    return;
                }

                self.cache  = self.fill64();
                self.index += 8;
                self.left   = 64;
            }
        }

        impl <'a> BitRead<'a> for $name<'a> {
            fn new(buffer: &'a[u8]) -> $name<'a> {
                let mut reader = $name {
                    buffer: buffer,
                    index: 0,
                    cache: 0,
                    left: 0
                };

                reader.refill64();
                return reader;
            }
            #[inline]
            fn consumed(&self) -> usize {
                self.index * 8 - self.left
            }

            #[inline]
            fn available(&self) -> usize {
                (self.buffer.len() - self.index) * 8 + self.left
            }

            #[inline]
            fn skip_bits(&mut self, mut n:usize) -> () {
                if self.left < n {
                    n -= self.left;
                    if n > 64 {
                        let skip = n / 8;

                        n -= skip * 8;
                        self.index += skip;
                    }
                    self.skip_rem(n);
                    self.refill64();
                }

                self.skip_rem(n);
            }

        }
    }
}

#[macro_export]
macro_rules! little_endian_reader {
    {$name: ident} => {
        endian_reader!{ $name }

        impl <'a> BitReadEndian for $name<'a> {
            #[inline]
            fn peek_val(&mut self, n:usize) -> u64 {
                let v = self.cache & ((1u64 << n) - 1);

                return v;
            }
            #[inline]
            fn skip_rem(&mut self, n:usize) -> () {
                self.cache = self.cache >> n;
                self.left = self.left.saturating_sub(n);
            }
            #[inline]
            fn merge_val(msp:u64, lsp:u64, msb:usize, _:usize) -> u64 {
                msp << msb | lsp
            }
        }
    }
}

little_endian_reader!{ BitReadLE }

impl <'a> BitReadFill for BitReadLE<'a> {
    #[inline]
    fn can_refill(&self) -> bool {
        self.index + 8 <= self.buffer.len()
    }
    #[inline(always)]
    fn fill32(&self) -> u64 {
        get_u32l(&self.buffer[self.index..]) as u64
    }
    #[inline(always)]
    fn fill64(&self) -> u64 {
        get_u64l(&self.buffer[self.index..])
    }
}

#[macro_export]
macro_rules! big_endian_reader {
    {$name: ident} => {
        endian_reader!{ $name }

        impl <'a> BitReadEndian for $name<'a> {
            #[inline]
            fn peek_val(&mut self, n:usize) -> u64 {
                self.cache >> (64 - n)
            }
            #[inline]
            fn skip_rem(&mut self, n:usize) -> () {
                self.cache = self.cache << n;
                self.left = self.left.saturating_sub(n);
            }
            #[inline]
            fn merge_val(msp:u64, lsp:u64, _:usize, lsb:usize) -> u64 {
                msp | lsp << lsb
            }
        }
    }
}

big_endian_reader!{ BitReadBE }

impl <'a> BitReadFill for BitReadBE<'a> {
    #[inline]
    fn can_refill(&self) -> bool {
        self.index + 8 <= self.buffer.len()
    }
    #[inline(always)]
    fn fill32(&self) -> u64 {
        get_u32b(&self.buffer[self.index..]) as u64
    }
    #[inline(always)]
    fn fill64(&self) -> u64 {
        get_u64b(&self.buffer[self.index..])
    }
}

#[cfg(test)]
mod test {
    pub const CHECKBOARD0101: [u8; 128] = [0b01010101; 128];
    pub const CHECKBOARD0011: [u8; 128] = [0b00110011; 128];

    mod le {
        use super::super::*;
        use super::*;

        #[test]
        fn get_bit() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.get_bit());
            assert!(!reader.get_bit());
        }

        #[test]
        fn get_bits_64() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.get_bits_64(1) == 1);
            assert!(reader.get_bits_64(2) == 2);
            assert!(reader.get_bits_64(4) == 10);
            assert!(reader.get_bits_64(1) == 0);
            assert!(reader.get_bits_64(8) == 85);
        }

        #[test]
        fn peek_bits_64() {
            let reader = BitReadLE {
                buffer: &CHECKBOARD0101,
                index: 0,
                cache: 0,
                left: 0
            };

            assert!(reader.peek_bits_64(1) == 1);
            assert!(reader.peek_bits_64(1) == 1);
            assert!(reader.peek_bits_64(2) == 1);
            assert!(reader.peek_bits_64(2) == 1);
        }

        #[test]
        fn get_bits_32() {
            let mut reader = BitReadLE {
                buffer: &CHECKBOARD0101,
                index: 0,
                cache: 0,
                left: 0
            };

            assert!(reader.get_bits_32(1) == 1);
            assert!(reader.get_bits_32(2) == 2);
            assert!(reader.get_bits_32(4) == 10);
            assert!(reader.get_bits_32(1) == 0);
            assert!(reader.get_bits_32(8) == 85);
        }

        #[test]
        fn peek_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.peek_bits_32(1) == 1);
            assert!(reader.peek_bits_32(1) == 1);
            assert!(reader.peek_bits_32(2) == 1);
            assert!(reader.peek_bits_32(2) == 1);
        }

        #[test]
        fn skip_bits() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            reader.skip_bits(0);
            assert!(reader.peek_bits_32(1) == 1);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 1);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 1);
        }

        #[test]
        fn align_bits() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadLE::new(b);

            reader.align_bits();
            assert!(reader.get_bits_64(3) == 3);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
            reader.skip_bits(1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
        }

        #[test]
        fn overread() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadLE::new(b);

            reader.skip_bits(128*8+2);
            reader.get_bits_64(6);
        }
    }
    mod be {
        use super::super::*;
        use super::*;

        #[test]
        fn get_bit() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(!reader.get_bit());
            assert!(reader.get_bit());
        }

        #[test]
        fn get_bits_64() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.get_bits_64(1) == 0);
            assert!(reader.get_bits_64(2) == 2);
            assert!(reader.get_bits_64(4) == 10);
            assert!(reader.get_bits_64(1) == 1);
            assert!(reader.get_bits_64(8) == 85);
        }

        #[test]
        fn peek_bits_64() {
            let b = &CHECKBOARD0101;
            let reader = BitReadBE::new(b);

            assert!(reader.peek_bits_64(1) == 0);
            assert!(reader.peek_bits_64(1) == 0);
            assert!(reader.peek_bits_64(2) == 1);
            assert!(reader.peek_bits_64(2) == 1);
        }

        #[test]
        fn get_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.get_bits_32(1) == 0);
            assert!(reader.get_bits_32(2) == 2);
            assert!(reader.get_bits_32(4) == 10);
            assert!(reader.get_bits_32(1) == 1);
            assert!(reader.get_bits_32(8) == 85);
        }

        #[test]
        fn peek_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.peek_bits_32(1) == 0);
            assert!(reader.peek_bits_32(1) == 0);
            assert!(reader.peek_bits_32(2) == 1);
            assert!(reader.peek_bits_32(2) == 1);
        }

        #[test]
        fn skip_bits() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            reader.skip_bits(0);
            assert!(reader.peek_bits_32(1) == 0);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 0);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 0);
        }

        #[test]
        fn align_bits() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadBE::new(b);

            reader.align_bits();
            assert!(reader.get_bits_64(3) == 1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
            reader.skip_bits(1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
        }

        #[test]
        fn overread() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadBE::new(b);

            reader.skip_bits(128*8+2);
            reader.get_bits_64(6);
        }
    }
}
<commit_msg>bitstream: Do the short refill correctly<commit_after>use byteread::*;

pub trait BitReadEndian {
    fn peek_val(&mut self, n:usize) -> u64;
    fn merge_val(msp:u64, lsp:u64, msb:usize, lsb:usize) -> u64;
    fn build_cache(cache:u64, refill:u64, cache_size:usize) -> u64;
    fn skip_rem(&mut self, n:usize) -> ();
}

pub trait BitReadFill {
    fn can_refill(&self) -> bool;
    fn fill32(&self) -> u64;
    fn fill64(&self) -> u64;
}

pub trait BitReadInternal : BitReadEndian + BitReadFill {
    #[inline]
    fn left(&self) -> usize;
    fn refill32(&mut self) -> ();
    fn refill64(&mut self) -> ();

    fn get_val(&mut self, n:usize) -> u64 {
        let ret = self.peek_val(n);

        self.skip_rem(n);

        return ret;
    }

}

pub trait BitRead<'a>: BitReadInternal+Copy {
    fn new(&'a[u8]) -> Self;
    fn consumed(&self) -> usize;
    fn available(&self) -> usize;

    fn skip_bits(&mut self, size : usize) -> ();

    #[inline]
    fn get_bit(&mut self) -> bool {
        if self.left() <= 0 {
            self.refill64();
        }

        self.get_val(1) != 0
    }

    #[inline]
    fn get_bits_64(&mut self, mut n:usize) -> u64 {
        let mut left = 0;
        let mut ret = 0;

        if n == 0 {
            return 0;
        }

        if self.left() < n {
            n   -= self.left();
            left = self.left();
            ret  = self.get_val(left);
            self.refill64();
        }

        Self::merge_val(self.get_val(n), ret, left, n)
    }

    #[inline]
    fn get_bits_32(&mut self, n:usize) -> u32 {
        if n == 0 {
            return 0;
        }

        if self.left() <= n {
            self.refill32();
        }

        return self.get_val(n) as u32;
    }

    #[inline]
    fn peek_bit(&mut self) -> bool {
        let mut tmp = *self;

        tmp.get_bit()
    }

    #[inline]
    fn peek_bits_32(&mut self, n:usize) -> u32 {
        let mut tmp = *self;

        tmp.get_bits_32(n)
    }

    #[inline]
    fn peek_bits_64(&self, n:usize) -> u64 {
        let mut tmp = *self;

        tmp.get_bits_64(n)
    }

    #[inline]
    fn align_bits(&mut self) -> () {
        let left = self.left() & 7;

        self.skip_bits(left);
    }
}

#[macro_export]
macro_rules! endian_reader {
    {$name: ident} => {
        #[derive(Debug, Clone, Copy)]
        pub struct $name<'a> {
            buffer : &'a[u8], /// read buffer, 8-bytes padded
            index : usize,
            cache : u64,
            left : usize,
        }
        impl <'a> BitReadInternal for $name<'a> {
            #[inline]
            fn left(&self) -> usize {
                self.left
            }
            #[inline]
            fn refill32(&mut self) -> () {
                if !self.can_refill() {
                    return;
                }
                let val = self.fill32();

                self.cache  = Self::build_cache(self.cache, val, self.left);

                self.index += 4;
                self.left  += 32;
            }
            #[inline]
            fn refill64(&mut self) -> () {
                if !self.can_refill() {
                    return;
                }

                self.cache  = self.fill64();
                self.index += 8;
                self.left   = 64;
            }
        }

        impl <'a> BitRead<'a> for $name<'a> {
            fn new(buffer: &'a[u8]) -> $name<'a> {
                let mut reader = $name {
                    buffer: buffer,
                    index: 0,
                    cache: 0,
                    left: 0
                };

                reader.refill64();
                return reader;
            }
            #[inline]
            fn consumed(&self) -> usize {
                self.index * 8 - self.left
            }

            #[inline]
            fn available(&self) -> usize {
                (self.buffer.len() - self.index) * 8 + self.left
            }

            #[inline]
            fn skip_bits(&mut self, mut n:usize) -> () {
                if self.left < n {
                    n -= self.left;
                    if n > 64 {
                        let skip = n / 8;

                        n -= skip * 8;
                        self.index += skip;
                    }
                    self.skip_rem(n);
                    self.refill64();
                }

                self.skip_rem(n);
            }

        }
    }
}

#[macro_export]
macro_rules! little_endian_reader {
    {$name: ident} => {
        endian_reader!{ $name }

        impl <'a> BitReadEndian for $name<'a> {
            #[inline]
            fn peek_val(&mut self, n:usize) -> u64 {
                let v = self.cache & ((1u64 << n) - 1);

                return v;
            }
            #[inline]
            fn skip_rem(&mut self, n:usize) -> () {
                self.cache = self.cache >> n;
                self.left = self.left.saturating_sub(n);
            }
            #[inline]
            fn merge_val(msp:u64, lsp:u64, msb:usize, _:usize) -> u64 {
                msp << msb | lsp
            }
            #[inline]
            fn build_cache(cache:u64, refill:u64, cache_size:usize) -> u64 {
                cache | refill << cache_size
            }
        }
    }
}

little_endian_reader!{ BitReadLE }

impl <'a> BitReadFill for BitReadLE<'a> {
    #[inline]
    fn can_refill(&self) -> bool {
        self.index + 8 <= self.buffer.len()
    }
    #[inline(always)]
    fn fill32(&self) -> u64 {
        get_u32l(&self.buffer[self.index..]) as u64
    }
    #[inline(always)]
    fn fill64(&self) -> u64 {
        get_u64l(&self.buffer[self.index..])
    }
}

#[macro_export]
macro_rules! big_endian_reader {
    {$name: ident} => {
        endian_reader!{ $name }

        impl <'a> BitReadEndian for $name<'a> {
            #[inline]
            fn peek_val(&mut self, n:usize) -> u64 {
                self.cache >> (64 - n)
            }
            #[inline]
            fn skip_rem(&mut self, n:usize) -> () {
                self.cache = self.cache << n;
                self.left = self.left.saturating_sub(n);
            }
            #[inline]
            fn merge_val(msp:u64, lsp:u64, _:usize, lsb:usize) -> u64 {
                msp | lsp << lsb
            }
            #[inline]
            fn build_cache(cache:u64, refill:u64, cache_size:usize) -> u64 {
                cache | refill << (32 - cache_size)
            }
        }
    }
}

big_endian_reader!{ BitReadBE }

impl <'a> BitReadFill for BitReadBE<'a> {
    #[inline]
    fn can_refill(&self) -> bool {
        self.index + 8 <= self.buffer.len()
    }
    #[inline(always)]
    fn fill32(&self) -> u64 {
        get_u32b(&self.buffer[self.index..]) as u64
    }
    #[inline(always)]
    fn fill64(&self) -> u64 {
        get_u64b(&self.buffer[self.index..])
    }
}

#[cfg(test)]
mod test {
    pub const CHECKBOARD0101: [u8; 128] = [0b01010101; 128];
    pub const CHECKBOARD0011: [u8; 128] = [0b00110011; 128];

    mod le {
        use super::super::*;
        use super::*;

        #[test]
        fn get_bit() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.get_bit());
            assert!(!reader.get_bit());
        }

        #[test]
        fn get_bits_64() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.get_bits_64(1) == 1);
            assert!(reader.get_bits_64(2) == 2);
            assert!(reader.get_bits_64(4) == 10);
            assert!(reader.get_bits_64(1) == 0);
            assert!(reader.get_bits_64(8) == 85);
            assert_eq!(reader.get_bits_64(32), 1431655765);
        }

        #[test]
        fn peek_bits_64() {
            let reader = BitReadLE {
                buffer: &CHECKBOARD0101,
                index: 0,
                cache: 0,
                left: 0
            };

            assert!(reader.peek_bits_64(1) == 1);
            assert!(reader.peek_bits_64(1) == 1);
            assert!(reader.peek_bits_64(2) == 1);
            assert!(reader.peek_bits_64(2) == 1);
        }

        #[test]
        fn get_bits_32() {
            let mut reader = BitReadLE {
                buffer: &CHECKBOARD0101,
                index: 0,
                cache: 0,
                left: 0
            };

            assert!(reader.get_bits_32(1) == 1);
            assert!(reader.get_bits_32(2) == 2);
            assert!(reader.get_bits_32(4) == 10);
            assert!(reader.get_bits_32(1) == 0);
            assert!(reader.get_bits_32(8) == 85);

            assert_eq!(reader.get_bits_32(1), 1);
            for _ in 0..8 {
                assert_eq!(reader.get_bits_32(8), 170);
            }
        }

        #[test]
        fn peek_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.peek_bits_32(1) == 1);
            assert!(reader.peek_bits_32(1) == 1);
            assert!(reader.peek_bits_32(2) == 1);
            assert!(reader.peek_bits_32(2) == 1);
        }

        #[test]
        fn skip_bits() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            reader.skip_bits(0);
            assert!(reader.peek_bits_32(1) == 1);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 1);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 1);
        }

        #[test]
        fn align_bits() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadLE::new(b);

            reader.align_bits();
            assert!(reader.get_bits_64(3) == 3);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
            reader.skip_bits(1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
        }

        #[test]
        fn overread() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadLE::new(b);

            reader.skip_bits(128*8+2);
            reader.get_bits_64(6);
        }
    }
    mod be {
        use super::super::*;
        use super::*;

        #[test]
        fn get_bit() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(!reader.get_bit());
            assert!(reader.get_bit());
        }

        #[test]
        fn get_bits_64() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.get_bits_64(1) == 0);
            assert!(reader.get_bits_64(2) == 2);
            assert!(reader.get_bits_64(4) == 10);
            assert!(reader.get_bits_64(1) == 1);
            assert!(reader.get_bits_64(8) == 85);
            assert_eq!(reader.get_bits_64(32), 1431655765);
        }

        #[test]
        fn peek_bits_64() {
            let b = &CHECKBOARD0101;
            let reader = BitReadBE::new(b);

            assert!(reader.peek_bits_64(1) == 0);
            assert!(reader.peek_bits_64(1) == 0);
            assert!(reader.peek_bits_64(2) == 1);
            assert!(reader.peek_bits_64(2) == 1);
        }

        #[test]
        fn get_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.get_bits_32(1) == 0);
            assert!(reader.get_bits_32(2) == 2);
            assert!(reader.get_bits_32(4) == 10);
            assert!(reader.get_bits_32(1) == 1);
            assert!(reader.get_bits_32(8) == 85);

            assert_eq!(reader.get_bits_32(1), 0);
            for _ in 0..8 {
                assert_eq!(reader.get_bits_32(8), 170);
            }
        }

        #[test]
        fn peek_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.peek_bits_32(1) == 0);
            assert!(reader.peek_bits_32(1) == 0);
            assert!(reader.peek_bits_32(2) == 1);
            assert!(reader.peek_bits_32(2) == 1);
        }

        #[test]
        fn skip_bits() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            reader.skip_bits(0);
            assert!(reader.peek_bits_32(1) == 0);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 0);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 0);
        }

        #[test]
        fn align_bits() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadBE::new(b);

            reader.align_bits();
            assert!(reader.get_bits_64(3) == 1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
            reader.skip_bits(1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
        }

        #[test]
        fn overread() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadBE::new(b);

            reader.skip_bits(128*8+2);
            reader.get_bits_64(6);
        }
    }
}
<|endoftext|>


<commit_before>use super::*;
use super::super::packets::{self, Encodable, Decodable};
use super::super::status_translator;
use crc::crc32;
use byteorder::{self, BigEndian, ByteOrder};
use std::collections;
use std::net;
use std::thread;
use std::io;
use std::sync::mpsc;
use mio;
use mio::udp;

const SOCKET_TOKEN: mio::Token = mio::Token(0);
const TIMEOUT1000: u32 = 1; //happens every second.
const TIMEOUT200: u32 = 2; //Happens every 200 MS.

pub enum MioHandlerCommand {
    DoCall(Box<fn(&mut MioHandler)>),
}

/*This exists so that we can lend these buffers around too, not just the socket.
If we implement just for the socket, then everything has to figure out its own buffer.*/
pub struct MioSocketState<'a> {
    socket: &'a udp::UdpSocket,
    incoming_packet_buffer: [u8; 1000],
    outgoing_packet_buffer: [u8; 1000],
}

pub struct MioHandler<'a> {
    socket_state: MioSocketState<'a>,
    connections: collections::HashMap<net::SocketAddr, Connection>,
    next_connection_id: u64,
    status_translator: status_translator::StatusTranslator,
}

impl<'a> MioHandler<'a> {
    pub fn new(socket: &'a udp::UdpSocket)->MioHandler<'a> {
        MioHandler {
            socket_state: MioSocketState {
                socket: socket,
                incoming_packet_buffer: [0u8; 1000],
                outgoing_packet_buffer: [0u8; 1000],
            },
            connections: collections::HashMap::new(),
            next_connection_id: 1,
            status_translator: status_translator::StatusTranslator::new(true, packets::PROTOCOL_VERSION, &[""; 0]),
        }
    }

    fn got_packet(&mut self, size: usize, address: net::SocketAddr) {
        if size == 0 {return;}
        let maybe_packet = {
            let slice = &self.socket_state.incoming_packet_buffer[0..size];
            let computed_checksum = crc32::checksum_castagnoli(&slice[4..]);
            let expected_checksum = BigEndian::read_u32(&slice[..4]);
            if computed_checksum != expected_checksum {Err(packets::PacketDecodingError::Invalid)}
            else {packets::decode_packet(&slice[4..])}
        };
        if let Err(_) = maybe_packet {return;}
        let packet = maybe_packet.unwrap();
        if let Some(ref mut conn) = self.connections.get_mut(&address) {
            if conn.handle_incoming_packet(&packet, &mut self.socket_state) {return;}
        }
        match packet {
            packets ::Packet::Connect => {
                if let Some(_) = self.connections.get(&address) {return;}
                let id = self.next_connection_id;
                self.next_connection_id += 1;
                let conn = Connection::new(address, id);
                self.connections.insert(address, conn);
                self.socket_state.send(&packets::Packet::Connected(id), address);
            },
            packets::Packet::StatusRequest(ref req) => {
                self.socket_state.send(&packets::Packet::StatusResponse(self.status_translator.translate(req)), address);
            },
            _ => {}
        }
    }
}

impl<'a> PacketSender for MioSocketState<'a> {
    fn send(&mut self, packet: &packets::Packet, address: net::SocketAddr)->bool {
        if let Ok(size) = packets::encode_packet(packet, &mut self.outgoing_packet_buffer[4..]) {
            let checksum = crc32::checksum_castagnoli(&self.outgoing_packet_buffer[4..size]);
            BigEndian::write_u32(&mut self.outgoing_packet_buffer[..4], checksum);
            if let Ok(Some(sent_bytes)) = self.socket.send_to(&self.outgoing_packet_buffer[..4+size], &address) {
                if sent_bytes == 4+size {return true;}
                else {return false;}
            }
            else {return false;}
        }
        else {return false;};
    }
}

impl<'a> mio::Handler for MioHandler<'a> {
    type Timeout = u32;
    type Message = MioHandlerCommand;

    fn ready(&mut self, event_loop: &mut mio::EventLoop<Self>, token: mio::Token, events: mio::EventSet) {
        //We only have one socket, so can avoid the match on the token.
        if events.is_error() {
            //We need to do something sensible here, probably a callback with whatever state we can get.
        }
        if events.is_readable() {
            let result = self.socket_state.socket.recv_from(&mut self.socket_state.incoming_packet_buffer);
            if let Ok(Some((size, address))) = result {
                self.got_packet(size, address);
            }
        }
    }

    fn timeout(&mut self, event_loop: &mut mio::EventLoop<Self>, timeout: Self::Timeout) {
        //Rust isn't smart enough to realize that closures only borrow a field, so we pull it out here to satisfy the borrow checker.
        let sender = &mut self.socket_state;
        let rereg = match timeout {
            TIMEOUT200 => {
                self.connections.iter_mut().map(|x| x.1.tick200(sender));
                200
            },
            TIMEOUT1000 => {
                self.connections.iter_mut().map(|x| x.1.tick1000(sender));
                1000
            },
            _ => panic!("Unrecognized timeout. This should never happen.")
        };
        event_loop.timeout_ms(timeout, rereg).unwrap();
    }
}

fn mio_server_thread(address: net::SocketAddr, notify_created: mpsc::Sender<Result<mio::Sender<MioHandlerCommand>, io::Error>>) {
    let maybe_socket = match address {
        net::SocketAddr::V4(_) => udp::UdpSocket::v4(),
        net::SocketAddr::V6(_) => udp::UdpSocket::v6()
    };
    if let Err(what) = maybe_socket {
        notify_created.send(Err(what)).unwrap();
        return;
    }
    let socket = maybe_socket.unwrap();
    if let  Err(what) = socket.bind(&address) {
        notify_created.send(Err(what)).unwrap();
        return;
    }
    let maybe_loop  = mio::EventLoop::new();
    if let Err(what) = maybe_loop {
        notify_created.send(Err(what)).unwrap();
        return;
    }
    let mut event_loop = maybe_loop.unwrap();
    let mut handler = MioHandler::new(&socket);
    if let Err(what)  = event_loop.register(&socket, SOCKET_TOKEN, mio::EventSet::all(), mio::PollOpt::all()) {
        notify_created.send(Err(what)).unwrap();
        return;
    }
    let timer_error = Err(io::Error::new(io::ErrorKind::Other, "Couldn't create the timer."));
    if let Err(_) = event_loop.timeout_ms(TIMEOUT1000, 1000) {
        notify_created.send(timer_error).unwrap();
        return;
    }
    if let Err(_) = event_loop.timeout_ms(TIMEOUT200, 200) {
        notify_created.send(timer_error);
        return;
    }
    let sender = event_loop.channel();
    notify_created.send(Ok(sender));
    event_loop.run(&mut handler);
}

pub struct MioServer {
    thread: thread::JoinHandle<()>,
    sender: mio::Sender<MioHandlerCommand>,
}

impl MioServer {
    fn new(address: net::SocketAddr)->Result<MioServer, io::Error> {
        let (sender, receiver) = mpsc::channel();
        let join_handle = thread::spawn(move || mio_server_thread(address, sender));
        let message_sender = try!(receiver.recv().unwrap());
        Ok(MioServer {
            thread: join_handle,
            sender: message_sender,
        })
    }
}
<commit_msg>Use an enum for the types of timeout, thus avoiding the need for a panic.<commit_after>use super::*;
use super::super::packets::{self, Encodable, Decodable};
use super::super::status_translator;
use crc::crc32;
use byteorder::{self, BigEndian, ByteOrder};
use std::collections;
use std::net;
use std::thread;
use std::io;
use std::sync::mpsc;
use mio;
use mio::udp;

const SOCKET_TOKEN: mio::Token = mio::Token(0);

#[derive(Debug, Copy, Clone)]
pub enum TimeoutTypes {
    Timeout1000,
    Timeout200,
}

pub enum MioHandlerCommand {
    DoCall(Box<fn(&mut MioHandler)>),
}

/*This exists so that we can lend these buffers around too, not just the socket.
If we implement just for the socket, then everything has to figure out its own buffer.*/
pub struct MioSocketState<'a> {
    socket: &'a udp::UdpSocket,
    incoming_packet_buffer: [u8; 1000],
    outgoing_packet_buffer: [u8; 1000],
}

pub struct MioHandler<'a> {
    socket_state: MioSocketState<'a>,
    connections: collections::HashMap<net::SocketAddr, Connection>,
    next_connection_id: u64,
    status_translator: status_translator::StatusTranslator,
}

impl<'a> MioHandler<'a> {
    pub fn new(socket: &'a udp::UdpSocket)->MioHandler<'a> {
        MioHandler {
            socket_state: MioSocketState {
                socket: socket,
                incoming_packet_buffer: [0u8; 1000],
                outgoing_packet_buffer: [0u8; 1000],
            },
            connections: collections::HashMap::new(),
            next_connection_id: 1,
            status_translator: status_translator::StatusTranslator::new(true, packets::PROTOCOL_VERSION, &[""; 0]),
        }
    }

    fn got_packet(&mut self, size: usize, address: net::SocketAddr) {
        if size == 0 {return;}
        let maybe_packet = {
            let slice = &self.socket_state.incoming_packet_buffer[0..size];
            let computed_checksum = crc32::checksum_castagnoli(&slice[4..]);
            let expected_checksum = BigEndian::read_u32(&slice[..4]);
            if computed_checksum != expected_checksum {Err(packets::PacketDecodingError::Invalid)}
            else {packets::decode_packet(&slice[4..])}
        };
        if let Err(_) = maybe_packet {return;}
        let packet = maybe_packet.unwrap();
        if let Some(ref mut conn) = self.connections.get_mut(&address) {
            if conn.handle_incoming_packet(&packet, &mut self.socket_state) {return;}
        }
        match packet {
            packets ::Packet::Connect => {
                if let Some(_) = self.connections.get(&address) {return;}
                let id = self.next_connection_id;
                self.next_connection_id += 1;
                let conn = Connection::new(address, id);
                self.connections.insert(address, conn);
                self.socket_state.send(&packets::Packet::Connected(id), address);
            },
            packets::Packet::StatusRequest(ref req) => {
                self.socket_state.send(&packets::Packet::StatusResponse(self.status_translator.translate(req)), address);
            },
            _ => {}
        }
    }
}

impl<'a> PacketSender for MioSocketState<'a> {
    fn send(&mut self, packet: &packets::Packet, address: net::SocketAddr)->bool {
        if let Ok(size) = packets::encode_packet(packet, &mut self.outgoing_packet_buffer[4..]) {
            let checksum = crc32::checksum_castagnoli(&self.outgoing_packet_buffer[4..size]);
            BigEndian::write_u32(&mut self.outgoing_packet_buffer[..4], checksum);
            if let Ok(Some(sent_bytes)) = self.socket.send_to(&self.outgoing_packet_buffer[..4+size], &address) {
                if sent_bytes == 4+size {return true;}
                else {return false;}
            }
            else {return false;}
        }
        else {return false;};
    }
}

impl<'a> mio::Handler for MioHandler<'a> {
    type Timeout = TimeoutTypes;
    type Message = MioHandlerCommand;

    fn ready(&mut self, event_loop: &mut mio::EventLoop<Self>, token: mio::Token, events: mio::EventSet) {
        //We only have one socket, so can avoid the match on the token.
        if events.is_error() {
            //We need to do something sensible here, probably a callback with whatever state we can get.
        }
        if events.is_readable() {
            let result = self.socket_state.socket.recv_from(&mut self.socket_state.incoming_packet_buffer);
            if let Ok(Some((size, address))) = result {
                self.got_packet(size, address);
            }
        }
    }

    fn timeout(&mut self, event_loop: &mut mio::EventLoop<Self>, timeout: Self::Timeout) {
        //Rust isn't smart enough to realize that closures only borrow a field, so we pull it out here to satisfy the borrow checker.
        let sender = &mut self.socket_state;
        let rereg = match timeout {
            TimeoutTypes::Timeout200 => {
                self.connections.iter_mut().map(|x| x.1.tick200(sender));
                200
            },
            TimeoutTypes::Timeout1000 => {
                self.connections.iter_mut().map(|x| x.1.tick1000(sender));
                1000
            },
        };
        event_loop.timeout_ms(timeout, rereg).unwrap();
    }
}

fn mio_server_thread(address: net::SocketAddr, notify_created: mpsc::Sender<Result<mio::Sender<MioHandlerCommand>, io::Error>>) {
    let maybe_socket = match address {
        net::SocketAddr::V4(_) => udp::UdpSocket::v4(),
        net::SocketAddr::V6(_) => udp::UdpSocket::v6()
    };
    if let Err(what) = maybe_socket {
        notify_created.send(Err(what)).unwrap();
        return;
    }
    let socket = maybe_socket.unwrap();
    if let  Err(what) = socket.bind(&address) {
        notify_created.send(Err(what)).unwrap();
        return;
    }
    let maybe_loop  = mio::EventLoop::new();
    if let Err(what) = maybe_loop {
        notify_created.send(Err(what)).unwrap();
        return;
    }
    let mut event_loop = maybe_loop.unwrap();
    let mut handler = MioHandler::new(&socket);
    if let Err(what)  = event_loop.register(&socket, SOCKET_TOKEN, mio::EventSet::all(), mio::PollOpt::all()) {
        notify_created.send(Err(what)).unwrap();
        return;
    }
    let timer_error = Err(io::Error::new(io::ErrorKind::Other, "Couldn't create the timer."));
    if let Err(_) = event_loop.timeout_ms(TimeoutTypes::Timeout1000, 1000) {
        notify_created.send(timer_error).unwrap();
        return;
    }
    if let Err(_) = event_loop.timeout_ms(TimeoutTypes::Timeout200, 200) {
        notify_created.send(timer_error);
        return;
    }
    let sender = event_loop.channel();
    notify_created.send(Ok(sender));
    event_loop.run(&mut handler);
}

pub struct MioServer {
    thread: thread::JoinHandle<()>,
    sender: mio::Sender<MioHandlerCommand>,
}

impl MioServer {
    fn new(address: net::SocketAddr)->Result<MioServer, io::Error> {
        let (sender, receiver) = mpsc::channel();
        let join_handle = thread::spawn(move || mio_server_thread(address, sender));
        let message_sender = try!(receiver.recv().unwrap());
        Ok(MioServer {
            thread: join_handle,
            sender: message_sender,
        })
    }
}
<|endoftext|>


<commit_before>//! A chat server that broadcasts a message to all connections.

extern crate tokio_core;
extern crate futures;

use std::collections::HashMap;
use std::rc::Rc;
use std::cell::RefCell;
use std::iter;
use std::env;
use std::io::BufReader;

use tokio_core::net::TcpListener;
use tokio_core::reactor::Core;
use tokio_core::io::{self, Io};

use futures::stream::{self, Stream};
use futures::Future;

fn main() {
    let addr = env::args().nth(1).unwrap_or("127.0.0.1:8080".to_string());
    let addr = addr.parse().unwrap();
    // We are single-threaded, so we can just use Rc and RefCell.
    let connections = Rc::new(RefCell::new(HashMap::new()));

    let mut core = Core::new().unwrap();
    let handle = core.handle();
    let socket = TcpListener::bind(&addr, &handle).unwrap();
    println!("Listening on: {}", addr);

    let connections = connections.clone();
    let future = socket.incoming().for_each(move |(stream, addr)| {
        let connections = connections.clone();
        let handle_inner = handle.clone();
        // We create a new future in which we create all other futures.
        // This makes `stream` be bound on the `lazy` future's task, allowing
        // `ReadHalf` and `WriteHalf` to be shared between inner futures.
        handle.spawn_fn(move || {
            println!("New Connection: {}", addr);
            let (reader, writer) = stream.split();
            // channel to send messages to this connection from other futures
            let (tx, rx) = tokio_core::channel::channel(&handle_inner).unwrap();
            // add sender to hashmap of all current connections
            connections.borrow_mut().insert(addr, tx);

            let connections_inner = connections.clone();
            // https://users.rust-lang.org/t/loop-futures-for-client-handling/6950/2
            // We have an endless loop reading from a client.
            // In order to fuse the reading and writing futures in the end, we need to have the same
            // output type. Therefore we use `(Option<BufReader<ReadHalf<TcpStream>>>,
            // Option<WriteHalf<TcpStream>>)`.
            let reader = BufReader::new(reader);
            let socket_reader = stream::iter::<_, _, std::io::Error>(iter::repeat(()).map(Ok)).fold((Some(reader),None), move |(reader, _), _| {
                let reader = reader.unwrap();
                let connections = connections_inner.clone();
                // read and parse length prefix
                io::read_until(reader, '\n' as u8, vec![])
                    .and_then(|(reader, vec)| futures::lazy(|| {
                        // EOF was hit without reading a delimiter
                        if vec.len() == 0 {
                            futures::failed((std::io::Error::new(std::io::ErrorKind::BrokenPipe, "Broken Pipe"))).boxed()
                        } else {
                            futures::finished((reader, vec)).boxed()
                        }
                    }))
                    // convert bytes into string
                    .map(|(reader, vec)| (reader, String::from_utf8(vec).unwrap()))
                    .and_then(move |(reader, message)| {
                        println!("{}: {:?}", addr, message);
                        // For each open connection except the sender, send the string via the channel
                        for tx in connections.borrow_mut().iter().filter(|&(&k,_)| k != addr).map(|(_,v)| v) {
                            tx.send(message.clone()).unwrap();
                        }
                        futures::finished((Some(reader),None))
                    })
            });

            // Whenever we receive a string on the Receiver, we write it to `WriteHalf<TcpStream>`.
            let socket_writer = rx.fold((None, Some(writer)), move |(_, writer), msg| {
                let writer = writer.unwrap();
                io::write_all(writer, msg.into_bytes()).map(|(writer, _)| (None, Some(writer))).boxed()
            });

            socket_reader.select(socket_writer)
                .then(move |_| {
                    connections.borrow_mut().remove(&addr);
                    println!("Connection {:?} closed.", addr);
                    Ok(())
                })
        });
        Ok(())
    });

    // exectue server
    core.run(future).unwrap();
}
<commit_msg>fix(chat): Implement alexcrichton's suggestions<commit_after>//! A chat server that broadcasts a message to all connections.

extern crate tokio_core;
extern crate futures;

use std::collections::HashMap;
use std::rc::Rc;
use std::cell::RefCell;
use std::iter;
use std::env;
use std::io::{Error, ErrorKind, BufReader};

use tokio_core::net::TcpListener;
use tokio_core::reactor::Core;
use tokio_core::io::{self, Io};

use futures::stream::{self, Stream};
use futures::Future;

fn main() {
    let addr = env::args().nth(1).unwrap_or("127.0.0.1:8080".to_string());
    let addr = addr.parse().unwrap();
    // We are single-threaded, so we can just use Rc and RefCell.
    let connections = Rc::new(RefCell::new(HashMap::new()));

    let mut core = Core::new().unwrap();
    let handle = core.handle();
    let socket = TcpListener::bind(&addr, &handle).unwrap();
    println!("Listening on: {}", addr);

    let future = socket.incoming().for_each(move |(stream, addr)| {
        let connections = connections.clone();
        let handle_inner = handle.clone();
        // We create a new future in which we create all other futures.
        // This makes `stream` be bound on the outer future's task, allowing
        // `ReadHalf` and `WriteHalf` to be shared between inner futures.
        handle.spawn_fn(move || {
            println!("New Connection: {}", addr);
            let (reader, writer) = stream.split();
            // channel to send messages to this connection from other futures
            let (tx, rx) = tokio_core::channel::channel(&handle_inner).unwrap();
            // add sender to hashmap of all current connections
            connections.borrow_mut().insert(addr, tx);

            let reader = BufReader::new(reader);
            let connections_inner = connections.clone();
            // https://users.rust-lang.org/t/loop-futures-for-client-handling/6950/2
            // First we need to get an infinite iterator
            let iter = stream::iter::<_, _, std::io::Error>(iter::repeat(()).map(Ok));
            // Then we fold it as infinite loop
            let socket_reader = iter.fold(reader, move |reader, _| {
                let connections = connections_inner.clone();
                // read line
                let amt = io::read_until(reader, '\n' as u8, vec![]);
                // check if we hit EOF and need to close the connection
                let amt = amt.and_then(|(reader, vec)| {
                    // EOF was hit without reading a delimiter
                    if vec.len() == 0 {
                        let err = Error::new(ErrorKind::BrokenPipe, "Broken Pipe");
                        futures::failed(err).boxed()
                    } else {
                        futures::finished((reader, vec)).boxed()
                    }
                });
                // convert bytes into string
                let amt = amt.map(|(reader, vec)| (reader, String::from_utf8(vec)));
                amt.and_then(move |(reader, message)| {
                    println!("{}: {:?}", addr, message);
                    let conns = connections.borrow_mut();
                    if let Ok(msg) = message {
                        // For each open connection except the sender, send the string
                        // via the channel
                        let iter = conns.iter().filter(|&(&k,_)| k != addr).map(|(_,v)| v);
                        for tx in iter {
                            tx.send(msg.clone()).unwrap();
                        }
                    } else {
                        let tx = conns.get(&addr).unwrap();
                        tx.send("You didn't send valid UTF-8.".to_string()).unwrap();
                    }
                    futures::finished(reader)
                })
            });

            // Whenever we receive a string on the Receiver, we write it to `WriteHalf<TcpStream>`.
            let socket_writer = rx.fold(writer, |writer, msg| {
                let amt = io::write_all(writer, msg.into_bytes());
                let amt = amt.map(|(writer, _)| writer);
                amt
            });

            // In order to fuse the reading and writing futures in the end, we need to have the
            // same output type. Therefore we use `(Option<BufReader<ReadHalf<TcpStream>>>,
            // Option<WriteHalf<TcpStream>>)`.
            let socket_reader = socket_reader.map(|reader| (Some(reader), None));
            let socket_writer = socket_writer.map(|writer| (None, Some(writer)));

            let amt = socket_reader.select(socket_writer);
            amt.then(move |_| {
                connections.borrow_mut().remove(&addr);
                println!("Connection {:?} closed.", addr);
                Ok(())
            })
        });
        Ok(())
    });

    // exectue server
    core.run(future).unwrap();
}

<|endoftext|>