<commit_before>//! A fast general-purpose object arena.
//!
//! This crate contains `VecArena`, which is an allocator that can hold objects of only one type.
//! It can allocate space for new objects and reclaim space upon removal of objects. The amount
//! of space to hold objects is dynamically expanded as needed. Inserting objects into an arena
//! is amortized `O(1)`, and removing is `O(1)`.
// TODO: An example of a doubly linked list?

use bitmap::Bitmap;
use std::marker::PhantomData;
use std::mem;
use std::ops::{Index, IndexMut};
use std::ptr;

mod bitmap;

// TODO: fn insert_near / nearest_vacant
// TODO: DoubleEndedIterator

/// An arena that can insert and remove objects of a single type.
///
/// `VecArena<T>` is in many ways like `Vec<Option<T>>` because it holds an array of slots for
/// storing objects. A slot can be either occupied or vacant. Inserting a new object into an arena
/// involves finding a vacant slot and placing the object into the slot. Removing an object means
/// taking it out of the slot and marking it as vacant.
///
/// Internally, a bitmap is used instead of `Option`s to conserve space and improve cache
/// performance. Every object access makes sure that the accessed object really exists, otherwise
/// it panics.
// TODO: a bunch of examples, see the docs for Vec for inspiration.
pub struct VecArena<T> {
    // TODO: Docs for these fields
    slots: *const T,
    bitmap: Bitmap,
    marker: PhantomData<T>,
}

impl<T> VecArena<T> {
    /// Constructs a new, empty arena.
    ///
    /// The arena will not allocate until objects are inserted into it.
    pub fn new() -> Self {
        let slots = {
            let mut vec = Vec::with_capacity(0);
            let ptr = vec.as_mut_ptr();
            mem::forget(vec);
            ptr
        };
        VecArena {
            slots: slots,
            bitmap: Bitmap::new(),
            marker: PhantomData,
        }
    }

    /// Constructs a new, empty arena with the specified capacity (number of slots).
    ///
    /// The arena will be able to hold exactly `capacity` objects without reallocating.
    /// If `capacity` is 0, the arena will not allocate.
    pub fn with_capacity(cap: usize) -> Self {
        let mut arena = Self::new();
        arena.reserve_exact(cap);
        arena
    }

    /// Returns the number of objects the arena can hold without reallocating.
    /// In other words, this is the number of slots.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.bitmap.len()
    }

    /// Returns the number of objects in the arena.
    /// In other words, this is the number of occupied slots.
    #[inline]
    pub fn len(&self) -> usize {
        self.bitmap.occupied()
    }

    /// Returns `true` if the arena holds no objects.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Returns `true` if the slot at `index` is vacant.
    #[inline]
    pub fn is_vacant(&self, index: usize) -> bool {
        !self.bitmap.is_occupied(index)
    }

    /// Returns `true` if the slot at `index` is occupied.
    #[inline]
    pub fn is_occupied(&self, index: usize) -> bool {
        self.bitmap.is_occupied(index)
    }

    /// Inserts an object into the arena and returns the slot index in which it was stored.
    /// The arena will reallocate if it's full.
    pub fn insert(&mut self, object: T) -> usize {
        if self.bitmap.occupied() == self.bitmap.len() {
            self.double();
        }

        let index = self.bitmap.acquire();
        unsafe {
            ptr::write(self.slots.offset(index as isize) as *mut T, object);
        }
        index
    }

    /// Inserts an object into the vacant slot at `index`.
    ///
    /// # Panics
    ///
    /// Panics if `index` is out of bounds or the slot is already occupied.
    pub fn insert_at(&mut self, index: usize, object: T) -> usize {
        self.bitmap.acquire_at(index);
        unsafe {
            ptr::write(self.slots.offset(index as isize) as *mut T, object);
        }
        index
    }

    /// Removes the object stored at `index` from the arena and returns it.
    ///
    /// # Panics
    ///
    /// Panics if `index` is out of bounds or the slot is already vacant.
    pub fn remove(&mut self, index: usize) -> T {
        // `release` will panic if the index is out of bounds or the slot is already vacant.
        self.bitmap.release(index);

        unsafe {
            ptr::read(self.slots.offset(index as isize) as *mut T)
        }
    }

    /// Clears the arena, removing and dropping all objects it holds. Keeps the allocated memory
    /// for reuse.
    pub fn clear(&mut self) {
        let mut arena = VecArena::with_capacity(self.capacity());
        mem::swap(self, &mut arena);
    }

    /// Returns a reference to the object stored at `index`.
    ///
    /// # Panics
    ///
    /// Panics if `index` is out of bounds.
    #[inline]
    pub fn get(&self, index: usize) -> Option<&T> {
        self.guard_index(index);
        if self.is_occupied(index) {
            unsafe {
                Some(self.get_unchecked(index))
            }
        } else {
            None
        }
    }

    /// Returns a mutable reference to the object stored at `index`.
    ///
    /// # Panics
    ///
    /// Panics if `index` is out of bounds.
    #[inline]
    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
        self.guard_index(index);
        if self.is_occupied(index) {
            unsafe {
                Some(self.get_unchecked_mut(index))
            }
        } else {
            None
        }
    }

    /// Reserves capacity for at least `additional` more elements to be inserted. The arena may
    /// reserve more space to avoid frequent reallocations.
    ///
    /// # Panics
    ///
    /// Panics if the new capacity overflows `usize`.
    pub fn reserve(&mut self, additional: usize) {
        let len = self.bitmap.len();
        self.bitmap.reserve(additional);
        unsafe {
            self.reallocate(len);
        }
    }

    /// Reserves the minimum capacity for exactly `additional` more elements to be inserted.
    ///
    /// Note that the allocator may give the arena more space than it requests.
    ///
    /// # Panics
    ///
    /// Panics if the new capacity overflows `usize`.
    pub fn reserve_exact(&mut self, additional: usize) {
        let len = self.bitmap.len();
        self.bitmap.reserve_exact(additional);
        unsafe {
            self.reallocate(len);
        }
    }

    /// Reallocates the object array because the bitmap was resized.
    unsafe fn reallocate(&mut self, old_len: usize) {
        let new_len = self.bitmap.len();

        // Allocate a new array.
        let mut vec = Vec::with_capacity(new_len);
        let ptr = vec.as_mut_ptr();
        mem::forget(vec);

        // Copy data into the new array.
        ptr::copy_nonoverlapping(self.slots, ptr, old_len);

        // Deallocate the old array.
        Vec::from_raw_parts(self.slots as *mut T, 0, old_len);

        self.slots = ptr;
    }

    /// Doubles the capacity of the arena.
    #[cold]
    #[inline(never)]
    fn double(&mut self) {
        let len = self.bitmap.len();
        let elem_size = mem::size_of::<T>();

        let new_len = if len == 0 {
            if elem_size.checked_mul(4).is_some() {
                4
            } else {
                1
            }
        } else {
            len.checked_mul(2).expect("len overflow")
        };

        self.reserve_exact(new_len - len);
    }

    #[inline]
    pub fn iter(&self) -> Iter<T> {
        Iter {
            arena: self,
            index: 0,
        }
    }

    #[inline]
    pub fn iter_mut(&mut self) -> IterMut<T> {
        IterMut {
            slots: self.slots as *mut T,
            bitmap: &self.bitmap,
            index: 0,
            marker: PhantomData,
        }
    }

    #[inline]
    pub fn drain(&mut self) -> Drain<T> {
        Drain {
            arena: self,
            index: 0,
        }
    }

    /// Returns a reference to the object at `index`, without bounds checking nor checking that
    /// the object exists.
    #[inline]
    pub unsafe fn get_unchecked(&self, index: usize) -> &T {
        &*self.slots.offset(index as isize)
    }

    /// Returns a mutable reference to the object at `index`, without bounds checking nor checking
    /// that the object exists.
    #[inline]
    pub unsafe fn get_unchecked_mut(&mut self, index: usize) -> &mut T {
        &mut *(self.slots.offset(index as isize) as *mut T)
    }

    /// Panics if `index` is out of bounds.
    #[inline]
    fn guard_index(&self, index: usize) {
        assert!(index < self.bitmap.len(), "`index` out of bounds");
    }

    /// Panics if `index` is out of bounds or the slot is vacant.
    #[inline]
    fn guard_occupied(&self, index: usize) {
        // `is_occupied` will panic if the index is out of bounds.
        assert!(self.bitmap.is_occupied(index), "vacant slot at `index`");
    }
}

impl<T> Drop for VecArena<T> {
    fn drop(&mut self) {
        unsafe {
            // Drop all objects in the arena.
            for index in self.bitmap.iter() {
                ptr::drop_in_place(self.slots.offset(index as isize) as *mut T);
            }

            // Deallocate the old array.
            Vec::from_raw_parts(self.slots as *mut T, 0, self.bitmap.len());
        }
    }
}

impl<T> Index<usize> for VecArena<T> {
    type Output = T;

    #[inline]
    fn index(&self, index: usize) -> &T {
        self.guard_occupied(index);
        unsafe {
            self.get_unchecked(index)
        }
    }
}

impl<T> IndexMut<usize> for VecArena<T> {
    #[inline]
    fn index_mut(&mut self, index: usize) -> &mut T {
        self.guard_occupied(index);
        unsafe {
            self.get_unchecked_mut(index)
        }
    }
}

impl<T> Default for VecArena<T> {
    fn default() -> Self {
        VecArena::new()
    }
}

impl<T: Clone> Clone for VecArena<T> {
    fn clone(&self) -> Self {
        let mut arena = VecArena::with_capacity(self.capacity());
        for index in self.bitmap.iter() {
            let clone = unsafe {
                self.get_unchecked(index).clone()
            };
            arena.insert_at(index, clone);
        }
        arena
    }
}

pub struct IntoIter<T> {
    slots: *const T,
    bitmap: Bitmap,
    index: usize,
    marker: PhantomData<T>,
}

impl<T> Drop for IntoIter<T> {
    fn drop(&mut self) {
        while self.next().is_some() {}
    }
}

impl<T> Iterator for IntoIter<T> {
    type Item = (usize, T);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.bitmap.next_occupied(self.index).map(|index| {
            self.index = index + 1;
            unsafe {
                self.bitmap.release(index);
                (index, ptr::read(self.slots.offset(index as isize)))
            }
        })
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let lower = 0;
        let upper = self.bitmap.len() - self.index;
        (lower, Some(upper))
    }
}

impl<T> IntoIterator for VecArena<T> {
    type Item = (usize, T);
    type IntoIter = IntoIter<T>;

    fn into_iter(mut self) -> Self::IntoIter {
        let iter = IntoIter {
            slots: self.slots,
            bitmap: mem::replace(&mut self.bitmap, Bitmap::new()),
            index: 0,
            marker: PhantomData,
        };
        mem::forget(self);
        iter
    }
}

pub struct Iter<'a, T: 'a> {
    arena: &'a VecArena<T>,
    index: usize,
}

impl<'a, T> Iterator for Iter<'a, T> {
    type Item = (usize, &'a T);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.arena.bitmap.next_occupied(self.index).map(|index| {
            self.index = index + 1;
            unsafe {
                (index, self.arena.get_unchecked(index))
            }
        })
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let lower = 0;
        let upper = self.arena.capacity() - self.index;
        (lower, Some(upper))
    }
}

pub struct IterMut<'a, T: 'a> {
    slots: *mut T,
    bitmap: &'a Bitmap,
    index: usize,
    marker: PhantomData<&'a mut T>,
}

impl<'a, T> Iterator for IterMut<'a, T> {
    type Item = (usize, &'a mut T);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.bitmap.next_occupied(self.index).map(|index| {
            self.index = index + 1;
            unsafe {
                (index, &mut *self.slots.offset(index as isize))
            }
        })
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let lower = 0;
        let upper = self.bitmap.len() - self.index;
        (lower, Some(upper))
    }
}

pub struct Drain<'a, T: 'a> {
    arena: &'a mut VecArena<T>,
    index: usize,
}

