<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use utils;

use std::path::{Path, PathBuf};
use std::collections::HashMap;

use syntax::ast;
use syntax::codemap;
use syntax::parse::parser;

/// List all the files containing modules of a crate.
/// If a file is used twice in a crate, it appears only once.
pub fn list_files<'a>(krate: &'a ast::Crate,
                      codemap: &codemap::CodeMap)
                      -> HashMap<PathBuf, &'a ast::Mod> {
    let mut result = HashMap::new();
    let root_filename: PathBuf = codemap.span_to_filename(krate.span).into();
    list_submodules(&krate.module,
                    root_filename.parent().unwrap(),
                    codemap,
                    &mut result);
    result.insert(root_filename, &krate.module);
    result
}

/// Recursively list all external modules included in a module.
fn list_submodules<'a>(module: &'a ast::Mod,
                       search_dir: &Path,
                       codemap: &codemap::CodeMap,
                       result: &mut HashMap<PathBuf, &'a ast::Mod>) {
    debug!("list_submodules: search_dir: {:?}", search_dir);
    for item in &module.items {
        if let ast::ItemKind::Mod(ref sub_mod) = item.node {
            if !utils::contains_skip(&item.attrs) {
                let is_internal = codemap.span_to_filename(item.span) ==
                                  codemap.span_to_filename(sub_mod.inner);
                let dir_path = if is_internal {
                    search_dir.join(&item.ident.to_string())
                } else {
                    let mod_path = module_file(item.ident, &item.attrs, search_dir, codemap);
                    let dir_path = mod_path.parent().unwrap().to_owned();
                    result.insert(mod_path, sub_mod);
                    dir_path
                };
                list_submodules(sub_mod, &dir_path, codemap, result);
            }
        }
    }
}

/// Find the file corresponding to an external mod
fn module_file(id: ast::Ident,
               attrs: &[ast::Attribute],
               dir_path: &Path,
               codemap: &codemap::CodeMap)
               -> PathBuf {
    if let Some(path) = parser::Parser::submod_path_from_attr(attrs, dir_path) {
        return path;
    }

    match parser::Parser::default_submod_path(id, dir_path, codemap).result {
        Ok(parser::ModulePathSuccess { path, .. }) => path,
        Err(_) => panic!("Couldn't find module {}", id),
    }
}
<commit_msg>Make list_files deterministic (Closes #1141) (#1393)<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use utils;

use std::path::{Path, PathBuf};
use std::collections::BTreeMap;

use syntax::ast;
use syntax::codemap;
use syntax::parse::parser;

/// List all the files containing modules of a crate.
/// If a file is used twice in a crate, it appears only once.
pub fn list_files<'a>(krate: &'a ast::Crate,
                      codemap: &codemap::CodeMap)
                      -> BTreeMap<PathBuf, &'a ast::Mod> {
    let mut result = BTreeMap::new(); // Enforce file order determinism
    let root_filename: PathBuf = codemap.span_to_filename(krate.span).into();
    list_submodules(&krate.module,
                    root_filename.parent().unwrap(),
                    codemap,
                    &mut result);
    result.insert(root_filename, &krate.module);
    result
}

/// Recursively list all external modules included in a module.
fn list_submodules<'a>(module: &'a ast::Mod,
                       search_dir: &Path,
                       codemap: &codemap::CodeMap,
                       result: &mut BTreeMap<PathBuf, &'a ast::Mod>) {
    debug!("list_submodules: search_dir: {:?}", search_dir);
    for item in &module.items {
        if let ast::ItemKind::Mod(ref sub_mod) = item.node {
            if !utils::contains_skip(&item.attrs) {
                let is_internal = codemap.span_to_filename(item.span) ==
                                  codemap.span_to_filename(sub_mod.inner);
                let dir_path = if is_internal {
                    search_dir.join(&item.ident.to_string())
                } else {
                    let mod_path = module_file(item.ident, &item.attrs, search_dir, codemap);
                    let dir_path = mod_path.parent().unwrap().to_owned();
                    result.insert(mod_path, sub_mod);
                    dir_path
                };
                list_submodules(sub_mod, &dir_path, codemap, result);
            }
        }
    }
}

/// Find the file corresponding to an external mod
fn module_file(id: ast::Ident,
               attrs: &[ast::Attribute],
               dir_path: &Path,
               codemap: &codemap::CodeMap)
               -> PathBuf {
    if let Some(path) = parser::Parser::submod_path_from_attr(attrs, dir_path) {
        return path;
    }

    match parser::Parser::default_submod_path(id, dir_path, codemap).result {
        Ok(parser::ModulePathSuccess { path, .. }) => path,
        Err(_) => panic!("Couldn't find module {}", id),
    }
}
<|endoftext|>


<commit_before>use sync::Arc;

use vec::Vec;
use vec_deque::VecDeque;

use cell::GlobalCell;

use sync::spin::{InterruptSpinLock, RwLock};

use thread::WaitQueue;

use net::Interface;

use hal::net::discover;

use net::Packet;

const MAX_QUEUE_SIZE: usize = 512;

pub static STACK: GlobalCell<StackImpl> = GlobalCell::new();

/// Uni.rs network stack
pub struct Stack;

impl Stack {
    #[doc(hidden)]
    pub fn init() {
        STACK.set(StackImpl::new());
    }

    #[inline]
    /// Enqueue a packet inside the `rx_queue` and notify `rx_wait`.
    ///
    /// Returns false if the `rx_queue` is full and therefor was not enqueued,
    /// true otherwise.
    ///
    /// Note: This is safe to be called from interrupt context. Indeed the
    /// `rx_queue` is not resizable so no allocation will be performed by this
    /// function.
    pub fn enqueue_rx_packet(intf: *const Arc<RwLock<Interface>>,
                             packet: Packet) -> bool {
        STACK.as_mut().enqueue_rx_packet(intf, packet)
    }
}

pub struct StackImpl {
    interfaces: Vec<Arc<RwLock<Interface>>>,
    /// Contains packets to be processed
    rx_queue: InterruptSpinLock<VecDeque<(*const Arc<RwLock<Interface>>, Packet)>>,
    /// Used to wait for packet to arrive in the rx_queue
    rx_wait: WaitQueue,
}

// rx_wait is Sync
// rx_queue is protected by a spin lock
unsafe impl Sync for StackImpl {}

impl StackImpl {
    pub fn new() -> Self {
        let intfs = discover();

        if intfs.is_empty() {
            println!("Warning: Uni.rs is built with network capabilities but no interface found");
        } else {
            println!("{} interface(s) discovered:", intfs.len());

            for i in &intfs {
                println!("  - {} ({})", i.read().name_ref(),
                         i.read().hw_addr_ref());
            }
        }

        StackImpl {
            interfaces: intfs,
            rx_queue: InterruptSpinLock::new(VecDeque::with_capacity(MAX_QUEUE_SIZE)),
            rx_wait: WaitQueue::new(),
        }
    }

    pub fn enqueue_rx_packet(&mut self, intf: *const Arc<RwLock<Interface>>,
                             packet: Packet) -> bool {
        let mut locked_rx_queue = self.rx_queue.lock();

        if locked_rx_queue.len() == MAX_QUEUE_SIZE {
            // Queue is full, we don't want to cause a reallocation in
            // interruption context. So we don't enqueue the packet
            return false;
        }

        // Enqueue the packet
        locked_rx_queue.push_back((intf, packet));

        // Wake up network thread
        self.rx_wait.unblock();

        true
    }
}
<commit_msg>net: Add Stack::interfaces()<commit_after>use sync::Arc;

use vec::Vec;
use vec_deque::VecDeque;

use cell::{GlobalCell, GlobalCellRef};

use sync::spin::{InterruptSpinLock, RwLock};

use thread::WaitQueue;

use net::Interface;

use hal::net::discover;

use net::Packet;

const MAX_QUEUE_SIZE: usize = 512;

pub static STACK: GlobalCell<StackImpl> = GlobalCell::new();

/// Uni.rs network stack
pub struct Stack;

impl Stack {
    #[doc(hidden)]
    pub fn init() {
        STACK.set(StackImpl::new());
    }

    /// Returns interfaces registered in the network stack
    pub fn interfaces<'a>() -> Interfaces<'a> {
        Interfaces {
            imp: STACK.as_ref(),
        }
    }

    #[inline]
    /// Enqueue a packet inside the `rx_queue` and notify `rx_wait`.
    ///
    /// Returns false if the `rx_queue` is full and therefor was not enqueued,
    /// true otherwise.
    ///
    /// Note: This is safe to be called from interrupt context. Indeed the
    /// `rx_queue` is not resizable so no allocation will be performed by this
    /// function.
    pub fn enqueue_rx_packet(intf: *const Arc<RwLock<Interface>>,
                             packet: Packet) -> bool {
        STACK.as_mut().enqueue_rx_packet(intf, packet)
    }
}

/// Reference over the list of interfaces detected
pub struct Interfaces<'a> {
    imp: GlobalCellRef<'a, StackImpl>,
}

impl<'a> Interfaces<'a> {
    #[inline]
    /// Returns the number of interfaces detected
    pub fn count(&self) -> usize {
        self.imp.interfaces.len()
    }

    #[inline]
    /// Returns the list of interfaces as an immutable slice
    pub fn as_slice(&self) -> &[Arc<RwLock<Interface>>] {
        &self.imp.interfaces[..]
    }
}

pub struct StackImpl {
    /// Interfaces registered
    interfaces: Vec<Arc<RwLock<Interface>>>,
    /// Contains packets to be processed
    rx_queue: InterruptSpinLock<VecDeque<(*const Arc<RwLock<Interface>>, Packet)>>,
    /// Used to wait for packet to arrive in the rx_queue
    rx_wait: WaitQueue,
}

