<commit_before>extern crate libc;
// export it for now

use libc::{c_int, c_char, c_void};
use std::ffi::CString;
use std::ptr;

pub mod ffi;
pub use ffi::{IPaddress, _TCPsocket, _SDLNet_SocketSet};

pub struct TCPsocket {
    opaquePtr: *mut _TCPsocket,
}
pub struct SocketSet {
    opaquePtr: *mut _SDLNet_SocketSet,
}

pub fn init() -> () { unsafe { ffi::SDLNet_Init(); } }
pub fn quit() -> () { unsafe { ffi::SDLNet_Quit(); } }
pub fn get_error() -> String {
    unsafe {
        let raw = &ffi::SDLNet_GetError();
        std::str::from_utf8(std::ffi::c_str_to_bytes(raw)).unwrap().to_string()
    }
}

pub fn become_host(port: u16) -> Option<ffi::IPaddress> {
    let mut address = ffi::IPaddress { host: 0, port: 0};
    let mut result = 0;
    unsafe {
        result = ffi::SDLNet_ResolveHost(&mut address, ptr::null() , port)
    }
    if (result == 0){
        Some(address)
    } else {
        None
    }
}

pub fn resolve_host(host: &str, port: u16) -> Option<ffi::IPaddress> {
    let mut address = ffi::IPaddress { host: 0, port: 0 };
    let mut result = 0;
    unsafe {
        result = ffi::SDLNet_ResolveHost(&mut address, CString::from_slice(host.as_bytes()).as_ptr(), port)
    }
    if (result == 0) {
        Some(address)
    } else {
        None
    }
}

pub fn resolve_ip(mut address: IPaddress) -> String {
    unsafe {
        let raw = &ffi::SDLNet_ResolveIP(&mut address);
        std::str::from_utf8(std::ffi::c_str_to_bytes(raw)).unwrap().to_string()
    }
}

pub fn tcp_open() -> Option<TCPsocket> {
    let mut address = ffi::IPaddress { host: 0, port: 0};
    unsafe {
        let socket = ffi::SDLNet_TCP_Open(&mut address);
        if socket as *const _TCPsocket != ptr::null() {
            Some(TCPsocket { opaquePtr: socket })
        } else {
            None
        }
    }
}

pub fn tcp_close(sock: TCPsocket) -> () {
    unsafe {
        ffi::SDLNet_TCP_Close(sock.opaquePtr)
    }
}

pub fn tcp_accept(server: TCPsocket) -> TCPsocket {
    unsafe {
        TCPsocket { opaquePtr: ffi::SDLNet_TCP_Accept(server.opaquePtr) }
    }
}

pub fn tcp_get_peer_address(sock: TCPsocket) -> Box<*mut IPaddress> {
    unsafe {
        Box::new(ffi::SDLNet_TCP_GetPeerAddress(sock.opaquePtr))
    }
}

pub fn tcp_send(sock: TCPsocket, data: &[u8]) -> () {
    unsafe {
        ffi::SDLNet_TCP_Send(sock.opaquePtr, data.as_ptr() as *const c_void, data.len() as i32);
    }
}

pub fn tcp_recv(sock: TCPsocket, maxlen: i32) -> Vec<u8> {
    let mut data: Vec<u8> = Vec::with_capacity(maxlen as usize);
    unsafe {
        let read_ammnt = ffi::SDLNet_TCP_Recv(sock.opaquePtr, data.as_mut_ptr() as *mut c_void, data.len() as i32);
        data.set_len(read_ammnt as usize);
    }
    data
}

pub fn alloc_socket_set(maxsockets: i32) -> SocketSet {
    unsafe {
        SocketSet { opaquePtr: ffi::SDLNet_AllocSocketSet(maxsockets) }
    }
}

pub fn free_socket_set(set: SocketSet) -> () {
    unsafe {
        ffi::SDLNet_FreeSocketSet(set.opaquePtr);
    }
}

pub fn add_socket<'a>(set: &'a SocketSet, sock: &'a TCPsocket) -> i32 {
    unsafe {
        ffi::SDLNet_AddSocket(set.opaquePtr, sock.opaquePtr)
    }
}

// Maybe should take in the box here as it may get deleted.... not sure
pub fn del_socket(set: SocketSet, sock: TCPsocket) -> i32 {
    unsafe {
        ffi::SDLNet_DelSocket(set.opaquePtr, sock.opaquePtr)
    }
}

pub fn check_sockets(set: SocketSet, timeout: u32) -> i32 {
    unsafe {
        ffi::SDLNet_CheckSockets(set.opaquePtr, timeout)
    }
}

pub fn socket_ready(sock: TCPsocket) -> i32 {
    unsafe {
        ffi::SDLNet_SocketReady(sock.opaquePtr)
    }
}
<commit_msg>Added more lifetimes.<commit_after>extern crate libc;
// export it for now

use libc::{c_int, c_char, c_void};
use std::ffi::CString;
use std::ptr;

pub mod ffi;
pub use ffi::{IPaddress, _TCPsocket, _SDLNet_SocketSet};

pub struct TCPsocket {
    opaquePtr: *mut _TCPsocket,
}
pub struct SocketSet {
    opaquePtr: *mut _SDLNet_SocketSet,
}

pub fn init() -> () { unsafe { ffi::SDLNet_Init(); } }
pub fn quit() -> () { unsafe { ffi::SDLNet_Quit(); } }
pub fn get_error() -> String {
    unsafe {
        let raw = &ffi::SDLNet_GetError();
        std::str::from_utf8(std::ffi::c_str_to_bytes(raw)).unwrap().to_string()
    }
}

pub fn become_host(port: u16) -> Option<ffi::IPaddress> {
    let mut address = ffi::IPaddress { host: 0, port: 0};
    let mut result = 0;
    unsafe {
        result = ffi::SDLNet_ResolveHost(&mut address, ptr::null() , port)
    }
    if (result == 0){
        Some(address)
    } else {
        None
    }
}

pub fn resolve_host(host: &str, port: u16) -> Option<ffi::IPaddress> {
    let mut address = ffi::IPaddress { host: 0, port: 0 };
    let mut result = 0;
    unsafe {
        result = ffi::SDLNet_ResolveHost(&mut address, CString::from_slice(host.as_bytes()).as_ptr(), port)
    }
    if (result == 0) {
        Some(address)
    } else {
        None
    }
}

pub fn resolve_ip(mut address: IPaddress) -> String {
    unsafe {
        let raw = &ffi::SDLNet_ResolveIP(&mut address);
        std::str::from_utf8(std::ffi::c_str_to_bytes(raw)).unwrap().to_string()
    }
}

pub fn tcp_open() -> Option<TCPsocket> {
    let mut address = ffi::IPaddress { host: 0, port: 0};
    unsafe {
        let socket = ffi::SDLNet_TCP_Open(&mut address);
        if socket as *const _TCPsocket != ptr::null() {
            Some(TCPsocket { opaquePtr: socket })
        } else {
            None
        }
    }
}

pub fn tcp_close(sock: TCPsocket) -> () {
    unsafe {
        ffi::SDLNet_TCP_Close(sock.opaquePtr)
    }
}

pub fn tcp_accept(server: TCPsocket) -> TCPsocket {
    unsafe {
        TCPsocket { opaquePtr: ffi::SDLNet_TCP_Accept(server.opaquePtr) }
    }
}

pub fn tcp_get_peer_address(sock: TCPsocket) -> Box<*mut IPaddress> {
    unsafe {
        Box::new(ffi::SDLNet_TCP_GetPeerAddress(sock.opaquePtr))
    }
}

pub fn tcp_send(sock: TCPsocket, data: &[u8]) -> () {
    unsafe {
        ffi::SDLNet_TCP_Send(sock.opaquePtr, data.as_ptr() as *const c_void, data.len() as i32);
    }
}

pub fn tcp_recv(sock: TCPsocket, maxlen: i32) -> Vec<u8> {
    let mut data: Vec<u8> = Vec::with_capacity(maxlen as usize);
    unsafe {
        let read_ammnt = ffi::SDLNet_TCP_Recv(sock.opaquePtr, data.as_mut_ptr() as *mut c_void, data.len() as i32);
        data.set_len(read_ammnt as usize);
    }
    data
}

pub fn alloc_socket_set(maxsockets: i32) -> SocketSet {
    unsafe {
        SocketSet { opaquePtr: ffi::SDLNet_AllocSocketSet(maxsockets) }
    }
}

pub fn free_socket_set(set: SocketSet) -> () {
    unsafe {
        ffi::SDLNet_FreeSocketSet(set.opaquePtr);
    }
}

pub fn add_socket<'a>(set: &'a SocketSet, sock: &'a TCPsocket) -> i32 {
    unsafe {
        ffi::SDLNet_AddSocket(set.opaquePtr, sock.opaquePtr)
    }
}

// Maybe should take in the box here as it may get deleted.... not sure
pub fn del_socket(set: SocketSet, sock: TCPsocket) -> i32 {
    unsafe {
        ffi::SDLNet_DelSocket(set.opaquePtr, sock.opaquePtr)
    }
}

pub fn check_sockets<'b>(set: &'b SocketSet, timeout: u32) -> i32 {
    unsafe {
        ffi::SDLNet_CheckSockets(set.opaquePtr, timeout)
    }
}

pub fn socket_ready<'c>(sock: &'c TCPsocket) -> i32 {
    unsafe {
        ffi::SDLNet_SocketReady(sock.opaquePtr)
    }
}
<|endoftext|>


<commit_before>//! Cicada is a bash-like Unix shell written in Rust.
//!
//! If you would like to use cicada as a regular shell,
//! please see details in [its repository](https://github.com/mitnk/cicada)
//!
//! Here is how to use cicada as a library:
//!
//! **Add cicada into Cargo.toml**
//!
//! ```
//! [dependencies]
//! cicada = "0.8.0"
//! ```
//!
//! **Use cicada functions**
//!
//! ```no-run
//! extern crate cicada;
//!
//! fn main() {
//!     let tokens = cicada::cmd_to_tokens("echo 'hi yoo' | `which wc`");
//!     assert_eq!(tokens.len(), 4);
//!
//!     assert_eq!(tokens[0].0, "");
//!     assert_eq!(tokens[0].1, "echo");
//!
//!     assert_eq!(tokens[1].0, "'");
//!     assert_eq!(tokens[1].1, "hi yoo");
//!
//!     assert_eq!(tokens[2].0, "");
//!     assert_eq!(tokens[2].1, "|");
//!
//!     assert_eq!(tokens[3].0, "`");
//!     assert_eq!(tokens[3].1, "which wc");
//!
//!     let out1 = cicada::run("ls Cargo.toml foo");
//!     assert_eq!(out1.status, 1);
//!     assert_eq!(out1.stdout, "Cargo.toml\n");
//!     assert_eq!(out1.stderr, "ls: foo: No such file or directory\n");
//!
//!     let out2 = cicada::run("ls | wc");
//!     assert_eq!(out2.status, 0);
//!     assert_eq!(out2.stdout, "       4       4      33\n");
//! }
//! ```
//!
#![allow(dead_code)]
#![allow(unknown_lints)]
// #![feature(tool_lints)]
extern crate errno;
extern crate exec;
extern crate glob;
extern crate libc;
extern crate linefeed;
extern crate nix;
extern crate regex;
extern crate sqlite;
extern crate time;

#[macro_use]
extern crate nom;

mod types;

#[macro_use]
mod tools;

mod builtins;
mod execute;
mod history;
mod jobc;
mod libs;
mod parsers;
mod shell;

/// Represents an error calling `exec`.
pub use types::CommandResult;

/// Parse a command to tokens.
pub fn cmd_to_tokens(cmd: &str) -> Vec<(String, String)> {
    return parsers::parser_line::cmd_to_tokens(cmd);
}

