<commit_before>//! TCP stream for communicating with shadowsocks' proxy server

use std::{
    io::{self, ErrorKind},
    pin::Pin,
    task::{self, Poll},
};

use bytes::{BufMut, BytesMut};
use futures::ready;
use log::trace;
use once_cell::sync::Lazy;
use tokio::{
    io::{AsyncRead, AsyncWrite, ReadBuf},
    net::TcpStream,
    time,
};

use crate::{
    config::ServerConfig,
    context::SharedContext,
    net::{ConnectOpts, TcpStream as OutboundTcpStream},
    relay::{
        socks5::Address,
        tcprelay::crypto_io::{CryptoStream, CryptoStreamReadHalf, CryptoStreamWriteHalf},
    },
};

/// A stream for sending / receiving data stream from remote server via shadowsocks' proxy server
pub struct ProxyClientStream<S> {
    stream: CryptoStream<S>,
    addr: Option<Address>,
    context: SharedContext,
}

static DEFAULT_CONNECT_OPTS: Lazy<ConnectOpts> = Lazy::new(Default::default);

impl ProxyClientStream<TcpStream> {
    /// Connect to target `addr` via shadowsocks' server configured by `svr_cfg`
    pub async fn connect<A>(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: A,
    ) -> io::Result<ProxyClientStream<TcpStream>>
    where
        A: Into<Address>,
    {
        ProxyClientStream::connect_with_opts(context, svr_cfg, addr, &DEFAULT_CONNECT_OPTS).await
    }

    /// Connect to target `addr` via shadowsocks' server configured by `svr_cfg`
    pub async fn connect_with_opts<A>(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: A,
        opts: &ConnectOpts,
    ) -> io::Result<ProxyClientStream<TcpStream>>
    where
        A: Into<Address>,
    {
        ProxyClientStream::connect_with_opts_map(context, svr_cfg, addr, opts, |s| s).await
    }
}

impl<S> ProxyClientStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    /// Connect to target `addr` via shadowsocks' server configured by `svr_cfg`, maps `TcpStream` to customized stream with `map_fn`
    pub async fn connect_map<A, F>(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: A,
        map_fn: F,
    ) -> io::Result<ProxyClientStream<S>>
    where
        A: Into<Address>,
        F: FnOnce(TcpStream) -> S,
    {
        ProxyClientStream::connect_with_opts_map(context, svr_cfg, addr, &DEFAULT_CONNECT_OPTS, map_fn).await
    }

    /// Connect to target `addr` via shadowsocks' server configured by `svr_cfg`, maps `TcpStream` to customized stream with `map_fn`
    pub async fn connect_with_opts_map<A, F>(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: A,
        opts: &ConnectOpts,
        map_fn: F,
    ) -> io::Result<ProxyClientStream<S>>
    where
        A: Into<Address>,
        F: FnOnce(TcpStream) -> S,
    {
        let stream = match svr_cfg.timeout() {
            Some(d) => {
                match time::timeout(
                    d,
                    OutboundTcpStream::connect_server_with_opts(&context, svr_cfg.external_addr(), opts),
                )
                .await
                {
                    Ok(Ok(s)) => s,
                    Ok(Err(e)) => return Err(e),
                    Err(..) => {
                        return Err(io::Error::new(
                            ErrorKind::TimedOut,
                            format!("connect {} timeout", svr_cfg.addr()),
                        ))
                    }
                }
            }
            None => OutboundTcpStream::connect_server_with_opts(&context, svr_cfg.external_addr(), opts).await?,
        };

        trace!(
            "connected tcp remote {} (outbound: {}) with {:?}",
            svr_cfg.addr(),
            svr_cfg.external_addr(),
            opts
        );

        Ok(ProxyClientStream::from_stream(
            context,
            map_fn(stream.into()),
            svr_cfg,
            addr,
        ))
    }

    /// Create a `ProxyClientStream` with a connected `stream` to a shadowsocks' server
    ///
    /// NOTE: `stream` must be connected to the server with the same configuration as `svr_cfg`, otherwise strange errors would occurs
    pub fn from_stream<A>(context: SharedContext, stream: S, svr_cfg: &ServerConfig, addr: A) -> ProxyClientStream<S>
    where
        A: Into<Address>,
    {
        let addr = addr.into();
        let stream = CryptoStream::from_stream(&context, stream, svr_cfg.method(), svr_cfg.key());

        ProxyClientStream {
            stream,
            addr: Some(addr),
            context,
        }
    }

    /// Get reference to the underlying stream
    pub fn get_ref(&self) -> &S {
        self.stream.get_ref()
    }

    /// Get mutable reference to the underlying stream
    pub fn get_mut(&mut self) -> &mut S {
        self.stream.get_mut()
    }

    /// Consumes the `ProxyClientStream` and return the underlying stream
    pub fn into_inner(self) -> S {
        self.stream.into_inner()
    }
}

impl<S> AsyncRead for ProxyClientStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    #[inline]
    fn poll_read(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        let context = unsafe { &*(self.context.as_ref() as *const _) };
        self.stream.poll_read_decrypted(cx, context, buf)
    }
}

impl<S> AsyncWrite for ProxyClientStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    fn poll_write(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll<Result<usize, io::Error>> {
        if self.addr.is_none() {
            // For all subsequence calls, just proxy it to self.stream
            return self.stream.poll_write_encrypted(cx, buf);
        }

        let addr = self.addr.take().unwrap();
        let addr_length = addr.serialized_len();

        let mut buffer = BytesMut::with_capacity(addr_length + buf.len());
        addr.write_to_buf(&mut buffer);
        buffer.put_slice(buf);

        ready!(self.stream.poll_write_encrypted(cx, &buffer))?;

        // NOTE:
        // poll_write will return Ok(0) if buf.len() == 0
        // But for the first call, this function will eventually send the handshake packet (IV/Salt + ADDR) to the remote address.
        //
        // https://github.com/shadowsocks/shadowsocks-rust/issues/232
        //
        // For protocols that requires *Server Hello* message, like FTP, clients won't send anything to the server until server sends handshake messages.
        // This could be achieved by calling poll_write with an empty input buffer.

        Ok(buf.len()).into()
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), io::Error>> {
        self.stream.poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), io::Error>> {
        self.stream.poll_shutdown(cx)
    }
}

impl<S> ProxyClientStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    pub fn into_split(self) -> (ProxyClientStreamReadHalf<S>, ProxyClientStreamWriteHalf<S>) {
        let (reader, writer) = self.stream.into_split();
        (
            ProxyClientStreamReadHalf {
                reader,
                context: self.context,
            },
            ProxyClientStreamWriteHalf {
                writer,
                addr: self.addr,
            },
        )
    }
}

pub struct ProxyClientStreamReadHalf<S> {
    reader: CryptoStreamReadHalf<S>,
    context: SharedContext,
}

impl<S> AsyncRead for ProxyClientStreamReadHalf<S>
where
    S: AsyncRead + Unpin,
{
    #[inline]
    fn poll_read(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        let context = unsafe { &*(self.context.as_ref() as *const _) };
        self.reader.poll_read_decrypted(cx, context, buf)
    }
}

pub struct ProxyClientStreamWriteHalf<S> {
    writer: CryptoStreamWriteHalf<S>,
    addr: Option<Address>,
}

impl<S> AsyncWrite for ProxyClientStreamWriteHalf<S>
where
    S: AsyncWrite + Unpin,
{
    fn poll_write(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll<Result<usize, io::Error>> {
        if self.addr.is_none() {
            // For all subsequence calls, just proxy it to self.writer
            return self.writer.poll_write_encrypted(cx, buf);
        }

        let addr = self.addr.take().unwrap();
        let addr_length = addr.serialized_len();

        let mut buffer = BytesMut::with_capacity(addr_length + buf.len());
        addr.write_to_buf(&mut buffer);
        buffer.put_slice(buf);

        ready!(self.writer.poll_write_encrypted(cx, &buffer))?;

        // NOTE:
        // poll_write will return Ok(0) if buf.len() == 0
        // But for the first call, this function will eventually send the handshake packet (IV/Salt + ADDR) to the remote address.
        //
        // https://github.com/shadowsocks/shadowsocks-rust/issues/232
        //
        // For protocols that requires *Server Hello* message, like FTP, clients won't send anything to the server until server sends handshake messages.
        // This could be achieved by calling poll_write with an empty input buffer.

        Ok(buf.len()).into()
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), io::Error>> {
        self.writer.poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), io::Error>> {
        self.writer.poll_shutdown(cx)
    }
}
<commit_msg>fixed bug. ProxyClientStream write may missed the addr if send() returns EAGAIN<commit_after>//! TCP stream for communicating with shadowsocks' proxy server

use std::{
    io::{self, ErrorKind},
    pin::Pin,
    task::{self, Poll},
};

use bytes::{BufMut, BytesMut};
use futures::ready;
use log::trace;
use once_cell::sync::Lazy;
use tokio::{
    io::{AsyncRead, AsyncWrite, ReadBuf},
    net::TcpStream,
    time,
};

use crate::{
    config::ServerConfig,
    context::SharedContext,
    net::{ConnectOpts, TcpStream as OutboundTcpStream},
    relay::{
        socks5::Address,
        tcprelay::crypto_io::{CryptoStream, CryptoStreamReadHalf, CryptoStreamWriteHalf},
    },
};

/// A stream for sending / receiving data stream from remote server via shadowsocks' proxy server
pub struct ProxyClientStream<S> {
    stream: CryptoStream<S>,
    addr: Option<Address>,
    context: SharedContext,
}

static DEFAULT_CONNECT_OPTS: Lazy<ConnectOpts> = Lazy::new(Default::default);

impl ProxyClientStream<TcpStream> {
    /// Connect to target `addr` via shadowsocks' server configured by `svr_cfg`
    pub async fn connect<A>(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: A,
    ) -> io::Result<ProxyClientStream<TcpStream>>
    where
        A: Into<Address>,
    {
        ProxyClientStream::connect_with_opts(context, svr_cfg, addr, &DEFAULT_CONNECT_OPTS).await
    }

    /// Connect to target `addr` via shadowsocks' server configured by `svr_cfg`
    pub async fn connect_with_opts<A>(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: A,
        opts: &ConnectOpts,
    ) -> io::Result<ProxyClientStream<TcpStream>>
    where
        A: Into<Address>,
    {
        ProxyClientStream::connect_with_opts_map(context, svr_cfg, addr, opts, |s| s).await
    }
}

impl<S> ProxyClientStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    /// Connect to target `addr` via shadowsocks' server configured by `svr_cfg`, maps `TcpStream` to customized stream with `map_fn`
    pub async fn connect_map<A, F>(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: A,
        map_fn: F,
    ) -> io::Result<ProxyClientStream<S>>
    where
        A: Into<Address>,
        F: FnOnce(TcpStream) -> S,
    {
        ProxyClientStream::connect_with_opts_map(context, svr_cfg, addr, &DEFAULT_CONNECT_OPTS, map_fn).await
    }

    /// Connect to target `addr` via shadowsocks' server configured by `svr_cfg`, maps `TcpStream` to customized stream with `map_fn`
    pub async fn connect_with_opts_map<A, F>(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: A,
        opts: &ConnectOpts,
        map_fn: F,
    ) -> io::Result<ProxyClientStream<S>>
    where
        A: Into<Address>,
        F: FnOnce(TcpStream) -> S,
    {
        let stream = match svr_cfg.timeout() {
            Some(d) => {
                match time::timeout(
                    d,
                    OutboundTcpStream::connect_server_with_opts(&context, svr_cfg.external_addr(), opts),
                )
                .await
                {
                    Ok(Ok(s)) => s,
                    Ok(Err(e)) => return Err(e),
                    Err(..) => {
                        return Err(io::Error::new(
                            ErrorKind::TimedOut,
                            format!("connect {} timeout", svr_cfg.addr()),
                        ))
                    }
                }
            }
            None => OutboundTcpStream::connect_server_with_opts(&context, svr_cfg.external_addr(), opts).await?,
        };

        trace!(
            "connected tcp remote {} (outbound: {}) with {:?}",
            svr_cfg.addr(),
            svr_cfg.external_addr(),
            opts
        );

        Ok(ProxyClientStream::from_stream(
            context,
            map_fn(stream.into()),
            svr_cfg,
            addr,
        ))
    }

    /// Create a `ProxyClientStream` with a connected `stream` to a shadowsocks' server
    ///
    /// NOTE: `stream` must be connected to the server with the same configuration as `svr_cfg`, otherwise strange errors would occurs
    pub fn from_stream<A>(context: SharedContext, stream: S, svr_cfg: &ServerConfig, addr: A) -> ProxyClientStream<S>
    where
        A: Into<Address>,
    {
        let addr = addr.into();
        let stream = CryptoStream::from_stream(&context, stream, svr_cfg.method(), svr_cfg.key());

        ProxyClientStream {
            stream,
            addr: Some(addr),
            context,
        }
    }

    /// Get reference to the underlying stream
    pub fn get_ref(&self) -> &S {
        self.stream.get_ref()
    }

    /// Get mutable reference to the underlying stream
    pub fn get_mut(&mut self) -> &mut S {
        self.stream.get_mut()
    }

    /// Consumes the `ProxyClientStream` and return the underlying stream
    pub fn into_inner(self) -> S {
        self.stream.into_inner()
    }
}

