<commit_before>use crate::rparser::*;
use der_parser::der::der_read_element_header;
use kerberos_parser::krb5_parser;
use kerberos_parser::krb5::{EncryptionType,PAType};

pub struct KerberosUDPBuilder {}
impl RBuilder for KerberosUDPBuilder {
    fn new(&self) -> Box<RParser> { Box::new(KerberosParserUDP::new(b"Kerberos/UDP")) }
    fn probe(&self, i:&[u8]) -> bool { kerberos_probe_udp(i) }
}

pub struct KerberosParserUDP<'a> {
    _name: Option<&'a[u8]>,
    req_id: u8,
}

pub fn to_hex_string(bytes: &[u8]) -> String {
    let mut s = String::new();
    for &b in bytes {
        s.push_str(&format!("{:02X}", b));
    }
    s
}

impl<'a> RParser for KerberosParserUDP<'a> {
    fn parse(&mut self, i: &[u8], _direction: u8) -> u32 {
        match der_read_element_header(i) {
            Ok((_rem,hdr)) => {
                // Kerberos messages start with an APPLICATION header
                if hdr.class != 0b01 { return 1; }
                debug!("hdr: {:?}", hdr);
                match hdr.tag.0 {
                    10 => {
                        self.req_id = hdr.tag.0;
                        let res = krb5_parser::parse_as_req(i);
                        debug!("AS-REQ: {:?}", res);
                        if let Ok((_,kdc_req)) = res {
                            debug!("AS-REQ cname: {:?}", kdc_req.req_body.cname);
                            debug!("AS-REQ realm: {:?}", kdc_req.req_body.realm);
                            debug!("AS-REQ sname: {:?}", kdc_req.req_body.sname);
                            debug!("AS-REQ addrs: {:?}", kdc_req.req_body.addresses);
                            match kdc_req.req_body.kdc_options.as_slice() {
                                Ok(s)  => debug!("AS-REQ kdc_options: {}", to_hex_string(s)),
                                Err(_) => (),
                            }
                        };
                    },
                    11 => {
                        let res = krb5_parser::parse_as_rep(i);
                        debug!("AS-REP: {:?}", res);
                        if let Ok((_,kdc_rep)) = res {
                            debug!("AS-REP cname: {:?}", kdc_rep.cname);
                            debug!("AS-REP crealm: {:?}", kdc_rep.crealm);
                            debug!("AS-REP ticket.sname: {:?}", kdc_rep.ticket.sname);
                            debug!("AS-REP ticket.realm: {:?}", kdc_rep.ticket.realm);
                            debug!("AS-REP enc-part.etype: {:?}", kdc_rep.enc_part.etype);
                            if test_weak_crypto(kdc_rep.enc_part.etype) {
                                debug!("AS-REP weak crypto");
                            }
                        };
                        self.req_id = 0;
                    },
                    12 => {
                        self.req_id = hdr.tag.0;
                        let res = krb5_parser::parse_tgs_req(i);
                        debug!("TGS-REQ: {:?}", res);
                        if let Ok((_,kdc_req)) = res {
                            debug!("TGS-REQ cname: {:?}", kdc_req.req_body.cname);
                            debug!("TGS-REQ realm: {:?}", kdc_req.req_body.realm);
                            debug!("TGS-REQ sname: {:?}", kdc_req.req_body.sname);
                            for ref padata in kdc_req.padata.iter() {
                                // SCLogInfo!("TGS-REQ padata: {:?}", padata);
                                if padata.padata_type == PAType::PA_TGS_REQ {
                                    match krb5_parser::parse_ap_req(padata.padata_value) {
                                        Ok((_,ap_req)) => {
                                            // SCLogInfo!("parse_ap_req: {:?}",ap_req);
                                            match ap_req.ap_options.as_slice() {
                                                Ok(s)  => debug!("TGS-REQ AP-REQ ap_options: {}", to_hex_string(s)),
                                                Err(_) => (),
                                            }
                                            debug!("TGS-REQ AP-REQ ticket.sname: {:?}", ap_req.ticket.sname);
                                            debug!("TGS-REQ AP-REQ ticket.realm: {:?}", ap_req.ticket.realm);
                                        }
                                        x => debug!("ERROR parse_ap_req: {:?}",x),
                                    }
                                }
                            }
                        };
                    },
                    13 => {
                        let res = krb5_parser::parse_tgs_rep(i);
                        debug!("TGS-REP: {:?}", res);
                        if let Ok((_,kdc_rep)) = res {
                            debug!("TGS-REP cname: {:?}", kdc_rep.cname);
                            debug!("TGS-REP ticket sname: {:?}", kdc_rep.ticket.sname);
                            debug!("TGS-REP ticket.realm: {:?}", kdc_rep.ticket.realm);
                            if test_weak_crypto(kdc_rep.enc_part.etype) {
                                debug!("AS-REP weak crypto");
                            }
                        };
                        self.req_id = 0;
                    },
                    14 => {
                        self.req_id = hdr.tag.0;
                        let res = krb5_parser::parse_ap_req(i);
                        debug!("AP-REQ: {:?}", res);
                        if let Ok((_,ap_rep)) = res {
                            debug!("AP-REQ ticket sname: {:?}", ap_rep.ticket.sname);
                            debug!("AP-REQ ticket.realm: {:?}", ap_rep.ticket.realm);
                        };
                    },
                    15 => {
                        let res = krb5_parser::parse_ap_rep(i);
                        debug!("AP-REP {:?}", res);
                        self.req_id = 0;
                    },
                    30 => {
                        let res = krb5_parser::parse_krb_error(i);
                        debug!("KRB-ERROR: {:?}", res);
                        debug!("KRB-ERROR: failed request: {:?}", self.req_id);
                        if let Ok((_,error)) = res {
                            debug!("KRB-ERROR: cname: {:?}", error.cname);
                            debug!("KRB-ERROR: realm: {:?}", error.realm);
                            debug!("KRB-ERROR: sname: {:?}", error.sname);
                            debug!("KRB-ERROR: error_code: {:?}", error.error_code);
                        };
                        self.req_id = 0;
                    },
                    _ => debug!("unknown/unsupported tag {}", hdr.tag),
                }
                R_STATUS_OK
            },
            _ => R_STATUS_FAIL
        }
    }
}

impl<'a> KerberosParserUDP<'a> {
    pub fn new(name: &'a[u8]) -> KerberosParserUDP<'a> {
        KerberosParserUDP{
            _name: Some(name),
            req_id: 0,
        }
    }
}

/// Return true if Kerberos `EncryptionType` is weak
pub fn test_weak_crypto(alg:EncryptionType) -> bool {
    match alg {
        EncryptionType::AES128_CTS_HMAC_SHA1_96 |
        EncryptionType::AES256_CTS_HMAC_SHA1_96 |
        EncryptionType::AES128_CTS_HMAC_SHA256_128 |
        EncryptionType::AES256_CTS_HMAC_SHA384_192 |
        EncryptionType::CAMELLIA128_CTS_CMAC |
        EncryptionType::CAMELLIA256_CTS_CMAC => false,
        _ => true, // all other ciphers are weak or deprecated
    }
}

pub fn kerberos_probe_udp(i: &[u8]) -> bool {
    if i.len() <= 10 { return false; }
    match der_read_element_header(i) {
        Ok((rem,hdr)) => {
            // Kerberos messages start with an APPLICATION header
            if hdr.class != 0b01 { return false; }
            // Tag number should be <= 30
            if hdr.tag.0 >= 30 { return false; }
            // Kerberos messages contain sequences
            if rem.is_empty() || rem[0] != 0x30 { return false; }
            // Check kerberos version
            // debug!("hdr: {:?}", hdr);
            if let Ok((rem,_hdr)) = der_read_element_header(rem) {
                if rem.len() > 5 {
                    match (rem[2],rem[3],rem[4]) {
                        // Encoding of DER integer 5 (version)
                        (2,1,5) => { return true; },
                        _       => (),
                    }
                }
            }
            false
        },
        _ => false,
    }
}
<commit_msg>Update to der-parser 2.1.0 (tag is 32 bits)<commit_after>use crate::rparser::*;
use der_parser::der::der_read_element_header;
use kerberos_parser::krb5_parser;
use kerberos_parser::krb5::{EncryptionType,PAType};

pub struct KerberosUDPBuilder {}
impl RBuilder for KerberosUDPBuilder {
    fn new(&self) -> Box<RParser> { Box::new(KerberosParserUDP::new(b"Kerberos/UDP")) }
    fn probe(&self, i:&[u8]) -> bool { kerberos_probe_udp(i) }
}

pub struct KerberosParserUDP<'a> {
    _name: Option<&'a[u8]>,
    req_id: u32,
}

pub fn to_hex_string(bytes: &[u8]) -> String {
    let mut s = String::new();
    for &b in bytes {
        s.push_str(&format!("{:02X}", b));
    }
    s
}