/// Run a command or a pipeline.
pub fn run(line: &str) -> CommandResult {
    execute::run(line)
}
<commit_msg>update library docs - fix<commit_after>//! Cicada is a bash-like Unix shell written in Rust.
//!
//! If you would like to use cicada as a regular shell,
//! please see details in [its repository](https://github.com/mitnk/cicada)
//!
//! Here is how to use cicada as a library:
//!
//! **Add cicada into Cargo.toml**
//!
//! ```no-run
//! [dependencies]
//! cicada = "0.8.0"
//! ```
//!
//! **Use cicada functions**
//!
//! ```no-run
//! extern crate cicada;
//!
//! fn main() {
//!     let tokens = cicada::cmd_to_tokens("echo 'hi yoo' | `which wc`");
//!     assert_eq!(tokens.len(), 4);
//!
//!     assert_eq!(tokens[0].0, "");
//!     assert_eq!(tokens[0].1, "echo");
//!
//!     assert_eq!(tokens[1].0, "'");
//!     assert_eq!(tokens[1].1, "hi yoo");
//!
//!     assert_eq!(tokens[2].0, "");
//!     assert_eq!(tokens[2].1, "|");
//!
//!     assert_eq!(tokens[3].0, "`");
//!     assert_eq!(tokens[3].1, "which wc");
//!
//!     let out1 = cicada::run("ls Cargo.toml foo");
//!     assert_eq!(out1.status, 1);
//!     assert_eq!(out1.stdout, "Cargo.toml\n");
//!     assert_eq!(out1.stderr, "ls: foo: No such file or directory\n");
//!
//!     let out2 = cicada::run("ls | wc");
//!     assert_eq!(out2.status, 0);
//!     assert_eq!(out2.stdout, "       4       4      33\n");
//! }
//! ```
//!
#![allow(dead_code)]
#![allow(unknown_lints)]
// #![feature(tool_lints)]
extern crate errno;
extern crate exec;
extern crate glob;
extern crate libc;
extern crate linefeed;
extern crate nix;
extern crate regex;
extern crate sqlite;
extern crate time;

#[macro_use]
extern crate nom;

mod types;

#[macro_use]
mod tools;

mod builtins;
mod execute;
mod history;
mod jobc;
mod libs;
mod parsers;
mod shell;

/// Represents an error calling `exec`.
pub use types::CommandResult;

/// Parse a command to tokens.
pub fn cmd_to_tokens(cmd: &str) -> Vec<(String, String)> {
    return parsers::parser_line::cmd_to_tokens(cmd);
}

/// Run a command or a pipeline.
pub fn run(line: &str) -> CommandResult {
    execute::run(line)
}
<|endoftext|>


<commit_before>mod raft_log;
pub mod storage;
mod raft;
mod progress;
mod errors;
mod log_unstable;
mod status;
pub mod raw_node;

pub use self::storage::{RaftState, Storage};
pub use self::errors::{Result, Error, StorageError};
pub use self::raft::{Raft, StateRole, Config, INVALID_ID, SoftState};
pub use self::raft_log::{RaftLog, NO_LIMIT};
pub use self::raw_node::{Ready, RawNode, Peer, is_empty_snap};
pub use self::status::Status;
pub use self::log_unstable::Unstable;
pub use self::progress::{Inflights, Progress, ProgressState};
<commit_msg>*: merge master and fix conflict.<commit_after>mod raft_log;
pub mod storage;
mod raft;
mod progress;
mod errors;
mod log_unstable;
mod status;
pub mod raw_node;

pub use self::storage::{RaftState, Storage};
pub use self::errors::{Result, Error, StorageError};
pub use self::raft::{Raft, StateRole, Config, INVALID_ID, SoftState};
pub use self::raft_log::{RaftLog, NO_LIMIT};
pub use self::raw_node::{Ready, RawNode, Peer, is_empty_snap, SnapshotStatus};
pub use self::status::Status;
pub use self::log_unstable::Unstable;
pub use self::progress::{Inflights, Progress, ProgressState};
<|endoftext|>


<commit_before>use std::time::{Duration, Instant};

use control::cio;
use torrent::Peer;
use util::{random_sample, FHashSet, UHashMap};

pub struct Choker {
    unchoked: Vec<usize>,
    interested: FHashSet<usize>,
    last_updated: Instant,
}

#[derive(Debug, PartialEq)]
pub struct SwapRes {
    pub choked: usize,
    pub unchoked: usize,
}

impl Choker {
    pub fn new() -> Choker {
        Choker {
            unchoked: Vec::with_capacity(5),
            interested: FHashSet::default(),
            last_updated: Instant::now(),
        }
    }

    pub fn add_peer<T: cio::CIO>(&mut self, peer: &mut Peer<T>) {
        if self.unchoked.len() < 5 {
            self.unchoked.push(peer.id());
            peer.flush();
            peer.unchoke();
        } else {
            self.interested.insert(peer.id());
        }
    }

    fn unchoke_random<T: cio::CIO>(&mut self, peers: &mut UHashMap<Peer<T>>) -> Option<usize> {
        if let Some(random_id) = random_sample(self.interested.iter()).cloned() {
            let mut peer = peers.get_mut(&random_id).unwrap();
            self.interested.remove(&random_id);
            self.add_peer(&mut peer);
            Some(random_id)
        } else {
            None
        }
    }

    pub fn remove_peer<T: cio::CIO>(
        &mut self,
        peer: &mut Peer<T>,
        peers: &mut UHashMap<Peer<T>>,
    ) -> Option<SwapRes> {
        if let Some(idx) = self.unchoked.iter().position(|&id| id == peer.id()) {
            self.unchoked.remove(idx);
            peer.choke();
            self.unchoke_random(peers).map(|unchoked| SwapRes {
                choked: peer.id(),
                unchoked,
            })
        } else {
            self.interested.remove(&peer.id());
            None
        }
    }

    fn update_timer(&mut self) -> Result<(), ()> {
        if self.last_updated.elapsed() < Duration::from_secs(10)
            || self.unchoked.len() < 5
            || self.interested.is_empty()
        {
            Err(())
        } else {
            self.last_updated = Instant::now();
            Ok(())
        }
    }

    pub fn update_upload<T: cio::CIO>(&mut self, peers: &mut UHashMap<Peer<T>>) -> Option<SwapRes> {
        if self.update_timer().is_err() {
            return None;
        }
        if self.interested.is_empty() {
            return None;
        }
        let (slowest, _) = self.unchoked.iter().enumerate().fold(
            (0, ::std::u32::MAX),
            |(slowest, min), (idx, id)| {
                let (ul, _) = peers.get_mut(id).unwrap().flush();
                if ul < min {
                    (idx, ul)
                } else {
                    (slowest, min)
                }
            },
        );
        Some(self.swap_peer(slowest, peers))
    }

    pub fn update_download<T: cio::CIO>(
        &mut self,
        peers: &mut UHashMap<Peer<T>>,
    ) -> Option<SwapRes> {
        if self.update_timer().is_err() {
            return None;
        }

        let (slowest, _) = self.unchoked.iter().enumerate().fold(
            (0, ::std::u32::MAX),
            |(slowest, min), (idx, id)| {
                match peers.get_mut(id).map(|peer| peer.flush()) {
                    Some((_, dl)) if dl < min =>
                        (idx, dl),
                    _ =>
                        (slowest, min),
                }
            },
        );
        Some(self.swap_peer(slowest, peers))
    }

    fn swap_peer<T: cio::CIO>(&mut self, idx: usize, peers: &mut UHashMap<Peer<T>>) -> SwapRes {
        let id = self.unchoked.remove(idx);
        {
            let peer = peers.get_mut(&id).unwrap();
            peer.choke();
        }

        // Unchoke one random interested peer
        let r = SwapRes {
            choked: id,
            unchoked: self.unchoke_random(peers).unwrap(),
        };
        self.interested.insert(id);
        r
    }
}

#[cfg(test)]
mod tests {
    use super::{Choker, SwapRes};
    use std::time::{Duration, Instant};
    use torrent::{Bitfield, Peer};
    use util::UHashMap;

    #[test]
    fn test_add_peers() {
        let mut c = Choker::new();
        for i in 0..6 {
            let mut p = Peer::test(i, 0, 0, 0, Bitfield::new(1));
            // Since the socket is a dummy
            c.add_peer(&mut p);
        }
        assert_eq!(c.unchoked.len(), 5);
        assert_eq!(c.interested.len(), 1);
    }

    #[test]
    fn test_remove_peers() {
        let mut c = Choker::new();
        let mut v = Vec::new();
        let mut h = UHashMap::default();
        for i in 0..6 {
            let mut p = Peer::test_from_stats(i, 0, 0);
            c.add_peer(&mut p);
            v.push(p);
            // Semi copy
            let pc = Peer::test_from_stats(i, 0, 0);
            h.insert(i, pc);
        }
        assert_eq!(c.unchoked.contains(&v[0].id()), true);
        assert_eq!(
            c.remove_peer(&mut v[0], &mut h),
            Some(SwapRes {
                choked: v[0].id(),
                unchoked: 5,
            })
        );
        assert_eq!(c.unchoked.contains(&v[0].id()), false);
    }

    #[test]
    fn test_update_upload() {
        let mut c = Choker::new();
        let mut h = UHashMap::default();
        assert_eq!(c.update_upload(&mut h).is_none(), true);
        for i in 0..6 {
            let mut p = Peer::test_from_stats(i, i as u32, 6 - i as u32);
            c.add_peer(&mut p);
            h.insert(i, p);
        }
        assert_eq!(c.update_upload(&mut h).is_none(), true);
        c.last_updated = Instant::now() - Duration::from_secs(11);
        let res = c.update_upload(&mut h).unwrap();
        assert_eq!(res.choked, 0);
        assert_eq!(res.unchoked, 5);
    }

    #[test]
    fn test_update_download() {
        let mut c = Choker::new();
        let mut h = UHashMap::default();
        assert_eq!(c.update_download(&mut h).is_none(), true);
        for i in 0..6 {
            let mut p = Peer::test_from_stats(i, 6 - i as u32, i as u32);
            c.add_peer(&mut p);
            h.insert(i, p);
        }
        assert_eq!(c.update_download(&mut h).is_none(), true);
        c.last_updated = Instant::now() - Duration::from_secs(11);
        let res = c.update_download(&mut h).unwrap();
        assert_eq!(res.choked, 0);
        assert_eq!(res.unchoked, 5);
    }
}
<commit_msg>Prevent panics on missing peers in choker<commit_after>use std::time::{Duration, Instant};

use control::cio;
use torrent::Peer;
use util::{random_sample, FHashSet, UHashMap};

pub struct Choker {
    unchoked: Vec<usize>,
    interested: FHashSet<usize>,
    last_updated: Instant,
}

#[derive(Debug, PartialEq)]
pub struct SwapRes {
    pub choked: usize,
    pub unchoked: usize,
}

impl Choker {
    pub fn new() -> Choker {
        Choker {
            unchoked: Vec::with_capacity(5),
            interested: FHashSet::default(),
            last_updated: Instant::now(),
        }
    }

    pub fn add_peer<T: cio::CIO>(&mut self, peer: &mut Peer<T>) {
        if self.unchoked.len() < 5 {
            self.unchoked.push(peer.id());
            peer.flush();
            peer.unchoke();
        } else {
            self.interested.insert(peer.id());
        }
    }

    fn unchoke_random<T: cio::CIO>(&mut self, peers: &mut UHashMap<Peer<T>>) -> Option<usize> {
        if let Some(random_id) = random_sample(self.interested.iter()).cloned() {
            peers.get_mut(&random_id).map(|mut peer| {
                self.interested.remove(&random_id);
                self.add_peer(&mut peer);
                random_id
            })
        } else {
            None
        }
    }

