<commit_before>use std::collections::HashMap;
use std::fs::File;
use std::io::Read;

use crate::stacked_borrows::Tag;
use crate::*;

pub struct FileHandler {
    files: HashMap<i32, File>,
    flags: HashMap<i32, i32>,
    low: i32,
}

impl Default for FileHandler {
    fn default() -> Self {
        FileHandler {
            files: Default::default(),
            flags: Default::default(),
            // 0, 1 and 2 are reserved for stdin, stdout and stderr
            low: 3,
        }
    }
}

impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    fn open(
        &mut self,
        path_op: OpTy<'tcx, Tag>,
        flag_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`open` not available when isolation is enabled")
        }

        let flag = this.read_scalar(flag_op)?.to_i32()?;

        if flag != this.eval_libc_i32("O_RDONLY")? && flag != this.eval_libc_i32("O_CLOEXEC")? {
            throw_unsup_format!("Unsupported flag {:#x}", flag);
        }

        let path_bytes = this
            .memory()
            .read_c_str(this.read_scalar(path_op)?.not_undef()?)?;
        let path = std::str::from_utf8(path_bytes)
            .map_err(|_| err_unsup_format!("{:?} is not a valid utf-8 string", path_bytes))?;

        match File::open(path) {
            Ok(file) => {
                let mut fh = &mut this.machine.file_handler;
                fh.low += 1;
                fh.files.insert(fh.low, file);
                fh.flags.insert(fh.low, flag);
                Ok(fh.low)
            }

            Err(e) => {
                this.machine.last_error = e.raw_os_error().unwrap() as u32;
                Ok(-1)
            }
        }
    }

    fn fcntl(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        cmd_op: OpTy<'tcx, Tag>,
        arg_op: Option<OpTy<'tcx, Tag>>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`open` not available when isolation is enabled")
        }

        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let cmd = this.read_scalar(cmd_op)?.to_i32()?;

        if cmd == this.eval_libc_i32("F_SETFD")? {
            // This does not affect the file itself. Certain flags might require changing the file
            // or the way it is accessed somehow.
            let flag = this.read_scalar(arg_op.unwrap())?.to_i32()?;
            // The only usage of this in stdlib at the moment is to enable the `FD_CLOEXEC` flag.
            let fd_cloexec = this.eval_libc_i32("FD_CLOEXEC")?;
            if let Some(old_flag) = this.machine.file_handler.flags.get_mut(&fd) {
                if flag ^ *old_flag == fd_cloexec {
                    *old_flag = flag;
                } else {
                    throw_unsup_format!("Unsupported arg {:#x} for `F_SETFD`", flag);
                }
            }
            Ok(0)
        } else if cmd == this.eval_libc_i32("F_GETFD")? {
            if let Some(flag) = this.machine.file_handler.flags.get(&fd) {
                Ok(*flag)
            } else {
                this.machine.last_error = this.eval_libc_i32("EBADF")? as u32;
                Ok(-1)
            }
        } else {
            throw_unsup_format!("Unsupported command {:#x}", cmd);
        }
    }

    fn close(&mut self, fd_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`open` not available when isolation is enabled")
        }

        let fd = this.read_scalar(fd_op)?.to_i32()?;

        if let Some(file) = this.machine.file_handler.files.remove(&fd) {
            match file.sync_all() {
                Ok(()) => Ok(0),
                Err(e) => {
                    this.machine.last_error = e.raw_os_error().unwrap() as u32;
                    Ok(-1)
                }
            }
        } else {
            this.machine.last_error = this.eval_libc_i32("EBADF")? as u32;
            Ok(-1)
        }
    }

    fn read(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        buf_op: OpTy<'tcx, Tag>,
        count_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i64> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`open` not available when isolation is enabled")
        }

        let tcx = &{ this.tcx.tcx };

        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let buf = this.force_ptr(this.read_scalar(buf_op)?.not_undef()?)?;
        let count = this.read_scalar(count_op)?.to_usize(&*this.tcx)?;

        if let Some(file) = this.machine.file_handler.files.get_mut(&fd) {
            let mut bytes = vec![0; count as usize];
            match file.read(&mut bytes) {
                Ok(read_bytes) => {
                    bytes.truncate(read_bytes);

                    this.memory_mut()
                        .get_mut(buf.alloc_id)?
                        .write_bytes(tcx, buf, &bytes)?;

                    Ok(read_bytes as i64)
                }
                Err(e) => {
                    this.machine.last_error = e.raw_os_error().unwrap() as u32;
                    Ok(-1)
                }
            }
        } else {
            this.machine.last_error = this.eval_libc_i32("EBADF")? as u32;
            Ok(-1)
        }
    }
}
<commit_msg>Add FileHandle struct<commit_after>use std::collections::HashMap;
use std::fs::File;
use std::io::Read;

use crate::stacked_borrows::Tag;
use crate::*;

struct FileHandle {
    file: File,
    flag: i32,
}

pub struct FileHandler {
    handles: HashMap<i32, FileHandle>,
    low: i32,
}

impl Default for FileHandler {
    fn default() -> Self {
        FileHandler {
            handles: Default::default(),
            // 0, 1 and 2 are reserved for stdin, stdout and stderr
            low: 3,
        }
    }
}

impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    fn open(
        &mut self,
        path_op: OpTy<'tcx, Tag>,
        flag_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`open` not available when isolation is enabled")
        }

        let flag = this.read_scalar(flag_op)?.to_i32()?;

        if flag != this.eval_libc_i32("O_RDONLY")? && flag != this.eval_libc_i32("O_CLOEXEC")? {
            throw_unsup_format!("Unsupported flag {:#x}", flag);
        }

        let path_bytes = this
            .memory()
            .read_c_str(this.read_scalar(path_op)?.not_undef()?)?;
        let path = std::str::from_utf8(path_bytes)
            .map_err(|_| err_unsup_format!("{:?} is not a valid utf-8 string", path_bytes))?;

        match File::open(path) {
            Ok(file) => {
                let mut fh = &mut this.machine.file_handler;
                fh.low += 1;
                fh.handles.insert(fh.low, FileHandle{ file, flag});
                Ok(fh.low)
            }

            Err(e) => {
                this.machine.last_error = e.raw_os_error().unwrap() as u32;
                Ok(-1)
            }
        }
    }

    fn fcntl(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        cmd_op: OpTy<'tcx, Tag>,
        arg_op: Option<OpTy<'tcx, Tag>>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`open` not available when isolation is enabled")
        }

        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let cmd = this.read_scalar(cmd_op)?.to_i32()?;

        if cmd == this.eval_libc_i32("F_SETFD")? {
            // This does not affect the file itself. Certain flags might require changing the file
            // or the way it is accessed somehow.
            let flag = this.read_scalar(arg_op.unwrap())?.to_i32()?;
            // The only usage of this in stdlib at the moment is to enable the `FD_CLOEXEC` flag.
            let fd_cloexec = this.eval_libc_i32("FD_CLOEXEC")?;
            if let Some(FileHandle{ flag: old_flag, .. }) = this.machine.file_handler.handles.get_mut(&fd) {
                if flag ^ *old_flag == fd_cloexec {
                    *old_flag = flag;
                } else {
                    throw_unsup_format!("Unsupported arg {:#x} for `F_SETFD`", flag);
                }
            }
            Ok(0)
        } else if cmd == this.eval_libc_i32("F_GETFD")? {
            if let Some(handle) = this.machine.file_handler.handles.get(&fd) {
                Ok(handle.flag)
            } else {
                this.machine.last_error = this.eval_libc_i32("EBADF")? as u32;
                Ok(-1)
            }
        } else {
            throw_unsup_format!("Unsupported command {:#x}", cmd);
        }
    }

    fn close(&mut self, fd_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`open` not available when isolation is enabled")
        }

        let fd = this.read_scalar(fd_op)?.to_i32()?;

        if let Some(handle) = this.machine.file_handler.handles.remove(&fd) {
            match handle.file.sync_all() {
                Ok(()) => Ok(0),
                Err(e) => {
                    this.machine.last_error = e.raw_os_error().unwrap() as u32;
                    Ok(-1)
                }
            }
        } else {
            this.machine.last_error = this.eval_libc_i32("EBADF")? as u32;
            Ok(-1)
        }
    }

    fn read(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        buf_op: OpTy<'tcx, Tag>,
        count_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i64> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`open` not available when isolation is enabled")
        }

        let tcx = &{ this.tcx.tcx };

        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let buf = this.force_ptr(this.read_scalar(buf_op)?.not_undef()?)?;
        let count = this.read_scalar(count_op)?.to_usize(&*this.tcx)?;

        if let Some(FileHandle { file, ..}) = this.machine.file_handler.handles.get_mut(&fd) {
            let mut bytes = vec![0; count as usize];
            match file.read(&mut bytes) {
                Ok(read_bytes) => {
                    bytes.truncate(read_bytes);

                    this.memory_mut()
                        .get_mut(buf.alloc_id)?
                        .write_bytes(tcx, buf, &bytes)?;

                    Ok(read_bytes as i64)
                }
                Err(e) => {
                    this.machine.last_error = e.raw_os_error().unwrap() as u32;
                    Ok(-1)
                }
            }
        } else {
            this.machine.last_error = this.eval_libc_i32("EBADF")? as u32;
            Ok(-1)
        }
    }
}
<|endoftext|>


<commit_before>use std::cmp::Ordering;
use std::fmt;
use std::hash::{self, Hash};
use std::marker::PhantomData;
use std::mem::{self, ManuallyDrop};
use std::ops;
use std::ptr;

#[cfg(not(feature = "std"))]
use alloc::alloc::{self, Layout};
#[cfg(feature = "std")]
use std::alloc::{self, Layout};

#[cfg(feature = "coerce")]
use std::marker::Unsize;
#[cfg(feature = "coerce")]
use std::ops::CoerceUnsized;

#[cfg(feature = "coerce")]
impl<T: ?Sized + Unsize<U>, U: ?Sized, Space> CoerceUnsized<SmallBox<U, Space>>
    for SmallBox<T, Space>
{}

/// Box value on stack or heap depending on its size
///
/// This macro is similar to `SmallBox::new`, which is used to create a new `Smallbox` instance,
/// but relaxing the constraint of `T: Sized`.
/// In order to do that, this macro will check the coersion rules between type `T` and the
/// expression type. This macro invokes a complier error for any invalid type coersion.
///
/// You can think that it has the signature of `smallbox!<U: Sized, T: ?Sized>(val: U) -> SmallBox<T, Space>`
///
/// # Example
///
/// ```
/// #[macro_use]
/// extern crate smallbox;
///
/// # fn main() {
/// use smallbox::SmallBox;
/// use smallbox::space::*;
///
/// let small: SmallBox<[usize], S4> = smallbox!([0usize; 2]);
/// let large: SmallBox<[usize], S4> = smallbox!([1usize; 8]);
///
/// assert_eq!(small.len(), 2);
/// assert_eq!(large[7], 1);
///
/// assert!(large.is_heap() == true);
/// # }
/// ```
#[macro_export]
macro_rules! smallbox {
    ( $e: expr ) => {{
        let val = $e;
        let ptr = &val as *const _;
        #[allow(unsafe_code)]
        unsafe {
            $crate::SmallBox::new_unchecked(val, ptr)
        }
    }};
}

/// An optimized box that store value on stack or heap depending on its size
pub struct SmallBox<T: ?Sized, Space> {
    space: ManuallyDrop<Space>,
    ptr: *const T,
    _phantom: PhantomData<T>,
}

impl<T: ?Sized, Space> SmallBox<T, Space> {
    /// Box value on stack or heap depending on its size.
    ///
    /// # Example
    ///
    /// ```
    /// use smallbox::SmallBox;
    /// use smallbox::space::*;
    ///
    /// let small: SmallBox<_, S4> = SmallBox::new([0usize; 2]);
    /// let large: SmallBox<_, S4> = SmallBox::new([1usize; 8]);
    ///
    /// assert_eq!(small.len(), 2);
    /// assert_eq!(large[7], 1);
    ///
    /// assert!(large.is_heap() == true);
    /// ```
    pub fn new(val: T) -> SmallBox<T, Space>
    where
        T: Sized,
    {
        smallbox!(val)
    }