impl<'a> RParser for KerberosParserUDP<'a> {
    fn parse(&mut self, i: &[u8], _direction: u8) -> u32 {
        match der_read_element_header(i) {
            Ok((_rem,hdr)) => {
                // Kerberos messages start with an APPLICATION header
                if hdr.class != 0b01 { return 1; }
                debug!("hdr: {:?}", hdr);
                match hdr.tag.0 {
                    10 => {
                        self.req_id = hdr.tag.0;
                        let res = krb5_parser::parse_as_req(i);
                        debug!("AS-REQ: {:?}", res);
                        if let Ok((_,kdc_req)) = res {
                            debug!("AS-REQ cname: {:?}", kdc_req.req_body.cname);
                            debug!("AS-REQ realm: {:?}", kdc_req.req_body.realm);
                            debug!("AS-REQ sname: {:?}", kdc_req.req_body.sname);
                            debug!("AS-REQ addrs: {:?}", kdc_req.req_body.addresses);
                            match kdc_req.req_body.kdc_options.as_slice() {
                                Ok(s)  => debug!("AS-REQ kdc_options: {}", to_hex_string(s)),
                                Err(_) => (),
                            }
                        };
                    },
                    11 => {
                        let res = krb5_parser::parse_as_rep(i);
                        debug!("AS-REP: {:?}", res);
                        if let Ok((_,kdc_rep)) = res {
                            debug!("AS-REP cname: {:?}", kdc_rep.cname);
                            debug!("AS-REP crealm: {:?}", kdc_rep.crealm);
                            debug!("AS-REP ticket.sname: {:?}", kdc_rep.ticket.sname);
                            debug!("AS-REP ticket.realm: {:?}", kdc_rep.ticket.realm);
                            debug!("AS-REP enc-part.etype: {:?}", kdc_rep.enc_part.etype);
                            if test_weak_crypto(kdc_rep.enc_part.etype) {
                                debug!("AS-REP weak crypto");
                            }
                        };
                        self.req_id = 0;
                    },
                    12 => {
                        self.req_id = hdr.tag.0;
                        let res = krb5_parser::parse_tgs_req(i);
                        debug!("TGS-REQ: {:?}", res);
                        if let Ok((_,kdc_req)) = res {
                            debug!("TGS-REQ cname: {:?}", kdc_req.req_body.cname);
                            debug!("TGS-REQ realm: {:?}", kdc_req.req_body.realm);
                            debug!("TGS-REQ sname: {:?}", kdc_req.req_body.sname);
                            for ref padata in kdc_req.padata.iter() {
                                // SCLogInfo!("TGS-REQ padata: {:?}", padata);
                                if padata.padata_type == PAType::PA_TGS_REQ {
                                    match krb5_parser::parse_ap_req(padata.padata_value) {
                                        Ok((_,ap_req)) => {
                                            // SCLogInfo!("parse_ap_req: {:?}",ap_req);
                                            match ap_req.ap_options.as_slice() {
                                                Ok(s)  => debug!("TGS-REQ AP-REQ ap_options: {}", to_hex_string(s)),
                                                Err(_) => (),
                                            }
                                            debug!("TGS-REQ AP-REQ ticket.sname: {:?}", ap_req.ticket.sname);
                                            debug!("TGS-REQ AP-REQ ticket.realm: {:?}", ap_req.ticket.realm);
                                        }
                                        x => debug!("ERROR parse_ap_req: {:?}",x),
                                    }
                                }
                            }
                        };
                    },
                    13 => {
                        let res = krb5_parser::parse_tgs_rep(i);
                        debug!("TGS-REP: {:?}", res);
                        if let Ok((_,kdc_rep)) = res {
                            debug!("TGS-REP cname: {:?}", kdc_rep.cname);
                            debug!("TGS-REP ticket sname: {:?}", kdc_rep.ticket.sname);
                            debug!("TGS-REP ticket.realm: {:?}", kdc_rep.ticket.realm);
                            if test_weak_crypto(kdc_rep.enc_part.etype) {
                                debug!("AS-REP weak crypto");
                            }
                        };
                        self.req_id = 0;
                    },
                    14 => {
                        self.req_id = hdr.tag.0;
                        let res = krb5_parser::parse_ap_req(i);
                        debug!("AP-REQ: {:?}", res);
                        if let Ok((_,ap_rep)) = res {
                            debug!("AP-REQ ticket sname: {:?}", ap_rep.ticket.sname);
                            debug!("AP-REQ ticket.realm: {:?}", ap_rep.ticket.realm);
                        };
                    },
                    15 => {
                        let res = krb5_parser::parse_ap_rep(i);
                        debug!("AP-REP {:?}", res);
                        self.req_id = 0;
                    },
                    30 => {
                        let res = krb5_parser::parse_krb_error(i);
                        debug!("KRB-ERROR: {:?}", res);
                        debug!("KRB-ERROR: failed request: {:?}", self.req_id);
                        if let Ok((_,error)) = res {
                            debug!("KRB-ERROR: cname: {:?}", error.cname);
                            debug!("KRB-ERROR: realm: {:?}", error.realm);
                            debug!("KRB-ERROR: sname: {:?}", error.sname);
                            debug!("KRB-ERROR: error_code: {:?}", error.error_code);
                        };
                        self.req_id = 0;
                    },
                    _ => debug!("unknown/unsupported tag {}", hdr.tag),
                }
                R_STATUS_OK
            },
            _ => R_STATUS_FAIL
        }
    }
}

impl<'a> KerberosParserUDP<'a> {
    pub fn new(name: &'a[u8]) -> KerberosParserUDP<'a> {
        KerberosParserUDP{
            _name: Some(name),
            req_id: 0,
        }
    }
}

/// Return true if Kerberos `EncryptionType` is weak
pub fn test_weak_crypto(alg:EncryptionType) -> bool {
    match alg {
        EncryptionType::AES128_CTS_HMAC_SHA1_96 |
        EncryptionType::AES256_CTS_HMAC_SHA1_96 |
        EncryptionType::AES128_CTS_HMAC_SHA256_128 |
        EncryptionType::AES256_CTS_HMAC_SHA384_192 |
        EncryptionType::CAMELLIA128_CTS_CMAC |
        EncryptionType::CAMELLIA256_CTS_CMAC => false,
        _ => true, // all other ciphers are weak or deprecated
    }
}

pub fn kerberos_probe_udp(i: &[u8]) -> bool {
    if i.len() <= 10 { return false; }
    match der_read_element_header(i) {
        Ok((rem,hdr)) => {
            // Kerberos messages start with an APPLICATION header
            if hdr.class != 0b01 { return false; }
            // Tag number should be <= 30
            if hdr.tag.0 >= 30 { return false; }
            // Kerberos messages contain sequences
            if rem.is_empty() || rem[0] != 0x30 { return false; }
            // Check kerberos version
            // debug!("hdr: {:?}", hdr);
            if let Ok((rem,_hdr)) = der_read_element_header(rem) {
                if rem.len() > 5 {
                    match (rem[2],rem[3],rem[4]) {
                        // Encoding of DER integer 5 (version)
                        (2,1,5) => { return true; },
                        _       => (),
                    }
                }
            }
            false
        },
        _ => false,
    }
}
<|endoftext|>


<commit_before>//! Relay for TCP server that running on the server side

use std::io::{self, ErrorKind};
use std::rc::Rc;
use std::time::Duration;

use config::ServerConfig;

use relay::{boxed_future, BoxIoFuture};
use relay::Context;
use relay::dns_resolver::resolve;
use relay::socks5::Address;
use relay::tcprelay::crypto_io::{DecryptedRead, EncryptedWrite};

use futures::{self, Future};
use futures::stream::Stream;

use tokio_core::net::{TcpListener, TcpStream};
use tokio_io::AsyncRead;
use tokio_io::io::{ReadHalf, WriteHalf};

use super::{proxy_handshake, try_timeout, tunnel, DecryptedHalf, EncryptedHalfFut, TcpStreamConnect};

/// Context for doing handshake with client
pub struct TcpRelayClientHandshake {
    s: TcpStream,
    svr_cfg: Rc<ServerConfig>,
}

impl TcpRelayClientHandshake {
    /// Doing handshake with client
    pub fn handshake(self) -> BoxIoFuture<TcpRelayClientPending> {
        let TcpRelayClientHandshake { s, svr_cfg } = self;

        let peer_addr = s.peer_addr().expect("Failed to get peer addr for client");
        debug!("Handshaking with peer {}", peer_addr);

        let timeout = *svr_cfg.timeout();
        let fut = proxy_handshake(s, svr_cfg).and_then(move |(r_fut, w_fut)| {
            r_fut.and_then(move |r| {
                               let fut = Address::read_from(r).map_err(move |_| {
                                   io::Error::new(ErrorKind::Other,
                                   format!("failed to decode Address, may be wrong method or key, peer: {}", peer_addr))
                               });
                               Context::with(|ctx| try_timeout(fut, timeout, ctx.handle()))
                           })
                 .map(move |(r, addr)| TcpRelayClientPending { r: r,
                                                               addr: addr,
                                                               w: w_fut,
                                                               timeout: timeout, })
        });
        boxed_future(fut)
    }
}

/// Context for connecting remote
pub struct TcpRelayClientPending {
    r: DecryptedHalf,
    addr: Address,
    w: EncryptedHalfFut,
    timeout: Option<Duration>,
}

impl TcpRelayClientPending {
    /// Connect to the remote server
    #[inline]
    fn connect_remote(addr: Address, timeout: Option<Duration>) -> BoxIoFuture<TcpStream> {
        debug!("Connecting to remote {}", addr);

        match addr {
            Address::SocketAddress(saddr) => {
                Context::with(move |ctx| {
                                  if ctx.forbidden_ip().contains(&saddr.ip()) {
                                      let err = io::Error::new(ErrorKind::Other,
                                                               format!("{} is forbidden, failed to connect {}",
                                                                       saddr.ip(),
                                                                       saddr));
                                      return boxed_future(futures::done(Err(err)));
                                  }

                                  let conn = TcpStream::connect(&saddr, ctx.handle());
                                  try_timeout(conn, timeout, ctx.handle())
                              })
            }
            Address::DomainNameAddress(dname, port) => {
                let fut = Context::with(move |ctx| {
                    let handle = ctx.handle().clone();
                    try_timeout(resolve(dname.as_str(), port, true), timeout, &handle).and_then(move |addrs| {
                        let conn = TcpStreamConnect::new(addrs.into_iter(), &handle);
                        try_timeout(conn, timeout, &handle)
                    })
                });
                boxed_future(fut)
            }
        }
    }