// rx_wait is Sync
// rx_queue is protected by a spin lock
unsafe impl Sync for StackImpl {}

impl StackImpl {
    pub fn new() -> Self {
        let intfs = discover();

        if intfs.is_empty() {
            println!("Warning: Uni.rs is built with network capabilities but no interface found");
        } else {
            println!("{} interface(s) discovered:", intfs.len());

            for i in &intfs {
                println!("  - {} ({})", i.read().name_ref(),
                         i.read().hw_addr_ref());
            }
        }

        StackImpl {
            interfaces: intfs,
            rx_queue: InterruptSpinLock::new(VecDeque::with_capacity(MAX_QUEUE_SIZE)),
            rx_wait: WaitQueue::new(),
        }
    }

    pub fn enqueue_rx_packet(&mut self, intf: *const Arc<RwLock<Interface>>,
                             packet: Packet) -> bool {
        let mut locked_rx_queue = self.rx_queue.lock();

        if locked_rx_queue.len() == MAX_QUEUE_SIZE {
            // Queue is full, we don't want to cause a reallocation in
            // interruption context. So we don't enqueue the packet
            return false;
        }

        // Enqueue the packet
        locked_rx_queue.push_back((intf, packet));

        // Wake up network thread
        self.rx_wait.unblock();

        true
    }
}
<|endoftext|>


<commit_before>// The MIT License (MIT)

// Copyright (c) 2015 Rustcc Developers

// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

use std::io;
use std::net::{ToSocketAddrs, SocketAddr};
use std::ops::{Deref, DerefMut};
use std::convert::From;
use std::iter::Iterator;

use mio::{self, Interest};

use processor::Processor;

#[derive(Debug)]
pub struct TcpSocket(::mio::tcp::TcpSocket);

impl TcpSocket {
    /// Returns a new, unbound, non-blocking, IPv4 socket
    pub fn v4() -> io::Result<TcpSocket> {
        Ok(TcpSocket(try!(::mio::tcp::TcpSocket::v4())))
    }

    /// Returns a new, unbound, non-blocking, IPv6 socket
    pub fn v6() -> io::Result<TcpSocket> {
        Ok(TcpSocket(try!(::mio::tcp::TcpSocket::v6())))
    }

    pub fn connect<A: ToSocketAddrs>(addr: A) -> io::Result<(TcpStream, bool)> {
        super::each_addr(addr, |a| {
            match a {
                &SocketAddr::V4(..) => try!(TcpSocket::v4()).0.connect(a),
                &SocketAddr::V6(..) => try!(TcpSocket::v6()).0.connect(a),
            }
        }).map(|(stream, complete)| (TcpStream(stream), complete))
    }

    pub fn listen(self, backlog: usize) -> io::Result<TcpListener> {
        Ok(TcpListener(try!(self.0.listen(backlog))))
    }
}

impl Deref for TcpSocket {
    type Target = ::mio::tcp::TcpSocket;

    fn deref(&self) -> &::mio::tcp::TcpSocket {
        &self.0
    }
}

impl DerefMut for TcpSocket {
    fn deref_mut(&mut self) -> &mut ::mio::tcp::TcpSocket {
        &mut self.0
    }
}

#[derive(Debug)]
pub struct TcpListener(::mio::tcp::TcpListener);

impl TcpListener {
    pub fn bind<A: ToSocketAddrs>(addr: A) -> io::Result<TcpListener> {
        super::each_addr(addr, ::mio::tcp::TcpListener::bind).map(TcpListener)
    }

    pub fn accept(&self) -> io::Result<TcpStream> {
        match self.0.accept() {
            Ok(None) => {
                debug!("accept WouldBlock; going to register into eventloop");
            },
            Ok(Some(stream)) => {
                return Ok(TcpStream(stream));
            },
            Err(err) => {
                return Err(err);
            }
        }

        loop {
            try!(Processor::current().wait_event(&self.0, Interest::readable()));

            match self.0.accept() {
                Ok(None) => {
                    warn!("accept WouldBlock; Coroutine was awaked by readable event");
                },
                Ok(Some(stream)) => {
                    return Ok(TcpStream(stream));
                },
                Err(err) => {
                    return Err(err);
                }
            }
        }
    }

    pub fn try_clone(&self) -> io::Result<TcpListener> {
        Ok(TcpListener(try!(self.0.try_clone())))
    }

    pub fn incoming<'a>(&'a self) -> Incoming<'a> {
        Incoming(self)
    }
}

impl Deref for TcpListener {
    type Target = ::mio::tcp::TcpListener;

    fn deref(&self) -> &::mio::tcp::TcpListener {
        &self.0
    }
}

impl DerefMut for TcpListener {
    fn deref_mut(&mut self) -> &mut ::mio::tcp::TcpListener {
        &mut self.0
    }
}

pub struct Incoming<'a>(&'a TcpListener);

impl<'a> Iterator for Incoming<'a> {
    type Item = io::Result<TcpStream>;

    fn next(&mut self) -> Option<io::Result<TcpStream>> {
        Some(self.0.accept())
    }
}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Shutdown {
    /// Further receptions will be disallowed.
    Read,
    /// Further  transmissions will be disallowed.
    Write,
    /// Further receptions and transmissions will be disallowed.
    Both,
}

impl From<Shutdown> for mio::tcp::Shutdown {
    fn from(shutdown: Shutdown) -> mio::tcp::Shutdown {
        match shutdown {
            Shutdown::Read => mio::tcp::Shutdown::Read,
            Shutdown::Write => mio::tcp::Shutdown::Write,
            Shutdown::Both => mio::tcp::Shutdown::Both,
        }
    }
}

#[derive(Debug)]
pub struct TcpStream(mio::tcp::TcpStream);

impl TcpStream {
    pub fn connect<A: ToSocketAddrs>(addr: A) -> io::Result<TcpStream> {
        match TcpSocket::connect(addr) {
            Ok((stream, completed)) => {
                if !completed {
                    try!(Processor::current().wait_event(&stream.0, Interest::writable()));
                    try!(stream.take_socket_error());
                }

                Ok(stream)
            },
            Err(err) => Err(err)
        }
    }

    pub fn peer_addr(&self) -> io::Result<SocketAddr> {
        self.0.peer_addr()
    }

    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.0.local_addr()
    }

    pub fn try_clone(&self) -> io::Result<TcpStream> {
        let stream = try!(self.0.try_clone());

        Ok(TcpStream(stream))
    }

    pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
        self.0.shutdown(From::from(how))
    }

    pub fn take_socket_error(&self) -> io::Result<()> {
        self.0.take_socket_error()
    }
}

impl io::Read for TcpStream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        use mio::TryRead;

        match self.0.try_read(buf) {
            Ok(None) => {
                debug!("TcpStream read WouldBlock");
            },
            Ok(Some(0)) => {
                debug!("TcpStream read 0 bytes; may be EOF");
                return Ok(0);
            },
            Ok(Some(len)) => {
                debug!("TcpStream read {} bytes", len);
                return Ok(len);
            },
            Err(err) => {
                return Err(err);
            }
        }

        loop {
            debug!("Read: Going to register event");
            try!(Processor::current().wait_event(&self.0, Interest::readable()));
            debug!("Read: Got read event");

            match self.0.try_read(buf) {
                Ok(None) => {
                    warn!("TcpStream read WouldBlock");
                },
                Ok(Some(0)) => {
                    debug!("TcpStream read 0 bytes; may be EOF");
                    return Ok(0);
                },
                Ok(Some(len)) => {
                    debug!("TcpStream read {} bytes", len);
                    return Ok(len);
                },
                Err(err) => {
                    return Err(err);
                }
            }
        }
    }
}

