<commit_before>use mio::Timeout as MioTimeout;
use mio::tcp::{TcpListener, TcpStream};
use mio::util::Slab;

use std::borrow::Borrow;
use std::io;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::cmp::min;
use std::default::Default;

use messages::RawMessage;
use super::connection::{Connection, IncomingConnection, OutgoingConnection};
use super::{Timeout, InternalTimeout, EventLoop, EventHandler, Event};

pub use mio::{EventSet, PollOpt, Token};

pub type PeerId = Token;

const SERVER_ID: PeerId = Token(1);
const RECONNECT_GROW_FACTOR: f32 = 2.0;

#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
pub struct NetworkConfiguration {
    // TODO: think more about config parameters
    pub max_incoming_connections: usize,
    pub max_outgoing_connections: usize,
    pub tcp_nodelay: bool,
    pub tcp_keep_alive: Option<u32>,
    pub tcp_reconnect_timeout: u64,
    pub tcp_reconnect_timeout_max: u64,
}

impl Default for NetworkConfiguration {
    fn default() -> NetworkConfiguration {
        NetworkConfiguration {
            max_incoming_connections: 128,
            max_outgoing_connections: 128,
            tcp_keep_alive: None,
            tcp_nodelay: false,
            tcp_reconnect_timeout: 500,
            tcp_reconnect_timeout_max: 600000,
        }
    }
}

// TODO Implement generic ConnectionPool struct to avoid copy paste.
// Write proper code to configure outgoing streams
pub struct Network {
    config: NetworkConfiguration,
    listen_address: SocketAddr,
    listener: Option<TcpListener>,

    incoming: Slab<IncomingConnection>,
    outgoing: Slab<OutgoingConnection>,
    // FIXME addresses only needs for outgoing connections
    addresses: HashMap<SocketAddr, PeerId>,

    reconnects: HashMap<SocketAddr, MioTimeout>,
}

enum PeerKind {
    Server,
    Incoming,
    Outgoing,
}

fn make_io_error<T: Borrow<str>>(s: T) -> io::Error {
    io::Error::new(io::ErrorKind::Other, s.borrow())
}

impl Network {
    pub fn with_config(address: SocketAddr, config: NetworkConfiguration) -> Network {
        Network {
            config: config,
            listen_address: address,
            listener: None,

            incoming: Slab::new_starting_at(Token(2), config.max_incoming_connections),
            outgoing: Slab::new_starting_at(Token(2 + config.max_incoming_connections),
                                            config.max_outgoing_connections),
            addresses: HashMap::new(),

            reconnects: HashMap::new(),
        }
    }

    pub fn address(&self) -> &SocketAddr {
        &self.listen_address
    }

    // TODO use error trait
    pub fn bind<H: EventHandler>(&mut self, event_loop: &mut EventLoop<H>) -> io::Result<()> {
        if self.listener.is_some() {
            return Err(make_io_error("Already binded"));
        }
        let listener = TcpListener::bind(&self.listen_address)?;
        event_loop.register(&listener, SERVER_ID, EventSet::readable(), PollOpt::edge())?;
        self.listener = Some(listener);
        Ok(())
    }

    // TODO: Use ticks for fast reregistering sockets
    // TODO: Implement Connections collection with (re)registering
    pub fn io<H: EventHandler>(&mut self,
                               event_loop: &mut EventLoop<H>,
                               handler: &mut H,
                               id: PeerId,
                               set: EventSet)
                               -> io::Result<()> {

        match self.peer_kind(id) {
            PeerKind::Server => {
                // Accept new connections
                // FIXME: Fail-safe accepting of new connections?
                let pair = match self.listener {
                    Some(ref listener) => listener.accept()?,
                    None => None,
                };
                if let Some((mut stream, address)) = pair {
                    self.configure_stream(&mut stream)?;
                    let peer = IncomingConnection::new(stream, address);
                    self.add_incoming_connection(event_loop, peer)?;

                    trace!("{}: Accepted incoming connection from {} id: {}",
                           self.address(),
                           address,
                           id.0);
                }
                return Ok(());
            }
            PeerKind::Incoming => {
                if !self.incoming.contains(id) {
                    return Ok(());
                }

                if set.is_hup() | set.is_error() {
                    trace!("{}: incoming connection with addr {} closed",
                           self.address(),
                           self.incoming[id].address());

                    self.remove_incoming_connection(event_loop, id);
                    return Ok(());
                }

                if set.is_readable() {
                    loop {
                        match self.incoming[id].try_read() {
                            Ok(Some(buf)) => {
                                let msg = RawMessage::new(buf);
                                handler.handle_event(Event::Incoming(msg));
                            }
                            Ok(None) => return Ok(()),
                            Err(e) => {
                                self.remove_incoming_connection(event_loop, id);
                                return Err(e);
                            }
                        };
                    }
                }
            }
            PeerKind::Outgoing => {
                if !self.outgoing.contains(id) {
                    return Ok(());
                }

                if set.is_hup() | set.is_error() {
                    let address = *self.outgoing[id].address();

                    trace!("{}: outgoing connection with addr {} closed",
                           self.address(),
                           self.outgoing[id].address());

                    self.remove_outgoing_connection(event_loop, id);
                    if !self.reconnects.contains_key(&address) {
                        handler.handle_event(Event::Disconnected(address));
                    }
                    return Ok(());
                }

                if set.is_writable() {
                    let address = *self.outgoing[id].address();
                    let r = {
                        // Write data into socket
                        self.outgoing[id].try_write()?;
                        event_loop.reregister(self.outgoing[id].socket(),
                                        id,
                                        self.outgoing[id].interest(),
                                        PollOpt::edge())?;
                        Ok(())
                    };
                    if let Err(e) = r {
                        self.remove_outgoing_connection(event_loop, id);
                        handler.handle_event(Event::Disconnected(address));
                        return Err(e);
                    }

                    if self.mark_connected(event_loop, id) {
                        trace!("{}: Handle connected with={}", self.address(), address);
                        handler.handle_event(Event::Connected(address));
                    }
                }
            }
        }
        Ok(())
    }

    pub fn tick<H: EventHandler>(&mut self, _: &mut EventLoop<H>) {}

    pub fn send_to<H: EventHandler>(&mut self,
                                    event_loop: &mut EventLoop<H>,
                                    address: &SocketAddr,
                                    message: RawMessage)
                                    -> io::Result<()> {
        match self.get_outgoing_peer(address) {
            Ok(id) => {
                self.outgoing[id]
                    .send(message)
                    .and_then(|_| {
                        event_loop.reregister(self.outgoing[id].socket(),
                                        id,
                                        self.outgoing[id].interest(),
                                        PollOpt::edge())?;
                        self.mark_connected(event_loop, id);
                        Ok(())
                    })
                    .or_else(|e| {
                        self.remove_outgoing_connection(event_loop, id);
                        Err(e)
                    })
            }
            Err(e) => Err(e),
        }
    }

    pub fn connect<H: EventHandler>(&mut self,
                                    event_loop: &mut EventLoop<H>,
                                    address: &SocketAddr)
                                    -> io::Result<()> {
        if !self.is_connected(address) {
            self.add_reconnect_request(event_loop, *address)?;

            let mut stream = TcpStream::connect(address)?;
            self.configure_stream(&mut stream)?;
            let peer = OutgoingConnection::new(stream, *address);
            let id = self.add_outgoing_connection(event_loop, peer)?;

            trace!("{}: Establish connection with {}, id: {}",
                   self.address(),
                   address,
                   id.0);
        }
        Ok(())
    }

    pub fn is_connected(&self, address: &SocketAddr) -> bool {
        self.addresses.contains_key(address)
    }

    pub fn handle_timeout<H: EventHandler>(&mut self,
                                           event_loop: &mut EventLoop<H>,
                                           timeout: InternalTimeout) {
        match timeout {
            InternalTimeout::Reconnect(addr, delay) => {
                if self.reconnects.contains_key(&addr) {
                    trace!("Try to reconnect with delay {}", delay);

                    if let Err(e) = self.connect(event_loop, &addr) {
                        error!("{}: Unable to create connection to addr {}, error: {:?}",
                               self.address(),
                               addr,
                               e);
                    }

                    let delay = min((delay as f32 * RECONNECT_GROW_FACTOR) as u64,
                                    self.config.tcp_reconnect_timeout_max);

                    if let Err(e) = self.add_reconnect_timeout(event_loop, addr, delay) {
                        error!("{}: Unable to add timeout, error: {:?}", self.address(), e);
                    }
                }
            }
        }
    }

    fn peer_kind(&self, id: PeerId) -> PeerKind {
        if id == SERVER_ID {
            PeerKind::Server
        } else if id.0 >= (2 + self.config.max_incoming_connections) {
            PeerKind::Outgoing
        } else {
            PeerKind::Incoming
        }
    }

    fn get_outgoing_peer(&self, addr: &SocketAddr) -> io::Result<PeerId> {
        if let Some(id) = self.addresses.get(addr) {
            return Ok(*id);
        };
        Err(make_io_error(format!("{}: Outgoing peer not found {}", self.address(), addr)))
    }

    fn add_incoming_connection<H: EventHandler>(&mut self,
                                                event_loop: &mut EventLoop<H>,
                                                connection: IncomingConnection)
                                                -> io::Result<PeerId> {
        let address = *connection.address();
        let id = self.incoming
            .insert(connection)
            .map_err(|_| make_io_error("Maximum incoming onnections"))?;
        self.addresses.insert(address, id);

        let r = event_loop.register(self.incoming[id].socket(),
                                    id,
                                    self.incoming[id].interest(),
                                    PollOpt::edge());

        if let Err(e) = r {
            self.remove_incoming_connection(event_loop, id);
            return Err(e);
        }
        Ok(id)
    }

    fn add_outgoing_connection<H: EventHandler>(&mut self,
                                                event_loop: &mut EventLoop<H>,
                                                connection: OutgoingConnection)
                                                -> io::Result<PeerId> {
        let address = *connection.address();
        let id = self.outgoing
            .insert(connection)
            .map_err(|_| make_io_error("Maximum outgoing onnections"))?;
        self.addresses.insert(address, id);

        let r = event_loop.register(self.outgoing[id].socket(),
                                    id,
                                    self.outgoing[id].interest() | EventSet::writable(),
                                    PollOpt::edge());

        if let Err(e) = r {
            self.remove_outgoing_connection(event_loop, id);
            return Err(e);
        }
        Ok(id)
    }