    /// Connect to the remote server
    pub fn connect(self) -> BoxIoFuture<TcpRelayClientConnected> {
        let addr = self.addr.clone();
        let client_pair = (self.r, self.w);
        let timeout = self.timeout;
        let fut = TcpRelayClientPending::connect_remote(self.addr, self.timeout);
        let fut = fut.map(move |stream| TcpRelayClientConnected { server: stream.split(),
                                                                  client: client_pair,
                                                                  addr: addr,
                                                                  timeout: timeout, });
        Box::new(fut)
    }
}

/// Context for extablishing tunnel
pub struct TcpRelayClientConnected {
    server: (ReadHalf<TcpStream>, WriteHalf<TcpStream>),
    client: (DecryptedHalf, EncryptedHalfFut),
    addr: Address,
    timeout: Option<Duration>,
}

impl TcpRelayClientConnected {
    /// Establish tunnel
    #[inline]
    pub fn tunnel(self) -> BoxIoFuture<()> {
        let (svr_r, svr_w) = self.server;
        let (r, w_fut) = self.client;
        let timeout = self.timeout;

        tunnel(self.addr,
               r.copy_timeout_opt(svr_w, self.timeout),
               w_fut.and_then(move |w| w.copy_timeout_opt(svr_r, timeout)))
    }
}

/// Runs the server
pub fn run() -> Box<Future<Item = (), Error = io::Error>> {
    let mut fut: Option<Box<Future<Item = (), Error = io::Error>>> = None;

    Context::with(|ctx| {
                      let config = ctx.config();

                      for svr_cfg in &config.server {
                          let listener = {
                              let addr = svr_cfg.addr();
                              let addr = addr.listen_addr();

                              let listener = TcpListener::bind(&addr, ctx.handle())
                    .unwrap_or_else(|err| panic!("Failed to listen, {}", err));

                              info!("ShadowSocks TCP Listening on {}", addr);
                              listener
                          };

                          let svr_cfg = Rc::new(svr_cfg.clone());
                          let listening = listener.incoming()
                                                  .for_each(move |(socket, addr)| {
                                                      let server_cfg = svr_cfg.clone();

                                                      trace!("Got connection, addr: {}", addr);
                                                      trace!("Picked proxy server: {:?}", server_cfg);

                                                      let client = TcpRelayClientHandshake { s: socket,
                                                                                svr_cfg: server_cfg, };

                                                      let fut =
                                                          client.handshake()
                                                                .and_then(|c| c.connect())
                                                                .and_then(|c| c.tunnel())
                                                                .map_err(move |err| {
                                                                             error!("Failed to handle client ({}): {}",
                                                                                addr, err);
                                                                         });

                                                      Context::with(|ctx| ctx.handle().spawn(fut));
                                                      Ok(())
                                                  })
                                                  .map_err(|err| {
                                                               error!("Server run failed: {}", err);
                                                               err
                                                           });

                          fut = Some(match fut.take() {
                                         Some(fut) => {
                                             Box::new(fut.join(listening).map(|_| ()))
                                             as Box<Future<Item = (), Error = io::Error>>
                                         }
                                         None => Box::new(listening) as Box<Future<Item = (), Error = io::Error>>,
                                     })
                      }

                      fut.expect("Must have at least one server")
                  })
}
<commit_msg>[#96] Do not panic if peer_addr() failed after accept<commit_after>//! Relay for TCP server that running on the server side

use std::io::{self, ErrorKind};
use std::rc::Rc;
use std::time::Duration;

use config::ServerConfig;

use relay::{boxed_future, BoxIoFuture};
use relay::Context;
use relay::dns_resolver::resolve;
use relay::socks5::Address;
use relay::tcprelay::crypto_io::{DecryptedRead, EncryptedWrite};

use futures::{self, Future};
use futures::stream::Stream;

use tokio_core::net::{TcpListener, TcpStream};
use tokio_io::AsyncRead;
use tokio_io::io::{ReadHalf, WriteHalf};

use super::{proxy_handshake, try_timeout, tunnel, DecryptedHalf, EncryptedHalfFut, TcpStreamConnect};

/// Context for doing handshake with client
pub struct TcpRelayClientHandshake {
    s: TcpStream,
    svr_cfg: Rc<ServerConfig>,
}

impl TcpRelayClientHandshake {
    /// Doing handshake with client
    pub fn handshake(self) -> BoxIoFuture<TcpRelayClientPending> {
        let TcpRelayClientHandshake { s, svr_cfg } = self;

        let fut = futures::lazy(move || s.peer_addr().map(|p| (s, p))).and_then(|(s, peer_addr)| {
            debug!("Handshaking with peer {}", peer_addr);

            let timeout = *svr_cfg.timeout();
            proxy_handshake(s, svr_cfg).and_then(move |(r_fut, w_fut)| {
                r_fut
                    .and_then(move |r| {
                        let fut = Address::read_from(r).map_err(move |_| {
                            io::Error::new(
                                ErrorKind::Other,
                                format!(
                                    "failed to decode Address, may be wrong method or key, peer: {}",
                                    peer_addr
                                ),
                            )
                        });
                        Context::with(|ctx| try_timeout(fut, timeout, ctx.handle()))
                    })
                    .map(move |(r, addr)| TcpRelayClientPending {
                        r: r,
                        addr: addr,
                        w: w_fut,
                        timeout: timeout,
                    })
            })
        });
        boxed_future(fut)
    }
}

/// Context for connecting remote
pub struct TcpRelayClientPending {
    r: DecryptedHalf,
    addr: Address,
    w: EncryptedHalfFut,
    timeout: Option<Duration>,
}

impl TcpRelayClientPending {
    /// Connect to the remote server
    #[inline]
    fn connect_remote(addr: Address, timeout: Option<Duration>) -> BoxIoFuture<TcpStream> {
        debug!("Connecting to remote {}", addr);

        match addr {
            Address::SocketAddress(saddr) => Context::with(move |ctx| {
                if ctx.forbidden_ip().contains(&saddr.ip()) {
                    let err = io::Error::new(
                        ErrorKind::Other,
                        format!("{} is forbidden, failed to connect {}", saddr.ip(), saddr),
                    );
                    return boxed_future(futures::done(Err(err)));
                }

                let conn = TcpStream::connect(&saddr, ctx.handle());
                try_timeout(conn, timeout, ctx.handle())
            }),
            Address::DomainNameAddress(dname, port) => {
                let fut = Context::with(move |ctx| {
                    let handle = ctx.handle().clone();
                    try_timeout(resolve(dname.as_str(), port, true), timeout, &handle).and_then(move |addrs| {
                        let conn = TcpStreamConnect::new(addrs.into_iter(), &handle);
                        try_timeout(conn, timeout, &handle)
                    })
                });
                boxed_future(fut)
            }
        }
    }

    /// Connect to the remote server
    pub fn connect(self) -> BoxIoFuture<TcpRelayClientConnected> {
        let addr = self.addr.clone();
        let client_pair = (self.r, self.w);
        let timeout = self.timeout;
        let fut = TcpRelayClientPending::connect_remote(self.addr, self.timeout);
        let fut = fut.map(move |stream| TcpRelayClientConnected {
            server: stream.split(),
            client: client_pair,
            addr: addr,
            timeout: timeout,
        });
        Box::new(fut)
    }
}

/// Context for extablishing tunnel
pub struct TcpRelayClientConnected {
    server: (ReadHalf<TcpStream>, WriteHalf<TcpStream>),
    client: (DecryptedHalf, EncryptedHalfFut),
    addr: Address,
    timeout: Option<Duration>,
}

impl TcpRelayClientConnected {
    /// Establish tunnel
    #[inline]
    pub fn tunnel(self) -> BoxIoFuture<()> {
        let (svr_r, svr_w) = self.server;
        let (r, w_fut) = self.client;
        let timeout = self.timeout;

        tunnel(
            self.addr,
            r.copy_timeout_opt(svr_w, self.timeout),
            w_fut.and_then(move |w| w.copy_timeout_opt(svr_r, timeout)),
        )
    }
}

/// Runs the server
pub fn run() -> Box<Future<Item = (), Error = io::Error>> {
    let mut fut: Option<Box<Future<Item = (), Error = io::Error>>> = None;

    Context::with(|ctx| {
        let config = ctx.config();

        for svr_cfg in &config.server {
            let listener = {
                let addr = svr_cfg.addr();
                let addr = addr.listen_addr();

                let listener =
                    TcpListener::bind(&addr, ctx.handle()).unwrap_or_else(|err| panic!("Failed to listen, {}", err));

                info!("ShadowSocks TCP Listening on {}", addr);
                listener
            };

            let svr_cfg = Rc::new(svr_cfg.clone());
            let listening = listener
                .incoming()
                .for_each(move |(socket, addr)| {
                    let server_cfg = svr_cfg.clone();

                    trace!("Got connection, addr: {}", addr);
                    trace!("Picked proxy server: {:?}", server_cfg);

                    let client = TcpRelayClientHandshake {
                        s: socket,
                        svr_cfg: server_cfg,
                    };

                    let fut = client
                        .handshake()
                        .and_then(|c| c.connect())
                        .and_then(|c| c.tunnel())
                        .map_err(move |err| {
                            error!("Failed to handle client ({}): {}", addr, err);
                        });

                    Context::with(|ctx| ctx.handle().spawn(fut));
                    Ok(())
                })
                .map_err(|err| {
                    error!("Server run failed: {}", err);
                    err
                });

            fut = Some(match fut.take() {
                Some(fut) => Box::new(fut.join(listening).map(|_| ())) as Box<Future<Item = (), Error = io::Error>>,
                None => Box::new(listening) as Box<Future<Item = (), Error = io::Error>>,
            })
        }

        fut.expect("Must have at least one server")
    })
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

//! Viewports and scissor boxes.
//!
//! There are two different concepts to determine where things will be drawn:
//!
//! - The viewport is the region of the image which corresponds to the
//!   vertex coordinates `-1.0` to `1.0`.
//! - Any pixel outside of the scissor box will be discarded.
//!
//! In other words modifying the viewport will stretch the image, while modifying the scissor
//! box acts like a filter.
//!
//! It is legal and sensible to use a viewport that is larger than the target image or that
//! only partially overlaps the target image.
//!
//! # Multiple viewports
//!
//! In most situations, you only need a single viewport and a single scissor box.
//!
//! If, however, you use a geometry shader, you can specify multiple viewports and scissor boxes.
//! Then in your geometry shader you can specify in which viewport and scissor box the primitive
//! should be written to. In GLSL this is done by writing to the special variable
//! `gl_ViewportIndex`.
//!
//! If you don't use a geometry shader or use a geometry shader where don't set which viewport to
//! use, then the first viewport and scissor box will be used.
//!
//! # Dynamic and fixed
//!
//! Vulkan allows four different setups:
//!
//! - The state of both the viewports and scissor boxes is known at pipeline creation.
//! - The state of viewports is known at pipeline creation, but the state of scissor boxes is
//!   only known when submitting the draw command.
//! - The state of scissor boxes is known at pipeline creation, but the state of viewports is
//!   only known when submitting the draw command.
//! - The state of both the viewports and scissor boxes is only known when submitting the
//!   draw command.
//!
//! In all cases the number of viewports and scissor boxes must be the same.
//!
use std::ops::Range;
use vk;

/// List of viewports and scissors that are used when creating a graphics pipeline object.
///
/// Note that the number of viewports and scissors must be the same.
#[derive(Debug, Clone)]
pub enum ViewportsState {
    /// The state is known in advance.
    Fixed {
        /// State of the viewports and scissors.
        data: Vec<(Viewport, Scissor)>,
    },

    /// The state of scissors is known in advance, but the state of viewports is dynamic and will
    /// bet set when drawing.
    ///
    /// Note that the number of viewports and scissors must be the same.
    DynamicViewports {
        /// State of the scissors.
        scissors: Vec<Scissor>,
    },

    /// The state of viewports is known in advance, but the state of scissors is dynamic and will
    /// bet set when drawing.
    ///
    /// Note that the number of viewports and scissors must be the same.
    DynamicScissors {
        /// State of the viewports
        viewports: Vec<Viewport>,
    },

    /// The state of both the viewports and scissors is dynamic and will be set when drawing.
    Dynamic {
        /// Number of viewports and scissors.
        num: u32,
    },
}

impl ViewportsState {
    /// Returns true if the state of the viewports is dynamic.
    pub fn dynamic_viewports(&self) -> bool {
        match *self {
            ViewportsState::Fixed { .. } => false,
            ViewportsState::DynamicViewports { .. } => true,
            ViewportsState::DynamicScissors { .. } => false,
            ViewportsState::Dynamic { .. } => true,
        }
    }

    /// Returns true if the state of the scissors is dynamic.
    pub fn dynamic_scissors(&self) -> bool {
        match *self {
            ViewportsState::Fixed { .. } => false,
            ViewportsState::DynamicViewports { .. } => false,
            ViewportsState::DynamicScissors { .. } => true,
            ViewportsState::Dynamic { .. } => true,
        }
    }

    /// Returns the number of viewports and scissors.
    pub fn num_viewports(&self) -> u32 {
        match *self {
            ViewportsState::Fixed { ref data } => data.len() as u32,
            ViewportsState::DynamicViewports { ref scissors } => scissors.len() as u32,
            ViewportsState::DynamicScissors { ref viewports } => viewports.len() as u32,
            ViewportsState::Dynamic { num } => num,
        }
    }
}

/// State of a single viewport.
// FIXME: check that:
//        x + width must be less than or equal to viewportBoundsRange[0]
//        y + height must be less than or equal to viewportBoundsRange[1]
#[derive(Debug, Clone)]
pub struct Viewport {
    /// Coordinates in pixels of the top-left hand corner of the viewport.
    pub origin: [f32; 2],

    /// Dimensions in pixels of the viewport.
    pub dimensions: [f32; 2],

    /// Minimum and maximum values of the depth.
    ///
    /// The values `-1.0` to `1.0` of each vertex's Z coordinate will be mapped to this
    /// `depth_range` before being compared to the existing depth value.
    ///
    /// This is equivalents to `glDepthRange` in OpenGL.
    pub depth_range: Range<f32>,
}

#[doc(hidden)]
impl Into<vk::Viewport> for Viewport {
    #[inline]
    fn into(self) -> vk::Viewport {
        vk::Viewport {
            x: self.origin[0],
            y: self.origin[1],
            width: self.dimensions[0],
            height: self.dimensions[1],
            minDepth: self.depth_range.start,
            maxDepth: self.depth_range.end,
        }
    }
}

/// State of a single scissor box.
// FIXME: add a check:
//      Evaluation of (offset.x + extent.width) must not cause a signed integer addition overflow
//      Evaluation of (offset.y + extent.height) must not cause a signed integer addition overflow
#[derive(Debug, Copy, Clone)]
pub struct Scissor {
    /// Coordinates in pixels of the top-left hand corner of the box.
    pub origin: [i32; 2],

    /// Dimensions in pixels of the box.
    pub dimensions: [u32; 2],
}

impl Scissor {
    /// Defines a scissor box that it outside of the image.
    #[inline]
    pub fn irrelevant() -> Scissor {
        Scissor {
            origin: [0, 0],
            dimensions: [0x7fffffff, 0x7fffffff],
        }
    }
}

impl Default for Scissor {
    #[inline]
    fn default() -> Scissor {
        Scissor::irrelevant()
    }
}

#[doc(hidden)]
impl Into<vk::Rect2D> for Scissor {
    #[inline]
    fn into(self) -> vk::Rect2D {
        vk::Rect2D {
            offset: vk::Offset2D {
                x: self.origin[0],
                y: self.origin[1],
            },
            extent: vk::Extent2D {
                width: self.dimensions[0],
                height: self.dimensions[1],
            },
        }
    }
}
<commit_msg>Small documentation fix in viewport<commit_after>// Copyright (c) 2016 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

//! Viewports and scissor boxes.
//!
//! There are two different concepts to determine where things will be drawn:
//!
//! - The viewport is the region of the image which corresponds to the
//!   vertex coordinates `-1.0` to `1.0`.
//! - Any pixel outside of the scissor box will be discarded.
//!
//! In other words modifying the viewport will stretch the image, while modifying the scissor
//! box acts like a filter.
//!
//! It is legal and sensible to use a viewport that is larger than the target image or that
//! only partially overlaps the target image.
//!
//! # Multiple viewports
//!
//! In most situations, you only need a single viewport and a single scissor box.
//!
//! If, however, you use a geometry shader, you can specify multiple viewports and scissor boxes.
//! Then in your geometry shader you can specify in which viewport and scissor box the primitive
//! should be written to. In GLSL this is done by writing to the special variable
//! `gl_ViewportIndex`.
//!
//! If you don't use a geometry shader or use a geometry shader where don't set which viewport to
//! use, then the first viewport and scissor box will be used.
//!
//! # Dynamic and fixed
//!
//! Vulkan allows four different setups:
//!
//! - The state of both the viewports and scissor boxes is known at pipeline creation.
//! - The state of viewports is known at pipeline creation, but the state of scissor boxes is
//!   only known when submitting the draw command.
//! - The state of scissor boxes is known at pipeline creation, but the state of viewports is
//!   only known when submitting the draw command.
//! - The state of both the viewports and scissor boxes is only known when submitting the
//!   draw command.
//!
//! In all cases the number of viewports and scissor boxes must be the same.
//!
use std::ops::Range;
use vk;

/// List of viewports and scissors that are used when creating a graphics pipeline object.
///
/// Note that the number of viewports and scissors must be the same.
#[derive(Debug, Clone)]
pub enum ViewportsState {
    /// The state is known in advance.
    Fixed {
        /// State of the viewports and scissors.
        data: Vec<(Viewport, Scissor)>,
    },

    /// The state of scissors is known in advance, but the state of viewports is dynamic and will
    /// bet set when drawing.
    ///
    /// Note that the number of viewports and scissors must be the same.
    DynamicViewports {
        /// State of the scissors.
        scissors: Vec<Scissor>,
    },

    /// The state of viewports is known in advance, but the state of scissors is dynamic and will
    /// bet set when drawing.
    ///
    /// Note that the number of viewports and scissors must be the same.
    DynamicScissors {
        /// State of the viewports
        viewports: Vec<Viewport>,
    },

    /// The state of both the viewports and scissors is dynamic and will be set when drawing.
    Dynamic {
        /// Number of viewports and scissors.
        num: u32,
    },
}

impl ViewportsState {
    /// Returns true if the state of the viewports is dynamic.
    pub fn dynamic_viewports(&self) -> bool {
        match *self {
            ViewportsState::Fixed { .. } => false,
            ViewportsState::DynamicViewports { .. } => true,
            ViewportsState::DynamicScissors { .. } => false,
            ViewportsState::Dynamic { .. } => true,
        }
    }

    /// Returns true if the state of the scissors is dynamic.
    pub fn dynamic_scissors(&self) -> bool {
        match *self {
            ViewportsState::Fixed { .. } => false,
            ViewportsState::DynamicViewports { .. } => false,
            ViewportsState::DynamicScissors { .. } => true,
            ViewportsState::Dynamic { .. } => true,
        }
    }

    /// Returns the number of viewports and scissors.
    pub fn num_viewports(&self) -> u32 {
        match *self {
            ViewportsState::Fixed { ref data } => data.len() as u32,
            ViewportsState::DynamicViewports { ref scissors } => scissors.len() as u32,
            ViewportsState::DynamicScissors { ref viewports } => viewports.len() as u32,
            ViewportsState::Dynamic { num } => num,
        }
    }
}

/// State of a single viewport.
// FIXME: check that:
//        x + width must be less than or equal to viewportBoundsRange[0]
//        y + height must be less than or equal to viewportBoundsRange[1]
#[derive(Debug, Clone)]
pub struct Viewport {
    /// Coordinates in pixels of the top-left hand corner of the viewport.
    pub origin: [f32; 2],

    /// Dimensions in pixels of the viewport.
    pub dimensions: [f32; 2],

    /// Minimum and maximum values of the depth.
    ///
    /// The values `0.0` to `1.0` of each vertex's Z coordinate will be mapped to this
    /// `depth_range` before being compared to the existing depth value.
    ///
    /// This is equivalents to `glDepthRange` in OpenGL, except that OpenGL uses the Z coordinate
    /// range from `-1.0` to `1.0` instead.
    pub depth_range: Range<f32>,
}

#[doc(hidden)]
impl Into<vk::Viewport> for Viewport {
    #[inline]
    fn into(self) -> vk::Viewport {
        vk::Viewport {
            x: self.origin[0],
            y: self.origin[1],
            width: self.dimensions[0],
            height: self.dimensions[1],
            minDepth: self.depth_range.start,
            maxDepth: self.depth_range.end,
        }
    }
}

/// State of a single scissor box.
// FIXME: add a check:
//      Evaluation of (offset.x + extent.width) must not cause a signed integer addition overflow
//      Evaluation of (offset.y + extent.height) must not cause a signed integer addition overflow
#[derive(Debug, Copy, Clone)]
pub struct Scissor {
    /// Coordinates in pixels of the top-left hand corner of the box.
    pub origin: [i32; 2],

    /// Dimensions in pixels of the box.
    pub dimensions: [u32; 2],
}

impl Scissor {
    /// Defines a scissor box that it outside of the image.
    #[inline]
    pub fn irrelevant() -> Scissor {
        Scissor {
            origin: [0, 0],
            dimensions: [0x7fffffff, 0x7fffffff],
        }
    }
}

impl Default for Scissor {
    #[inline]
    fn default() -> Scissor {
        Scissor::irrelevant()
    }
}

#[doc(hidden)]
impl Into<vk::Rect2D> for Scissor {
    #[inline]
    fn into(self) -> vk::Rect2D {
        vk::Rect2D {
            offset: vk::Offset2D {
                x: self.origin[0],
                y: self.origin[1],
            },
            extent: vk::Extent2D {
                width: self.dimensions[0],
                height: self.dimensions[1],
            },
        }
    }
}
<|endoftext|>


<commit_before>//! Tests for the `cargo sync-lockfile` command.

use cargo_test_support::registry::Package;
use cargo_test_support::rustc_host;
use cargo_test_support::{basic_manifest, cross_compile, project};

#[cargo_test]
fn no_deps() {
    let p = project()
        .file("src/main.rs", "mod a; fn main() {}")
        .file("src/a.rs", "")
        .build();

    p.cargo("sync-lockfile").with_stdout("").run();
}

#[cargo_test]
fn sync_all_platform_dependencies_when_no_target_is_given() {
    if cross_compile::disabled() {
        return;
    }

    Package::new("d1", "1.2.3")
        .file("Cargo.toml", &basic_manifest("d1", "1.2.3"))
        .file("src/lib.rs", "")
        .publish();

    Package::new("d2", "0.1.2")
        .file("Cargo.toml", &basic_manifest("d2", "0.1.2"))
        .file("src/lib.rs", "")
        .publish();

    let target = cross_compile::alternate();
    let host = rustc_host();
    let p = project()
        .file(
            "Cargo.toml",
            &format!(
                r#"
            [package]
            name = "foo"
            version = "0.0.1"
            authors = []

            [target.{host}.dependencies]
            d1 = "1.2.3"

            [target.{target}.dependencies]
            d2 = "0.1.2"
        "#,
                host = host,
                target = target
            ),
        )
        .file("src/lib.rs", "")
        .build();

    p.cargo("sync-lockfile")
        .with_stderr_contains("[DOWNLOADED] d1 v1.2.3 [..]")
        .with_stderr_contains("[DOWNLOADED] d2 v0.1.2 [..]")
        .run();
}

#[cargo_test]
fn sync_platform_specific_dependencies() {
    if cross_compile::disabled() {
        return;
    }

    Package::new("d1", "1.2.3")
        .file("Cargo.toml", &basic_manifest("d1", "1.2.3"))
        .file("src/lib.rs", "")
        .publish();

    Package::new("d2", "0.1.2")
        .file("Cargo.toml", &basic_manifest("d2", "0.1.2"))
        .file("src/lib.rs", "")
        .publish();

    let target = cross_compile::alternate();
    let host = rustc_host();
    let p = project()
        .file(
            "Cargo.toml",
            &format!(
                r#"
            [package]
            name = "foo"
            version = "0.0.1"
            authors = []

            [target.{host}.dependencies]
            d1 = "1.2.3"

            [target.{target}.dependencies]
            d2 = "0.1.2"
        "#,
                host = host,
                target = target
            ),
        )
        .file("src/lib.rs", "")
        .build();

    p.cargo("sync-lockfile --target")
        .arg(&host)
        .with_stderr_contains("[DOWNLOADED] d1 v1.2.3 [..]")
        .with_stderr_does_not_contain("[DOWNLOADED] d2 v0.1.2 [..]")
        .run();

    p.cargo("sync-lockfile --target")
        .arg(&target)
        .with_stderr_contains("[DOWNLOADED] d2 v0.1.2[..]")
        .with_stderr_does_not_contain("[DOWNLOADED] d1 v1.2.3 [..]")
        .run();
}

#[cargo_test]
fn sync_warning() {
    let p = project()
        .file(
            "Cargo.toml",
            r#"
            [package]
            name = "foo"
            version = "1.0.0"
            misspelled = "wut"
            "#,
        )
        .file("src/lib.rs", "")
        .build();
    p.cargo("sync-lockfile")
        .with_stderr("[WARNING] unused manifest key: package.misspelled")
        .run();
}
<commit_msg>Fix for testsuite<commit_after>//! Tests for the `cargo sync-lockfile` command.

use cargo_test_support::project;

#[cargo_test]
fn sync_after_generate() {
    let p = project().file("src/main.rs", "fn main() {}").build();
    p.cargo("generate-lockfile").run();
    let lock1 = p.read_lockfile();

    // add a dep
    p.change_file(
        "Cargo.toml",
        r#"
            [package]
            name = "foo"
            authors = []
            version = "0.0.2"
        "#,
    );
    p.cargo("sync-lockfile").run();
    let lock2 = p.read_lockfile();
    assert_ne!(lock1, lock2);
    assert!(lock1.contains("0.0.1"));
    assert!(lock2.contains("0.0.2"));
    assert!(!lock1.contains("0.0.2"));
    assert!(!lock2.contains("0.0.1"));
}
<|endoftext|>


<commit_before>extern crate rand;
extern crate time;
extern crate timely;
extern crate differential_dataflow;

use std::hash::Hash;

use timely::dataflow::*;
use timely::dataflow::operators::*;
use timely::progress::timestamp::RootTimestamp;

use rand::{Rng, SeedableRng, StdRng};

use differential_dataflow::Collection;
use differential_dataflow::operators::*;
use differential_dataflow::operators::join::JoinUnsigned;
use differential_dataflow::operators::group::GroupUnsigned;
use differential_dataflow::collection::LeastUpperBound;

type Node = usize;
type Edge = (Node, Node);

fn main() {

    let users: usize = std::env::args().nth(1).unwrap().parse().unwrap();
    let topics: usize = std::env::args().nth(2).unwrap().parse().unwrap();
    let batch: usize = std::env::args().nth(3).unwrap().parse().unwrap();
    let inspect: bool = std::env::args().nth(4).unwrap() == "inspect";

    // define a new computational scope, in which to run BFS
    timely::execute_from_args(std::env::args().skip(5), move |computation| {
        let start = time::precise_time_s();

        // define BFS dataflow; return handles to roots and edges inputs
        let (mut tweets, mut queries, probe) = computation.scoped(|scope| {

            // entries corresponding to (@username, @mention, #topic), but a (u32, u32, u32) instead.
            let (tweet_input, tweets) = scope.new_input();
            let tweets = Collection::new(tweets);

            // determine connected components based on mentions.
            let labels = connected_components(&tweets.map(|(u,m,_)| (u,m)));

            // produce pairs (label, topic) for each topic.
            let label_topics = tweets.map(|(u,_,t)| (u,t))
                                     .join(&labels)
                                     .map(|(_,t,l)| (l,t));

            // group by (l,t) and emit a count for each.
            let counts = label_topics.map(|x| (x,()))
                                     .group(|_,s,t| {
                                         t.push((s.next().unwrap().1, 1))
                                     });

            let k = 5;
            // retain the k largest counts. negate first to exploit ordering.
            let topk = counts.map(|((l,t), c)| (l, (-c, t)))
                             .group(move |_,s,t| {
                                 t.extend(s.take(k).map(|(&(_,t),_)| (t,1)));
                             });

            // entries corresponding to a @username, but a u32 instead.
            let (query_input, queries) = scope.new_input();
            let queries = Collection::new(queries);

            let label_query = queries.map(|q| (q,()))
                                     .join(&labels)
                                     .map(|(q, _, l)| (l,q));

            let mut query_topics = label_query.join(&topk);

            if !inspect {
                query_topics = query_topics.filter(|_| false);
            }

            let probe = query_topics.consolidate_by(|&(_,q,_)| q)
                                    .inspect(|&((l,q,t),w)| println!("\t(query: {},\tlabel: {},\ttopic:{}\t(weight: {})", q, l, t, w))
                                    .probe();

            (tweet_input, query_input, probe.0)
        });

        let tweet_seed: &[_] = &[0, 1, 2, computation.index()];
        let mut tweet_rng1: StdRng = SeedableRng::from_seed(tweet_seed);    // rng for edge additions
        let mut tweet_rng2: StdRng = SeedableRng::from_seed(tweet_seed);    // rng for edge deletions

        let query_seed: &[_] = &[1, 2, 3, computation.index()];
        let mut query_rng1: StdRng = SeedableRng::from_seed(query_seed);    // rng for edge additions
        let mut query_rng2: StdRng = SeedableRng::from_seed(query_seed);    // rng for edge deletions

        println!("performing AppealingDataflow with {} users, {} topics:", users, topics);

        for _ in 0 .. users/computation.peers() {
            tweets.send(((tweet_rng1.gen_range(0, users),
                          tweet_rng1.gen_range(0, users),
                          tweet_rng1.gen_range(0, topics)),1));
        }

        if computation.index() == 0 {
            queries.send((query_rng1.gen_range(0, users),1));
        }

        println!("loaded; elapsed: {}s", time::precise_time_s() - start);

        tweets.advance_to(1);
        queries.advance_to(1);
        while probe.le(&RootTimestamp::new(0)) { computation.step(); }

        println!("stable; elapsed: {}s", time::precise_time_s() - start);

        if batch > 0 {
            let mut changes = Vec::new();
            for wave in 0.. {
                let mut batch = batch / computation.peers();
                if computation.index() < batch % computation.peers() {
                    batch += 1;
                }

                for _ in 0..batch {
                    changes.push(((tweet_rng1.gen_range(0, users),
                                   tweet_rng1.gen_range(0, users),
                                   tweet_rng1.gen_range(0, topics)), 1));
                    changes.push(((tweet_rng2.gen_range(0, users),
                                   tweet_rng2.gen_range(0, users),
                                   tweet_rng2.gen_range(0, topics)),-1));
                }

                let start = time::precise_time_s();
                let round = *tweets.epoch();
                for change in changes.drain(..) {
                    tweets.send(change);
                }
                if computation.index() == 0 {
                    queries.send((query_rng1.gen_range(0, users), 1));
                    queries.send((query_rng2.gen_range(0, users),-1));
                }

                tweets.advance_to(round + 1);
                queries.advance_to(round + 1);

                while probe.le(&RootTimestamp::new(round)) { computation.step(); }

                if computation.index() == 0 {
                    println!("wave {}: avg {}", wave, (time::precise_time_s() - start) / (batch as f64));
                }
            }
        }
    });
}

fn connected_components<G: Scope>(edges: &Collection<G, Edge>) -> Collection<G, (Node, Node)>
    where G::Timestamp: LeastUpperBound+Hash {

    // each edge (x,y) means that we need at least a label for the min of x and y.
    let nodes = edges.map_in_place(|pair| {
                        let min = std::cmp::min(pair.0, pair.1);
                        *pair = (min, min);
                     })
                     .consolidate_by(|x| x.0);

    // each edge should exist in both directions.
    let edges = edges.map_in_place(|x| ::std::mem::swap(&mut x.0, &mut x.1))
                     .concat(&edges);

    // don't actually use these labels, just grab the type
    nodes.filter(|_| false)
         .iterate(|inner| {
             let edges = edges.enter(&inner.scope());
             let nodes = nodes.enter_at(&inner.scope(), |r| 256 * (64 - (r.0).0.leading_zeros() as u64));

            inner.join_map_u(&edges, |_k,l,d| (*d,*l))
                 .concat(&nodes)
                 .group_u(|_, s, t| { t.push((*s.peek().unwrap().0, 1)); } )
         })
}
<commit_msg>added dataflow example<commit_after>extern crate rand;
extern crate time;
extern crate timely;
extern crate differential_dataflow;

use std::hash::Hash;

use timely::dataflow::*;
use timely::dataflow::operators::*;
use timely::progress::timestamp::RootTimestamp;

use rand::{Rng, SeedableRng, StdRng};

use differential_dataflow::Collection;
use differential_dataflow::operators::*;
use differential_dataflow::operators::join::JoinUnsigned;
use differential_dataflow::operators::group::GroupUnsigned;
use differential_dataflow::collection::LeastUpperBound;

type Node = usize;
type Edge = (Node, Node);

fn main() {

    let users: usize = std::env::args().nth(1).unwrap().parse().unwrap();
    let topics: usize = std::env::args().nth(2).unwrap().parse().unwrap();
    let batch: usize = std::env::args().nth(3).unwrap().parse().unwrap();
    let inspect: bool = std::env::args().nth(4).unwrap() == "inspect";

    // define a new computational scope, in which to run BFS
    timely::execute_from_args(std::env::args().skip(5), move |computation| {
        let start = time::precise_time_s();

        // define BFS dataflow; return handles to roots and edges inputs
        let (mut tweets, mut queries, probe) = computation.scoped(|scope| {

            // entries corresponding to (@username, @mention, #topic), but a (u32, u32, u32) instead.
            let (tweet_input, tweets) = scope.new_input();
            let tweets = Collection::new(tweets);

            // determine connected components based on mentions.
            let labels = connected_components(&tweets.map(|(u,m,_)| (u,m)));

            // produce pairs (label, topic) for each topic.
            let label_topics = tweets.map(|(u,_,t)| (u,t))
                                     .join(&labels)
                                     .map(|(_,t,l)| (l,t));

            // group by (l,t) and emit a count for each.
            let counts = label_topics.map(|x| (x,()))
                                     .group(|_,s,t| {
                                         t.push((s.next().unwrap().1, 1))
                                     });

            let k = 5;
            // retain the k largest counts. negate first to exploit ordering.
            let topk = counts.map(|((l,t), c)| (l, (-c, t)))
                             .group(move |_,s,t| {
                                 t.extend(s.take(k).map(|(&(_,t),_)| (t,1)));
                             });

            // entries corresponding to a @username, but a u32 instead.
            let (query_input, queries) = scope.new_input();
            let queries = Collection::new(queries);

            let label_query = queries.map(|q| (q,()))
                                     .join(&labels)
                                     .map(|(q,_,l)| (l,q));

            let mut query_topics = label_query.join(&topk);

            if !inspect {
                query_topics = query_topics.filter(|_| false);
            }

            let probe = query_topics.consolidate_by(|&(_,q,_)| q)
                                    .inspect(|&((l,q,t),w)| println!("\t(query: {},\tlabel: {},\ttopic:{}\t(weight: {})", q, l, t, w))
                                    .probe();

            (tweet_input, query_input, probe.0)
        });

        let tweet_seed: &[_] = &[0, 1, 2, computation.index()];
        let mut tweet_rng1: StdRng = SeedableRng::from_seed(tweet_seed);    // rng for edge additions
        let mut tweet_rng2: StdRng = SeedableRng::from_seed(tweet_seed);    // rng for edge deletions

        let query_seed: &[_] = &[1, 2, 3, computation.index()];
        let mut query_rng1: StdRng = SeedableRng::from_seed(query_seed);    // rng for edge additions
        let mut query_rng2: StdRng = SeedableRng::from_seed(query_seed);    // rng for edge deletions

        println!("performing AppealingDataflow with {} users, {} topics:", users, topics);

        for _ in 0 .. users/computation.peers() {
            tweets.send(((tweet_rng1.gen_range(0, users),
                          tweet_rng1.gen_range(0, users),
                          tweet_rng1.gen_range(0, topics)),1));
        }

        if computation.index() == 0 {
            queries.send((query_rng1.gen_range(0, users),1));
        }

        println!("loaded; elapsed: {}s", time::precise_time_s() - start);

        tweets.advance_to(1);
        queries.advance_to(1);
        while probe.le(&RootTimestamp::new(0)) { computation.step(); }

        println!("stable; elapsed: {}s", time::precise_time_s() - start);

        if batch > 0 {
            let mut changes = Vec::new();
            for wave in 0.. {
                let mut my_batch = batch / computation.peers();
                if computation.index() < (batch % computation.peers()) {
                    my_batch += 1;
                }

                for _ in 0..my_batch {
                    changes.push(((tweet_rng1.gen_range(0, users),
                                   tweet_rng1.gen_range(0, users),
                                   tweet_rng1.gen_range(0, topics)), 1));
                    changes.push(((tweet_rng2.gen_range(0, users),
                                   tweet_rng2.gen_range(0, users),
                                   tweet_rng2.gen_range(0, topics)),-1));
                }

                let start = time::precise_time_s();
                let round = *tweets.epoch();
                for change in changes.drain(..) {
                    tweets.send(change);
                }
                if computation.index() == 0 {
                    queries.send((query_rng1.gen_range(0, users), 1));
                    queries.send((query_rng2.gen_range(0, users),-1));
                }

                tweets.advance_to(round + 1);
                queries.advance_to(round + 1);

                while probe.le(&RootTimestamp::new(round)) { computation.step(); }

                if computation.index() == 0 {
                    println!("wave {}: avg {}", wave, (time::precise_time_s() - start) / (batch as f64));
                }
            }
        }
    });
}

fn connected_components<G: Scope>(edges: &Collection<G, Edge>) -> Collection<G, (Node, Node)>
    where G::Timestamp: LeastUpperBound+Hash {

    // each edge (x,y) means that we need at least a label for the min of x and y.
    let nodes = edges.map_in_place(|pair| {
                        let min = std::cmp::min(pair.0, pair.1);
                        *pair = (min, min);
                     })
                     .consolidate_by(|x| x.0);

    // each edge should exist in both directions.
    let edges = edges.map_in_place(|x| ::std::mem::swap(&mut x.0, &mut x.1))
                     .concat(&edges);

    // don't actually use these labels, just grab the type
    nodes.filter(|_| false)
         .iterate(|inner| {
             let edges = edges.enter(&inner.scope());
             let nodes = nodes.enter_at(&inner.scope(), |r| 256 * (64 - (r.0).0.leading_zeros() as u64));

            inner.join_map_u(&edges, |_k,l,d| (*d,*l))
                 .concat(&nodes)
                 .group_u(|_, s, t| { t.push((*s.peek().unwrap().0, 1)); } )
         })
}
<|endoftext|>


<commit_before>use futures_core::stream::Stream;
use futures_core::task::{Context, Poll};
use futures_sink::Sink;
use core::fmt;
use core::pin::Pin;

use crate::lock::BiLock;

/// A `Stream` part of the split pair
#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct SplitStream<S>(BiLock<S>);

impl<S> Unpin for SplitStream<S> {}

impl<S: Unpin> SplitStream<S> {
    /// Attempts to put the two "halves" of a split `Stream + Sink` back
    /// together. Succeeds only if the `SplitStream<S>` and `SplitSink<S>` are
    /// a matching pair originating from the same call to `Stream::split`.
    pub fn reunite<Item>(self, other: SplitSink<S, Item>) -> Result<S, ReuniteError<S, Item>>
        where S: Sink<Item>,
    {
        other.reunite(self)
    }
}

impl<S: Stream> Stream for SplitStream<S> {
    type Item = S::Item;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<S::Item>> {
        ready!(self.0.poll_lock(cx)).as_pin_mut().poll_next(cx)
    }
}

#[allow(bad_style)]
fn SplitSink<S: Sink<Item>, Item>(lock: BiLock<S>) -> SplitSink<S, Item> {
    SplitSink {
        lock,
        slot: None,
    }
}

/// A `Sink` part of the split pair
#[derive(Debug)]
#[must_use = "sinks do nothing unless polled"]
pub struct SplitSink<S, Item> {
    lock: BiLock<S>,
    slot: Option<Item>,
}

impl<S, Item> Unpin for SplitSink<S, Item> {}

impl<S: Sink<Item> + Unpin, Item> SplitSink<S, Item> {
    /// Attempts to put the two "halves" of a split `Stream + Sink` back
    /// together. Succeeds only if the `SplitStream<S>` and `SplitSink<S>` are
    /// a matching pair originating from the same call to `Stream::split`.
    pub fn reunite(self, other: SplitStream<S>) -> Result<S, ReuniteError<S, Item>> {
        self.lock.reunite(other.0).map_err(|err| {
            ReuniteError(SplitSink(err.0), SplitStream(err.1))
        })
    }
}

impl<S: Sink<Item>, Item> Sink<Item> for SplitSink<S, Item> {
    type Error = S::Error;

    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> {
        loop {
            if self.slot.is_none() {
                return Poll::Ready(Ok(()));
            }
            ready!(self.as_mut().poll_flush(cx))?;
        }
    }

    fn start_send(mut self: Pin<&mut Self>, item: Item) -> Result<(), S::Error> {
        self.slot = Some(item);
        Ok(())
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> {
        let this = &mut *self;
        let mut inner = ready!(this.lock.poll_lock(cx));
        if this.slot.is_some() {
            ready!(inner.as_pin_mut().poll_ready(cx))?;
            inner.as_pin_mut().start_send(this.slot.take().unwrap())?;
        }
        inner.as_pin_mut().poll_flush(cx)
    }

    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> {
        let this = &mut *self;
        let mut inner = ready!(this.lock.poll_lock(cx));
        if this.slot.is_some() {
            ready!(inner.as_pin_mut().poll_ready(cx))?;
            inner.as_pin_mut().start_send(this.slot.take().unwrap())?
        }
        inner.as_pin_mut().poll_close(cx)
    }
}

pub(super) fn split<S: Stream + Sink<Item>, Item>(s: S) -> (SplitSink<S, Item>, SplitStream<S>) {
    let (a, b) = BiLock::new(s);
    let read = SplitStream(a);
    let write = SplitSink(b);
    (write, read)
}

/// Error indicating a `SplitSink<S>` and `SplitStream<S>` were not two halves
/// of a `Stream + Split`, and thus could not be `reunite`d.
pub struct ReuniteError<T, Item>(pub SplitSink<T, Item>, pub SplitStream<T>);

impl<T, Item> fmt::Debug for ReuniteError<T, Item> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_tuple("ReuniteError")
            .field(&"...")
            .finish()
    }
}

impl<T, Item> fmt::Display for ReuniteError<T, Item> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "tried to reunite a SplitStream and SplitSink that don't form a pair")
    }
}