impl io::Write for TcpStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        use mio::TryWrite;

        match self.0.try_write(buf) {
            Ok(None) => {
                debug!("TcpStream write WouldBlock");
            },
            Ok(Some(0)) => {
                debug!("TcpStream write 0 bytes; may be EOF");
                return Ok(0);
            },
            Ok(Some(len)) => {
                debug!("TcpStream written {} bytes", len);
                return Ok(len);
            },
            Err(err) => {
                return Err(err)
            }
        }

        loop {
            debug!("Write: Going to register event");
            try!(Processor::current().wait_event(&self.0, Interest::writable()));
            debug!("Write: Got write event");

            match self.0.try_write(buf) {
                Ok(None) => {
                    warn!("TcpStream write WouldBlock");
                },
                Ok(Some(0)) => {
                    debug!("TcpStream write 0 bytes; may be EOF");
                    return Ok(0);
                },
                Ok(Some(len)) => {
                    debug!("TcpStream written {} bytes", len);
                    return Ok(len);
                },
                Err(err) => {
                    return Err(err)
                }
            }
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl Deref for TcpStream {
    type Target = ::mio::tcp::TcpStream;

    fn deref(&self) -> &::mio::tcp::TcpStream {
        &self.0
    }
}

impl DerefMut for TcpStream {
    fn deref_mut(&mut self) -> &mut ::mio::tcp::TcpStream {
        &mut self.0
    }
}
<commit_msg>lower the log level<commit_after>// The MIT License (MIT)

// Copyright (c) 2015 Rustcc Developers

// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

use std::io;
use std::net::{ToSocketAddrs, SocketAddr};
use std::ops::{Deref, DerefMut};
use std::convert::From;
use std::iter::Iterator;

use mio::{self, Interest};

use processor::Processor;

#[derive(Debug)]
pub struct TcpSocket(::mio::tcp::TcpSocket);

impl TcpSocket {
    /// Returns a new, unbound, non-blocking, IPv4 socket
    pub fn v4() -> io::Result<TcpSocket> {
        Ok(TcpSocket(try!(::mio::tcp::TcpSocket::v4())))
    }

    /// Returns a new, unbound, non-blocking, IPv6 socket
    pub fn v6() -> io::Result<TcpSocket> {
        Ok(TcpSocket(try!(::mio::tcp::TcpSocket::v6())))
    }

    pub fn connect<A: ToSocketAddrs>(addr: A) -> io::Result<(TcpStream, bool)> {
        super::each_addr(addr, |a| {
            match a {
                &SocketAddr::V4(..) => try!(TcpSocket::v4()).0.connect(a),
                &SocketAddr::V6(..) => try!(TcpSocket::v6()).0.connect(a),
            }
        }).map(|(stream, complete)| (TcpStream(stream), complete))
    }

    pub fn listen(self, backlog: usize) -> io::Result<TcpListener> {
        Ok(TcpListener(try!(self.0.listen(backlog))))
    }
}

impl Deref for TcpSocket {
    type Target = ::mio::tcp::TcpSocket;

    fn deref(&self) -> &::mio::tcp::TcpSocket {
        &self.0
    }
}

impl DerefMut for TcpSocket {
    fn deref_mut(&mut self) -> &mut ::mio::tcp::TcpSocket {
        &mut self.0
    }
}

#[derive(Debug)]
pub struct TcpListener(::mio::tcp::TcpListener);

impl TcpListener {
    pub fn bind<A: ToSocketAddrs>(addr: A) -> io::Result<TcpListener> {
        super::each_addr(addr, ::mio::tcp::TcpListener::bind).map(TcpListener)
    }

    pub fn accept(&self) -> io::Result<TcpStream> {
        match self.0.accept() {
            Ok(None) => {
                debug!("accept WouldBlock; going to register into eventloop");
            },
            Ok(Some(stream)) => {
                return Ok(TcpStream(stream));
            },
            Err(err) => {
                return Err(err);
            }
        }

        loop {
            try!(Processor::current().wait_event(&self.0, Interest::readable()));

            match self.0.accept() {
                Ok(None) => {
                    warn!("accept WouldBlock; Coroutine was awaked by readable event");
                },
                Ok(Some(stream)) => {
                    return Ok(TcpStream(stream));
                },
                Err(err) => {
                    return Err(err);
                }
            }
        }
    }

    pub fn try_clone(&self) -> io::Result<TcpListener> {
        Ok(TcpListener(try!(self.0.try_clone())))
    }

    pub fn incoming<'a>(&'a self) -> Incoming<'a> {
        Incoming(self)
    }
}

impl Deref for TcpListener {
    type Target = ::mio::tcp::TcpListener;

    fn deref(&self) -> &::mio::tcp::TcpListener {
        &self.0
    }
}

impl DerefMut for TcpListener {
    fn deref_mut(&mut self) -> &mut ::mio::tcp::TcpListener {
        &mut self.0
    }
}

pub struct Incoming<'a>(&'a TcpListener);

impl<'a> Iterator for Incoming<'a> {
    type Item = io::Result<TcpStream>;

    fn next(&mut self) -> Option<io::Result<TcpStream>> {
        Some(self.0.accept())
    }
}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Shutdown {
    /// Further receptions will be disallowed.
    Read,
    /// Further  transmissions will be disallowed.
    Write,
    /// Further receptions and transmissions will be disallowed.
    Both,
}

impl From<Shutdown> for mio::tcp::Shutdown {
    fn from(shutdown: Shutdown) -> mio::tcp::Shutdown {
        match shutdown {
            Shutdown::Read => mio::tcp::Shutdown::Read,
            Shutdown::Write => mio::tcp::Shutdown::Write,
            Shutdown::Both => mio::tcp::Shutdown::Both,
        }
    }
}

#[derive(Debug)]
pub struct TcpStream(mio::tcp::TcpStream);

impl TcpStream {
    pub fn connect<A: ToSocketAddrs>(addr: A) -> io::Result<TcpStream> {
        match TcpSocket::connect(addr) {
            Ok((stream, completed)) => {
                if !completed {
                    try!(Processor::current().wait_event(&stream.0, Interest::writable()));
                    try!(stream.take_socket_error());
                }

                Ok(stream)
            },
            Err(err) => Err(err)
        }
    }

    pub fn peer_addr(&self) -> io::Result<SocketAddr> {
        self.0.peer_addr()
    }

    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.0.local_addr()
    }

    pub fn try_clone(&self) -> io::Result<TcpStream> {
        let stream = try!(self.0.try_clone());

        Ok(TcpStream(stream))
    }

    pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
        self.0.shutdown(From::from(how))
    }

    pub fn take_socket_error(&self) -> io::Result<()> {
        self.0.take_socket_error()
    }
}

impl io::Read for TcpStream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        use mio::TryRead;

        match self.0.try_read(buf) {
            Ok(None) => {
                debug!("TcpStream read WouldBlock");
            },
            Ok(Some(0)) => {
                debug!("TcpStream read 0 bytes; may be EOF");
                return Ok(0);
            },
            Ok(Some(len)) => {
                debug!("TcpStream read {} bytes", len);
                return Ok(len);
            },
            Err(err) => {
                return Err(err);
            }
        }

        loop {
            debug!("Read: Going to register event");
            try!(Processor::current().wait_event(&self.0, Interest::readable()));
            debug!("Read: Got read event");

            match self.0.try_read(buf) {
                Ok(None) => {
                    debug!("TcpStream read WouldBlock");
                },
                Ok(Some(0)) => {
                    debug!("TcpStream read 0 bytes; may be EOF");
                    return Ok(0);
                },
                Ok(Some(len)) => {
                    debug!("TcpStream read {} bytes", len);
                    return Ok(len);
                },
                Err(err) => {
                    return Err(err);
                }
            }
        }
    }
}