    fn remove_incoming_connection<H: EventHandler>(&mut self,
                                                   event_loop: &mut EventLoop<H>,
                                                   id: PeerId) {
        let addr = *self.incoming[id].address();
        self.addresses.remove(&addr);
        if let Some(connection) = self.incoming.remove(id) {
            if let Err(e) = event_loop.deregister(connection.socket()) {
                error!("{}: Unable to deregister incoming connection, id: {}, error: {:?}",
                       self.address(),
                       id.0,
                       e);
            }
        }
    }

    fn remove_outgoing_connection<H: EventHandler>(&mut self,
                                                   event_loop: &mut EventLoop<H>,
                                                   id: PeerId) {
        let addr = *self.outgoing[id].address();
        self.addresses.remove(&addr);
        if let Some(connection) = self.outgoing.remove(id) {
            if let Err(e) = event_loop.deregister(connection.socket()) {
                error!("{}: Unable to deregister outgoing connection, id: {}, error: {:?}",
                       self.address(),
                       id.0,
                       e);
            }
        }
    }

    fn configure_stream(&self, stream: &mut TcpStream) -> io::Result<()> {
        stream.take_socket_error()?;
        stream.set_keepalive(self.config.tcp_keep_alive)?;
        stream.set_nodelay(self.config.tcp_nodelay)?;
        stream.take_socket_error()
    }

    fn add_reconnect_request<H: EventHandler>(&mut self,
                                              event_loop: &mut EventLoop<H>,
                                              address: SocketAddr)
                                              -> io::Result<()> {
        if !self.reconnects.contains_key(&address) {
            let delay = self.config.tcp_reconnect_timeout;
            return self.add_reconnect_timeout(event_loop, address, delay);
        }
        Ok(())
    }

    fn add_reconnect_timeout<H: EventHandler>(&mut self,
                                              event_loop: &mut EventLoop<H>,
                                              address: SocketAddr,
                                              delay: u64)
                                              -> io::Result<()> {
        trace!("{}: Add reconnect timeout to={}, delay={}",
               self.address(),
               address,
               delay);
        let reconnect = Timeout::Internal(InternalTimeout::Reconnect(address, delay));
        let timeout = event_loop.timeout_ms(reconnect, delay)
            .map_err(|e| make_io_error(format!("A mio error occured {:?}", e)))?;
        self.reconnects.insert(address, timeout);
        Ok(())
    }

    fn mark_connected<H: EventHandler>(&mut self,
                                       event_loop: &mut EventLoop<H>,
                                       id: Token)
                                       -> bool {
        let address = *self.outgoing[id].address();
        self.clear_reconnect_request(event_loop, &address)
    }

    fn clear_reconnect_request<H: EventHandler>(&mut self,
                                                event_loop: &mut EventLoop<H>,
                                                addr: &SocketAddr)
                                                -> bool {
        if let Some(timeout) = self.reconnects.remove(addr) {
            debug!("{}: Clear reconnect timeout to={}", self.address(), addr);
            event_loop.clear_timeout(timeout);
            return true;
        }
        false
    }
}

#[cfg(test)]
mod tests {}
<commit_msg>Remove debug logs<commit_after>use mio::Timeout as MioTimeout;
use mio::tcp::{TcpListener, TcpStream};
use mio::util::Slab;

use std::borrow::Borrow;
use std::io;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::cmp::min;
use std::default::Default;

use messages::RawMessage;
use super::connection::{Connection, IncomingConnection, OutgoingConnection};
use super::{Timeout, InternalTimeout, EventLoop, EventHandler, Event};

pub use mio::{EventSet, PollOpt, Token};

pub type PeerId = Token;

const SERVER_ID: PeerId = Token(1);
const RECONNECT_GROW_FACTOR: f32 = 2.0;

#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
pub struct NetworkConfiguration {
    // TODO: think more about config parameters
    pub max_incoming_connections: usize,
    pub max_outgoing_connections: usize,
    pub tcp_nodelay: bool,
    pub tcp_keep_alive: Option<u32>,
    pub tcp_reconnect_timeout: u64,
    pub tcp_reconnect_timeout_max: u64,
}

impl Default for NetworkConfiguration {
    fn default() -> NetworkConfiguration {
        NetworkConfiguration {
            max_incoming_connections: 128,
            max_outgoing_connections: 128,
            tcp_keep_alive: None,
            tcp_nodelay: false,
            tcp_reconnect_timeout: 500,
            tcp_reconnect_timeout_max: 600000,
        }
    }
}

// TODO Implement generic ConnectionPool struct to avoid copy paste.
// Write proper code to configure outgoing streams
pub struct Network {
    config: NetworkConfiguration,
    listen_address: SocketAddr,
    listener: Option<TcpListener>,

    incoming: Slab<IncomingConnection>,
    outgoing: Slab<OutgoingConnection>,
    // FIXME addresses only needs for outgoing connections
    addresses: HashMap<SocketAddr, PeerId>,

    reconnects: HashMap<SocketAddr, MioTimeout>,
}

enum PeerKind {
    Server,
    Incoming,
    Outgoing,
}

fn make_io_error<T: Borrow<str>>(s: T) -> io::Error {
    io::Error::new(io::ErrorKind::Other, s.borrow())
}

impl Network {
    pub fn with_config(address: SocketAddr, config: NetworkConfiguration) -> Network {
        Network {
            config: config,
            listen_address: address,
            listener: None,

            incoming: Slab::new_starting_at(Token(2), config.max_incoming_connections),
            outgoing: Slab::new_starting_at(Token(2 + config.max_incoming_connections),
                                            config.max_outgoing_connections),
            addresses: HashMap::new(),

            reconnects: HashMap::new(),
        }
    }

    pub fn address(&self) -> &SocketAddr {
        &self.listen_address
    }

    // TODO use error trait
    pub fn bind<H: EventHandler>(&mut self, event_loop: &mut EventLoop<H>) -> io::Result<()> {
        if self.listener.is_some() {
            return Err(make_io_error("Already binded"));
        }
        let listener = TcpListener::bind(&self.listen_address)?;
        event_loop.register(&listener, SERVER_ID, EventSet::readable(), PollOpt::edge())?;
        self.listener = Some(listener);
        Ok(())
    }

    // TODO: Use ticks for fast reregistering sockets
    // TODO: Implement Connections collection with (re)registering
    pub fn io<H: EventHandler>(&mut self,
                               event_loop: &mut EventLoop<H>,
                               handler: &mut H,
                               id: PeerId,
                               set: EventSet)
                               -> io::Result<()> {

        match self.peer_kind(id) {
            PeerKind::Server => {
                // Accept new connections
                // FIXME: Fail-safe accepting of new connections?
                let pair = match self.listener {
                    Some(ref listener) => listener.accept()?,
                    None => None,
                };
                if let Some((mut stream, address)) = pair {
                    self.configure_stream(&mut stream)?;
                    let peer = IncomingConnection::new(stream, address);
                    self.add_incoming_connection(event_loop, peer)?;

                    trace!("{}: Accepted incoming connection from {} id: {}",
                           self.address(),
                           address,
                           id.0);
                }
                return Ok(());
            }
            PeerKind::Incoming => {
                if !self.incoming.contains(id) {
                    return Ok(());
                }

                if set.is_hup() | set.is_error() {
                    trace!("{}: incoming connection with addr {} closed",
                           self.address(),
                           self.incoming[id].address());

                    self.remove_incoming_connection(event_loop, id);
                    return Ok(());
                }

                if set.is_readable() {
                    loop {
                        match self.incoming[id].try_read() {
                            Ok(Some(buf)) => {
                                let msg = RawMessage::new(buf);
                                handler.handle_event(Event::Incoming(msg));
                            }
                            Ok(None) => return Ok(()),
                            Err(e) => {
                                self.remove_incoming_connection(event_loop, id);
                                return Err(e);
                            }
                        };
                    }
                }
            }
            PeerKind::Outgoing => {
                if !self.outgoing.contains(id) {
                    return Ok(());
                }

                if set.is_hup() | set.is_error() {
                    let address = *self.outgoing[id].address();

                    trace!("{}: outgoing connection with addr {} closed",
                           self.address(),
                           self.outgoing[id].address());

                    self.remove_outgoing_connection(event_loop, id);
                    if !self.reconnects.contains_key(&address) {
                        handler.handle_event(Event::Disconnected(address));
                    }
                    return Ok(());
                }

                if set.is_writable() {
                    let address = *self.outgoing[id].address();
                    let r = {
                        // Write data into socket
                        self.outgoing[id].try_write()?;
                        event_loop.reregister(self.outgoing[id].socket(),
                                        id,
                                        self.outgoing[id].interest(),
                                        PollOpt::edge())?;
                        Ok(())
                    };
                    if let Err(e) = r {
                        self.remove_outgoing_connection(event_loop, id);
                        handler.handle_event(Event::Disconnected(address));
                        return Err(e);
                    }

                    if self.mark_connected(event_loop, id) {
                        trace!("{}: Handle connected with={}", self.address(), address);
                        handler.handle_event(Event::Connected(address));
                    }
                }
            }
        }
        Ok(())
    }

    pub fn tick<H: EventHandler>(&mut self, _: &mut EventLoop<H>) {}

    pub fn send_to<H: EventHandler>(&mut self,
                                    event_loop: &mut EventLoop<H>,
                                    address: &SocketAddr,
                                    message: RawMessage)
                                    -> io::Result<()> {
        match self.get_outgoing_peer(address) {
            Ok(id) => {
                self.outgoing[id]
                    .send(message)
                    .and_then(|_| {
                        event_loop.reregister(self.outgoing[id].socket(),
                                        id,
                                        self.outgoing[id].interest(),
                                        PollOpt::edge())?;
                        self.mark_connected(event_loop, id);
                        Ok(())
                    })
                    .or_else(|e| {
                        self.remove_outgoing_connection(event_loop, id);
                        Err(e)
                    })
            }
            Err(e) => Err(e),
        }
    }

    pub fn connect<H: EventHandler>(&mut self,
                                    event_loop: &mut EventLoop<H>,
                                    address: &SocketAddr)
                                    -> io::Result<()> {
        if !self.is_connected(address) {
            self.add_reconnect_request(event_loop, *address)?;

            let mut stream = TcpStream::connect(address)?;
            self.configure_stream(&mut stream)?;
            let peer = OutgoingConnection::new(stream, *address);
            let id = self.add_outgoing_connection(event_loop, peer)?;

            trace!("{}: Establish connection with {}, id: {}",
                   self.address(),
                   address,
                   id.0);
        }
        Ok(())
    }