#[cfg(feature = "std")]
impl<T: core::any::Any, Item> std::error::Error for ReuniteError<T, Item> {}
<commit_msg>Fix SplitSink performance problem caused by flushing inner sink every time<commit_after>use futures_core::stream::Stream;
use futures_core::task::{Context, Poll};
use futures_sink::Sink;
use core::fmt;
use core::pin::Pin;

use crate::lock::BiLock;

/// A `Stream` part of the split pair
#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct SplitStream<S>(BiLock<S>);

impl<S> Unpin for SplitStream<S> {}

impl<S: Unpin> SplitStream<S> {
    /// Attempts to put the two "halves" of a split `Stream + Sink` back
    /// together. Succeeds only if the `SplitStream<S>` and `SplitSink<S>` are
    /// a matching pair originating from the same call to `Stream::split`.
    pub fn reunite<Item>(self, other: SplitSink<S, Item>) -> Result<S, ReuniteError<S, Item>>
        where S: Sink<Item>,
    {
        other.reunite(self)
    }
}

impl<S: Stream> Stream for SplitStream<S> {
    type Item = S::Item;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<S::Item>> {
        ready!(self.0.poll_lock(cx)).as_pin_mut().poll_next(cx)
    }
}

#[allow(bad_style)]
fn SplitSink<S: Sink<Item>, Item>(lock: BiLock<S>) -> SplitSink<S, Item> {
    SplitSink {
        lock,
        slot: None,
    }
}