impl io::Write for TcpStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        use mio::TryWrite;

        match self.0.try_write(buf) {
            Ok(None) => {
                debug!("TcpStream write WouldBlock");
            },
            Ok(Some(0)) => {
                debug!("TcpStream write 0 bytes; may be EOF");
                return Ok(0);
            },
            Ok(Some(len)) => {
                debug!("TcpStream written {} bytes", len);
                return Ok(len);
            },
            Err(err) => {
                return Err(err)
            }
        }

        loop {
            debug!("Write: Going to register event");
            try!(Processor::current().wait_event(&self.0, Interest::writable()));
            debug!("Write: Got write event");

            match self.0.try_write(buf) {
                Ok(None) => {
                    debug!("TcpStream write WouldBlock");
                },
                Ok(Some(0)) => {
                    debug!("TcpStream write 0 bytes; may be EOF");
                    return Ok(0);
                },
                Ok(Some(len)) => {
                    debug!("TcpStream written {} bytes", len);
                    return Ok(len);
                },
                Err(err) => {
                    return Err(err)
                }
            }
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl Deref for TcpStream {
    type Target = ::mio::tcp::TcpStream;

    fn deref(&self) -> &::mio::tcp::TcpStream {
        &self.0
    }
}

impl DerefMut for TcpStream {
    fn deref_mut(&mut self) -> &mut ::mio::tcp::TcpStream {
        &mut self.0
    }
}
<|endoftext|>


<commit_before>use address::Address;
use messages::outgoing;
use mio;
use node;
use routing_table::{InsertOutcome, InsertionResult, RoutingTable};
use std::collections::HashMap;
use std::io::Cursor;
use std::net::{UdpSocket, ToSocketAddrs, SocketAddr};
use std::sync::mpsc;
use std::thread;
use stun;
use transaction::{TransactionId, TransactionIdGenerator};

#[derive(Debug)]
pub enum Event {
    ReceivedPacket(Address, Vec<u8>),
    Shutdown
}

pub enum OneshotTask {
    Incoming(Vec<u8>),
    StartBootstrap,
    SendPacket(Address, Vec<u8>),
    Shutdown
}

pub enum ScheduledTask {
    ContinueBootstrap,
    ContinueHealthCheck,
    ContinueRefresh
}

enum TableAction {
    Bootstrap(mio::Timeout),
    HealthCheck(mio::Timeout),
    RefreshBucket(mio::Timeout)
}

pub type TaskSender = mio::Sender<OneshotTask>;

enum Status {
    Bootstrapping,
    Idle
}

pub struct Network {
    host: SocketAddr,
    routing_table: RoutingTable,
    self_node: Box<node::Node + 'static>,
    transaction_ids: TransactionIdGenerator,
    status: Status,
    pending_actions: HashMap<TransactionId, TableAction>,
    event_listeners: Vec<mpsc::Sender<Event>>,
    shutdown_signals: Vec<mpsc::Sender<mpsc::Sender<()>>>
}

impl Network {
    pub fn new<T: ToSocketAddrs>(self_address: Address, host: T, routers: Vec<Box<node::Node>>) -> Network {
        let host = host.to_socket_addrs().unwrap().next().unwrap();
        let mapped_host = stun::get_mapped_address(host).unwrap();
        let self_node = node::UdpNode::new(self_address, mapped_host);
        let routing_table = RoutingTable::new(8, self_address, routers);

        Network {
            host: host,
            routing_table: routing_table,
            self_node: Box::new(self_node),
            transaction_ids: TransactionIdGenerator::new(),
            status: Status::Idle,
            pending_actions: HashMap::new(),
            event_listeners: vec![],
            shutdown_signals: vec![]
        }
    }

    pub fn run(mut self) -> TaskSender {
        let mut event_loop_config = mio::EventLoopConfig::new();
        event_loop_config.notify_capacity(16384);
        let mut event_loop = mio::EventLoop::configured(event_loop_config).unwrap();

        let shutdown_signal = create_incoming_udp_channel(self.host, event_loop.channel());
        self.shutdown_signals.push(shutdown_signal);

        event_loop.channel().send(OneshotTask::StartBootstrap).unwrap();
        info!("Running server at {:?}", self.self_node);
        let mut handler = Handler::new(self);
        let task_sender = event_loop.channel();
        thread::spawn(move|| event_loop.run(&mut handler).unwrap());
        task_sender
    }

    pub fn register_event_listener(&mut self, event_listener: mpsc::Sender<Event>) {
        self.event_listeners.push(event_listener);
    }

    fn handle_incoming(&mut self, data: Vec<u8>, event_loop: &mut mio::EventLoop<Handler>) {
        use messages::incoming::{Message, Query, Response, self};
        let mut data = Cursor::new(data);
        let message = incoming::parse_from_reader(&mut data).unwrap();

        match message {
            Message::Query(transaction_id, origin, query) => {
                let origin_address = origin.address();

                match query {
                    Query::FindNode(target) => {
                        let response: Vec<u8> = outgoing::create_find_node_response(
                                transaction_id,
                                &self.self_node,
                                self.routing_table.nearest_live_nodes_to(&target, false));
                        origin.send(response);
                    },
                    Query::Packet(payload) => {
                        for listener in &self.event_listeners {
                            listener.send(Event::ReceivedPacket(origin_address, payload.clone())).unwrap();
                        }
                        let response = outgoing::create_packet_response(
                            transaction_id, &self.self_node);
                        origin.send(response);
                    },
                    Query::Ping => {
                        let response = outgoing::create_ping_response(
                            transaction_id,
                            &self.self_node);
                        origin.send(response);
                    }
                }

                self.insert_node(origin).unwrap();
                if let Some(origin) = self.routing_table.find_node(&origin_address) {
                    origin.received_query(transaction_id);
                }
            }

            Message::Response(transaction_id, origin, response) => {
                let origin_address = origin.address();
                match response {
                    Response::FindNode(mut nodes) => {
                        let mut encounted_new_node = false;
                        let mut tail = vec![origin];
                        for node in nodes.drain(..).chain(tail.drain(..)) {
                            match self.insert_node(node) {
                                Ok(InsertOutcome::Inserted) => {
                                    encounted_new_node = true;
                                }
                                Err(error) => { panic!(error) }
                                _ => { }
                            }
                        }

                        {
                            if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_response(transaction_id);
                            }
                        }

                        match self.pending_actions.remove(&transaction_id) {
                            Some(TableAction::Bootstrap(timeout)) => {
                                event_loop.clear_timeout(timeout);
                                if encounted_new_node {
                                    self.continue_bootstrap(event_loop);
                                } else {
                                    self.continue_health_check(event_loop);
                                    self.continue_refresh(event_loop)
                                }
                            }
                            Some(TableAction::HealthCheck(_)) => {
                                // Does not clear timeout. Let's Healthcheck proceed at regular
                                // interval.
                            },
                            Some(TableAction::RefreshBucket(_)) => {
                            }
                            None => { }
                        }

                        match self.status {
                            Status::Bootstrapping => {
                                self.status = Status::Idle;
                            }
                            _ => { }
                        }
                    }
                    Response::Packet => {
                        if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_response(transaction_id);
                        }
                    }
                    Response::Ping => {
                        if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_response(transaction_id);
                        }
                    }
                }
            }
        }
    }

    fn start_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        self.status = Status::Bootstrapping;
        self.continue_bootstrap(event_loop);
    }

    fn continue_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let address = &self.self_node.address();
        let transaction_id = self.find_node(address);
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueBootstrap, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::Bootstrap(timeout));
    }

    fn continue_health_check(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let transaction_id = self.health_check();
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueHealthCheck, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::HealthCheck(timeout));
    }

    fn continue_refresh(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueRefresh, 1000).unwrap();

        match self.refresh_bucket() {
            Some(transaction_id) => {
                self.pending_actions.insert(transaction_id, TableAction::RefreshBucket(timeout));
            }
            None => { }
        }
    }

    fn find_node(&mut self, address: &Address) -> TransactionId {
        let transaction_id = self.transaction_ids.generate();
        let query = outgoing::create_find_node_query(
            transaction_id,
            &self.self_node,
            address);
        for node in self.routing_table.nearest() {
            node.sent_query(transaction_id);
            node.send(query.clone());
        }
        transaction_id
    }

    fn health_check(&mut self) -> TransactionId {
        let transaction_id = self.transaction_ids.generate();
        for node in self.routing_table.questionable_nodes().iter_mut() {
            let query = outgoing::create_ping_query(
                transaction_id, &self.self_node);
            node.sent_query(transaction_id);
            node.send(query);
        }
        transaction_id
    }

    fn refresh_bucket(&mut self) -> Option<TransactionId> {
        if let Some(address) = self.address_to_find_for_refresh() {
            Some(self.find_node(&address))
        } else {
            None
        }
    }

    fn address_to_find_for_refresh(&self) -> Option<Address> {
        match self.routing_table.bucket_needing_refresh() {
            Some(bucket) => {
                Some(bucket.random_address_in_space())
            }
            None => None
        }
    }

    fn insert_node(&mut self, node: Box<node::Node>) -> InsertionResult {
        self.routing_table.insert(node, &self.self_node, &mut self.transaction_ids)
    }

    fn send_packet(&mut self, recipient: Address, payload: Vec<u8>, _event_loop: &mut mio::EventLoop<Handler>) {
        for node in self.routing_table.nearest_live_nodes_to(&recipient, false) {
            let transaction_id = self.transaction_ids.generate();
            let query = outgoing::create_packet_query(
                transaction_id, &self.self_node, payload.clone());
            node.sent_query(transaction_id);
            node.send(query.clone());
        }
    }

    fn shutdown(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let sentinels: Vec<mpsc::Receiver<()>> = self.shutdown_signals.iter().map(|signal| {
            let (sentinal_sender, sentinal_receiver) = mpsc::channel();
            signal.send(sentinal_sender).expect("couldn't send shutdown signal");
            sentinal_receiver
        }).collect();
        for sentinel in sentinels {
            sentinel.recv().unwrap();
        }
        event_loop.shutdown();
        for listener in &self.event_listeners {
            listener.send(Event::Shutdown).unwrap();
        }
    }
}

struct Handler {
    network: Network
}

impl Handler {
    fn new(network: Network) -> Handler {
        Handler {
            network: network
        }
    }
}

impl mio::Handler for Handler {
    type Timeout = ScheduledTask;
    type Message = OneshotTask;

    fn notify(&mut self, event_loop: &mut mio::EventLoop<Handler>, task: OneshotTask) {
        match task {
            OneshotTask::Incoming(data) => self.network.handle_incoming(data, event_loop),
            OneshotTask::StartBootstrap => self.network.start_bootstrap(event_loop),
            OneshotTask::SendPacket(recipient, payload) =>
                self.network.send_packet(recipient, payload, event_loop),
            OneshotTask::Shutdown => self.network.shutdown(event_loop)
        }
    }

    fn timeout(&mut self, event_loop: &mut mio::EventLoop<Handler>, timeout: ScheduledTask) {
        match timeout {
            ScheduledTask::ContinueBootstrap => self.network.continue_bootstrap(event_loop),
            ScheduledTask::ContinueHealthCheck => self.network.continue_health_check(event_loop),
            ScheduledTask::ContinueRefresh => self.network.continue_refresh(event_loop)
        }
    }
}

// TODO: I'd rather have some kind of select on two channels instead of this nested non-blocking
// recv. It only works because of the 10ms timeout on the socket.
fn create_incoming_udp_channel(host: SocketAddr, sender: TaskSender) -> mpsc::Sender<mpsc::Sender<()>> {
    use std::time::Duration;

    let (shutdown_sender, shutdown_receiver) = mpsc::channel::<mpsc::Sender<()>>();

    thread::spawn(move || {
        let host = ("0.0.0.0", host.port());
        let socket = UdpSocket::bind(host).unwrap();
        socket.set_read_timeout(Some(Duration::from_millis(10))).unwrap();

        loop {
            match shutdown_receiver.try_recv() {
                Err(mpsc::TryRecvError::Empty) => {
                    let mut buf = [0; 4096];
                    match socket.recv(&mut buf) {
                        Ok(size) => {
                            sender.send(OneshotTask::Incoming(buf[..size].iter().cloned().collect()))
                                .unwrap_or_else(|err| info!("Couldn't handling incoming: {:?}", err));
                        }
                        Err(_) => {
                            //warn!("Error receiving from server: {}", e)
                        }
                    }
                }

                Err(mpsc::TryRecvError::Disconnected) => {
                    // The other end of the shutdown_receiver was dropped
                    break
                }

                Ok(sentinel) => {
                    sentinel.send(()).unwrap();
                    break;
                }
            }
        }

    });
    shutdown_sender
}
<commit_msg>Always insert origin nodes when receiving messages<commit_after>use address::Address;
use messages::outgoing;
use mio;
use node;
use routing_table::{InsertOutcome, InsertionResult, RoutingTable};
use std::collections::HashMap;
use std::io::Cursor;
use std::net::{UdpSocket, ToSocketAddrs, SocketAddr};
use std::sync::mpsc;
use std::thread;
use stun;
use transaction::{TransactionId, TransactionIdGenerator};

#[derive(Debug)]
pub enum Event {
    ReceivedPacket(Address, Vec<u8>),
    Shutdown
}

pub enum OneshotTask {
    Incoming(Vec<u8>),
    StartBootstrap,
    SendPacket(Address, Vec<u8>),
    Shutdown
}

pub enum ScheduledTask {
    ContinueBootstrap,
    ContinueHealthCheck,
    ContinueRefresh
}

enum TableAction {
    Bootstrap(mio::Timeout),
    HealthCheck(mio::Timeout),
    RefreshBucket(mio::Timeout)
}

pub type TaskSender = mio::Sender<OneshotTask>;

enum Status {
    Bootstrapping,
    Idle
}