impl<S> AsyncRead for ProxyClientStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    #[inline]
    fn poll_read(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        let context = unsafe { &*(self.context.as_ref() as *const _) };
        self.stream.poll_read_decrypted(cx, context, buf)
    }
}

impl<S> AsyncWrite for ProxyClientStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    fn poll_write(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll<Result<usize, io::Error>> {
        match self.addr {
            None => {
                // For all subsequence calls, just proxy it to self.stream
                return self.stream.poll_write_encrypted(cx, buf);
            }
            Some(ref addr) => {
                let addr_length = addr.serialized_len();

                let mut buffer = BytesMut::with_capacity(addr_length + buf.len());
                addr.write_to_buf(&mut buffer);
                buffer.put_slice(buf);

                ready!(self.stream.poll_write_encrypted(cx, &buffer))?;

                // fallthrough. take the self.addr out
            }
        }

        let _ = self.addr.take();

        // NOTE:
        // poll_write will return Ok(0) if buf.len() == 0
        // But for the first call, this function will eventually send the handshake packet (IV/Salt + ADDR) to the remote address.
        //
        // https://github.com/shadowsocks/shadowsocks-rust/issues/232
        //
        // For protocols that requires *Server Hello* message, like FTP, clients won't send anything to the server until server sends handshake messages.
        // This could be achieved by calling poll_write with an empty input buffer.

        Ok(buf.len()).into()
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), io::Error>> {
        self.stream.poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), io::Error>> {
        self.stream.poll_shutdown(cx)
    }
}

impl<S> ProxyClientStream<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    pub fn into_split(self) -> (ProxyClientStreamReadHalf<S>, ProxyClientStreamWriteHalf<S>) {
        let (reader, writer) = self.stream.into_split();
        (
            ProxyClientStreamReadHalf {
                reader,
                context: self.context,
            },
            ProxyClientStreamWriteHalf {
                writer,
                addr: self.addr,
            },
        )
    }
}

pub struct ProxyClientStreamReadHalf<S> {
    reader: CryptoStreamReadHalf<S>,
    context: SharedContext,
}

impl<S> AsyncRead for ProxyClientStreamReadHalf<S>
where
    S: AsyncRead + Unpin,
{
    #[inline]
    fn poll_read(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        let context = unsafe { &*(self.context.as_ref() as *const _) };
        self.reader.poll_read_decrypted(cx, context, buf)
    }
}

pub struct ProxyClientStreamWriteHalf<S> {
    writer: CryptoStreamWriteHalf<S>,
    addr: Option<Address>,
}

impl<S> AsyncWrite for ProxyClientStreamWriteHalf<S>
where
    S: AsyncWrite + Unpin,
{
    fn poll_write(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll<Result<usize, io::Error>> {
        if self.addr.is_none() {
            // For all subsequence calls, just proxy it to self.writer
            return self.writer.poll_write_encrypted(cx, buf);
        }

        let addr = self.addr.take().unwrap();
        let addr_length = addr.serialized_len();

        let mut buffer = BytesMut::with_capacity(addr_length + buf.len());
        addr.write_to_buf(&mut buffer);
        buffer.put_slice(buf);

        ready!(self.writer.poll_write_encrypted(cx, &buffer))?;

        // NOTE:
        // poll_write will return Ok(0) if buf.len() == 0
        // But for the first call, this function will eventually send the handshake packet (IV/Salt + ADDR) to the remote address.
        //
        // https://github.com/shadowsocks/shadowsocks-rust/issues/232
        //
        // For protocols that requires *Server Hello* message, like FTP, clients won't send anything to the server until server sends handshake messages.
        // This could be achieved by calling poll_write with an empty input buffer.

        Ok(buf.len()).into()
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), io::Error>> {
        self.writer.poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), io::Error>> {
        self.writer.poll_shutdown(cx)
    }
}
<|endoftext|>


<commit_before>use itertools::Itertools;

use model;

pub type CV = f64;

#[allow(unused_parens)]
pub fn pmf(pmfs: &[model::expressionset::PMF]) -> model::diffexp::PMF {
    let meanvar = model::pmf::MeanVar::new(pmfs);

    let mut pmf = meanvar.iter().map(|e| model::pmf::Entry {
        value: e.standard_deviation().log2() - e.mean.log2(),
        prob: e.prob
    }).collect_vec();
    pmf.sort_by(|a, b| a.value.partial_cmp(&b.value).unwrap());

    model::diffexp::PMF::new(pmf)
}

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

    use itertools::Itertools;
    use bio::stats::logprobs;

    use model;
    use io;

    const GENE: &'static str = "COL5A1";

    fn setup() -> Box<model::readout::Model> {
        model::readout::new_model(16, 4, 0.04, 0.1, 4, io::codebook::Reader::from_file("evaluation/codebook/140genesData.1.txt", 4).unwrap().codebook())
    }

    #[test]
    fn test_pmf() {
        let readout = setup();
        let pmfs1 = [
            model::expression::pmf(GENE, 5, 5, &readout, 100),
            model::expression::pmf(GENE, 5, 5, &readout, 100),
            model::expression::pmf(GENE, 5, 5, &readout, 100),
            model::expression::pmf(GENE, 5, 5, &readout, 100)
        ];
        let pmfs2 = [
            model::expression::pmf(GENE, 50, 50, &readout, 100),
            model::expression::pmf(GENE, 50, 50, &readout, 100),
            model::expression::pmf(GENE, 50, 50, &readout, 100),
            model::expression::pmf(GENE, 50, 50, &readout, 100)
        ];
        let pmf1 = model::expressionset::pmf(&pmfs1);
        let pmf2 = model::expressionset::pmf(&pmfs2);
        //println!("{} {}", pmfs1[0].expected_value(), pmfs2[0].expected_value());

        let pmf = pmf(&[pmf1, pmf2]);

        let total = logprobs::sum(&pmf.iter().map(|fc| fc.prob).collect_vec());

        assert!(total <= 0.0);
        assert_relative_eq!(total, 0.0, epsilon = 0.0002);
        assert_relative_eq!(2.0f64.powf(pmf.expected_value()), 1.14, epsilon = 0.01);
    }
}
<commit_msg>Cleanup.<commit_after>use itertools::Itertools;

use model;

pub type CV = f64;

pub fn pmf(pmfs: &[model::expressionset::PMF]) -> model::diffexp::PMF {
    let meanvar = model::pmf::MeanVar::new(pmfs);

    let mut pmf = meanvar.iter().map(|e| model::pmf::Entry {
        value: e.standard_deviation().log2() - e.mean.log2(),
        prob: e.prob
    }).collect_vec();
    pmf.sort_by(|a, b| a.value.partial_cmp(&b.value).unwrap());

    model::diffexp::PMF::new(pmf)
}

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

    use itertools::Itertools;
    use bio::stats::logprobs;

    use model;
    use io;

    const GENE: &'static str = "COL5A1";

    fn setup() -> Box<model::readout::Model> {
        model::readout::new_model(16, 4, 0.04, 0.1, 4, io::codebook::Reader::from_file("evaluation/codebook/140genesData.1.txt", 4).unwrap().codebook())
    }

    #[test]
    fn test_pmf() {
        let readout = setup();
        let pmfs1 = [
            model::expression::pmf(GENE, 5, 5, &readout, 100),
            model::expression::pmf(GENE, 5, 5, &readout, 100),
            model::expression::pmf(GENE, 5, 5, &readout, 100),
            model::expression::pmf(GENE, 5, 5, &readout, 100)
        ];
        let pmfs2 = [
            model::expression::pmf(GENE, 50, 50, &readout, 100),
            model::expression::pmf(GENE, 50, 50, &readout, 100),
            model::expression::pmf(GENE, 50, 50, &readout, 100),
            model::expression::pmf(GENE, 50, 50, &readout, 100)
        ];
        let pmf1 = model::expressionset::pmf(&pmfs1);
        let pmf2 = model::expressionset::pmf(&pmfs2);
        //println!("{} {}", pmfs1[0].expected_value(), pmfs2[0].expected_value());

        let pmf = pmf(&[pmf1, pmf2]);

        let total = logprobs::sum(&pmf.iter().map(|fc| fc.prob).collect_vec());

        assert!(total <= 0.0);
        assert_relative_eq!(total, 0.0, epsilon = 0.0002);
        assert_relative_eq!(2.0f64.powf(pmf.expected_value()), 1.14, epsilon = 0.01);
    }
}
<|endoftext|>


<commit_before>#![feature(no_panic_pow)]
#![feature(test)]

extern crate test;

use can_dbc::{ByteOrder, Message, MessageId, Signal, ValueDescription, DBC, MultiplexIndicator};
use codegen::{Enum, Function, Impl, Scope, Struct};
use heck::{CamelCase, ShoutySnakeCase};
use log::warn;