    #[doc(hidden)]
    pub unsafe fn new_unchecked<U>(val: U, ptr: *const T) -> SmallBox<T, Space>
    where
        U: Sized,
    {
        let mut space = ManuallyDrop::new(mem::uninitialized::<Space>());

        let (ptr_addr, ptr_copy): (*const u8, *mut U) = if mem::size_of::<U>()
            > mem::size_of::<Space>()
            || mem::align_of::<U>() > mem::align_of::<Space>()
        {
            // Heap
            let layout = Layout::new::<U>();
            let heap_ptr = alloc::alloc(layout);

            (heap_ptr, heap_ptr as *mut U)
        } else {
            // Stack
            (ptr::null(), mem::transmute(&mut space))
        };

        let mut ptr = ptr;
        let ptr_ptr = &mut ptr as *mut _ as *mut usize;
        ptr_ptr.write(ptr_addr as usize);

        ptr::copy_nonoverlapping(&val, ptr_copy, 1);

        mem::forget(val);

        SmallBox {
            space,
            ptr,
            _phantom: PhantomData,
        }
    }

    /// Change the capacity of `SmallBox`
    ///
    /// This method may move stack-allocated data to heap
    /// if the inline space is not sufficient. Once the data
    /// is stored on heap, it'll never be moved again.
    ///
    /// # Example
    ///
    /// ```
    /// use smallbox::SmallBox;
    /// use smallbox::space::{S2, S4};
    ///
    /// let s: SmallBox::<_, S4> = SmallBox::new([0usize; 4]);
    /// let m: SmallBox::<_, S2> = s.resize();
    /// ```
    pub fn resize<ToSpace>(self) -> SmallBox<T, ToSpace> {
        unsafe {
            let mut space = ManuallyDrop::new(mem::uninitialized::<ToSpace>());

            let ptr = if self.is_heap() {
                // don't touch anything if the data is already on heap
                self.ptr
            } else {
                let ptr: &T = &*self;
                // original data is on stack
                let (ptr, ptr_copy): (*const T, *mut u8) = if mem::size_of_val::<T>(ptr)
                    > mem::size_of::<ToSpace>()
                    || mem::align_of_val::<T>(ptr) > mem::align_of::<ToSpace>()
                {
                    // but we have to move it to heap
                    let layout = Layout::for_value::<T>(ptr);
                    let heap_ptr = alloc::alloc(layout) as *mut u8;

                    let mut ptr: *const T = ptr;
                    let ptr_ptr = &mut ptr as *mut _ as *mut usize;
                    ptr_ptr.write(heap_ptr as usize);

                    (ptr, heap_ptr)
                } else {
                    // still store it on stack
                    (self.ptr, mem::transmute(&mut space))
                };

                ptr::copy_nonoverlapping(
                    &self.space as *const _ as *const u8,
                    ptr_copy,
                    mem::size_of_val::<T>(&*self),
                );

                ptr
            };

            mem::forget(self);

            SmallBox {
                space,
                ptr,
                _phantom: PhantomData,
            }
        }
    }

    /// Returns true if the data is heap-allocated
    pub fn is_heap(&self) -> bool {
        !self.ptr.is_null()
    }

    unsafe fn as_ptr(&self) -> *const T {
        let mut ptr = self.ptr;

        if !self.is_heap() {
            let ptr_ptr = &mut ptr as *mut _ as *mut usize;
            ptr_ptr.write(mem::transmute(&self.space));
        }

        ptr
    }
}

impl<T: ?Sized, Space> ops::Deref for SmallBox<T, Space> {
    type Target = T;

    fn deref(&self) -> &T {
        unsafe { &*self.as_ptr() }
    }
}

impl<T: ?Sized, Space> ops::DerefMut for SmallBox<T, Space> {
    fn deref_mut(&mut self) -> &mut T {
        unsafe { &mut *(self.as_ptr() as *const _ as *mut _) }
    }
}

impl<T: ?Sized, Space> ops::Drop for SmallBox<T, Space> {
    fn drop(&mut self) {
        unsafe {
            let layout = Layout::for_value::<T>(&*self);
            ptr::drop_in_place::<T>(&mut **self);
            if self.is_heap() {
                alloc::dealloc(self.ptr as *mut u8, layout);
            }
        }
    }
}

impl<T: ?Sized + fmt::Display, Space> fmt::Display for SmallBox<T, Space> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&**self, f)
    }
}

impl<T: ?Sized + fmt::Debug, Space> fmt::Debug for SmallBox<T, Space> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&**self, f)
    }
}

impl<T: ?Sized, Space> fmt::Pointer for SmallBox<T, Space> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // It's not possible to extract the inner Unique directly from the Box,
        // instead we cast it to a *const which aliases the Unique
        let ptr: *const T = &**self;
        fmt::Pointer::fmt(&ptr, f)
    }
}

impl<T: ?Sized + PartialEq, Space> PartialEq for SmallBox<T, Space> {
    #[inline]
    fn eq(&self, other: &SmallBox<T, Space>) -> bool {
        PartialEq::eq(&**self, &**other)
    }
    #[inline]
    fn ne(&self, other: &SmallBox<T, Space>) -> bool {
        PartialEq::ne(&**self, &**other)
    }
}

impl<T: ?Sized + PartialOrd, Space> PartialOrd for SmallBox<T, Space> {
    #[inline]
    fn partial_cmp(&self, other: &SmallBox<T, Space>) -> Option<Ordering> {
        PartialOrd::partial_cmp(&**self, &**other)
    }
    #[inline]
    fn lt(&self, other: &SmallBox<T, Space>) -> bool {
        PartialOrd::lt(&**self, &**other)
    }
    #[inline]
    fn le(&self, other: &SmallBox<T, Space>) -> bool {
        PartialOrd::le(&**self, &**other)
    }
    #[inline]
    fn ge(&self, other: &SmallBox<T, Space>) -> bool {
        PartialOrd::ge(&**self, &**other)
    }
    #[inline]
    fn gt(&self, other: &SmallBox<T, Space>) -> bool {
        PartialOrd::gt(&**self, &**other)
    }
}

impl<T: ?Sized + Ord, Space> Ord for SmallBox<T, Space> {
    #[inline]
    fn cmp(&self, other: &SmallBox<T, Space>) -> Ordering {
        Ord::cmp(&**self, &**other)
    }
}

impl<T: ?Sized + Eq, Space> Eq for SmallBox<T, Space> {}

impl<T: ?Sized + Hash, Space> Hash for SmallBox<T, Space> {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        (**self).hash(state);
    }
}

unsafe impl<T: ?Sized + Send, Space> Send for SmallBox<T, Space> {}
unsafe impl<T: ?Sized + Sync, Space> Sync for SmallBox<T, Space> {}

#[cfg(test)]
mod tests {
    use super::SmallBox;
    use space::*;
    use std::any::Any;

    #[test]
    fn test_basic() {
        let stacked: SmallBox<usize, S1> = SmallBox::new(1234usize);
        assert!(*stacked == 1234);

        let heaped: SmallBox<(usize, usize), S1> = SmallBox::new((0, 1));
        assert!(*heaped == (0, 1));
    }

    #[test]
    fn test_new_unchecked() {
        let val = [0usize, 1];
        let ptr = &val as *const _;

        unsafe {
            let stacked: SmallBox<[usize], S2> = SmallBox::new_unchecked(val, ptr);
            assert!(*stacked == [0, 1]);
            assert!(!stacked.is_heap());
        }

        let val = [0usize, 1, 2];
        let ptr = &val as *const _;

        unsafe {
            let heaped: SmallBox<Any, S2> = SmallBox::new_unchecked(val, ptr);
            assert!(heaped.is_heap());

            if let Some(array) = heaped.downcast_ref::<[usize; 3]>() {
                assert_eq!(*array, [0, 1, 2]);
            } else {
                unreachable!();
            }
        }
    }

    #[test]
    #[deny(unsafe_code)]
    fn test_macro() {
        let stacked: SmallBox<Any, S1> = smallbox!(1234usize);
        if let Some(num) = stacked.downcast_ref::<usize>() {
            assert_eq!(*num, 1234);
        } else {
            unreachable!();
        }

        let heaped: SmallBox<Any, S1> = smallbox!([0usize, 1]);
        if let Some(array) = heaped.downcast_ref::<[usize; 2]>() {
            assert_eq!(*array, [0, 1]);
        } else {
            unreachable!();
        }

        let is_even: SmallBox<Fn(u8) -> bool, S1> = smallbox!(|num: u8| num % 2 == 0);
        assert!(!is_even(5));
        assert!(is_even(6));
    }

    #[test]
    #[cfg(feature = "coerce")]
    fn test_coerce() {
        let stacked: SmallBox<Any, S1> = SmallBox::new(1234usize);
        if let Some(num) = stacked.downcast_ref::<usize>() {
            assert_eq!(*num, 1234);
        } else {
            unreachable!();
        }

        let heaped: SmallBox<Any, S1> = SmallBox::new([0usize, 1]);
        if let Some(array) = heaped.downcast_ref::<[usize; 2]>() {
            assert_eq!(*array, [0, 1]);
        } else {
            unreachable!();
        }
    }

    #[test]
    fn test_drop() {
        use std::cell::Cell;

        #[derive(Debug)]
        struct Struct<'a>(&'a Cell<bool>);
        impl<'a> Drop for Struct<'a> {
            fn drop(&mut self) {
                self.0.set(true);
            }
        }

        let flag = Cell::new(false);

        let val: SmallBox<_, S2> = SmallBox::new(Struct(&flag));

        assert!(flag.get() == false);
        drop(val);
        assert!(flag.get() == true);
    }

    #[test]
    fn test_dont_drop_space() {
        struct NoDrop(S1);
        impl Drop for NoDrop {
            fn drop(&mut self) {
                unreachable!();
            }
        }

        drop(SmallBox::<_, NoDrop>::new([true]));
    }

    #[test]
    fn test_oversize() {
        let fit = SmallBox::<_, S1>::new([0usize; 1]);
        let oversize = SmallBox::<_, S1>::new([0usize; 2]);
        assert!(!fit.is_heap());
        assert!(oversize.is_heap());
    }

    #[test]
    fn test_resize() {
        let m = SmallBox::<_, S4>::new([0usize; 2]);
        let l = m.resize::<S8>();
        assert!(!l.is_heap());
        let m = l.resize::<S4>();
        assert!(!m.is_heap());
        let s = m.resize::<S2>();
        assert!(!s.is_heap());
        let xs = s.resize::<S1>();
        assert!(xs.is_heap());
        let m = xs.resize::<S4>();
        assert!(m.is_heap());
    }

    #[test]
    fn test_zst() {
        let zst = SmallBox::<_, S1>::new([0usize; 0]);
        assert_eq!(*zst, [0usize; 0]);
    }
}
<commit_msg>impl clone<commit_after>use std::cmp::Ordering;
use std::fmt;
use std::hash::{self, Hash};
use std::marker::PhantomData;
use std::mem::{self, ManuallyDrop};
use std::ops;
use std::ptr;

#[cfg(not(feature = "std"))]
use alloc::alloc::{self, Layout};
#[cfg(feature = "std")]
use std::alloc::{self, Layout};

#[cfg(feature = "coerce")]
use std::marker::Unsize;
#[cfg(feature = "coerce")]
use std::ops::CoerceUnsized;

#[cfg(feature = "coerce")]
impl<T: ?Sized + Unsize<U>, U: ?Sized, Space> CoerceUnsized<SmallBox<U, Space>>
    for SmallBox<T, Space>
{}