pub struct Network {
    host: SocketAddr,
    routing_table: RoutingTable,
    self_node: Box<node::Node + 'static>,
    transaction_ids: TransactionIdGenerator,
    status: Status,
    pending_actions: HashMap<TransactionId, TableAction>,
    event_listeners: Vec<mpsc::Sender<Event>>,
    shutdown_signals: Vec<mpsc::Sender<mpsc::Sender<()>>>
}

impl Network {
    pub fn new<T: ToSocketAddrs>(self_address: Address, host: T, routers: Vec<Box<node::Node>>) -> Network {
        let host = host.to_socket_addrs().unwrap().next().unwrap();
        let mapped_host = stun::get_mapped_address(host).unwrap();
        let self_node = node::UdpNode::new(self_address, mapped_host);
        let routing_table = RoutingTable::new(8, self_address, routers);

        Network {
            host: host,
            routing_table: routing_table,
            self_node: Box::new(self_node),
            transaction_ids: TransactionIdGenerator::new(),
            status: Status::Idle,
            pending_actions: HashMap::new(),
            event_listeners: vec![],
            shutdown_signals: vec![]
        }
    }

    pub fn run(mut self) -> TaskSender {
        let mut event_loop_config = mio::EventLoopConfig::new();
        event_loop_config.notify_capacity(16384);
        let mut event_loop = mio::EventLoop::configured(event_loop_config).unwrap();

        let shutdown_signal = create_incoming_udp_channel(self.host, event_loop.channel());
        self.shutdown_signals.push(shutdown_signal);

        event_loop.channel().send(OneshotTask::StartBootstrap).unwrap();
        info!("Running server at {:?}", self.self_node);
        let mut handler = Handler::new(self);
        let task_sender = event_loop.channel();
        thread::spawn(move|| event_loop.run(&mut handler).unwrap());
        task_sender
    }

    pub fn register_event_listener(&mut self, event_listener: mpsc::Sender<Event>) {
        self.event_listeners.push(event_listener);
    }

    fn handle_incoming(&mut self, data: Vec<u8>, event_loop: &mut mio::EventLoop<Handler>) {
        use messages::incoming::{Message, Query, Response, self};
        let mut data = Cursor::new(data);
        let message = incoming::parse_from_reader(&mut data).unwrap();

        match message {
            Message::Query(transaction_id, origin, query) => {
                let origin_address = origin.address();

                match query {
                    Query::FindNode(target) => {
                        let response: Vec<u8> = outgoing::create_find_node_response(
                                transaction_id,
                                &self.self_node,
                                self.routing_table.nearest_live_nodes_to(&target, false));
                        origin.send(response);
                    },
                    Query::Packet(payload) => {
                        for listener in &self.event_listeners {
                            listener.send(Event::ReceivedPacket(origin_address, payload.clone())).unwrap();
                        }
                        let response = outgoing::create_packet_response(
                            transaction_id, &self.self_node);
                        origin.send(response);
                    },
                    Query::Ping => {
                        let response = outgoing::create_ping_response(
                            transaction_id,
                            &self.self_node);
                        origin.send(response);
                    }
                }

                // Always insert the origin node
                self.insert_node(origin).unwrap();
                if let Some(origin) = self.routing_table.find_node(&origin_address) {
                    origin.received_query(transaction_id);
                }
            }

            Message::Response(transaction_id, origin, response) => {
                let origin_address = origin.address();
                // Always insert the origin node
                let mut encounted_new_node = self.insert_node(origin) == Ok(InsertOutcome::Inserted);

                match response {
                    Response::FindNode(mut nodes) => {
                        for node in nodes.drain(..) {
                            match self.insert_node(node) {
                                Ok(InsertOutcome::Inserted) => {
                                    encounted_new_node = true;
                                }
                                Err(error) => { panic!(error) }
                                _ => { }
                            }
                        }

                        {
                            if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_response(transaction_id);
                            }
                        }

                        match self.pending_actions.remove(&transaction_id) {
                            Some(TableAction::Bootstrap(timeout)) => {
                                event_loop.clear_timeout(timeout);
                                if encounted_new_node {
                                    self.continue_bootstrap(event_loop);
                                } else {
                                    self.continue_health_check(event_loop);
                                    self.continue_refresh(event_loop)
                                }
                            }
                            Some(TableAction::HealthCheck(_)) => {
                                // Does not clear timeout. Lets Healthcheck proceed at regular
                                // interval.
                            },
                            Some(TableAction::RefreshBucket(_)) => {
                            }
                            None => { }
                        }

                        match self.status {
                            Status::Bootstrapping => {
                                self.status = Status::Idle;
                            }
                            _ => { }
                        }
                    }
                    Response::Packet => {
                        if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_response(transaction_id);
                        }
                    }
                    Response::Ping => {
                        if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_response(transaction_id);
                        }
                    }
                }
            }
        }
    }

    fn start_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        self.status = Status::Bootstrapping;
        self.continue_bootstrap(event_loop);
    }

    fn continue_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let address = &self.self_node.address();
        let transaction_id = self.find_node(address);
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueBootstrap, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::Bootstrap(timeout));
    }

    fn continue_health_check(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let transaction_id = self.health_check();
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueHealthCheck, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::HealthCheck(timeout));
    }

    fn continue_refresh(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueRefresh, 1000).unwrap();

        match self.refresh_bucket() {
            Some(transaction_id) => {
                self.pending_actions.insert(transaction_id, TableAction::RefreshBucket(timeout));
            }
            None => { }
        }
    }

    fn find_node(&mut self, address: &Address) -> TransactionId {
        let transaction_id = self.transaction_ids.generate();
        let query = outgoing::create_find_node_query(
            transaction_id,
            &self.self_node,
            address);
        for node in self.routing_table.nearest() {
            node.sent_query(transaction_id);
            node.send(query.clone());
        }
        transaction_id
    }

    fn health_check(&mut self) -> TransactionId {
        let transaction_id = self.transaction_ids.generate();
        for node in self.routing_table.questionable_nodes().iter_mut() {
            let query = outgoing::create_ping_query(
                transaction_id, &self.self_node);
            node.sent_query(transaction_id);
            node.send(query);
        }
        transaction_id
    }

    fn refresh_bucket(&mut self) -> Option<TransactionId> {
        if let Some(address) = self.address_to_find_for_refresh() {
            Some(self.find_node(&address))
        } else {
            None
        }
    }

    fn address_to_find_for_refresh(&self) -> Option<Address> {
        match self.routing_table.bucket_needing_refresh() {
            Some(bucket) => {
                Some(bucket.random_address_in_space())
            }
            None => None
        }
    }

    fn insert_node(&mut self, node: Box<node::Node>) -> InsertionResult {
        self.routing_table.insert(node, &self.self_node, &mut self.transaction_ids)
    }

    fn send_packet(&mut self, recipient: Address, payload: Vec<u8>, _event_loop: &mut mio::EventLoop<Handler>) {
        for node in self.routing_table.nearest_live_nodes_to(&recipient, false) {
            let transaction_id = self.transaction_ids.generate();
            let query = outgoing::create_packet_query(
                transaction_id, &self.self_node, payload.clone());
            node.sent_query(transaction_id);
            node.send(query.clone());
        }
    }

    fn shutdown(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let sentinels: Vec<mpsc::Receiver<()>> = self.shutdown_signals.iter().map(|signal| {
            let (sentinal_sender, sentinal_receiver) = mpsc::channel();
            signal.send(sentinal_sender).expect("couldn't send shutdown signal");
            sentinal_receiver
        }).collect();
        for sentinel in sentinels {
            sentinel.recv().unwrap();
        }
        event_loop.shutdown();
        for listener in &self.event_listeners {
            listener.send(Event::Shutdown).unwrap();
        }
    }
}

struct Handler {
    network: Network
}

impl Handler {
    fn new(network: Network) -> Handler {
        Handler {
            network: network
        }
    }
}

impl mio::Handler for Handler {
    type Timeout = ScheduledTask;
    type Message = OneshotTask;

    fn notify(&mut self, event_loop: &mut mio::EventLoop<Handler>, task: OneshotTask) {
        match task {
            OneshotTask::Incoming(data) => self.network.handle_incoming(data, event_loop),
            OneshotTask::StartBootstrap => self.network.start_bootstrap(event_loop),
            OneshotTask::SendPacket(recipient, payload) =>
                self.network.send_packet(recipient, payload, event_loop),
            OneshotTask::Shutdown => self.network.shutdown(event_loop)
        }
    }

    fn timeout(&mut self, event_loop: &mut mio::EventLoop<Handler>, timeout: ScheduledTask) {
        match timeout {
            ScheduledTask::ContinueBootstrap => self.network.continue_bootstrap(event_loop),
            ScheduledTask::ContinueHealthCheck => self.network.continue_health_check(event_loop),
            ScheduledTask::ContinueRefresh => self.network.continue_refresh(event_loop)
        }
    }
}

// TODO: I'd rather have some kind of select on two channels instead of this nested non-blocking
// recv. It only works because of the 10ms timeout on the socket.
fn create_incoming_udp_channel(host: SocketAddr, sender: TaskSender) -> mpsc::Sender<mpsc::Sender<()>> {
    use std::time::Duration;

    let (shutdown_sender, shutdown_receiver) = mpsc::channel::<mpsc::Sender<()>>();

    thread::spawn(move || {
        let host = ("0.0.0.0", host.port());
        let socket = UdpSocket::bind(host).unwrap();
        socket.set_read_timeout(Some(Duration::from_millis(10))).unwrap();

        loop {
            match shutdown_receiver.try_recv() {
                Err(mpsc::TryRecvError::Empty) => {
                    let mut buf = [0; 4096];
                    match socket.recv(&mut buf) {
                        Ok(size) => {
                            sender.send(OneshotTask::Incoming(buf[..size].iter().cloned().collect()))
                                .unwrap_or_else(|err| info!("Couldn't handling incoming: {:?}", err));
                        }
                        Err(_) => {
                            //warn!("Error receiving from server: {}", e)
                        }
                    }
                }

                Err(mpsc::TryRecvError::Disconnected) => {
                    // The other end of the shutdown_receiver was dropped
                    break
                }

                Ok(sentinel) => {
                    sentinel.send(()).unwrap();
                    break;
                }
            }
        }

    });
    shutdown_sender
}
<|endoftext|>