/// A `Sink` part of the split pair
#[derive(Debug)]
#[must_use = "sinks do nothing unless polled"]
pub struct SplitSink<S, Item> {
    lock: BiLock<S>,
    slot: Option<Item>,
}

impl<S, Item> Unpin for SplitSink<S, Item> {}

impl<S: Sink<Item> + Unpin, Item> SplitSink<S, Item> {
    /// Attempts to put the two "halves" of a split `Stream + Sink` back
    /// together. Succeeds only if the `SplitStream<S>` and `SplitSink<S>` are
    /// a matching pair originating from the same call to `Stream::split`.
    pub fn reunite(self, other: SplitStream<S>) -> Result<S, ReuniteError<S, Item>> {
        self.lock.reunite(other.0).map_err(|err| {
            ReuniteError(SplitSink(err.0), SplitStream(err.1))
        })
    }
}

impl<S: Sink<Item>, Item> SplitSink<S, Item> {
    fn poll_flush_slot(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> {
        let this = &mut *self;
        let mut inner = ready!(this.lock.poll_lock(cx));
        if this.slot.is_some() {
            ready!(inner.as_pin_mut().poll_ready(cx))?;
            Poll::Ready(inner.as_pin_mut().start_send(this.slot.take().unwrap()))
        } else {
            Poll::Ready(Ok(()))
        }
    }
}

impl<S: Sink<Item>, Item> Sink<Item> for SplitSink<S, Item> {
    type Error = S::Error;

    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> {
        loop {
            if self.slot.is_none() {
                return Poll::Ready(Ok(()));
            }
            ready!(self.as_mut().poll_flush_slot(cx))?;
        }
    }