impl<'a, T> Iterator for Drain<'a, T> {
    type Item = (usize, T);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.arena.bitmap.next_occupied(self.index).map(|index| {
            self.index = index + 1;
            unsafe {
                self.arena.bitmap.release(index);
                (index, ptr::read(self.arena.slots.offset(index as isize)))
            }
        })
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let lower = 0;
        let upper = self.arena.capacity() - self.index;
        (lower, Some(upper))
    }
}

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

    // TODO: test variance: https://github.com/rust-lang/rust/pull/30998/files

    #[test]
    fn it_works() {
        let mut arena = VecArena::new();
        for i in 0..10 {
            assert_eq!(arena.insert(()), i);
            assert!(arena[i] == ());
        }
        for i in 0..10 {
            assert!(arena[i] == ());
        }
    }
}
<commit_msg>Clean up the public API<commit_after>//! A fast general-purpose object arena.
//!
//! This crate contains `VecArena`, which is an allocator that can hold objects of only one type.
//! It can allocate space for new objects and reclaim space upon removal of objects. The amount
//! of space to hold objects is dynamically expanded as needed. Inserting objects into an arena
//! is amortized `O(1)`, and removing is `O(1)`.
// TODO: An example of a doubly linked list?

use bitmap::Bitmap;
use std::marker::PhantomData;
use std::mem;
use std::ops::{Index, IndexMut};
use std::ptr;

mod bitmap;

/// An arena that can insert and remove objects of a single type.
///
/// `VecArena<T>` resembles `Vec<Option<T>>` a lot because it holds an array of slots for
/// storing objects. A slot can be either occupied or vacant. Inserting a new object into an arena
/// involves finding a vacant slot and placing the object into the slot. Removing an object means
/// taking it out of the slot and marking it as vacant.
///
/// Internally, a bitmap is used instead of `Option`s to conserve space and improve cache
/// performance. Every indexing operation checks that an object really exists at the specified
/// index, otherwise it panics.
// TODO: Explain that arena[index] is equivalent to arena.get(index).unwrap() or
//       arena.get_mut(index).unwrap()
// TODO: a bunch of examples, see the docs for Vec for inspiration.
pub struct VecArena<T> {
    // TODO: Docs for these fields
    slots: *const T,
    bitmap: Bitmap,
    marker: PhantomData<T>,
}

impl<T> VecArena<T> {
    /// Constructs a new, empty arena.
    ///
    /// The arena will not allocate until objects are inserted into it.
    pub fn new() -> Self {
        let slots = {
            let mut vec = Vec::with_capacity(0);
            let ptr = vec.as_mut_ptr();
            mem::forget(vec);
            ptr
        };
        VecArena {
            slots: slots,
            bitmap: Bitmap::new(),
            marker: PhantomData,
        }
    }

    /// Constructs a new, empty arena with the specified capacity (number of slots).
    ///
    /// The arena will be able to hold exactly `capacity` objects without reallocating.
    /// If `capacity` is 0, the arena will not allocate.
    pub fn with_capacity(cap: usize) -> Self {
        let mut arena = Self::new();
        arena.reserve_exact(cap);
        arena
    }

    /// Returns the number of objects the arena can hold without reallocating.
    /// In other words, this is the number of slots.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.bitmap.len()
    }

    /// Returns the number of objects in the arena.
    /// In other words, this is the number of occupied slots.
    #[inline]
    pub fn len(&self) -> usize {
        self.bitmap.occupied()
    }

    /// Returns `true` if the arena holds no objects.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Inserts an object into the arena and returns the slot index in which it was stored.
    /// The arena will reallocate if it's full.
    pub fn insert(&mut self, object: T) -> usize {
        if self.bitmap.occupied() == self.bitmap.len() {
            self.double();
        }

        let index = self.bitmap.acquire();
        unsafe {
            ptr::write(self.slots.offset(index as isize) as *mut T, object);
        }
        index
    }

    /// Inserts an object into the vacant slot at `index`.
    ///
    /// # Panics
    ///
    /// Panics if `index` is out of bounds or the slot is already occupied.
    pub fn insert_at(&mut self, index: usize, object: T) -> usize {
        self.bitmap.acquire_at(index);
        unsafe {
            ptr::write(self.slots.offset(index as isize) as *mut T, object);
        }
        index
    }

    /// Removes the object stored at `index` from the arena and returns it.
    ///
    /// # Panics
    ///
    /// Panics if `index` is out of bounds or the slot is already vacant.
    pub fn remove(&mut self, index: usize) -> T {
        // `release` will panic if the index is out of bounds or the slot is already vacant.
        self.bitmap.release(index);

        unsafe {
            ptr::read(self.slots.offset(index as isize) as *mut T)
        }
    }

    /// Clears the arena, removing and dropping all objects it holds. Keeps the allocated memory
    /// for reuse.
    pub fn clear(&mut self) {
        let mut arena = VecArena::with_capacity(self.capacity());
        mem::swap(self, &mut arena);
    }

    /// Returns a reference to the object stored at `index`.
    ///
    /// # Panics
    ///
    /// Panics if `index` is out of bounds.
    #[inline]
    pub fn get(&self, index: usize) -> Option<&T> {
        if self.bitmap.is_occupied(index) {
            unsafe { Some(self.get_unchecked(index)) }
        } else {
            None
        }
    }

    /// Returns a mutable reference to the object stored at `index`.
    ///
    /// # Panics
    ///
    /// Panics if `index` is out of bounds.
    #[inline]
    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
        if self.bitmap.is_occupied(index) {
            unsafe { Some(self.get_unchecked_mut(index)) }
        } else {
            None
        }
    }

    /// Reserves capacity for at least `additional` more elements to be inserted. The arena may
    /// reserve more space to avoid frequent reallocations.
    ///
    /// # Panics
    ///
    /// Panics if the new capacity overflows `usize`.
    pub fn reserve(&mut self, additional: usize) {
        let len = self.bitmap.len();
        self.bitmap.reserve(additional);
        unsafe {
            self.reallocate(len);
        }
    }

    /// Reserves the minimum capacity for exactly `additional` more elements to be inserted.
    ///
    /// Note that the allocator may give the arena more space than it requests.
    ///
    /// # Panics
    ///
    /// Panics if the new capacity overflows `usize`.
    pub fn reserve_exact(&mut self, additional: usize) {
        let len = self.bitmap.len();
        self.bitmap.reserve_exact(additional);
        unsafe {
            self.reallocate(len);
        }
    }

    /// Reallocates the object array because the bitmap was resized.
    unsafe fn reallocate(&mut self, old_len: usize) {
        let new_len = self.bitmap.len();

        // Allocate a new array.
        let mut vec = Vec::with_capacity(new_len);
        let ptr = vec.as_mut_ptr();
        mem::forget(vec);

        // Copy data into the new array.
        ptr::copy_nonoverlapping(self.slots, ptr, old_len);

        // Deallocate the old array.
        Vec::from_raw_parts(self.slots as *mut T, 0, old_len);

        self.slots = ptr;
    }

    /// Doubles the capacity of the arena.
    #[cold]
    #[inline(never)]
    fn double(&mut self) {
        let len = self.bitmap.len();
        let elem_size = mem::size_of::<T>();

        let new_len = if len == 0 {
            if elem_size.checked_mul(4).is_some() {
                4
            } else {
                1
            }
        } else {
            len.checked_mul(2).expect("len overflow")
        };

        self.reserve_exact(new_len - len);
    }

    #[inline]
    pub fn iter(&self) -> Iter<T> {
        Iter {
            arena: self,
            index: 0,
        }
    }

    #[inline]
    pub fn iter_mut(&mut self) -> IterMut<T> {
        IterMut {
            slots: self.slots as *mut T,
            bitmap: &self.bitmap,
            index: 0,
            marker: PhantomData,
        }
    }

    #[inline]
    pub fn drain(&mut self) -> Drain<T> {
        Drain {
            arena: self,
            index: 0,
        }
    }

    /// Returns a reference to the object at `index`, without bounds checking nor checking that
    /// the object exists.
    #[inline]
    pub unsafe fn get_unchecked(&self, index: usize) -> &T {
        &*self.slots.offset(index as isize)
    }

    /// Returns a mutable reference to the object at `index`, without bounds checking nor checking
    /// that the object exists.
    #[inline]
    pub unsafe fn get_unchecked_mut(&mut self, index: usize) -> &mut T {
        &mut *(self.slots.offset(index as isize) as *mut T)
    }
}

impl<T> Drop for VecArena<T> {
    fn drop(&mut self) {
        unsafe {
            // Drop all objects in the arena.
            for index in self.bitmap.iter() {
                ptr::drop_in_place(self.slots.offset(index as isize) as *mut T);
            }

            // Deallocate the old array.
            Vec::from_raw_parts(self.slots as *mut T, 0, self.bitmap.len());
        }
    }
}

impl<T> Index<usize> for VecArena<T> {
    type Output = T;

    #[inline]
    fn index(&self, index: usize) -> &T {
        self.get(index).expect("vacant slot at `index`")
    }
}

impl<T> IndexMut<usize> for VecArena<T> {
    #[inline]
    fn index_mut(&mut self, index: usize) -> &mut T {
        self.get_mut(index).expect("vacant slot at `index`")
    }
}

impl<T> Default for VecArena<T> {
    fn default() -> Self {
        VecArena::new()
    }
}

impl<T: Clone> Clone for VecArena<T> {
    fn clone(&self) -> Self {
        let mut arena = VecArena::with_capacity(self.capacity());
        for index in self.bitmap.iter() {
            let clone = unsafe {
                self.get_unchecked(index).clone()
            };
            arena.insert_at(index, clone);
        }
        arena
    }
}

pub struct IntoIter<T> {
    slots: *const T,
    bitmap: Bitmap,
    index: usize,
    marker: PhantomData<T>,
}

impl<T> Drop for IntoIter<T> {
    fn drop(&mut self) {
        while self.next().is_some() {}
    }
}

impl<T> Iterator for IntoIter<T> {
    type Item = (usize, T);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.bitmap.next_occupied(self.index).map(|index| {
            self.index = index + 1;
            unsafe {
                self.bitmap.release(index);
                (index, ptr::read(self.slots.offset(index as isize)))
            }
        })
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let lower = 0;
        let upper = self.bitmap.len() - self.index;
        (lower, Some(upper))
    }
}

impl<T> IntoIterator for VecArena<T> {
    type Item = (usize, T);
    type IntoIter = IntoIter<T>;

    fn into_iter(mut self) -> Self::IntoIter {
        let iter = IntoIter {
            slots: self.slots,
            bitmap: mem::replace(&mut self.bitmap, Bitmap::new()),
            index: 0,
            marker: PhantomData,
        };
        mem::forget(self);
        iter
    }
}

pub struct Iter<'a, T: 'a> {
    arena: &'a VecArena<T>,
    index: usize,
}

impl<'a, T> Iterator for Iter<'a, T> {
    type Item = (usize, &'a T);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.arena.bitmap.next_occupied(self.index).map(|index| {
            self.index = index + 1;
            unsafe {
                (index, self.arena.get_unchecked(index))
            }
        })
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let lower = 0;
        let upper = self.arena.capacity() - self.index;
        (lower, Some(upper))
    }
}

pub struct IterMut<'a, T: 'a> {
    slots: *mut T,
    bitmap: &'a Bitmap,
    index: usize,
    marker: PhantomData<&'a mut T>,
}

impl<'a, T> Iterator for IterMut<'a, T> {
    type Item = (usize, &'a mut T);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.bitmap.next_occupied(self.index).map(|index| {
            self.index = index + 1;
            unsafe {
                (index, &mut *self.slots.offset(index as isize))
            }
        })
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let lower = 0;
        let upper = self.bitmap.len() - self.index;
        (lower, Some(upper))
    }
}

pub struct Drain<'a, T: 'a> {
    arena: &'a mut VecArena<T>,
    index: usize,
}