<commit_before>use std::cell::UnsafeCell;

use {
    log::warn,
    rand::{
        rngs::{adapter::ReseedingRng, OsRng},
        CryptoRng, Rng, RngCore, SeedableRng,
    },
    rand_chacha::ChaCha20Core as Core,
    rand_distr::{Distribution, Gamma},
    std::sync::atomic::{AtomicUsize, Ordering::Relaxed},
};

/// This function is useful for inducing random jitter into our atomic
/// operations, shaking out more possible interleavings quickly. It gets
/// fully eliminated by the compiler in non-test code.
pub fn debug_delay() {
    use std::thread;
    use std::time::Duration;

    static GLOBAL_DELAYS: AtomicUsize = AtomicUsize::new(0);

    thread_local!(
        static LOCAL_DELAYS: std::cell::RefCell<usize> = std::cell::RefCell::new(0)
    );

    let global_delays = GLOBAL_DELAYS.fetch_add(1, Relaxed);
    let local_delays = LOCAL_DELAYS.with(|ld| {
        let mut ld = ld.borrow_mut();
        let old = *ld;
        *ld = std::cmp::max(global_delays + 1, *ld + 1);
        old
    });

    if global_delays == local_delays {
        // no other threads seem to be
        // calling this, so we may as
        // well skip it
        return;
    }

    let mut rng = if let Some(rng) = try_thread_rng() {
        rng
    } else {
        warn!("already destroyed TLS when this debug delay was called");
        return;
    };

    let intensity: f64 = std::env::var("SLED_LOCK_FREE_DELAY_INTENSITY")
        .unwrap_or_else(|_| "100.0".into())
        .parse()
        .expect(
            "SLED_LOCK_FREE_DELAY_INTENSITY must be set to a \
             float (ideally between 1-1,000,000)",
        );

    if rng.gen_bool(1. / 1000.) {
        let gamma = Gamma::new(0.3, 1_000.0 * intensity).unwrap();
        let duration = gamma.sample(&mut try_thread_rng().unwrap());
        thread::sleep(Duration::from_micros(duration as u64));
    }

    if rng.gen::<bool>() {
        thread::yield_now();
    }
}

#[derive(Copy, Clone, Debug)]
pub struct ThreadRng {
    // use of raw pointer implies type is neither Send nor Sync
    rng: *mut ReseedingRng<Core, OsRng>,
}

const THREAD_RNG_RESEED_THRESHOLD: u64 = 1024 * 64;

thread_local!(
    static THREAD_RNG_KEY: UnsafeCell<ReseedingRng<Core, OsRng>> = {
        let r = Core::from_rng(OsRng).unwrap_or_else(|err|
                panic!("could not initialize thread_rng: {}", err));
        let rng = ReseedingRng::new(r,
                                    THREAD_RNG_RESEED_THRESHOLD,
                                    OsRng);
        UnsafeCell::new(rng)
    }
);

/// Access a thread-rng that may have been destroyed.
fn try_thread_rng() -> Option<ThreadRng> {
    THREAD_RNG_KEY.try_with(|t| ThreadRng { rng: t.get() }).ok()
}

impl RngCore for ThreadRng {
    #[inline(always)]
    fn next_u32(&mut self) -> u32 {
        unsafe { (*self.rng).next_u32() }
    }

    #[inline(always)]
    fn next_u64(&mut self) -> u64 {
        unsafe { (*self.rng).next_u64() }
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        unsafe { (*self.rng).fill_bytes(dest) }
    }

    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
        self.fill_bytes(dest);
        Ok(())
    }
}

impl CryptoRng for ThreadRng {}
<commit_msg>Use cached and lazily-accessed lock-free delay intensity<commit_after>use std::cell::UnsafeCell;

use {
    log::warn,
    rand::{
        rngs::{adapter::ReseedingRng, OsRng},
        CryptoRng, Rng, RngCore, SeedableRng,
    },
    rand_chacha::ChaCha20Core as Core,
    rand_distr::{Distribution, Gamma},
    std::sync::atomic::{AtomicUsize, Ordering::Relaxed},
};

use crate::Lazy;

/// This function is useful for inducing random jitter into our atomic
/// operations, shaking out more possible interleavings quickly. It gets
/// fully eliminated by the compiler in non-test code.
pub fn debug_delay() {
    use std::thread;
    use std::time::Duration;

    static GLOBAL_DELAYS: AtomicUsize = AtomicUsize::new(0);

    thread_local!(
        static LOCAL_DELAYS: std::cell::RefCell<usize> = std::cell::RefCell::new(0)
    );

    let global_delays = GLOBAL_DELAYS.fetch_add(1, Relaxed);
    let local_delays = LOCAL_DELAYS.with(|ld| {
        let mut ld = ld.borrow_mut();
        let old = *ld;
        *ld = std::cmp::max(global_delays + 1, *ld + 1);
        old
    });

    if global_delays == local_delays {
        // no other threads seem to be
        // calling this, so we may as
        // well skip it
        return;
    }

    let mut rng = if let Some(rng) = try_thread_rng() {
        rng
    } else {
        warn!("already destroyed TLS when this debug delay was called");
        return;
    };

    static INTENSITY: Lazy<f64, fn() -> f64> = Lazy::new(|| {
        std::env::var("SLED_LOCK_FREE_DELAY_INTENSITY")
            .unwrap_or_else(|_| "100.0".into())
            .parse()
            .expect(
                "SLED_LOCK_FREE_DELAY_INTENSITY must be set to a \
                 float (ideally between 1-1,000,000)",
            )
    });

    if rng.gen_bool(1. / 1000.) {
        let gamma = Gamma::new(0.3, 1_000.0 * *INTENSITY).unwrap();
        let duration = gamma.sample(&mut try_thread_rng().unwrap());
        thread::sleep(Duration::from_micros(duration as u64));
    }

    if rng.gen::<bool>() {
        thread::yield_now();
    }
}

#[derive(Copy, Clone, Debug)]
pub struct ThreadRng {
    // use of raw pointer implies type is neither Send nor Sync
    rng: *mut ReseedingRng<Core, OsRng>,
}

const THREAD_RNG_RESEED_THRESHOLD: u64 = 1024 * 64;

thread_local!(
    static THREAD_RNG_KEY: UnsafeCell<ReseedingRng<Core, OsRng>> = {
        let r = Core::from_rng(OsRng).unwrap_or_else(|err|
                panic!("could not initialize thread_rng: {}", err));
        let rng = ReseedingRng::new(r,
                                    THREAD_RNG_RESEED_THRESHOLD,
                                    OsRng);
        UnsafeCell::new(rng)
    }
);

/// Access a thread-rng that may have been destroyed.
fn try_thread_rng() -> Option<ThreadRng> {
    THREAD_RNG_KEY.try_with(|t| ThreadRng { rng: t.get() }).ok()
}

impl RngCore for ThreadRng {
    #[inline(always)]
    fn next_u32(&mut self) -> u32 {
        unsafe { (*self.rng).next_u32() }
    }

    #[inline(always)]
    fn next_u64(&mut self) -> u64 {
        unsafe { (*self.rng).next_u64() }
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        unsafe { (*self.rng).fill_bytes(dest) }
    }

    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
        self.fill_bytes(dest);
        Ok(())
    }
}

impl CryptoRng for ThreadRng {}
<|endoftext|>


<commit_before>//! FIXME: write short doc here

use ra_parser::{Token, TokenSource};
use ra_syntax::{lex_single_valid_syntax_kind, SmolStr, SyntaxKind, SyntaxKind::*, T};
use std::cell::{Cell, Ref, RefCell};
use tt::buffer::{Cursor, TokenBuffer};

#[derive(Debug, Clone, Eq, PartialEq)]
struct TtToken {
    pub kind: SyntaxKind,
    pub is_joint_to_next: bool,
    pub text: SmolStr,
}

pub(crate) struct SubtreeTokenSource<'a> {
    cached_cursor: Cell<Cursor<'a>>,
    cached: RefCell<Vec<Option<TtToken>>>,
    curr: (Token, usize),
}

impl<'a> SubtreeTokenSource<'a> {
    // Helper function used in test
    #[cfg(test)]
    pub fn text(&self) -> SmolStr {
        match *self.get(self.curr.1) {
            Some(ref tt) => tt.text.clone(),
            _ => SmolStr::new(""),
        }
    }
}