    fn start_send(mut self: Pin<&mut Self>, item: Item) -> Result<(), S::Error> {
        self.slot = Some(item);
        Ok(())
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> {
        let this = &mut *self;
        let mut inner = ready!(this.lock.poll_lock(cx));
        if this.slot.is_some() {
            ready!(inner.as_pin_mut().poll_ready(cx))?;
            inner.as_pin_mut().start_send(this.slot.take().unwrap())?;
        }
        inner.as_pin_mut().poll_flush(cx)
    }

    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> {
        let this = &mut *self;
        let mut inner = ready!(this.lock.poll_lock(cx));
        if this.slot.is_some() {
            ready!(inner.as_pin_mut().poll_ready(cx))?;
            inner.as_pin_mut().start_send(this.slot.take().unwrap())?
        }
        inner.as_pin_mut().poll_close(cx)
    }
}

pub(super) fn split<S: Stream + Sink<Item>, Item>(s: S) -> (SplitSink<S, Item>, SplitStream<S>) {
    let (a, b) = BiLock::new(s);
    let read = SplitStream(a);
    let write = SplitSink(b);
    (write, read)
}

/// Error indicating a `SplitSink<S>` and `SplitStream<S>` were not two halves
/// of a `Stream + Split`, and thus could not be `reunite`d.
pub struct ReuniteError<T, Item>(pub SplitSink<T, Item>, pub SplitStream<T>);

impl<T, Item> fmt::Debug for ReuniteError<T, Item> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_tuple("ReuniteError")
            .field(&"...")
            .finish()
    }
}