/// Box value on stack or heap depending on its size
///
/// This macro is similar to `SmallBox::new`, which is used to create a new `Smallbox` instance,
/// but relaxing the constraint of `T: Sized`.
/// In order to do that, this macro will check the coersion rules between type `T` and the
/// expression type. This macro invokes a complier error for any invalid type coersion.
///
/// You can think that it has the signature of `smallbox!<U: Sized, T: ?Sized>(val: U) -> SmallBox<T, Space>`
///
/// # Example
///
/// ```
/// #[macro_use]
/// extern crate smallbox;
///
/// # fn main() {
/// use smallbox::SmallBox;
/// use smallbox::space::*;
///
/// let small: SmallBox<[usize], S4> = smallbox!([0usize; 2]);
/// let large: SmallBox<[usize], S4> = smallbox!([1usize; 8]);
///
/// assert_eq!(small.len(), 2);
/// assert_eq!(large[7], 1);
///
/// assert!(large.is_heap() == true);
/// # }
/// ```
#[macro_export]
macro_rules! smallbox {
    ( $e: expr ) => {{
        let val = $e;
        let ptr = &val as *const _;
        #[allow(unsafe_code)]
        unsafe {
            $crate::SmallBox::new_unchecked(val, ptr)
        }
    }};
}

/// An optimized box that store value on stack or heap depending on its size
pub struct SmallBox<T: ?Sized, Space> {
    space: ManuallyDrop<Space>,
    ptr: *const T,
    _phantom: PhantomData<T>,
}

impl<T: ?Sized, Space> SmallBox<T, Space> {
    /// Box value on stack or heap depending on its size.
    ///
    /// # Example
    ///
    /// ```
    /// use smallbox::SmallBox;
    /// use smallbox::space::*;
    ///
    /// let small: SmallBox<_, S4> = SmallBox::new([0usize; 2]);
    /// let large: SmallBox<_, S4> = SmallBox::new([1usize; 8]);
    ///
    /// assert_eq!(small.len(), 2);
    /// assert_eq!(large[7], 1);
    ///
    /// assert!(large.is_heap() == true);
    /// ```
    pub fn new(val: T) -> SmallBox<T, Space>
    where
        T: Sized,
    {
        smallbox!(val)
    }

    #[doc(hidden)]
    pub unsafe fn new_unchecked<U>(val: U, ptr: *const T) -> SmallBox<T, Space>
    where
        U: Sized,
    {
        let result = Self::new_copy(&val, ptr);
        mem::forget(val);
        result
    }

    unsafe fn new_copy<U>(val: &U, ptr: *const T) -> SmallBox<T, Space>
    where
        U: ?Sized,
    {
        let mut space = ManuallyDrop::new(mem::uninitialized::<Space>());

        let (ptr_addr, ptr_copy): (*const u8, *mut u8) = if mem::size_of_val::<U>(val)
            > mem::size_of::<Space>()
            || mem::align_of_val::<U>(val) > mem::align_of::<Space>()
        {
            // Heap
            let layout = Layout::for_value::<U>(val);
            let heap_ptr = alloc::alloc(layout);

            (heap_ptr, heap_ptr)
        } else {
            // Stack
            (ptr::null(), mem::transmute(&mut space))
        };

        let mut ptr = ptr;
        let ptr_ptr = &mut ptr as *mut _ as *mut usize;
        ptr_ptr.write(ptr_addr as usize);

        ptr::copy_nonoverlapping(
            val as *const _ as *const u8,
            ptr_copy,
            mem::size_of_val::<U>(val),
        );

        SmallBox {
            space,
            ptr,
            _phantom: PhantomData,
        }
    }

    /// Change the capacity of `SmallBox`
    ///
    /// This method may move stack-allocated data to heap
    /// if the inline space is not sufficient. Once the data
    /// is stored on heap, it'll never be moved again.
    ///
    /// # Example
    ///
    /// ```
    /// use smallbox::SmallBox;
    /// use smallbox::space::{S2, S4};
    ///
    /// let s: SmallBox::<_, S4> = SmallBox::new([0usize; 4]);
    /// let m: SmallBox::<_, S2> = s.resize();
    /// ```
    pub fn resize<ToSpace>(self) -> SmallBox<T, ToSpace> {
        unsafe {
            let result = if self.is_heap() {
                // don't touch anything if the data is already on heap
                let mut space = ManuallyDrop::new(mem::uninitialized::<ToSpace>());
                SmallBox {
                    space,
                    ptr: self.ptr,
                    _phantom: PhantomData,
                }
            } else {
                let val: &T = &*self;
                SmallBox::<T, ToSpace>::new_copy(val, val as *const T)
            };

            mem::forget(self);

            result
        }
    }

    /// Returns true if the data is heap-allocated
    pub fn is_heap(&self) -> bool {
        !self.ptr.is_null()
    }

    unsafe fn as_ptr(&self) -> *const T {
        let mut ptr = self.ptr;

        if !self.is_heap() {
            let ptr_ptr = &mut ptr as *mut _ as *mut usize;
            ptr_ptr.write(mem::transmute(&self.space));
        }

        ptr
    }
}

impl<T: ?Sized, Space> ops::Deref for SmallBox<T, Space> {
    type Target = T;

    fn deref(&self) -> &T {
        unsafe { &*self.as_ptr() }
    }
}

impl<T: ?Sized, Space> ops::DerefMut for SmallBox<T, Space> {
    fn deref_mut(&mut self) -> &mut T {
        unsafe { &mut *(self.as_ptr() as *const _ as *mut _) }
    }
}

impl<T: ?Sized, Space> ops::Drop for SmallBox<T, Space> {
    fn drop(&mut self) {
        unsafe {
            let layout = Layout::for_value::<T>(&*self);
            ptr::drop_in_place::<T>(&mut **self);
            if self.is_heap() {
                alloc::dealloc(self.ptr as *mut u8, layout);
            }
        }
    }
}

impl<T: Clone, Space> Clone for SmallBox<T, Space>
where
    T: Sized,
{
    fn clone(&self) -> Self {
        let val: &T = &*self;
        SmallBox::new(val.clone())
    }
}

impl<T: ?Sized + fmt::Display, Space> fmt::Display for SmallBox<T, Space> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&**self, f)
    }
}

impl<T: ?Sized + fmt::Debug, Space> fmt::Debug for SmallBox<T, Space> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&**self, f)
    }
}

impl<T: ?Sized, Space> fmt::Pointer for SmallBox<T, Space> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // It's not possible to extract the inner Unique directly from the Box,
        // instead we cast it to a *const which aliases the Unique
        let ptr: *const T = &**self;
        fmt::Pointer::fmt(&ptr, f)
    }
}

impl<T: ?Sized + PartialEq, Space> PartialEq for SmallBox<T, Space> {
    #[inline]
    fn eq(&self, other: &SmallBox<T, Space>) -> bool {
        PartialEq::eq(&**self, &**other)
    }
    #[inline]
    fn ne(&self, other: &SmallBox<T, Space>) -> bool {
        PartialEq::ne(&**self, &**other)
    }
}

impl<T: ?Sized + PartialOrd, Space> PartialOrd for SmallBox<T, Space> {
    #[inline]
    fn partial_cmp(&self, other: &SmallBox<T, Space>) -> Option<Ordering> {
        PartialOrd::partial_cmp(&**self, &**other)
    }
    #[inline]
    fn lt(&self, other: &SmallBox<T, Space>) -> bool {
        PartialOrd::lt(&**self, &**other)
    }
    #[inline]
    fn le(&self, other: &SmallBox<T, Space>) -> bool {
        PartialOrd::le(&**self, &**other)
    }
    #[inline]
    fn ge(&self, other: &SmallBox<T, Space>) -> bool {
        PartialOrd::ge(&**self, &**other)
    }
    #[inline]
    fn gt(&self, other: &SmallBox<T, Space>) -> bool {
        PartialOrd::gt(&**self, &**other)
    }
}

impl<T: ?Sized + Ord, Space> Ord for SmallBox<T, Space> {
    #[inline]
    fn cmp(&self, other: &SmallBox<T, Space>) -> Ordering {
        Ord::cmp(&**self, &**other)
    }
}

impl<T: ?Sized + Eq, Space> Eq for SmallBox<T, Space> {}

impl<T: ?Sized + Hash, Space> Hash for SmallBox<T, Space> {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        (**self).hash(state);
    }
}

unsafe impl<T: ?Sized + Send, Space> Send for SmallBox<T, Space> {}
unsafe impl<T: ?Sized + Sync, Space> Sync for SmallBox<T, Space> {}

#[cfg(test)]
mod tests {
    use super::SmallBox;
    use space::*;
    use std::any::Any;

    #[test]
    fn test_basic() {
        let stacked: SmallBox<usize, S1> = SmallBox::new(1234usize);
        assert!(*stacked == 1234);

        let heaped: SmallBox<(usize, usize), S1> = SmallBox::new((0, 1));
        assert!(*heaped == (0, 1));
    }

    #[test]
    fn test_new_unchecked() {
        let val = [0usize, 1];
        let ptr = &val as *const _;

        unsafe {
            let stacked: SmallBox<[usize], S2> = SmallBox::new_unchecked(val, ptr);
            assert!(*stacked == [0, 1]);
            assert!(!stacked.is_heap());
        }

        let val = [0usize, 1, 2];
        let ptr = &val as *const _;

        unsafe {
            let heaped: SmallBox<Any, S2> = SmallBox::new_unchecked(val, ptr);
            assert!(heaped.is_heap());

            if let Some(array) = heaped.downcast_ref::<[usize; 3]>() {
                assert_eq!(*array, [0, 1, 2]);
            } else {
                unreachable!();
            }
        }
    }

    #[test]
    #[deny(unsafe_code)]
    fn test_macro() {
        let stacked: SmallBox<Any, S1> = smallbox!(1234usize);
        if let Some(num) = stacked.downcast_ref::<usize>() {
            assert_eq!(*num, 1234);
        } else {
            unreachable!();
        }

        let heaped: SmallBox<Any, S1> = smallbox!([0usize, 1]);
        if let Some(array) = heaped.downcast_ref::<[usize; 2]>() {
            assert_eq!(*array, [0, 1]);
        } else {
            unreachable!();
        }

        let is_even: SmallBox<Fn(u8) -> bool, S1> = smallbox!(|num: u8| num % 2 == 0);
        assert!(!is_even(5));
        assert!(is_even(6));
    }

    #[test]
    #[cfg(feature = "coerce")]
    fn test_coerce() {
        let stacked: SmallBox<Any, S1> = SmallBox::new(1234usize);
        if let Some(num) = stacked.downcast_ref::<usize>() {
            assert_eq!(*num, 1234);
        } else {
            unreachable!();
        }

        let heaped: SmallBox<Any, S1> = SmallBox::new([0usize, 1]);
        if let Some(array) = heaped.downcast_ref::<[usize; 2]>() {
            assert_eq!(*array, [0, 1]);
        } else {
            unreachable!();
        }
    }

    #[test]
    fn test_drop() {
        use std::cell::Cell;

        #[derive(Debug, Clone)]
        struct Struct<'a>(&'a Cell<bool>);
        impl<'a> Drop for Struct<'a> {
            fn drop(&mut self) {
                self.0.set(true);
            }
        }

        let flag = Cell::new(false);
        let val: SmallBox<_, S2> = SmallBox::new(Struct(&flag));
        assert!(flag.get() == false);

        drop(val);
        assert!(flag.get() == true);
    }

    #[test]
    fn test_dont_drop_space() {
        struct NoDrop(S1);
        impl Drop for NoDrop {
            fn drop(&mut self) {
                unreachable!();
            }
        }

        drop(SmallBox::<_, NoDrop>::new([true]));
    }

    #[test]
    fn test_oversize() {
        let fit = SmallBox::<_, S1>::new([0usize; 1]);
        let oversize = SmallBox::<_, S1>::new([0usize; 2]);
        assert!(!fit.is_heap());
        assert!(oversize.is_heap());
    }

    #[test]
    fn test_resize() {
        let m = SmallBox::<_, S4>::new([0usize; 2]);
        let l = m.resize::<S8>();
        assert!(!l.is_heap());
        let m = l.resize::<S4>();
        assert!(!m.is_heap());
        let s = m.resize::<S2>();
        assert!(!s.is_heap());
        let xs = s.resize::<S1>();
        assert!(xs.is_heap());
        let m = xs.resize::<S4>();
        assert!(m.is_heap());
    }

    #[test]
    fn test_clone() {
        let stacked: SmallBox<[usize; 2], S2> = smallbox!([0usize, 1]);
        assert_eq!(stacked, stacked.clone())
    }

    #[test]
    fn test_zst() {
        let zst = SmallBox::<_, S1>::new([0usize; 0]);
        assert_eq!(*zst, [0usize; 0]);
    }
}
<|endoftext|>