impl<'a, T> Iterator for Drain<'a, T> {
    type Item = (usize, T);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.arena.bitmap.next_occupied(self.index).map(|index| {
            self.index = index + 1;
            unsafe {
                self.arena.bitmap.release(index);
                (index, ptr::read(self.arena.slots.offset(index as isize)))
            }
        })
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let lower = 0;
        let upper = self.arena.capacity() - self.index;
        (lower, Some(upper))
    }
}

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

    // TODO: test variance: https://github.com/rust-lang/rust/pull/30998/files

    #[test]
    fn it_works() {
        let mut arena = VecArena::new();
        for i in 0..10 {
            assert_eq!(arena.insert(()), i);
            assert!(arena[i] == ());
        }
        for i in 0..10 {
            assert!(arena[i] == ());
        }
    }
}
<|endoftext|>


<commit_before>extern crate nom_lua53;

use nom_lua53::{parse_all, ParseResult, Statement, Exp};
use nom_lua53::op::BinOp;

#[derive(PartialEq,Debug)]
pub enum LuaValue {
    Nil,
    Integer(isize),
    Float(f64),
    Boolean(bool),
    Str(String),
}

#[derive(PartialEq,Eq,Debug)]
pub enum LuaError {
    TypeError(String),
    ArithmeticError(String),
}

use LuaError::*;

type Result<T> = std::result::Result<T, LuaError>;

// What do we say? We say "Merci Basile!"
fn lit_to_string(string: &nom_lua53::string::StringLit) -> String {
    return String::from_utf8_lossy(&(string.0)).to_string();
}

fn eval_arithmetic(left_op: &Box<Exp>, right_op: &Box<Exp>,
                   integer: fn(isize, isize) -> Result<LuaValue>,
                   float: fn(f64, f64) -> Result<LuaValue>) -> Result<LuaValue> {

    let left_op = eval_expr(&left_op)?;
    let right_op = eval_expr(&right_op)?;

    match left_op {
        LuaValue::Integer(i) => match right_op {
            LuaValue::Integer(j) => integer(i, j),
            LuaValue::Float(j) => float(i as f64, j),
            _ => Err(TypeError("Trying to do arithmetic on non-numerical values.".to_owned()))
        },
        LuaValue::Float(i) => match right_op {
            LuaValue::Integer(j) => float(i, j as f64),
            LuaValue::Float(j) => float(i, j),
            _ => Err(TypeError("Trying to do arithmetic on non-numerical values.".to_owned()))
        },
        _ => Err(TypeError("Trying to do arithmetic on non-numerical values.".to_owned()))
    }
}

fn eval_binary_expr(left_op: &Box<Exp>, right_op: &Box<Exp>, operator: &BinOp) -> Result<LuaValue> {
    // We cannot yet evaluate the operands as some binary operators are used to shortcut
    // evaluation.
    match *operator {
        BinOp::Plus => eval_arithmetic(left_op, right_op,
                                       |i, j| Ok(LuaValue::Integer(i+j)),
                                       |i, j| Ok(LuaValue::Float(i+j))),
        BinOp::Minus => eval_arithmetic(left_op, right_op,
                                        |i, j| Ok(LuaValue::Integer(i-j)),
                                        |i, j| Ok(LuaValue::Float(i-j))),
        BinOp::Mul => eval_arithmetic(left_op, right_op,
                                      |i, j| Ok(LuaValue::Integer(i*j)),
                                      |i, j| Ok(LuaValue::Float(i*j))),
        BinOp::Div => eval_arithmetic(left_op, right_op,
                                      |i, j| if j == 0 {
                                          Err(ArithmeticError("Dividing by 0".to_owned()))
                                      } else {
                                          Ok(LuaValue::Float((i as f64)/(j as f64)))
                                      },
                                      |i, j| if j == 0. {
                                          Err(ArithmeticError("Dividing by 0".to_owned()))
                                      } else {
                                          Ok(LuaValue::Float(i/j))
                                      }),
        BinOp::Mod => eval_arithmetic(left_op, right_op,
                                      |i, j| if j == 0 {
                                          Err(ArithmeticError("Dividing by 0".to_owned()))
                                      } else {
                                          Ok(LuaValue::Integer(i % j))
                                      },
                                      |i, j| if j == 0. {
                                          Err(ArithmeticError("Dividing by 0".to_owned()))
                                      } else {
                                          Ok(LuaValue::Float(i - (i/j).floor()*j))
                                      }),
        BinOp::IntDiv => eval_arithmetic(left_op, right_op,
                                         |i, j| if j == 0 {
                                             Err(ArithmeticError("Dividing by 0".to_owned()))
                                         } else {
                                             Ok(LuaValue::Integer(i/j))
                                         },
                                         |i, j| if j == 0. {
                                             Err(ArithmeticError("Dividing by 0".to_owned()))
                                         } else {
                                             Ok(LuaValue::Integer((i/j).floor() as isize))
                                         }),
        BinOp::Pow => eval_arithmetic(left_op, right_op,
                                         |i, j| Ok(LuaValue::Float((i as f64).powf(j as f64))),
                                         |i, j| Ok(LuaValue::Float(i.powf(j)))),
        _ => Ok(LuaValue::Nil)
    }
}

fn eval_expr(expr: &Exp) -> Result<LuaValue> {
    match *expr {
        Exp::Nil => Ok(LuaValue::Nil),
        Exp::Bool(val) => Ok(LuaValue::Boolean(val)),
        Exp::Num(val) => Ok(match val {
            nom_lua53::num::Numeral::Float(fl) => LuaValue::Float(fl),
            nom_lua53::num::Numeral::Int(i) => LuaValue::Integer(i),
        }),
        Exp::BinExp(ref left, ref op, ref right) => eval_binary_expr(left, right, &op),
        Exp::Str(ref s) => Ok(LuaValue::Str(lit_to_string(s))),
        _ => Ok(LuaValue::Nil),
    }
}

pub fn eval_file(input: &[u8]) {
    match parse_all(input) {
        ParseResult::Done(blk) => {
            for stmt in blk.stmts {
                println!("{:?}", stmt);
                match stmt {
                    Statement::LVarAssign(ass) => {
                        let values = ass.vals.expect("There should be some values. Why isn't there any value?!");
                        for (var, val) in ass.vars.iter().zip(values.iter()) {
                            let computed_value = eval_expr(val);
                            println!("Assigning {:?} to {:?}", computed_value, var);
                        }
                    }
                    Statement::Assignment(ass) => {
                        println!("Assigning {:?} to {:?}", ass.vals, ass.vars);
                    }
                    _ => {}
                }
            }
        }

        ParseResult::Error(rest, ss) => {
            println!("Error. statements == {:#?}", ss);
            println!("rest == '{}'", String::from_utf8_lossy(rest));
        }
    }
}

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

    use nom_lua53::{parse_all, ParseResult, Statement, Exp};
    use nom_lua53::num::Numeral;
    use nom_lua53::op::BinOp;

    #[test]
    fn test_addition() {
        // 1. + -1. == 0.
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Float(1.0))), &Box::new(Exp::Num(Numeral::Float(-1.0))), &BinOp::Plus).unwrap();
        assert_eq!(res, LuaValue::Float(0.));

        // 1 + -1. == 0.
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(1))), &Box::new(Exp::Num(Numeral::Float(-1.0))), &BinOp::Plus).unwrap();
        assert_eq!( res, LuaValue::Float(0.));
        // 1 + 3 == 4
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(1))), &Box::new(Exp::Num(Numeral::Int(3))), &BinOp::Plus).unwrap();
        assert_eq!(res, LuaValue::Integer(4));
    }

    #[test]
    fn test_arithmetic_types() {
        eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(1))), &Box::new(Exp::Num(Numeral::Float(-1.0))), &BinOp::Plus).unwrap();

        eval_binary_expr(&Box::new(Exp::Num(Numeral::Float(1.))), &Box::new(Exp::Num(Numeral::Int(-1))), &BinOp::Plus).unwrap();

        let res = eval_binary_expr(&Box::new(Exp::Bool(true)), &Box::new(Exp::Num(Numeral::Int(-1))), &BinOp::Plus).unwrap_err();
        assert!(match res { TypeError(_) => true, _ => false })
    }

    #[test]
    fn test_substraction() {
        // 1. - -1. == 2.
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Float(1.0))), &Box::new(Exp::Num(Numeral::Float(-1.0))), &BinOp::Minus).unwrap();
        assert_eq!(res, LuaValue::Float(2.));

        // 1 - -1. == 2.
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(1))), &Box::new(Exp::Num(Numeral::Float(-1.0))), &BinOp::Minus).unwrap();
        assert_eq!(res, LuaValue::Float(2.));

        // 1 - 3 == -2
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(1))), &Box::new(Exp::Num(Numeral::Int(3))), &BinOp::Minus).unwrap();
        assert_eq!(res, LuaValue::Integer(-2));
    }

    #[test]
    fn test_multiplication() {
        // 1.5 * 2.5. == 3.75.
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Float(1.5))), &Box::new(Exp::Num(Numeral::Float(2.5))), &BinOp::Mul).unwrap();
        assert_eq!(res, LuaValue::Float(3.75));

        // -2 * 1.25 == -2.5.
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(-2))), &Box::new(Exp::Num(Numeral::Float(1.25))), &BinOp::Mul).unwrap();
        assert_eq!(res, LuaValue::Float(-2.5));

        // 10 * 3 == 30
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(10))), &Box::new(Exp::Num(Numeral::Int(3))), &BinOp::Mul).unwrap();
        assert_eq!(res, LuaValue::Integer(30));
    }

    #[test]
    fn test_division() {
        // 1.5 / 0.5. == 3.0
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Float(1.5))), &Box::new(Exp::Num(Numeral::Float(0.5))), &BinOp::Div).unwrap();
        assert_eq!(res, LuaValue::Float(3.));

        // 3 / -2. == -1.5
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(3))), &Box::new(Exp::Num(Numeral::Float(-2.0))), &BinOp::Div).unwrap();
        assert_eq!(res, LuaValue::Float(-1.5));

        // 3 / 2 == 1.5
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(3))), &Box::new(Exp::Num(Numeral::Int(2))), &BinOp::Div).unwrap();
        assert_eq!(res, LuaValue::Float(1.5));
    }

    #[test]
    fn test_division_null_error() {
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(1))), &Box::new(Exp::Num(Numeral::Int(0))), &BinOp::Div).unwrap_err();
        assert!(match res { ArithmeticError(_) => true, _ => false });

        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Float(1.))), &Box::new(Exp::Num(Numeral::Float(0.))), &BinOp::Div).unwrap_err();
        assert!(match res { ArithmeticError(_) => true, _ => false });
    }

    #[test]
    fn test_intdiv() {
        // 1.5 // 0.5. == 3
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Float(1.5))), &Box::new(Exp::Num(Numeral::Float(0.5))), &BinOp::IntDiv).unwrap();
        assert_eq!(res, LuaValue::Integer(3));

        // 3 // -2. == -2
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(3))), &Box::new(Exp::Num(Numeral::Float(-2.0))), &BinOp::IntDiv).unwrap();
        assert_eq!(res, LuaValue::Integer(-2));

        // 3 // 2 == 1
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(3))), &Box::new(Exp::Num(Numeral::Int(2))), &BinOp::IntDiv).unwrap();
        assert_eq!(res, LuaValue::Integer(1));
    }
}
<commit_msg>ALso test null-div for the integer division<commit_after>extern crate nom_lua53;

use nom_lua53::{parse_all, ParseResult, Statement, Exp};
use nom_lua53::op::BinOp;

#[derive(PartialEq,Debug)]
pub enum LuaValue {
    Nil,
    Integer(isize),
    Float(f64),
    Boolean(bool),
    Str(String),
}

#[derive(PartialEq,Eq,Debug)]
pub enum LuaError {
    TypeError(String),
    ArithmeticError(String),
}

use LuaError::*;

type Result<T> = std::result::Result<T, LuaError>;