use std::fmt::Write;

#[cfg(test)]
mod tests {
    use byteorder;
    use byteorder::{ByteOrder, LE};
    use test::{black_box, Bencher};

    #[bench]
    fn bench_read_signal(b: &mut Bencher) {
        const byte_payload: &[u8] = &[
            0x4, 0x2, 0xA, 0xA, 0xF, 0xF, 0xE, 0xE, 0xD, 0xD, 0xA, 0xA, 0xF, 0xF, 0xD, 0xD,
        ];

        b.iter(|| {
            let frame_payload: u64 = LE::read_u64(byte_payload);
            let bit_msk_const = 2u64.saturating_pow(8 as u32) - 1;
            let factor: f64 = test::black_box(2.0);
            let start_bit: u64 = test::black_box(8);
            let offset: f64 = test::black_box(10.0);

            (((frame_payload >> start_bit) & bit_msk_const) as f64) * factor + offset
        });
    }
}

/// Character that is prefixed before type names that are
/// are not starting with an alphabetic char.
const PREFIX_CHAR: char = 'X';

/// Character that is used to replace invalid characters
/// in type names.
const REPLACEMENT_CHAR: char = 'X';

/// Suffix that is append to the raw signal function
const RAW_FN_SUFFIX: &str = "raw_value";

type Result<T> = std::result::Result<T, std::fmt::Error>;

#[derive(Debug)]
pub struct DbccOpt {
    /// Should tokio SocketCan BCM streams be generated.
    /// This requires the `tokio-socketcan-bcm` crate.
    pub with_tokio: bool,
}

pub trait TypeName: ToOwned {
    fn to_type_name(&self) -> Self::Owned;
}

impl TypeName for str {
    fn to_type_name(&self) -> String {
        let mut out = String::with_capacity(self.len() + 1);
        let mut chars = self.chars();
        if let Some(first) = chars.next() {
            if !first.is_alphabetic() && first != '_' {
                warn!("string: {} is prefixed with `{}`", self, PREFIX_CHAR);
                out.push(PREFIX_CHAR);
            }
            out.push(first);
        }

        while let Some(chr) = chars.next() {
            if chr.is_digit(10) || chr.is_alphabetic() || chr == '_' {
                out.push(chr);
            } else {
                warn!(
                    "`{}` character in string: {} is replaced by `{}`",
                    chr, self, REPLACEMENT_CHAR
                );
                out.push(REPLACEMENT_CHAR);
            }
        }

        out
    }
}

fn to_enum_name(message_id: &MessageId, signal_name: &str) -> String {
    format!("{}{}", &signal_name.to_camel_case(), message_id.0)
}

pub fn signal_enum(val_desc: &ValueDescription) -> Option<Enum> {
    if let ValueDescription::Signal {
        ref message_id,
        ref signal_name,
        ref value_descriptions,
    } = val_desc
    {
        let mut sig_enum = Enum::new(&to_enum_name(message_id, signal_name));
        sig_enum.allow("dead_code");
        sig_enum.vis("pub");
        sig_enum.repr("u64");
        sig_enum.derive("Debug");
        sig_enum.derive("Clone");
        sig_enum.derive("Copy");
        sig_enum.derive("PartialEq");
        sig_enum.derive("Eq");
        for desc in value_descriptions {
            sig_enum.new_variant(&desc.b().to_camel_case().to_type_name());
        }
        sig_enum.new_variant("XValue(u64)");
        return Some(sig_enum);
    }
    None
}

pub fn signal_enum_impl_from(val_desc: &ValueDescription) -> Option<Impl> {
    if let ValueDescription::Signal {
        ref message_id,
        ref signal_name,
        ref value_descriptions,
    } = val_desc
    {
        let enum_name = to_enum_name(message_id, signal_name);
        let mut enum_impl = Impl::new(codegen::Type::new(&enum_name));
        enum_impl.impl_trait("From<u64>");

        let from_fn = enum_impl.new_fn("from");
        from_fn.allow("dead_code");
        from_fn.arg("val", codegen::Type::new("u64"));
        from_fn.ret(codegen::Type::new("Self"));

        let mut matching = String::new();
        write!(&mut matching, "match val {{\n").unwrap();
        for value_description in value_descriptions {
            write!(
                &mut matching,
                "    {} => {}::{},\n",
                value_description.a(),
                enum_name,
                value_description.b().to_camel_case().to_type_name()
            )
            .unwrap();
        }
        write!(
            &mut matching,
            "    value => {}::XValue(value),\n",
            enum_name
        )
        .unwrap();
        write!(&mut matching, "}}").unwrap();

        from_fn.line(matching);

        return Some(enum_impl);
    }
    None
}

pub fn signal_fn_raw(dbc: &DBC, signal: &Signal, message_id: &MessageId) -> Result<Function> {
    let raw_fn_name = format!("{}_{}", signal.name().to_lowercase(), RAW_FN_SUFFIX);

    let mut signal_fn = codegen::Function::new(&raw_fn_name);
    signal_fn.allow("dead_code");
    signal_fn.vis("pub");
    signal_fn.arg_ref_self();

    let signal_return_type = signal_return_type(signal);
    signal_fn.ret(codegen::Type::new(&signal_return_type));

    let default_signal_comment = format!("Read {} signal from can frame", signal.name());
    let signal_comment = dbc
        .signal_comment(message_id, signal.name())
        .unwrap_or(&default_signal_comment);

    let signal_unit = if !signal.unit().is_empty() {
        format!("\nUnit: {}", signal.unit())
    } else {
        String::default()
    };

    signal_fn.doc(&format!("{}{}", signal_comment, signal_unit));

    let read_byte_order = match signal.byte_order() {
        ByteOrder::LittleEndian => "let frame_payload: u64 = LE::read_u64(&self.frame_payload);",
        ByteOrder::BigEndian => "let  frame_payload: u64 = BE::read_u64(&self.frame_payload);",
    };
    signal_fn.line(read_byte_order);

    let bit_msk_const = 2u64.saturating_pow(*signal.signal_size() as u32) - 1;
    let signal_shift = shift_amount(
        *signal.byte_order(),
        *signal.start_bit(),
        *signal.signal_size(),
    );

    let calc = calc_raw(signal, signal_return_type, signal_shift, bit_msk_const)?;
    signal_fn.line(calc);

    Ok(signal_fn)
}

pub fn signal_fn_enum(signal: &Signal, enum_type: String) -> Result<Function> {
    let mut signal_fn = codegen::Function::new(&signal.name().to_lowercase());
    signal_fn.allow("dead_code");
    signal_fn.vis("pub");
    signal_fn.arg_ref_self();

    signal_fn.ret(enum_type.clone());

    let raw_fn_name = format!("{}_{}", signal.name().to_lowercase(), RAW_FN_SUFFIX);

    signal_fn.line(format!(
        "{}::from(self.{}() as u64)",
        enum_type, raw_fn_name
    ));

    Ok(signal_fn)
}

fn calc_raw(
    signal: &Signal,
    signal_return_type: String,
    signal_shift: u64,
    bit_msk_const: u64,
) -> Result<String> {
    let mut calc = String::new();

    // No shift required if start_bit == 0
    let shift = if signal_shift != 0 {
        format!("(frame_payload >> {})", signal_shift)
    } else {
        format!("frame_payload")
    };

    write!(&mut calc, "({} & {:#X})", shift, bit_msk_const)?;

    if *signal.signal_size() != 1 {
        write!(&mut calc, " as {}", signal_return_type)?;
    }

    if *signal.factor() != 1.0 {
        write!(&mut calc, " * {:.6}", signal.factor())?;
    }

    if *signal.offset() != 0.0 && *signal.signal_size() <= 32 {
        write!(&mut calc, " + {}f32", signal.offset())?;
    } else if *signal.offset() != 0.0 {
        write!(&mut calc, " + {}f64", signal.offset())?;
    }

    // boolean signal
    if *signal.signal_size() == 1 {
        write!(&mut calc, " == 1")?;
    }

    Ok(calc)
}

fn signal_return_type(signal: &Signal) -> String {
    match signal.signal_size() {
        _ if *signal.signal_size() == 1 => "bool".to_string(),
        _ if *signal.signal_size() > 1 && *signal.signal_size() <= 32 => "f32".to_string(),
        _ => "f64".to_string(),
    }
}

fn shift_amount(byte_order: ByteOrder, start_bit: u64, signal_size: u64) -> u64 {
    match byte_order {
        ByteOrder::LittleEndian => start_bit,
        ByteOrder::BigEndian => 64 - signal_size - ((start_bit / 8) * 8 + (7 - (start_bit % 8))),
    }
}

fn message_const(message: &Message) -> String {
    format!(
        "#[allow(dead_code)]\npub const MESSAGE_ID_{}: u32 = {};",
        message.message_name().to_shouty_snake_case(),
        message.message_id().0
    )
}

fn message_struct(dbc: &DBC, message: &Message) -> Struct {
    let mut message_struct = Struct::new(&message.message_name().to_camel_case());
    if let Some(message_comment) = dbc.message_comment(message.message_id()) {
        message_struct.doc(message_comment);
    }
    message_struct.allow("dead_code");
    message_struct.derive("Debug");
    message_struct.vis("pub");
    message_struct.field("frame_payload", "Vec<u8>");
    message_struct
}

fn message_impl(opt: &DbccOpt, dbc: &DBC, message: &Message) -> Result<Impl> {
    let mut msg_impl = Impl::new(codegen::Type::new(&message.message_name().to_camel_case()));

    let new_fn = msg_impl.new_fn("new");
    new_fn.allow("dead_code");
    new_fn.vis("pub");
    new_fn.arg("mut frame_payload", codegen::Type::new("Vec<u8>"));
    new_fn.line("frame_payload.resize(8, 0);");
    new_fn.line(format!(
        "{} {{ frame_payload }}",
        message.message_name().to_camel_case()
    ));
    new_fn.ret(codegen::Type::new(&message.message_name().to_camel_case()));

    if opt.with_tokio {
        msg_impl.push_fn(message_stream(message));
    }

    for signal in message.signals() {

        if *signal.multiplexer_indicator() != MultiplexIndicator::Plain {
            warn!("Multiplexed signals are currently not supported, the message `{}` signal `{}` will be skipped", message.message_name(), signal.name());
            continue;
        }

        msg_impl.push_fn(signal_fn_raw(dbc, signal, message.message_id())?);

        // Check if this signal can be turned into an enum
        let enum_type = dbc
            .value_descriptions_for_signal(message.message_id(), signal.name())
            .map(|_| to_enum_name(message.message_id(), signal.name()));
        if let Some(enum_type) = enum_type {
            msg_impl.push_fn(signal_fn_enum(signal, enum_type)?);
        }
    }

    Ok(msg_impl)
}