impl<T, Item> fmt::Display for ReuniteError<T, Item> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "tried to reunite a SplitStream and SplitSink that don't form a pair")
    }
}

#[cfg(feature = "std")]
impl<T: core::any::Any, Item> std::error::Error for ReuniteError<T, Item> {}
<|endoftext|>


<commit_before>use std::convert::{TryFrom};
use either::Either;

use crate::values::AsValueRef;
use crate::values::{FunctionValue, PointerValue, AnyValue};

use llvm_sys::prelude::LLVMValueRef;
use llvm_sys::core::{LLVMGetTypeKind, LLVMGetElementType, LLVMTypeOf, LLVMGetReturnType};
use llvm_sys::LLVMTypeKind;

/// A value that can be called with the [`Builder::build_call`] instruction.
///
/// In practice, the `F : Into<CallableValue<'ctx>>` bound of [`Builder::build_call`] means it is
/// possible to pass a [`FunctionValue`] to `build_call` directly. It will be implicitly converted
/// into a `CallableValue`.
///
/// ```no_run
/// use inkwell::context::Context;
///
/// // A simple function which calls itself:
/// let context = Context::create();
/// let module = context.create_module("ret");
/// let builder = context.create_builder();
/// let i32_type = context.i32_type();
/// let fn_type = i32_type.fn_type(&[i32_type.into()], false);
/// let fn_value = module.add_function("ret", fn_type, None);
/// let entry = context.append_basic_block(fn_value, "entry");
/// let i32_arg = fn_value.get_first_param().unwrap();
///
/// builder.position_at_end(entry);
///
/// let ret_val = builder.build_call(fn_value, &[i32_arg], "call")
///     .try_as_basic_value()
///     .left()
///     .unwrap();
///
/// builder.build_return(Some(&ret_val));
/// ```
///
/// A [`PointerValue`] cannot be implicitly converted to a `CallableValue` because the pointer may
/// point to a non-function value. Instead we can use [`TryFrom`] and [`TryInto`] to handle this
/// failure case explicitly.
///
/// ```no_run
/// use std::convert::TryFrom;
/// use inkwell::context::Context;
/// use inkwell::values::CallableValue;
///
/// // A simple function which calls itself:
/// let context = Context::create();
/// let module = context.create_module("ret");
/// let builder = context.create_builder();
/// let i32_type = context.i32_type();
/// let fn_type = i32_type.fn_type(&[i32_type.into()], false);
/// let fn_value = module.add_function("ret", fn_type, None);
/// let entry = context.append_basic_block(fn_value, "entry");
/// let i32_arg = fn_value.get_first_param().unwrap();
///
/// builder.position_at_end(entry);
///
/// // take a pointer to the function value
/// let fn_pointer_value = fn_value.as_global_value().as_pointer_value();
///
/// // convert that pointer value into a callable value
/// // explicitly handling the failure case (here with `unwrap`)
/// let callable_value = CallableValue::try_from(fn_pointer_value).unwrap();
///
/// let ret_val = builder.build_call(callable_value, &[i32_arg], "call")
///     .try_as_basic_value()
///     .left()
///     .unwrap();
///
/// builder.build_return(Some(&ret_val));
/// ```
#[derive(Debug)]
pub struct CallableValue<'ctx>(Either<FunctionValue<'ctx>, PointerValue<'ctx>>);