<commit_before>// Extern crates.
use graphics::*;
use input::keyboard;
use current::{ Usage, UseCurrent };

// Local crate.
use action;
use settings;
use object;
use object::Object;
use text;
use game_state;

pub fn current_cam() -> Usage<'static, Cam> { UseCurrent }
pub fn current_game_state()
    -> Usage<'static, game_state::GameState> { UseCurrent }
pub fn current_objects() -> Usage<'static, Vec<Object>> { UseCurrent }
pub fn current_index() -> Usage<'static, Index> { UseCurrent }
pub fn current_settings() -> Usage<'static, Settings> { UseCurrent }

pub fn app() {
    use std::cell::RefCell;
    use current::Current;

    let cam = Cam([0.0, 0.0]);
    let game_state = game_state::Play;
    let objects: Vec<Object> = Vec::new();
    let index = Index::new();
    let settings = Settings::new();

    let cam = RefCell::new(cam);
    let game_state = RefCell::new(game_state);
    let objects = RefCell::new(objects);
    let index = RefCell::new(index);
    let settings = RefCell::new(settings);

    let cam_guard = cam.set_current();
    let game_state_guard = game_state.set_current();
    let objects_guard = objects.set_current();
    let index_guard = index.set_current();
    let settings_guard = settings.set_current();

    start();

    drop(cam_guard);
    drop(game_state_guard);
    drop(objects_guard);
    drop(index_guard);
    drop(settings_guard);
}

fn start() {
    use event::{ RenderEvent, UpdateEvent, PressEvent, ReleaseEvent };
    use input;

    load();
    for e in ::events() {
        ::swap_backend(&e);
        e.render(|args| {
            ::render(args);
            ::fps_tick();
        });
        e.update(|args| {
            update(args.dt);
        });
        e.press(|button| {
            if let input::Keyboard(key) = button {
                key_press(key);
            }
        });
        e.release(|button| {
            if let input::Keyboard(key) = button {
                key_release(key);
            }
        });
    }
}

pub struct Cam(pub [f64, ..2]);

impl Cam {
    fn pos(&self) -> [f64, ..2] {
        let Cam(pos) = *self;
        pos
    }

    fn set(&mut self, val: [f64, ..2]) {
        *self = Cam(val);
    }
}

pub struct Index {
    pub player: Option<uint>,
    pub blood_bar: Option<uint>,
    pub air_bar: Option<uint>,
}

impl Index {
    pub fn new() -> Index {
        Index {
            player: None,
            blood_bar: None,
            air_bar: None
        }
    }
}

pub struct Settings {
    // Tells where the surface is.
    surface_y: Option<f64>,
    camera_follow_percentage: Option<f64>,
}

impl Settings {
    pub fn new() -> Settings {
        Settings {
            surface_y: None,
            camera_follow_percentage: None,
        }
    }
}

pub fn render<B: BackEnd<I>, I: ImageSize>(c: &Context, gl: &mut B) {
    let c = &c.reset();

    // Get camera coordinates.
    let [cam_x, mut cam_y] = current_cam().pos();
    cam_y += 0.4;

    // Render surface.
    let surface_y = current_settings().surface_y.unwrap();
    c.rect(-1.0, surface_y - cam_y, 2.0, 0.05).color(settings::BLUE).draw(gl);

    // Render objects in layers.
    let cam = &c.trans(-cam_x, -cam_y);
    for i in range(0u, settings::NUMBER_OF_LAYERS) {
        for obj in current_objects().iter() {
            if obj.layer == i { obj.render(cam, c, gl); }
        }
    }

    let text_c = c.flip_v();
    let text_c = text_c.zoom(0.0025);
    match *current_game_state() {
        game_state::Win => {
            let pos = settings::YOU_WIN_POS;
            text::text(settings::YOU_WIN_TEXT,
                &text_c
                .trans(pos[0], pos[1])
                .color(settings::YOU_WIN_TEXT_COLOR)
            , gl);
        },
        game_state::Lose => {
            let pos = settings::YOU_LOSE_POS;
            text::text(settings::YOU_LOSE_TEXT,
                &text_c
                .trans(pos[0], pos[1])
                .color(settings::YOU_LOSE_TEXT_COLOR)
            , gl);
        },
        game_state::Play => {

        },
    }
}

pub fn update(dt: f64) {
    update_objects(dt);
    fill_air();
    win();
    lose();
    show_blood();
    show_air();
    follow_player(dt);
}

pub fn load() {
    current_cam().set(settings::INITIAL_CAMERA_POS);
    current_settings().camera_follow_percentage = Some(settings::CAMERA_FOLLOW_PERCENTAGE);
    current_settings().surface_y = Some(settings::SURFACE_Y);
    *current_game_state() = settings::INITIAL_GAME_STATE;

    // Add player.
    current_objects().push(Object::player(
        settings::ORIGIN,
        settings::BLUE,
        settings::PLAYER_INITIAL_BLOOD,
        settings::PLAYER_INITIAL_AIR,
        [settings::PLAYER_ACCELERATION_LEFT, settings::PLAYER_ACCELERATION_RIGHT],
        [settings::PLAYER_ACCELERATION_UP, settings::PLAYER_ACCELERATION_DOWN]
    ));
    current_objects().push(Object::bar_background());
    current_index().player = Some(0);

    // Add blood and air bar.
    add_bars();

    // Add air bottles.
    add_air_bottles();

    // Add snakes.
    add_snakes();
}

pub fn key_press(key: keyboard::Key) {
    if *current_game_state() != game_state::Play { return; }

    match (key, current_index().player) {
        (keyboard::Right, Some(player_index)) => {
            current_objects()[player_index].move_right();
        },
        (keyboard::Up, Some(player_index)) => {
            current_objects()[player_index].move_up();
        },
        (keyboard::Left, Some(player_index)) => {
            current_objects()[player_index].move_left();
        },
        (keyboard::Down, Some(player_index)) => {
            current_objects()[player_index].move_down();
        },
        _ => {},
    }
}

pub fn key_release(key: keyboard::Key) {
    if key == keyboard::Return || key == keyboard::Space {
        match *current_game_state() {
            game_state::Win | game_state::Lose => restart(),
            _ => {},
        }
    }
}

pub fn add_bars() {
    let objects = &mut *current_objects();
    objects.push(Object::bar(
        settings::AIR_BAR_POS,
        "air",
        settings::AIR_BAR_TEXT_COLOR,
        settings::AIR_BAR_BACKGROUND_COLOR,
        settings::AIR_BAR_BAR_COLOR,
        settings::AIR_BAR_INITIAL_VALUE
    ));
    current_index().air_bar = Some(objects.len() - 1);
    objects.push(Object::bar(
        settings::BLOOD_BAR_POS,
        "blood",
        settings::BLOOD_BAR_TEXT_COLOR,
        settings::BLOOD_BAR_BACKGROUND_COLOR,
        settings::BLOOD_BAR_BAR_COLOR,
        settings::BLOOD_BAR_INITIAL_VALUE
    ));
    current_index().blood_bar = Some(objects.len() - 1);
}

pub fn add_snakes() {
    if settings::SNAKE_1_ADD { current_objects().push(
        Object::snake(settings::SNAKE_1_POS, settings::SNAKE_1_SETTINGS)); }
    if settings::SNAKE_2_ADD { current_objects().push(
        Object::snake(settings::SNAKE_2_POS, settings::SNAKE_2_SETTINGS)); }
}

pub fn add_air_bottles() {
    let air_bottles = settings::AIR_BOTTLE_POS;
    let n = air_bottles.len() / 2;
    let objects = &mut *current_objects();
    for i in range(0, n) {
        objects.push(Object::air_bottle([air_bottles[i * 2], air_bottles[i * 2 + 1]]));
    }
}

fn follow_player(dt: f64) {
    // Make camera follow player.
    let camera_pos = current_cam().pos();
    let camera_follow_percentage = current_settings().camera_follow_percentage.unwrap();
    let player_pos = player_pos();
    let (dx, dy) = (player_pos[0] - camera_pos[0], player_pos[1] - camera_pos[1]);
    let dx = camera_follow_percentage * dt * dx;
    let dy = camera_follow_percentage * dt * dy;
    current_cam().set([camera_pos[0] + dx, camera_pos[1] + dy]);
}

fn show_blood() {
    if current_index().blood_bar == None { return; }
    // Show blood.
    let player_blood = player_blood();
    let blood_bar_index = current_index().blood_bar.unwrap();
    let objects = &mut *current_objects();
    let blood_bar = objects.get_mut(blood_bar_index).unwrap();
    match blood_bar.data {
        object::BarData(ref mut bar) => {
            bar.value = player_blood;
        },
        _ => {},
    }
}

fn show_air() {
    if current_index().air_bar == None { return; }
    // Show air.
    let player_air = player_air();
    let air_bar_index = current_index().air_bar.unwrap();
    let objects = &mut *current_objects();
    let ref mut air_bar = objects.get_mut(air_bar_index).unwrap();
    match air_bar.data {
        object::BarData(ref mut bar) => {
            bar.value = player_air;
        },
        _ => {},
    }
}

fn win() {
    let player_pos = player_pos();
    // When player reaches surface, win.
    if player_pos[1] >= current_settings().surface_y.unwrap() {
        *current_game_state() = game_state::Win;
        return;
    }
}

fn lose() {
    let blood = player_blood();
    let air = player_air();
    if blood < 0.0 || air < 0.0 {
        *current_game_state() = game_state::Lose;
        return;
    }
}

fn restart() {
    load();
}

fn update_objects(dt: f64) {
    if *current_game_state() != game_state::Play { return; }

    // Update states of objects.
    let player_pos = player_pos();
    let mut attack_damage: f64 = 0.0;
    for obj in current_objects().iter_mut() {
        match obj.update(dt, player_pos) {
            action::Passive => {},
            action::Attack(attack) => { attack_damage += attack; },
        }
    }
    // Bite player.
    bite_player(attack_damage);

    // Decrease the player's air with time.
    let air = player_air();
    set_player_air(air - dt * settings::PLAYER_LOSE_AIR_SPEED);
}

fn player_pos() -> [f64, ..2] {
    let player_index = current_index().player.unwrap();
    current_objects()[player_index].pos
}

fn player_blood() -> f64 {
    let player_index = current_index().player.unwrap();
    current_objects()[player_index].blood().unwrap()
}

fn player_air() -> f64 {
    let player_index = current_index().player.unwrap();
    current_objects()[player_index].air().unwrap()
}

fn set_player_air(val: f64) {
    let player_index = current_index().player.unwrap();
    *current_objects()[player_index].air_mut().unwrap() = val;
}

fn bite_player(damage: f64) {
    let player_index = current_index().player.unwrap();
    current_objects()[player_index].player_mut().unwrap().bite(damage);
}

fn fill_air() {
    let player_pos = player_pos();
    let mut air = player_air();
    for obj in current_objects().iter_mut() {
        let pos = obj.pos;
        match obj.air_bottle_mut() {
            None => {},
            Some(air_bottle) => {
                let dx = pos[0] - player_pos[0];
                let dy = pos[1] - player_pos[1];
                let d = dx * dx + dy * dy;
                if d <= settings::AIR_BOTTLE_RADIUS * settings::AIR_BOTTLE_RADIUS {
                    air += air_bottle.fill_up;
                    air_bottle.fill_up = 0.0;
                }
            },
        }
    }

    air = if air > 1.0 { 1.0 } else { air };
    set_player_air(air);
}
<commit_msg>Moved functions inside other functions<commit_after>// Extern crates.
use graphics::*;
use input::keyboard;
use current::{ Usage, UseCurrent };