/// Generate message stream using socketcan's Broadcast Manager filters via socketcan-tokio.
fn message_stream(message: &Message) -> Function {
    let mut stream_fn = codegen::Function::new("stream");
    stream_fn.allow("dead_code");
    stream_fn.vis("pub");

    stream_fn.arg("can_interface", codegen::Type::new("&str"));
    stream_fn.arg("ival1", codegen::Type::new("&std::time::Duration"));
    stream_fn.arg("ival2", codegen::Type::new("&std::time::Duration"));

    let ret = format!(
        "std::io::Result<Box<Stream<Item = {}, Error = std::io::Error> + Sync + Send>>",
        message.message_name().to_camel_case()
    );
    stream_fn.ret(ret);

    stream_fn.line("let socket = BCMSocket::open_nb(&can_interface)?;");
    stream_fn.line(format!(
        "let message_id = CANMessageId::try_from({} as u32).unwrap();",
        message.message_id().0.to_string()
    ));
    stream_fn.line("let frame_stream = socket.filter_id_incoming_frames(message_id, ival1.clone(), ival2.clone())?;");
    stream_fn.line(format!(
        "let f = frame_stream.map(|frame| {}::new(frame.data().to_vec()));",
        message.message_name().to_camel_case()
    ));
    stream_fn.line("Ok(Box::new(f))");

    stream_fn
}

pub fn can_reader(opt: &DbccOpt, dbc: &DBC) -> Result<Scope> {
    let mut scope = Scope::new();
    scope.import("byteorder", "{ByteOrder, LE, BE}");

    if opt.with_tokio {
        scope.import("tokio_socketcan_bcm", "{CANMessageId, BCMSocket}");
        scope.import("futures", "Stream");
        scope.import("std::convert", "TryFrom");
    }

    for message in dbc.messages() {
        scope.raw(&message_const(message));
    }

    for value_description in dbc.value_descriptions() {
        if let Some(signal_enum) = signal_enum(value_description) {
            scope.push_enum(signal_enum);
        }

        if let Some(enum_impl) = signal_enum_impl_from(value_description) {
            scope.push_impl(enum_impl);
        }
    }

    for message in dbc.messages() {
        scope.push_struct(message_struct(&dbc, message));
        scope.push_impl(message_impl(opt, &dbc, message)?);
    }

    Ok(scope)
}
<commit_msg>Replace unwraps for generated CANMessageIDs<commit_after>#![feature(no_panic_pow)]
#![feature(test)]

extern crate test;

use can_dbc::{ByteOrder, Message, MessageId, Signal, ValueDescription, DBC, MultiplexIndicator};
use codegen::{Enum, Function, Impl, Scope, Struct};
use heck::{CamelCase, ShoutySnakeCase};
use log::warn;
use socketcan::{SFF_MASK, EFF_MASK};

use std::fmt::Write;

#[cfg(test)]
mod tests {
    use byteorder;
    use byteorder::{ByteOrder, LE};
    use test::{black_box, Bencher};

    #[bench]
    fn bench_read_signal(b: &mut Bencher) {
        const byte_payload: &[u8] = &[
            0x4, 0x2, 0xA, 0xA, 0xF, 0xF, 0xE, 0xE, 0xD, 0xD, 0xA, 0xA, 0xF, 0xF, 0xD, 0xD,
        ];

        b.iter(|| {
            let frame_payload: u64 = LE::read_u64(byte_payload);
            let bit_msk_const = 2u64.saturating_pow(8 as u32) - 1;
            let factor: f64 = test::black_box(2.0);
            let start_bit: u64 = test::black_box(8);
            let offset: f64 = test::black_box(10.0);

            (((frame_payload >> start_bit) & bit_msk_const) as f64) * factor + offset
        });
    }
}

/// Character that is prefixed before type names that are
/// are not starting with an alphabetic char.
const PREFIX_CHAR: char = 'X';

/// Character that is used to replace invalid characters
/// in type names.
const REPLACEMENT_CHAR: char = 'X';

/// Suffix that is append to the raw signal function
const RAW_FN_SUFFIX: &str = "raw_value";

type Result<T> = std::result::Result<T, std::fmt::Error>;

#[derive(Debug)]
pub struct DbccOpt {
    /// Should tokio SocketCan BCM streams be generated.
    /// This requires the `tokio-socketcan-bcm` crate.
    pub with_tokio: bool,
}

pub trait TypeName: ToOwned {
    fn to_type_name(&self) -> Self::Owned;
}

impl TypeName for str {
    fn to_type_name(&self) -> String {
        let mut out = String::with_capacity(self.len() + 1);
        let mut chars = self.chars();
        if let Some(first) = chars.next() {
            if !first.is_alphabetic() && first != '_' {
                warn!("string: {} is prefixed with `{}`", self, PREFIX_CHAR);
                out.push(PREFIX_CHAR);
            }
            out.push(first);
        }

        while let Some(chr) = chars.next() {
            if chr.is_digit(10) || chr.is_alphabetic() || chr == '_' {
                out.push(chr);
            } else {
                warn!(
                    "`{}` character in string: {} is replaced by `{}`",
                    chr, self, REPLACEMENT_CHAR
                );
                out.push(REPLACEMENT_CHAR);
            }
        }

        out
    }
}

fn to_enum_name(message_id: &MessageId, signal_name: &str) -> String {
    format!("{}{}", &signal_name.to_camel_case(), message_id.0)
}

pub fn signal_enum(val_desc: &ValueDescription) -> Option<Enum> {
    if let ValueDescription::Signal {
        ref message_id,
        ref signal_name,
        ref value_descriptions,
    } = val_desc
    {
        let mut sig_enum = Enum::new(&to_enum_name(message_id, signal_name));
        sig_enum.allow("dead_code");
        sig_enum.vis("pub");
        sig_enum.repr("u64");
        sig_enum.derive("Debug");
        sig_enum.derive("Clone");
        sig_enum.derive("Copy");
        sig_enum.derive("PartialEq");
        sig_enum.derive("Eq");
        for desc in value_descriptions {
            sig_enum.new_variant(&desc.b().to_camel_case().to_type_name());
        }
        sig_enum.new_variant("XValue(u64)");
        return Some(sig_enum);
    }
    None
}

pub fn signal_enum_impl_from(val_desc: &ValueDescription) -> Option<Impl> {
    if let ValueDescription::Signal {
        ref message_id,
        ref signal_name,
        ref value_descriptions,
    } = val_desc
    {
        let enum_name = to_enum_name(message_id, signal_name);
        let mut enum_impl = Impl::new(codegen::Type::new(&enum_name));
        enum_impl.impl_trait("From<u64>");

        let from_fn = enum_impl.new_fn("from");
        from_fn.allow("dead_code");
        from_fn.arg("val", codegen::Type::new("u64"));
        from_fn.ret(codegen::Type::new("Self"));

        let mut matching = String::new();
        write!(&mut matching, "match val {{\n").unwrap();
        for value_description in value_descriptions {
            write!(
                &mut matching,
                "    {} => {}::{},\n",
                value_description.a(),
                enum_name,
                value_description.b().to_camel_case().to_type_name()
            )
            .unwrap();
        }
        write!(
            &mut matching,
            "    value => {}::XValue(value),\n",
            enum_name
        )
        .unwrap();
        write!(&mut matching, "}}").unwrap();

        from_fn.line(matching);

        return Some(enum_impl);
    }
    None
}

pub fn signal_fn_raw(dbc: &DBC, signal: &Signal, message_id: &MessageId) -> Result<Function> {
    let raw_fn_name = format!("{}_{}", signal.name().to_lowercase(), RAW_FN_SUFFIX);

    let mut signal_fn = codegen::Function::new(&raw_fn_name);
    signal_fn.allow("dead_code");
    signal_fn.vis("pub");
    signal_fn.arg_ref_self();

    let signal_return_type = signal_return_type(signal);
    signal_fn.ret(codegen::Type::new(&signal_return_type));

    let default_signal_comment = format!("Read {} signal from can frame", signal.name());
    let signal_comment = dbc
        .signal_comment(message_id, signal.name())
        .unwrap_or(&default_signal_comment);

    let signal_unit = if !signal.unit().is_empty() {
        format!("\nUnit: {}", signal.unit())
    } else {
        String::default()
    };

    signal_fn.doc(&format!("{}{}", signal_comment, signal_unit));

    let read_byte_order = match signal.byte_order() {
        ByteOrder::LittleEndian => "let frame_payload: u64 = LE::read_u64(&self.frame_payload);",
        ByteOrder::BigEndian => "let  frame_payload: u64 = BE::read_u64(&self.frame_payload);",
    };
    signal_fn.line(read_byte_order);

    let bit_msk_const = 2u64.saturating_pow(*signal.signal_size() as u32) - 1;
    let signal_shift = shift_amount(
        *signal.byte_order(),
        *signal.start_bit(),
        *signal.signal_size(),
    );

    let calc = calc_raw(signal, signal_return_type, signal_shift, bit_msk_const)?;
    signal_fn.line(calc);

    Ok(signal_fn)
}

pub fn signal_fn_enum(signal: &Signal, enum_type: String) -> Result<Function> {
    let mut signal_fn = codegen::Function::new(&signal.name().to_lowercase());
    signal_fn.allow("dead_code");
    signal_fn.vis("pub");
    signal_fn.arg_ref_self();

    signal_fn.ret(enum_type.clone());

    let raw_fn_name = format!("{}_{}", signal.name().to_lowercase(), RAW_FN_SUFFIX);

    signal_fn.line(format!(
        "{}::from(self.{}() as u64)",
        enum_type, raw_fn_name
    ));

    Ok(signal_fn)
}

fn calc_raw(
    signal: &Signal,
    signal_return_type: String,
    signal_shift: u64,
    bit_msk_const: u64,
) -> Result<String> {
    let mut calc = String::new();

    // No shift required if start_bit == 0
    let shift = if signal_shift != 0 {
        format!("(frame_payload >> {})", signal_shift)
    } else {
        format!("frame_payload")
    };

    write!(&mut calc, "({} & {:#X})", shift, bit_msk_const)?;

    if *signal.signal_size() != 1 {
        write!(&mut calc, " as {}", signal_return_type)?;
    }

    if *signal.factor() != 1.0 {
        write!(&mut calc, " * {:.6}", signal.factor())?;
    }

    if *signal.offset() != 0.0 && *signal.signal_size() <= 32 {
        write!(&mut calc, " + {}f32", signal.offset())?;
    } else if *signal.offset() != 0.0 {
        write!(&mut calc, " + {}f64", signal.offset())?;
    }

    // boolean signal
    if *signal.signal_size() == 1 {
        write!(&mut calc, " == 1")?;
    }

    Ok(calc)
}

fn signal_return_type(signal: &Signal) -> String {
    match signal.signal_size() {
        _ if *signal.signal_size() == 1 => "bool".to_string(),
        _ if *signal.signal_size() > 1 && *signal.signal_size() <= 32 => "f32".to_string(),
        _ => "f64".to_string(),
    }
}

fn shift_amount(byte_order: ByteOrder, start_bit: u64, signal_size: u64) -> u64 {
    match byte_order {
        ByteOrder::LittleEndian => start_bit,
        ByteOrder::BigEndian => 64 - signal_size - ((start_bit / 8) * 8 + (7 - (start_bit % 8))),
    }
}