impl<'a> SubtreeTokenSource<'a> {
    pub fn new(buffer: &'a TokenBuffer) -> SubtreeTokenSource<'a> {
        let cursor = buffer.begin();

        let mut res = SubtreeTokenSource {
            curr: (Token { kind: EOF, is_jointed_to_next: false }, 0),
            cached_cursor: Cell::new(cursor),
            cached: RefCell::new(Vec::with_capacity(10)),
        };
        res.curr = (res.mk_token(0), 0);
        res
    }

    fn mk_token(&self, pos: usize) -> Token {
        match *self.get(pos) {
            Some(ref tt) => Token { kind: tt.kind, is_jointed_to_next: tt.is_joint_to_next },
            None => Token { kind: EOF, is_jointed_to_next: false },
        }
    }

    fn get(&self, pos: usize) -> Ref<Option<TtToken>> {
        if pos < self.cached.borrow().len() {
            return Ref::map(self.cached.borrow(), |c| &c[pos]);
        }

        {
            let mut cached = self.cached.borrow_mut();
            while pos >= cached.len() {
                let cursor = self.cached_cursor.get();
                if cursor.eof() {
                    cached.push(None);
                    continue;
                }

                match cursor.token_tree() {
                    Some(tt::TokenTree::Leaf(leaf)) => {
                        cached.push(Some(convert_leaf(&leaf)));
                        self.cached_cursor.set(cursor.bump());
                    }
                    Some(tt::TokenTree::Subtree(subtree)) => {
                        self.cached_cursor.set(cursor.subtree().unwrap());
                        cached.push(Some(convert_delim(subtree.delimiter_kind(), false)));
                    }
                    None => {
                        if let Some(subtree) = cursor.end() {
                            cached.push(Some(convert_delim(subtree.delimiter_kind(), true)));
                            self.cached_cursor.set(cursor.bump());
                        }
                    }
                }
            }
        }

        Ref::map(self.cached.borrow(), |c| &c[pos])
    }
}

impl<'a> TokenSource for SubtreeTokenSource<'a> {
    fn current(&self) -> Token {
        self.curr.0
    }

    /// Lookahead n token
    fn lookahead_nth(&self, n: usize) -> Token {
        self.mk_token(self.curr.1 + n)
    }

    /// bump cursor to next token
    fn bump(&mut self) {
        if self.current().kind == EOF {
            return;
        }

        self.curr = (self.mk_token(self.curr.1 + 1), self.curr.1 + 1);
    }

    /// Is the current token a specified keyword?
    fn is_keyword(&self, kw: &str) -> bool {
        match *self.get(self.curr.1) {
            Some(ref t) => t.text == *kw,
            _ => false,
        }
    }
}

fn convert_delim(d: Option<tt::DelimiterKind>, closing: bool) -> TtToken {
    let (kinds, texts) = match d {
        Some(tt::DelimiterKind::Parenthesis) => ([T!['('], T![')']], "()"),
        Some(tt::DelimiterKind::Brace) => ([T!['{'], T!['}']], "{}"),
        Some(tt::DelimiterKind::Bracket) => ([T!['['], T![']']], "[]"),
        None => ([L_DOLLAR, R_DOLLAR], ""),
    };

    let idx = closing as usize;
    let kind = kinds[idx];
    let text = if !texts.is_empty() { &texts[idx..texts.len() - (1 - idx)] } else { "" };
    TtToken { kind, is_joint_to_next: false, text: SmolStr::new(text) }
}

fn convert_literal(l: &tt::Literal) -> TtToken {
    let kind = lex_single_valid_syntax_kind(&l.text)
        .filter(|kind| kind.is_literal())
        .unwrap_or_else(|| match l.text.as_ref() {
            "true" => T![true],
            "false" => T![false],
            _ => panic!("Fail to convert given literal {:#?}", &l),
        });

    TtToken { kind, is_joint_to_next: false, text: l.text.clone() }
}

fn convert_ident(ident: &tt::Ident) -> TtToken {
    let kind = if ident.text.starts_with('\'') {
        LIFETIME
    } else {
        SyntaxKind::from_keyword(ident.text.as_str()).unwrap_or(IDENT)
    };

    TtToken { kind, is_joint_to_next: false, text: ident.text.clone() }
}

fn convert_punct(p: tt::Punct) -> TtToken {
    let kind = SyntaxKind::from_char(p.char).unwrap();
    let text = {
        let mut buf = [0u8; 4];
        let s: &str = p.char.encode_utf8(&mut buf);
        SmolStr::new(s)
    };
    TtToken { kind, is_joint_to_next: p.spacing == tt::Spacing::Joint, text }
}

fn convert_leaf(leaf: &tt::Leaf) -> TtToken {
    match leaf {
        tt::Leaf::Literal(l) => convert_literal(l),
        tt::Leaf::Ident(ident) => convert_ident(ident),
        tt::Leaf::Punct(punct) => convert_punct(*punct),
    }
}
<commit_msg>ra_mbe: convert_literal now works with malformed tokens<commit_after>//! FIXME: write short doc here

use ra_parser::{Token, TokenSource};
use ra_syntax::{lex_single_syntax_kind, SmolStr, SyntaxKind, SyntaxKind::*, T};
use std::cell::{Cell, Ref, RefCell};
use tt::buffer::{Cursor, TokenBuffer};

#[derive(Debug, Clone, Eq, PartialEq)]
struct TtToken {
    pub kind: SyntaxKind,
    pub is_joint_to_next: bool,
    pub text: SmolStr,
}

pub(crate) struct SubtreeTokenSource<'a> {
    cached_cursor: Cell<Cursor<'a>>,
    cached: RefCell<Vec<Option<TtToken>>>,
    curr: (Token, usize),
}

impl<'a> SubtreeTokenSource<'a> {
    // Helper function used in test
    #[cfg(test)]
    pub fn text(&self) -> SmolStr {
        match *self.get(self.curr.1) {
            Some(ref tt) => tt.text.clone(),
            _ => SmolStr::new(""),
        }
    }
}

impl<'a> SubtreeTokenSource<'a> {
    pub fn new(buffer: &'a TokenBuffer) -> SubtreeTokenSource<'a> {
        let cursor = buffer.begin();

        let mut res = SubtreeTokenSource {
            curr: (Token { kind: EOF, is_jointed_to_next: false }, 0),
            cached_cursor: Cell::new(cursor),
            cached: RefCell::new(Vec::with_capacity(10)),
        };
        res.curr = (res.mk_token(0), 0);
        res
    }

    fn mk_token(&self, pos: usize) -> Token {
        match *self.get(pos) {
            Some(ref tt) => Token { kind: tt.kind, is_jointed_to_next: tt.is_joint_to_next },
            None => Token { kind: EOF, is_jointed_to_next: false },
        }
    }

    fn get(&self, pos: usize) -> Ref<Option<TtToken>> {
        if pos < self.cached.borrow().len() {
            return Ref::map(self.cached.borrow(), |c| &c[pos]);
        }

        {
            let mut cached = self.cached.borrow_mut();
            while pos >= cached.len() {
                let cursor = self.cached_cursor.get();
                if cursor.eof() {
                    cached.push(None);
                    continue;
                }

                match cursor.token_tree() {
                    Some(tt::TokenTree::Leaf(leaf)) => {
                        cached.push(Some(convert_leaf(&leaf)));
                        self.cached_cursor.set(cursor.bump());
                    }
                    Some(tt::TokenTree::Subtree(subtree)) => {
                        self.cached_cursor.set(cursor.subtree().unwrap());
                        cached.push(Some(convert_delim(subtree.delimiter_kind(), false)));
                    }
                    None => {
                        if let Some(subtree) = cursor.end() {
                            cached.push(Some(convert_delim(subtree.delimiter_kind(), true)));
                            self.cached_cursor.set(cursor.bump());
                        }
                    }
                }
            }
        }

        Ref::map(self.cached.borrow(), |c| &c[pos])
    }
}

impl<'a> TokenSource for SubtreeTokenSource<'a> {
    fn current(&self) -> Token {
        self.curr.0
    }

    /// Lookahead n token
    fn lookahead_nth(&self, n: usize) -> Token {
        self.mk_token(self.curr.1 + n)
    }

    /// bump cursor to next token
    fn bump(&mut self) {
        if self.current().kind == EOF {
            return;
        }

        self.curr = (self.mk_token(self.curr.1 + 1), self.curr.1 + 1);
    }

    /// Is the current token a specified keyword?
    fn is_keyword(&self, kw: &str) -> bool {
        match *self.get(self.curr.1) {
            Some(ref t) => t.text == *kw,
            _ => false,
        }
    }
}

fn convert_delim(d: Option<tt::DelimiterKind>, closing: bool) -> TtToken {
    let (kinds, texts) = match d {
        Some(tt::DelimiterKind::Parenthesis) => ([T!['('], T![')']], "()"),
        Some(tt::DelimiterKind::Brace) => ([T!['{'], T!['}']], "{}"),
        Some(tt::DelimiterKind::Bracket) => ([T!['['], T![']']], "[]"),
        None => ([L_DOLLAR, R_DOLLAR], ""),
    };

    let idx = closing as usize;
    let kind = kinds[idx];
    let text = if !texts.is_empty() { &texts[idx..texts.len() - (1 - idx)] } else { "" };
    TtToken { kind, is_joint_to_next: false, text: SmolStr::new(text) }
}

fn convert_literal(l: &tt::Literal) -> TtToken {
    let kind = lex_single_syntax_kind(&l.text)
        .map(|(kind, _error)| kind)
        .filter(|kind| kind.is_literal())
        .unwrap_or_else(|| match l.text.as_ref() {
            "true" => T![true],
            "false" => T![false],
            _ => panic!("Fail to convert given literal {:#?}", &l),
        });

    TtToken { kind, is_joint_to_next: false, text: l.text.clone() }
}

fn convert_ident(ident: &tt::Ident) -> TtToken {
    let kind = if ident.text.starts_with('\'') {
        LIFETIME
    } else {
        SyntaxKind::from_keyword(ident.text.as_str()).unwrap_or(IDENT)
    };

    TtToken { kind, is_joint_to_next: false, text: ident.text.clone() }
}

fn convert_punct(p: tt::Punct) -> TtToken {
    let kind = SyntaxKind::from_char(p.char).unwrap();
    let text = {
        let mut buf = [0u8; 4];
        let s: &str = p.char.encode_utf8(&mut buf);
        SmolStr::new(s)
    };
    TtToken { kind, is_joint_to_next: p.spacing == tt::Spacing::Joint, text }
}

fn convert_leaf(leaf: &tt::Leaf) -> TtToken {
    match leaf {
        tt::Leaf::Literal(l) => convert_literal(l),
        tt::Leaf::Ident(ident) => convert_ident(ident),
        tt::Leaf::Punct(punct) => convert_punct(*punct),
    }
}
<|endoftext|>


<commit_before>mod add;
mod build;
mod cmin;
mod fmt;
mod init;
mod list;
mod run;
mod tmin;

pub use self::{
    add::Add, build::Build, cmin::Cmin, fmt::Fmt, init::Init, list::List, run::Run, tmin::Tmin,
};

use std::fmt as stdfmt;
use std::str::FromStr;
use structopt::StructOpt;

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Sanitizer {
    Address,
    Leak,
    Memory,
    Thread,
    None,
}

impl stdfmt::Display for Sanitizer {
    fn fmt(&self, f: &mut stdfmt::Formatter) -> stdfmt::Result {
        write!(
            f,
            "{}",
            match self {
                Sanitizer::Address => "address",
                Sanitizer::Leak => "leak",
                Sanitizer::Memory => "memory",
                Sanitizer::Thread => "thread",
                Sanitizer::None => "",
            }
        )
    }
}

impl FromStr for Sanitizer {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "address" => Ok(Sanitizer::Address),
            "leak" => Ok(Sanitizer::Leak),
            "memory" => Ok(Sanitizer::Memory),
            "thread" => Ok(Sanitizer::Thread),
            "none" => Ok(Sanitizer::None),
            _ => Err(format!("unknown sanitizer: {}", s)),
        }
    }
}