// What do we say? We say "Merci Basile!"
fn lit_to_string(string: &nom_lua53::string::StringLit) -> String {
    return String::from_utf8_lossy(&(string.0)).to_string();
}

fn eval_arithmetic(left_op: &Box<Exp>, right_op: &Box<Exp>,
                   integer: fn(isize, isize) -> Result<LuaValue>,
                   float: fn(f64, f64) -> Result<LuaValue>) -> Result<LuaValue> {

    let left_op = eval_expr(&left_op)?;
    let right_op = eval_expr(&right_op)?;

    match left_op {
        LuaValue::Integer(i) => match right_op {
            LuaValue::Integer(j) => integer(i, j),
            LuaValue::Float(j) => float(i as f64, j),
            _ => Err(TypeError("Trying to do arithmetic on non-numerical values.".to_owned()))
        },
        LuaValue::Float(i) => match right_op {
            LuaValue::Integer(j) => float(i, j as f64),
            LuaValue::Float(j) => float(i, j),
            _ => Err(TypeError("Trying to do arithmetic on non-numerical values.".to_owned()))
        },
        _ => Err(TypeError("Trying to do arithmetic on non-numerical values.".to_owned()))
    }
}

fn eval_binary_expr(left_op: &Box<Exp>, right_op: &Box<Exp>, operator: &BinOp) -> Result<LuaValue> {
    // We cannot yet evaluate the operands as some binary operators are used to shortcut
    // evaluation.
    match *operator {
        BinOp::Plus => eval_arithmetic(left_op, right_op,
                                       |i, j| Ok(LuaValue::Integer(i+j)),
                                       |i, j| Ok(LuaValue::Float(i+j))),
        BinOp::Minus => eval_arithmetic(left_op, right_op,
                                        |i, j| Ok(LuaValue::Integer(i-j)),
                                        |i, j| Ok(LuaValue::Float(i-j))),
        BinOp::Mul => eval_arithmetic(left_op, right_op,
                                      |i, j| Ok(LuaValue::Integer(i*j)),
                                      |i, j| Ok(LuaValue::Float(i*j))),
        BinOp::Div => eval_arithmetic(left_op, right_op,
                                      |i, j| if j == 0 {
                                          Err(ArithmeticError("Dividing by 0".to_owned()))
                                      } else {
                                          Ok(LuaValue::Float((i as f64)/(j as f64)))
                                      },
                                      |i, j| if j == 0. {
                                          Err(ArithmeticError("Dividing by 0".to_owned()))
                                      } else {
                                          Ok(LuaValue::Float(i/j))
                                      }),
        BinOp::Mod => eval_arithmetic(left_op, right_op,
                                      |i, j| if j == 0 {
                                          Err(ArithmeticError("Dividing by 0".to_owned()))
                                      } else {
                                          Ok(LuaValue::Integer(i % j))
                                      },
                                      |i, j| if j == 0. {
                                          Err(ArithmeticError("Dividing by 0".to_owned()))
                                      } else {
                                          Ok(LuaValue::Float(i - (i/j).floor()*j))
                                      }),
        BinOp::IntDiv => eval_arithmetic(left_op, right_op,
                                         |i, j| if j == 0 {
                                             Err(ArithmeticError("Dividing by 0".to_owned()))
                                         } else {
                                             Ok(LuaValue::Integer(i/j))
                                         },
                                         |i, j| if j == 0. {
                                             Err(ArithmeticError("Dividing by 0".to_owned()))
                                         } else {
                                             Ok(LuaValue::Integer((i/j).floor() as isize))
                                         }),
        BinOp::Pow => eval_arithmetic(left_op, right_op,
                                         |i, j| Ok(LuaValue::Float((i as f64).powf(j as f64))),
                                         |i, j| Ok(LuaValue::Float(i.powf(j)))),
        _ => Ok(LuaValue::Nil)
    }
}

fn eval_expr(expr: &Exp) -> Result<LuaValue> {
    match *expr {
        Exp::Nil => Ok(LuaValue::Nil),
        Exp::Bool(val) => Ok(LuaValue::Boolean(val)),
        Exp::Num(val) => Ok(match val {
            nom_lua53::num::Numeral::Float(fl) => LuaValue::Float(fl),
            nom_lua53::num::Numeral::Int(i) => LuaValue::Integer(i),
        }),
        Exp::BinExp(ref left, ref op, ref right) => eval_binary_expr(left, right, &op),
        Exp::Str(ref s) => Ok(LuaValue::Str(lit_to_string(s))),
        _ => Ok(LuaValue::Nil),
    }
}

pub fn eval_file(input: &[u8]) {
    match parse_all(input) {
        ParseResult::Done(blk) => {
            for stmt in blk.stmts {
                println!("{:?}", stmt);
                match stmt {
                    Statement::LVarAssign(ass) => {
                        let values = ass.vals.expect("There should be some values. Why isn't there any value?!");
                        for (var, val) in ass.vars.iter().zip(values.iter()) {
                            let computed_value = eval_expr(val);
                            println!("Assigning {:?} to {:?}", computed_value, var);
                        }
                    }
                    Statement::Assignment(ass) => {
                        println!("Assigning {:?} to {:?}", ass.vals, ass.vars);
                    }
                    _ => {}
                }
            }
        }

        ParseResult::Error(rest, ss) => {
            println!("Error. statements == {:#?}", ss);
            println!("rest == '{}'", String::from_utf8_lossy(rest));
        }
    }
}

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

    use nom_lua53::{parse_all, ParseResult, Statement, Exp};
    use nom_lua53::num::Numeral;
    use nom_lua53::op::BinOp;

    #[test]
    fn test_addition() {
        // 1. + -1. == 0.
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Float(1.0))), &Box::new(Exp::Num(Numeral::Float(-1.0))), &BinOp::Plus).unwrap();
        assert_eq!(res, LuaValue::Float(0.));

        // 1 + -1. == 0.
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(1))), &Box::new(Exp::Num(Numeral::Float(-1.0))), &BinOp::Plus).unwrap();
        assert_eq!( res, LuaValue::Float(0.));
        // 1 + 3 == 4
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(1))), &Box::new(Exp::Num(Numeral::Int(3))), &BinOp::Plus).unwrap();
        assert_eq!(res, LuaValue::Integer(4));
    }

    #[test]
    fn test_arithmetic_types() {
        eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(1))), &Box::new(Exp::Num(Numeral::Float(-1.0))), &BinOp::Plus).unwrap();

        eval_binary_expr(&Box::new(Exp::Num(Numeral::Float(1.))), &Box::new(Exp::Num(Numeral::Int(-1))), &BinOp::Plus).unwrap();

        let res = eval_binary_expr(&Box::new(Exp::Bool(true)), &Box::new(Exp::Num(Numeral::Int(-1))), &BinOp::Plus).unwrap_err();
        assert!(match res { TypeError(_) => true, _ => false })
    }

    #[test]
    fn test_substraction() {
        // 1. - -1. == 2.
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Float(1.0))), &Box::new(Exp::Num(Numeral::Float(-1.0))), &BinOp::Minus).unwrap();
        assert_eq!(res, LuaValue::Float(2.));

        // 1 - -1. == 2.
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(1))), &Box::new(Exp::Num(Numeral::Float(-1.0))), &BinOp::Minus).unwrap();
        assert_eq!(res, LuaValue::Float(2.));

        // 1 - 3 == -2
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(1))), &Box::new(Exp::Num(Numeral::Int(3))), &BinOp::Minus).unwrap();
        assert_eq!(res, LuaValue::Integer(-2));
    }

    #[test]
    fn test_multiplication() {
        // 1.5 * 2.5. == 3.75.
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Float(1.5))), &Box::new(Exp::Num(Numeral::Float(2.5))), &BinOp::Mul).unwrap();
        assert_eq!(res, LuaValue::Float(3.75));

        // -2 * 1.25 == -2.5.
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(-2))), &Box::new(Exp::Num(Numeral::Float(1.25))), &BinOp::Mul).unwrap();
        assert_eq!(res, LuaValue::Float(-2.5));

        // 10 * 3 == 30
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(10))), &Box::new(Exp::Num(Numeral::Int(3))), &BinOp::Mul).unwrap();
        assert_eq!(res, LuaValue::Integer(30));
    }

    #[test]
    fn test_division() {
        // 1.5 / 0.5. == 3.0
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Float(1.5))), &Box::new(Exp::Num(Numeral::Float(0.5))), &BinOp::Div).unwrap();
        assert_eq!(res, LuaValue::Float(3.));

        // 3 / -2. == -1.5
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(3))), &Box::new(Exp::Num(Numeral::Float(-2.0))), &BinOp::Div).unwrap();
        assert_eq!(res, LuaValue::Float(-1.5));

        // 3 / 2 == 1.5
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(3))), &Box::new(Exp::Num(Numeral::Int(2))), &BinOp::Div).unwrap();
        assert_eq!(res, LuaValue::Float(1.5));
    }

    #[test]
    fn test_intdiv() {
        // 1.5 // 0.5. == 3
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Float(1.5))), &Box::new(Exp::Num(Numeral::Float(0.5))), &BinOp::IntDiv).unwrap();
        assert_eq!(res, LuaValue::Integer(3));

        // 3 // -2. == -2
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(3))), &Box::new(Exp::Num(Numeral::Float(-2.0))), &BinOp::IntDiv).unwrap();
        assert_eq!(res, LuaValue::Integer(-2));

        // 3 // 2 == 1
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(3))), &Box::new(Exp::Num(Numeral::Int(2))), &BinOp::IntDiv).unwrap();
        assert_eq!(res, LuaValue::Integer(1));
    }

    #[test]
    fn test_divisions_null_error() {
        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(1))), &Box::new(Exp::Num(Numeral::Int(0))), &BinOp::Div).unwrap_err();
        assert!(match res { ArithmeticError(_) => true, _ => false });

        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Float(1.))), &Box::new(Exp::Num(Numeral::Float(0.))), &BinOp::Div).unwrap_err();
        assert!(match res { ArithmeticError(_) => true, _ => false });

        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Int(1))), &Box::new(Exp::Num(Numeral::Int(0))), &BinOp::IntDiv).unwrap_err();
        assert!(match res { ArithmeticError(_) => true, _ => false });

        let res = eval_binary_expr(&Box::new(Exp::Num(Numeral::Float(1.))), &Box::new(Exp::Num(Numeral::Float(0.))), &BinOp::IntDiv).unwrap_err();
        assert!(match res { ArithmeticError(_) => true, _ => false });
    }
}
<|endoftext|>


<commit_before>//! Rust friendly bindings to the various *nix system functions.
//!
//! Modules are structured according to the C header file that they would be
//! defined in.
#![crate_name = "nix"]
#![cfg(unix)]
#![allow(non_camel_case_types)]
// latest bitflags triggers a rustc bug with cross-crate macro expansions causing dead_code
// warnings even though the macro expands into something with allow(dead_code)
#![allow(dead_code)]
#![deny(warnings)]

#[macro_use]
extern crate bitflags;

extern crate libc;

#[cfg(test)]
extern crate nix_test as nixtest;

// Re-export some libc constants
pub use libc::{c_int, c_void};

pub mod errno;
pub mod features;
pub mod fcntl;

#[cfg(any(target_os = "linux", target_os = "android"))]
pub mod mount;

#[cfg(any(target_os = "linux"))]
pub mod mqueue;

#[cfg(any(target_os = "linux", target_os = "macos"))]
pub mod poll;

pub mod net;

#[cfg(any(target_os = "linux", target_os = "android"))]
pub mod sched;

pub mod sys;
pub mod unistd;

/*
 *
 * ===== Result / Error =====
 *
 */

use libc::c_char;
use std::{ptr, result};
use std::ffi::CStr;
use std::path::{Path, PathBuf};
use std::os::unix::ffi::OsStrExt;
use std::io;
use std::fmt;
use std::error;
use libc::PATH_MAX;

pub type Result<T> = result::Result<T, Error>;

#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Error {
    Sys(errno::Errno),
    InvalidPath,
}