fn message_const(message: &Message) -> String {
    format!(
        "#[allow(dead_code)]\npub const MESSAGE_ID_{}: u32 = {};",
        message.message_name().to_shouty_snake_case(),
        message.message_id().0
    )
}

fn message_struct(dbc: &DBC, message: &Message) -> Struct {
    let mut message_struct = Struct::new(&message.message_name().to_camel_case());
    if let Some(message_comment) = dbc.message_comment(message.message_id()) {
        message_struct.doc(message_comment);
    }
    message_struct.allow("dead_code");
    message_struct.derive("Debug");
    message_struct.vis("pub");
    message_struct.field("frame_payload", "Vec<u8>");
    message_struct
}

fn message_impl(opt: &DbccOpt, dbc: &DBC, message: &Message) -> Result<Impl> {
    let mut msg_impl = Impl::new(codegen::Type::new(&message.message_name().to_camel_case()));

    let new_fn = msg_impl.new_fn("new");
    new_fn.allow("dead_code");
    new_fn.vis("pub");
    new_fn.arg("mut frame_payload", codegen::Type::new("Vec<u8>"));
    new_fn.line("frame_payload.resize(8, 0);");
    new_fn.line(format!(
        "{} {{ frame_payload }}",
        message.message_name().to_camel_case()
    ));
    new_fn.ret(codegen::Type::new(&message.message_name().to_camel_case()));

    if opt.with_tokio {
        msg_impl.push_fn(message_stream(message));
    }

    for signal in message.signals() {

        if *signal.multiplexer_indicator() != MultiplexIndicator::Plain {
            warn!("Multiplexed signals are currently not supported, the message `{}` signal `{}` will be skipped", message.message_name(), signal.name());
            continue;
        }

        msg_impl.push_fn(signal_fn_raw(dbc, signal, message.message_id())?);

        // Check if this signal can be turned into an enum
        let enum_type = dbc
            .value_descriptions_for_signal(message.message_id(), signal.name())
            .map(|_| to_enum_name(message.message_id(), signal.name()));
        if let Some(enum_type) = enum_type {
            msg_impl.push_fn(signal_fn_enum(signal, enum_type)?);
        }
    }

    Ok(msg_impl)
}

/// Generate message stream using socketcan's Broadcast Manager filters via socketcan-tokio.
fn message_stream(message: &Message) -> Function {
    let mut stream_fn = codegen::Function::new("stream");
    stream_fn.allow("dead_code");
    stream_fn.vis("pub");

    stream_fn.arg("can_interface", codegen::Type::new("&str"));
    stream_fn.arg("ival1", codegen::Type::new("&std::time::Duration"));
    stream_fn.arg("ival2", codegen::Type::new("&std::time::Duration"));

    let ret = format!(
        "std::io::Result<Box<Stream<Item = {}, Error = std::io::Error> + Sync + Send>>",
        message.message_name().to_camel_case()
    );
    stream_fn.ret(ret);

    stream_fn.line("let socket = BCMSocket::open_nb(&can_interface)?;");

    let message_id = match message.message_id().0 {
            0...SFF_MASK => format!("let message_id = CANMessageId::SFF({} as u16);", message.message_id().0.to_string()),
            SFF_MASK...EFF_MASK => format!("let message_id = CANMessageId::EFF({} as u16);", message.message_id().0.to_string()),
            _ => panic!(format!("can message identifier `{}` exceeds max size", message.message_id().0)),
    };
    stream_fn.line(message_id);

    stream_fn.line("let frame_stream = socket.filter_id_incoming_frames(message_id, ival1.clone(), ival2.clone())?;");
    stream_fn.line(format!(
        "let f = frame_stream.map(|frame| {}::new(frame.data().to_vec()));",
        message.message_name().to_camel_case()
    ));
    stream_fn.line("Ok(Box::new(f))");

    stream_fn
}

pub fn can_reader(opt: &DbccOpt, dbc: &DBC) -> Result<Scope> {
    let mut scope = Scope::new();
    scope.import("byteorder", "{ByteOrder, LE, BE}");

    if opt.with_tokio {
        scope.import("tokio_socketcan_bcm", "{CANMessageId, BCMSocket}");
        scope.import("futures", "Stream");
        scope.import("std::convert", "TryFrom");
    }

    for message in dbc.messages() {
        scope.raw(&message_const(message));
    }

    for value_description in dbc.value_descriptions() {
        if let Some(signal_enum) = signal_enum(value_description) {
            scope.push_enum(signal_enum);
        }

        if let Some(enum_impl) = signal_enum_impl_from(value_description) {
            scope.push_impl(enum_impl);
        }
    }

    for message in dbc.messages() {
        scope.push_struct(message_struct(&dbc, message));
        scope.push_impl(message_impl(opt, &dbc, message)?);
    }

    Ok(scope)
}
<|endoftext|>


<commit_before>pub mod cli;

use serde::{Serialize};

use std::hash::Hash;
use std::fmt;
use integeriser::{Integeriser, HashIntegeriser};
use std::collections::{HashSet, BTreeMap, BTreeSet};

use pmcfg::{PMCFG, PMCFGRule, VarT};
use dyck::multiple::{Bracket, MultipleDyckLanguage};
use openfsa::fsa::{Automaton, generator};
use log_domain::LogDomain;
use dyck;

use util::partition::Partition;

pub mod bracket_fragment;
pub mod filter_automaton;
pub mod generator_automaton;
use cs_representation::bracket_fragment::BracketFragment;
use cs_representation::generator_automaton::{GeneratorAutomaton, Delta};
use cs_representation::filter_automaton::{FilterAutomaton};

use std::fmt::{Display, Formatter, Error};

/// A derivation tree of PMCFG rules.
#[derive(PartialEq, Debug)]
pub struct Derivation<'a, N: 'a, T: 'a>(BTreeMap<Vec<usize>, &'a PMCFGRule<N, T, LogDomain<f32>>>);

// impl<'a, N, T> Derivation<'a, N, T>
// where
//     N: 'a,
//     T: 'a
// {
//     pub fn weight(&self) -> LogDomain<f32> {
//         let mut dweight = LogDomain::one();
//         let &Derivation(ref map) = self;

//         for &&PMCFGRule{ weight, .. } in map.values() {
//             dweight = dweight * weight;
//         }

//         dweight
//     }
// }

impl<'a, N: 'a + fmt::Display, T: 'a + fmt::Display> fmt::Display for Derivation<'a, N, T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut buffer = String::new();
        let &Derivation(ref tree) = self;

        for (pos, rule) in tree {
            if !pos.is_empty() {
                for _ in 0..(pos.len() - 1) {
                    buffer.push('\t');
                }
                buffer.push_str(" ⮡ ");
            }
            buffer.push_str(format!("{}\n", rule).as_str());
        }
        write!(f, "{}", buffer)
    }
}

/// A mutliple context-free grammar.
#[derive(Clone, Debug)]
pub struct MCFG<N, T, W> {
    rules: Vec<PMCFGRule<N, T, W>>,
    initial: N,
}

impl<N, T> From<PMCFG<N, T, LogDomain<f64>>> for MCFG<N, T, LogDomain<f32>> {
    fn from(grammar: PMCFG<N, T, LogDomain<f64>>) -> Self {
        let PMCFG { rules, mut initial, .. } = grammar;
        assert!(initial.len() == 1);

        MCFG {
            rules: rules
                .into_iter()
                .map(|r| match r {
                    PMCFGRule {
                        head,
                        tail,
                        composition,
                        weight,
                    } => PMCFGRule {
                        head,
                        tail,
                        composition,
                        weight: LogDomain::new(weight.value() as f32).unwrap(),
                    },
                })
                .collect(),
            initial: initial.remove(0),
        }
    }
}

/// The index of a bracket in cs representation.
/// Assumes integerized rules.
#[derive(PartialEq, Eq, Hash, Clone, Debug, PartialOrd, Ord, Serialize, Deserialize)]
pub enum BracketContent<T> {
    Terminal(T),
    Component(usize, usize),
    Variable(usize, usize, usize),
}

impl<T> Display for BracketContent<T>
where T: Display
{
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        match *self {
            BracketContent::Terminal(ref t) => write!(f, "_{{{}}}", t),
            BracketContent::Component(rule_id, comp_id) => write!(f, "_{}^{}", rule_id, comp_id),
            BracketContent::Variable(rule_id, i, j) => write!(f, "_{{{},{}}}^{}", rule_id, i, j)
        }
    }
}

/// A cs representation of a grammar contains a generator automaton and a Dyck language.
#[derive(Debug, Serialize, Deserialize)]
pub struct CSRepresentation<N, T, F>
where
    N: Eq + Hash + Clone,
    T: Ord + Hash + Clone,
    F: FilterAutomaton<T> + Serialize
{
    generator: Automaton<BracketFragment<T>>,
    dyck: Partition<BracketContent<T>>,

    // rules are integerized in CSRepresentation::new()
    // TODO: maybe integerize each `BracketContent` ?
    rules: HashIntegeriser<PMCFGRule<N, T, LogDomain<f32>>>,

    // saves all brackets σ with h(σ)(ε) ≠ 0
    filter: F,
}

fn get_partition<N, T, F>(rules: &HashIntegeriser<PMCFGRule<N, T, F>>) -> Vec<BTreeSet<BracketContent<T>>>
where
    T: Clone + Hash + Eq + Ord,
    N: Clone + Hash + Eq + Ord,
    F: Clone
{
    let mut cells = Vec::new();
    let mut terminals = HashSet::new();

    for rule_id in 0..(rules.size()) {
        let rule = rules.find_value(rule_id).unwrap();
        cells.push(
            (0..(rule.composition.composition.len())).map(|j| BracketContent::Component(rule_id, j)).collect()
        );
        cells.extend(
            {
                let mut succs = vec![BTreeSet::new();rule.tail.len()];
                for symbol in rule.composition.composition.iter().flat_map(|x| (*x).iter()) {
                    match *symbol {
                        VarT::T(ref t) => {
                            if !terminals.contains(t) {
                                terminals.insert(t.clone());
                            }
                        },
                        VarT::Var(i, j) => {
                            succs[i].insert(BracketContent::Variable(rule_id, i, j));
                        }
                    }
                }
                succs.into_iter()
            }
        )
    }
    cells.extend(
        terminals.into_iter().map(|t| vec![BracketContent::Terminal(t)].into_iter().collect())
    );

    cells
}

impl<N, T, F> CSRepresentation<N, T, F>
where
    N: Eq + Hash + Clone + ::std::fmt::Debug,
    T: Ord + Eq + Hash + Clone + ::std::fmt::Debug,
    F: Serialize + FilterAutomaton<T>
{
    ///
    pub fn new<S, G: Into<MCFG<N, T, LogDomain<f32>>>>(strat: S, grammar: G) -> Self
    where
        S: GeneratorAutomaton,
        N: Ord + Clone + Eq,
        T: Ord + Clone + Eq
    {
        let MCFG{ rules, initial } = grammar.into();
        let mut rulemap = HashIntegeriser::new();
        for rule in rules {
            rulemap.integerise(rule);
        }
        let cells = get_partition(&rulemap);
        let filter = F::new(&rulemap);

        CSRepresentation{
            generator: strat.convert(&rulemap, initial),
            dyck: Partition::new(cells).unwrap(),
            rules: rulemap,
            filter
        }
    }

    /// Produces a `CSGenerator` for a Chomsky-Schützenberger characterization and a `word`.
    pub fn generate(&self, word: &[T], n: usize) -> CSGenerator<T, N> {
        CSGenerator {
            candidates: self.generator.intersect(&self.filter.fsa(word, &self.generator)).generate(n),
            checker: MultipleDyckLanguage::new(&self.dyck),
            rules: &self.rules,
        }
    }
}