// Local crate.
use action;
use settings;
use object;
use object::Object;
use text;
use game_state;

pub fn current_cam() -> Usage<'static, Cam> { UseCurrent }
pub fn current_game_state()
    -> Usage<'static, game_state::GameState> { UseCurrent }
pub fn current_objects() -> Usage<'static, Vec<Object>> { UseCurrent }
pub fn current_index() -> Usage<'static, Index> { UseCurrent }
pub fn current_settings() -> Usage<'static, Settings> { UseCurrent }

pub fn app() {
    use std::cell::RefCell;
    use current::Current;

    let cam = Cam([0.0, 0.0]);
    let game_state = game_state::Play;
    let objects: Vec<Object> = Vec::new();
    let index = Index::new();
    let settings = Settings::new();

    let cam = RefCell::new(cam);
    let game_state = RefCell::new(game_state);
    let objects = RefCell::new(objects);
    let index = RefCell::new(index);
    let settings = RefCell::new(settings);

    let cam_guard = cam.set_current();
    let game_state_guard = game_state.set_current();
    let objects_guard = objects.set_current();
    let index_guard = index.set_current();
    let settings_guard = settings.set_current();

    start();

    drop(cam_guard);
    drop(game_state_guard);
    drop(objects_guard);
    drop(index_guard);
    drop(settings_guard);
}

fn start() {
    use event::{ RenderEvent, UpdateEvent, PressEvent, ReleaseEvent };
    use input;

    load();
    for e in ::events() {
        ::swap_backend(&e);
        e.render(|args| {
            ::render(args);
            ::fps_tick();
        });
        e.update(|args| {
            update(args.dt);
        });
        e.press(|button| {
            if let input::Keyboard(key) = button {
                key_press(key);
            }
        });
        e.release(|button| {
            if let input::Keyboard(key) = button {
                key_release(key);
            }
        });
    }
}

pub struct Cam(pub [f64, ..2]);

impl Cam {
    fn pos(&self) -> [f64, ..2] {
        let Cam(pos) = *self;
        pos
    }

    fn set(&mut self, val: [f64, ..2]) {
        *self = Cam(val);
    }
}

pub struct Index {
    pub player: Option<uint>,
    pub blood_bar: Option<uint>,
    pub air_bar: Option<uint>,
}

impl Index {
    pub fn new() -> Index {
        Index {
            player: None,
            blood_bar: None,
            air_bar: None
        }
    }
}

pub struct Settings {
    // Tells where the surface is.
    surface_y: Option<f64>,
    camera_follow_percentage: Option<f64>,
}

impl Settings {
    pub fn new() -> Settings {
        Settings {
            surface_y: None,
            camera_follow_percentage: None,
        }
    }
}

pub fn render<B: BackEnd<I>, I: ImageSize>(c: &Context, gl: &mut B) {
    let c = &c.reset();

    // Get camera coordinates.
    let [cam_x, mut cam_y] = current_cam().pos();
    cam_y += 0.4;

    // Render surface.
    let surface_y = current_settings().surface_y.unwrap();
    c.rect(-1.0, surface_y - cam_y, 2.0, 0.05).color(settings::BLUE).draw(gl);

    // Render objects in layers.
    let cam = &c.trans(-cam_x, -cam_y);
    for i in range(0u, settings::NUMBER_OF_LAYERS) {
        for obj in current_objects().iter() {
            if obj.layer == i { obj.render(cam, c, gl); }
        }
    }

    let text_c = c.flip_v();
    let text_c = text_c.zoom(0.0025);
    match *current_game_state() {
        game_state::Win => {
            let pos = settings::YOU_WIN_POS;
            text::text(settings::YOU_WIN_TEXT,
                &text_c
                .trans(pos[0], pos[1])
                .color(settings::YOU_WIN_TEXT_COLOR)
            , gl);
        },
        game_state::Lose => {
            let pos = settings::YOU_LOSE_POS;
            text::text(settings::YOU_LOSE_TEXT,
                &text_c
                .trans(pos[0], pos[1])
                .color(settings::YOU_LOSE_TEXT_COLOR)
            , gl);
        },
        game_state::Play => {

        },
    }
}

pub fn update(dt: f64) {

    fn player_air() -> f64 {
        let player_index = current_index().player.unwrap();
        current_objects()[player_index].air().unwrap()
    }

    fn set_player_air(val: f64) {
        let player_index = current_index().player.unwrap();
        *current_objects()[player_index].air_mut().unwrap() = val;
    }

    fn bite_player(damage: f64) {
        let player_index = current_index().player.unwrap();
        current_objects()[player_index].player_mut().unwrap().bite(damage);
    }

    fn player_blood() -> f64 {
        let player_index = current_index().player.unwrap();
        current_objects()[player_index].blood().unwrap()
    }

    fn player_pos() -> [f64, ..2] {
        let player_index = current_index().player.unwrap();
        current_objects()[player_index].pos
    }

    fn update_objects(dt: f64) {
        if *current_game_state() != game_state::Play { return; }

        // Update states of objects.
        let player_pos = player_pos();
        let mut attack_damage: f64 = 0.0;
        for obj in current_objects().iter_mut() {
            match obj.update(dt, player_pos) {
                action::Passive => {},
                action::Attack(attack) => { attack_damage += attack; },
            }
        }
        // Bite player.
        bite_player(attack_damage);

        // Decrease the player's air with time.
        let air = player_air();
        set_player_air(air - dt * settings::PLAYER_LOSE_AIR_SPEED);
    }

    update_objects(dt);

    fn fill_air() {
        let player_pos = player_pos();
        let mut air = player_air();
        for obj in current_objects().iter_mut() {
            let pos = obj.pos;
            match obj.air_bottle_mut() {
                None => {},
                Some(air_bottle) => {
                    let dx = pos[0] - player_pos[0];
                    let dy = pos[1] - player_pos[1];
                    let d = dx * dx + dy * dy;
                    if d <= settings::AIR_BOTTLE_RADIUS * settings::AIR_BOTTLE_RADIUS {
                        air += air_bottle.fill_up;
                        air_bottle.fill_up = 0.0;
                    }
                },
            }
        }

        air = if air > 1.0 { 1.0 } else { air };
        set_player_air(air);
    }

    fill_air();

    fn win() {
        let player_pos = player_pos();
        // When player reaches surface, win.
        if player_pos[1] >= current_settings().surface_y.unwrap() {
            *current_game_state() = game_state::Win;
            return;
        }
    }

    win();

    fn lose() {
        let blood = player_blood();
        let air = player_air();
        if blood < 0.0 || air < 0.0 {
            *current_game_state() = game_state::Lose;
            return;
        }
    }

    lose();

    fn show_blood() {
        if current_index().blood_bar == None { return; }
        // Show blood.
        let player_blood = player_blood();
        let blood_bar_index = current_index().blood_bar.unwrap();
        let objects = &mut *current_objects();
        let blood_bar = objects.get_mut(blood_bar_index).unwrap();
        match blood_bar.data {
            object::BarData(ref mut bar) => {
                bar.value = player_blood;
            },
            _ => {},
        }
    }

    show_blood();

    fn show_air() {
        if current_index().air_bar == None { return; }
        // Show air.
        let player_air = player_air();
        let air_bar_index = current_index().air_bar.unwrap();
        let objects = &mut *current_objects();
        let ref mut air_bar = objects.get_mut(air_bar_index).unwrap();
        match air_bar.data {
            object::BarData(ref mut bar) => {
                bar.value = player_air;
            },
            _ => {},
        }
    }

    show_air();

    fn follow_player(dt: f64) {
        // Make camera follow player.
        let camera_pos = current_cam().pos();
        let camera_follow_percentage = current_settings().camera_follow_percentage.unwrap();
        let player_pos = player_pos();
        let (dx, dy) = (player_pos[0] - camera_pos[0], player_pos[1] - camera_pos[1]);
        let dx = camera_follow_percentage * dt * dx;
        let dy = camera_follow_percentage * dt * dy;
        current_cam().set([camera_pos[0] + dx, camera_pos[1] + dy]);
    }

    follow_player(dt);
}

pub fn load() {
    current_cam().set(settings::INITIAL_CAMERA_POS);
    current_settings().camera_follow_percentage = Some(settings::CAMERA_FOLLOW_PERCENTAGE);
    current_settings().surface_y = Some(settings::SURFACE_Y);
    *current_game_state() = settings::INITIAL_GAME_STATE;

    // Add player.
    current_objects().push(Object::player(
        settings::ORIGIN,
        settings::BLUE,
        settings::PLAYER_INITIAL_BLOOD,
        settings::PLAYER_INITIAL_AIR,
        [settings::PLAYER_ACCELERATION_LEFT, settings::PLAYER_ACCELERATION_RIGHT],
        [settings::PLAYER_ACCELERATION_UP, settings::PLAYER_ACCELERATION_DOWN]
    ));
    current_objects().push(Object::bar_background());
    current_index().player = Some(0);

    fn add_bars() {
        let objects = &mut *current_objects();
        objects.push(Object::bar(
            settings::AIR_BAR_POS,
            "air",
            settings::AIR_BAR_TEXT_COLOR,
            settings::AIR_BAR_BACKGROUND_COLOR,
            settings::AIR_BAR_BAR_COLOR,
            settings::AIR_BAR_INITIAL_VALUE
        ));
        current_index().air_bar = Some(objects.len() - 1);
        objects.push(Object::bar(
            settings::BLOOD_BAR_POS,
            "blood",
            settings::BLOOD_BAR_TEXT_COLOR,
            settings::BLOOD_BAR_BACKGROUND_COLOR,
            settings::BLOOD_BAR_BAR_COLOR,
            settings::BLOOD_BAR_INITIAL_VALUE
        ));
        current_index().blood_bar = Some(objects.len() - 1);
    }

    // Add blood and air bar.
    add_bars();

    fn add_air_bottles() {
        let air_bottles = settings::AIR_BOTTLE_POS;
        let n = air_bottles.len() / 2;
        let objects = &mut *current_objects();
        for i in range(0, n) {
            objects.push(Object::air_bottle([air_bottles[i * 2], air_bottles[i * 2 + 1]]));
        }
    }

    // Add air bottles.
    add_air_bottles();

    fn add_snakes() {
        if settings::SNAKE_1_ADD { current_objects().push(
            Object::snake(settings::SNAKE_1_POS, settings::SNAKE_1_SETTINGS)); }
        if settings::SNAKE_2_ADD { current_objects().push(
            Object::snake(settings::SNAKE_2_POS, settings::SNAKE_2_SETTINGS)); }
    }

    // Add snakes.
    add_snakes();
}

fn restart() {
    load();
}

pub fn key_press(key: keyboard::Key) {
    if *current_game_state() != game_state::Play { return; }

    match (key, current_index().player) {
        (keyboard::Right, Some(player_index)) => {
            current_objects()[player_index].move_right();
        },
        (keyboard::Up, Some(player_index)) => {
            current_objects()[player_index].move_up();
        },
        (keyboard::Left, Some(player_index)) => {
            current_objects()[player_index].move_left();
        },
        (keyboard::Down, Some(player_index)) => {
            current_objects()[player_index].move_down();
        },
        _ => {},
    }
}

pub fn key_release(key: keyboard::Key) {
    if key == keyboard::Return || key == keyboard::Space {
        match *current_game_state() {
            game_state::Win | game_state::Lose => restart(),
            _ => {},
        }
    }
}
<|endoftext|>


<commit_before>use crate::{error, NetworkConfig};
/// This manages the discovery and management of peers.
///
/// Currently using discv5 for peer discovery.
///
use futures::prelude::*;
use libp2p::core::swarm::{
    ConnectedPoint, NetworkBehaviour, NetworkBehaviourAction, PollParameters,
};
use libp2p::core::{identity::Keypair, Multiaddr, PeerId, ProtocolsHandler};
use libp2p::discv5::{Discv5, Discv5Event};
use libp2p::enr::{Enr, EnrBuilder, NodeId};
use libp2p::multiaddr::Protocol;
use slog::{debug, info, o, warn};
use std::collections::HashSet;
use std::fs::File;
use std::io::prelude::*;
use std::str::FromStr;
use std::time::{Duration, Instant};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_timer::Delay;