    pub fn remove_peer<T: cio::CIO>(
        &mut self,
        peer: &mut Peer<T>,
        peers: &mut UHashMap<Peer<T>>,
    ) -> Option<SwapRes> {
        if let Some(idx) = self.unchoked.iter().position(|&id| id == peer.id()) {
            self.unchoked.remove(idx);
            peer.choke();
            self.unchoke_random(peers).map(|unchoked| SwapRes {
                choked: peer.id(),
                unchoked,
            })
        } else {
            self.interested.remove(&peer.id());
            None
        }
    }

    fn update_timer(&mut self) -> Result<(), ()> {
        if self.last_updated.elapsed() < Duration::from_secs(10)
            || self.unchoked.len() < 5
            || self.interested.is_empty()
        {
            Err(())
        } else {
            self.last_updated = Instant::now();
            Ok(())
        }
    }

    pub fn update_upload<T: cio::CIO>(&mut self, peers: &mut UHashMap<Peer<T>>) -> Option<SwapRes> {
        if self.update_timer().is_err() {
            return None;
        }
        if self.interested.is_empty() {
            return None;
        }
        let (slowest, _) = self.unchoked.iter().enumerate().fold(
            (0, ::std::u32::MAX),
            |(slowest, min), (idx, id)| {
                match peers.get_mut(id).map(Peer::flush) {
                    Some((ul, _)) if ul < min => (idx, ul),
                    _ => (slowest, min),
                }
            }
        );
        Some(self.swap_peer(slowest, peers))
    }

    pub fn update_download<T: cio::CIO>(
        &mut self,
        peers: &mut UHashMap<Peer<T>>,
    ) -> Option<SwapRes> {
        if self.update_timer().is_err() {
            return None;
        }

        let (slowest, _) = self.unchoked.iter().enumerate().fold(
            (0, ::std::u32::MAX),
            |(slowest, min), (idx, id)| {
                match peers.get_mut(id).map(Peer::flush) {
                    Some((_, dl)) if dl < min =>
                        (idx, dl),
                    _ =>
                        (slowest, min),
                }
            },
        );
        Some(self.swap_peer(slowest, peers))
    }

    fn swap_peer<T: cio::CIO>(&mut self, idx: usize, peers: &mut UHashMap<Peer<T>>) -> SwapRes {
        let id = self.unchoked.remove(idx);
        {
            peers.get_mut(&id).map(Peer::choke);
        }

        // Unchoke one random interested peer
        let r = SwapRes {
            choked: id,
            unchoked: self.unchoke_random(peers).unwrap(),
        };
        self.interested.insert(id);
        r
    }
}

#[cfg(test)]
mod tests {
    use super::{Choker, SwapRes};
    use std::time::{Duration, Instant};
    use torrent::{Bitfield, Peer};
    use util::UHashMap;

    #[test]
    fn test_add_peers() {
        let mut c = Choker::new();
        for i in 0..6 {
            let mut p = Peer::test(i, 0, 0, 0, Bitfield::new(1));
            // Since the socket is a dummy
            c.add_peer(&mut p);
        }
        assert_eq!(c.unchoked.len(), 5);
        assert_eq!(c.interested.len(), 1);
    }

    #[test]
    fn test_remove_peers() {
        let mut c = Choker::new();
        let mut v = Vec::new();
        let mut h = UHashMap::default();
        for i in 0..6 {
            let mut p = Peer::test_from_stats(i, 0, 0);
            c.add_peer(&mut p);
            v.push(p);
            // Semi copy
            let pc = Peer::test_from_stats(i, 0, 0);
            h.insert(i, pc);
        }
        assert_eq!(c.unchoked.contains(&v[0].id()), true);
        assert_eq!(
            c.remove_peer(&mut v[0], &mut h),
            Some(SwapRes {
                choked: v[0].id(),
                unchoked: 5,
            })
        );
        assert_eq!(c.unchoked.contains(&v[0].id()), false);
    }

    #[test]
    fn test_update_upload() {
        let mut c = Choker::new();
        let mut h = UHashMap::default();
        assert_eq!(c.update_upload(&mut h).is_none(), true);
        for i in 0..6 {
            let mut p = Peer::test_from_stats(i, i as u32, 6 - i as u32);
            c.add_peer(&mut p);
            h.insert(i, p);
        }
        assert_eq!(c.update_upload(&mut h).is_none(), true);
        c.last_updated = Instant::now() - Duration::from_secs(11);
        let res = c.update_upload(&mut h).unwrap();
        assert_eq!(res.choked, 0);
        assert_eq!(res.unchoked, 5);
    }

    #[test]
    fn test_update_download() {
        let mut c = Choker::new();
        let mut h = UHashMap::default();
        assert_eq!(c.update_download(&mut h).is_none(), true);
        for i in 0..6 {
            let mut p = Peer::test_from_stats(i, 6 - i as u32, i as u32);
            c.add_peer(&mut p);
            h.insert(i, p);
        }
        assert_eq!(c.update_download(&mut h).is_none(), true);
        c.last_updated = Instant::now() - Duration::from_secs(11);
        let res = c.update_download(&mut h).unwrap();
        assert_eq!(res.choked, 0);
        assert_eq!(res.unchoked, 5);
    }
}
<|endoftext|>


<commit_before>use libc::{c_void,c_char,c_int,c_long,size_t};
use native;
use std::c_str::CString;
use std::io;
use std::mem::transmute;
use std::ptr::{null, null_mut};
use std;

pub struct Config {
    pub port: u16,
    pub threads: uint
}

impl Config {
    pub fn default() -> Config {
        Config { port: 8888, threads: 50 }
    }
}

#[link(name = "civetweb", kind = "static")]
extern {
    fn mg_start(callbacks: *const MgCallbacks, user_data: *mut c_void,
                options: *const *mut c_char) -> *mut MgContext;
    fn mg_stop(context: *mut MgContext);
    fn mg_set_request_handler(context: *mut MgContext, uri: *const c_char,
                              handler: MgRequestHandler, data: *mut c_void);
    fn mg_read(connection: *mut MgConnection, buf: *mut c_void,
               len: size_t) -> c_int;
    fn mg_write(connection: *mut MgConnection, data: *const c_void,
                len: size_t) -> c_int;
    fn mg_get_header(connection: *mut MgConnection,
                     name: *const c_char) -> *const c_char;
    fn mg_get_request_info(connection: *mut MgConnection) -> *mut MgRequestInfo;
}

enum MgContext {}

pub struct Server<T>(*mut MgContext, Box<ServerCallback<T>>);

pub struct ServerCallback<T> {
    callback: fn(&mut Connection, &T) -> Result<(), ()>,
    param: T,
}

impl<T: Sync> ServerCallback<T> {
    pub fn new(callback: fn(&mut Connection, &T) -> Result<(), ()>,
               param: T) -> ServerCallback<T> {
        ServerCallback { callback: callback, param: param }
    }
}

impl<T: 'static + Sync> Server<T> {
    fn as_ptr(&self) -> *mut MgContext {
        let Server(context, _) = *self; context
    }

    pub fn start(options: Config,
                 callback: ServerCallback<T>) -> io::IoResult<Server<T>> {
        let Config { port, threads } = options;
        let options = vec!(
            "listening_ports".to_c_str(), port.to_string().to_c_str(),
            "num_threads".to_c_str(), threads.to_string().to_c_str(),
        );
        let mut ptrs: Vec<*const c_char> = options.iter().map(|a| {
            a.as_ptr()
        }).collect();
        ptrs.push(0 as *const c_char);

        let context = start(ptrs.as_ptr() as *const _);
        // TODO: fill in this error
        if context.is_null() { return Err(io::standard_error(io::OtherIoError)) }

        let uri = "**".to_c_str();
        let mut callback = box callback;
        unsafe {
            mg_set_request_handler(context, uri.as_ptr(),
                                   raw_handler::<T>,
                                   &mut *callback as *mut _ as *mut c_void);
        }
        Ok(Server(context, callback))
    }
}

#[unsafe_destructor]
impl<T: 'static + Sync> Drop for Server<T> {
    fn drop(&mut self) {
        unsafe { mg_stop(self.as_ptr()) }
    }
}

fn raw_handler<T: 'static>(conn: *mut MgConnection, param: *mut c_void) -> int {
    let callback: &ServerCallback<T> = unsafe { transmute(param) };
    let task = native::task::new((0, std::uint::MAX));
    let mut result = None;

    task.run(|| {
        let mut connection = Connection(conn);
        result = Some((callback.callback)(&mut connection, &callback.param));
    }).destroy();

    match result {
        None => 0,
        Some(Err(_)) => 0,
        Some(Ok(_)) => 1
    }
}

pub enum MgConnection {}

pub struct Connection(*mut MgConnection);

impl Connection {
    fn unwrap(&self) -> *mut MgConnection {
        match *self { Connection(conn) => conn }
    }
}

type MgRequestHandler = fn(*mut MgConnection, *mut c_void) -> int;

#[repr(C)]
struct MgHeader {
    name: *const c_char,
    value: *const c_char
}

pub struct Header<'a>(*mut MgHeader);

impl<'a> Header<'a> {
    fn as_ref(&self) -> &'a MgHeader {
        match *self { Header(header) => unsafe { &*header } }
    }

    pub fn name(&self) -> Option<&'a str> {
        to_slice(self.as_ref(), |header| header.name)
    }

    pub fn value(&self) -> Option<&'a str> {
        to_slice(self.as_ref(), |header| header.value)
    }
}

#[repr(C)]
struct MgRequestInfo {
    request_method: *const c_char,
    uri: *const c_char,
    http_version: *const c_char,
    query_string: *const c_char,
    remote_user: *const c_char,
    remote_ip: c_long,
    remote_port: c_int,
    is_ssl: c_int,

    user_data: *mut c_void,
    conn_data: *mut c_void,

    num_headers: c_int,
    headers: [MgHeader, ..64]
}

pub struct RequestInfo<'a>(*mut MgRequestInfo);

impl<'a> RequestInfo<'a> {
    pub fn as_ref<'a>(&'a self) -> &'a MgRequestInfo {
        match *self { RequestInfo(info) => unsafe { &*info } }
    }

    fn as_ptr(&self) -> *mut MgRequestInfo {
        match *self { RequestInfo(info) => info }
    }

    pub fn method<'a>(&'a self) -> Option<&'a str> {
        to_slice(self.as_ref(), |info| info.request_method)
    }

    pub fn url<'a>(&'a self) -> Option<&'a str> {
        to_slice(self.as_ref(), |info| info.uri)
    }

    pub fn http_version<'a>(&'a self) -> Option<&'a str> {
        to_slice(self.as_ref(), |info| info.http_version)
    }

    pub fn query_string<'a>(&'a self) -> Option<&'a str> {
        to_slice(self.as_ref(), |info| info.query_string)
    }

    pub fn remote_ip(&self) -> int {
        self.as_ref().remote_ip as int
    }

    pub fn is_ssl(&self) -> bool {
        self.as_ref().is_ssl != 0
    }
}

#[repr(C)]
struct MgCallbacks {
    begin_request: *const c_void,
    end_request: *const c_void,
    log_message: *const c_void,
    init_ssl: *const c_void,
    websocket_connect: *const c_void,
    websocket_ready: *const c_void,
    websocket_data: *const c_void,
    connection_close: *const c_void,
    open_file: *const c_void,
    init_lua: *const c_void,
    upload: *const c_void,
    http_error: *const c_void
}

impl MgCallbacks {
    fn new() -> MgCallbacks {
        MgCallbacks {
            begin_request: null(),
            end_request: null(),
            log_message: null(),
            init_ssl: null(),
            websocket_connect: null(),
            websocket_ready: null(),
            websocket_data: null(),
            connection_close: null(),
            open_file: null(),
            init_lua: null(),
            upload: null(),
            http_error: null()
        }
    }
}

fn to_slice<'a, T>(obj: &'a T,
                   callback: |&'a T| -> *const c_char) -> Option<&'a str> {
    let chars = callback(obj);

    if unsafe { chars.is_null() || *chars == 0 } {
        return None;
    }

    let c_string = unsafe { CString::new(chars, false) };
    let len = c_string.len();

    unsafe { Some(transmute(std::raw::Slice { data: chars, len: len })) }
}