fn from_brackets<N, T>(
    rules: &HashIntegeriser<PMCFGRule<N, T, LogDomain<f32>>>,
    word: Vec<Bracket<BracketContent<T>>>,
) -> Derivation<N, T>
where
    N: Hash + Eq + Clone,
    T: Hash + Eq + Clone
{
    let mut tree = BTreeMap::new();
    let mut pos = Vec::new();

    for sigma in word {
        match sigma {
            Bracket::Open(BracketContent::Component(rule_id, _)) => {
                tree.insert(pos.clone(), rules.find_value(rule_id).unwrap());
            }
            Bracket::Open(BracketContent::Variable(_, i, _)) => {
                pos.push(i);
            }
            Bracket::Close(BracketContent::Variable(_, _, _)) => {
                pos.pop();
            }
            _ => (),
        }
    }

    Derivation(tree)
}

/// Iterates Dyck words that represent a derivation for a word according to
/// the Chomsky-Schützenberger characterization of an MCFG.
pub struct CSGenerator<'a, T: 'a + PartialEq + Hash + Clone + Eq + Ord + fmt::Debug, N: 'a + Hash + Eq> {
    candidates: generator::BatchGenerator<BracketFragment<T>>,

    checker: MultipleDyckLanguage<BracketContent<T>>,
    rules: &'a HashIntegeriser<PMCFGRule<N, T, LogDomain<f32>>>,
}

impl<'a, N, T> Iterator for CSGenerator<'a, T, N>
where
    T: PartialEq + Hash + Clone + Eq + Ord + fmt::Debug,
    N: Hash + Eq + Clone + fmt::Debug,
{
    type Item = Derivation<'a, N, T>;

    fn next(&mut self) -> Option<Derivation<'a, N, T>> {
        let &mut CSGenerator {
            ref mut candidates,
            ref checker,
            rules,
        } = self;

        let mut cans = 0;
        let mut dycks = 0;

        for candidate_batch in candidates {
            for (candidate_fragments, _) in candidate_batch {
                let candidate: Vec<Delta<T>> = BracketFragment::concat_owned(candidate_fragments);
                cans += 1;
                if dyck::recognize(&candidate){
                    eprintln!("{:?}", candidate);
                    dycks += 1;
                    if checker.recognize(&candidate) {
                        eprintln!("found after {} candidates, where {} were dyck words", cans, dycks);
                        return Some(from_brackets(rules, candidate));
                    }
                }
            }
        }
        None
    }
}

#[cfg(test)]
mod test {

    #[test]
    fn csrep() {
        use VarT;
        use PMCFGRule;
        use Composition;
        use super::CSRepresentation;
        use super::LogDomain;
        use super::Derivation;
        use super::MCFG;
        use cs_representation::generator_automaton::naive::{NaiveGeneratorAutomaton};
        use cs_representation::filter_automaton::naive::{NaiveFilterAutomaton};

        let grammar = MCFG {
            initial: "S",
            rules: vec![
                PMCFGRule {
                    head: "S",
                    tail: vec!["S", "S"],
                    composition: Composition {
                        composition: vec![vec![VarT::Var(0, 0), VarT::Var(1, 0)]],
                    },
                    weight: LogDomain::new(0.3f32).unwrap(),
                },
                PMCFGRule {
                    head: "S",
                    tail: vec![],
                    composition: Composition { composition: vec![vec![VarT::T('A')]] },
                    weight: LogDomain::new(0.7f32).unwrap(),
                },
            ],
        };
        let d1 = Derivation(vec![(vec![], &grammar.rules[1])].into_iter().collect());
        let d2 = Derivation(
            vec![
                (vec![], &grammar.rules[0]),
                (vec![0], &grammar.rules[1]),
                (vec![1], &grammar.rules[1]),
            ].into_iter()
                .collect(),
        );

        assert_eq!(
            CSRepresentation::<&str, char, NaiveFilterAutomaton<char>>::new(NaiveGeneratorAutomaton, grammar.clone())
                .generate(&['A'], 2usize)
                .next(),
            Some(d1)
        );
        assert_eq!(
            CSRepresentation::<&str, char, NaiveFilterAutomaton<char>>::new(NaiveGeneratorAutomaton, grammar.clone())
                .generate(&['A', 'A'], 2usize)
                .next(),
            Some(d2)
        );
    }
}
<commit_msg>avoid tsa run<commit_after>pub mod cli;

use serde::{Serialize};

use std::hash::Hash;
use std::fmt;
use integeriser::{Integeriser, HashIntegeriser};
use std::collections::{HashSet, BTreeMap, BTreeSet};

use pmcfg::{PMCFG, PMCFGRule, VarT};
use dyck::multiple::{Bracket, MultipleDyckLanguage};
use openfsa::fsa::{Automaton, generator};
use log_domain::LogDomain;
use dyck;

use util::partition::Partition;

pub mod bracket_fragment;
pub mod filter_automaton;
pub mod generator_automaton;
use cs_representation::bracket_fragment::BracketFragment;
use cs_representation::generator_automaton::{GeneratorAutomaton, Delta};
use cs_representation::filter_automaton::{FilterAutomaton};

use std::fmt::{Display, Formatter, Error};

/// A derivation tree of PMCFG rules.
#[derive(PartialEq, Debug)]
pub struct Derivation<'a, N: 'a, T: 'a>(BTreeMap<Vec<usize>, &'a PMCFGRule<N, T, LogDomain<f32>>>);

// impl<'a, N, T> Derivation<'a, N, T>
// where
//     N: 'a,
//     T: 'a
// {
//     pub fn weight(&self) -> LogDomain<f32> {
//         let mut dweight = LogDomain::one();
//         let &Derivation(ref map) = self;

//         for &&PMCFGRule{ weight, .. } in map.values() {
//             dweight = dweight * weight;
//         }

//         dweight
//     }
// }

impl<'a, N: 'a + fmt::Display, T: 'a + fmt::Display> fmt::Display for Derivation<'a, N, T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut buffer = String::new();
        let &Derivation(ref tree) = self;

        for (pos, rule) in tree {
            if !pos.is_empty() {
                for _ in 0..(pos.len() - 1) {
                    buffer.push('\t');
                }
                buffer.push_str(" ⮡ ");
            }
            buffer.push_str(format!("{}\n", rule).as_str());
        }
        write!(f, "{}", buffer)
    }
}

/// A mutliple context-free grammar.
#[derive(Clone, Debug)]
pub struct MCFG<N, T, W> {
    rules: Vec<PMCFGRule<N, T, W>>,
    initial: N,
}

impl<N, T> From<PMCFG<N, T, LogDomain<f64>>> for MCFG<N, T, LogDomain<f32>> {
    fn from(grammar: PMCFG<N, T, LogDomain<f64>>) -> Self {
        let PMCFG { rules, mut initial, .. } = grammar;
        assert!(initial.len() == 1);

        MCFG {
            rules: rules
                .into_iter()
                .map(|r| match r {
                    PMCFGRule {
                        head,
                        tail,
                        composition,
                        weight,
                    } => PMCFGRule {
                        head,
                        tail,
                        composition,
                        weight: LogDomain::new(weight.value() as f32).unwrap(),
                    },
                })
                .collect(),
            initial: initial.remove(0),
        }
    }
}

/// The index of a bracket in cs representation.
/// Assumes integerized rules.
#[derive(PartialEq, Eq, Hash, Clone, Debug, PartialOrd, Ord, Serialize, Deserialize)]
pub enum BracketContent<T> {
    Terminal(T),
    Component(usize, usize),
    Variable(usize, usize, usize),
}

impl<T> Display for BracketContent<T>
where T: Display
{
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        match *self {
            BracketContent::Terminal(ref t) => write!(f, "_{{{}}}", t),
            BracketContent::Component(rule_id, comp_id) => write!(f, "_{}^{}", rule_id, comp_id),
            BracketContent::Variable(rule_id, i, j) => write!(f, "_{{{},{}}}^{}", rule_id, i, j)
        }
    }
}

/// A cs representation of a grammar contains a generator automaton and a Dyck language.
#[derive(Debug, Serialize, Deserialize)]
pub struct CSRepresentation<N, T, F>
where
    N: Eq + Hash + Clone,
    T: Ord + Hash + Clone,
    F: FilterAutomaton<T> + Serialize
{
    generator: Automaton<BracketFragment<T>>,
    dyck: Partition<BracketContent<T>>,

    // rules are integerized in CSRepresentation::new()
    // TODO: maybe integerize each `BracketContent` ?
    rules: HashIntegeriser<PMCFGRule<N, T, LogDomain<f32>>>,

    // saves all brackets σ with h(σ)(ε) ≠ 0
    filter: F,
}

fn get_partition<N, T, F>(rules: &HashIntegeriser<PMCFGRule<N, T, F>>) -> Vec<BTreeSet<BracketContent<T>>>
where
    T: Clone + Hash + Eq + Ord,
    N: Clone + Hash + Eq + Ord,
    F: Clone
{
    let mut cells = Vec::new();
    let mut terminals = HashSet::new();

    for rule_id in 0..(rules.size()) {
        let rule = rules.find_value(rule_id).unwrap();
        cells.push(
            (0..(rule.composition.composition.len())).map(|j| BracketContent::Component(rule_id, j)).collect()
        );
        cells.extend(
            {
                let mut succs = vec![BTreeSet::new();rule.tail.len()];
                for symbol in rule.composition.composition.iter().flat_map(|x| (*x).iter()) {
                    match *symbol {
                        VarT::T(ref t) => {
                            if !terminals.contains(t) {
                                terminals.insert(t.clone());
                            }
                        },
                        VarT::Var(i, j) => {
                            succs[i].insert(BracketContent::Variable(rule_id, i, j));
                        }
                    }
                }
                succs.into_iter()
            }
        )
    }
    cells.extend(
        terminals.into_iter().map(|t| vec![BracketContent::Terminal(t)].into_iter().collect())
    );

    cells
}