/// Maximum seconds before searching for extra peers.
const MAX_TIME_BETWEEN_PEER_SEARCHES: u64 = 60;
/// Initial delay between peer searches.
const INITIAL_SEARCH_DELAY: u64 = 5;
/// Local ENR storage filename.
const ENR_FILENAME: &str = "enr.dat";

/// Lighthouse discovery behaviour. This provides peer management and discovery using the Discv5
/// libp2p protocol.
pub struct Discovery<TSubstream> {
    /// The peers currently connected to libp2p streams.
    connected_peers: HashSet<PeerId>,

    /// The target number of connected peers on the libp2p interface.
    max_peers: usize,

    /// The delay between peer discovery searches.
    peer_discovery_delay: Delay,

    /// Tracks the last discovery delay. The delay is doubled each round until the max
    /// time is reached.
    past_discovery_delay: u64,

    /// The TCP port for libp2p. Used to convert an updated IP address to a multiaddr. Note: This
    /// assumes that the external TCP port is the same as the internal TCP port if behind a NAT.
    //TODO: Improve NAT handling limit the above restriction
    tcp_port: u16,

    /// The discovery behaviour used to discover new peers.
    discovery: Discv5<TSubstream>,

    /// Logger for the discovery behaviour.
    log: slog::Logger,
}

impl<TSubstream> Discovery<TSubstream> {
    pub fn new(
        local_key: &Keypair,
        config: &NetworkConfig,
        log: &slog::Logger,
    ) -> error::Result<Self> {
        let log = log.new(o!("Service" => "Libp2p-Discovery"));

        // checks if current ENR matches that found on disk
        let local_enr = load_enr(local_key, config, &log)?;

        info!(log, "Local ENR: {}", local_enr.to_base64());
        debug!(log, "Local Node Id: {}", local_enr.node_id());

        let mut discovery = Discv5::new(local_enr, local_key.clone(), config.listen_address)
            .map_err(|e| format!("Discv5 service failed: {:?}", e))?;

        // Add bootnodes to routing table
        for bootnode_enr in config.boot_nodes.clone() {
            debug!(
                log,
                "Adding node to routing table: {}",
                bootnode_enr.node_id()
            );
            discovery.add_enr(bootnode_enr);
        }

        Ok(Self {
            connected_peers: HashSet::new(),
            max_peers: config.max_peers,
            peer_discovery_delay: Delay::new(Instant::now()),
            past_discovery_delay: INITIAL_SEARCH_DELAY,
            tcp_port: config.libp2p_port,
            discovery,
            log,
        })
    }

    /// Manually search for peers. This restarts the discovery round, sparking multiple rapid
    /// queries.
    pub fn discover_peers(&mut self) {
        self.past_discovery_delay = INITIAL_SEARCH_DELAY;
        self.find_peers();
    }

    /// Add an Enr to the routing table of the discovery mechanism.
    pub fn add_enr(&mut self, enr: Enr) {
        self.discovery.add_enr(enr);
    }

    /// Search for new peers using the underlying discovery mechanism.
    fn find_peers(&mut self) {
        // pick a random NodeId
        let random_node = NodeId::random();
        debug!(self.log, "Searching for peers...");
        self.discovery.find_node(random_node);

        // update the time until next discovery
        let delay = {
            if self.past_discovery_delay < MAX_TIME_BETWEEN_PEER_SEARCHES {
                self.past_discovery_delay *= 2;
                self.past_discovery_delay
            } else {
                MAX_TIME_BETWEEN_PEER_SEARCHES
            }
        };
        self.peer_discovery_delay
            .reset(Instant::now() + Duration::from_secs(delay));
    }
}

// Redirect all behaviour events to underlying discovery behaviour.
impl<TSubstream> NetworkBehaviour for Discovery<TSubstream>
where
    TSubstream: AsyncRead + AsyncWrite,
{
    type ProtocolsHandler = <Discv5<TSubstream> as NetworkBehaviour>::ProtocolsHandler;
    type OutEvent = <Discv5<TSubstream> as NetworkBehaviour>::OutEvent;

    fn new_handler(&mut self) -> Self::ProtocolsHandler {
        NetworkBehaviour::new_handler(&mut self.discovery)
    }

    fn addresses_of_peer(&mut self, peer_id: &PeerId) -> Vec<Multiaddr> {
        // Let discovery track possible known peers.
        self.discovery.addresses_of_peer(peer_id)
    }

    fn inject_connected(&mut self, peer_id: PeerId, _endpoint: ConnectedPoint) {
        self.connected_peers.insert(peer_id);
    }

    fn inject_disconnected(&mut self, peer_id: &PeerId, _endpoint: ConnectedPoint) {
        self.connected_peers.remove(peer_id);
    }

    fn inject_replaced(
        &mut self,
        _peer_id: PeerId,
        _closed: ConnectedPoint,
        _opened: ConnectedPoint,
    ) {
        // discv5 doesn't implement
    }

    fn inject_node_event(
        &mut self,
        _peer_id: PeerId,
        _event: <Self::ProtocolsHandler as ProtocolsHandler>::OutEvent,
    ) {
        // discv5 doesn't implement
    }

    fn poll(
        &mut self,
        params: &mut impl PollParameters,
    ) -> Async<
        NetworkBehaviourAction<
            <Self::ProtocolsHandler as ProtocolsHandler>::InEvent,
            Self::OutEvent,
        >,
    > {
        // search of peers if it is time
        loop {
            match self.peer_discovery_delay.poll() {
                Ok(Async::Ready(_)) => {
                    if self.connected_peers.len() < self.max_peers {
                        self.find_peers();
                    }
                }
                Ok(Async::NotReady) => break,
                Err(e) => {
                    warn!(self.log, "Discovery peer search failed: {:?}", e);
                }
            }
        }

        // Poll discovery
        loop {
            match self.discovery.poll(params) {
                Async::Ready(NetworkBehaviourAction::GenerateEvent(event)) => {
                    match event {
                        Discv5Event::Discovered(_enr) => {
                            // not concerned about FINDNODE results, rather the result of an entire
                            // query.
                        }
                        Discv5Event::SocketUpdated(socket) => {
                            info!(self.log, "Address updated"; "IP" => format!("{}",socket.ip()));
                            let mut address = Multiaddr::from(socket.ip());
                            address.push(Protocol::Tcp(self.tcp_port));
                            return Async::Ready(NetworkBehaviourAction::ReportObservedAddr {
                                address,
                            });
                        }
                        Discv5Event::FindNodeResult { closer_peers, .. } => {
                            debug!(self.log, "Discv5 query found {} peers", closer_peers.len());
                            if closer_peers.is_empty() {
                                debug!(self.log, "Discv5 random query yielded empty results");
                            }
                            for peer_id in closer_peers {
                                // if we need more peers, attempt a connection
                                if self.connected_peers.len() < self.max_peers
                                    && self.connected_peers.get(&peer_id).is_none()
                                {
                                    debug!(self.log, "Discv5: Peer discovered"; "Peer"=> format!("{:?}", peer_id));
                                    return Async::Ready(NetworkBehaviourAction::DialPeer {
                                        peer_id,
                                    });
                                }
                            }
                        }
                        _ => {}
                    }
                }
                // discv5 does not output any other NetworkBehaviourAction
                Async::Ready(_) => {}
                Async::NotReady => break,
            }
        }
        Async::NotReady
    }
}

/// Loads an ENR from file if it exists and matches the current NodeId and sequence number. If none
/// exists, generates a new one.
///
/// If an ENR exists, with the same NodeId and IP addresses, we use the disk-generated one as it's
/// ENR sequence will be equal or higher than a newly generated one.
fn load_enr(
    local_key: &Keypair,
    config: &NetworkConfig,
    log: &slog::Logger,
) -> Result<Enr, String> {
    // Build the local ENR.
    // Note: Discovery should update the ENR record's IP to the external IP as seen by the
    // majority of our peers.
    let mut local_enr = EnrBuilder::new()
        .ip(config.discovery_address.into())
        .tcp(config.libp2p_port)
        .udp(config.discovery_port)
        .build(&local_key)
        .map_err(|e| format!("Could not build Local ENR: {:?}", e))?;

    let enr_f = config.network_dir.join(ENR_FILENAME);
    if let Ok(mut enr_file) = File::open(enr_f.clone()) {
        let mut enr_string = String::new();
        match enr_file.read_to_string(&mut enr_string) {
            Err(_) => debug!(log, "Could not read ENR from file"),
            Ok(_) => {
                match Enr::from_str(&enr_string) {
                    Ok(enr) => {
                        debug!(log, "ENR found in file: {:?}", enr_f);

                        if enr.node_id() == local_enr.node_id() {
                            if enr.ip() == config.discovery_address.into()
                                && enr.tcp() == Some(config.libp2p_port)
                                && enr.udp() == Some(config.discovery_port)
                            {
                                debug!(log, "ENR loaded from file");
                                // the stored ENR has the same configuration, use it
                                return Ok(enr);
                            }

                            // same node id, different configuration - update the sequence number
                            let new_seq_no = enr.seq().checked_add(1).ok_or_else(|| "ENR sequence number on file is too large. Remove it to generate a new NodeId")?;
                            local_enr.set_seq(new_seq_no, local_key).map_err(|e| {
                                format!("Could not update ENR sequence number: {:?}", e)
                            })?;
                            debug!(log, "ENR sequence number increased to: {}", new_seq_no);
                        }
                    }
                    Err(e) => {
                        warn!(log, "ENR from file could not be decoded: {:?}", e);
                    }
                }
            }
        }
    }

    // write ENR to disk
    let _ = std::fs::create_dir_all(&config.network_dir);
    match File::create(enr_f.clone())
        .and_then(|mut f| f.write_all(&local_enr.to_base64().as_bytes()))
    {
        Ok(_) => {
            debug!(log, "ENR written to disk");
        }
        Err(e) => {
            warn!(
                log,
                "Could not write ENR to file: {:?}. Error: {}", enr_f, e
            );
        }
    }
    Ok(local_enr)
}
<commit_msg>Typo fixes<commit_after>use crate::{error, NetworkConfig};
/// This manages the discovery and management of peers.
///
/// Currently using discv5 for peer discovery.
///
use futures::prelude::*;
use libp2p::core::swarm::{
    ConnectedPoint, NetworkBehaviour, NetworkBehaviourAction, PollParameters,
};
use libp2p::core::{identity::Keypair, Multiaddr, PeerId, ProtocolsHandler};
use libp2p::discv5::{Discv5, Discv5Event};
use libp2p::enr::{Enr, EnrBuilder, NodeId};
use libp2p::multiaddr::Protocol;
use slog::{debug, info, o, warn};
use std::collections::HashSet;
use std::fs::File;
use std::io::prelude::*;
use std::str::FromStr;
use std::time::{Duration, Instant};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_timer::Delay;

/// Maximum seconds before searching for extra peers.
const MAX_TIME_BETWEEN_PEER_SEARCHES: u64 = 60;
/// Initial delay between peer searches.
const INITIAL_SEARCH_DELAY: u64 = 5;
/// Local ENR storage filename.
const ENR_FILENAME: &str = "enr.dat";

/// Lighthouse discovery behaviour. This provides peer management and discovery using the Discv5
/// libp2p protocol.
pub struct Discovery<TSubstream> {
    /// The peers currently connected to libp2p streams.
    connected_peers: HashSet<PeerId>,

    /// The target number of connected peers on the libp2p interface.
    max_peers: usize,

    /// The delay between peer discovery searches.
    peer_discovery_delay: Delay,

    /// Tracks the last discovery delay. The delay is doubled each round until the max
    /// time is reached.
    past_discovery_delay: u64,