pub fn start(options: *const *mut c_char) -> *mut MgContext {
    unsafe { mg_start(&MgCallbacks::new(), null_mut(), options) }
}

pub fn read(conn: &Connection, buf: &mut [u8]) -> i32 {
    unsafe { mg_read(conn.unwrap(), buf.as_mut_ptr() as *mut c_void, buf.len() as size_t) }
}

pub fn write(conn: &Connection, bytes: &[u8]) -> i32 {
    let c_bytes = bytes.as_ptr() as *const c_void;
    unsafe { mg_write(conn.unwrap(), c_bytes, bytes.len() as size_t) }
}

pub fn get_header<'a>(conn: &'a Connection, string: &str) -> Option<&'a str> {
    let string = string.to_c_str();

    unsafe {
        to_slice(conn, |conn| mg_get_header(conn.unwrap(), string.as_ptr()))
    }
}

pub fn get_request_info<'a>(conn: &'a Connection) -> Option<RequestInfo<'a>> {
    unsafe {
        let info = mg_get_request_info(conn.unwrap());
        if info.is_null() {
            None
        } else {
            Some(RequestInfo(info))
        }
    }
}

pub fn get_headers<'a>(conn: &'a Connection) -> Vec<Header<'a>> {
    match get_request_info(conn) {
        Some(info) => unsafe {
            (*info.as_ptr()).headers.iter_mut().map(|h| Header(h)).collect()
        },
        None => vec!()
    }
}
<commit_msg>Update to rust master<commit_after>use libc::{c_void,c_char,c_int,c_long,size_t};
use native;
use std::c_str::CString;
use std::io;
use std::mem::transmute;
use std::ptr::{null, null_mut};
use std;

pub struct Config {
    pub port: u16,
    pub threads: uint
}

impl Config {
    pub fn default() -> Config {
        Config { port: 8888, threads: 50 }
    }
}

#[link(name = "civetweb", kind = "static")]
extern {
    fn mg_start(callbacks: *const MgCallbacks, user_data: *mut c_void,
                options: *const *mut c_char) -> *mut MgContext;
    fn mg_stop(context: *mut MgContext);
    fn mg_set_request_handler(context: *mut MgContext, uri: *const c_char,
                              handler: MgRequestHandler, data: *mut c_void);
    fn mg_read(connection: *mut MgConnection, buf: *mut c_void,
               len: size_t) -> c_int;
    fn mg_write(connection: *mut MgConnection, data: *const c_void,
                len: size_t) -> c_int;
    fn mg_get_header(connection: *mut MgConnection,
                     name: *const c_char) -> *const c_char;
    fn mg_get_request_info(connection: *mut MgConnection) -> *mut MgRequestInfo;
}

enum MgContext {}

pub struct Server<T>(*mut MgContext, Box<ServerCallback<T>>);

pub struct ServerCallback<T> {
    callback: fn(&mut Connection, &T) -> Result<(), ()>,
    param: T,
}

impl<T: Sync> ServerCallback<T> {
    pub fn new(callback: fn(&mut Connection, &T) -> Result<(), ()>,
               param: T) -> ServerCallback<T> {
        ServerCallback { callback: callback, param: param }
    }
}

impl<T: 'static + Sync> Server<T> {
    fn as_ptr(&self) -> *mut MgContext {
        let Server(context, _) = *self; context
    }

    pub fn start(options: Config,
                 callback: ServerCallback<T>) -> io::IoResult<Server<T>> {
        let Config { port, threads } = options;
        let options = vec!(
            "listening_ports".to_c_str(), port.to_string().to_c_str(),
            "num_threads".to_c_str(), threads.to_string().to_c_str(),
        );
        let mut ptrs: Vec<*const c_char> = options.iter().map(|a| {
            a.as_ptr()
        }).collect();
        ptrs.push(0 as *const c_char);

        let context = start(ptrs.as_ptr() as *const _);
        // TODO: fill in this error
        if context.is_null() { return Err(io::standard_error(io::OtherIoError)) }

        let uri = "**".to_c_str();
        let mut callback = box callback;
        unsafe {
            mg_set_request_handler(context, uri.as_ptr(),
                                   raw_handler::<T>,
                                   &mut *callback as *mut _ as *mut c_void);
        }
        Ok(Server(context, callback))
    }
}

#[unsafe_destructor]
impl<T: 'static + Sync> Drop for Server<T> {
    fn drop(&mut self) {
        unsafe { mg_stop(self.as_ptr()) }
    }
}

fn raw_handler<T: 'static>(conn: *mut MgConnection, param: *mut c_void) -> int {
    let callback: &ServerCallback<T> = unsafe { transmute(param) };
    let task = native::task::new((0, std::uint::MAX), 0);
    let mut result = None;

    task.run(|| {
        let mut connection = Connection(conn);
        result = Some((callback.callback)(&mut connection, &callback.param));
    }).destroy();

    match result {
        None => 0,
        Some(Err(_)) => 0,
        Some(Ok(_)) => 1
    }
}

pub enum MgConnection {}

pub struct Connection(*mut MgConnection);

impl Connection {
    fn unwrap(&self) -> *mut MgConnection {
        match *self { Connection(conn) => conn }
    }
}

type MgRequestHandler = fn(*mut MgConnection, *mut c_void) -> int;

#[repr(C)]
struct MgHeader {
    name: *const c_char,
    value: *const c_char
}

pub struct Header<'a>(*mut MgHeader);

impl<'a> Header<'a> {
    fn as_ref(&self) -> &'a MgHeader {
        match *self { Header(header) => unsafe { &*header } }
    }

    pub fn name(&self) -> Option<&'a str> {
        to_slice(self.as_ref(), |header| header.name)
    }

    pub fn value(&self) -> Option<&'a str> {
        to_slice(self.as_ref(), |header| header.value)
    }
}

#[repr(C)]
struct MgRequestInfo {
    request_method: *const c_char,
    uri: *const c_char,
    http_version: *const c_char,
    query_string: *const c_char,
    remote_user: *const c_char,
    remote_ip: c_long,
    remote_port: c_int,
    is_ssl: c_int,

    user_data: *mut c_void,
    conn_data: *mut c_void,

    num_headers: c_int,
    headers: [MgHeader, ..64]
}

pub struct RequestInfo<'a>(*mut MgRequestInfo);

impl<'a> RequestInfo<'a> {
    pub fn as_ref<'a>(&'a self) -> &'a MgRequestInfo {
        match *self { RequestInfo(info) => unsafe { &*info } }
    }

    fn as_ptr(&self) -> *mut MgRequestInfo {
        match *self { RequestInfo(info) => info }
    }

    pub fn method<'a>(&'a self) -> Option<&'a str> {
        to_slice(self.as_ref(), |info| info.request_method)
    }

    pub fn url<'a>(&'a self) -> Option<&'a str> {
        to_slice(self.as_ref(), |info| info.uri)
    }

    pub fn http_version<'a>(&'a self) -> Option<&'a str> {
        to_slice(self.as_ref(), |info| info.http_version)
    }

    pub fn query_string<'a>(&'a self) -> Option<&'a str> {
        to_slice(self.as_ref(), |info| info.query_string)
    }

    pub fn remote_ip(&self) -> int {
        self.as_ref().remote_ip as int
    }

    pub fn is_ssl(&self) -> bool {
        self.as_ref().is_ssl != 0
    }
}

#[repr(C)]
struct MgCallbacks {
    begin_request: *const c_void,
    end_request: *const c_void,
    log_message: *const c_void,
    init_ssl: *const c_void,
    websocket_connect: *const c_void,
    websocket_ready: *const c_void,
    websocket_data: *const c_void,
    connection_close: *const c_void,
    open_file: *const c_void,
    init_lua: *const c_void,
    upload: *const c_void,
    http_error: *const c_void
}

impl MgCallbacks {
    fn new() -> MgCallbacks {
        MgCallbacks {
            begin_request: null(),
            end_request: null(),
            log_message: null(),
            init_ssl: null(),
            websocket_connect: null(),
            websocket_ready: null(),
            websocket_data: null(),
            connection_close: null(),
            open_file: null(),
            init_lua: null(),
            upload: null(),
            http_error: null()
        }
    }
}

fn to_slice<'a, T>(obj: &'a T,
                   callback: |&'a T| -> *const c_char) -> Option<&'a str> {
    let chars = callback(obj);

    if unsafe { chars.is_null() || *chars == 0 } {
        return None;
    }

    let c_string = unsafe { CString::new(chars, false) };
    let len = c_string.len();

    unsafe { Some(transmute(std::raw::Slice { data: chars, len: len })) }
}

pub fn start(options: *const *mut c_char) -> *mut MgContext {
    unsafe { mg_start(&MgCallbacks::new(), null_mut(), options) }
}

pub fn read(conn: &Connection, buf: &mut [u8]) -> i32 {
    unsafe { mg_read(conn.unwrap(), buf.as_mut_ptr() as *mut c_void, buf.len() as size_t) }
}

pub fn write(conn: &Connection, bytes: &[u8]) -> i32 {
    let c_bytes = bytes.as_ptr() as *const c_void;
    unsafe { mg_write(conn.unwrap(), c_bytes, bytes.len() as size_t) }
}

pub fn get_header<'a>(conn: &'a Connection, string: &str) -> Option<&'a str> {
    let string = string.to_c_str();

    unsafe {
        to_slice(conn, |conn| mg_get_header(conn.unwrap(), string.as_ptr()))
    }
}

pub fn get_request_info<'a>(conn: &'a Connection) -> Option<RequestInfo<'a>> {
    unsafe {
        let info = mg_get_request_info(conn.unwrap());
        if info.is_null() {
            None
        } else {
            Some(RequestInfo(info))
        }
    }
}

pub fn get_headers<'a>(conn: &'a Connection) -> Vec<Header<'a>> {
    match get_request_info(conn) {
        Some(info) => unsafe {
            (*info.as_ptr()).headers.iter_mut().map(|h| Header(h)).collect()
        },
        None => vec!()
    }
}
<|endoftext|>


<commit_before>//! A built-in simulator, with referee support.
//!
//! ```no_run
//! use std::process::Command;
//! use roboime_next::prelude::*;
//! use roboime_next::{ai, sim};
//!
//! let mut sim = sim::Builder::new()
//!     .initial_formation(true)
//!     .build();
//!
//! let ai = ai::Builder::new(Command::new("./demo-ai"))
//!     .color(Yellow)
//!     .build().unwrap();
//!
//! let mut ai = ai.init(&sim).unwrap();
//!
//! loop {
//!     let cmd = ai.update(&sim).unwrap();
//!     sim.step(cmd, 0.016_666_667);
//!     // sleep maybe
//! }
//! ```
use std::collections::BTreeMap;
use std::collections::btree_map;
use ::prelude::*;
use ::game;

#[derive(Clone, Debug, PartialEq)]
pub struct Geometry {
}

const FIELD_LENGTH: f32    = 9.010;
const FIELD_WIDTH: f32     = 6.010;
const CENTER_DIAMETER: f32 = 1.000;
const DEFENSE_RADIUS: f32  = 1.000;
const DEFENSE_STRETCH: f32 = 0.500;
const GOAL_WIDTH: f32      = 1.000;
const ROBOT_RADIUS: f32    = 0.090;
const ROBOT_MOUTH: f32     = 0.500;
const BALL_RADIUS: f32     = 0.023;
const BALL_FRICT_LOSS: f32 = 1.000;