impl<'ctx> AsValueRef for CallableValue<'ctx> {
    fn as_value_ref(&self) -> LLVMValueRef {
        use either::Either::*;

        match self.0 {
            Left(function) => function.as_value_ref(),
            Right(pointer) => pointer.as_value_ref(),
        }
    }
}

impl<'ctx> AnyValue<'ctx> for CallableValue<'ctx> {}

impl<'ctx> CallableValue<'ctx> {
    pub(crate) fn returns_void(&self) -> bool {
        let return_type = unsafe {
            LLVMGetTypeKind(LLVMGetReturnType(LLVMGetElementType(LLVMTypeOf(self.as_value_ref()))))
        };

        matches!(return_type, LLVMTypeKind::LLVMVoidTypeKind)
    }
}

impl<'ctx> From<FunctionValue<'ctx>> for CallableValue<'ctx> {
    fn from(value: FunctionValue<'ctx>) -> Self {
        Self(Either::Left(value))
    }
}

impl<'ctx> TryFrom<PointerValue<'ctx>> for CallableValue<'ctx> {
    type Error = ();

    fn try_from(value: PointerValue<'ctx>) -> Result<Self, Self::Error> {
        // If using a pointer value, we must validate it's a valid function ptr
        let value_ref = value.as_value_ref();
        let ty_kind = unsafe { LLVMGetTypeKind(LLVMGetElementType(LLVMTypeOf(value_ref))) };
        let is_a_fn_ptr = matches!(ty_kind, LLVMTypeKind::LLVMFunctionTypeKind);

        if is_a_fn_ptr {
            Ok(Self(Either::Right(value)))
        } else {
            Err(())
        }
    }
}
<commit_msg>fix documentation imports<commit_after>use std::convert::TryFrom;
use either::Either;

use crate::values::AsValueRef;
use crate::values::{FunctionValue, PointerValue, AnyValue};

use llvm_sys::prelude::LLVMValueRef;
use llvm_sys::core::{LLVMGetTypeKind, LLVMGetElementType, LLVMTypeOf, LLVMGetReturnType};
use llvm_sys::LLVMTypeKind;

// imports used in documentation
#[allow(unused_imports)]
use crate::builder::Builder;
#[allow(unused_imports)]
use std::convert::TryInto;

/// A value that can be called with the [`Builder::build_call`] instruction.
///
/// In practice, the `F : Into<CallableValue<'ctx>>` bound of [`Builder::build_call`] means it is
/// possible to pass a [`FunctionValue`] to `build_call` directly. It will be implicitly converted
/// into a `CallableValue`.
///
/// ```no_run
/// use inkwell::context::Context;
///
/// // A simple function which calls itself:
/// let context = Context::create();
/// let module = context.create_module("ret");
/// let builder = context.create_builder();
/// let i32_type = context.i32_type();
/// let fn_type = i32_type.fn_type(&[i32_type.into()], false);
/// let fn_value = module.add_function("ret", fn_type, None);
/// let entry = context.append_basic_block(fn_value, "entry");
/// let i32_arg = fn_value.get_first_param().unwrap();
///
/// builder.position_at_end(entry);
///
/// let ret_val = builder.build_call(fn_value, &[i32_arg], "call")
///     .try_as_basic_value()
///     .left()
///     .unwrap();
///
/// builder.build_return(Some(&ret_val));
/// ```
///
/// A [`PointerValue`] cannot be implicitly converted to a `CallableValue` because the pointer may
/// point to a non-function value. Instead we can use [`TryFrom`] and [`TryInto`] to handle this
/// failure case explicitly.
///
/// ```no_run
/// use std::convert::TryFrom;
/// use inkwell::context::Context;
/// use inkwell::values::CallableValue;
///
/// // A simple function which calls itself:
/// let context = Context::create();
/// let module = context.create_module("ret");
/// let builder = context.create_builder();
/// let i32_type = context.i32_type();
/// let fn_type = i32_type.fn_type(&[i32_type.into()], false);
/// let fn_value = module.add_function("ret", fn_type, None);
/// let entry = context.append_basic_block(fn_value, "entry");
/// let i32_arg = fn_value.get_first_param().unwrap();
///
/// builder.position_at_end(entry);
///
/// // take a pointer to the function value
/// let fn_pointer_value = fn_value.as_global_value().as_pointer_value();
///
/// // convert that pointer value into a callable value
/// // explicitly handling the failure case (here with `unwrap`)
/// let callable_value = CallableValue::try_from(fn_pointer_value).unwrap();
///
/// let ret_val = builder.build_call(callable_value, &[i32_arg], "call")
///     .try_as_basic_value()
///     .left()
///     .unwrap();
///
/// builder.build_return(Some(&ret_val));
/// ```
#[derive(Debug)]
pub struct CallableValue<'ctx>(Either<FunctionValue<'ctx>, PointerValue<'ctx>>);

impl<'ctx> AsValueRef for CallableValue<'ctx> {
    fn as_value_ref(&self) -> LLVMValueRef {
        use either::Either::*;

        match self.0 {
            Left(function) => function.as_value_ref(),
            Right(pointer) => pointer.as_value_ref(),
        }
    }
}

impl<'ctx> AnyValue<'ctx> for CallableValue<'ctx> {}

impl<'ctx> CallableValue<'ctx> {
    pub(crate) fn returns_void(&self) -> bool {
        let return_type = unsafe {
            LLVMGetTypeKind(LLVMGetReturnType(LLVMGetElementType(LLVMTypeOf(self.as_value_ref()))))
        };

        matches!(return_type, LLVMTypeKind::LLVMVoidTypeKind)
    }
}

impl<'ctx> From<FunctionValue<'ctx>> for CallableValue<'ctx> {
    fn from(value: FunctionValue<'ctx>) -> Self {
        Self(Either::Left(value))
    }
}

impl<'ctx> TryFrom<PointerValue<'ctx>> for CallableValue<'ctx> {
    type Error = ();

    fn try_from(value: PointerValue<'ctx>) -> Result<Self, Self::Error> {
        // If using a pointer value, we must validate it's a valid function ptr
        let value_ref = value.as_value_ref();
        let ty_kind = unsafe { LLVMGetTypeKind(LLVMGetElementType(LLVMTypeOf(value_ref))) };
        let is_a_fn_ptr = matches!(ty_kind, LLVMTypeKind::LLVMFunctionTypeKind);

        if is_a_fn_ptr {
            Ok(Self(Either::Right(value)))
        } else {
            Err(())
        }
    }
}
<|endoftext|>