    /// The TCP port for libp2p. Used to convert an updated IP address to a multiaddr. Note: This
    /// assumes that the external TCP port is the same as the internal TCP port if behind a NAT.
    //TODO: Improve NAT handling limit the above restriction
    tcp_port: u16,

    /// The discovery behaviour used to discover new peers.
    discovery: Discv5<TSubstream>,

    /// Logger for the discovery behaviour.
    log: slog::Logger,
}

impl<TSubstream> Discovery<TSubstream> {
    pub fn new(
        local_key: &Keypair,
        config: &NetworkConfig,
        log: &slog::Logger,
    ) -> error::Result<Self> {
        let log = log.new(o!("Service" => "Libp2p-Discovery"));

        // checks if current ENR matches that found on disk
        let local_enr = load_enr(local_key, config, &log)?;

        info!(log, "Local ENR: {}", local_enr.to_base64());
        debug!(log, "Local Node Id: {}", local_enr.node_id());

        let mut discovery = Discv5::new(local_enr, local_key.clone(), config.listen_address)
            .map_err(|e| format!("Discv5 service failed: {:?}", e))?;

        // Add bootnodes to routing table
        for bootnode_enr in config.boot_nodes.clone() {
            debug!(
                log,
                "Adding node to routing table: {}",
                bootnode_enr.node_id()
            );
            discovery.add_enr(bootnode_enr);
        }

        Ok(Self {
            connected_peers: HashSet::new(),
            max_peers: config.max_peers,
            peer_discovery_delay: Delay::new(Instant::now()),
            past_discovery_delay: INITIAL_SEARCH_DELAY,
            tcp_port: config.libp2p_port,
            discovery,
            log,
        })
    }

    /// Manually search for peers. This restarts the discovery round, sparking multiple rapid
    /// queries.
    pub fn discover_peers(&mut self) {
        self.past_discovery_delay = INITIAL_SEARCH_DELAY;
        self.find_peers();
    }

    /// Add an Enr to the routing table of the discovery mechanism.
    pub fn add_enr(&mut self, enr: Enr) {
        self.discovery.add_enr(enr);
    }

    /// Search for new peers using the underlying discovery mechanism.
    fn find_peers(&mut self) {
        // pick a random NodeId
        let random_node = NodeId::random();
        debug!(self.log, "Searching for peers...");
        self.discovery.find_node(random_node);

        // update the time until next discovery
        let delay = {
            if self.past_discovery_delay < MAX_TIME_BETWEEN_PEER_SEARCHES {
                self.past_discovery_delay *= 2;
                self.past_discovery_delay
            } else {
                MAX_TIME_BETWEEN_PEER_SEARCHES
            }
        };
        self.peer_discovery_delay
            .reset(Instant::now() + Duration::from_secs(delay));
    }
}

// Redirect all behaviour events to underlying discovery behaviour.
impl<TSubstream> NetworkBehaviour for Discovery<TSubstream>
where
    TSubstream: AsyncRead + AsyncWrite,
{
    type ProtocolsHandler = <Discv5<TSubstream> as NetworkBehaviour>::ProtocolsHandler;
    type OutEvent = <Discv5<TSubstream> as NetworkBehaviour>::OutEvent;

    fn new_handler(&mut self) -> Self::ProtocolsHandler {
        NetworkBehaviour::new_handler(&mut self.discovery)
    }

    fn addresses_of_peer(&mut self, peer_id: &PeerId) -> Vec<Multiaddr> {
        // Let discovery track possible known peers.
        self.discovery.addresses_of_peer(peer_id)
    }

    fn inject_connected(&mut self, peer_id: PeerId, _endpoint: ConnectedPoint) {
        self.connected_peers.insert(peer_id);
    }

    fn inject_disconnected(&mut self, peer_id: &PeerId, _endpoint: ConnectedPoint) {
        self.connected_peers.remove(peer_id);
    }

    fn inject_replaced(
        &mut self,
        _peer_id: PeerId,
        _closed: ConnectedPoint,
        _opened: ConnectedPoint,
    ) {
        // discv5 doesn't implement
    }

    fn inject_node_event(
        &mut self,
        _peer_id: PeerId,
        _event: <Self::ProtocolsHandler as ProtocolsHandler>::OutEvent,
    ) {
        // discv5 doesn't implement
    }

    fn poll(
        &mut self,
        params: &mut impl PollParameters,
    ) -> Async<
        NetworkBehaviourAction<
            <Self::ProtocolsHandler as ProtocolsHandler>::InEvent,
            Self::OutEvent,
        >,
    > {
        // search for peers if it is time
        loop {
            match self.peer_discovery_delay.poll() {
                Ok(Async::Ready(_)) => {
                    if self.connected_peers.len() < self.max_peers {
                        self.find_peers();
                    }
                }
                Ok(Async::NotReady) => break,
                Err(e) => {
                    warn!(self.log, "Discovery peer search failed: {:?}", e);
                }
            }
        }

        // Poll discovery
        loop {
            match self.discovery.poll(params) {
                Async::Ready(NetworkBehaviourAction::GenerateEvent(event)) => {
                    match event {
                        Discv5Event::Discovered(_enr) => {
                            // not concerned about FINDNODE results, rather the result of an entire
                            // query.
                        }
                        Discv5Event::SocketUpdated(socket) => {
                            info!(self.log, "Address updated"; "IP" => format!("{}",socket.ip()));
                            let mut address = Multiaddr::from(socket.ip());
                            address.push(Protocol::Tcp(self.tcp_port));
                            return Async::Ready(NetworkBehaviourAction::ReportObservedAddr {
                                address,
                            });
                        }
                        Discv5Event::FindNodeResult { closer_peers, .. } => {
                            debug!(self.log, "Discv5 query found {} peers", closer_peers.len());
                            if closer_peers.is_empty() {
                                debug!(self.log, "Discv5 random query yielded empty results");
                            }
                            for peer_id in closer_peers {
                                // if we need more peers, attempt a connection
                                if self.connected_peers.len() < self.max_peers
                                    && self.connected_peers.get(&peer_id).is_none()
                                {
                                    debug!(self.log, "Discv5: Peer discovered"; "Peer"=> format!("{:?}", peer_id));
                                    return Async::Ready(NetworkBehaviourAction::DialPeer {
                                        peer_id,
                                    });
                                }
                            }
                        }
                        _ => {}
                    }
                }
                // discv5 does not output any other NetworkBehaviourAction
                Async::Ready(_) => {}
                Async::NotReady => break,
            }
        }
        Async::NotReady
    }
}

/// Loads an ENR from file if it exists and matches the current NodeId and sequence number. If none
/// exists, generates a new one.
///
/// If an ENR exists, with the same NodeId and IP address, we use the disk-generated one as its
/// ENR sequence will be equal or higher than a newly generated one.
fn load_enr(
    local_key: &Keypair,
    config: &NetworkConfig,
    log: &slog::Logger,
) -> Result<Enr, String> {
    // Build the local ENR.
    // Note: Discovery should update the ENR record's IP to the external IP as seen by the
    // majority of our peers.
    let mut local_enr = EnrBuilder::new()
        .ip(config.discovery_address.into())
        .tcp(config.libp2p_port)
        .udp(config.discovery_port)
        .build(&local_key)
        .map_err(|e| format!("Could not build Local ENR: {:?}", e))?;

    let enr_f = config.network_dir.join(ENR_FILENAME);
    if let Ok(mut enr_file) = File::open(enr_f.clone()) {
        let mut enr_string = String::new();
        match enr_file.read_to_string(&mut enr_string) {
            Err(_) => debug!(log, "Could not read ENR from file"),
            Ok(_) => {
                match Enr::from_str(&enr_string) {
                    Ok(enr) => {
                        debug!(log, "ENR found in file: {:?}", enr_f);

                        if enr.node_id() == local_enr.node_id() {
                            if enr.ip() == config.discovery_address.into()
                                && enr.tcp() == Some(config.libp2p_port)
                                && enr.udp() == Some(config.discovery_port)
                            {
                                debug!(log, "ENR loaded from file");
                                // the stored ENR has the same configuration, use it
                                return Ok(enr);
                            }

                            // same node id, different configuration - update the sequence number
                            let new_seq_no = enr.seq().checked_add(1).ok_or_else(|| "ENR sequence number on file is too large. Remove it to generate a new NodeId")?;
                            local_enr.set_seq(new_seq_no, local_key).map_err(|e| {
                                format!("Could not update ENR sequence number: {:?}", e)
                            })?;
                            debug!(log, "ENR sequence number increased to: {}", new_seq_no);
                        }
                    }
                    Err(e) => {
                        warn!(log, "ENR from file could not be decoded: {:?}", e);
                    }
                }
            }
        }
    }

    // write ENR to disk
    let _ = std::fs::create_dir_all(&config.network_dir);
    match File::create(enr_f.clone())
        .and_then(|mut f| f.write_all(&local_enr.to_base64().as_bytes()))
    {
        Ok(_) => {
            debug!(log, "ENR written to disk");
        }
        Err(e) => {
            warn!(
                log,
                "Could not write ENR to file: {:?}. Error: {}", enr_f, e
            );
        }
    }
    Ok(local_enr)
}
<|endoftext|>


<commit_before>//! gpx handles parsing of GPX elements.

use errors::*;
use std::io::Read;
use xml::reader::XmlEvent;
use geo::Bbox;
use chrono::{DateTime, Utc};

use parser::bounds;
use parser::time;
use parser::string;
use parser::track;
use parser::metadata;
use parser::waypoint;
use parser::Context;

use Gpx;
use GpxVersion;
use Link;
use Person;
use Metadata;

enum ParseEvent {
    StartAuthor,      // GPX 1.0
    StartBounds,      // GPX 1.0
    StartDescription, // GPX 1.0
    StartEmail,       // GPX 1.0
    StartKeywords,    // GPX 1.0
    StartMetadata,
    StartName, // GPX 1.0
    StartTime, // GPX 1.0
    StartTrack,
    StartUrl,     // GPX 1.0
    StartUrlname, // GPX 1.0
    StartWaypoint,
    Ignore,
    EndGpx,
}

/// Convert the version string to the version enum
fn version_string_to_version(version_str: &str) -> Result<GpxVersion> {
    match version_str {
        "1.0" => Ok(GpxVersion::Gpx10),
        "1.1" => Ok(GpxVersion::Gpx11),
        version => Err(Error::from(format!("Unknown version {}", version))),
    }
}