impl<N, T, F> CSRepresentation<N, T, F>
where
    N: Eq + Hash + Clone + ::std::fmt::Debug,
    T: Ord + Eq + Hash + Clone + ::std::fmt::Debug,
    F: Serialize + FilterAutomaton<T>
{
    ///
    pub fn new<S, G: Into<MCFG<N, T, LogDomain<f32>>>>(strat: S, grammar: G) -> Self
    where
        S: GeneratorAutomaton,
        N: Ord + Clone + Eq,
        T: Ord + Clone + Eq
    {
        let MCFG{ rules, initial } = grammar.into();
        let mut rulemap = HashIntegeriser::new();
        for rule in rules {
            rulemap.integerise(rule);
        }
        let cells = get_partition(&rulemap);
        let filter = F::new(&rulemap);

        CSRepresentation{
            generator: strat.convert(&rulemap, initial),
            dyck: Partition::new(cells).unwrap(),
            rules: rulemap,
            filter
        }
    }

    /// Produces a `CSGenerator` for a Chomsky-Schützenberger characterization and a `word`.
    pub fn generate(&self, word: &[T], n: usize) -> CSGenerator<T, N> {
        CSGenerator {
            candidates: self.generator.intersect(&self.filter.fsa(word, &self.generator)).generate(n),
            checker: MultipleDyckLanguage::new(&self.dyck),
            rules: &self.rules,
        }
    }
}

fn from_brackets<N, T>(
    rules: &HashIntegeriser<PMCFGRule<N, T, LogDomain<f32>>>,
    word: Vec<Bracket<BracketContent<T>>>,
) -> Option<Derivation<N, T>>
where
    N: Hash + Eq + Clone,
    T: Hash + Eq + Clone
{
    let mut tree = BTreeMap::new();
    let mut pos = Vec::new();

    for sigma in word {
        match sigma {
            Bracket::Open(BracketContent::Component(rule_id, _)) => {
                let rule_at_pos = tree.entry(pos.clone()).or_insert(rule_id);
                if rule_at_pos != &rule_id {
                    return None;
                }
            }
            Bracket::Open(BracketContent::Variable(rule_id, i, _)) => {
                pos.push(i);
            }
            Bracket::Close(BracketContent::Variable(_, _, _)) => {
                pos.pop();
            }
            _ => (),
        }
    }

    Some(Derivation(tree.into_iter().map(|(pos, i)| (pos, rules.find_value(i).unwrap())).collect()))
}

/// Iterates Dyck words that represent a derivation for a word according to
/// the Chomsky-Schützenberger characterization of an MCFG.
pub struct CSGenerator<'a, T: 'a + PartialEq + Hash + Clone + Eq + Ord + fmt::Debug, N: 'a + Hash + Eq> {
    candidates: generator::BatchGenerator<BracketFragment<T>>,

    checker: MultipleDyckLanguage<BracketContent<T>>,
    rules: &'a HashIntegeriser<PMCFGRule<N, T, LogDomain<f32>>>,
}

impl<'a, N, T> Iterator for CSGenerator<'a, T, N>
where
    T: PartialEq + Hash + Clone + Eq + Ord + fmt::Debug,
    N: Hash + Eq + Clone + fmt::Debug,
{
    type Item = Derivation<'a, N, T>;

    fn next(&mut self) -> Option<Derivation<'a, N, T>> {
        let &mut CSGenerator {
            ref mut candidates,
            ref checker,
            rules,
        } = self;

        let mut cans = 0;
        let mut dycks = 0;

        for candidate_batch in candidates {
            eprintln!("starting batch");
            for (candidate_fragments, _) in candidate_batch {
                let candidate: Vec<Delta<T>> = BracketFragment::concat_owned(candidate_fragments);
                cans += 1;
                if dyck::recognize(&candidate){
                    dycks += 1;
                    if let Some(derivation) = from_brackets(rules, candidate) {
                        eprintln!("found after {} candidates, where {} were dyck words", cans, dycks);
                        return Some(derivation);
                    }
                }
            }
            eprintln!("finshed batch");
        }
        None
    }
}

#[cfg(test)]
mod test {

    #[test]
    fn csrep() {
        use VarT;
        use PMCFGRule;
        use Composition;
        use super::CSRepresentation;
        use super::LogDomain;
        use super::Derivation;
        use super::MCFG;
        use cs_representation::generator_automaton::naive::{NaiveGeneratorAutomaton};
        use cs_representation::filter_automaton::naive::{NaiveFilterAutomaton};

        let grammar = MCFG {
            initial: "S",
            rules: vec![
                PMCFGRule {
                    head: "S",
                    tail: vec!["S", "S"],
                    composition: Composition {
                        composition: vec![vec![VarT::Var(0, 0), VarT::Var(1, 0)]],
                    },
                    weight: LogDomain::new(0.3f32).unwrap(),
                },
                PMCFGRule {
                    head: "S",
                    tail: vec![],
                    composition: Composition { composition: vec![vec![VarT::T('A')]] },
                    weight: LogDomain::new(0.7f32).unwrap(),
                },
            ],
        };
        let d1 = Derivation(vec![(vec![], &grammar.rules[1])].into_iter().collect());
        let d2 = Derivation(
            vec![
                (vec![], &grammar.rules[0]),
                (vec![0], &grammar.rules[1]),
                (vec![1], &grammar.rules[1]),
            ].into_iter()
                .collect(),
        );

        assert_eq!(
            CSRepresentation::<&str, char, NaiveFilterAutomaton<char>>::new(NaiveGeneratorAutomaton, grammar.clone())
                .generate(&['A'], 2usize)
                .next(),
            Some(d1)
        );
        assert_eq!(
            CSRepresentation::<&str, char, NaiveFilterAutomaton<char>>::new(NaiveGeneratorAutomaton, grammar.clone())
                .generate(&['A', 'A'], 2usize)
                .next(),
            Some(d2)
        );
    }
}
<|endoftext|>


<commit_before>//
// FlightVars
// Copyright (c) 2015, 2016 Alvaro Polo
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::io;
use std::net;

use comm::*;

pub struct TcpTransport {
    listener: net::TcpListener
}

impl TcpTransport {
    fn bind<A: net::ToSocketAddrs>(addr: A) -> io::Result<TcpTransport> {
        Ok(TcpTransport { listener: try!(net::TcpListener::bind(addr)) })
    }
}

impl Transport for TcpTransport {
    type Input = net::TcpStream;
    type Output = net::TcpStream;
    type Listener = net::TcpListener;

    fn listener(&mut self) -> &mut net::TcpListener {
        &mut self.listener
    }
}

impl Listen<net::TcpStream, net::TcpStream> for net::TcpListener {
    fn listen(&mut self) -> io::Result<(net::TcpStream, net::TcpStream)> {
        let (conn, _) = try!(self.accept());
        let input = try!(conn.try_clone());
        let output = conn;
        Ok((input, output))
    }
}

impl ShutdownInterruption for net::TcpListener {
    type Int = TcpInterruptor;

    fn shutdown_interruption(&mut self) -> TcpInterruptor {
        TcpInterruptor::from_listener(self)
    }
}

impl ShutdownInterruption for net::TcpStream {
    type Int = TcpInterruptor;

    fn shutdown_interruption(&mut self) -> TcpInterruptor {
        TcpInterruptor::from_stream(self)
    }
}

#[cfg(unix)]
pub mod unix {
    use std::net;
    use std::os::unix::io::{AsRawFd, RawFd};

    use libc;

    use comm::*;

    pub struct TcpInterruptor {
        fd: RawFd
    }

    impl TcpInterruptor {
        pub fn from_listener(listener: &net::TcpListener) -> TcpInterruptor {
            TcpInterruptor { fd: listener.as_raw_fd() }
        }

        pub fn from_stream(stream: &net::TcpStream) -> TcpInterruptor {
            TcpInterruptor { fd: stream.as_raw_fd() }
        }
    }

    impl Interrupt for TcpInterruptor {
        fn interrupt(self) {
            unsafe { libc::close(self.fd); }
        }
    }
}

#[cfg(unix)]
pub use self::unix::*;

#[cfg(test)]
mod tests {
    use std::io;
    use std::io::{BufRead, Write};
    use std::net;
    use std::thread;
    use std::time;

    use comm::{ShutdownHandle, Transport};

    #[test]
    fn should_wait_conn() {
        let child = thread::spawn(|| {
            let mut listener = net::TcpListener::bind("127.0.0.1:1234").unwrap();
            let (r, mut w) = listener.wait_conn().unwrap();
            let mut input = io::BufReader::new(r);
            let mut line = String::new();
            input.read_line(&mut line).unwrap();
            assert_eq!(line, "Hello server\n");
            w.write_all(b"Hello client\n").unwrap();
        });
        thread::sleep(time::Duration::from_millis(25));
        let mut conn = net::TcpStream::connect("127.0.0.1:1234").unwrap();
        conn.write_all(b"Hello server\n").unwrap();
        let mut input = io::BufReader::new(conn);
        let mut line = String::new();
        input.read_line(&mut line).unwrap();
        assert_eq!(line, "Hello client\n");
        child.join().unwrap();
    }

    #[test]
    fn should_close_from_shutdown_handle() {
        let mut listener = net::TcpListener::bind("127.0.0.1:1235").unwrap();
        let handle = listener.shutdown_handle();
        let child = thread::spawn(move || {
            assert_eq!(listener.wait_conn().unwrap_err().kind(), io::ErrorKind::ConnectionAborted);
        });
        thread::sleep(time::Duration::from_millis(25));
        handle.shutdown();
        child.join().unwrap();
    }
}
<commit_msg>Fix broken tests in `comm::tcp`<commit_after>//
// FlightVars
// Copyright (c) 2015, 2016 Alvaro Polo
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::io;
use std::net;

use comm::*;

pub struct TcpTransport {
    listener: net::TcpListener
}

impl TcpTransport {
    pub fn bind<A: net::ToSocketAddrs>(addr: A) -> io::Result<TcpTransport> {
        Ok(TcpTransport { listener: try!(net::TcpListener::bind(addr)) })
    }
}

impl Transport for TcpTransport {
    type Input = net::TcpStream;
    type Output = net::TcpStream;
    type Listener = net::TcpListener;

    fn listener(&mut self) -> &mut net::TcpListener {
        &mut self.listener
    }
}

impl Listen<net::TcpStream, net::TcpStream> for net::TcpListener {
    fn listen(&mut self) -> io::Result<(net::TcpStream, net::TcpStream)> {
        let (conn, _) = try!(self.accept());
        let input = try!(conn.try_clone());
        let output = conn;
        Ok((input, output))
    }
}

impl ShutdownInterruption for net::TcpListener {
    type Int = TcpInterruptor;

    fn shutdown_interruption(&mut self) -> TcpInterruptor {
        TcpInterruptor::from_listener(self)
    }
}

impl ShutdownInterruption for net::TcpStream {
    type Int = TcpInterruptor;

    fn shutdown_interruption(&mut self) -> TcpInterruptor {
        TcpInterruptor::from_stream(self)
    }
}

#[cfg(unix)]
pub mod unix {
    use std::net;
    use std::os::unix::io::{AsRawFd, RawFd};

    use libc;

    use comm::*;

    pub struct TcpInterruptor {
        fd: RawFd
    }

    impl TcpInterruptor {
        pub fn from_listener(listener: &net::TcpListener) -> TcpInterruptor {
            TcpInterruptor { fd: listener.as_raw_fd() }
        }

        pub fn from_stream(stream: &net::TcpStream) -> TcpInterruptor {
            TcpInterruptor { fd: stream.as_raw_fd() }
        }
    }