impl Error {
    pub fn from_errno(errno: errno::Errno) -> Error {
        Error::Sys(errno)
    }

    pub fn last() -> Error {
        Error::Sys(errno::Errno::last())
    }

    pub fn invalid_argument() -> Error {
        Error::Sys(errno::EINVAL)
    }

    pub fn errno(&self) -> errno::Errno {
        match *self {
            Error::Sys(errno) => errno,
            Error::InvalidPath => errno::Errno::EINVAL,
        }
    }
}

impl From<errno::Errno> for Error {
    fn from(errno: errno::Errno) -> Error { Error::from_errno(errno) }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match self {
            &Error::InvalidPath => "Invalid path",
            &Error::Sys(ref errno) => errno.desc(),
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Error::InvalidPath => write!(f, "Invalid path"),
            &Error::Sys(errno) => write!(f, "{:?}: {}", errno, errno.desc()),
        }
    }
}

impl From<Error> for io::Error {
    fn from(err: Error) -> Self {
        match err {
            Error::InvalidPath => io::Error::new(io::ErrorKind::InvalidInput, err),
            Error::Sys(errno) => io::Error::from_raw_os_error(errno as i32),
        }
    }
}

pub trait NixPath {
    fn len(&self) -> usize;

    fn with_nix_path<T, F>(&self, f: F) -> Result<T>
        where F: FnOnce(&CStr) -> T;
}

impl NixPath for CStr {
    fn len(&self) -> usize {
        self.to_bytes().len()
    }

    fn with_nix_path<T, F>(&self, f: F) -> Result<T>
            where F: FnOnce(&CStr) -> T {
        // Equivalence with the [u8] impl.
        if self.len() >= PATH_MAX as usize {
            return Err(Error::InvalidPath);
        }

        Ok(f(self))
    }
}

impl NixPath for [u8] {
    fn len(&self) -> usize {
        self.len()
    }

    fn with_nix_path<T, F>(&self, f: F) -> Result<T>
            where F: FnOnce(&CStr) -> T {
        let mut buf = [0u8; PATH_MAX as usize];

        if self.len() >= PATH_MAX as usize {
            return Err(Error::InvalidPath);
        }

        match self.iter().position(|b| *b == 0) {
            Some(_) => Err(Error::InvalidPath),
            None => {
                unsafe {
                    // TODO: Replace with bytes::copy_memory. rust-lang/rust#24028
                    ptr::copy_nonoverlapping(self.as_ptr(), buf.as_mut_ptr(), self.len());
                    Ok(f(CStr::from_ptr(buf.as_ptr() as *const c_char)))
                }

            }
        }
    }
}

impl NixPath for Path {
    fn len(&self) -> usize {
        self.as_os_str().as_bytes().len()
    }

    fn with_nix_path<T, F>(&self, f: F) -> Result<T> where F: FnOnce(&CStr) -> T {
        self.as_os_str().as_bytes().with_nix_path(f)
    }
}

impl NixPath for PathBuf {
    fn len(&self) -> usize {
        self.as_os_str().as_bytes().len()
    }

    fn with_nix_path<T, F>(&self, f: F) -> Result<T> where F: FnOnce(&CStr) -> T {
        self.as_os_str().as_bytes().with_nix_path(f)
    }
}

#[inline]
pub fn from_ffi(res: libc::c_int) -> Result<()> {
    if res != 0 {
        return Err(Error::Sys(errno::Errno::last()));
    }

    Ok(())
}
<commit_msg>Implement NixPath for str and OsStr<commit_after>//! Rust friendly bindings to the various *nix system functions.
//!
//! Modules are structured according to the C header file that they would be
//! defined in.
#![crate_name = "nix"]
#![cfg(unix)]
#![allow(non_camel_case_types)]
// latest bitflags triggers a rustc bug with cross-crate macro expansions causing dead_code
// warnings even though the macro expands into something with allow(dead_code)
#![allow(dead_code)]
#![deny(warnings)]

#[macro_use]
extern crate bitflags;

extern crate libc;

#[cfg(test)]
extern crate nix_test as nixtest;

// Re-export some libc constants
pub use libc::{c_int, c_void};

pub mod errno;
pub mod features;
pub mod fcntl;

#[cfg(any(target_os = "linux", target_os = "android"))]
pub mod mount;

#[cfg(any(target_os = "linux"))]
pub mod mqueue;

#[cfg(any(target_os = "linux", target_os = "macos"))]
pub mod poll;

pub mod net;

#[cfg(any(target_os = "linux", target_os = "android"))]
pub mod sched;

pub mod sys;
pub mod unistd;

/*
 *
 * ===== Result / Error =====
 *
 */

use libc::c_char;
use std::{ptr, result};
use std::ffi::{CStr, OsStr};
use std::path::{Path, PathBuf};
use std::os::unix::ffi::OsStrExt;
use std::io;
use std::fmt;
use std::error;
use libc::PATH_MAX;

pub type Result<T> = result::Result<T, Error>;

#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Error {
    Sys(errno::Errno),
    InvalidPath,
}

impl Error {
    pub fn from_errno(errno: errno::Errno) -> Error {
        Error::Sys(errno)
    }

    pub fn last() -> Error {
        Error::Sys(errno::Errno::last())
    }

    pub fn invalid_argument() -> Error {
        Error::Sys(errno::EINVAL)
    }

    pub fn errno(&self) -> errno::Errno {
        match *self {
            Error::Sys(errno) => errno,
            Error::InvalidPath => errno::Errno::EINVAL,
        }
    }
}

impl From<errno::Errno> for Error {
    fn from(errno: errno::Errno) -> Error { Error::from_errno(errno) }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match self {
            &Error::InvalidPath => "Invalid path",
            &Error::Sys(ref errno) => errno.desc(),
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Error::InvalidPath => write!(f, "Invalid path"),
            &Error::Sys(errno) => write!(f, "{:?}: {}", errno, errno.desc()),
        }
    }
}

impl From<Error> for io::Error {
    fn from(err: Error) -> Self {
        match err {
            Error::InvalidPath => io::Error::new(io::ErrorKind::InvalidInput, err),
            Error::Sys(errno) => io::Error::from_raw_os_error(errno as i32),
        }
    }
}

pub trait NixPath {
    fn len(&self) -> usize;

    fn with_nix_path<T, F>(&self, f: F) -> Result<T>
        where F: FnOnce(&CStr) -> T;
}

impl NixPath for str {
    fn len(&self) -> usize {
        NixPath::len(OsStr::new(self))
    }

    fn with_nix_path<T, F>(&self, f: F) -> Result<T>
        where F: FnOnce(&CStr) -> T {
            OsStr::new(self).with_nix_path(f)
        }
}

impl NixPath for OsStr {
    fn len(&self) -> usize {
        self.as_bytes().len()
    }

    fn with_nix_path<T, F>(&self, f: F) -> Result<T>
        where F: FnOnce(&CStr) -> T {
            self.as_bytes().with_nix_path(f)
        }
}

impl NixPath for CStr {
    fn len(&self) -> usize {
        self.to_bytes().len()
    }

    fn with_nix_path<T, F>(&self, f: F) -> Result<T>
            where F: FnOnce(&CStr) -> T {
        // Equivalence with the [u8] impl.
        if self.len() >= PATH_MAX as usize {
            return Err(Error::InvalidPath);
        }

        Ok(f(self))
    }
}

impl NixPath for [u8] {
    fn len(&self) -> usize {
        self.len()
    }

    fn with_nix_path<T, F>(&self, f: F) -> Result<T>
            where F: FnOnce(&CStr) -> T {
        let mut buf = [0u8; PATH_MAX as usize];

        if self.len() >= PATH_MAX as usize {
            return Err(Error::InvalidPath);
        }

        match self.iter().position(|b| *b == 0) {
            Some(_) => Err(Error::InvalidPath),
            None => {
                unsafe {
                    // TODO: Replace with bytes::copy_memory. rust-lang/rust#24028
                    ptr::copy_nonoverlapping(self.as_ptr(), buf.as_mut_ptr(), self.len());
                    Ok(f(CStr::from_ptr(buf.as_ptr() as *const c_char)))
                }

            }
        }
    }
}

impl NixPath for Path {
    fn len(&self) -> usize {
        NixPath::len(self.as_os_str())
    }

    fn with_nix_path<T, F>(&self, f: F) -> Result<T> where F: FnOnce(&CStr) -> T {
        self.as_os_str().with_nix_path(f)
    }
}

impl NixPath for PathBuf {
    fn len(&self) -> usize {
        NixPath::len(self.as_os_str())
    }

    fn with_nix_path<T, F>(&self, f: F) -> Result<T> where F: FnOnce(&CStr) -> T {
        self.as_os_str().with_nix_path(f)
    }
}

#[inline]
pub fn from_ffi(res: libc::c_int) -> Result<()> {
    if res != 0 {
        return Err(Error::Sys(errno::Errno::last()));
    }

    Ok(())
}
<|endoftext|>


<commit_before>// Copyright 2015 Ilkka Rauta
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! BitReader is a helper type to extract strings of bits from a slice of bytes.
//!
//! Here is how you read first a single bit, then three bits and finally four bits from a byte
//! buffer:
//!
//! ```
//! use bitreader::BitReader;
//!
//! let slice_of_u8 = &[0b1000_1111];
//! let mut reader = BitReader::new(slice_of_u8);
//!
//! // You probably should use try! or some other error handling mechanism in real code if the
//! // length of the input is not known in advance.
//! let a_single_bit = reader.read_u8(1).unwrap(); // 1
//! let more_bits = reader.read_u8(3).unwrap(); // 0
//! let last_bits_of_byte = reader.read_u8(4).unwrap(); // 0b1111
//! ```
//! You can naturally read bits from longer buffer of data than just a single byte.
//!
//! As you read bits, the internal cursor of BitReader moves on along the stream of bits. Little
//! endian format is assumed when reading the multi-byte values. BitReader supports reading maximum
//! of 64 bits at a time (with read_u64). Reading signed values directly is not supported at the
//! moment.
//!
//! The reads do not need to be aligned in any particular way.
//!
//! Reading zero bits is a no-op.
//!
//! You can also skip over a number of bits, in which case there is no arbitrary small limits like
//! when reading the values to a variable. However, you can not seek past the end of the slice,
//! either when reading or when skipping bits.
//!
//! Note that the code will likely not work correctly if the slice is longer than 2^61 bytes, but
//! exceeding that should be pretty unlikely. Let's get back to this when people read exabytes of
//! information one bit at a time.

use std::fmt;
use std::error::Error;
use std::result;

/// BitReader reads data from a byte slice at the granularity of a single bit.
pub struct BitReader<'a> {
    bytes: &'a [u8],
    /// Position from the start of the slice, counted as bits instead of bytes
    position: u64,
}