    pub fn is_connected(&self, address: &SocketAddr) -> bool {
        self.addresses.contains_key(address)
    }

    pub fn handle_timeout<H: EventHandler>(&mut self,
                                           event_loop: &mut EventLoop<H>,
                                           timeout: InternalTimeout) {
        match timeout {
            InternalTimeout::Reconnect(addr, delay) => {
                if self.reconnects.contains_key(&addr) {
                    trace!("Try to reconnect with delay {}", delay);

                    if let Err(e) = self.connect(event_loop, &addr) {
                        error!("{}: Unable to create connection to addr {}, error: {:?}",
                               self.address(),
                               addr,
                               e);
                    }

                    let delay = min((delay as f32 * RECONNECT_GROW_FACTOR) as u64,
                                    self.config.tcp_reconnect_timeout_max);

                    if let Err(e) = self.add_reconnect_timeout(event_loop, addr, delay) {
                        error!("{}: Unable to add timeout, error: {:?}", self.address(), e);
                    }
                }
            }
        }
    }

    fn peer_kind(&self, id: PeerId) -> PeerKind {
        if id == SERVER_ID {
            PeerKind::Server
        } else if id.0 >= (2 + self.config.max_incoming_connections) {
            PeerKind::Outgoing
        } else {
            PeerKind::Incoming
        }
    }

    fn get_outgoing_peer(&self, addr: &SocketAddr) -> io::Result<PeerId> {
        if let Some(id) = self.addresses.get(addr) {
            return Ok(*id);
        };
        Err(make_io_error(format!("{}: Outgoing peer not found {}", self.address(), addr)))
    }

    fn add_incoming_connection<H: EventHandler>(&mut self,
                                                event_loop: &mut EventLoop<H>,
                                                connection: IncomingConnection)
                                                -> io::Result<PeerId> {
        let address = *connection.address();
        let id = self.incoming
            .insert(connection)
            .map_err(|_| make_io_error("Maximum incoming onnections"))?;
        self.addresses.insert(address, id);

        let r = event_loop.register(self.incoming[id].socket(),
                                    id,
                                    self.incoming[id].interest(),
                                    PollOpt::edge());

        if let Err(e) = r {
            self.remove_incoming_connection(event_loop, id);
            return Err(e);
        }
        Ok(id)
    }

    fn add_outgoing_connection<H: EventHandler>(&mut self,
                                                event_loop: &mut EventLoop<H>,
                                                connection: OutgoingConnection)
                                                -> io::Result<PeerId> {
        let address = *connection.address();
        let id = self.outgoing
            .insert(connection)
            .map_err(|_| make_io_error("Maximum outgoing onnections"))?;
        self.addresses.insert(address, id);

        let r = event_loop.register(self.outgoing[id].socket(),
                                    id,
                                    self.outgoing[id].interest() | EventSet::writable(),
                                    PollOpt::edge());

        if let Err(e) = r {
            self.remove_outgoing_connection(event_loop, id);
            return Err(e);
        }
        Ok(id)
    }

    fn remove_incoming_connection<H: EventHandler>(&mut self,
                                                   event_loop: &mut EventLoop<H>,
                                                   id: PeerId) {
        let addr = *self.incoming[id].address();
        self.addresses.remove(&addr);
        if let Some(connection) = self.incoming.remove(id) {
            if let Err(e) = event_loop.deregister(connection.socket()) {
                error!("{}: Unable to deregister incoming connection, id: {}, error: {:?}",
                       self.address(),
                       id.0,
                       e);
            }
        }
    }

    fn remove_outgoing_connection<H: EventHandler>(&mut self,
                                                   event_loop: &mut EventLoop<H>,
                                                   id: PeerId) {
        let addr = *self.outgoing[id].address();
        self.addresses.remove(&addr);
        if let Some(connection) = self.outgoing.remove(id) {
            if let Err(e) = event_loop.deregister(connection.socket()) {
                error!("{}: Unable to deregister outgoing connection, id: {}, error: {:?}",
                       self.address(),
                       id.0,
                       e);
            }
        }
    }

    fn configure_stream(&self, stream: &mut TcpStream) -> io::Result<()> {
        stream.take_socket_error()?;
        stream.set_keepalive(self.config.tcp_keep_alive)?;
        stream.set_nodelay(self.config.tcp_nodelay)?;
        stream.take_socket_error()
    }

    fn add_reconnect_request<H: EventHandler>(&mut self,
                                              event_loop: &mut EventLoop<H>,
                                              address: SocketAddr)
                                              -> io::Result<()> {
        if !self.reconnects.contains_key(&address) {
            let delay = self.config.tcp_reconnect_timeout;
            return self.add_reconnect_timeout(event_loop, address, delay);
        }
        Ok(())
    }

    fn add_reconnect_timeout<H: EventHandler>(&mut self,
                                              event_loop: &mut EventLoop<H>,
                                              address: SocketAddr,
                                              delay: u64)
                                              -> io::Result<()> {
        trace!("{}: Add reconnect timeout to={}, delay={}",
               self.address(),
               address,
               delay);
        let reconnect = Timeout::Internal(InternalTimeout::Reconnect(address, delay));
        let timeout = event_loop.timeout_ms(reconnect, delay)
            .map_err(|e| make_io_error(format!("A mio error occured {:?}", e)))?;
        self.reconnects.insert(address, timeout);
        Ok(())
    }

    fn mark_connected<H: EventHandler>(&mut self,
                                       event_loop: &mut EventLoop<H>,
                                       id: Token)
                                       -> bool {
        let address = *self.outgoing[id].address();
        self.clear_reconnect_request(event_loop, &address)
    }

    fn clear_reconnect_request<H: EventHandler>(&mut self,
                                                event_loop: &mut EventLoop<H>,
                                                addr: &SocketAddr)
                                                -> bool {
        if let Some(timeout) = self.reconnects.remove(addr) {
            trace!("{}: Clear reconnect timeout to={}", self.address(), addr);
            event_loop.clear_timeout(timeout);
            return true;
        }
        false
    }
}

#[cfg(test)]
mod tests {}
<|endoftext|>


<commit_before>// Copyright (c) 2020 Jason Ish
//
// 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.

extern crate nom;

use std::io::BufRead;

use nom::{
    bytes::complete::tag,
    character::complete::multispace0,
    error::{ErrorKind, ParseError},
    multi::many0,
    sequence::{preceded, tuple},
    IResult,
};

static WHITESPACE: &str = " \t\r\n";

#[derive(Debug, PartialEq)]
enum InternalError<I> {
    NoListEnd,
    Nom(I, ErrorKind),
}

impl<I> ParseError<I> for InternalError<I> {
    fn from_error_kind(input: I, kind: ErrorKind) -> Self {
        InternalError::Nom(input, kind)
    }

    fn append(_input: I, _kind: ErrorKind, other: Self) -> Self {
        other
    }
}

#[derive(Debug, PartialEq, Default)]
pub struct RuleHeader {
    pub action: String,
    pub proto: String,
    pub src_addr: String,
    pub src_port: String,
    pub direction: String,
    pub dst_addr: String,
    pub dst_port: String,
}

#[derive(Debug, PartialEq)]
pub struct RuleOption {
    pub key: String,
    pub val: Option<String>,
}

#[derive(Debug)]
pub struct TokenizedRule {
    pub header: RuleHeader,
    pub options: Vec<RuleOption>,
    pub original: String,
}

fn parse_header_token(input: &str) -> IResult<&str, &str, InternalError<&str>> {
    preceded(multispace0, nom::bytes::complete::is_not(WHITESPACE))(input)
}

/// Get the next rule option key (name).
///
/// This will return the next set of characters terminated by ; or : with all
/// leading whitespace removed.
fn get_option_key(input: &str) -> IResult<&str, &str, InternalError<&str>> {
    preceded(multispace0, nom::bytes::complete::is_not(";:"))(input)
}

/// Get the next option value.
///
/// This will return the next sequence of characters terminated by ';',
/// while removing any instances of an escaped ';'.
fn get_option_value(input: &str) -> IResult<&str, String, InternalError<&str>> {
    let mut output = Vec::new();
    let mut escaped = false;
    let mut end = 0;
    for (i, c) in input.chars().enumerate() {
        end = i;
        if c == '\\' {
            escaped = true;
        } else if escaped {
            if c == ';' {
                output.push(c);
            } else {
                output.push('\\');
                output.push(c);
            }
            escaped = false;
        } else if c == ';' {
            // Eat the ';'.
            end += 1;
            break;
        } else {
            output.push(c);
        }
    }
    let (_, rem) = input.split_at(end);
    Ok((rem, output.into_iter().collect()))
}

/// Parse a list as a token.
///
/// Returns the complete list as a single token, handling nested
/// lists, and spaces between listen items.
///
/// As this doesn't use nom, we make it look like a nom parser,
/// and use nom error types that are close enough.
fn parse_list_token(input: &str) -> IResult<&str, &str, InternalError<&str>> {
    // Peek into the input to make sure it looks like the start of a list.
    let input = preceded(
        multispace0,
        nom::combinator::peek(nom::bytes::complete::tag("[")),
    )(input)?
    .0;
    let mut in_list = 0;
    let mut end = 0;

    for (i, c) in input.chars().enumerate() {
        match c {
            '[' => in_list += 1,
            ']' => in_list -= 1,
            _ => {}
        }

        if in_list == 0 {
            end = i + 1;
            break;
        }
    }

    // If we got here and are still in the list, then we ran out of input.
    if in_list > 0 {
        return Err(nom::Err::Error(InternalError::NoListEnd));
    }

    let (list, rem) = input.split_at(end);
    Ok((rem, list))
}

fn parse_header(input: &str) -> IResult<&str, RuleHeader, InternalError<&str>> {
    let maybe_list = &nom::branch::alt((parse_list_token, parse_header_token));
    let (rem, (action, proto, src_addr, src_port, direction, dst_addr, dst_port)) = tuple((
        // Action.
        parse_header_token,
        // Proto.
        maybe_list,
        // Source address.
        maybe_list,
        // Source port.
        maybe_list,
        // Direction.
        parse_header_token,
        // Destination address.
        maybe_list,
        // Destination port.
        maybe_list,
    ))(input)?;
    Ok((
        rem,
        RuleHeader {
            action: String::from(action),
            proto: String::from(proto),
            src_addr: String::from(src_addr),
            src_port: String::from(src_port),
            direction: String::from(direction),
            dst_addr: String::from(dst_addr),
            dst_port: String::from(dst_port),
        },
    ))
}