    impl Interrupt for TcpInterruptor {
        fn interrupt(self) {
            unsafe { libc::close(self.fd); }
        }
    }
}

#[cfg(unix)]
pub use self::unix::*;

#[cfg(test)]
mod tests {
    use std::io;
    use std::io::{BufRead, Write};
    use std::net;
    use std::thread;
    use std::time;

    use comm::*;

    #[test]
    fn should_wait_conn() {
        let child = thread::spawn(|| {
            let mut listener = net::TcpListener::bind("127.0.0.1:1234").unwrap();
            let (r, mut w) = listener.listen().unwrap();
            let mut input = io::BufReader::new(r);
            let mut line = String::new();
            input.read_line(&mut line).unwrap();
            assert_eq!(line, "Hello server\n");
            w.write_all(b"Hello client\n").unwrap();
        });
        thread::sleep(time::Duration::from_millis(25));
        let mut conn = net::TcpStream::connect("127.0.0.1:1234").unwrap();
        conn.write_all(b"Hello server\n").unwrap();
        let mut input = io::BufReader::new(conn);
        let mut line = String::new();
        input.read_line(&mut line).unwrap();
        assert_eq!(line, "Hello client\n");
        child.join().unwrap();
    }

    #[test]
    fn should_close_from_shutdown_handle() {
        let mut listener = net::TcpListener::bind("127.0.0.1:1235").unwrap();
        let interruption = listener.shutdown_interruption();
        let child = thread::spawn(move || {
            assert_eq!(listener.listen().unwrap_err().kind(), io::ErrorKind::ConnectionAborted);
        });
        thread::sleep(time::Duration::from_millis(25));
        interruption.interrupt();
        child.join().unwrap();
    }
}
<|endoftext|>


<commit_before>extern crate hyper;

use std::sync::mpsc::{Receiver, channel};
use std::thread;
use args::Config;
use run_info::RunInfo;
use misc::split_number;
use tokio_core::reactor::Core;
use chrono::{Duration, Local, DateTime};
use hyper::{Client, Url};
use std::str::FromStr;
use hyper::client::HttpConnector;

use futures::{Future, stream};
use futures::future::{ok, loop_fn, Loop};
use futures::Stream;
use futures_utils::stopwatch;
use request_result::RequestResult;
use errors::*;

/// Delegates and manages all the work.
pub struct Boss {
    /// Number of requests successfully completed
    pub requests_completed: usize,
    /// Number of active connections
    pub connections: usize,
    num_threads: usize,
}

impl Boss {
    pub fn new(num_threads: usize) -> Self {
        Boss {
            requests_completed: 0,
            connections: 0,
            num_threads: num_threads,
        }
    }

    pub fn start_workforce(&self, config: &Config) -> Result<RunInfo> {
        let (tx, rx) = channel();
        let desired_connections_per_worker_iter = split_number(config.num_connections,
                                                               self.num_threads);
        // start num_threads workers
        for desired_connections_per_worker in desired_connections_per_worker_iter {
            let tx = tx.clone();
            let url = config.url.clone();
            let duration = config.duration;
            let timeout = config.timeout.to_std()?;
            let hyper_url = Url::from_str(&url)?;
            let start_time = Local::now();
            let wanted_end_time = start_time + duration;
            thread::spawn(move || {
                // TODO: how to use error_chain on tokio-core?
                let mut core = Core::new().expect("Failed to create Tokio core");
                let hyper_client = hyper::Client::configure()
                    .keep_alive_timeout(Some(timeout))
                    .build(&core.handle());

                let iterator = (0..desired_connections_per_worker).map(|_| {
                    create_looping_worker(duration,
                                          hyper_url.clone(),
                                          &hyper_client,
                                          wanted_end_time)
                });
                let future = stream::futures_unordered(iterator)
                    .fold(RunInfo::new(duration), |mut runinfo_acc, runinfo| {
                        runinfo_acc.merge(&runinfo);
                        ok(runinfo_acc)
                    });
                // TODO: how to use error_chain with tokio?
                let res = core.run(future).expect("Failed to run Tokio Core");

                tx.send(res)
            });
        }

        Ok(self.collect_workers(rx, config.duration))
    }

    /// Collects information from all workers
    fn collect_workers(&self, rx: Receiver<RunInfo>, run_duration: Duration) -> RunInfo {
        rx.iter()
            .take(self.num_threads)
            .fold(RunInfo::new(run_duration), |mut runinfo_acc, runinfo| {
                runinfo_acc.merge(&runinfo);
                runinfo_acc
            })
    }
}

fn create_looping_worker<'a>(duration: Duration,
                             url: Url,
                             hyper_client: &'a Client<HttpConnector>,
                             wanted_end_time: DateTime<Local>)
                             -> impl Future<Item = RunInfo, Error = ()> + 'a {
    let runinfo = RunInfo::new(duration);
    loop_fn((runinfo), move |mut runinfo| {
        let url = url.clone();
        send_request(url, hyper_client).then(move |res| {
            match res {
                // on request success
                Ok(request_res) => {
                    // update histogram
                    request_res.latency
                        .num_nanoseconds()
                        .map(|x| runinfo.histogram.increment(x as u64).ok());
                    runinfo.requests_completed += 1;
                    ok(loop_iter(runinfo, wanted_end_time))
                }
                // on request failure
                Err(_) => {
                    runinfo.num_failed_requests += 1;
                    ok(loop_iter(runinfo, wanted_end_time))
                }
            }
        })
    })
}

/// Determines what the next loop iteration should be; Continue or break.
fn loop_iter(state: RunInfo,
             wanted_end_time: DateTime<Local>)
             -> Loop<RunInfo, RunInfo> {
    let now_time = Local::now();
    if now_time < wanted_end_time {
        Loop::Continue(state)
    } else {
        Loop::Break(state)
    }
}

fn send_request<C>(url: Url,
                   hyper_client: &Client<C>)
                   -> impl Future<Item = RequestResult, Error = hyper::Error>
    where C: hyper::client::Connect
{
    stopwatch(hyper_client.get(url)).and_then(|(_, latency)| {
        let request_result = RequestResult::new(latency);
        Ok(request_result)
    })
}
<commit_msg>send result to channel<commit_after>extern crate hyper;

use std::sync::mpsc::{Receiver, channel};
use std::thread;
use args::Config;
use run_info::RunInfo;
use misc::split_number;
use tokio_core::reactor::Core;
use chrono::{Duration, Local, DateTime};
use hyper::{Client, Url};
use std::str::FromStr;
use hyper::client::HttpConnector;

use futures::{Future, stream};
use futures::future::{ok, loop_fn, Loop};
use futures::Stream;
use futures_utils::stopwatch;
use request_result::RequestResult;
use errors::*;

/// Delegates and manages all the work.
pub struct Boss {
    /// Number of requests successfully completed
    pub requests_completed: usize,
    /// Number of active connections
    pub connections: usize,
    num_threads: usize,
}

impl Boss {
    pub fn new(num_threads: usize) -> Self {
        Boss {
            requests_completed: 0,
            connections: 0,
            num_threads: num_threads,
        }
    }

    pub fn start_workforce(&self, config: &Config) -> Result<RunInfo> {
        let (tx, rx) = channel();
        let desired_connections_per_worker_iter = split_number(config.num_connections,
                                                               self.num_threads);
        // start num_threads workers
        for desired_connections_per_worker in desired_connections_per_worker_iter {
            let tx = tx.clone();
            let url = config.url.clone();
            let duration = config.duration;
            let timeout = config.timeout.to_std()?;
            let hyper_url = Url::from_str(&url)?;
            let start_time = Local::now();
            let wanted_end_time = start_time + duration;
            thread::spawn(move || {
                // TODO: how to use error_chain on tokio-core?
                let mut core = Core::new().expect("Failed to create Tokio core");
                let hyper_client = hyper::Client::configure()
                    .keep_alive_timeout(Some(timeout))
                    .build(&core.handle());

                // // create infinite stream of requests by unfolding
                // let request_stream = stream::unfold(0u32, |state| {
                //     let fut = send_request(hyper_url.clone(), &hyper_client);
                //     // create a wrapper around the future, so we can pass the state.
                //     // in this case, we do not need the state, so just pass it on unmodified
                //     let fut_wrapper = fut.map(move |request_result| (request_result, state));
                //     Some(fut_wrapper)
                // });

                // TODO: make infinite
                let iterator = (0..desired_connections_per_worker).map(|_| {
                    create_looping_worker(duration,
                                          hyper_url.clone(),
                                          &hyper_client,
                                          wanted_end_time)
                });
                // create stream of requests
                let request_stream = stream::futures_unordered(iterator);
                let future = request_stream
                    .for_each(|request_result| {
                        let _ = tx.send(request_result);
                        Ok(())
                    });
                // TODO: how to use error_chain with tokio?
                let _ = core.run(future).expect("Failed to run Tokio Core");
            });
        }

        Ok(self.collect_workers(rx, config.duration))
    }

    /// Collects information from all workers
    fn collect_workers(&self, rx: Receiver<RunInfo>, run_duration: Duration) -> RunInfo {
        let start_time = Local::now();
        let wanted_end_time = start_time + run_duration;
        rx.iter()
            .take(100)
            // .take_while(|_| Local::now() < wanted_end_time)
            .fold(RunInfo::new(run_duration), |mut runinfo_acc, runinfo| {
                runinfo_acc.merge(&runinfo);
                runinfo_acc
            })
    }
}

fn create_looping_worker<'a>(duration: Duration,
                             url: Url,
                             hyper_client: &'a Client<HttpConnector>,
                             wanted_end_time: DateTime<Local>)
                             -> impl Future<Item = RunInfo, Error = ()> + 'a {
    let runinfo = RunInfo::new(duration);
    loop_fn((runinfo), move |mut runinfo| {
        let url = url.clone();
        send_request(url, hyper_client).then(move |res| {
            match res {
                // on request success
                Ok(request_res) => {
                    // update histogram
                    request_res.latency
                        .num_nanoseconds()
                        .map(|x| runinfo.histogram.increment(x as u64).ok());
                    runinfo.requests_completed += 1;
                    ok(loop_iter(runinfo, wanted_end_time))
                }
                // on request failure
                Err(_) => {
                    runinfo.num_failed_requests += 1;
                    ok(loop_iter(runinfo, wanted_end_time))
                }
            }
        })
    })
}

/// Determines what the next loop iteration should be; Continue or break.
fn loop_iter(state: RunInfo,
             wanted_end_time: DateTime<Local>)
             -> Loop<RunInfo, RunInfo> {
    let now_time = Local::now();
    if now_time < wanted_end_time {
        Loop::Continue(state)
    } else {
        Loop::Break(state)
    }
}

fn send_request<C>(url: Url,
                   hyper_client: &Client<C>)
                   -> impl Future<Item = RequestResult, Error = hyper::Error>
    where C: hyper::client::Connect
{
    stopwatch(hyper_client.get(url)).and_then(|(_, latency)| {
        let request_result = RequestResult::new(latency);
        Ok(request_result)
    })
}
<|endoftext|>