impl<'a> BitReader<'a> {
    /// Construct a new BitReader from a byte slice. The returned reader lives at most as long as
    /// the slice given to is valid.
    pub fn new(bytes: &'a [u8]) -> BitReader<'a> {
        BitReader {
            bytes: bytes,
            position: 0
        }
    }

    /// Read at most 8 bits into a u8.
    pub fn read_u8(&mut self, bit_count: u8) -> Result<u8> {
        let value = try!(self.read_value(bit_count, 8));
        Ok((value & 0xff) as u8)
    }

    /// Read at most 16 bits into a u16.
    pub fn read_u16(&mut self, bit_count: u8) -> Result<u16> {
        let value = try!(self.read_value(bit_count, 16));
        Ok((value & 0xffff) as u16)
    }

    /// Read at most 32 bits into a u32.
    pub fn read_u32(&mut self, bit_count: u8) -> Result<u32> {
        let value = try!(self.read_value(bit_count, 32));
        Ok((value & 0xffffffff) as u32)
    }

    /// Read at most 64 bits into a u64.
    pub fn read_u64(&mut self, bit_count: u8) -> Result<u64> {
        let value = try!(self.read_value(bit_count, 64));
        Ok(value)
    }

    /// Skip arbitrary number of bits. However, you can skip at most to the end of the byte slice.
    pub fn skip(&mut self, bit_count: u32) -> Result<()> {
        let end_position = self.position + bit_count as u64;
        if end_position > self.bytes.len() as u64 * 8 {
            return Err(BitReaderError::NotEnoughData);
        }
        self.position = end_position;
        Ok(())
    }

    /// Returns the position of the cursor, or how many bits have been read so far.
    pub fn position(&self) -> u64 {
        self.position
    }

    /// Helper to make sure the "bit cursor" is exactly at the beginning of a byte, or at specific
    /// multi-byte alignment position.
    ///
    /// For example `reader.is_aligned(1)` returns true if exactly n bytes, or n * 8 bits, has been
    /// read. Similarly, `reader.is_aligned(4)` returns true if exactly n * 32 bits, or n 4-byte
    /// sequences has been read.
    ///
    /// This function can be used to validate the data is being read properly, for example by
    /// adding invocations wrapped into `debug_assert!()` to places where it is known the data
    /// should be n-byte aligned.
    pub fn is_aligned(&self, alignment_bytes: u32) -> bool {
        self.position % (alignment_bytes as u64 * 8) == 0
    }

    fn read_value(&mut self, bit_count: u8, maximum_count: u8) -> Result<u64> {
        if bit_count == 0 {
            return Ok(0);
        }
        if bit_count > maximum_count {
            return Err(BitReaderError::TooManyBitsForType);
        }
        let start_position = self.position;
        let end_position = self.position + bit_count as u64;
        if end_position > self.bytes.len() as u64 * 8 {
            return Err(BitReaderError::NotEnoughData);
        }

        let mut value: u64 = 0;

        for i in start_position..end_position {
            let byte_index = (i / 8) as usize;
            let byte = self.bytes[byte_index];
            let shift = 7 - (i % 8);
            let bit = (byte >> shift) as u64 & 1;
            value = (value << 1) | bit;
        }

        self.position = end_position;
        Ok(value)
    }
}

/// Result type for those BitReader operations that can fail.
pub type Result<T> = result::Result<T, BitReaderError>;

/// Error enumeration of BitReader errors.
#[derive(Debug,PartialEq)]
pub enum BitReaderError {
    /// Requested more bits than there are left in the byte slice at the current position.
    NotEnoughData,
    /// Requested more bits than the returned variable can hold, for example more than 8 bits when
    /// reading into a u8.
    TooManyBitsForType
}

impl Error for BitReaderError {
    fn description(&self) -> &str {
        match *self {
            BitReaderError::NotEnoughData => "Requested more bits than the byte slice has left",
            BitReaderError::TooManyBitsForType => "Requested more bits than the requested integer type can hold",
        }
    }
}

impl fmt::Display for BitReaderError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        self.description().fmt(fmt)
    }
}

#[test]
fn read_buffer() {
    let bytes = &[
        0b1011_0101, 0b0110_1010, 0b1010_1100, 0b1001_1001,
        0b1001_1001, 0b1001_1001, 0b1001_1001, 0b1110_0111,
    ];

    let mut reader = BitReader::new(bytes);

    assert_eq!(reader.read_u8(1).unwrap(), 0b1);
    assert_eq!(reader.read_u8(1).unwrap(), 0b0);
    assert_eq!(reader.read_u8(2).unwrap(), 0b11);

    assert_eq!(reader.read_u8(4).unwrap(), 0b0101);

    assert!(reader.is_aligned(1));

    assert_eq!(reader.read_u8(3).unwrap(), 0b11);
    assert_eq!(reader.read_u16(10).unwrap(), 0b01_0101_0101);
    assert_eq!(reader.read_u8(3).unwrap(), 0b100);

    assert!(reader.is_aligned(1));

    assert_eq!(reader.read_u32(32).unwrap(), 0b1001_1001_1001_1001_1001_1001_1001_1001);

    assert_eq!(reader.read_u8(4).unwrap(), 0b1110);
    assert_eq!(reader.read_u8(3).unwrap(), 0b011);
    assert_eq!(reader.read_u8(1).unwrap(), 0b1);

    // Could also be 8 at this point!
    assert!(reader.is_aligned(4));
}

#[test]
fn try_all_sizes() {
    let bytes = &[
        0x4a, 0x1e, 0x39, 0xbb, 0xd0, 0x07, 0xca, 0x9a,
        0xa6, 0xba, 0x25, 0x52, 0x6f, 0x0a, 0x6a, 0xba,
    ];

    let mut reader = BitReader::new(bytes);
    assert_eq!(reader.read_u64(64).unwrap(), 0x4a1e39bbd007ca9a);
    assert_eq!(reader.read_u64(64).unwrap(), 0xa6ba25526f0a6aba);

    let mut reader = BitReader::new(bytes);
    assert_eq!(reader.read_u32(32).unwrap(), 0x4a1e39bb);
    assert_eq!(reader.read_u32(32).unwrap(), 0xd007ca9a);
    assert_eq!(reader.read_u32(32).unwrap(), 0xa6ba2552);
    assert_eq!(reader.read_u32(32).unwrap(), 0x6f0a6aba);

    let mut reader = BitReader::new(bytes);
    assert_eq!(reader.read_u16(16).unwrap(), 0x4a1e);
    assert_eq!(reader.read_u16(16).unwrap(), 0x39bb);
    assert_eq!(reader.read_u16(16).unwrap(), 0xd007);
    assert_eq!(reader.read_u16(16).unwrap(), 0xca9a);
    assert_eq!(reader.read_u16(16).unwrap(), 0xa6ba);
    assert_eq!(reader.read_u16(16).unwrap(), 0x2552);
    assert_eq!(reader.read_u16(16).unwrap(), 0x6f0a);
    assert_eq!(reader.read_u16(16).unwrap(), 0x6aba);

    let mut reader = BitReader::new(&bytes[..]);
    for byte in bytes {
        assert_eq!(reader.read_u8(8).unwrap(), *byte);
    }
}

#[test]
fn skipping_and_zero_reads() {
    let bytes = &[
        0b1011_0101, 0b1110_1010, 0b1010_1100,
    ];

    let mut reader = BitReader::new(bytes);

    assert_eq!(reader.read_u8(4).unwrap(), 0b1011);
    // Reading zero bits should be a no-op
    assert_eq!(reader.read_u8(0).unwrap(), 0b0);
    assert_eq!(reader.read_u8(4).unwrap(), 0b0101);
    reader.skip(3).unwrap(); // 0b111
    assert_eq!(reader.read_u16(10).unwrap(), 0b0101010101);
    assert_eq!(reader.read_u8(3).unwrap(), 0b100);
}

#[test]
fn errors() {
    let bytes = &[
        0b1011_0101, 0b1110_1010, 0b1010_1100,
    ];

    let mut reader = BitReader::new(bytes);
    assert_eq!(reader.read_u8(4).unwrap(), 0b1011);
    assert_eq!(reader.read_u8(9).unwrap_err(), BitReaderError::TooManyBitsForType);
    // If an error happens, it should be possible to resume as if nothing had happened
    assert_eq!(reader.read_u8(4).unwrap(), 0b0101);

    let mut reader = BitReader::new(bytes);
    assert_eq!(reader.read_u8(4).unwrap(), 0b1011);
    // Same with this error
    assert_eq!(reader.read_u32(21).unwrap_err(), BitReaderError::NotEnoughData);
    assert_eq!(reader.read_u8(4).unwrap(), 0b0101);
}
<commit_msg>Add support for reading signed values.<commit_after>// Copyright 2015 Ilkka Rauta
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! BitReader is a helper type to extract strings of bits from a slice of bytes.
//!
//! Here is how you read first a single bit, then three bits and finally four bits from a byte
//! buffer:
//!
//! ```
//! use bitreader::BitReader;
//!
//! let slice_of_u8 = &[0b1000_1111];
//! let mut reader = BitReader::new(slice_of_u8);
//!
//! // You probably should use try! or some other error handling mechanism in real code if the
//! // length of the input is not known in advance.
//! let a_single_bit = reader.read_u8(1).unwrap(); // 1
//! let more_bits = reader.read_u8(3).unwrap(); // 0
//! let last_bits_of_byte = reader.read_u8(4).unwrap(); // 0b1111
//! ```
//! You can naturally read bits from longer buffer of data than just a single byte.
//!
//! As you read bits, the internal cursor of BitReader moves on along the stream of bits. Little
//! endian format is assumed when reading the multi-byte values. BitReader supports reading maximum
//! of 64 bits at a time (with read_u64). Reading signed values directly is not supported at the
//! moment.
//!
//! The reads do not need to be aligned in any particular way.
//!
//! Reading zero bits is a no-op.
//!
//! You can also skip over a number of bits, in which case there is no arbitrary small limits like
//! when reading the values to a variable. However, you can not seek past the end of the slice,
//! either when reading or when skipping bits.
//!
//! Note that the code will likely not work correctly if the slice is longer than 2^61 bytes, but
//! exceeding that should be pretty unlikely. Let's get back to this when people read exabytes of
//! information one bit at a time.

use std::fmt;
use std::error::Error;
use std::result;

/// BitReader reads data from a byte slice at the granularity of a single bit.
pub struct BitReader<'a> {
    bytes: &'a [u8],
    /// Position from the start of the slice, counted as bits instead of bytes
    position: u64,
}