fn parse_option(input: &str) -> IResult<&str, RuleOption, InternalError<&str>> {
    // Eat any leading space, then parse up to a : or a ;.
    let (input, _) = multispace0(input)?;
    let (input, key) = get_option_key(input)?;
    let (input, sep) = nom::character::complete::one_of(";:")(input)?;
    if sep == ';' {
        return Ok((
            input,
            RuleOption {
                key: String::from(key),
                val: None,
            },
        ));
    }

    // Drop any leading whitespace then parse up to ';', accounting for
    // escaped occurrences of ';'.
    let (input, _) = multispace0(input)?;
    let (input, val) = get_option_value(input)?;

    Ok((
        input,
        RuleOption {
            key: String::from(key),
            val: Some(strip_quotes(&val)),
        },
    ))
}

/// Remove quotes from a string, but preserve any escaped quotes.
fn strip_quotes(input: &str) -> String {
    let mut escaped = false;
    let mut out: Vec<char> = Vec::new();

    for c in input.chars() {
        if escaped {
            out.push(c);
            escaped = false;
        } else {
            match c {
                '"' => {}
                '\\' => {
                    escaped = true;
                }
                _ => {
                    out.push(c);
                }
            }
        }
    }

    out.iter().collect()
}

fn internal_parse_rule(input: &str) -> IResult<&str, TokenizedRule, InternalError<&str>> {
    let original = String::from(input);
    let (input, header) = parse_header(input)?;
    let (input, _) = preceded(multispace0, tag("("))(input)?;
    let (input, options) = many0(parse_option)(input)?;
    let (input, _) = preceded(multispace0, tag(")"))(input)?;

    Ok((
        input,
        TokenizedRule {
            header,
            options,
            original,
        },
    ))
}

pub fn parse_rule(input: &str) -> anyhow::Result<TokenizedRule> {
    match internal_parse_rule(input) {
        Ok((_, rule)) => return Ok(rule),
        Err(err) => return Err(anyhow::anyhow!(err.to_string())),
    }
}