#[derive(Clone, Debug, PartialEq)]
pub struct Ball {
    //pub pos: Vec3d,
    //pub vel: Vec3d,
    pub pos: Vec2d,
    pub vel: Vec2d,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Robot {
    pub pos: Vec2d,
    pub vel: Vec2d,
    pub w: f32,
    pub vw: f32,
}

impl Robot {
    #[inline]
    pub fn new() -> Robot {
        Robot { pos: Vec2d(0.0, 0.0), w: 0.0, vel: Vec2d(0.0, 0.0), vw: 0.0 }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum Referee {
    Halt,
    Stop,
    Normal,
    PreKickoff(Color),
    Kickoff(Color),
    PrePenalty(Color),
    Penalty(Color),
    DirectFree(Color),
    IndirectFree(Color),
    Timeout(Color),
    Goal(Color),
}

#[derive(Clone, Debug, PartialEq)]
pub enum RefereeState {
    Idle,
    PlacingBall(f32),
}

#[derive(Clone, Debug, PartialEq)]
pub struct TeamInfo {
    name: String,
    score: u8,
    goalie: u8,
    //yellow_card_times: BTreeMap<u8, f32>,
    //yellow_cards: BTreeMap<u8, u8>,
    //red_cards: BTreeMap<u8, u8>,
}

impl Default for TeamInfo {
    fn default() -> TeamInfo {
        TeamInfo {
            name: "".to_string(),
            score: 0,
            goalie: 0,
            //yellow_card_times: BTreeMap::new(),
            //yellow_cards: BTreeMap::new(),
            //red_cards: BTreeMap::new(),
        }
    }
}

trait BTreeMapSimExt {
    fn initial_formation(&mut self, count: u8, color: Color, side: Side);
}

impl BTreeMapSimExt for BTreeMap<Id, Robot> {
    fn initial_formation(&mut self, count: u8, color: Color, side: Side) {
        use std::f32::consts::PI;

        let w_0 = if side.is_right() { 0.0 } else { PI };
        let w_delta = 2.0 * PI / (count as f32) * if side.is_right() { 1.0 } else { -1.0 };
        let x_offset = (CENTER_DIAMETER / 4.0 + FIELD_LENGTH / 4.0 - DEFENSE_RADIUS / 2.0) * if side.is_right() { 1.0 } else { -1.0 };
        //let radius = FIELD_LENGTH / 8.0;
        let radius = (FIELD_LENGTH / 2.0 - CENTER_DIAMETER / 2.0 - DEFENSE_RADIUS) / 3.0 - ROBOT_RADIUS;

        for i in 0..count {
            let robot_id = Id(color, i);
            let robot = self.entry(robot_id).or_insert(Robot::new());
            let w = w_0 + (i as f32) * w_delta;
            robot.pos.0 = radius * w.cos() + x_offset;
            robot.pos.1 = radius * w.sin();
            robot.w = w + PI;
        }
    }
}

#[derive(Debug, Clone)]
pub struct Builder {
    team_side: TeamSide,
    initial_formation: bool,
}

impl Builder {
    pub fn new() -> Builder {
        Builder {
            team_side: Default::default(),
            initial_formation: false,
        }
    }

    pub fn team_side(&mut self, team_side: TeamSide) -> &mut Builder {
        self.team_side = team_side;
        self
    }

    pub fn initial_formation(&mut self, initial_formation: bool) -> &mut Builder {
        self.initial_formation = initial_formation;
        self
    }

    pub fn build(&self) -> State {
        let mut robots = BTreeMap::new();
        if self.initial_formation {
            if self.team_side.yellow_is_left() {
                robots.initial_formation(6, Yellow, Left);
                robots.initial_formation(6, Blue, Right);
            } else {
                robots.initial_formation(6, Yellow, Right);
                robots.initial_formation(6, Blue, Left);
            }
        }
        State {
            initial_timestamp: 0.0, // TODO
            counter: 0,
            timestamp: 0.0,
            //ball: Ball { pos: Vec2d(0.0, 0.0), vel: Vec2d(0.0, 0.0) },
            ball: Ball { pos: Vec2d(0.0, 0.05), vel: Vec2d(-4.0, 0.0) },
            robots: robots,
            last_ball_touch: None,
            team_side: self.team_side,
            referee: Referee::Normal,
            referee_state: RefereeState::Idle,
            info_yellow: Default::default(),
            info_blue: Default::default(),
            geometry: Geometry {}
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct State {
    pub initial_timestamp: f64,
    pub counter: u64,
    pub timestamp: f32,
    pub ball: Ball,
    pub robots: BTreeMap<Id, Robot>,
    pub last_ball_touch: Option<Id>,
    pub team_side: TeamSide,
    pub referee: Referee,
    pub referee_state: RefereeState,
    pub info_yellow: TeamInfo,
    pub info_blue: TeamInfo,
    pub geometry: Geometry,
}

impl State {
    pub fn step(&mut self, command: game::Command, timestep: f32) {
        self.counter += 1;
        self.timestamp += timestep;
        let &mut State {
            ref mut ball,
            ref mut robots,
            ref mut last_ball_touch,
            ..
        } = self;
        let mut d_time_ball = timestep;
        let d_time = timestep;

        // XXX: overly simplified physics ahead
        for (robot_id, robot) in robots.iter_mut() {
            let robot_command = {
                if robot_id.color() == command.color {
                    command.robots.get(&robot_id.id())
                } else {
                    None
                }
            };

            if let Some(robot_command) = robot_command {
                let v_tangent = robot_command.v_tangent;
                let v_normal  = robot_command.v_normal;
                let v_angular = robot_command.v_angular;
                robot.vel = Vec2d(v_tangent, v_normal).rotate(-robot.w);
                robot.vw = v_angular;
            }

            // detect collision with ball
            let r = ROBOT_RADIUS + BALL_RADIUS;

            // Bhāskara:
            let a = (robot.vel - ball.vel).norm2();
            let b = 2.0 * ((robot.pos - ball.pos) * (robot.vel - ball.vel));
            let c = (robot.pos - ball.pos).norm2() - r * r;
            let delta = b * b - 4.0 * a * c;

            if delta >= 0.0 && a != 0.0 {
                let tc = (-b - delta.sqrt()) * 0.5 / a;
                if 0.0 <= tc && tc <= timestep {
                    use std::f32::consts::PI;

                    debug!("collision: ball and robot {:?}", robot_id);
                    *last_ball_touch = Some(*robot_id);

                    ball.pos += ball.vel * tc;
                    let robot_pos = robot.pos + robot.vel * tc;
                    let cw = (ball.pos - robot_pos).angle();
                    let vel = ball.vel - robot.vel;
                    let vw = vel.angle() - PI;
                    ball.vel = robot.vel + Vec2d(vel.norm(), 0.0).rotate(2.0 * cw - vw);

                    let rw = robot.w + robot.vw * tc;
                    if rw.abs_sub(cw) < ROBOT_MOUTH {
                        if let Some(robot_command) = robot_command {
                            use ::game::RobotAction::*;
                            if let Kick(force) = robot_command.action {
                                ball.vel += Vec2d(force, 0.0).rotate(rw);
                            }
                            if let ChipKick(force) = robot_command.action {
                                // TODO: z force
                                ball.vel += Vec2d(force, 0.0).rotate(rw);
                            }
                            // TODO: dribble effect
                        }
                    }

                    d_time_ball -= tc;
                }
            }

            robot.pos += robot.vel * d_time;
            robot.w   += robot.vw  * d_time;

            // TODO: effect of robot_command.action
        }

        // update ball pos
        ball.pos += ball.vel * d_time_ball ;
        ball.vel *= 1.0 - BALL_FRICT_LOSS * timestep;
    }
}

impl<'a> game::State<'a> for State {
    type Ball = &'a Ball;
    type Robot = (Id, &'a Robot);
    type Robots = Iter<'a>;
    type Geometry = &'a Geometry;

    fn counter(&self) -> u64 {
        self.counter
    }

    fn timestamp(&self) -> f64 {
        self.initial_timestamp + self.timestamp as f64
    }

    fn ball(&'a self) -> Self::Ball {
        &self.ball
    }

    fn robot(&'a self, id: Id) -> Option<Self::Robot> {
        self.robots.get(&id).map(|r| (id, r))
    }

    fn robots(&'a self) -> Self::Robots {
        Iter {
            filter: None,
            robots: self.robots.iter()
        }
    }

    fn robots_team(&'a self, color: Color) -> Self::Robots {
        Iter {
            filter: Some(color),
            robots: self.robots.iter()
        }
    }

    fn geometry(&'a self) -> Self::Geometry {
        &self.geometry
    }
}

pub struct Iter<'a> {
    filter: Option<Color>,
    robots: btree_map::Iter<'a, Id, Robot>,
}

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

    fn next(&mut self) -> Option<Self::Item> {
        (match self.filter {
            None => self.robots.next(),
            Some(color) => self.robots.find(|&(&id, _)| id.color() == color),
        }).map(|(id, robot)| (*id, robot))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = match self.filter {
            None => self.robots.len(),
            Some(color) => self.robots.clone().filter(|&(&id, _)| id.color() == color).count(),
        };
        (len, Some(len))
    }
}

impl<'a> ExactSizeIterator for Iter<'a> {}

impl<'a> game::Ball for &'a Ball {
    fn pos(&self) -> Vec2d { self.pos }
    fn vel(&self) -> Vec2d { self.vel }
}

impl<'a> game::Robot for (Id, &'a Robot) {
    fn id(&self) -> Id { self.0 }
    fn pos(&self) -> Vec2d { self.1.pos }
    fn vel(&self) -> Vec2d { self.1.vel }
    fn w(&self) -> f32 { self.1.w }
    fn vw(&self) -> f32 { self.1.vw }
}

impl<'a> game::Geometry for &'a Geometry {
    fn field_length(&self) -> f32 { FIELD_LENGTH }
    fn field_width(&self) -> f32 { FIELD_WIDTH }
    fn goal_width(&self) -> f32 { GOAL_WIDTH }
    fn center_circle_radius(&self) -> f32 { CENTER_DIAMETER / 2.0 }
    fn defense_radius(&self) -> f32 { DEFENSE_RADIUS }
    fn defense_stretch(&self) -> f32 { DEFENSE_STRETCH }
}
<commit_msg>Fix kick/chip activation detection.<commit_after>//! A built-in simulator, with referee support.
//!
//! ```no_run
//! use std::process::Command;
//! use roboime_next::prelude::*;
//! use roboime_next::{ai, sim};
//!
//! let mut sim = sim::Builder::new()
//!     .initial_formation(true)
//!     .build();
//!
//! let ai = ai::Builder::new(Command::new("./demo-ai"))
//!     .color(Yellow)
//!     .build().unwrap();
//!
//! let mut ai = ai.init(&sim).unwrap();
//!
//! loop {
//!     let cmd = ai.update(&sim).unwrap();
//!     sim.step(cmd, 0.016_666_667);
//!     // sleep maybe
//! }
//! ```
use std::collections::BTreeMap;
use std::collections::btree_map;
use ::prelude::*;
use ::game;

#[derive(Clone, Debug, PartialEq)]
pub struct Geometry {
}

const FIELD_LENGTH: f32    = 9.010;
const FIELD_WIDTH: f32     = 6.010;
const CENTER_DIAMETER: f32 = 1.000;
const DEFENSE_RADIUS: f32  = 1.000;
const DEFENSE_STRETCH: f32 = 0.500;
const GOAL_WIDTH: f32      = 1.000;
const ROBOT_RADIUS: f32    = 0.090;
const ROBOT_MOUTH: f32     = 0.500;
const BALL_RADIUS: f32     = 0.023;
const BALL_FRICT_LOSS: f32 = 1.000;

#[derive(Clone, Debug, PartialEq)]
pub struct Ball {
    //pub pos: Vec3d,
    //pub vel: Vec3d,
    pub pos: Vec2d,
    pub vel: Vec2d,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Robot {
    pub pos: Vec2d,
    pub vel: Vec2d,
    pub w: f32,
    pub vw: f32,
}

impl Robot {
    #[inline]
    pub fn new() -> Robot {
        Robot { pos: Vec2d(0.0, 0.0), w: 0.0, vel: Vec2d(0.0, 0.0), vw: 0.0 }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum Referee {
    Halt,
    Stop,
    Normal,
    PreKickoff(Color),
    Kickoff(Color),
    PrePenalty(Color),
    Penalty(Color),
    DirectFree(Color),
    IndirectFree(Color),
    Timeout(Color),
    Goal(Color),
}

#[derive(Clone, Debug, PartialEq)]
pub enum RefereeState {
    Idle,
    PlacingBall(f32),
}

#[derive(Clone, Debug, PartialEq)]
pub struct TeamInfo {
    name: String,
    score: u8,
    goalie: u8,
    //yellow_card_times: BTreeMap<u8, f32>,
    //yellow_cards: BTreeMap<u8, u8>,
    //red_cards: BTreeMap<u8, u8>,
}

impl Default for TeamInfo {
    fn default() -> TeamInfo {
        TeamInfo {
            name: "".to_string(),
            score: 0,
            goalie: 0,
            //yellow_card_times: BTreeMap::new(),
            //yellow_cards: BTreeMap::new(),
            //red_cards: BTreeMap::new(),
        }
    }
}

trait BTreeMapSimExt {
    fn initial_formation(&mut self, count: u8, color: Color, side: Side);
}

impl BTreeMapSimExt for BTreeMap<Id, Robot> {
    fn initial_formation(&mut self, count: u8, color: Color, side: Side) {
        use std::f32::consts::PI;

        let w_0 = if side.is_right() { 0.0 } else { PI };
        let w_delta = 2.0 * PI / (count as f32) * if side.is_right() { 1.0 } else { -1.0 };
        let x_offset = (CENTER_DIAMETER / 4.0 + FIELD_LENGTH / 4.0 - DEFENSE_RADIUS / 2.0) * if side.is_right() { 1.0 } else { -1.0 };
        //let radius = FIELD_LENGTH / 8.0;
        let radius = (FIELD_LENGTH / 2.0 - CENTER_DIAMETER / 2.0 - DEFENSE_RADIUS) / 3.0 - ROBOT_RADIUS;

        for i in 0..count {
            let robot_id = Id(color, i);
            let robot = self.entry(robot_id).or_insert(Robot::new());
            let w = w_0 + (i as f32) * w_delta;
            robot.pos.0 = radius * w.cos() + x_offset;
            robot.pos.1 = radius * w.sin();
            robot.w = w + PI;
        }
    }
}

#[derive(Debug, Clone)]
pub struct Builder {
    team_side: TeamSide,
    initial_formation: bool,
}

impl Builder {
    pub fn new() -> Builder {
        Builder {
            team_side: Default::default(),
            initial_formation: false,
        }
    }

    pub fn team_side(&mut self, team_side: TeamSide) -> &mut Builder {
        self.team_side = team_side;
        self
    }

    pub fn initial_formation(&mut self, initial_formation: bool) -> &mut Builder {
        self.initial_formation = initial_formation;
        self
    }

    pub fn build(&self) -> State {
        let mut robots = BTreeMap::new();
        if self.initial_formation {
            if self.team_side.yellow_is_left() {
                robots.initial_formation(6, Yellow, Left);
                robots.initial_formation(6, Blue, Right);
            } else {
                robots.initial_formation(6, Yellow, Right);
                robots.initial_formation(6, Blue, Left);
            }
        }
        State {
            initial_timestamp: 0.0, // TODO
            counter: 0,
            timestamp: 0.0,
            //ball: Ball { pos: Vec2d(0.0, 0.0), vel: Vec2d(0.0, 0.0) },
            ball: Ball { pos: Vec2d(0.0, 0.05), vel: Vec2d(-4.0, 0.0) },
            robots: robots,
            last_ball_touch: None,
            team_side: self.team_side,
            referee: Referee::Normal,
            referee_state: RefereeState::Idle,
            info_yellow: Default::default(),
            info_blue: Default::default(),
            geometry: Geometry {}
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct State {
    pub initial_timestamp: f64,
    pub counter: u64,
    pub timestamp: f32,
    pub ball: Ball,
    pub robots: BTreeMap<Id, Robot>,
    pub last_ball_touch: Option<Id>,
    pub team_side: TeamSide,
    pub referee: Referee,
    pub referee_state: RefereeState,
    pub info_yellow: TeamInfo,
    pub info_blue: TeamInfo,
    pub geometry: Geometry,
}

impl State {
    pub fn step(&mut self, command: game::Command, timestep: f32) {
        self.counter += 1;
        self.timestamp += timestep;
        let &mut State {
            ref mut ball,
            ref mut robots,
            ref mut last_ball_touch,
            ..
        } = self;
        let mut d_time_ball = timestep;
        let d_time = timestep;

        // XXX: overly simplified physics ahead
        for (robot_id, robot) in robots.iter_mut() {
            let robot_command = {
                if robot_id.color() == command.color {
                    command.robots.get(&robot_id.id())
                } else {
                    None
                }
            };

            if let Some(robot_command) = robot_command {
                let v_tangent = robot_command.v_tangent;
                let v_normal  = robot_command.v_normal;
                let v_angular = robot_command.v_angular;
                robot.vel = Vec2d(v_tangent, v_normal).rotate(-robot.w);
                robot.vw = v_angular;
            }

            // detect collision with ball
            let r = ROBOT_RADIUS + BALL_RADIUS;

            // Bhāskara:
            let a = (robot.vel - ball.vel).norm2();
            let b = 2.0 * ((robot.pos - ball.pos) * (robot.vel - ball.vel));
            let c = (robot.pos - ball.pos).norm2() - r * r;
            let delta = b * b - 4.0 * a * c;

            if delta >= 0.0 && a != 0.0 {
                let tc = (-b - delta.sqrt()) * 0.5 / a;
                if 0.0 <= tc && tc <= timestep {
                    use std::f32::consts::PI;

                    debug!("collision: ball and robot {:?}", robot_id);
                    *last_ball_touch = Some(*robot_id);

                    ball.pos += ball.vel * tc;
                    let robot_pos = robot.pos + robot.vel * tc;
                    let cw = (ball.pos - robot_pos).angle();
                    let vel = ball.vel - robot.vel;
                    let vw = vel.angle() - PI;
                    ball.vel = robot.vel + Vec2d(vel.norm(), 0.0).rotate(2.0 * cw - vw);

                    fn angle_abs_sub(a: f32, b: f32) -> f32 {
                        (((a % PI) - (b % PI)) % PI).abs()
                    }

                    let rw = robot.w + robot.vw * tc;
                    if angle_abs_sub(rw, cw) < ROBOT_MOUTH / 2.0 {
                        if let Some(robot_command) = robot_command {
                            use ::game::RobotAction::*;
                            if let Kick(force) = robot_command.action {
                                ball.vel += Vec2d(force, 0.0).rotate(rw);
                            }
                            if let ChipKick(force) = robot_command.action {
                                // TODO: z force
                                ball.vel += Vec2d(force, 0.0).rotate(rw);
                            }
                            // TODO: dribble effect
                        }
                    }

                    d_time_ball -= tc;
                }
            }

            robot.pos += robot.vel * d_time;
            robot.w   += robot.vw  * d_time;

            // TODO: effect of robot_command.action
        }

        // update ball pos
        ball.pos += ball.vel * d_time_ball ;
        ball.vel *= 1.0 - BALL_FRICT_LOSS * timestep;
    }
}

impl<'a> game::State<'a> for State {
    type Ball = &'a Ball;
    type Robot = (Id, &'a Robot);
    type Robots = Iter<'a>;
    type Geometry = &'a Geometry;

    fn counter(&self) -> u64 {
        self.counter
    }

    fn timestamp(&self) -> f64 {
        self.initial_timestamp + self.timestamp as f64
    }

    fn ball(&'a self) -> Self::Ball {
        &self.ball
    }

    fn robot(&'a self, id: Id) -> Option<Self::Robot> {
        self.robots.get(&id).map(|r| (id, r))
    }

    fn robots(&'a self) -> Self::Robots {
        Iter {
            filter: None,
            robots: self.robots.iter()
        }
    }

    fn robots_team(&'a self, color: Color) -> Self::Robots {
        Iter {
            filter: Some(color),
            robots: self.robots.iter()
        }
    }

    fn geometry(&'a self) -> Self::Geometry {
        &self.geometry
    }
}

pub struct Iter<'a> {
    filter: Option<Color>,
    robots: btree_map::Iter<'a, Id, Robot>,
}

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

    fn next(&mut self) -> Option<Self::Item> {
        (match self.filter {
            None => self.robots.next(),
            Some(color) => self.robots.find(|&(&id, _)| id.color() == color),
        }).map(|(id, robot)| (*id, robot))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = match self.filter {
            None => self.robots.len(),
            Some(color) => self.robots.clone().filter(|&(&id, _)| id.color() == color).count(),
        };
        (len, Some(len))
    }
}

impl<'a> ExactSizeIterator for Iter<'a> {}

impl<'a> game::Ball for &'a Ball {
    fn pos(&self) -> Vec2d { self.pos }
    fn vel(&self) -> Vec2d { self.vel }
}

impl<'a> game::Robot for (Id, &'a Robot) {
    fn id(&self) -> Id { self.0 }
    fn pos(&self) -> Vec2d { self.1.pos }
    fn vel(&self) -> Vec2d { self.1.vel }
    fn w(&self) -> f32 { self.1.w }
    fn vw(&self) -> f32 { self.1.vw }
}

impl<'a> game::Geometry for &'a Geometry {
    fn field_length(&self) -> f32 { FIELD_LENGTH }
    fn field_width(&self) -> f32 { FIELD_WIDTH }
    fn goal_width(&self) -> f32 { GOAL_WIDTH }
    fn center_circle_radius(&self) -> f32 { CENTER_DIAMETER / 2.0 }
    fn defense_radius(&self) -> f32 { DEFENSE_RADIUS }
    fn defense_stretch(&self) -> f32 { DEFENSE_STRETCH }
}
<|endoftext|>


<commit_before>use std::io::{self, Write, Read, Cursor};
use std::net::{TcpStream, Shutdown};
use std::time::Duration;
use std::borrow::BorrowMut;
use byteorder::{BigEndian, WriteBytesExt};
use bincode::rustc_serialize::{encode, decode};
use bincode::SizeLimit;
use enum_primitive::FromPrimitive;
use {Function, Result, ExceptionCode, Error, Coil, binary, Client};

const MODBUS_PROTOCOL_TCP: u16 = 0x0000;
const MODBUS_TCP_DEFAULT_PORT: u16 = 502;
const MODBUS_HEADER_SIZE: usize = 7;
const MODBUS_MAX_READ_COUNT: usize = 0x7d;
const MODBUS_MAX_WRITE_COUNT: usize = 0x79;

#[derive(RustcEncodable, RustcDecodable)]
#[repr(packed)]
struct Header {
    tid: u16,
    pid: u16,
    len: u16,
    uid: u8,
}

impl Header {
    fn new(transport: &mut Transport, len: u16) -> Header {
        Header {
            tid: transport.new_tid(),
            pid: MODBUS_PROTOCOL_TCP,
            len: len - MODBUS_HEADER_SIZE as u16,
            uid: transport.uid,
        }
    }
}

/// Context object which holds state for all modbus operations.
pub struct Transport {
    tid: u16,
    uid: u8,
    stream: TcpStream,
}

impl Transport {
    /// Create a new context context object and connect it to `addr` on modbus-tcp default
    /// port (502)
    pub fn new(addr: &str) -> io::Result<Transport> {
        Self::new_with_port(addr, MODBUS_TCP_DEFAULT_PORT)
    }

    /// Create a new context object and connect it to `addr` on port `port`
    pub fn new_with_port(addr: &str, port: u16) -> io::Result<Transport> {
        match TcpStream::connect((addr, port)) {
            Ok(s) => {
                // set some sane tcp socket options
                let t = Duration::from_secs(5);
                try!(s.set_read_timeout(Some(t)));
                try!(s.set_write_timeout(Some(t)));
                //                try!(s.set_nodelay(true));
                //                try!(s.set_keepalive(None));
                Ok(Transport {
                    tid: 0,
                    uid: 1,
                    stream: s,
                })
            }
            Err(e) => Err(e),
        }
    }

    // Create a new transaction Id, incrementing the previous one.
    // The Id is wrapping around if the Id reaches `u16::MAX`.
    fn new_tid(&mut self) -> u16 {
        self.tid = self.tid.wrapping_add(1);
        self.tid
    }

    fn read(self: &mut Self, fun: Function) -> Result<Vec<u8>> {
        let packed_size = |v: u16| {
            v / 8 +
            if v % 8 > 0 {
                1
            } else {
                0
            }
        };
        let (addr, count, expected_bytes) = match fun {
            Function::ReadCoils(a, c) => (a, c, packed_size(c) as usize),
            Function::ReadDiscreteInputs(a, c) => (a, c, packed_size(c) as usize),
            Function::ReadHoldingRegisters(a, c) => (a, c, 2 * c as usize),
            Function::ReadInputRegisters(a, c) => (a, c, 2 * c as usize),
            _ => panic!("Unexpected modbus function"),
        };

        if count < 1 || count as usize > MODBUS_MAX_READ_COUNT {
            return Err(Error::InvalidData);
        }

        let header = Header::new(self, MODBUS_HEADER_SIZE as u16 + 6u16);
        let mut buff = try!(encode(&header, SizeLimit::Infinite));
        try!(buff.write_u8(fun.code()));
        try!(buff.write_u16::<BigEndian>(addr));
        try!(buff.write_u16::<BigEndian>(count));

        match self.stream.write_all(&buff[..]) {
            Ok(_s) => {
                let mut reply = vec![0; MODBUS_HEADER_SIZE + expected_bytes + 2];
                match self.stream.read(&mut reply) {
                    Ok(_s) => {
                        let resp_hd = try!(decode(&reply[..MODBUS_HEADER_SIZE]));
                        try!(Transport::validate_response_header(&header, &resp_hd));
                        try!(Transport::validate_response_code(&buff, &reply[..]));
                        Transport::get_reply_data(&reply, expected_bytes)
                    }
                    Err(e) => Err(Error::Io(e)),
                }
            }
            Err(e) => Err(Error::Io(e)),
        }
    }

    fn validate_response_header(req: &Header, resp: &Header) -> Result<()> {
        if req.tid != resp.tid || resp.pid != MODBUS_PROTOCOL_TCP {
            Err(Error::InvalidResponse)
        } else {
            Ok(())
        }
    }

    fn validate_response_code(req: &[u8], resp: &[u8]) -> Result<()> {
        if req[7] + 0x80 == resp[7] {
            let code = ExceptionCode::from_u8(resp[8]).unwrap();
            Err(Error::Exception(code))
        } else if req[7] != resp[7] {
            Err(Error::InvalidResponse)
        } else {
            Ok(())
        }
    }

    fn get_reply_data(reply: &[u8], expected_bytes: usize) -> Result<Vec<u8>> {
        if reply[8] as usize != expected_bytes ||
           reply.len() != MODBUS_HEADER_SIZE + expected_bytes + 2 {
            Err(Error::InvalidData)
        } else {
            let mut d = Vec::new();
            d.extend(reply[MODBUS_HEADER_SIZE + 2..].iter());
            Ok(d)
        }
    }

    fn write_single(self: &mut Self, fun: Function) -> Result<()> {
        let (addr, value) = match fun {
            Function::WriteSingleCoil(a, v) => (a, v),
            Function::WriteSingleRegister(a, v) => (a, v),
            _ => panic!("Unexpected modbus function"),
        };

        let mut buff = vec![0; MODBUS_HEADER_SIZE];  // Header gets filled in later
        try!(buff.write_u8(fun.code()));
        try!(buff.write_u16::<BigEndian>(addr));
        try!(buff.write_u16::<BigEndian>(value));
        self.write(&mut buff)
    }

    fn write_multiple(self: &mut Self, fun: Function) -> Result<()> {
        let (addr, quantity, values) = match fun {
            Function::WriteMultipleCoils(a, q, v) => (a, q, v),
            Function::WriteMultipleRegisters(a, q, v) => (a, q, v),
            _ => panic!("Unexpected modbus function"),
        };

        let mut buff = vec![0; MODBUS_HEADER_SIZE];  // Header gets filled in later
        try!(buff.write_u8(fun.code()));
        try!(buff.write_u16::<BigEndian>(addr));
        try!(buff.write_u16::<BigEndian>(quantity));
        try!(buff.write_u8(values.len() as u8));
        for v in values {
            try!(buff.write_u8(*v));
        }
        self.write(&mut buff)
    }

    fn write(self: &mut Self, buff: &mut [u8]) -> Result<()> {
        if buff.len() < 1 || buff.len() > MODBUS_MAX_WRITE_COUNT {
            return Err(Error::InvalidData);
        }
        let header = Header::new(self, buff.len() as u16 + 1u16);
        let head_buff = try!(encode(&header, SizeLimit::Infinite));
        {
            let mut start = Cursor::new(buff.borrow_mut());
            try!(start.write(&head_buff[..]));
        }
        match self.stream.write_all(&buff[..]) {
            Ok(_s) => {
                let reply = &mut [0; 12];
                match self.stream.read(reply) {
                    Ok(_s) => {
                        let resp_hd = try!(decode(&reply[..MODBUS_HEADER_SIZE]));
                        try!(Transport::validate_response_header(&header, &resp_hd));
                        Transport::validate_response_code(&buff, reply)
                    }
                    Err(e) => Err(Error::Io(e)),
                }
            }
            Err(e) => Err(Error::Io(e)),
        }
    }

}

impl Drop for Transport {
    fn drop(&mut self) {
        self.stream.shutdown(Shutdown::Both).unwrap();
    }
}

impl Client for Transport {
    /// Read `count` bits starting at address `addr`.
    fn read_coils(self: &mut Self, addr: u16, count: u16) -> Result<Vec<Coil>> {
        let bytes = try!(self.read(Function::ReadCoils(addr, count)));
        Ok(binary::unpack_bits(&bytes, count))
    }

    /// Read `count` input bits starting at address `addr`.
    fn read_discrete_inputs(self: &mut Self, addr: u16, count: u16) -> Result<Vec<Coil>> {
        let bytes = try!(self.read(Function::ReadDiscreteInputs(addr, count)));
        Ok(binary::unpack_bits(&bytes, count))
    }

    /// Read `count` 16bit registers starting at address `addr`.
    fn read_holding_registers(self: &mut Self, addr: u16, count: u16) -> Result<Vec<u16>> {
        let bytes = try!(self.read(Function::ReadHoldingRegisters(addr, count)));
        binary::pack_bytes(&bytes[..])
    }

    /// Read `count` 16bit input registers starting at address `addr`.
    fn read_input_registers(self: &mut Self, addr: u16, count: u16) -> Result<Vec<u16>> {
        let bytes = try!(self.read(Function::ReadInputRegisters(addr, count)));
        binary::pack_bytes(&bytes[..])
    }

    /// Write a single coil (bit) to address `addr`.
    fn write_single_coil(self: &mut Self, addr: u16, value: Coil) -> Result<()> {
        self.write_single(Function::WriteSingleCoil(addr, value.code()))
    }

    /// Write a single 16bit register to address `addr`.
    fn write_single_register(self: &mut Self, addr: u16, value: u16) -> Result<()> {
        self.write_single(Function::WriteSingleRegister(addr, value))
    }

    /// Write a multiple coils (bits) starting at address `addr`.
    fn write_multiple_coils(self: &mut Self, addr: u16, values: &[Coil]) -> Result<()> {
        let bytes = binary::pack_bits(values);
        self.write_multiple(Function::WriteMultipleCoils(addr, values.len() as u16, &bytes[..]))
    }

    /// Write a multiple 16bit registers starting at address `addr`.
    fn write_multiple_registers(self: &mut Self, addr: u16, values: &[u16]) -> Result<()> {
        let bytes = binary::unpack_bytes(values);
        self.write_multiple(Function::WriteMultipleRegisters(addr, values.len() as u16, &bytes[..]))
    }
}
<commit_msg>do not use unwrap in 'validate_response_code' function<commit_after>use std::io::{self, Write, Read, Cursor};
use std::net::{TcpStream, Shutdown};
use std::time::Duration;
use std::borrow::BorrowMut;
use byteorder::{BigEndian, WriteBytesExt};
use bincode::rustc_serialize::{encode, decode};
use bincode::SizeLimit;
use enum_primitive::FromPrimitive;
use {Function, Result, ExceptionCode, Error, Coil, binary, Client};

const MODBUS_PROTOCOL_TCP: u16 = 0x0000;
const MODBUS_TCP_DEFAULT_PORT: u16 = 502;
const MODBUS_HEADER_SIZE: usize = 7;
const MODBUS_MAX_READ_COUNT: usize = 0x7d;
const MODBUS_MAX_WRITE_COUNT: usize = 0x79;

#[derive(RustcEncodable, RustcDecodable)]
#[repr(packed)]
struct Header {
    tid: u16,
    pid: u16,
    len: u16,
    uid: u8,
}

impl Header {
    fn new(transport: &mut Transport, len: u16) -> Header {
        Header {
            tid: transport.new_tid(),
            pid: MODBUS_PROTOCOL_TCP,
            len: len - MODBUS_HEADER_SIZE as u16,
            uid: transport.uid,
        }
    }
}

/// Context object which holds state for all modbus operations.
pub struct Transport {
    tid: u16,
    uid: u8,
    stream: TcpStream,
}

impl Transport {
    /// Create a new context context object and connect it to `addr` on modbus-tcp default
    /// port (502)
    pub fn new(addr: &str) -> io::Result<Transport> {
        Self::new_with_port(addr, MODBUS_TCP_DEFAULT_PORT)
    }

    /// Create a new context object and connect it to `addr` on port `port`
    pub fn new_with_port(addr: &str, port: u16) -> io::Result<Transport> {
        match TcpStream::connect((addr, port)) {
            Ok(s) => {
                // set some sane tcp socket options
                let t = Duration::from_secs(5);
                try!(s.set_read_timeout(Some(t)));
                try!(s.set_write_timeout(Some(t)));
                //                try!(s.set_nodelay(true));
                //                try!(s.set_keepalive(None));
                Ok(Transport {
                    tid: 0,
                    uid: 1,
                    stream: s,
                })
            }
            Err(e) => Err(e),
        }
    }

    // Create a new transaction Id, incrementing the previous one.
    // The Id is wrapping around if the Id reaches `u16::MAX`.
    fn new_tid(&mut self) -> u16 {
        self.tid = self.tid.wrapping_add(1);
        self.tid
    }

    fn read(self: &mut Self, fun: Function) -> Result<Vec<u8>> {
        let packed_size = |v: u16| {
            v / 8 +
            if v % 8 > 0 {
                1
            } else {
                0
            }
        };
        let (addr, count, expected_bytes) = match fun {
            Function::ReadCoils(a, c) => (a, c, packed_size(c) as usize),
            Function::ReadDiscreteInputs(a, c) => (a, c, packed_size(c) as usize),
            Function::ReadHoldingRegisters(a, c) => (a, c, 2 * c as usize),
            Function::ReadInputRegisters(a, c) => (a, c, 2 * c as usize),
            _ => panic!("Unexpected modbus function"),
        };

        if count < 1 || count as usize > MODBUS_MAX_READ_COUNT {
            return Err(Error::InvalidData);
        }

        let header = Header::new(self, MODBUS_HEADER_SIZE as u16 + 6u16);
        let mut buff = try!(encode(&header, SizeLimit::Infinite));
        try!(buff.write_u8(fun.code()));
        try!(buff.write_u16::<BigEndian>(addr));
        try!(buff.write_u16::<BigEndian>(count));

        match self.stream.write_all(&buff[..]) {
            Ok(_s) => {
                let mut reply = vec![0; MODBUS_HEADER_SIZE + expected_bytes + 2];
                match self.stream.read(&mut reply) {
                    Ok(_s) => {
                        let resp_hd = try!(decode(&reply[..MODBUS_HEADER_SIZE]));
                        try!(Transport::validate_response_header(&header, &resp_hd));
                        try!(Transport::validate_response_code(&buff, &reply[..]));
                        Transport::get_reply_data(&reply, expected_bytes)
                    }
                    Err(e) => Err(Error::Io(e)),
                }
            }
            Err(e) => Err(Error::Io(e)),
        }
    }

    fn validate_response_header(req: &Header, resp: &Header) -> Result<()> {
        if req.tid != resp.tid || resp.pid != MODBUS_PROTOCOL_TCP {
            Err(Error::InvalidResponse)
        } else {
            Ok(())
        }
    }

    fn validate_response_code(req: &[u8], resp: &[u8]) -> Result<()> {
        if req[7] + 0x80 == resp[7] {
            match ExceptionCode::from_u8(resp[8]) {
                Some(code) => Err(Error::Exception(code)),
                None => Err(Error::InvalidResponse)
            }
        } else if req[7] != resp[7] {
            Err(Error::InvalidResponse)
        } else {
            Ok(())
        }
    }

    fn get_reply_data(reply: &[u8], expected_bytes: usize) -> Result<Vec<u8>> {
        if reply[8] as usize != expected_bytes ||
           reply.len() != MODBUS_HEADER_SIZE + expected_bytes + 2 {
            Err(Error::InvalidData)
        } else {
            let mut d = Vec::new();
            d.extend(reply[MODBUS_HEADER_SIZE + 2..].iter());
            Ok(d)
        }
    }

    fn write_single(self: &mut Self, fun: Function) -> Result<()> {
        let (addr, value) = match fun {
            Function::WriteSingleCoil(a, v) => (a, v),
            Function::WriteSingleRegister(a, v) => (a, v),
            _ => panic!("Unexpected modbus function"),
        };

        let mut buff = vec![0; MODBUS_HEADER_SIZE];  // Header gets filled in later
        try!(buff.write_u8(fun.code()));
        try!(buff.write_u16::<BigEndian>(addr));
        try!(buff.write_u16::<BigEndian>(value));
        self.write(&mut buff)
    }

    fn write_multiple(self: &mut Self, fun: Function) -> Result<()> {
        let (addr, quantity, values) = match fun {
            Function::WriteMultipleCoils(a, q, v) => (a, q, v),
            Function::WriteMultipleRegisters(a, q, v) => (a, q, v),
            _ => panic!("Unexpected modbus function"),
        };

        let mut buff = vec![0; MODBUS_HEADER_SIZE];  // Header gets filled in later
        try!(buff.write_u8(fun.code()));
        try!(buff.write_u16::<BigEndian>(addr));
        try!(buff.write_u16::<BigEndian>(quantity));
        try!(buff.write_u8(values.len() as u8));
        for v in values {
            try!(buff.write_u8(*v));
        }
        self.write(&mut buff)
    }

    fn write(self: &mut Self, buff: &mut [u8]) -> Result<()> {
        if buff.len() < 1 || buff.len() > MODBUS_MAX_WRITE_COUNT {
            return Err(Error::InvalidData);
        }
        let header = Header::new(self, buff.len() as u16 + 1u16);
        let head_buff = try!(encode(&header, SizeLimit::Infinite));
        {
            let mut start = Cursor::new(buff.borrow_mut());
            try!(start.write(&head_buff[..]));
        }
        match self.stream.write_all(&buff[..]) {
            Ok(_s) => {
                let reply = &mut [0; 12];
                match self.stream.read(reply) {
                    Ok(_s) => {
                        let resp_hd = try!(decode(&reply[..MODBUS_HEADER_SIZE]));
                        try!(Transport::validate_response_header(&header, &resp_hd));
                        Transport::validate_response_code(&buff, reply)
                    }
                    Err(e) => Err(Error::Io(e)),
                }
            }
            Err(e) => Err(Error::Io(e)),
        }
    }

}

impl Drop for Transport {
    fn drop(&mut self) {
        self.stream.shutdown(Shutdown::Both).unwrap();
    }
}

impl Client for Transport {
    /// Read `count` bits starting at address `addr`.
    fn read_coils(self: &mut Self, addr: u16, count: u16) -> Result<Vec<Coil>> {
        let bytes = try!(self.read(Function::ReadCoils(addr, count)));
        Ok(binary::unpack_bits(&bytes, count))
    }

    /// Read `count` input bits starting at address `addr`.
    fn read_discrete_inputs(self: &mut Self, addr: u16, count: u16) -> Result<Vec<Coil>> {
        let bytes = try!(self.read(Function::ReadDiscreteInputs(addr, count)));
        Ok(binary::unpack_bits(&bytes, count))
    }

    /// Read `count` 16bit registers starting at address `addr`.
    fn read_holding_registers(self: &mut Self, addr: u16, count: u16) -> Result<Vec<u16>> {
        let bytes = try!(self.read(Function::ReadHoldingRegisters(addr, count)));
        binary::pack_bytes(&bytes[..])
    }

    /// Read `count` 16bit input registers starting at address `addr`.
    fn read_input_registers(self: &mut Self, addr: u16, count: u16) -> Result<Vec<u16>> {
        let bytes = try!(self.read(Function::ReadInputRegisters(addr, count)));
        binary::pack_bytes(&bytes[..])
    }

    /// Write a single coil (bit) to address `addr`.
    fn write_single_coil(self: &mut Self, addr: u16, value: Coil) -> Result<()> {
        self.write_single(Function::WriteSingleCoil(addr, value.code()))
    }

    /// Write a single 16bit register to address `addr`.
    fn write_single_register(self: &mut Self, addr: u16, value: u16) -> Result<()> {
        self.write_single(Function::WriteSingleRegister(addr, value))
    }

    /// Write a multiple coils (bits) starting at address `addr`.
    fn write_multiple_coils(self: &mut Self, addr: u16, values: &[Coil]) -> Result<()> {
        let bytes = binary::pack_bits(values);
        self.write_multiple(Function::WriteMultipleCoils(addr, values.len() as u16, &bytes[..]))
    }

    /// Write a multiple 16bit registers starting at address `addr`.
    fn write_multiple_registers(self: &mut Self, addr: u16, values: &[u16]) -> Result<()> {
        let bytes = binary::unpack_bytes(values);
        self.write_multiple(Function::WriteMultipleRegisters(addr, values.len() as u16, &bytes[..]))
    }
}
<|endoftext|>


<commit_before>//! Feeds back the input stream directly into the output stream
//!
//! You can play and pause the streams by pressing space key
use nannou::prelude::*;
use nannou_audio as audio;
use nannou_audio::Buffer;
use ringbuf::{Consumer, Producer, RingBuffer};

fn main() {
    nannou::app(model).run();
}

struct Model {
    in_stream: audio::Stream<InputModel>,
    out_stream: audio::Stream<OutputModel>,
}

struct InputModel {
    pub producer: Producer<[f32; 2]>,
}

struct OutputModel {
    pub consumer: Consumer<[f32; 2]>,
}

fn model(app: &App) -> Model {
    // Create a window to receive key pressed events.
    app.new_window()
        .key_pressed(key_pressed)
        .view(view)
        .build()
        .unwrap();

    // Initialise the audio host so we can spawn an audio stream.
    let audio_host = audio::Host::new();

    // Create a ring buffer and split it into producer and consumer
    let ring_buffer = RingBuffer::<[f32; 2]>::new(1024 * 2); // Add some latency
    let (prod, cons) = ring_buffer.split();

    // Create input model and input stream using that model
    let in_model = InputModel { producer: prod };
    let in_stream = audio_host
        .new_input_stream(in_model)
        .capture(pass_in)
        .build()
        .unwrap();

    // Create output model and output stream using that model
    let out_model = OutputModel { consumer: cons };
    let out_stream = audio_host
        .new_output_stream(out_model)
        .render(pass_out)
        .build()
        .unwrap();

    Model {
        in_stream,
        out_stream,
    }
}

fn pass_in(model: &mut InputModel, buffer: &Buffer) {
    for frame in buffer.frames() {
        model.producer.push([frame[0], frame[1]]).unwrap();
    }
}

fn pass_out(model: &mut OutputModel, buffer: &mut Buffer) {
    for frame in buffer.frames_mut() {
        let recording_frame = match model.consumer.pop() {
            Some(f) => f,
            None => [0.0, 0.0],
        };
        for (sample, recording_sample) in frame.iter_mut().zip(&recording_frame) {
            *sample = *recording_sample * 0.5;
        }
    }
}

fn key_pressed(_app: &App, model: &mut Model, key: Key) {
    match key {
        Key::Space => {
            if model.in_stream.is_paused() {
                model.in_stream.play().unwrap();
                model.out_stream.play().unwrap();
            } else {
                model.in_stream.pause().unwrap();
                model.out_stream.pause().unwrap();
            }
        }
        _ => {}
    }
}

fn view(_app: &App, _model: &Model, frame: Frame) {
    frame.clear(DIMGRAY);
}
<commit_msg>Make models channel count indepentent<commit_after>//! Feeds back the input stream directly into the output stream
//!
//! You can play and pause the streams by pressing space key
use nannou::prelude::*;
use nannou_audio as audio;
use nannou_audio::Buffer;
use ringbuf::{Consumer, Producer, RingBuffer};

fn main() {
    nannou::app(model).run();
}

struct Model {
    in_stream: audio::Stream<InputModel>,
    out_stream: audio::Stream<OutputModel>,
}

struct InputModel {
    pub producer: Producer<f32>,
}

struct OutputModel {
    pub consumer: Consumer<f32>,
}

fn model(app: &App) -> Model {
    // Create a window to receive key pressed events.
    app.new_window()
        .key_pressed(key_pressed)
        .view(view)
        .build()
        .unwrap();

    // Initialise the audio host so we can spawn an audio stream.
    let audio_host = audio::Host::new();

    // Create a ring buffer and split it into producer and consumer
    let ring_buffer = RingBuffer::<f32>::new(1024 * 2); // Add some latency
    let (prod, cons) = ring_buffer.split();

    // Create input model and input stream using that model
    let in_model = InputModel { producer: prod };
    let in_stream = audio_host
        .new_input_stream(in_model)
        .capture(pass_in)
        .build()
        .unwrap();

    // Create output model and output stream using that model
    let out_model = OutputModel { consumer: cons };
    let out_stream = audio_host
        .new_output_stream(out_model)
        .render(pass_out)
        .build()
        .unwrap();

    Model {
        in_stream,
        out_stream,
    }
}

fn pass_in(model: &mut InputModel, buffer: &Buffer) {
    for frame in buffer.frames() {
        for sample in frame {
            model.producer.push(*sample).unwrap();
        }
    }
}

fn pass_out(model: &mut OutputModel, buffer: &mut Buffer) {
    for frame in buffer.frames_mut() {
        for sample in frame {
            let recorded_sample = match model.consumer.pop() {
                Some(f) => f,
                None => 0.0,
            };
            *sample = recorded_sample;
        }
    }
}

fn key_pressed(_app: &App, model: &mut Model, key: Key) {
    match key {
        Key::Space => {
            if model.in_stream.is_paused() {
                model.in_stream.play().unwrap();
                model.out_stream.play().unwrap();
            } else {
                model.in_stream.pause().unwrap();
                model.out_stream.pause().unwrap();
            }
        }
        _ => {}
    }
}

fn view(_app: &App, _model: &Model, frame: Frame) {
    frame.clear(DIMGRAY);
}
<|endoftext|>