/// consume consumes an entire GPX element.
pub fn consume<R: Read>(context: &mut Context<R>) -> Result<Gpx> {
    let mut gpx: Gpx = Default::default();

    let mut author: Option<String> = None;
    let mut url: Option<String> = None;
    let mut urlname: Option<String> = None;
    let mut email: Option<String> = None;
    let mut time: Option<DateTime<Utc>> = None;
    let mut bounds: Option<Bbox<f64>> = None;
    let mut name: Option<String> = None;
    let mut description: Option<String> = None;
    let mut keywords: Option<String> = None;

    loop {
        // Peep into the reader and see what type of event is next. Based on
        // that information, we'll either forward the event to a downstream
        // module or take the information for ourselves.
        let event: Result<ParseEvent> = {
            if let Some(next) = context.reader.peek() {
                match next {
                    &Ok(XmlEvent::StartElement {
                        ref name,
                        ref attributes,
                        ..
                    }) => match name.local_name.as_ref() {
                        "metadata" if context.version != GpxVersion::Gpx10 => {
                            Ok(ParseEvent::StartMetadata)
                        }
                        "trk" => Ok(ParseEvent::StartTrack),
                        "wpt" => Ok(ParseEvent::StartWaypoint),
                        "gpx" => {
                            if let Ok(version) = attributes
                                .iter()
                                .filter(|attr| attr.name.local_name == "version")
                                .nth(0)
                                .ok_or("no version found".to_owned())
                            {
                                gpx.version = version_string_to_version(&version.value)?;
                                context.version = gpx.version;
                                Ok(ParseEvent::Ignore)
                            } else {
                                Err(Error::from(ErrorKind::InvalidElementLacksAttribute(
                                    "version",
                                    "gpx",
                                )))
                            }
                        }
                        "time" if context.version == GpxVersion::Gpx10 => Ok(ParseEvent::StartTime),
                        "bounds" if context.version == GpxVersion::Gpx10 => {
                            Ok(ParseEvent::StartBounds)
                        }
                        "author" if context.version == GpxVersion::Gpx10 => {
                            Ok(ParseEvent::StartAuthor)
                        }
                        "email" if context.version == GpxVersion::Gpx10 => {
                            Ok(ParseEvent::StartEmail)
                        }
                        "url" if context.version == GpxVersion::Gpx10 => Ok(ParseEvent::StartUrl),
                        "urlname" if context.version == GpxVersion::Gpx10 => {
                            Ok(ParseEvent::StartUrlname)
                        }
                        "name" if context.version == GpxVersion::Gpx10 => Ok(ParseEvent::StartName),
                        "description" if context.version == GpxVersion::Gpx10 => {
                            Ok(ParseEvent::StartDescription)
                        }
                        "keywords" if context.version == GpxVersion::Gpx10 => {
                            Ok(ParseEvent::StartKeywords)
                        }
                        child => Err(Error::from(ErrorKind::InvalidChildElement(
                            String::from(child),
                            "gpx",
                        )))?,
                    },

                    &Ok(XmlEvent::EndElement { .. }) => Ok(ParseEvent::EndGpx),

                    _ => Ok(ParseEvent::Ignore),
                }
            } else {
                break;
            }
        };

        match event.chain_err(|| Error::from("error while parsing gpx event"))? {
            ParseEvent::Ignore => {
                context.reader.next();
            }

            ParseEvent::StartAuthor => {
                author = Some(string::consume(context)?);
            }

            ParseEvent::StartBounds => {
                bounds = Some(bounds::consume(context)?);
            }

            ParseEvent::StartEmail => {
                email = Some(string::consume(context)?);
            }

            ParseEvent::StartMetadata => {
                gpx.metadata = Some(metadata::consume(context)?);
            }

            ParseEvent::StartName => {
                name = Some(string::consume(context)?);
            }

            ParseEvent::StartTime => {
                time = Some(time::consume(context)?);
            }

            ParseEvent::StartUrl => {
                url = Some(string::consume(context)?);
            }

            ParseEvent::StartUrlname => {
                urlname = Some(string::consume(context)?);
            }

            ParseEvent::StartDescription => {
                description = Some(string::consume(context)?);
            }

            ParseEvent::StartKeywords => {
                keywords = Some(string::consume(context)?);
            }

            ParseEvent::StartTrack => {
                gpx.tracks.push(track::consume(context)?);
            }

            ParseEvent::StartWaypoint => {
                gpx.waypoints.push(waypoint::consume(context)?);
            }

            ParseEvent::EndGpx => {
                if gpx.version == GpxVersion::Gpx10 {
                    let mut metadata: Metadata = Default::default();
                    metadata.name = name;
                    metadata.time = time;
                    metadata.bounds = bounds;
                    let mut person: Person = Default::default();
                    person.name = author;
                    if let Some(url) = url {
                        let mut link: Link = Default::default();
                        link.href = url;
                        link.text = urlname;
                        person.link = Some(link);
                    }
                    person.email = email;
                    metadata.author = Some(person);
                    metadata.keywords = keywords;
                    metadata.description = description;
                    gpx.metadata = Some(metadata);
                }
                context.reader.next();

                return Ok(gpx);
            }
        }
    }

    return Err("no end tag for gpx".into());
}

#[cfg(test)]
mod tests {
    use std::io::BufReader;
    use xml::reader::EventReader;
    use geo::Point;

    use GpxVersion;
    use parser::Context;
    use super::consume;

    #[test]
    fn consume_gpx() {
        let gpx = consume!("<gpx version=\"1.1\"></gpx>", GpxVersion::Unknown);

        assert!(gpx.is_ok());
    }

    #[test]
    fn consume_gpx_no_version() {
        let gpx = consume!("<gpx></gpx>", GpxVersion::Unknown);

        assert!(gpx.is_err());
    }

    #[test]
    fn consume_gpx_version_error() {
        let gpx = consume!("<gpx version=\"1.2\"></gpx>", GpxVersion::Unknown);

        assert!(gpx.is_err());
    }

    #[test]
    fn consume_gpx_full() {
        let gpx = consume!(
            "
            <gpx version=\"1.0\">
                <time>2016-03-27T18:57:55Z</time>
                <bounds minlat=\"45.487064362\" minlon=\"-74.031837463\" maxlat=\"45.701225281\" maxlon=\"-73.586273193\"></bounds>
                <trk></trk>
                <wpt lat=\"1.23\" lon=\"2.34\"></wpt>
                <wpt lon=\"10.256\" lat=\"-81.324\">
                    <time>2001-10-26T19:32:52+00:00</time>
                </wpt>
            </gpx>
            ",
            GpxVersion::Unknown
        );

        assert!(gpx.is_ok());
        let gpx = gpx.unwrap();

        assert_eq!(gpx.version, GpxVersion::Gpx10);
        assert_eq!(gpx.tracks.len(), 1);

        assert_eq!(gpx.waypoints.len(), 2);

        let wpt = &gpx.waypoints[1];
        assert_eq!(wpt.point(), Point::new(10.256, -81.324));
    }
}
<commit_msg>Rewrite gpx::consume cleaner and more strict - Do not accept arbitrary amounts of opening gpx tags, we require exactly   one at the beginning and do not allow further ones - Do not accept any closing tag as tag that closes gpx, require name==gpx - Use new custom errors InvalidClosingTag and MissingClosingTag - Avoid indirection of ParseEvents, handle XmlEvents more directly<commit_after>//! gpx handles parsing of GPX elements.

use errors::*;
use std::io::Read;
use xml::reader::XmlEvent;
use geo::Bbox;
use chrono::{DateTime, Utc};

use parser::bounds;
use parser::time;
use parser::string;
use parser::track;
use parser::metadata;
use parser::waypoint;
use parser::Context;
use parser::verify_starting_tag;

use Gpx;
use GpxVersion;
use Link;
use Person;
use Metadata;

/// Convert the version string to the version enum
fn version_string_to_version(version_str: &str) -> Result<GpxVersion> {
    match version_str {
        "1.0" => Ok(GpxVersion::Gpx10),
        "1.1" => Ok(GpxVersion::Gpx11),
        version => Err(Error::from(format!("Unknown version {}", version))),
    }
}

/// consume consumes an entire GPX element.
pub fn consume<R: Read>(context: &mut Context<R>) -> Result<Gpx> {
    let mut gpx: Gpx = Default::default();

    let mut author: Option<String> = None;
    let mut url: Option<String> = None;
    let mut urlname: Option<String> = None;
    let mut email: Option<String> = None;
    let mut time: Option<DateTime<Utc>> = None;
    let mut bounds: Option<Bbox<f64>> = None;
    let mut gpx_name: Option<String> = None;
    let mut description: Option<String> = None;
    let mut keywords: Option<String> = None;

    // First we consume the gpx tag and its attributes
    let attributes = verify_starting_tag(context, "gpx")?;
    let version = attributes
        .iter()
        .filter(|attr| attr.name.local_name == "version")
        .nth(0)
        .ok_or(ErrorKind::InvalidElementLacksAttribute("version", "gpx"))?;
    gpx.version = version_string_to_version(&version.value)?;
    context.version = gpx.version;

    loop {
        let next_event = {
            if let Some(next) = context.reader.peek() {
                next.clone()
            } else {
                break;
            }
        };

        match next_event.chain_err(|| Error::from("error while parsing gpx event"))? {
            XmlEvent::StartElement { ref name, .. } => match name.local_name.as_ref() {
                "metadata" if context.version != GpxVersion::Gpx10 => {
                    gpx.metadata = Some(metadata::consume(context)?);
                }
                "trk" => {
                    gpx.tracks.push(track::consume(context)?);
                }
                "wpt" => {
                    gpx.waypoints.push(waypoint::consume(context)?);
                }
                "time" if context.version == GpxVersion::Gpx10 => {
                    time = Some(time::consume(context)?);
                }
                "bounds" if context.version == GpxVersion::Gpx10 => {
                    bounds = Some(bounds::consume(context)?);
                }
                "author" if context.version == GpxVersion::Gpx10 => {
                    author = Some(string::consume(context)?);
                }
                "email" if context.version == GpxVersion::Gpx10 => {
                    email = Some(string::consume(context)?);
                }
                "url" if context.version == GpxVersion::Gpx10 => {
                    url = Some(string::consume(context)?);
                }
                "urlname" if context.version == GpxVersion::Gpx10 => {
                    urlname = Some(string::consume(context)?);
                }
                "name" if context.version == GpxVersion::Gpx10 => {
                    gpx_name = Some(string::consume(context)?);
                }
                "description" if context.version == GpxVersion::Gpx10 => {
                    description = Some(string::consume(context)?);
                }
                "keywords" if context.version == GpxVersion::Gpx10 => {
                    keywords = Some(string::consume(context)?);
                }
                child => {
                    bail!(ErrorKind::InvalidChildElement(String::from(child), "gpx"));
                }
            },
            XmlEvent::EndElement { name } => {
                ensure!(
                    name.local_name == "gpx",
                    ErrorKind::InvalidClosingTag(name.local_name.clone(), "gpx")
                );
                if gpx.version == GpxVersion::Gpx10 {
                    let mut metadata: Metadata = Default::default();
                    metadata.name = gpx_name;
                    metadata.time = time;
                    metadata.bounds = bounds;
                    let mut person: Person = Default::default();
                    person.name = author;
                    if let Some(url) = url {
                        let mut link: Link = Default::default();
                        link.href = url;
                        link.text = urlname;
                        person.link = Some(link);
                    }
                    person.email = email;
                    metadata.author = Some(person);
                    metadata.keywords = keywords;
                    metadata.description = description;
                    gpx.metadata = Some(metadata);
                }
                context.reader.next();

                return Ok(gpx);
            }
            _ => {
                context.reader.next(); //consume and ignore this event
            }
        }
    }

    bail!(ErrorKind::MissingClosingTag("gpx"));
}

#[cfg(test)]
mod tests {
    use std::io::BufReader;
    use xml::reader::EventReader;
    use geo::Point;

    use GpxVersion;
    use parser::Context;
    use super::consume;

    #[test]
    fn consume_gpx() {
        let gpx = consume!("<gpx version=\"1.1\"></gpx>", GpxVersion::Unknown);

        assert!(gpx.is_ok());
    }

    #[test]
    fn consume_gpx_no_version() {
        let gpx = consume!("<gpx></gpx>", GpxVersion::Unknown);

        assert!(gpx.is_err());
    }

    #[test]
    fn consume_gpx_version_error() {
        let gpx = consume!("<gpx version=\"1.2\"></gpx>", GpxVersion::Unknown);

        assert!(gpx.is_err());
    }

    #[test]
    fn consume_gpx_full() {
        let gpx = consume!(
            "
            <gpx version=\"1.0\">
                <time>2016-03-27T18:57:55Z</time>
                <bounds minlat=\"45.487064362\" minlon=\"-74.031837463\" maxlat=\"45.701225281\" maxlon=\"-73.586273193\"></bounds>
                <trk></trk>
                <wpt lat=\"1.23\" lon=\"2.34\"></wpt>
                <wpt lon=\"10.256\" lat=\"-81.324\">
                    <time>2001-10-26T19:32:52+00:00</time>
                </wpt>
            </gpx>
            ",
            GpxVersion::Unknown
        );

        println!("########################### {:?}", gpx);
        assert!(gpx.is_ok());
        let gpx = gpx.unwrap();

        assert_eq!(gpx.version, GpxVersion::Gpx10);
        assert_eq!(gpx.tracks.len(), 1);

        assert_eq!(gpx.waypoints.len(), 2);

        let wpt = &gpx.waypoints[1];
        assert_eq!(wpt.point(), Point::new(10.256, -81.324));
    }
}
<|endoftext|>