#[derive(Clone, Debug, StructOpt)]
pub struct BuildOptions {
    #[structopt(short = "D", long = "dev", conflicts_with = "release")]
    /// Build artifacts in development mode, without optimizations
    pub dev: bool,

    #[structopt(short = "O", long = "release", conflicts_with = "dev")]
    /// Build artifacts in release mode, with optimizations
    pub release: bool,

    #[structopt(short = "a", long = "debug-assertions")]
    /// Build artifacts with debug assertions and overflow checks enabled (default if not -O)
    pub debug_assertions: bool,

    /// Build target with verbose output from `cargo build`
    #[structopt(short = "v", long = "verbose")]
    pub verbose: bool,

    #[structopt(long = "no-default-features")]
    /// Build artifacts with default Cargo features disabled
    pub no_default_features: bool,

    #[structopt(
        long = "all-features",
        conflicts_with = "no-default-features",
        conflicts_with = "features"
    )]
    /// Build artifacts with all Cargo features enabled
    pub all_features: bool,

    #[structopt(long = "features")]
    /// Build artifacts with given Cargo feature enabled
    pub features: Option<String>,

    #[structopt(
        short = "s",
        long = "sanitizer",
        possible_values(&["address", "leak", "memory", "thread", "none"]),
        default_value = "address",
    )]
    /// Use a specific sanitizer
    pub sanitizer: Sanitizer,

    #[structopt(
        name = "triple",
        long = "target",
        default_value(crate::utils::default_target())
    )]
    /// Target triple of the fuzz target
    pub triple: String,

    #[structopt(short = "Z", value_name = "FLAG")]
    /// Unstable (nightly-only) flags to Cargo
    pub unstable_flags: Vec<String>,
}

impl stdfmt::Display for BuildOptions {
    fn fmt(&self, f: &mut stdfmt::Formatter) -> stdfmt::Result {
        if self.dev {
            write!(f, " -D")?;
        }

        if self.release {
            write!(f, " -O")?;
        }

        if self.debug_assertions {
            write!(f, " -a")?;
        }

        if self.verbose {
            write!(f, " -v")?;
        }

        if self.no_default_features {
            write!(f, " --no-default-features")?;
        }

        if self.all_features {
            write!(f, " --all-features")?;
        }

        if let Some(feature) = &self.features {
            write!(f, " --feature={}", feature)?;
        }

        if self.sanitizer != Sanitizer::None {
            write!(f, " --sanitizer={}", self.sanitizer)?;
        }

        if self.triple != crate::utils::default_target() {
            write!(f, " --target={}", self.triple)?;
        }

        if !self.unstable_flags.is_empty() {
            write!(f, " -Z '{}'", self.unstable_flags.join(" "))?;
        }

        Ok(())
    }
}
<commit_msg>Add unit test for BuildOptions::fmt<commit_after>mod add;
mod build;
mod cmin;
mod fmt;
mod init;
mod list;
mod run;
mod tmin;

pub use self::{
    add::Add, build::Build, cmin::Cmin, fmt::Fmt, init::Init, list::List, run::Run, tmin::Tmin,
};

use std::fmt as stdfmt;
use std::str::FromStr;
use structopt::StructOpt;

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Sanitizer {
    Address,
    Leak,
    Memory,
    Thread,
    None,
}

impl stdfmt::Display for Sanitizer {
    fn fmt(&self, f: &mut stdfmt::Formatter) -> stdfmt::Result {
        write!(
            f,
            "{}",
            match self {
                Sanitizer::Address => "address",
                Sanitizer::Leak => "leak",
                Sanitizer::Memory => "memory",
                Sanitizer::Thread => "thread",
                Sanitizer::None => "",
            }
        )
    }
}

impl FromStr for Sanitizer {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "address" => Ok(Sanitizer::Address),
            "leak" => Ok(Sanitizer::Leak),
            "memory" => Ok(Sanitizer::Memory),
            "thread" => Ok(Sanitizer::Thread),
            "none" => Ok(Sanitizer::None),
            _ => Err(format!("unknown sanitizer: {}", s)),
        }
    }
}

#[derive(Clone, Debug, StructOpt, PartialEq)]
pub struct BuildOptions {
    #[structopt(short = "D", long = "dev", conflicts_with = "release")]
    /// Build artifacts in development mode, without optimizations
    pub dev: bool,

    #[structopt(short = "O", long = "release", conflicts_with = "dev")]
    /// Build artifacts in release mode, with optimizations
    pub release: bool,

    #[structopt(short = "a", long = "debug-assertions")]
    /// Build artifacts with debug assertions and overflow checks enabled (default if not -O)
    pub debug_assertions: bool,

    /// Build target with verbose output from `cargo build`
    #[structopt(short = "v", long = "verbose")]
    pub verbose: bool,

    #[structopt(long = "no-default-features")]
    /// Build artifacts with default Cargo features disabled
    pub no_default_features: bool,

    #[structopt(
        long = "all-features",
        conflicts_with = "no-default-features",
        conflicts_with = "features"
    )]
    /// Build artifacts with all Cargo features enabled
    pub all_features: bool,

    #[structopt(long = "features")]
    /// Build artifacts with given Cargo feature enabled
    pub features: Option<String>,

    #[structopt(
        short = "s",
        long = "sanitizer",
        possible_values(&["address", "leak", "memory", "thread", "none"]),
        default_value = "address",
    )]
    /// Use a specific sanitizer
    pub sanitizer: Sanitizer,

    #[structopt(
        name = "triple",
        long = "target",
        default_value(crate::utils::default_target())
    )]
    /// Target triple of the fuzz target
    pub triple: String,

    #[structopt(short = "Z", value_name = "FLAG")]
    /// Unstable (nightly-only) flags to Cargo
    pub unstable_flags: Vec<String>,
}

impl stdfmt::Display for BuildOptions {
    fn fmt(&self, f: &mut stdfmt::Formatter) -> stdfmt::Result {
        if self.dev {
            write!(f, " -D")?;
        }

        if self.release {
            write!(f, " -O")?;
        }

        if self.debug_assertions {
            write!(f, " -a")?;
        }

        if self.verbose {
            write!(f, " -v")?;
        }

        if self.no_default_features {
            write!(f, " --no-default-features")?;
        }

        if self.all_features {
            write!(f, " --all-features")?;
        }

        if let Some(feature) = &self.features {
            write!(f, " --features={}", feature)?;
        }

        match self.sanitizer {
            Sanitizer::None => write!(f, " --sanitizer=none")?,
            Sanitizer::Address => {}
            _ => write!(f, " --sanitizer={}", self.sanitizer)?,
        }

        if self.triple != crate::utils::default_target() {
            write!(f, " --target={}", self.triple)?;
        }

        for flag in &self.unstable_flags {
            write!(f, " -Z{}", flag)?;
        }

        Ok(())
    }
}

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

    #[test]
    fn display_build_options() {
        let default_opts = BuildOptions {
            dev: false,
            release: false,
            debug_assertions: false,
            verbose: false,
            no_default_features: false,
            all_features: false,
            features: None,
            sanitizer: Sanitizer::Address,
            triple: String::from(crate::utils::default_target()),
            unstable_flags: Vec::new(),
        };

        let opts = vec![
            default_opts.clone(),
            BuildOptions {
                dev: true,
                ..default_opts.clone()
            },
            BuildOptions {
                release: true,
                ..default_opts.clone()
            },
            BuildOptions {
                debug_assertions: true,
                ..default_opts.clone()
            },
            BuildOptions {
                verbose: true,
                ..default_opts.clone()
            },
            BuildOptions {
                no_default_features: true,
                ..default_opts.clone()
            },
            BuildOptions {
                all_features: true,
                ..default_opts.clone()
            },
            BuildOptions {
                features: Some(String::from("features")),
                ..default_opts.clone()
            },
            BuildOptions {
                sanitizer: Sanitizer::None,
                ..default_opts.clone()
            },
            BuildOptions {
                triple: String::from("custom_triple"),
                ..default_opts.clone()
            },
            BuildOptions {
                unstable_flags: vec![String::from("unstable"), String::from("flags")],
                ..default_opts
            },
        ];

        for case in opts {
            assert_eq!(
                case,
                BuildOptions::from_clap(
                    &BuildOptions::clap().get_matches_from(case.to_string().split(' '))
                )
            );
        }
    }
}
<|endoftext|>