impl<'a> BitReader<'a> {
    /// Construct a new BitReader from a byte slice. The returned reader lives at most as long as
    /// the slice given to is valid.
    pub fn new(bytes: &'a [u8]) -> BitReader<'a> {
        BitReader {
            bytes: bytes,
            position: 0
        }
    }

    /// Read at most 8 bits into a u8.
    pub fn read_u8(&mut self, bit_count: u8) -> Result<u8> {
        let value = try!(self.read_value(bit_count, 8));
        Ok((value & 0xff) as u8)
    }

    /// Read at most 16 bits into a u16.
    pub fn read_u16(&mut self, bit_count: u8) -> Result<u16> {
        let value = try!(self.read_value(bit_count, 16));
        Ok((value & 0xffff) as u16)
    }

    /// Read at most 32 bits into a u32.
    pub fn read_u32(&mut self, bit_count: u8) -> Result<u32> {
        let value = try!(self.read_value(bit_count, 32));
        Ok((value & 0xffffffff) as u32)
    }

    /// Read at most 64 bits into a u64.
    pub fn read_u64(&mut self, bit_count: u8) -> Result<u64> {
        let value = try!(self.read_value(bit_count, 64));
        Ok(value)
    }

    /// Read at most 8 bits into a i8.
    /// Assumes the bits are stored in two's complement format.
    pub fn read_i8(&mut self, bit_count: u8) -> Result<i8> {
        let value = try!(self.read_signed_value(bit_count, 8));
        Ok((value & 0xff) as i8)
    }

    /// Read at most 16 bits into a i16.
    /// Assumes the bits are stored in two's complement format.
    pub fn read_i16(&mut self, bit_count: u8) -> Result<i16> {
        let value = try!(self.read_signed_value(bit_count, 16));
        Ok((value & 0xffff) as i16)
    }

    /// Read at most 32 bits into a i32.
    /// Assumes the bits are stored in two's complement format.
    pub fn read_i32(&mut self, bit_count: u8) -> Result<i32> {
        let value = try!(self.read_signed_value(bit_count, 32));
        Ok((value & 0xffffffff) as i32)
    }

    /// Read at most 64 bits into a i64.
    /// Assumes the bits are stored in two's complement format.
    pub fn read_i64(&mut self, bit_count: u8) -> Result<i64> {
        let value = try!(self.read_signed_value(bit_count, 64));
        Ok(value)
    }

    /// Skip arbitrary number of bits. However, you can skip at most to the end of the byte slice.
    pub fn skip(&mut self, bit_count: u32) -> Result<()> {
        let end_position = self.position + bit_count as u64;
        if end_position > self.bytes.len() as u64 * 8 {
            return Err(BitReaderError::NotEnoughData);
        }
        self.position = end_position;
        Ok(())
    }

    /// Returns the position of the cursor, or how many bits have been read so far.
    pub fn position(&self) -> u64 {
        self.position
    }

    /// Helper to make sure the "bit cursor" is exactly at the beginning of a byte, or at specific
    /// multi-byte alignment position.
    ///
    /// For example `reader.is_aligned(1)` returns true if exactly n bytes, or n * 8 bits, has been
    /// read. Similarly, `reader.is_aligned(4)` returns true if exactly n * 32 bits, or n 4-byte
    /// sequences has been read.
    ///
    /// This function can be used to validate the data is being read properly, for example by
    /// adding invocations wrapped into `debug_assert!()` to places where it is known the data
    /// should be n-byte aligned.
    pub fn is_aligned(&self, alignment_bytes: u32) -> bool {
        self.position % (alignment_bytes as u64 * 8) == 0
    }

    fn read_signed_value(&mut self, bit_count: u8, maximum_count: u8) -> Result<i64> {
        let unsigned = try!(self.read_value(bit_count, maximum_count));
        // Fill the bits above the requested bits with all ones or all zeros,
        // depending on the sign bit.
        let sign_bit = unsigned >> (bit_count - 1) & 1;
        let high_bits = if sign_bit == 1 { -1 } else { 0 };
        Ok(high_bits << bit_count | unsigned as i64)
    }

    fn read_value(&mut self, bit_count: u8, maximum_count: u8) -> Result<u64> {
        if bit_count == 0 {
            return Ok(0);
        }
        if bit_count > maximum_count {
            return Err(BitReaderError::TooManyBitsForType);
        }
        let start_position = self.position;
        let end_position = self.position + bit_count as u64;
        if end_position > self.bytes.len() as u64 * 8 {
            return Err(BitReaderError::NotEnoughData);
        }

        let mut value: u64 = 0;

        for i in start_position..end_position {
            let byte_index = (i / 8) as usize;
            let byte = self.bytes[byte_index];
            let shift = 7 - (i % 8);
            let bit = (byte >> shift) as u64 & 1;
            value = (value << 1) | bit;
        }

        self.position = end_position;
        Ok(value)
    }
}

/// Result type for those BitReader operations that can fail.
pub type Result<T> = result::Result<T, BitReaderError>;

/// Error enumeration of BitReader errors.
#[derive(Debug,PartialEq)]
pub enum BitReaderError {
    /// Requested more bits than there are left in the byte slice at the current position.
    NotEnoughData,
    /// Requested more bits than the returned variable can hold, for example more than 8 bits when
    /// reading into a u8.
    TooManyBitsForType
}

impl Error for BitReaderError {
    fn description(&self) -> &str {
        match *self {
            BitReaderError::NotEnoughData => "Requested more bits than the byte slice has left",
            BitReaderError::TooManyBitsForType => "Requested more bits than the requested integer type can hold",
        }
    }
}

impl fmt::Display for BitReaderError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        self.description().fmt(fmt)
    }
}

#[test]
fn read_buffer() {
    let bytes = &[
        0b1011_0101, 0b0110_1010, 0b1010_1100, 0b1001_1001,
        0b1001_1001, 0b1001_1001, 0b1001_1001, 0b1110_0111,
    ];

    let mut reader = BitReader::new(bytes);

    assert_eq!(reader.read_u8(1).unwrap(), 0b1);
    assert_eq!(reader.read_u8(1).unwrap(), 0b0);
    assert_eq!(reader.read_u8(2).unwrap(), 0b11);

    assert_eq!(reader.read_u8(4).unwrap(), 0b0101);

    assert!(reader.is_aligned(1));

    assert_eq!(reader.read_u8(3).unwrap(), 0b11);
    assert_eq!(reader.read_u16(10).unwrap(), 0b01_0101_0101);
    assert_eq!(reader.read_u8(3).unwrap(), 0b100);

    assert!(reader.is_aligned(1));

    assert_eq!(reader.read_u32(32).unwrap(), 0b1001_1001_1001_1001_1001_1001_1001_1001);

    assert_eq!(reader.read_u8(4).unwrap(), 0b1110);
    assert_eq!(reader.read_u8(3).unwrap(), 0b011);
    assert_eq!(reader.read_u8(1).unwrap(), 0b1);

    // Could also be 8 at this point!
    assert!(reader.is_aligned(4));
}

#[test]
fn try_all_sizes() {
    let bytes = &[
        0x4a, 0x1e, 0x39, 0xbb, 0xd0, 0x07, 0xca, 0x9a,
        0xa6, 0xba, 0x25, 0x52, 0x6f, 0x0a, 0x6a, 0xba,
    ];

    let mut reader = BitReader::new(bytes);
    assert_eq!(reader.read_u64(64).unwrap(), 0x4a1e39bbd007ca9a);
    assert_eq!(reader.read_u64(64).unwrap(), 0xa6ba25526f0a6aba);

    let mut reader = BitReader::new(bytes);
    assert_eq!(reader.read_u32(32).unwrap(), 0x4a1e39bb);
    assert_eq!(reader.read_u32(32).unwrap(), 0xd007ca9a);
    assert_eq!(reader.read_u32(32).unwrap(), 0xa6ba2552);
    assert_eq!(reader.read_u32(32).unwrap(), 0x6f0a6aba);

    let mut reader = BitReader::new(bytes);
    assert_eq!(reader.read_u16(16).unwrap(), 0x4a1e);
    assert_eq!(reader.read_u16(16).unwrap(), 0x39bb);
    assert_eq!(reader.read_u16(16).unwrap(), 0xd007);
    assert_eq!(reader.read_u16(16).unwrap(), 0xca9a);
    assert_eq!(reader.read_u16(16).unwrap(), 0xa6ba);
    assert_eq!(reader.read_u16(16).unwrap(), 0x2552);
    assert_eq!(reader.read_u16(16).unwrap(), 0x6f0a);
    assert_eq!(reader.read_u16(16).unwrap(), 0x6aba);

    let mut reader = BitReader::new(&bytes[..]);
    for byte in bytes {
        assert_eq!(reader.read_u8(8).unwrap(), *byte);
    }
}

#[test]
fn skipping_and_zero_reads() {
    let bytes = &[
        0b1011_0101, 0b1110_1010, 0b1010_1100,
    ];

    let mut reader = BitReader::new(bytes);

    assert_eq!(reader.read_u8(4).unwrap(), 0b1011);
    // Reading zero bits should be a no-op
    assert_eq!(reader.read_u8(0).unwrap(), 0b0);
    assert_eq!(reader.read_u8(4).unwrap(), 0b0101);
    reader.skip(3).unwrap(); // 0b111
    assert_eq!(reader.read_u16(10).unwrap(), 0b0101010101);
    assert_eq!(reader.read_u8(3).unwrap(), 0b100);
}

#[test]
fn errors() {
    let bytes = &[
        0b1011_0101, 0b1110_1010, 0b1010_1100,
    ];

    let mut reader = BitReader::new(bytes);
    assert_eq!(reader.read_u8(4).unwrap(), 0b1011);
    assert_eq!(reader.read_u8(9).unwrap_err(), BitReaderError::TooManyBitsForType);
    // If an error happens, it should be possible to resume as if nothing had happened
    assert_eq!(reader.read_u8(4).unwrap(), 0b0101);

    let mut reader = BitReader::new(bytes);
    assert_eq!(reader.read_u8(4).unwrap(), 0b1011);
    // Same with this error
    assert_eq!(reader.read_u32(21).unwrap_err(), BitReaderError::NotEnoughData);
    assert_eq!(reader.read_u8(4).unwrap(), 0b0101);
}

#[test]
fn signed_values() {
    let from = -2048;
    let to = 2048;
    for x in from..to {
        let bytes = &[
            (x >> 8) as u8,
            x as u8,
        ];
        let mut reader = BitReader::new(bytes);
        reader.skip(4).unwrap();
        let val = reader.read_i16(12).unwrap();
        assert_eq!(val, x);
    }
}
<|endoftext|>


<commit_before>//! OpenCC binding for Rust
//!
//! You should install [OpenCC 1.0.x](https://github.com/BYVoid/OpenCC) library first.
//!
//! Supported configurations:
//!
//! * `s2t.json` Simplified Chinese to Traditional Chinese
//! * `t2s.json` Traditional Chinese to Simplified Chinese
//! * `s2tw.json` Simplified Chinese to Traditional Chinese (Taiwan Standard)
//! * `tw2s.json` Traditional Chinese (Taiwan Standard) to Simplified Chinese
//! * `s2hk.json` Simplified Chinese to Traditional Chinese (Hong Kong Standard)
//! * `hk2s.json` Traditional Chinese (Hong Kong Standard) to Simplified Chinese
//! * `s2twp.json` Simplified Chinese to Traditional Chinese (Taiwan Standard) with Taiwanese idiom
//! * `tw2sp.json` Traditional Chinese (Taiwan Standard) to Simplified Chinese with Mainland Chinese idiom
//! * `t2tw.json` Traditional Chinese (OpenCC Standard) to Taiwan Standard
//! * `t2hk.json` Traditional Chinese (OpenCC Standard) to Hong Kong Standard
//!
//! # Examples
//! ```
//! extern crate opencc;
//!
//! use opencc::OpenCC;
//!
//! fn main() {
//!     let cc = OpenCC::new("t2s.json");
//!     println!("{}", cc.convert("乾坤一擲").unwrap());
//!     println!("{}", cc.convert("開放中文轉換").unwrap());
//! }
//! ```
extern crate libc;
use libc::{size_t, c_void, c_char};

use std::ffi::CStr;
use std::str;

#[link(name = "opencc")]
extern {
    fn opencc_open(config: *const c_char) -> *mut c_void;
    fn opencc_close(opencc: *mut c_void);
    fn opencc_convert_utf8(opencc: *mut c_void, text: *const c_char, length: size_t) -> *mut c_char;
    fn opencc_convert_utf8_free(text: *mut c_char);
}

pub struct OpenCC {
    pub config: String,
    libopencc: *mut c_void,
}

impl OpenCC {
    /// Constructs a new `OpenCC`
    ///
    /// # Examples
    /// ```
    /// let cc = opencc::OpenCC::new("t2s.json");
    /// ```
    pub fn new(config: &str) -> OpenCC {
        let _config = config.to_string();
        let config_ptr = _config.as_ptr();
        unsafe {
            OpenCC {
                config: _config,
                libopencc: opencc_open(config_ptr as *const c_char),
            }
        }
    }

    /// Convert a text
    ///
    /// # Examples
    /// ```
    /// let cc = opencc::OpenCC::new("t2s.json");
    /// cc.convert("乾坤一擲").unwrap();
    /// cc.convert("開放中文轉換").unwrap();
    /// ```
    pub fn convert(&self, text: &str) -> Option<String> {
        unsafe {
            let text_ptr = text.as_ptr();
            let c_ptr = opencc_convert_utf8(self.libopencc, text_ptr as *const c_char, text.len() as size_t);
            if c_ptr.is_null() {
                None
            } else {
                let c_str = CStr::from_ptr(c_ptr);
                let str_buf: String = str::from_utf8(c_str.to_bytes()).unwrap().to_owned();
                opencc_convert_utf8_free(c_ptr);
                Some(str_buf)
            }
        }
    }

    /// Close the underlying libopencc.
    /// Will be called automatically when the variable gets out of scope, should not call it
    /// directly by yourself.
    pub fn close(&self) {
        unsafe {
            opencc_close(self.libopencc);
        }
    }
}

impl Drop for OpenCC {
    /// Close the underlying libopencc when it been droped
    fn drop(&mut self) {
        self.close();
    }
}