/// Read the next rule from a reader.
///
/// This will actually return any line it reads, but will join together
/// multiline rules into a single line rule.
///
/// Ok(None) will be returned on EOF, and an error will be returned on read
/// error.
pub fn read_next_rule(input: &mut dyn BufRead) -> Result<Option<String>, std::io::Error> {
    let mut line = String::new();
    loop {
        let mut tmp = String::new();
        let n = input.read_line(&mut tmp)?;
        if n == 0 {
            return Ok(None);
        }

        let tmp = tmp.trim();

        if !tmp.ends_with('\\') {
            line.push_str(tmp);
            break;
        }

        line.push_str(&tmp[..tmp.len() - 1]);
    }
    Ok(Some(line))
}

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

    #[test]
    fn test_parse_token() {
        assert_eq!(parse_header_token("alert"), Ok(("", "alert")));
        assert_eq!(parse_header_token(" alert"), Ok(("", "alert")));
        assert_eq!(parse_header_token(" alert "), Ok((" ", "alert")));
        assert_eq!(parse_header_token("http_uri"), Ok(("", "http_uri")));
    }

    #[test]
    fn test_parse_quoted_string() {
        assert_eq!(
            strip_quotes(r#""some quoted \" string""#),
            r#"some quoted " string"#
        );
    }

    #[test]
    fn test_parse_list() {
        assert_eq!(parse_list_token("[1]"), Ok(("", "[1]")));
        assert_eq!(parse_list_token("[1,2,3]"), Ok(("", "[1,2,3]")));
        assert_eq!(parse_list_token(" [1,2,3]"), Ok(("", "[1,2,3]")));
        assert_eq!(
            parse_list_token(" [1,2,3,[a,b,c]]"),
            Ok(("", "[1,2,3,[a,b,c]]"))
        );

        assert!(parse_list_token("1,2,3]").is_err());
        assert!(parse_list_token("[1,2,3").is_err());

        assert!(parse_list_token("token").is_err());
    }

    #[test]
    fn test_parse_header() {
        let rule = parse_header("alert tcp any any -> any any");
        assert_eq!(
            rule,
            Ok((
                "",
                RuleHeader {
                    action: String::from("alert"),
                    proto: String::from("tcp"),
                    src_addr: String::from("any"),
                    src_port: String::from("any"),
                    direction: String::from("->"),
                    dst_addr: String::from("any"),
                    dst_port: String::from("any"),
                    ..Default::default()
                }
            ))
        );
    }

    #[test]
    fn test_parse_option_without_value() {
        assert_eq!(
            Ok((
                "",
                RuleOption {
                    key: String::from("http_uri"),
                    val: None,
                }
            )),
            parse_option("http_uri;")
        );
    }

    #[test]
    fn test_parse_option() {
        assert_eq!(
            parse_option("msg:value;"),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some("value".to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option("msg:value with spaces;"),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some("value with spaces".to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option("msg:terminated value with spaces;"),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some("terminated value with spaces".to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option(r#"msg: an escaped \; terminant; next_option;"#),
            Ok((
                " next_option;",
                RuleOption {
                    key: String::from("msg"),
                    val: Some("an escaped ; terminant".to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option(r#"msg: "A Quoted Message";"#),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some(r#"A Quoted Message"#.to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option(r#"msg: "A Quoted Message with escaped \" quotes.";"#),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some(r#"A Quoted Message with escaped " quotes."#.to_string()),
                }
            ))
        );
    }

    #[test]
    fn test_parse_rule() {
        let rule = parse_rule(r#"alert ip any any -> any any (msg:"Some Message with a \" Quote"; metadata: key, val; sid:1; rev:1;)"#).unwrap();
        assert_eq!(rule.header.action, "alert");
    }

    #[test]
    fn test_tokenize_option_value() {
        assert_eq!(
            get_option_value("one; two"),
            Ok((" two", String::from("one")))
        );

        assert_eq!(
            get_option_value("one\\;; two"),
            Ok((" two", String::from("one;")))
        );

        let optval = r#""/^(?:[a-zA-Z0-9_%+])*(?:[\x2c\x22\x27\x28]|\x252[c278])/PRi";"#;
        let expected = r#""/^(?:[a-zA-Z0-9_%+])*(?:[\x2c\x22\x27\x28]|\x252[c278])/PRi""#;
        assert_eq!(get_option_value(optval), Ok(("", String::from(expected))));
    }

    #[test]
    fn test_parse_from_reader() {
        let input = r#"alert ip any any -> any any (\
            msg:"TEST RULE"; sid:1; rev:1;)"#;
        let mut reader = input.as_bytes();
        let next = read_next_rule(&mut reader).unwrap();
        assert_eq!(
            next,
            Some(r#"alert ip any any -> any any (msg:"TEST RULE"; sid:1; rev:1;)"#.to_string())
        );
        assert_eq!(read_next_rule(&mut reader).unwrap(), None);
    }
}
<commit_msg>rule parser: parse disabled rules<commit_after>// Copyright (c) 2020 Jason Ish
//
// 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.

extern crate nom;

use std::io::BufRead;

use nom::{
    bytes::complete::tag,
    character::complete::multispace0,
    error::{ErrorKind, ParseError},
    multi::{many0, many0_count},
    sequence::{preceded, tuple},
    IResult,
};

static WHITESPACE: &str = " \t\r\n";

#[derive(Debug, PartialEq)]
enum InternalError<I> {
    NoListEnd,
    Nom(I, ErrorKind),
}

impl<I> ParseError<I> for InternalError<I> {
    fn from_error_kind(input: I, kind: ErrorKind) -> Self {
        InternalError::Nom(input, kind)
    }

    fn append(_input: I, _kind: ErrorKind, other: Self) -> Self {
        other
    }
}

#[derive(Debug, PartialEq, Default)]
pub struct RuleHeader {
    pub action: String,
    pub proto: String,
    pub src_addr: String,
    pub src_port: String,
    pub direction: String,
    pub dst_addr: String,
    pub dst_port: String,
}

#[derive(Debug, PartialEq)]
pub struct RuleOption {
    pub key: String,
    pub val: Option<String>,
}

#[derive(Debug)]
pub struct TokenizedRule {
    pub disabled: bool,
    pub header: RuleHeader,
    pub options: Vec<RuleOption>,
    pub original: String,
}

fn parse_header_token(input: &str) -> IResult<&str, &str, InternalError<&str>> {
    preceded(multispace0, nom::bytes::complete::is_not(WHITESPACE))(input)
}

/// Get the next rule option key (name).
///
/// This will return the next set of characters terminated by ; or : with all
/// leading whitespace removed.
fn get_option_key(input: &str) -> IResult<&str, &str, InternalError<&str>> {
    preceded(multispace0, nom::bytes::complete::is_not(";:"))(input)
}

/// Get the next option value.
///
/// This will return the next sequence of characters terminated by ';',
/// while removing any instances of an escaped ';'.
fn get_option_value(input: &str) -> IResult<&str, String, InternalError<&str>> {
    let mut output = Vec::new();
    let mut escaped = false;
    let mut end = 0;
    for (i, c) in input.chars().enumerate() {
        end = i;
        if c == '\\' {
            escaped = true;
        } else if escaped {
            if c == ';' {
                output.push(c);
            } else {
                output.push('\\');
                output.push(c);
            }
            escaped = false;
        } else if c == ';' {
            // Eat the ';'.
            end += 1;
            break;
        } else {
            output.push(c);
        }
    }
    let (_, rem) = input.split_at(end);
    Ok((rem, output.into_iter().collect()))
}

/// Parse a list as a token.
///
/// Returns the complete list as a single token, handling nested
/// lists, and spaces between listen items.
///
/// As this doesn't use nom, we make it look like a nom parser,
/// and use nom error types that are close enough.
fn parse_list_token(input: &str) -> IResult<&str, &str, InternalError<&str>> {
    // Peek into the input to make sure it looks like the start of a list.
    let input = preceded(
        multispace0,
        nom::combinator::peek(nom::bytes::complete::tag("[")),
    )(input)?
    .0;
    let mut in_list = 0;
    let mut end = 0;

    for (i, c) in input.chars().enumerate() {
        match c {
            '[' => in_list += 1,
            ']' => in_list -= 1,
            _ => {}
        }

        if in_list == 0 {
            end = i + 1;
            break;
        }
    }

    // If we got here and are still in the list, then we ran out of input.
    if in_list > 0 {
        return Err(nom::Err::Error(InternalError::NoListEnd));
    }

    let (list, rem) = input.split_at(end);
    Ok((rem, list))
}

fn parse_header(input: &str) -> IResult<&str, RuleHeader, InternalError<&str>> {
    let maybe_list = &nom::branch::alt((parse_list_token, parse_header_token));
    let (rem, (action, proto, src_addr, src_port, direction, dst_addr, dst_port)) = tuple((
        // Action.
        parse_header_token,
        // Proto.
        maybe_list,
        // Source address.
        maybe_list,
        // Source port.
        maybe_list,
        // Direction.
        parse_header_token,
        // Destination address.
        maybe_list,
        // Destination port.
        maybe_list,
    ))(input)?;
    Ok((
        rem,
        RuleHeader {
            action: String::from(action),
            proto: String::from(proto),
            src_addr: String::from(src_addr),
            src_port: String::from(src_port),
            direction: String::from(direction),
            dst_addr: String::from(dst_addr),
            dst_port: String::from(dst_port),
        },
    ))
}

fn parse_option(input: &str) -> IResult<&str, RuleOption, InternalError<&str>> {
    // Eat any leading space, then parse up to a : or a ;.
    let (input, _) = multispace0(input)?;
    let (input, key) = get_option_key(input)?;
    let (input, sep) = nom::character::complete::one_of(";:")(input)?;
    if sep == ';' {
        return Ok((
            input,
            RuleOption {
                key: String::from(key),
                val: None,
            },
        ));
    }

    // Drop any leading whitespace then parse up to ';', accounting for
    // escaped occurrences of ';'.
    let (input, _) = multispace0(input)?;
    let (input, val) = get_option_value(input)?;

    Ok((
        input,
        RuleOption {
            key: String::from(key),
            val: Some(strip_quotes(&val)),
        },
    ))
}

/// Remove quotes from a string, but preserve any escaped quotes.
fn strip_quotes(input: &str) -> String {
    let mut escaped = false;
    let mut out: Vec<char> = Vec::new();

    for c in input.chars() {
        if escaped {
            out.push(c);
            escaped = false;
        } else {
            match c {
                '"' => {}
                '\\' => {
                    escaped = true;
                }
                _ => {
                    out.push(c);
                }
            }
        }
    }

    out.iter().collect()
}

fn internal_parse_rule(input: &str) -> IResult<&str, TokenizedRule, InternalError<&str>> {
    let original = String::from(input);
    let (input, disabled) = preceded(multispace0, many0_count(tag("#")))(input)?;
    let (input, header) = parse_header(input)?;
    let (input, _) = preceded(multispace0, tag("("))(input)?;
    let (input, options) = many0(parse_option)(input)?;
    let (input, _) = preceded(multispace0, tag(")"))(input)?;

    Ok((
        input,
        TokenizedRule {
            disabled: disabled > 0,
            header,
            options,
            original,
        },
    ))
}

pub fn parse_rule(input: &str) -> anyhow::Result<TokenizedRule> {
    match internal_parse_rule(input) {
        Ok((_, rule)) => return Ok(rule),
        Err(err) => return Err(anyhow::anyhow!(err.to_string())),
    }
}

/// Read the next rule from a reader.
///
/// This will actually return any line it reads, but will join together
/// multiline rules into a single line rule.
///
/// Ok(None) will be returned on EOF, and an error will be returned on read
/// error.
pub fn read_next_rule(input: &mut dyn BufRead) -> Result<Option<String>, std::io::Error> {
    let mut line = String::new();
    loop {
        let mut tmp = String::new();
        let n = input.read_line(&mut tmp)?;
        if n == 0 {
            return Ok(None);
        }

        let tmp = tmp.trim();

        if !tmp.ends_with('\\') {
            line.push_str(tmp);
            break;
        }

        line.push_str(&tmp[..tmp.len() - 1]);
    }
    Ok(Some(line))
}

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

    #[test]
    fn test_parse_token() {
        assert_eq!(parse_header_token("alert"), Ok(("", "alert")));
        assert_eq!(parse_header_token(" alert"), Ok(("", "alert")));
        assert_eq!(parse_header_token(" alert "), Ok((" ", "alert")));
        assert_eq!(parse_header_token("http_uri"), Ok(("", "http_uri")));
    }

    #[test]
    fn test_parse_quoted_string() {
        assert_eq!(
            strip_quotes(r#""some quoted \" string""#),
            r#"some quoted " string"#
        );
    }

    #[test]
    fn test_parse_list() {
        assert_eq!(parse_list_token("[1]"), Ok(("", "[1]")));
        assert_eq!(parse_list_token("[1,2,3]"), Ok(("", "[1,2,3]")));
        assert_eq!(parse_list_token(" [1,2,3]"), Ok(("", "[1,2,3]")));
        assert_eq!(
            parse_list_token(" [1,2,3,[a,b,c]]"),
            Ok(("", "[1,2,3,[a,b,c]]"))
        );

        assert!(parse_list_token("1,2,3]").is_err());
        assert!(parse_list_token("[1,2,3").is_err());

        assert!(parse_list_token("token").is_err());
    }

    #[test]
    fn test_parse_header() {
        let rule = parse_header("alert tcp any any -> any any");
        assert_eq!(
            rule,
            Ok((
                "",
                RuleHeader {
                    action: String::from("alert"),
                    proto: String::from("tcp"),
                    src_addr: String::from("any"),
                    src_port: String::from("any"),
                    direction: String::from("->"),
                    dst_addr: String::from("any"),
                    dst_port: String::from("any"),
                    ..Default::default()
                }
            ))
        );
    }

    #[test]
    fn test_parse_option_without_value() {
        assert_eq!(
            Ok((
                "",
                RuleOption {
                    key: String::from("http_uri"),
                    val: None,
                }
            )),
            parse_option("http_uri;")
        );
    }

    #[test]
    fn test_parse_option() {
        assert_eq!(
            parse_option("msg:value;"),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some("value".to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option("msg:value with spaces;"),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some("value with spaces".to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option("msg:terminated value with spaces;"),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some("terminated value with spaces".to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option(r#"msg: an escaped \; terminant; next_option;"#),
            Ok((
                " next_option;",
                RuleOption {
                    key: String::from("msg"),
                    val: Some("an escaped ; terminant".to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option(r#"msg: "A Quoted Message";"#),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some(r#"A Quoted Message"#.to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option(r#"msg: "A Quoted Message with escaped \" quotes.";"#),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some(r#"A Quoted Message with escaped " quotes."#.to_string()),
                }
            ))
        );
    }

    #[test]
    fn test_parse_rule() {
        let rule = parse_rule(r#"alert ip any any -> any any (msg:"Some Message with a \" Quote"; metadata: key, val; sid:1; rev:1;)"#).unwrap();
        assert_eq!(rule.disabled, false);
        assert_eq!(rule.header.action, "alert");

        let rule = parse_rule(r#"  alert ip any any -> any any (msg:"Some Message with a \" Quote"; metadata: key, val; sid:1; rev:1;)"#).unwrap();
        assert_eq!(rule.disabled, false);
        assert_eq!(rule.header.action, "alert");

        let rule = parse_rule(r#"#alert ip any any -> any any (msg:"Some Message with a \" Quote"; metadata: key, val; sid:1; rev:1;)"#).unwrap();
        assert_eq!(rule.disabled, true);
        assert_eq!(rule.header.action, "alert");

        let rule = parse_rule(r#"# alert ip any any -> any any (msg:"Some Message with a \" Quote"; metadata: key, val; sid:1; rev:1;)"#).unwrap();
        assert_eq!(rule.disabled, true);
        assert_eq!(rule.header.action, "alert");

        let rule = parse_rule(r#"###   alert ip any any -> any any (msg:"Some Message with a \" Quote"; metadata: key, val; sid:1; rev:1;)"#).unwrap();
        assert_eq!(rule.disabled, true);
        assert_eq!(rule.header.action, "alert");
    }

    #[test]
    fn test_tokenize_option_value() {
        assert_eq!(
            get_option_value("one; two"),
            Ok((" two", String::from("one")))
        );

        assert_eq!(
            get_option_value("one\\;; two"),
            Ok((" two", String::from("one;")))
        );

        let optval = r#""/^(?:[a-zA-Z0-9_%+])*(?:[\x2c\x22\x27\x28]|\x252[c278])/PRi";"#;
        let expected = r#""/^(?:[a-zA-Z0-9_%+])*(?:[\x2c\x22\x27\x28]|\x252[c278])/PRi""#;
        assert_eq!(get_option_value(optval), Ok(("", String::from(expected))));
    }

    #[test]
    fn test_parse_from_reader() {
        let input = r#"alert ip any any -> any any (\
            msg:"TEST RULE"; sid:1; rev:1;)"#;
        let mut reader = input.as_bytes();
        let next = read_next_rule(&mut reader).unwrap();
        assert_eq!(
            next,
            Some(r#"alert ip any any -> any any (msg:"TEST RULE"; sid:1; rev:1;)"#.to_string())
        );
        assert_eq!(read_next_rule(&mut reader).unwrap(), None);
    }
}
<|endoftext|>


<commit_before>use core::rand::*;
use core::f64::sqrt;
use core::f64::pow;

#[deriving(Eq)]
pub struct Point {
    x : f64,
    y : f64
}

pub impl Point {
    fn new() -> Point {
        Point{x:0.0, y:0.0}
    }

    fn dist(&self, other : &Point) -> f64 {
        sqrt(pow(self.x-other.x,2.0)+pow(self.y-other.y,2.0))
    }
}

impl ToStr for Point {
    fn to_str(&self) -> ~str {
        fmt!("(%f, %f)", self.x as float, self.y as float)
    }
}

//
// * Génération aléatoire
//
impl Rand for Point {
    fn rand<R: Rng>(rng: &R) -> Point {
        Point { x: rng.gen::<f64>(), y: rng.gen::<f64>()}
    }
}

//
// * Opérateurs arithmétiques
//
// Addition
trait RhsOfAdd<Result> {
  fn add_lhs_to(&self, lhs: &Point) -> Result;
}
impl RhsOfAdd<Point> for Point {
  fn add_lhs_to(&self, lhs: &Point) -> Point {
    Point{x: (self.x+lhs.x), y:(self.y+lhs.y)}
   }
}
impl<Result, Rhs: RhsOfAdd<Result> > ops::Add<Rhs, Result> for Point {
  fn add(&self, rhs: &Rhs) -> Result {
    rhs.add_lhs_to(self)
  }
}
// Soustraction
trait RhsOfSub<Result> {
  fn sub_lhs_to(&self, lhs: &Point) -> Result;
}
impl RhsOfSub<Point> for Point {
  fn sub_lhs_to(&self, lhs: &Point) -> Point {
    Point{x: (-self.x+lhs.x), y:(-self.y+lhs.y)}
   }
}
impl RhsOfSub<Point> for f64 {
    fn sub_lhs_to(&self, lhs:&Point) -> Point {
        Point{x: lhs.x - *self, y: lhs.y-*self}
    }
}
impl<Result, Rhs: RhsOfSub<Result> > ops::Sub<Rhs, Result> for Point {
  fn sub(&self, rhs: &Rhs) -> Result {
    rhs.sub_lhs_to(self)
  }
}
// Multiplication
trait RhsOfMul<Result> {
  fn mul_lhs_to(&self, lhs: &Point) -> Result;
}
impl RhsOfMul<Point> for Point {
  fn mul_lhs_to(&self, lhs: &Point) -> Point {
    Point{x: (self.x*lhs.x), y:(self.y*lhs.y)}
   }
}
impl RhsOfMul<Point> for f64 {
  fn mul_lhs_to(&self, lhs: &Point) -> Point {
    Point{x: (self*lhs.x), y:(self*lhs.y)}
   }
}

impl<Result, Rhs: RhsOfMul<Result> > ops::Mul<Rhs, Result> for Point {
  fn mul(&self, rhs: &Rhs) -> Result {
    rhs.mul_lhs_to(self)
  }
}
// Division
trait RhsOfDiv<Result> {
  fn div_lhs_to(&self, lhs: &Point) -> Result;
}
impl RhsOfDiv<Point> for Point {
  fn div_lhs_to(&self, lhs: &Point) -> Point {
    Point{x: (lhs.x/self.x), y:(lhs.y/self.y)}
   }
}
impl RhsOfDiv<Point> for f64 {
  fn div_lhs_to(&self, lhs: &Point) -> Point {
    Point{x: (lhs.x/ *self), y:(lhs.y/ *self)}
   }
}
impl<Result, Rhs: RhsOfDiv<Result> > ops::Div<Rhs, Result> for Point {
  fn div(&self, rhs: &Rhs) -> Result {
    rhs.div_lhs_to(self)
  }
}

//
// * Test
//
#[test]
fn test_dist() {
    let pt1=@Point{ x: 1.0, y: 2.0};
    let pt2=@Point{ x: 2.0, y: 3.0};
    let distance=pt1.dist(pt2);

    assert!(distance==f64::sqrt(2.0));
}

#[test]
fn test_add() {
    let pt1=@Point{ x: 1.0, y: 2.0};
    let pt2=@Point{ x: -1.0, y: 3.0};
    let pt3=*pt1+*pt2;

    assert!((pt3.x==0.0)&&(pt3.y==5.0));
}

#[test]
fn test_sub() {
    let pt1=@Point{ x: 1.0, y: 2.0};
    let pt2=@Point{ x: -1.0, y: 3.0};
    let pt3=*pt1-*pt2;
    assert!((pt3.x==2.0)&&(pt3.y==-1.0));
}
fn test_sub_float() {
    let pt1=Point{ x: 1.0, y: 2.0};
    let pt3=pt1-0.5 as f64;
    assert!((pt3.x==0.5)&&(pt3.y==1.5));
}

#[test]
fn test_mult() {
    let pt1=@Point{ x: 1.0, y: 2.0};
    let pt2=@Point{ x: -1.0, y: 3.0};
    let pt3=(*pt1)*(*pt2);
    assert!((pt3.x==-1.0)&&(pt3.y==6.0));
    let pt4=(*pt1)*(2.0 as f64);
    assert!((pt4.x==2.0)&&(pt4.y==4.0));
}

#[test]
fn test_div() {
    let pt1=@Point{ x: 1.0, y: 2.0};
    let pt2=@Point{ x: -1.0, y: 3.0};
    let pt3=(*pt1)/(*pt2);
    assert!(pt3.dist(&Point{x:-1.0, y:2.0/3.0})==0.0);
    let pt4=(*pt1)/(2.0 as f64);
    assert!((pt4.x==0.5)&&(pt4.y==1.0));
}

#[test]
fn test_rand() {
    let pt1 : Point = rand::random();
    let pt2 : Point = rand::random();
    let pt3 : Point = rand::random();

    assert!(pt1 != pt2); assert!(pt1 != pt3); assert!(pt2 != pt3);
    assert!((pt1.x>=0.0)&&(pt1.y<=1.0));
}
<commit_msg>Addition d'un flottant à un point<commit_after>use core::rand::*;
use core::f64::sqrt;
use core::f64::pow;

#[deriving(Eq)]
pub struct Point {
    x : f64,
    y : f64
}

pub impl Point {
    fn new() -> Point {
        Point{x:0.0, y:0.0}
    }

    fn dist(&self, other : &Point) -> f64 {
        sqrt(pow(self.x-other.x,2.0)+pow(self.y-other.y,2.0))
    }
}

impl ToStr for Point {
    fn to_str(&self) -> ~str {
        fmt!("(%f, %f)", self.x as float, self.y as float)
    }
}

//
// * Génération aléatoire
//
impl Rand for Point {
    fn rand<R: Rng>(rng: &R) -> Point {
        Point { x: rng.gen::<f64>(), y: rng.gen::<f64>()}
    }
}

//
// * Opérateurs arithmétiques
//
// Addition
trait RhsOfAdd<Result> {
  fn add_lhs_to(&self, lhs: &Point) -> Result;
}
impl RhsOfAdd<Point> for Point {
  fn add_lhs_to(&self, lhs: &Point) -> Point {
    Point{x: (self.x+lhs.x), y:(self.y+lhs.y)}
   }
}
impl RhsOfAdd<Point> for f64 {
  fn add_lhs_to(&self, lhs: &Point) -> Point {
    Point{x: (*self+lhs.x), y:(*self+lhs.y)}
   }
}
impl<Result, Rhs: RhsOfAdd<Result> > ops::Add<Rhs, Result> for Point {
  fn add(&self, rhs: &Rhs) -> Result {
    rhs.add_lhs_to(self)
  }
}
// Soustraction
trait RhsOfSub<Result> {
  fn sub_lhs_to(&self, lhs: &Point) -> Result;
}
impl RhsOfSub<Point> for Point {
  fn sub_lhs_to(&self, lhs: &Point) -> Point {
    Point{x: (-self.x+lhs.x), y:(-self.y+lhs.y)}
   }
}
impl RhsOfSub<Point> for f64 {
    fn sub_lhs_to(&self, lhs:&Point) -> Point {
        Point{x: lhs.x - *self, y: lhs.y-*self}
    }
}
impl<Result, Rhs: RhsOfSub<Result> > ops::Sub<Rhs, Result> for Point {
  fn sub(&self, rhs: &Rhs) -> Result {
    rhs.sub_lhs_to(self)
  }
}
// Multiplication
trait RhsOfMul<Result> {
  fn mul_lhs_to(&self, lhs: &Point) -> Result;
}
impl RhsOfMul<Point> for Point {
  fn mul_lhs_to(&self, lhs: &Point) -> Point {
    Point{x: (self.x*lhs.x), y:(self.y*lhs.y)}
   }
}
impl RhsOfMul<Point> for f64 {
  fn mul_lhs_to(&self, lhs: &Point) -> Point {
    Point{x: (self*lhs.x), y:(self*lhs.y)}
   }
}

impl<Result, Rhs: RhsOfMul<Result> > ops::Mul<Rhs, Result> for Point {
  fn mul(&self, rhs: &Rhs) -> Result {
    rhs.mul_lhs_to(self)
  }
}
// Division
trait RhsOfDiv<Result> {
  fn div_lhs_to(&self, lhs: &Point) -> Result;
}
impl RhsOfDiv<Point> for Point {
  fn div_lhs_to(&self, lhs: &Point) -> Point {
    Point{x: (lhs.x/self.x), y:(lhs.y/self.y)}
   }
}
impl RhsOfDiv<Point> for f64 {
  fn div_lhs_to(&self, lhs: &Point) -> Point {
    Point{x: (lhs.x/ *self), y:(lhs.y/ *self)}
   }
}
impl<Result, Rhs: RhsOfDiv<Result> > ops::Div<Rhs, Result> for Point {
  fn div(&self, rhs: &Rhs) -> Result {
    rhs.div_lhs_to(self)
  }
}

//
// * Test
//
#[test]
fn test_dist() {
    let pt1=@Point{ x: 1.0, y: 2.0};
    let pt2=@Point{ x: 2.0, y: 3.0};
    let distance=pt1.dist(pt2);

    assert!(distance==f64::sqrt(2.0));
}

#[test]
fn test_add() {
    let pt1=@Point{ x: 1.0, y: 2.0};
    let pt2=@Point{ x: -1.0, y: 3.0};
    let pt3=*pt1+*pt2;

    assert!((pt3.x==0.0)&&(pt3.y==5.0));
}

#[test]
fn test_add_float() {
    let pt1=Point{ x: 1.0, y: 2.0};
    let pt3=pt1+1.0f64;

    assert!((pt3.x==2.0)&&(pt3.y==3.0));
}

#[test]
fn test_sub() {
    let pt1=@Point{ x: 1.0, y: 2.0};
    let pt2=@Point{ x: -1.0, y: 3.0};
    let pt3=*pt1-*pt2;
    assert!((pt3.x==2.0)&&(pt3.y==-1.0));
}
fn test_sub_float() {
    let pt1=Point{ x: 1.0, y: 2.0};
    let pt3=pt1-0.5 as f64;
    assert!((pt3.x==0.5)&&(pt3.y==1.5));
}

#[test]
fn test_mult() {
    let pt1=@Point{ x: 1.0, y: 2.0};
    let pt2=@Point{ x: -1.0, y: 3.0};
    let pt3=(*pt1)*(*pt2);
    assert!((pt3.x==-1.0)&&(pt3.y==6.0));
    let pt4=(*pt1)*(2.0 as f64);
    assert!((pt4.x==2.0)&&(pt4.y==4.0));
}

#[test]
fn test_div() {
    let pt1=@Point{ x: 1.0, y: 2.0};
    let pt2=@Point{ x: -1.0, y: 3.0};
    let pt3=(*pt1)/(*pt2);
    assert!(pt3.dist(&Point{x:-1.0, y:2.0/3.0})==0.0);
    let pt4=(*pt1)/(2.0 as f64);
    assert!((pt4.x==0.5)&&(pt4.y==1.0));
}

#[test]
fn test_rand() {
    let pt1 : Point = rand::random();
    let pt2 : Point = rand::random();
    let pt3 : Point = rand::random();

    assert!(pt1 != pt2); assert!(pt1 != pt3); assert!(pt2 != pt3);
    assert!((pt1.x>=0.0)&&(pt1.y<=1.0));
}
<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
// Copyright 2016-2017 The Rust Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The Cauchy distribution.

use rand::Rng;
use crate::{Distribution, Standard};
use crate::utils::Float;

/// The Cauchy distribution `Cauchy(median, scale)`.
///
/// This distribution has a density function:
/// `f(x) = 1 / (pi * scale * (1 + ((x - median) / scale)^2))`
///
/// # Example
///
/// ```
/// use rand_distr::{Cauchy, Distribution};
///
/// let cau = Cauchy::new(2.0, 5.0).unwrap();
/// let v = cau.sample(&mut rand::thread_rng());
/// println!("{} is from a Cauchy(2, 5) distribution", v);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct Cauchy<N> {
    median: N,
    scale: N,
}

/// Error type returned from `Cauchy::new`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    /// `scale <= 0` or `nan`.
    ScaleTooSmall,
}

impl<N: Float> Cauchy<N>
where Standard: Distribution<N>
{
    /// Construct a new `Cauchy` with the given shape parameters
    /// `median` the peak location and `scale` the scale factor.
    pub fn new(median: N, scale: N) -> Result<Cauchy<N>, Error> {
        if !(scale > N::from(0.0)) {
            return Err(Error::ScaleTooSmall);
        }
        Ok(Cauchy {
            median,
            scale
        })
    }
}

impl<N: Float> Distribution<N> for Cauchy<N>
where Standard: Distribution<N>
{
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> N {
        // sample from [0, 1)
        let x = Standard.sample(rng);
        // get standard cauchy random number
        // note that π/2 is not exactly representable, even if x=0.5 the result is finite
        let comp_dev = (N::pi() * x).tan();
        // shift and scale according to parameters
        self.median + self.scale * comp_dev
    }
}

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

    fn median(mut numbers: &mut [f64]) -> f64 {
        sort(&mut numbers);
        let mid = numbers.len() / 2;
        numbers[mid]
    }

    fn sort(numbers: &mut [f64]) {
        numbers.sort_by(|a, b| a.partial_cmp(b).unwrap());
    }

    #[test]
    fn test_cauchy_averages() {
        // NOTE: given that the variance and mean are undefined,
        // this test does not have any rigorous statistical meaning.
        let cauchy = Cauchy::new(10.0, 5.0).unwrap();
        let mut rng = crate::test::rng(123);
        let mut numbers: [f64; 1000] = [0.0; 1000];
        let mut sum = 0.0;
        for i in 0..1000 {
            numbers[i] = cauchy.sample(&mut rng);
            sum += numbers[i];
        }
        let median = median(&mut numbers);
        println!("Cauchy median: {}", median);
        assert!((median - 10.0).abs() < 0.4); // not 100% certain, but probable enough
        let mean = sum / 1000.0;
        println!("Cauchy mean: {}", mean);
        // for a Cauchy distribution the mean should not converge
        assert!((mean - 10.0).abs() > 0.4); // not 100% certain, but probable enough
    }

    #[test]
    #[should_panic]
    fn test_cauchy_invalid_scale_zero() {
        Cauchy::new(0.0, 0.0).unwrap();
    }

    #[test]
    #[should_panic]
    fn test_cauchy_invalid_scale_neg() {
        Cauchy::new(0.0, -10.0).unwrap();
    }

    #[test]
    fn value_stability() {
        fn test_samples<N: Float + core::fmt::Debug>(m: N, s: N, expected: &[N])
        where Standard: Distribution<N> {
            let distr = Cauchy::new(m, s).unwrap();
            let mut rng = crate::test::rng(353);
            let mut buf = [m; 4];
            for x in &mut buf {
                *x = rng.sample(&distr);
            }
            assert_eq!(buf, expected);
        }

        test_samples(100f64, 10.0, &[77.93369152808678, 90.1606912098641,
                125.31516221323625, 86.10217834773925]);
        test_samples(20f32, 10.0, &[27.175842, -2.0663052, 11.013268, 10.160688]);
    }
}
<commit_msg>rand_distr: work around tan func issue<commit_after>// Copyright 2018 Developers of the Rand project.
// Copyright 2016-2017 The Rust Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The Cauchy distribution.

use rand::Rng;
use crate::{Distribution, Standard};
use crate::utils::Float;

/// The Cauchy distribution `Cauchy(median, scale)`.
///
/// This distribution has a density function:
/// `f(x) = 1 / (pi * scale * (1 + ((x - median) / scale)^2))`
///
/// # Example
///
/// ```
/// use rand_distr::{Cauchy, Distribution};
///
/// let cau = Cauchy::new(2.0, 5.0).unwrap();
/// let v = cau.sample(&mut rand::thread_rng());
/// println!("{} is from a Cauchy(2, 5) distribution", v);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct Cauchy<N> {
    median: N,
    scale: N,
}

/// Error type returned from `Cauchy::new`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    /// `scale <= 0` or `nan`.
    ScaleTooSmall,
}

impl<N: Float> Cauchy<N>
where Standard: Distribution<N>
{
    /// Construct a new `Cauchy` with the given shape parameters
    /// `median` the peak location and `scale` the scale factor.
    pub fn new(median: N, scale: N) -> Result<Cauchy<N>, Error> {
        if !(scale > N::from(0.0)) {
            return Err(Error::ScaleTooSmall);
        }
        Ok(Cauchy {
            median,
            scale
        })
    }
}

impl<N: Float> Distribution<N> for Cauchy<N>
where Standard: Distribution<N>
{
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> N {
        // sample from [0, 1)
        let x = Standard.sample(rng);
        // get standard cauchy random number
        // note that π/2 is not exactly representable, even if x=0.5 the result is finite
        let comp_dev = (N::pi() * x).tan();
        // shift and scale according to parameters
        self.median + self.scale * comp_dev
    }
}

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

    fn median(mut numbers: &mut [f64]) -> f64 {
        sort(&mut numbers);
        let mid = numbers.len() / 2;
        numbers[mid]
    }

    fn sort(numbers: &mut [f64]) {
        numbers.sort_by(|a, b| a.partial_cmp(b).unwrap());
    }

    #[test]
    fn test_cauchy_averages() {
        // NOTE: given that the variance and mean are undefined,
        // this test does not have any rigorous statistical meaning.
        let cauchy = Cauchy::new(10.0, 5.0).unwrap();
        let mut rng = crate::test::rng(123);
        let mut numbers: [f64; 1000] = [0.0; 1000];
        let mut sum = 0.0;
        for i in 0..1000 {
            numbers[i] = cauchy.sample(&mut rng);
            sum += numbers[i];
        }
        let median = median(&mut numbers);
        println!("Cauchy median: {}", median);
        assert!((median - 10.0).abs() < 0.4); // not 100% certain, but probable enough
        let mean = sum / 1000.0;
        println!("Cauchy mean: {}", mean);
        // for a Cauchy distribution the mean should not converge
        assert!((mean - 10.0).abs() > 0.4); // not 100% certain, but probable enough
    }

    #[test]
    #[should_panic]
    fn test_cauchy_invalid_scale_zero() {
        Cauchy::new(0.0, 0.0).unwrap();
    }

    #[test]
    #[should_panic]
    fn test_cauchy_invalid_scale_neg() {
        Cauchy::new(0.0, -10.0).unwrap();
    }

    #[test]
    fn value_stability() {
        fn test_samples<N: Float + core::fmt::Debug>(m: N, s: N, expected: &[N])
        where Standard: Distribution<N> {
            let distr = Cauchy::new(m, s).unwrap();
            let mut rng = crate::test::rng(353);
            let mut buf = [m; 4];
            for x in &mut buf {
                *x = rng.sample(&distr);
            }
            assert_eq!(buf, expected);
        }

        // Warning: in a few cases, results vary slightly between different
        // platforms, presumably due to differences in precision of system impls
        // of the tan function. We work around this by avoiding these values.
        test_samples(100f64, 10.0, &[77.93369152808678, 90.1606912098641,
                125.31516221323625, 86.10217834773925]);
        test_samples(10f32, 7.0, &[15.023088, -5.446413, 3.7092876, 3.112482]);
    }
}
<|endoftext|>


<commit_before>#![allow(non_snake_case, non_upper_case_globals)]
// #![cfg_attr(test, feature(test))]
extern crate time;

use time::precise_time_s;
use std::mem;

///////////////////////////////////////////////////////////////////////////////
// 1. Calibration
///////////////////////////////////////////////////////////////////////////////
const aalpha:f64 = 0.33333333333;  // Elasticity of output w.r.t. capital
const bbeta:f64 = 0.95;           // Discount factor

// Productivity values
const vProductivity:[f64; 5] = [0.9792, 0.9896, 1.0000, 1.0106, 1.0212];

// Transition matrix
const mTransition:[[f64; 5]; 5] = [[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],
                                   [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],
                                   [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],
                                   [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],
                                   [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]];

// Dimensions to generate the grid of capital
const nGridCapital: usize = 17820;
const nGridProductivity: usize = 5;

fn solve(print: bool) -> f64 {
  /////////////////////////////////////////////////////////////////////////////
  // 2. Steady State
  /////////////////////////////////////////////////////////////////////////////

    let capitalSteadyState: f64 = (aalpha * bbeta).powf(1_f64 / (1_f64 - aalpha));
    let outputSteadyState: f64 = capitalSteadyState.powf(aalpha);
    let consumptionSteadyState: f64 = outputSteadyState - capitalSteadyState;

    if print {
        println!("Output = {}, Capital = {}, Consumption = {}",
                 outputSteadyState,
                 capitalSteadyState,
                 consumptionSteadyState);
    }

    let mut vGridCapital = [0f64; nGridCapital];

    for (i, val) in vGridCapital.iter_mut().enumerate() {
        *val = 0.5 * capitalSteadyState + 0.00001 * (i as f64)
    }

    // 3. Required matrices and vectors
    let mut mValueFunction = vec![[0f64; nGridProductivity]; nGridCapital];
    let mut mValueFunctionNew = mValueFunction.clone();
    let mut mPolicyFunction = mValueFunction.clone();
    let mut expectedValueFunction = mValueFunction.clone();

    // 4. We pre-build output for each point in the grid
    let mOutput = (0..nGridCapital)
                      .map(|nCapital| {
                          (0..nGridProductivity).map(|nProductivity|
        vProductivity[nProductivity]*vGridCapital[nCapital].powf(aalpha)
      ).collect::<Vec<f64>>()
                      })
                      .collect::<Vec<Vec<f64>>>();

  // 5. Main iteration
  // TODO: one could implement a macro for the multiple declarations
    let mut maxDifference = 10_f64;
    const tolerance: f64 = 0.0000001; // compiler warn: variable does not need to be mutable
    let mut capitalChoice: f64;

    let mut iteration = 0;

    while maxDifference > tolerance {
        for nProductivity in 0..nGridProductivity {
            let transitions = &mTransition[nProductivity];

            for nCapital in 0..nGridCapital {
                let value_fns = &mValueFunction[nCapital];

                let expected_value = transitions.iter()
                                                .zip(value_fns.iter())
                                                .fold(0.0f64, |acc, (transition, value_fn)| {
                    acc + (transition * value_fn)
                });

                expectedValueFunction[nCapital][nProductivity] = expected_value;
            }
        }

        for nProductivity in 0..nGridProductivity {

        // We start from previous choice (monotonicity of policy function)
            let mut gridCapitalNextPeriod = 0;

            for nCapital in 0..nGridCapital {

                let mut valueHighSoFar = -100000.0;

                for nCapitalNextPeriod in gridCapitalNextPeriod..nGridCapital {

                    let consumption = mOutput[nCapital][nProductivity] -
                                  vGridCapital[nCapitalNextPeriod];
                    let valueProvisional = (1_f64 - bbeta) * (consumption.ln()) +
                                       bbeta *
                                       expectedValueFunction[nCapitalNextPeriod][nProductivity];

                    if valueProvisional > valueHighSoFar {
                        valueHighSoFar = valueProvisional;
                        capitalChoice = vGridCapital[nCapitalNextPeriod];
                        gridCapitalNextPeriod = nCapitalNextPeriod;
                    } else {
                        break; // We break when we have achieved the max
                    }

                    mValueFunctionNew[nCapital][nProductivity] = valueHighSoFar;
                    mPolicyFunction[nCapital][nProductivity] = capitalChoice;
                }

            }

        }

        let mut diffHighSoFar = -100000.0;
        for nProductivity in 0..nGridProductivity {
            for nCapital in 0..nGridCapital {
                let diff = (mValueFunction[nCapital][nProductivity] -
                        mValueFunctionNew[nCapital][nProductivity])
                           .abs();
                if diff > diffHighSoFar {
                    diffHighSoFar = diff;
                }
            }
        }

        // swap buffers after the loop
        mem::swap(&mut mValueFunction, &mut mValueFunctionNew);

        maxDifference = diffHighSoFar;

        iteration += 1;
        if print && (iteration % 10 == 0 || iteration == 1) {
            println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
        }
    }

    if print {
        println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
    }

    mPolicyFunction[999][2]
}

fn main() {
    let sample = false;

    if sample {
        let mut samples = (0..5).map(|i| {
            let cpu0 = precise_time_s();
            let result = solve(false);
            let cpu1 = precise_time_s();

            assert_eq!(result, 0.1465491436956954);

            let diff = cpu1 - cpu0;
            println!("Sample #{}, Time: {}s", i + 1, diff);
            diff
        }).collect::<Vec<f64>>();

        samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
        println!("\nMedian time is = {}", samples[2]);
    } else {
        let cpu0 = precise_time_s();
        let result = solve(true);
        let cpu1 = precise_time_s();

        println!("My check = {}\n", result);
        assert_eq!(result, 0.1465491436956954);
        println!("Elapsed time is = {}", cpu1 - cpu0);
    }
}

// FIXME: Can't use bench as the loop is too slow (takes 8mins to get samples)
// #[cfg(test)]
// mod test {
//     extern crate test;
//     use self::test::Bencher;
//     use super::solve;

//     #[bench]
//     fn bench(b: &mut Bencher) {
//         b.iter(|| {
//             // use `test::black_box` to prevent compiler optimizations from disregarding
//             // unused values
//             test::black_box(solve(false));
//         });
//     }
// }
<commit_msg>Rewrite maxDifference loop<commit_after>#![allow(non_snake_case, non_upper_case_globals)]
// #![cfg_attr(test, feature(test))]
extern crate time;

use time::precise_time_s;
use std::mem;

///////////////////////////////////////////////////////////////////////////////
// 1. Calibration
///////////////////////////////////////////////////////////////////////////////
const aalpha:f64 = 0.33333333333;  // Elasticity of output w.r.t. capital
const bbeta:f64 = 0.95;           // Discount factor

// Productivity values
const vProductivity:[f64; 5] = [0.9792, 0.9896, 1.0000, 1.0106, 1.0212];

// Transition matrix
const mTransition:[[f64; 5]; 5] = [[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],
                                   [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],
                                   [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],
                                   [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],
                                   [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]];

// Dimensions to generate the grid of capital
const nGridCapital: usize = 17820;
const nGridProductivity: usize = 5;

fn solve(print: bool) -> f64 {
  /////////////////////////////////////////////////////////////////////////////
  // 2. Steady State
  /////////////////////////////////////////////////////////////////////////////

    let capitalSteadyState: f64 = (aalpha * bbeta).powf(1_f64 / (1_f64 - aalpha));
    let outputSteadyState: f64 = capitalSteadyState.powf(aalpha);
    let consumptionSteadyState: f64 = outputSteadyState - capitalSteadyState;

    if print {
        println!("Output = {}, Capital = {}, Consumption = {}",
                 outputSteadyState,
                 capitalSteadyState,
                 consumptionSteadyState);
    }

    let mut vGridCapital = [0f64; nGridCapital];

    for (i, val) in vGridCapital.iter_mut().enumerate() {
        *val = 0.5 * capitalSteadyState + 0.00001 * (i as f64)
    }

    // 3. Required matrices and vectors
    let mut mValueFunction = vec![[0f64; nGridProductivity]; nGridCapital];
    let mut mValueFunctionNew = mValueFunction.clone();
    let mut mPolicyFunction = mValueFunction.clone();
    let mut expectedValueFunction = mValueFunction.clone();

    // 4. We pre-build output for each point in the grid
    let mOutput = (0..nGridCapital)
                      .map(|nCapital| {
                          (0..nGridProductivity).map(|nProductivity|
        vProductivity[nProductivity]*vGridCapital[nCapital].powf(aalpha)
      ).collect::<Vec<f64>>()
                      })
                      .collect::<Vec<Vec<f64>>>();

  // 5. Main iteration
  // TODO: one could implement a macro for the multiple declarations
    let mut maxDifference = 10_f64;
    const tolerance: f64 = 0.0000001; // compiler warn: variable does not need to be mutable
    let mut capitalChoice: f64;

    let mut iteration = 0;

    while maxDifference > tolerance {
        for nProductivity in 0..nGridProductivity {
            let transitions = &mTransition[nProductivity];

            for nCapital in 0..nGridCapital {
                let value_fns = &mValueFunction[nCapital];

                let expected_value = transitions.iter()
                                                .zip(value_fns.iter())
                                                .fold(0.0f64, |acc, (transition, value_fn)| {
                    acc + (transition * value_fn)
                });

                expectedValueFunction[nCapital][nProductivity] = expected_value;
            }
        }

        for nProductivity in 0..nGridProductivity {

        // We start from previous choice (monotonicity of policy function)
            let mut gridCapitalNextPeriod = 0;

            for nCapital in 0..nGridCapital {

                let mut valueHighSoFar = -100000.0;

                for nCapitalNextPeriod in gridCapitalNextPeriod..nGridCapital {

                    let consumption = mOutput[nCapital][nProductivity] -
                                  vGridCapital[nCapitalNextPeriod];
                    let valueProvisional = (1_f64 - bbeta) * (consumption.ln()) +
                                       bbeta *
                                       expectedValueFunction[nCapitalNextPeriod][nProductivity];

                    if valueProvisional > valueHighSoFar {
                        valueHighSoFar = valueProvisional;
                        capitalChoice = vGridCapital[nCapitalNextPeriod];
                        gridCapitalNextPeriod = nCapitalNextPeriod;
                    } else {
                        break; // We break when we have achieved the max
                    }

                    mValueFunctionNew[nCapital][nProductivity] = valueHighSoFar;
                    mPolicyFunction[nCapital][nProductivity] = capitalChoice;
                }

            }

        }

        {
            let old_vals = mValueFunction.iter().flat_map(|arr| arr.iter());
            let new_vals = mValueFunctionNew.iter().flat_map(|arr| arr.iter());

            maxDifference = -100000.0;
            for diff in old_vals.zip(new_vals).map(|(old, new)| (old - new).abs()) {
                if diff > maxDifference {
                    maxDifference = diff
                }
            }
        }

        // swap buffers after the loop
        mem::swap(&mut mValueFunction, &mut mValueFunctionNew);

        iteration += 1;
        if print && (iteration % 10 == 0 || iteration == 1) {
            println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
        }
    }

    if print {
        println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
    }

    mPolicyFunction[999][2]
}

fn main() {
    let sample = false;

    if sample {
        let mut samples = (0..5).map(|i| {
            let cpu0 = precise_time_s();
            let result = solve(false);
            let cpu1 = precise_time_s();

            assert_eq!(result, 0.1465491436956954);

            let diff = cpu1 - cpu0;
            println!("Sample #{}, Time: {}s", i + 1, diff);
            diff
        }).collect::<Vec<f64>>();

        samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
        println!("\nMedian time is = {}", samples[2]);
    } else {
        let cpu0 = precise_time_s();
        let result = solve(true);
        let cpu1 = precise_time_s();

        println!("My check = {}\n", result);
        assert_eq!(result, 0.1465491436956954);
        println!("Elapsed time is = {}", cpu1 - cpu0);
    }
}

// FIXME: Can't use bench as the loop is too slow (takes 8mins to get samples)
// #[cfg(test)]
// mod test {
//     extern crate test;
//     use self::test::Bencher;
//     use super::solve;

//     #[bench]
//     fn bench(b: &mut Bencher) {
//         b.iter(|| {
//             // use `test::black_box` to prevent compiler optimizations from disregarding
//             // unused values
//             test::black_box(solve(false));
//         });
//     }
// }
<|endoftext|>