#[test]
fn test_convert() {
    let opencc = OpenCC::new("t2s.json");
    assert_eq!("乾坤一掷".to_string(), opencc.convert("乾坤一擲").unwrap());
    assert_eq!("开放中文转换".to_string(), opencc.convert("開放中文轉換").unwrap());
}
<commit_msg>Add is_closed function<commit_after>//! OpenCC binding for Rust
//!
//! You should install [OpenCC 1.0.x](https://github.com/BYVoid/OpenCC) library first.
//!
//! Supported configurations:
//!
//! * `s2t.json` Simplified Chinese to Traditional Chinese
//! * `t2s.json` Traditional Chinese to Simplified Chinese
//! * `s2tw.json` Simplified Chinese to Traditional Chinese (Taiwan Standard)
//! * `tw2s.json` Traditional Chinese (Taiwan Standard) to Simplified Chinese
//! * `s2hk.json` Simplified Chinese to Traditional Chinese (Hong Kong Standard)
//! * `hk2s.json` Traditional Chinese (Hong Kong Standard) to Simplified Chinese
//! * `s2twp.json` Simplified Chinese to Traditional Chinese (Taiwan Standard) with Taiwanese idiom
//! * `tw2sp.json` Traditional Chinese (Taiwan Standard) to Simplified Chinese with Mainland Chinese idiom
//! * `t2tw.json` Traditional Chinese (OpenCC Standard) to Taiwan Standard
//! * `t2hk.json` Traditional Chinese (OpenCC Standard) to Hong Kong Standard
//!
//! # Examples
//! ```
//! extern crate opencc;
//!
//! use opencc::OpenCC;
//!
//! fn main() {
//!     let cc = OpenCC::new("t2s.json");
//!     println!("{}", cc.convert("乾坤一擲").unwrap());
//!     println!("{}", cc.convert("開放中文轉換").unwrap());
//! }
//! ```
extern crate libc;
use libc::{size_t, c_void, c_char};

use std::ffi::CStr;
use std::str;

#[link(name = "opencc")]
extern {
    fn opencc_open(config: *const c_char) -> *mut c_void;
    fn opencc_close(opencc: *mut c_void);
    fn opencc_convert_utf8(opencc: *mut c_void, text: *const c_char, length: size_t) -> *mut c_char;
    fn opencc_convert_utf8_free(text: *mut c_char);
}

pub struct OpenCC {
    /// Configuration file
    pub config: String,
    _is_closed: bool,
    libopencc: *mut c_void,
}

impl OpenCC {
    /// Constructs a new `OpenCC`
    ///
    /// # Examples
    /// ```
    /// let cc = opencc::OpenCC::new("t2s.json");
    /// ```
    pub fn new(config: &str) -> OpenCC {
        let _config = config.to_string();
        let config_ptr = _config.as_ptr();
        unsafe {
            OpenCC {
                config: _config,
                _is_closed: false,
                libopencc: opencc_open(config_ptr as *const c_char),
            }
        }
    }

    /// Convert a text
    ///
    /// # Examples
    /// ```
    /// let cc = opencc::OpenCC::new("t2s.json");
    /// cc.convert("乾坤一擲").unwrap();
    /// cc.convert("開放中文轉換").unwrap();
    /// ```
    pub fn convert(&self, text: &str) -> Option<String> {
        if self._is_closed {
            None
        } else {
            unsafe {
                let text_ptr = text.as_ptr();
                let c_ptr = opencc_convert_utf8(self.libopencc, text_ptr as *const c_char, text.len() as size_t);
                if c_ptr.is_null() {
                    None
                } else {
                    let c_str = CStr::from_ptr(c_ptr);
                    let str_buf: String = str::from_utf8(c_str.to_bytes()).unwrap().to_owned();
                    opencc_convert_utf8_free(c_ptr);
                    Some(str_buf)
                }
            }
        }
    }

    /// Close the underlying libopencc.
    /// Will be called automatically when the variable gets out of scope.
    pub fn close(&mut self) {
        if !self._is_closed {
            unsafe {
                opencc_close(self.libopencc);
            }
            self._is_closed = true;
        }
    }

    /// Is the underlying libopencc closed
    pub fn is_closed(&self) -> bool {
        self._is_closed
    }
}

impl Drop for OpenCC {
    /// Close the underlying libopencc when it been droped
    fn drop(&mut self) {
        self.close();
    }
}

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

    #[test]
    fn test_simple_convert() {
        let cc = OpenCC::new("t2s.json");
        assert_eq!("乾坤一掷".to_string(), cc.convert("乾坤一擲").unwrap());
        assert_eq!("开放中文转换".to_string(), cc.convert("開放中文轉換").unwrap());
    }

    #[test]
    fn test_opencc_close() {
        let mut cc = OpenCC::new("t2s.json");
        assert_eq!("乾坤一掷".to_string(), cc.convert("乾坤一擲").unwrap());
        assert_eq!(false, cc.is_closed());
        cc.close();
        assert_eq!(true, cc.convert("開放中文轉換").is_none());
        assert_eq!(true, cc.is_closed());
    }
}
<|endoftext|>


<commit_before>#![recursion_limit = "1024"]
#[macro_use]
extern crate error_chain;

extern crate rustc_serialize;

use std::net::TcpStream;
use std::io::prelude::*;
use std::collections::HashMap;
use rustc_serialize::json;

pub mod errors;
use errors::*;

pub struct WhoIs<'a> {
    server: &'a str,
    follow: isize,
    new_whois: String,
    query: String,
}

impl<'a> WhoIs<'a> {
    pub fn new(x: &'a str) -> WhoIs<'a> {
        WhoIs {
            server: x,
            follow: 0,
            new_whois: String::new(),
            query: String::new(),
        }
    }
    ///
    ///
    ///  This function will get whois server from the `get_server` function decide the appropriate
    ///  query for the server and parse the whois data into JSON by calling parse_data()
    ///  If there is another whois server in the whois data then it calls 'parse_whois'
    ///  so it can get the whois data from that
    ///
    pub fn lookup(&mut self) -> Result<String> {
        let mut result = String::new();
        let mut server = self.new_whois.to_owned();
        let target = self.server.to_owned();
        let tld = match target.split(".").last() {
            Some(tld) => tld,
            None => return Err("Invalid URL?".into()),
        };
        if self.follow == 0 {
            self.query = match tld {
                "com" | "net" => "DOMAIN ".into(),
                _ => "".into(),
            };
            server = self.get_server(&tld).expect(&format!("Failed to get server for {}", tld));
        }
        let mut client = try!(TcpStream::connect((&*server, 43u16))
                                  .chain_err(|| "Could not connect to server!!"));

        try!(client.write_all(format!("{}{}\n", self.query, target).as_bytes())
                   .chain_err(|| "Could not write to client {}"));

        try!(client.read_to_string(&mut result)
                   .chain_err(|| "Failed to read to string"));

        if result.contains("Whois Server:") {
            self.query = "".into();
            self.follow += 1;                                             // If there is another Whois Server, take that server and pass it to
            Ok(self.parse_whois(&*result))                             // pass it to parse_whois
        } else {
            let clean = result.replace("http:", "").replace("https:", ""); // I'm splitting via ':' so the urls needs to be omitted
            self.parse_data(clean)
        }
    }
    fn parse_data(&self, result: String) -> Result<String> {
        let mut data = HashMap::new();
        for c in result.lines() {
            let mut line = c.split(':');
            let key = line.next().expect("Failed to get key");
            let value = match line.next() {
                Some(value) => value,
                None => continue,
            };
            data.insert(key, value.trim());
        }
        json::encode(&data).chain_err(|| "Could not encode data as json")
    }
    ///
    /// This function calls lookup() again if there is a another whois server
    ///
    ///
    fn parse_whois(&mut self, result: &str) -> String {
        let line = &result.lines()
                          .find(|i| i.contains("Whois Server:"))
                          .expect("Could not find wh");
        let target = line.split_whitespace().last().unwrap().to_owned();
        self.new_whois = target;
        self.lookup().expect("Failed lookup in parse_whois")
    }

    fn get_server(&self, target: &str) -> Result<String> {
        let mut result = String::new();
        let mut client = try!(TcpStream::connect("whois.iana.org:43")
                                  .chain_err(|| "Could not connect to server!"));

        try!(client.write_all(format!("{}\n", target).as_bytes())
                   .chain_err(|| "Could not write to client"));

        try!(client.read_to_string(&mut result).chain_err(|| "Failed to read result to string"));
        // println!("{:?}", result);
        let line = &result.lines().find(|i| i.starts_with("whois:")).expect("Couldnt get wh");
        let foo = line.split_whitespace().last().unwrap().to_owned();
        Ok(foo)
    }
}
<commit_msg>Minor comments changes<commit_after>#![recursion_limit = "1024"]
#[macro_use]
extern crate error_chain;

extern crate rustc_serialize;

use std::net::TcpStream;
use std::io::prelude::*;
use std::collections::HashMap;
use rustc_serialize::json;

pub mod errors;
use errors::*;

pub struct WhoIs<'a> {
    server: &'a str,
    follow: isize,
    new_whois: String,
    query: String,
}

impl<'a> WhoIs<'a> {
    pub fn new(x: &'a str) -> WhoIs<'a> {
        WhoIs {
            server: x,
            follow: 0,
            new_whois: String::new(),
            query: String::new(),
        }
    }
    ///  This function will get whois server from the `get_server` function, decide the appropriate
    ///  query for the server and parse the whois data into JSON by calling `parse_data()`.
    ///
    ///  If there is another whois server in the whois data then it calls `parse_whois`,
    ///  so it can get the whois data from that.
    pub fn lookup(&mut self) -> Result<String> {
        let mut result = String::new();
        let mut server = self.new_whois.to_owned();
        let target = self.server.to_owned();
        let tld = match target.split(".").last() {
            Some(tld) => tld,
            None => return Err("Invalid URL?".into()),
        };
        if self.follow == 0 {
            self.query = match tld {
                "com" | "net" => "DOMAIN ".into(),
                _ => "".into(),
            };
            server = self.get_server(&tld).expect(&format!("Failed to get server for {}", tld));
        }
        let mut client = try!(TcpStream::connect((&*server, 43u16))
                                  .chain_err(|| "Could not connect to server!!"));

        try!(client.write_all(format!("{}{}\n", self.query, target).as_bytes())
                   .chain_err(|| "Could not write to client {}"));

        try!(client.read_to_string(&mut result)
                   .chain_err(|| "Failed to read to string"));

        if result.contains("Whois Server:") {
            self.query = "".into();
            self.follow += 1; // If there is another Whois Server, take that server and pass it to.
            Ok(self.parse_whois(&*result)) // Pass it to parse_whois.
        } else {
            let clean = result.replace("http:", "").replace("https:", ""); // I'm splitting via ':' so the urls' protocols needs to be omitted.
            self.parse_data(clean)
        }
    }
    fn parse_data(&self, result: String) -> Result<String> {
        let mut data = HashMap::new();
        for c in result.lines() {
            let mut line = c.split(':');
            let key = line.next().expect("Failed to get key");
            let value = match line.next() {
                Some(value) => value,
                None => continue,
            };
            data.insert(key, value.trim());
        }
        json::encode(&data).chain_err(|| "Could not encode data as json")
    }

    /// This function calls `lookup()` again if there is a another whois server.
    fn parse_whois(&mut self, result: &str) -> String {
        let line = &result.lines()
                          .find(|i| i.contains("Whois Server:"))
                          .expect("Could not find wh");
        let target = line.split_whitespace().last().unwrap().to_owned();
        self.new_whois = target;
        self.lookup().expect("Failed lookup in parse_whois")
    }

    fn get_server(&self, target: &str) -> Result<String> {
        let mut result = String::new();
        let mut client = try!(TcpStream::connect("whois.iana.org:43")
                                  .chain_err(|| "Could not connect to server!"));

        try!(client.write_all(format!("{}\n", target).as_bytes())
                   .chain_err(|| "Could not write to client"));

        try!(client.read_to_string(&mut result).chain_err(|| "Failed to read result to string"));
        // println!("{:?}", result);
        let line = &result.lines().find(|i| i.starts_with("whois:")).expect("Couldnt get wh");
        let foo = line.split_whitespace().last().unwrap().to_owned();
        Ok(foo)
    }
}
<|endoftext|>