<commit_before>use std::{
    collections::{BTreeMap, HashMap},
    io::{self, ErrorKind},
    mem,
    net::{IpAddr, SocketAddr},
    pin::Pin,
    sync::Arc,
    task::{Context, Poll, Waker},
    time::Duration as StdDuration,
};

use log::{error, trace};
use parking_lot::Mutex as ParkingMutex;
use shadowsocks::{net::TcpSocketOpts, relay::socks5::Address};
use smoltcp::{
    iface::{Interface, InterfaceBuilder, Routes, SocketHandle},
    phy::{DeviceCapabilities, Medium},
    socket::{TcpSocket, TcpSocketBuffer, TcpState},
    storage::RingBuffer,
    time::{Duration, Instant},
    wire::{IpAddress, IpCidr, Ipv4Address, Ipv6Address, TcpPacket},
};
use tokio::{
    io::{AsyncRead, AsyncWrite, ReadBuf},
    sync::{mpsc, Notify},
    task::JoinHandle,
    time,
};

use crate::local::{
    context::ServiceContext,
    loadbalancing::PingBalancer,
    net::AutoProxyClientStream,
    utils::{establish_tcp_tunnel, to_ipv4_mapped},
};

use super::virt_device::VirtTunDevice;

// NOTE: Default value is taken from Linux
// recv: /proc/sys/net/ipv4/tcp_rmem 87380 bytes
// send: /proc/sys/net/ipv4/tcp_wmem 16384 bytes
const DEFAULT_TCP_SEND_BUFFER_SIZE: u32 = 16384;
const DEFAULT_TCP_RECV_BUFFER_SIZE: u32 = 87380;

struct TcpSocketControl {
    send_buffer: RingBuffer<'static, u8>,
    send_waker: Option<Waker>,
    recv_buffer: RingBuffer<'static, u8>,
    recv_waker: Option<Waker>,
    is_closed: bool,
}

struct TcpSocketManager {
    iface: Interface<'static, VirtTunDevice>,
    manager_notify: Arc<Notify>,
    sockets: HashMap<SocketHandle, Arc<ParkingMutex<TcpSocketControl>>>,
}

type SharedTcpSocketManager = Arc<ParkingMutex<TcpSocketManager>>;

struct TcpConnection {
    control: Arc<ParkingMutex<TcpSocketControl>>,
    manager_notify: Arc<Notify>,
}

impl Drop for TcpConnection {
    fn drop(&mut self) {
        let mut control = self.control.lock();
        control.is_closed = true;
    }
}

impl TcpConnection {
    fn new(socket: TcpSocket<'static>, manager: SharedTcpSocketManager, tcp_opts: &TcpSocketOpts) -> TcpConnection {
        let send_buffer_size = tcp_opts.send_buffer_size.unwrap_or(DEFAULT_TCP_SEND_BUFFER_SIZE);
        let recv_buffer_size = tcp_opts.recv_buffer_size.unwrap_or(DEFAULT_TCP_RECV_BUFFER_SIZE);

        let (control, manager_notify) = {
            let mut manager = manager.lock();
            let socket_handle = manager.iface.add_socket(socket);

            let control = Arc::new(ParkingMutex::new(TcpSocketControl {
                send_buffer: RingBuffer::new(vec![0u8; send_buffer_size as usize]),
                send_waker: None,
                recv_buffer: RingBuffer::new(vec![0u8; recv_buffer_size as usize]),
                recv_waker: None,
                is_closed: false,
            }));

            manager.sockets.insert(socket_handle.clone(), control.clone());
            (control, manager.manager_notify.clone())
        };

        TcpConnection {
            control,
            manager_notify,
        }
    }
}

impl AsyncRead for TcpConnection {
    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        let mut control = self.control.lock();

        // If socket is already closed, just return EOF directly.
        if control.is_closed {
            return Ok(()).into();
        }

        // Read from buffer

        if control.recv_buffer.is_empty() {
            // Nothing could be read. Wait for notify.
            if let Some(old_waker) = control.recv_waker.replace(cx.waker().clone()) {
                old_waker.wake();
            }

            return Poll::Pending;
        }

        let recv_buf = unsafe { mem::transmute::<_, &mut [u8]>(buf.unfilled_mut()) };
        let n = control.recv_buffer.dequeue_slice(recv_buf);
        buf.advance(n);

        self.manager_notify.notify_one();
        Ok(()).into()
    }
}

impl AsyncWrite for TcpConnection {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        let mut control = self.control.lock();
        if control.is_closed {
            return Err(io::ErrorKind::BrokenPipe.into()).into();
        }

        // Write to buffer

        if control.send_buffer.is_full() {
            if let Some(old_waker) = control.send_waker.replace(cx.waker().clone()) {
                old_waker.wake();
            }

            return Poll::Pending;
        }

        let n = control.send_buffer.enqueue_slice(buf);

        self.manager_notify.notify_one();
        Ok(n).into()
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Ok(()).into()
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        let mut control = self.control.lock();

        if control.is_closed {
            return Ok(()).into();
        }

        control.is_closed = true;
        if let Some(old_waker) = control.send_waker.replace(cx.waker().clone()) {
            old_waker.wake();
        }

        Poll::Pending
    }
}

pub struct TcpTun {
    context: Arc<ServiceContext>,
    manager: SharedTcpSocketManager,
    manager_handle: JoinHandle<()>,
    manager_notify: Arc<Notify>,
    balancer: PingBalancer,
    iface_rx: mpsc::UnboundedReceiver<Vec<u8>>,
    iface_tx: mpsc::Sender<Vec<u8>>,
}

impl Drop for TcpTun {
    fn drop(&mut self) {
        self.manager_handle.abort();
    }
}

impl TcpTun {
    pub fn new(context: Arc<ServiceContext>, balancer: PingBalancer, mtu: u32) -> TcpTun {
        let mut capabilities = DeviceCapabilities::default();
        capabilities.medium = Medium::Ip;
        capabilities.max_transmission_unit = mtu as usize;

        let (virt, iface_rx, iface_tx) = VirtTunDevice::new(capabilities);

        let iface_builder = InterfaceBuilder::new(virt, vec![]);
        let iface_ipaddrs = [
            IpCidr::new(IpAddress::v4(0, 0, 0, 1), 0),
            IpCidr::new(IpAddress::v6(0, 0, 0, 0, 0, 0, 0, 1), 0),
        ];
        let mut iface_routes = Routes::new(BTreeMap::new());
        iface_routes
            .add_default_ipv4_route(Ipv4Address::new(0, 0, 0, 1))
            .expect("IPv4 route");
        iface_routes
            .add_default_ipv6_route(Ipv6Address::new(0, 0, 0, 0, 0, 0, 0, 1))
            .expect("IPv6 route");
        let iface = iface_builder
            .any_ip(true)
            .ip_addrs(iface_ipaddrs)
            .routes(iface_routes)
            .finalize();

        let manager_notify = Arc::new(Notify::new());
        let manager = Arc::new(ParkingMutex::new(TcpSocketManager {
            iface,
            manager_notify: manager_notify.clone(),
            sockets: HashMap::new(),
        }));

        let manager_handle = {
            let manager = manager.clone();
            let manager_notify = manager_notify.clone();
            tokio::spawn(async move {
                loop {
                    let next_duration = {
                        let TcpSocketManager {
                            ref mut iface,
                            ref mut sockets,
                            ..
                        } = *(manager.lock());

                        let before_poll = Instant::now();
                        let updated_sockets = match iface.poll(before_poll) {
                            Ok(u) => u,
                            Err(err) => {
                                error!("VirtDevice::poll error: {}", err);
                                false
                            }
                        };

                        let after_poll = Instant::now();

                        if updated_sockets {
                            trace!("VirtDevice::poll costed {}", after_poll - before_poll);
                        }

                        // Check all the sockets' status
                        let mut sockets_to_remove = Vec::new();

                        for (socket_handle, control) in sockets.iter() {
                            let socket_handle = socket_handle.clone();
                            let socket = iface.get_socket::<TcpSocket>(socket_handle);
                            let mut control = control.lock();

                            #[inline]
                            fn close_socket_control(control: &mut TcpSocketControl) {
                                control.is_closed = true;
                                if let Some(waker) = control.send_waker.take() {
                                    waker.wake();
                                }
                                if let Some(waker) = control.recv_waker.take() {
                                    waker.wake();
                                }
                            }

                            if !socket.is_open() || socket.state() == TcpState::Closed {
                                sockets_to_remove.push(socket_handle);
                                close_socket_control(&mut *control);
                                continue;
                            }

                            if control.is_closed {
                                // Close the socket.
                                socket.close();
                                // sockets_to_remove.push(socket_handle);
                                // close_socket_control(&mut *control);
                                continue;
                            }

                            // Check if readable
                            if socket.can_recv() && !control.recv_buffer.is_full() {
                                let result = socket.recv(|buffer| {
                                    let n = control.recv_buffer.enqueue_slice(buffer);
                                    (n, ())
                                });

                                match result {
                                    Ok(..) => {
                                        if let Some(waker) = control.recv_waker.take() {
                                            waker.wake();
                                        }
                                    }
                                    Err(err) => {
                                        error!("socket recv error: {}", err);
                                        sockets_to_remove.push(socket_handle);
                                        close_socket_control(&mut *control);
                                    }
                                }
                            }

                            // Check if writable
                            if socket.can_send() && !control.send_buffer.is_empty() {
                                let result = socket.send(|buffer| {
                                    let n = control.send_buffer.dequeue_slice(buffer);
                                    (n, ())
                                });

                                match result {
                                    Ok(..) => {
                                        if let Some(waker) = control.send_waker.take() {
                                            waker.wake();
                                        }
                                    }
                                    Err(err) => {
                                        error!("socket send error: {}", err);
                                        sockets_to_remove.push(socket_handle);
                                        close_socket_control(&mut *control);
                                    }
                                }
                            }
                        }

                        for socket_handle in sockets_to_remove {
                            sockets.remove(&socket_handle);
                            iface.remove_socket(socket_handle);
                        }

                        let next_duration = iface.poll_delay(Instant::now()).unwrap_or(Duration::from_millis(50));

                        next_duration
                    };

                    tokio::select! {
                        _ = time::sleep(StdDuration::from(next_duration)) => {}
                        _ = manager_notify.notified() => {}
                    }
                }
            })
        };

        TcpTun {
            context,
            manager,
            manager_handle,
            manager_notify,
            balancer,
            iface_rx,
            iface_tx,
        }
    }

    pub async fn handle_packet(
        &mut self,
        src_addr: SocketAddr,
        dst_addr: SocketAddr,
        tcp_packet: &TcpPacket<&[u8]>,
    ) -> io::Result<()> {
        // TCP first handshake packet, create a new Connection
        if tcp_packet.syn() && !tcp_packet.ack() {
            let accept_opts = self.context.accept_opts();

            let send_buffer_size = accept_opts.tcp.send_buffer_size.unwrap_or(DEFAULT_TCP_SEND_BUFFER_SIZE);
            let recv_buffer_size = accept_opts.tcp.recv_buffer_size.unwrap_or(DEFAULT_TCP_RECV_BUFFER_SIZE);

            let mut socket = TcpSocket::new(
                TcpSocketBuffer::new(vec![0u8; recv_buffer_size as usize]),
                TcpSocketBuffer::new(vec![0u8; send_buffer_size as usize]),
            );
            socket.set_keep_alive(accept_opts.tcp.keepalive.map(From::from));

            if let Err(err) = socket.listen(dst_addr) {
                return Err(io::Error::new(ErrorKind::Other, err));
            }

            trace!("created TCP connection for {} <-> {}", src_addr, dst_addr);

            let connection = TcpConnection::new(socket, self.manager.clone(), &accept_opts.tcp);

            // establish a tunnel
            let context = self.context.clone();
            let balancer = self.balancer.clone();
            tokio::spawn(async move {
                if let Err(err) = handle_redir_client(context, balancer, connection, src_addr, dst_addr).await {
                    error!("TCP tunnel failure, {} <-> {}, error: {}", src_addr, dst_addr, err);
                }
            });
        }

        Ok(())
    }

    pub async fn drive_interface_state(&mut self, frame: &[u8]) {
        if let Err(..) = self.iface_tx.send(frame.to_vec()).await {
            panic!("interface send channel closed unexpectly");
        }

        // Wake up and poll the interface.
        self.manager_notify.notify_one();
    }

    pub async fn recv_packet(&mut self) -> Vec<u8> {
        match self.iface_rx.recv().await {
            Some(v) => v,
            None => unreachable!("channel closed unexpectedly"),
        }
    }
}

/// Established Client Transparent Proxy
///
/// This method must be called after handshaking with client (for example, socks5 handshaking)
async fn establish_client_tcp_redir<'a>(
    context: Arc<ServiceContext>,
    balancer: PingBalancer,
    mut stream: TcpConnection,
    peer_addr: SocketAddr,
    addr: &Address,
) -> io::Result<()> {
    let server = balancer.best_tcp_server();
    let svr_cfg = server.server_config();

    let mut remote = AutoProxyClientStream::connect(context, &server, addr).await?;

    establish_tcp_tunnel(svr_cfg, &mut stream, &mut remote, peer_addr, addr).await
}

async fn handle_redir_client(
    context: Arc<ServiceContext>,
    balancer: PingBalancer,
    s: TcpConnection,
    peer_addr: SocketAddr,
    mut daddr: SocketAddr,
) -> io::Result<()> {
    // Get forward address from socket
    //
    // Try to convert IPv4 mapped IPv6 address for dual-stack mode.
    if let SocketAddr::V6(ref a) = daddr {
        if let Some(v4) = to_ipv4_mapped(a.ip()) {
            daddr = SocketAddr::new(IpAddr::from(v4), a.port());
        }
    }
    let target_addr = Address::from(daddr);
    establish_client_tcp_redir(context, balancer, s, peer_addr, &target_addr).await
}
<commit_msg>copy send/recv buffer as much as possible<commit_after>use std::{
    collections::{BTreeMap, HashMap},
    io::{self, ErrorKind},
    mem,
    net::{IpAddr, SocketAddr},
    pin::Pin,
    sync::Arc,
    task::{Context, Poll, Waker},
    time::Duration as StdDuration,
};

use log::{error, trace};
use parking_lot::Mutex as ParkingMutex;
use shadowsocks::{net::TcpSocketOpts, relay::socks5::Address};
use smoltcp::{
    iface::{Interface, InterfaceBuilder, Routes, SocketHandle},
    phy::{DeviceCapabilities, Medium},
    socket::{TcpSocket, TcpSocketBuffer, TcpState},
    storage::RingBuffer,
    time::{Duration, Instant},
    wire::{IpAddress, IpCidr, Ipv4Address, Ipv6Address, TcpPacket},
};
use tokio::{
    io::{AsyncRead, AsyncWrite, ReadBuf},
    sync::{mpsc, Notify},
    task::JoinHandle,
    time,
};

use crate::local::{
    context::ServiceContext,
    loadbalancing::PingBalancer,
    net::AutoProxyClientStream,
    utils::{establish_tcp_tunnel, to_ipv4_mapped},
};

use super::virt_device::VirtTunDevice;

// NOTE: Default value is taken from Linux
// recv: /proc/sys/net/ipv4/tcp_rmem 87380 bytes
// send: /proc/sys/net/ipv4/tcp_wmem 16384 bytes
const DEFAULT_TCP_SEND_BUFFER_SIZE: u32 = 16384;
const DEFAULT_TCP_RECV_BUFFER_SIZE: u32 = 87380;

struct TcpSocketControl {
    send_buffer: RingBuffer<'static, u8>,
    send_waker: Option<Waker>,
    recv_buffer: RingBuffer<'static, u8>,
    recv_waker: Option<Waker>,
    is_closed: bool,
}

struct TcpSocketManager {
    iface: Interface<'static, VirtTunDevice>,
    manager_notify: Arc<Notify>,
    sockets: HashMap<SocketHandle, Arc<ParkingMutex<TcpSocketControl>>>,
}

type SharedTcpSocketManager = Arc<ParkingMutex<TcpSocketManager>>;

struct TcpConnection {
    control: Arc<ParkingMutex<TcpSocketControl>>,
    manager_notify: Arc<Notify>,
}

impl Drop for TcpConnection {
    fn drop(&mut self) {
        let mut control = self.control.lock();
        control.is_closed = true;
    }
}

impl TcpConnection {
    fn new(socket: TcpSocket<'static>, manager: SharedTcpSocketManager, tcp_opts: &TcpSocketOpts) -> TcpConnection {
        let send_buffer_size = tcp_opts.send_buffer_size.unwrap_or(DEFAULT_TCP_SEND_BUFFER_SIZE);
        let recv_buffer_size = tcp_opts.recv_buffer_size.unwrap_or(DEFAULT_TCP_RECV_BUFFER_SIZE);

        let (control, manager_notify) = {
            let mut manager = manager.lock();
            let socket_handle = manager.iface.add_socket(socket);

            let control = Arc::new(ParkingMutex::new(TcpSocketControl {
                send_buffer: RingBuffer::new(vec![0u8; send_buffer_size as usize]),
                send_waker: None,
                recv_buffer: RingBuffer::new(vec![0u8; recv_buffer_size as usize]),
                recv_waker: None,
                is_closed: false,
            }));

            manager.sockets.insert(socket_handle.clone(), control.clone());
            (control, manager.manager_notify.clone())
        };

        TcpConnection {
            control,
            manager_notify,
        }
    }
}

impl AsyncRead for TcpConnection {
    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        let mut control = self.control.lock();

        // If socket is already closed, just return EOF directly.
        if control.is_closed {
            return Ok(()).into();
        }

        // Read from buffer

        if control.recv_buffer.is_empty() {
            // Nothing could be read. Wait for notify.
            if let Some(old_waker) = control.recv_waker.replace(cx.waker().clone()) {
                old_waker.wake();
            }

            return Poll::Pending;
        }

        let recv_buf = unsafe { mem::transmute::<_, &mut [u8]>(buf.unfilled_mut()) };
        let n = control.recv_buffer.dequeue_slice(recv_buf);
        buf.advance(n);

        self.manager_notify.notify_one();
        Ok(()).into()
    }
}

impl AsyncWrite for TcpConnection {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        let mut control = self.control.lock();
        if control.is_closed {
            return Err(io::ErrorKind::BrokenPipe.into()).into();
        }

        // Write to buffer

        if control.send_buffer.is_full() {
            if let Some(old_waker) = control.send_waker.replace(cx.waker().clone()) {
                old_waker.wake();
            }

            return Poll::Pending;
        }

        let n = control.send_buffer.enqueue_slice(buf);

        self.manager_notify.notify_one();
        Ok(n).into()
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Ok(()).into()
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        let mut control = self.control.lock();

        if control.is_closed {
            return Ok(()).into();
        }

        control.is_closed = true;
        if let Some(old_waker) = control.send_waker.replace(cx.waker().clone()) {
            old_waker.wake();
        }

        Poll::Pending
    }
}

pub struct TcpTun {
    context: Arc<ServiceContext>,
    manager: SharedTcpSocketManager,
    manager_handle: JoinHandle<()>,
    manager_notify: Arc<Notify>,
    balancer: PingBalancer,
    iface_rx: mpsc::UnboundedReceiver<Vec<u8>>,
    iface_tx: mpsc::Sender<Vec<u8>>,
}

impl Drop for TcpTun {
    fn drop(&mut self) {
        self.manager_handle.abort();
    }
}

impl TcpTun {
    pub fn new(context: Arc<ServiceContext>, balancer: PingBalancer, mtu: u32) -> TcpTun {
        let mut capabilities = DeviceCapabilities::default();
        capabilities.medium = Medium::Ip;
        capabilities.max_transmission_unit = mtu as usize;

        let (virt, iface_rx, iface_tx) = VirtTunDevice::new(capabilities);

        let iface_builder = InterfaceBuilder::new(virt, vec![]);
        let iface_ipaddrs = [
            IpCidr::new(IpAddress::v4(0, 0, 0, 1), 0),
            IpCidr::new(IpAddress::v6(0, 0, 0, 0, 0, 0, 0, 1), 0),
        ];
        let mut iface_routes = Routes::new(BTreeMap::new());
        iface_routes
            .add_default_ipv4_route(Ipv4Address::new(0, 0, 0, 1))
            .expect("IPv4 route");
        iface_routes
            .add_default_ipv6_route(Ipv6Address::new(0, 0, 0, 0, 0, 0, 0, 1))
            .expect("IPv6 route");
        let iface = iface_builder
            .any_ip(true)
            .ip_addrs(iface_ipaddrs)
            .routes(iface_routes)
            .finalize();

        let manager_notify = Arc::new(Notify::new());
        let manager = Arc::new(ParkingMutex::new(TcpSocketManager {
            iface,
            manager_notify: manager_notify.clone(),
            sockets: HashMap::new(),
        }));

        let manager_handle = {
            let manager = manager.clone();
            let manager_notify = manager_notify.clone();
            tokio::spawn(async move {
                loop {
                    let next_duration = {
                        let TcpSocketManager {
                            ref mut iface,
                            ref mut sockets,
                            ..
                        } = *(manager.lock());

                        let before_poll = Instant::now();
                        let updated_sockets = match iface.poll(before_poll) {
                            Ok(u) => u,
                            Err(err) => {
                                error!("VirtDevice::poll error: {}", err);
                                false
                            }
                        };

                        let after_poll = Instant::now();

                        if updated_sockets {
                            trace!("VirtDevice::poll costed {}", after_poll - before_poll);
                        }

                        // Check all the sockets' status
                        let mut sockets_to_remove = Vec::new();

                        for (socket_handle, control) in sockets.iter() {
                            let socket_handle = socket_handle.clone();
                            let socket = iface.get_socket::<TcpSocket>(socket_handle);
                            let mut control = control.lock();

                            #[inline]
                            fn close_socket_control(control: &mut TcpSocketControl) {
                                control.is_closed = true;
                                if let Some(waker) = control.send_waker.take() {
                                    waker.wake();
                                }
                                if let Some(waker) = control.recv_waker.take() {
                                    waker.wake();
                                }
                            }

                            if !socket.is_open() || socket.state() == TcpState::Closed {
                                sockets_to_remove.push(socket_handle);
                                close_socket_control(&mut *control);
                                continue;
                            }

                            if control.is_closed {
                                // Close the socket.
                                socket.close();
                                // sockets_to_remove.push(socket_handle);
                                // close_socket_control(&mut *control);
                                continue;
                            }

                            // Check if readable
                            let mut has_received = false;
                            while socket.can_recv() && !control.recv_buffer.is_full() {
                                let result = socket.recv(|buffer| {
                                    let n = control.recv_buffer.enqueue_slice(buffer);
                                    (n, ())
                                });

                                match result {
                                    Ok(..) => {
                                        has_received = true;
                                    }
                                    Err(err) => {
                                        error!("socket recv error: {}", err);
                                        sockets_to_remove.push(socket_handle);
                                        close_socket_control(&mut *control);
                                        break;
                                    }
                                }
                            }

                            if has_received {
                                if let Some(waker) = control.recv_waker.take() {
                                    waker.wake();
                                }
                            }

                            // Check if writable
                            let mut has_sent = false;
                            while socket.can_send() && !control.send_buffer.is_empty() {
                                let result = socket.send(|buffer| {
                                    let n = control.send_buffer.dequeue_slice(buffer);
                                    (n, ())
                                });

                                match result {
                                    Ok(..) => {
                                        has_sent = true;
                                    }
                                    Err(err) => {
                                        error!("socket send error: {}", err);
                                        sockets_to_remove.push(socket_handle);
                                        close_socket_control(&mut *control);
                                        break;
                                    }
                                }
                            }

                            if has_sent {
                                if let Some(waker) = control.send_waker.take() {
                                    waker.wake();
                                }
                            }
                        }

                        for socket_handle in sockets_to_remove {
                            sockets.remove(&socket_handle);
                            iface.remove_socket(socket_handle);
                        }

                        let next_duration = iface.poll_delay(Instant::now()).unwrap_or(Duration::from_millis(50));

                        next_duration
                    };

                    tokio::select! {
                        _ = time::sleep(StdDuration::from(next_duration)) => {}
                        _ = manager_notify.notified() => {}
                    }
                }
            })
        };

        TcpTun {
            context,
            manager,
            manager_handle,
            manager_notify,
            balancer,
            iface_rx,
            iface_tx,
        }
    }

    pub async fn handle_packet(
        &mut self,
        src_addr: SocketAddr,
        dst_addr: SocketAddr,
        tcp_packet: &TcpPacket<&[u8]>,
    ) -> io::Result<()> {
        // TCP first handshake packet, create a new Connection
        if tcp_packet.syn() && !tcp_packet.ack() {
            let accept_opts = self.context.accept_opts();

            let send_buffer_size = accept_opts.tcp.send_buffer_size.unwrap_or(DEFAULT_TCP_SEND_BUFFER_SIZE);
            let recv_buffer_size = accept_opts.tcp.recv_buffer_size.unwrap_or(DEFAULT_TCP_RECV_BUFFER_SIZE);

            let mut socket = TcpSocket::new(
                TcpSocketBuffer::new(vec![0u8; recv_buffer_size as usize]),
                TcpSocketBuffer::new(vec![0u8; send_buffer_size as usize]),
            );
            socket.set_keep_alive(accept_opts.tcp.keepalive.map(From::from));

            if let Err(err) = socket.listen(dst_addr) {
                return Err(io::Error::new(ErrorKind::Other, err));
            }

            trace!("created TCP connection for {} <-> {}", src_addr, dst_addr);

            let connection = TcpConnection::new(socket, self.manager.clone(), &accept_opts.tcp);

            // establish a tunnel
            let context = self.context.clone();
            let balancer = self.balancer.clone();
            tokio::spawn(async move {
                if let Err(err) = handle_redir_client(context, balancer, connection, src_addr, dst_addr).await {
                    error!("TCP tunnel failure, {} <-> {}, error: {}", src_addr, dst_addr, err);
                }
            });
        }

        Ok(())
    }

    pub async fn drive_interface_state(&mut self, frame: &[u8]) {
        if let Err(..) = self.iface_tx.send(frame.to_vec()).await {
            panic!("interface send channel closed unexpectly");
        }

        // Wake up and poll the interface.
        self.manager_notify.notify_one();
    }

    pub async fn recv_packet(&mut self) -> Vec<u8> {
        match self.iface_rx.recv().await {
            Some(v) => v,
            None => unreachable!("channel closed unexpectedly"),
        }
    }
}

/// Established Client Transparent Proxy
///
/// This method must be called after handshaking with client (for example, socks5 handshaking)
async fn establish_client_tcp_redir<'a>(
    context: Arc<ServiceContext>,
    balancer: PingBalancer,
    mut stream: TcpConnection,
    peer_addr: SocketAddr,
    addr: &Address,
) -> io::Result<()> {
    let server = balancer.best_tcp_server();
    let svr_cfg = server.server_config();

    let mut remote = AutoProxyClientStream::connect(context, &server, addr).await?;

    establish_tcp_tunnel(svr_cfg, &mut stream, &mut remote, peer_addr, addr).await
}

async fn handle_redir_client(
    context: Arc<ServiceContext>,
    balancer: PingBalancer,
    s: TcpConnection,
    peer_addr: SocketAddr,
    mut daddr: SocketAddr,
) -> io::Result<()> {
    // Get forward address from socket
    //
    // Try to convert IPv4 mapped IPv6 address for dual-stack mode.
    if let SocketAddr::V6(ref a) = daddr {
        if let Some(v4) = to_ipv4_mapped(a.ip()) {
            daddr = SocketAddr::new(IpAddr::from(v4), a.port());
        }
    }
    let target_addr = Address::from(daddr);
    establish_client_tcp_redir(context, balancer, s, peer_addr, &target_addr).await
}
<|endoftext|>


<commit_before>use std::path::MAIN_SEPARATOR as SEP;
use support::{execs, project, ProjectBuilder};
use support::{COMPILING, RUNNING};
use hamcrest::{assert_that};

fn setup() {
}

fn verbose_output_for_target(lib: bool, p: &ProjectBuilder) -> String {
    let (target, kind) = match lib {
        true => ("lib", "lib"),
        false => ("main", "bin"),
    };
    format!("\
{compiling} {name} v{version} ({url})
{running} `rustc src{sep}{target}.rs --crate-name {name} --crate-type {kind} -g \
        --out-dir {dir}{sep}target{sep}debug \
        --emit=dep-info,link \
        -L dependency={dir}{sep}target{sep}debug \
        -L dependency={dir}{sep}target{sep}debug{sep}deps`
",
            running = RUNNING, compiling = COMPILING, sep = SEP,
            dir = p.root().display(), url = p.url(),
            target = target, kind = kind,
            name = "foo", version = "0.0.1")
}

fn verbose_output_for_target_with_args(lib: bool, p: &ProjectBuilder, args: &str) -> String {
    verbose_output_for_target(lib, p).replace(" -g ", &format!(" -g {} ", args))
}

test!(build_lib_for_foo {
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]

            name = "foo"
            version = "0.0.1"
            authors = ["wycats@example.com"]
        "#)
        .file("src/main.rs", r#"
            fn main() {}
        "#)
        .file("src/lib.rs", r#" "#);

    assert_that(p.cargo_process("rustc").arg("--lib").arg("-v"),
                execs()
                .with_status(0)
                .with_stdout(verbose_output_for_target(true, &p)));
});

test!(build_lib_and_allow_unstable_options {
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]

            name = "foo"
            version = "0.0.1"
            authors = ["wycats@example.com"]
        "#)
        .file("src/main.rs", r#"
            fn main() {}
        "#)
        .file("src/lib.rs", r#" "#);

    assert_that(p.cargo_process("rustc").arg("--lib").arg("-v")
                .arg("--").arg("-Z").arg("unstable-options"),
                execs()
                .with_status(0)
                .with_stdout(verbose_output_for_target_with_args(true, &p,
                                                                 "-Z unstable-options")));
});

test!(build_main_and_allow_unstable_options {
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]

            name = "foo"
            version = "0.0.1"
            authors = ["wycats@example.com"]
        "#)
        .file("src/main.rs", r#"
            fn main() {}
        "#)
        .file("src/lib.rs", r#" "#);

    assert_that(p.cargo_process("rustc").arg("-v").arg("--bin").arg("foo")
                .arg("--").arg("-Z").arg("unstable-options"),
                execs()
                .with_status(0)
                .with_stdout(&format!("\
{compiling} {name} v{version} ({url})
{running} `rustc src{sep}lib.rs --crate-name {name} --crate-type lib -g \
        --out-dir {dir}{sep}target{sep}debug \
        --emit=dep-info,link \
        -L dependency={dir}{sep}target{sep}debug \
        -L dependency={dir}{sep}target{sep}debug{sep}deps`
{running} `rustc src{sep}main.rs --crate-name {name} --crate-type bin -g \
        -Z unstable-options \
        --out-dir {dir}{sep}target{sep}debug \
        --emit=dep-info,link \
        -L dependency={dir}{sep}target{sep}debug \
        -L dependency={dir}{sep}target{sep}debug{sep}deps \
        --extern {name}={dir}{sep}target{sep}debug{sep}lib{name}.rlib`
",
            running = RUNNING, compiling = COMPILING, sep = SEP,
            dir = p.root().display(), url = p.url(),
            name = "foo", version = "0.0.1")));
});

test!(fails_when_trying_to_build_main_and_lib_with_args {
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]

            name = "foo"
            version = "0.0.1"
            authors = ["wycats@example.com"]
        "#)
        .file("src/main.rs", r#"
            fn main() {}
        "#)
        .file("src/lib.rs", r#" "#);

    assert_that(p.cargo_process("rustc").arg("-v")
                .arg("--").arg("-Z").arg("unstable-options"),
                execs()
                .with_status(101)
                .with_stderr("extra arguments to `rustc` can only be invoked for one target"));
});
<commit_msg>Cleanup tests - favor more explicit expected result instead of using helper functions<commit_after>use std::path::MAIN_SEPARATOR as SEP;
use support::{execs, project};
use support::{COMPILING, RUNNING};
use hamcrest::{assert_that};

fn setup() {
}

test!(build_lib_for_foo {
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]
            name = "foo"
            version = "0.0.1"
            authors = []
        "#)
        .file("src/main.rs", r#"
            fn main() {}
        "#)
        .file("src/lib.rs", r#" "#);

    assert_that(p.cargo_process("rustc").arg("--lib").arg("-v"),
                execs()
                .with_status(0)
                .with_stdout(format!("\
{compiling} foo v0.0.1 ({url})
{running} `rustc src{sep}lib.rs --crate-name foo --crate-type lib -g \
        --out-dir {dir}{sep}target{sep}debug \
        --emit=dep-info,link \
        -L dependency={dir}{sep}target{sep}debug \
        -L dependency={dir}{sep}target{sep}debug{sep}deps`
",
            running = RUNNING, compiling = COMPILING, sep = SEP,
            dir = p.root().display(), url = p.url())));
});

test!(build_lib_and_allow_unstable_options {
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]
            name = "foo"
            version = "0.0.1"
            authors = []
        "#)
        .file("src/main.rs", r#"
            fn main() {}
        "#)
        .file("src/lib.rs", r#" "#);

    assert_that(p.cargo_process("rustc").arg("--lib").arg("-v")
                .arg("--").arg("-Z").arg("unstable-options"),
                execs()
                .with_status(0)
                .with_stdout(format!("\
{compiling} foo v0.0.1 ({url})
{running} `rustc src{sep}lib.rs --crate-name foo --crate-type lib -g \
        -Z unstable-options \
        --out-dir {dir}{sep}target{sep}debug \
        --emit=dep-info,link \
        -L dependency={dir}{sep}target{sep}debug \
        -L dependency={dir}{sep}target{sep}debug{sep}deps`
",
            running = RUNNING, compiling = COMPILING, sep = SEP,
            dir = p.root().display(), url = p.url())))
});

test!(build_main_and_allow_unstable_options {
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]
            name = "foo"
            version = "0.0.1"
            authors = []
        "#)
        .file("src/main.rs", r#"
            fn main() {}
        "#)
        .file("src/lib.rs", r#" "#);

    assert_that(p.cargo_process("rustc").arg("-v").arg("--bin").arg("foo")
                .arg("--").arg("-Z").arg("unstable-options"),
                execs()
                .with_status(0)
                .with_stdout(&format!("\
{compiling} {name} v{version} ({url})
{running} `rustc src{sep}lib.rs --crate-name {name} --crate-type lib -g \
        --out-dir {dir}{sep}target{sep}debug \
        --emit=dep-info,link \
        -L dependency={dir}{sep}target{sep}debug \
        -L dependency={dir}{sep}target{sep}debug{sep}deps`
{running} `rustc src{sep}main.rs --crate-name {name} --crate-type bin -g \
        -Z unstable-options \
        --out-dir {dir}{sep}target{sep}debug \
        --emit=dep-info,link \
        -L dependency={dir}{sep}target{sep}debug \
        -L dependency={dir}{sep}target{sep}debug{sep}deps \
        --extern {name}={dir}{sep}target{sep}debug{sep}lib{name}.rlib`
",
            running = RUNNING, compiling = COMPILING, sep = SEP,
            dir = p.root().display(), url = p.url(),
            name = "foo", version = "0.0.1")));
});

test!(fails_when_trying_to_build_main_and_lib_with_args {
    let p = project("foo")
        .file("Cargo.toml", r#"
            [package]
            name = "foo"
            version = "0.0.1"
            authors = []
        "#)
        .file("src/main.rs", r#"
            fn main() {}
        "#)
        .file("src/lib.rs", r#" "#);

    assert_that(p.cargo_process("rustc").arg("-v")
                .arg("--").arg("-Z").arg("unstable-options"),
                execs()
                .with_status(101)
                .with_stderr("extra arguments to `rustc` can only be invoked for one target"));
});
<|endoftext|>


<commit_before>use nickel::{
    Params,
    BodyError,
};
use std;
use std::collections::HashMap;

macro_rules! validators {
    ( $([$($field:expr),*]: $vtype:ty => $($func:ident$args:tt),*; )+ ) => {
        let mut rules: Vec<ValidatorParams> = Vec::new();
        $(
            let fields = vec![$($field),*];
            let funcs: Vec<fn(ValidatorFuncArgs)> = vec![$($func),*];
            let validator_func_args = vec!( $(stringify!($args)),* );
            let validator_ref_type = stringify!($vtype);
            println!("{:#?}, {:#?}, {:#?}", validator_func_args, validator_ref_type, fields);
            for func in &funcs {
                let args = ValidatorFuncArgs::NoParams;
                println!("func:  [{:#?}]", func(args))
            }
            let rule = ValidatorParams {
                fields: fields,
                validators: funcs,
                type_ref: validator_ref_type,
                validators_arg: validator_func_args,
            };
            rules.push(rule);
        )+
        rules
    };
    ( @validator_args $(($(args:expr),*),)* ) => {
        $(
            vec!($(args),*)
        )*
    };
}

#[derive(Debug)]
pub struct Validators;

#[derive(Debug)]
pub enum ValidatorFuncArgs<'a> {
    NoParams,
    OneParam(&'a str),
    TwoParam(&'a str, &'a str),
}

pub struct ValidatorParams<'a> {
    pub fields: Vec<&'a str>,
    pub validators: Vec<fn(ValidatorFuncArgs)>,
    pub validators_arg: Vec<&'a str>,
    pub type_ref: &'a str,
}

pub type FormValidationErrors<'a> = HashMap<&'a str, Vec<FormValidationError>>;

#[derive(Debug, Clone)]
pub enum FormValidationError {
    FieldNotExist,
    ParseError
}

pub trait FormValidator {
    fn validate(form_data: &Params)
            -> Result<Self, FormValidationErrors>
            where Self: std::marker::Sized {
        let _validators = Self::validators();
        create_validator(&form_data);
        Ok(Self::fill_form(&form_data)?)
    }
    fn validators() -> Validators;
    fn fill_form(form_data: &Params)
            -> Result<Self, FormValidationErrors>
            where Self: std::marker::Sized;
}

fn create_validator(form_data: &Params) -> Result<(), FormValidationErrors> {
    let mut validation_errors: FormValidationErrors = HashMap::new();
    let field_name = "done.a";

    // Fetch fields withour any results
    // We iterate over all fields with `field_name`
    // Try parse with specific type
    // We gathering all possible errors, without skiping any posible error
    form_data.all(field_name)
        // If field not exist - add error
        .or_else(|| add_error(&mut validation_errors, field_name, FormValidationError::FieldNotExist) )
        .and_then(|v| {
            // Iterate over fields
            v.iter()
                .map(|x| {
                    // Type inference via Parse
                    x.parse::<i32>().ok()
                        // If parse Error - add error
                        .or_else(|| add_error(&mut validation_errors, field_name, FormValidationError::ParseError));
                    // Return same field
                    x
                }).cloned().collect::<String>();
            // Guaranty unwrap result
            Some(1)
        }).unwrap();
    println!("Errors: {:#?}", validation_errors);

    validators! (
        ["tst1", "tst2"]: i32 => required(), max();
        ["tst3", "tst4"]: i32 => max();
    );
    if validation_errors.is_empty() { Ok(()) } else { Err(validation_errors) }
}

fn add_error<'a, T>(validation_errors: &mut FormValidationErrors<'a>, field: &'a str, error: FormValidationError) -> Option<T> {
    if validation_errors.contains_key(&field) {
        let mut err = validation_errors.get_mut(&field).unwrap();
        err.push(error);
    } else {
        validation_errors.insert(&field, vec!(error));
    }
    None
}

fn required<'a>(_: ValidatorFuncArgs<'a>) {
    println!("required");
    // required template
}

fn max<'a>(_: ValidatorFuncArgs<'a>) {
    println!("max");
    // max template
}
<commit_msg>Added macro validator func args filling<commit_after>use nickel::{
    Params,
    BodyError,
};
use std;
use std::collections::HashMap;

macro_rules! validators {
    ( $([$($field:expr),*]: $vtype:ty => $($func:ident$args:tt),*; )+ ) => {
        let mut rules: Vec<ValidatorParams> = Vec::new();
        $(
            let fields = vec![$($field),*];
            let funcs: Vec<fn(ValidatorFuncArgs)> = vec![$($func),*];
            //let validator_func_args = vec!( $(stringify!($args)),* );
            let validator_func_args = vec!("-");
            let validator_ref_type = stringify!($vtype);
            println!("{:#?}, {:#?}, {:#?}", validator_func_args, validator_ref_type, fields);
            /*for func in &funcs {
                let args = ValidatorFuncArgs::NoParams;
                println!("func:  [{:#?}]", func(args))
            }*/
            let rule = ValidatorParams {
                fields: fields,
                validators: funcs,
                type_ref: validator_ref_type,
                validators_arg: validator_func_args,
            };
            rules.push(rule);
            validators! ( @validator_func $($args),* );
        )+
        rules
    };
    ( @validator_func $($args:tt),* ) => {
        $(
            println!("EXPR: {:#?}", validators! ( @validator_func_args $args ));
        )*
    };
    ( @validator_func_args () ) => {{
        ValidatorFuncArgs::NoParams
    }};
    ( @validator_func_args ($val:expr) ) => {{
        ValidatorFuncArgs::OneParam(&$val.to_string())
    }};
    ( @validator_func_args ($val1:expr, $val2:expr) ) => {{
        ValidatorFuncArgs::TwoParam(&$val1.to_string(), &$val2.to_string())
    }};
}

#[derive(Debug)]
pub struct Validators;

#[derive(Debug)]
pub enum ValidatorFuncArgs<'a> {
    NoParams,
    OneParam(&'a str),
    TwoParam(&'a str, &'a str),
}

pub struct ValidatorParams<'a> {
    pub fields: Vec<&'a str>,
    pub validators: Vec<fn(ValidatorFuncArgs)>,
    pub validators_arg: Vec<&'a str>,
    pub type_ref: &'a str,
}

pub type FormValidationErrors<'a> = HashMap<&'a str, Vec<FormValidationError>>;

#[derive(Debug, Clone)]
pub enum FormValidationError {
    FieldNotExist,
    ParseError
}

pub trait FormValidator {
    fn validate(form_data: &Params)
            -> Result<Self, FormValidationErrors>
            where Self: std::marker::Sized {
        let _validators = Self::validators();
        create_validator(&form_data);
        Ok(Self::fill_form(&form_data)?)
    }
    fn validators() -> Validators;
    fn fill_form(form_data: &Params)
            -> Result<Self, FormValidationErrors>
            where Self: std::marker::Sized;
}

fn create_validator(form_data: &Params) -> Result<(), FormValidationErrors> {
    let mut validation_errors: FormValidationErrors = HashMap::new();
    let field_name = "done.a";

    // Fetch fields withour any results
    // We iterate over all fields with `field_name`
    // Try parse with specific type
    // We gathering all possible errors, without skiping any posible error
    form_data.all(field_name)
        // If field not exist - add error
        .or_else(|| add_error(&mut validation_errors, field_name, FormValidationError::FieldNotExist) )
        .and_then(|v| {
            // Iterate over fields
            v.iter()
                .map(|x| {
                    // Type inference via Parse
                    x.parse::<i32>().ok()
                        // If parse Error - add error
                        .or_else(|| add_error(&mut validation_errors, field_name, FormValidationError::ParseError));
                    // Return same field
                    x
                }).cloned().collect::<String>();
            // Guaranty unwrap result
            Some(1)
        }).unwrap();
    println!("Errors: {:#?}", validation_errors);

    validators! (
        ["tst1", "tst2"]: i32 => required(), max(1);
        ["tst3", "tst4"]: i32 => max(1, "2+");
    );
    if validation_errors.is_empty() { Ok(()) } else { Err(validation_errors) }
}

fn add_error<'a, T>(validation_errors: &mut FormValidationErrors<'a>, field: &'a str, error: FormValidationError) -> Option<T> {
    if validation_errors.contains_key(&field) {
        let mut err = validation_errors.get_mut(&field).unwrap();
        err.push(error);
    } else {
        validation_errors.insert(&field, vec!(error));
    }
    None
}

fn required<'a>(_: ValidatorFuncArgs<'a>) {
    println!("required");
    // required template
}

fn max<'a>(_: ValidatorFuncArgs<'a>) {
    println!("max");
    // max template
}
<|endoftext|>


<commit_before>#![feature(test)]
extern crate test;

#[macro_use]
extern crate nom;

use nom::IResult;

use std::collections::HashMap;

fn is_alphabetic(chr: char) -> bool {
  (chr as u8 >= 0x41 && chr as u8 <= 0x5A) || (chr as u8 >= 0x61 && chr as u8 <= 0x7A)
}

fn is_digit(chr: char) -> bool {
  chr as u8 >= 0x30 && chr as u8 <= 0x39
}

fn is_alphanumeric(chr: char) -> bool {
  is_alphabetic(chr) || is_digit(chr)
}

fn is_space(chr: char) -> bool {
  chr == ' ' || chr == '\t'
}

fn is_line_ending_or_comment(chr: char) -> bool {
  chr == ';' || chr == '\n'
}

named!(alphanumeric<&str,&str>,         take_while_s!(is_alphanumeric));
named!(not_line_ending<&str,&str>,      is_not_s!("\r\n"));
named!(space<&str,&str>,                take_while_s!(is_space));
named!(space_or_line_ending<&str,&str>, is_a_s!(" \r\n"));

fn right_bracket(c: char) -> bool {
  c == ']'
}

named!(category     <&str, &str>,
  do_parse!(
          tag_s!("[")                 >>
    name: take_till_s!(right_bracket) >>
          tag_s!("]")                 >>
          opt!(space_or_line_ending)  >>
    (name)
  )
);

named!(key_value    <&str,(&str,&str)>,
  do_parse!(
    key: alphanumeric                              >>
         opt!(space)                               >>
         tag_s!("=")                               >>
         opt!(space)                               >>
    val: take_till_s!(is_line_ending_or_comment)   >>
         opt!(space)                               >>
         opt!(pair!(tag_s!(";"), not_line_ending)) >>
         opt!(space_or_line_ending)                >>
    (key, val)
  )
);

named!(keys_and_values_aggregator<&str, Vec<(&str,&str)> >, many0!(key_value));

fn keys_and_values(input: &str) -> IResult<&str, HashMap<&str, &str>> {
  match keys_and_values_aggregator(input) {
    Ok((i, tuple_vec)) => Ok((i, tuple_vec.into_iter().collect())),
    Err(e) => Err(e),
  }
}

named!(category_and_keys<&str,(&str,HashMap<&str,&str>)>,
  pair!(category, keys_and_values)
);

named!(categories_aggregator<&str, Vec<(&str, HashMap<&str,&str>)> >, many0!(category_and_keys));

fn categories(input: &str) -> IResult<&str, HashMap<&str, HashMap<&str, &str>>> {
  match categories_aggregator(input) {
    Ok((i, tuple_vec)) => Ok((i, tuple_vec.into_iter().collect())),
    Err(e) => Err(e),
  }
}

#[test]
fn parse_category_test() {
  let ini_file = "[category]

parameter=value
key = value2";

  let ini_without_category = "parameter=value
key = value2";

  let res = category(ini_file);
  println!("{:?}", res);
  match res {
    Ok((i, o)) => println!("i: {} | o: {:?}", i, o),
    _ => println!("error"),
  }

  assert_eq!(res, Ok((ini_without_category, "category")));
}

#[test]
fn parse_key_value_test() {
  let ini_file = "parameter=value
key = value2";

  let ini_without_key_value = "key = value2";

  let res = key_value(ini_file);
  println!("{:?}", res);
  match res {
    Ok((i, (o1, o2))) => println!("i: {} | o: ({:?},{:?})", i, o1, o2),
    _ => println!("error"),
  }

  assert_eq!(res, Ok((ini_without_key_value, ("parameter", "value"))));
}

#[test]
fn parse_key_value_with_space_test() {
  let ini_file = "parameter = value
key = value2";

  let ini_without_key_value = "key = value2";

  let res = key_value(ini_file);
  println!("{:?}", res);
  match res {
    Ok((i, (o1, o2))) => println!("i: {} | o: ({:?},{:?})", i, o1, o2),
    _ => println!("error"),
  }

  assert_eq!(res, Ok((ini_without_key_value, ("parameter", "value"))));
}

#[test]
fn parse_key_value_with_comment_test() {
  let ini_file = "parameter=value;abc
key = value2";

  let ini_without_key_value = "key = value2";

  let res = key_value(ini_file);
  println!("{:?}", res);
  match res {
    Ok((i, (o1, o2))) => println!("i: {} | o: ({:?},{:?})", i, o1, o2),
    _ => println!("error"),
  }

  assert_eq!(res, Ok((ini_without_key_value, ("parameter", "value"))));
}

#[test]
fn parse_multiple_keys_and_values_test() {
  let ini_file = "parameter=value;abc

key = value2

[category]";

  let ini_without_key_value = "[category]";

  let res = keys_and_values(ini_file);
  println!("{:?}", res);
  match res {
    Ok((i, ref o)) => println!("i: {} | o: {:?}", i, o),
    _ => println!("error"),
  }

  let mut expected: HashMap<&str, &str> = HashMap::new();
  expected.insert("parameter", "value");
  expected.insert("key", "value2");
  assert_eq!(res, Ok((ini_without_key_value, expected)));
}

#[test]
fn parse_category_then_multiple_keys_and_values_test() {
  //FIXME: there can be an empty line or a comment line after a category
  let ini_file = "[abcd]
parameter=value;abc

key = value2

[category]";

  let ini_after_parser = "[category]";

  let res = category_and_keys(ini_file);
  println!("{:?}", res);
  match res {
    Ok((i, ref o)) => println!("i: {} | o: {:?}", i, o),
    _ => println!("error"),
  }

  let mut expected_h: HashMap<&str, &str> = HashMap::new();
  expected_h.insert("parameter", "value");
  expected_h.insert("key", "value2");
  assert_eq!(res, Ok((ini_after_parser, ("abcd", expected_h))));
}

#[test]
fn parse_multiple_categories_test() {
  let ini_file = "[abcd]

parameter=value;abc

key = value2

[category]
parameter3=value3
key4 = value4
";

  let res = categories(ini_file);
  //println!("{:?}", res);
  match res {
    Ok((i, ref o)) => println!("i: {} | o: {:?}", i, o),
    _ => println!("error"),
  }

  let mut expected_1: HashMap<&str, &str> = HashMap::new();
  expected_1.insert("parameter", "value");
  expected_1.insert("key", "value2");
  let mut expected_2: HashMap<&str, &str> = HashMap::new();
  expected_2.insert("parameter3", "value3");
  expected_2.insert("key4", "value4");
  let mut expected_h: HashMap<&str, HashMap<&str, &str>> = HashMap::new();
  expected_h.insert("abcd", expected_1);
  expected_h.insert("category", expected_2);
  assert_eq!(res, Ok(("", expected_h)));
}

#[bench]
fn bench_ini_str(b: &mut test::Bencher) {
  let str = "[owner]
name=John Doe
organization=Acme Widgets Inc.

[database]
server=192.0.2.62
port=143
file=payroll.dat
";

  b.iter(|| categories(str).unwrap());
  b.bytes = str.len() as u64;
}
<commit_msg>use CompleteStr in the ini_str benchmark<commit_after>#![feature(test)]
extern crate test;

#[macro_use]
extern crate nom;

use nom::IResult;
use nom::types::CompleteStr;

use std::collections::HashMap;

fn is_alphabetic(chr: char) -> bool {
  (chr as u8 >= 0x41 && chr as u8 <= 0x5A) || (chr as u8 >= 0x61 && chr as u8 <= 0x7A)
}

fn is_digit(chr: char) -> bool {
  chr as u8 >= 0x30 && chr as u8 <= 0x39
}

fn is_alphanumeric(chr: char) -> bool {
  is_alphabetic(chr) || is_digit(chr)
}

fn is_space(chr: char) -> bool {
  chr == ' ' || chr == '\t'
}

fn is_line_ending_or_comment(chr: char) -> bool {
  chr == ';' || chr == '\n'
}

named!(alphanumeric<CompleteStr,CompleteStr>,         take_while_s!(is_alphanumeric));
named!(not_line_ending<CompleteStr,CompleteStr>,      is_not_s!("\r\n"));
named!(space<CompleteStr,CompleteStr>,                take_while_s!(is_space));
named!(space_or_line_ending<CompleteStr,CompleteStr>, is_a_s!(" \r\n"));

fn right_bracket(c: char) -> bool {
  c == ']'
}

named!(category     <CompleteStr, &str>,
  do_parse!(
          tag_s!("[")                 >>
    name: take_till_s!(right_bracket) >>
          tag_s!("]")                 >>
          opt!(space_or_line_ending)  >>
    (name.0)
  )
);

named!(key_value    <CompleteStr,(&str,&str)>,
  do_parse!(
    key: alphanumeric                              >>
         opt!(space)                               >>
         tag_s!("=")                               >>
         opt!(space)                               >>
    val: take_till_s!(is_line_ending_or_comment)   >>
         opt!(space)                               >>
         opt!(pair!(tag_s!(";"), not_line_ending)) >>
         opt!(space_or_line_ending)                >>
    (key.0, val.0)
  )
);

named!(keys_and_values_aggregator<CompleteStr, Vec<(&str, &str)> >, many0!(key_value));

fn keys_and_values(input: CompleteStr) -> IResult<CompleteStr, HashMap<&str, &str>> {
  match keys_and_values_aggregator(input) {
    Ok((i, tuple_vec)) => Ok((i, tuple_vec.into_iter().collect())),
    Err(e) => Err(e),
  }
}

named!(category_and_keys<CompleteStr,(&str,HashMap<&str,&str>)>,
  pair!(category, keys_and_values)
);

named!(categories_aggregator<CompleteStr, Vec<(&str, HashMap<&str,&str>)> >, many0!(category_and_keys));

fn categories(input: CompleteStr) -> IResult<CompleteStr, HashMap<&str, HashMap<&str, &str>>> {
  match categories_aggregator(input) {
    Ok((i, tuple_vec)) => Ok((i, tuple_vec.into_iter().collect())),
    Err(e) => Err(e),
  }
}

#[test]
fn parse_category_test() {
  let ini_file = CompleteStr(
    "[category]

parameter=value
key = value2",
  );

  let ini_without_category = CompleteStr(
    "parameter=value
key = value2",
  );

  let res = category(ini_file);
  println!("{:?}", res);
  match res {
    Ok((i, o)) => println!("i: {} | o: {:?}", i.0, o),
    _ => println!("error"),
  }

  assert_eq!(res, Ok((ini_without_category, "category")));
}

#[test]
fn parse_key_value_test() {
  let ini_file = CompleteStr(
    "parameter=value
key = value2",
  );

  let ini_without_key_value = CompleteStr("key = value2");

  let res = key_value(ini_file);
  println!("{:?}", res);
  match res {
    Ok((i, (o1, o2))) => println!("i: {} | o: ({:?},{:?})", i.0, o1, o2),
    _ => println!("error"),
  }

  assert_eq!(res, Ok((ini_without_key_value, ("parameter", "value"))));
}

#[test]
fn parse_key_value_with_space_test() {
  let ini_file = CompleteStr(
    "parameter = value
key = value2",
  );

  let ini_without_key_value = CompleteStr("key = value2");

  let res = key_value(ini_file);
  println!("{:?}", res);
  match res {
    Ok((i, (o1, o2))) => println!("i: {} | o: ({:?},{:?})", i.0, o1, o2),
    _ => println!("error"),
  }

  assert_eq!(res, Ok((ini_without_key_value, ("parameter", "value"))));
}

#[test]
fn parse_key_value_with_comment_test() {
  let ini_file = CompleteStr(
    "parameter=value;abc
key = value2",
  );

  let ini_without_key_value = CompleteStr("key = value2");

  let res = key_value(ini_file);
  println!("{:?}", res);
  match res {
    Ok((i, (o1, o2))) => println!("i: {} | o: ({:?},{:?})", i.0, o1, o2),
    _ => println!("error"),
  }

  assert_eq!(res, Ok((ini_without_key_value, ("parameter", "value"))));
}

#[test]
fn parse_multiple_keys_and_values_test() {
  let ini_file = CompleteStr(
    "parameter=value;abc

key = value2

[category]",
  );

  let ini_without_key_value = CompleteStr("[category]");

  let res = keys_and_values(ini_file);
  println!("{:?}", res);
  match res {
    Ok((i, ref o)) => println!("i: {} | o: {:?}", i.0, o),
    _ => println!("error"),
  }

  let mut expected: HashMap<&str, &str> = HashMap::new();
  expected.insert("parameter", "value");
  expected.insert("key", "value2");
  assert_eq!(res, Ok((ini_without_key_value, expected)));
}

#[test]
fn parse_category_then_multiple_keys_and_values_test() {
  //FIXME: there can be an empty line or a comment line after a category
  let ini_file = CompleteStr(
    "[abcd]
parameter=value;abc

key = value2

[category]",
  );

  let ini_after_parser = CompleteStr("[category]");

  let res = category_and_keys(ini_file);
  println!("{:?}", res);
  match res {
    Ok((i, ref o)) => println!("i: {} | o: {:?}", i.0, o),
    _ => println!("error"),
  }

  let mut expected_h: HashMap<&str, &str> = HashMap::new();
  expected_h.insert("parameter", "value");
  expected_h.insert("key", "value2");
  assert_eq!(res, Ok((ini_after_parser, ("abcd", expected_h))));
}

#[test]
fn parse_multiple_categories_test() {
  let ini_file = CompleteStr(
    "[abcd]

parameter=value;abc

key = value2

[category]
parameter3=value3
key4 = value4
",
  );

  let res = categories(ini_file);
  //println!("{:?}", res);
  match res {
    Ok((i, ref o)) => println!("i: {} | o: {:?}", i.0, o),
    _ => println!("error"),
  }

  let mut expected_1: HashMap<&str, &str> = HashMap::new();
  expected_1.insert("parameter", "value");
  expected_1.insert("key", "value2");
  let mut expected_2: HashMap<&str, &str> = HashMap::new();
  expected_2.insert("parameter3", "value3");
  expected_2.insert("key4", "value4");
  let mut expected_h: HashMap<&str, HashMap<&str, &str>> = HashMap::new();
  expected_h.insert("abcd", expected_1);
  expected_h.insert("category", expected_2);
  assert_eq!(res, Ok((CompleteStr(""), expected_h)));
}

#[bench]
fn bench_ini_str(b: &mut test::Bencher) {
  let s = "[owner]
name=John Doe
organization=Acme Widgets Inc.

[database]
server=192.0.2.62
port=143
file=payroll.dat
";

  b.iter(|| categories(CompleteStr(s)).unwrap());
  b.bytes = s.len() as u64;
}

<|endoftext|>


<commit_before>use criterion::{black_box, criterion_group, criterion_main, Bencher, Criterion};
use tokio::runtime::Runtime;

use zeromq::{prelude::*, PubSocket, SubSocket};

/// Binds n pubs and connects n subs to them
async fn setup(m_pubs: u8, n_subs: u8) -> (Vec<zeromq::PubSocket>, Vec<zeromq::SubSocket>) {
    let mut pubs = Vec::new();
    let mut subs: Vec<_> = (0..n_subs).map(|_| SubSocket::new()).collect();

    for _ in 0..m_pubs {
        let mut p = PubSocket::new();
        let bind_endpoint = p.bind("tcp://localhost:0").await.unwrap();
        for s in subs.iter_mut() {
            s.connect(bind_endpoint.clone()).await.unwrap();
        }
        pubs.push(p);
    }
    (pubs, subs)
}

fn criterion_benchmark(c: &mut Criterion) {
    let mut rt = Runtime::new().unwrap();
    let (mut pubs, mut subs) = rt.block_on(setup(4, 16));

    c.bench_function("m pubs and n subs messaging", |b| {
        b.iter(|| rt.block_on(f(&mut pubs, &mut subs)))
    });

    async fn f(pubs: &mut [PubSocket], subs: &mut [SubSocket]) {
        // Send/Recv messages
    }
}

criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);
<commit_msg>Added full pub_sub benchmark, but it currently hangs on recv<commit_after>use criterion::{black_box, criterion_group, criterion_main, Criterion};
use tokio::runtime::Runtime;

use zeromq::{prelude::*, PubSocket, SubSocket};

/// Binds n pubs and connects n subs to them
async fn setup(m_pubs: u8, n_subs: u8) -> (Vec<zeromq::PubSocket>, Vec<zeromq::SubSocket>) {
    let mut pubs = Vec::new();
    let mut subs: Vec<_> = (0..n_subs).map(|_| SubSocket::new()).collect();

    for _ in 0..m_pubs {
        let mut p = PubSocket::new();
        let bind_endpoint = p.bind("tcp://localhost:0").await.unwrap();
        println!("Bind endpoint: {}", bind_endpoint);
        for s in subs.iter_mut() {
            s.connect(bind_endpoint.clone()).await.unwrap();
        }
        pubs.push(p);
    }

    for s in subs.iter_mut() {
        s.subscribe("")
            .await
            .expect("Could not set subscription filter");
    }

    (pubs, subs)
}

fn criterion_benchmark(c: &mut Criterion) {
    let mut rt = Runtime::new().unwrap();

    const M_PUBS: u8 = 1;
    const N_SUBS: u8 = 16;
    let (pubs, subs) = rt.block_on(setup(M_PUBS, N_SUBS));
    assert_eq!(pubs.len(), usize::from(M_PUBS));
    assert_eq!(subs.len(), usize::from(N_SUBS));
    // Wrap in option to allow `f` to temporarily assume ownership
    let (mut pubs, mut subs) = (Some(pubs), Some(subs));

    c.bench_function("m pubs and n subs messaging", |b| {
        b.iter(|| rt.block_on(f(&mut pubs, &mut subs)))
    });

    async fn f(pubs: &mut Option<Vec<PubSocket>>, subs: &mut Option<Vec<SubSocket>>) {
        let owned_pubs = pubs.take().unwrap();
        let owned_subs = subs.take().unwrap();
        const MSGS_PER_PUB: u16 = 100;

        async fn send_msgs(mut pubs: Vec<PubSocket>) -> Vec<PubSocket> {
            for i in 0..MSGS_PER_PUB {
                for (j, p) in pubs.iter_mut().enumerate() {
                    p.send(format!("Sending {}th message from publisher {}", i, j).into())
                        .expect("Could not send message");
                }
            }
            println!("Sent all messages");
            pubs
        }
        async fn recv_msgs(mut subs: Vec<SubSocket>) -> Vec<SubSocket> {
            for i in 0..(u32::from(MSGS_PER_PUB) * u32::from(M_PUBS)) {
                for s in subs.iter_mut() {
                    println!("Attempting recv");
                    let msg = s.recv().await.expect("Could not recv");
                    black_box(msg);
                }
                println!("Received {}", i);
            }
            subs
        }
        println!("Starting send/recv loops");
        let owned_pubs = tokio::spawn(send_msgs(owned_pubs)).await.unwrap();
        tokio::time::delay_for(std::time::Duration::from_millis(1000)).await;
        let owned_subs = tokio::spawn(recv_msgs(owned_subs)).await.unwrap();
        // let (owned_pubs, owned_subs) = tokio::try_join!(pubs_handle,
        // subs_handle).unwrap();
        pubs.replace(owned_pubs);
        subs.replace(owned_subs);
    }
}

criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);
<|endoftext|>


<commit_before>//! A solver for dataflow problems.

use std::borrow::BorrowMut;
use std::ffi::OsString;
use std::fs;
use std::path::PathBuf;

use rustc_ast as ast;
use rustc_data_structures::work_queue::WorkQueue;
use rustc_graphviz as dot;
use rustc_hir::def_id::DefId;
use rustc_index::bit_set::BitSet;
use rustc_index::vec::{Idx, IndexVec};
use rustc_middle::mir::{self, traversal, BasicBlock};
use rustc_middle::ty::{self, TyCtxt};
use rustc_span::symbol::{sym, Symbol};

use super::fmt::DebugWithContext;
use super::graphviz;
use super::{
    visit_results, Analysis, Direction, GenKill, GenKillAnalysis, GenKillSet, JoinSemiLattice,
    ResultsCursor, ResultsVisitor,
};
use crate::util::pretty::dump_enabled;

/// A dataflow analysis that has converged to fixpoint.
pub struct Results<'tcx, A>
where
    A: Analysis<'tcx>,
{
    pub analysis: A,
    pub(super) entry_sets: IndexVec<BasicBlock, A::Domain>,
}

impl<A> Results<'tcx, A>
where
    A: Analysis<'tcx>,
{
    /// Creates a `ResultsCursor` that can inspect these `Results`.
    pub fn into_results_cursor(self, body: &'mir mir::Body<'tcx>) -> ResultsCursor<'mir, 'tcx, A> {
        ResultsCursor::new(body, self)
    }

    /// Gets the dataflow state for the given block.
    pub fn entry_set_for_block(&self, block: BasicBlock) -> &A::Domain {
        &self.entry_sets[block]
    }

    pub fn visit_with(
        &self,
        body: &'mir mir::Body<'tcx>,
        blocks: impl IntoIterator<Item = BasicBlock>,
        vis: &mut impl ResultsVisitor<'mir, 'tcx, FlowState = A::Domain>,
    ) {
        visit_results(body, blocks, self, vis)
    }

    pub fn visit_reachable_with(
        &self,
        body: &'mir mir::Body<'tcx>,
        vis: &mut impl ResultsVisitor<'mir, 'tcx, FlowState = A::Domain>,
    ) {
        let blocks = mir::traversal::reachable(body);
        visit_results(body, blocks.map(|(bb, _)| bb), self, vis)
    }

    pub fn visit_in_rpo_with(
        &self,
        body: &'mir mir::Body<'tcx>,
        vis: &mut impl ResultsVisitor<'mir, 'tcx, FlowState = A::Domain>,
    ) {
        let blocks = mir::traversal::reverse_postorder(body);
        visit_results(body, blocks.map(|(bb, _)| bb), self, vis)
    }
}

/// A solver for dataflow problems.
pub struct Engine<'a, 'tcx, A>
where
    A: Analysis<'tcx>,
{
    tcx: TyCtxt<'tcx>,
    body: &'a mir::Body<'tcx>,
    def_id: DefId,
    dead_unwinds: Option<&'a BitSet<BasicBlock>>,
    entry_sets: IndexVec<BasicBlock, A::Domain>,
    analysis: A,

    /// Cached, cumulative transfer functions for each block.
    apply_trans_for_block: Option<Box<dyn Fn(BasicBlock, &mut A::Domain)>>,
}

impl<A, D, T> Engine<'a, 'tcx, A>
where
    A: GenKillAnalysis<'tcx, Idx = T, Domain = D>,
    D: Clone + JoinSemiLattice + GenKill<T> + BorrowMut<BitSet<T>>,
    T: Idx,
{
    /// Creates a new `Engine` to solve a gen-kill dataflow problem.
    pub fn new_gen_kill(
        tcx: TyCtxt<'tcx>,
        body: &'a mir::Body<'tcx>,
        def_id: DefId,
        analysis: A,
    ) -> Self {
        // If there are no back-edges in the control-flow graph, we only ever need to apply the
        // transfer function for each block exactly once (assuming that we process blocks in RPO).
        //
        // In this case, there's no need to compute the block transfer functions ahead of time.
        if !body.is_cfg_cyclic() {
            return Self::new(tcx, body, def_id, analysis, None);
        }

        // Otherwise, compute and store the cumulative transfer function for each block.

        let identity = GenKillSet::identity(analysis.bottom_value(body).borrow().domain_size());
        let mut trans_for_block = IndexVec::from_elem(identity, body.basic_blocks());

        for (block, block_data) in body.basic_blocks().iter_enumerated() {
            let trans = &mut trans_for_block[block];
            A::Direction::gen_kill_effects_in_block(&analysis, trans, block, block_data);
        }

        let apply_trans = Box::new(move |bb: BasicBlock, state: &mut A::Domain| {
            trans_for_block[bb].apply(state.borrow_mut());
        });

        Self::new(tcx, body, def_id, analysis, Some(apply_trans as Box<_>))
    }
}

impl<A, D> Engine<'a, 'tcx, A>
where
    A: Analysis<'tcx, Domain = D>,
    D: Clone + JoinSemiLattice,
{
    /// Creates a new `Engine` to solve a dataflow problem with an arbitrary transfer
    /// function.
    ///
    /// Gen-kill problems should use `new_gen_kill`, which will coalesce transfer functions for
    /// better performance.
    pub fn new_generic(
        tcx: TyCtxt<'tcx>,
        body: &'a mir::Body<'tcx>,
        def_id: DefId,
        analysis: A,
    ) -> Self {
        Self::new(tcx, body, def_id, analysis, None)
    }

    fn new(
        tcx: TyCtxt<'tcx>,
        body: &'a mir::Body<'tcx>,
        def_id: DefId,
        analysis: A,
        apply_trans_for_block: Option<Box<dyn Fn(BasicBlock, &mut A::Domain)>>,
    ) -> Self {
        let bottom_value = analysis.bottom_value(body);
        let mut entry_sets = IndexVec::from_elem(bottom_value.clone(), body.basic_blocks());
        analysis.initialize_start_block(body, &mut entry_sets[mir::START_BLOCK]);

        if A::Direction::is_backward() && entry_sets[mir::START_BLOCK] != bottom_value {
            bug!("`initialize_start_block` is not yet supported for backward dataflow analyses");
        }

        Engine {
            analysis,
            tcx,
            body,
            def_id,
            dead_unwinds: None,
            entry_sets,
            apply_trans_for_block,
        }
    }

    /// Signals that we do not want dataflow state to propagate across unwind edges for these
    /// `BasicBlock`s.
    ///
    /// You must take care that `dead_unwinds` does not contain a `BasicBlock` that *can* actually
    /// unwind during execution. Otherwise, your dataflow results will not be correct.
    pub fn dead_unwinds(mut self, dead_unwinds: &'a BitSet<BasicBlock>) -> Self {
        self.dead_unwinds = Some(dead_unwinds);
        self
    }

    /// Computes the fixpoint for this dataflow problem and returns it.
    pub fn iterate_to_fixpoint(self) -> Results<'tcx, A>
    where
        A::Domain: DebugWithContext<A>,
    {
        let Engine {
            analysis,
            body,
            dead_unwinds,
            def_id,
            mut entry_sets,
            tcx,
            apply_trans_for_block,
            ..
        } = self;

        let mut dirty_queue: WorkQueue<BasicBlock> =
            WorkQueue::with_none(body.basic_blocks().len());

        if A::Direction::is_forward() {
            for (bb, _) in traversal::reverse_postorder(body) {
                dirty_queue.insert(bb);
            }
        } else {
            // Reverse post-order on the reverse CFG may generate a better iteration order for
            // backward dataflow analyses, but probably not enough to matter.
            for (bb, _) in traversal::postorder(body) {
                dirty_queue.insert(bb);
            }
        }

        let mut state = analysis.bottom_value(body);
        while let Some(bb) = dirty_queue.pop() {
            let bb_data = &body[bb];

            // Apply the block transfer function, using the cached one if it exists.
            state.clone_from(&entry_sets[bb]);
            match &apply_trans_for_block {
                Some(apply) => apply(bb, &mut state),
                None => A::Direction::apply_effects_in_block(&analysis, &mut state, bb, bb_data),
            }

            A::Direction::join_state_into_successors_of(
                &analysis,
                tcx,
                body,
                dead_unwinds,
                &mut state,
                (bb, bb_data),
                |target: BasicBlock, state: &A::Domain| {
                    let set_changed = entry_sets[target].join(state);
                    if set_changed {
                        dirty_queue.insert(target);
                    }
                },
            );
        }

        let results = Results { analysis, entry_sets };

        let res = write_graphviz_results(tcx, def_id, &body, &results);
        if let Err(e) = res {
            warn!("Failed to write graphviz dataflow results: {}", e);
        }

        results
    }
}

// Graphviz

/// Writes a DOT file containing the results of a dataflow analysis if the user requested it via
/// `rustc_mir` attributes.
fn write_graphviz_results<A>(
    tcx: TyCtxt<'tcx>,
    def_id: DefId,
    body: &mir::Body<'tcx>,
    results: &Results<'tcx, A>,
) -> std::io::Result<()>
where
    A: Analysis<'tcx>,
    A::Domain: DebugWithContext<A>,
{
    let attrs = match RustcMirAttrs::parse(tcx, def_id) {
        Ok(attrs) => attrs,

        // Invalid `rustc_mir` attrs are reported in `RustcMirAttrs::parse`
        Err(()) => return Ok(()),
    };

    let path = match attrs.output_path(A::NAME) {
        Some(path) => path,

        None if tcx.sess.opts.debugging_opts.dump_mir_dataflow
            && dump_enabled(tcx, A::NAME, def_id) =>
        {
            let mut path = PathBuf::from(&tcx.sess.opts.debugging_opts.dump_mir_dir);

            let item_name = ty::print::with_forced_impl_filename_line(|| {
                tcx.def_path(def_id).to_filename_friendly_no_crate()
            });
            path.push(format!("rustc.{}.{}.dot", item_name, A::NAME));
            path
        }

        None => return Ok(()),
    };

    let style = match attrs.formatter {
        Some(sym::two_phase) => graphviz::OutputStyle::BeforeAndAfter,
        _ => graphviz::OutputStyle::AfterOnly,
    };

    debug!("printing dataflow results for {:?} to {}", def_id, path.display());
    let mut buf = Vec::new();

    let graphviz = graphviz::Formatter::new(body, def_id, results, style);
    dot::render_opts(&graphviz, &mut buf, &[dot::RenderOption::Monospace])?;

    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent)?;
    }
    fs::write(&path, buf)?;

    Ok(())
}

#[derive(Default)]
struct RustcMirAttrs {
    basename_and_suffix: Option<PathBuf>,
    formatter: Option<Symbol>,
}

impl RustcMirAttrs {
    fn parse(tcx: TyCtxt<'tcx>, def_id: DefId) -> Result<Self, ()> {
        let attrs = tcx.get_attrs(def_id);

        let mut result = Ok(());
        let mut ret = RustcMirAttrs::default();

        let rustc_mir_attrs = attrs
            .iter()
            .filter(|attr| tcx.sess.check_name(attr, sym::rustc_mir))
            .flat_map(|attr| attr.meta_item_list().into_iter().flat_map(|v| v.into_iter()));

        for attr in rustc_mir_attrs {
            let attr_result = if attr.has_name(sym::borrowck_graphviz_postflow) {
                Self::set_field(&mut ret.basename_and_suffix, tcx, &attr, |s| {
                    let path = PathBuf::from(s.to_string());
                    match path.file_name() {
                        Some(_) => Ok(path),
                        None => {
                            tcx.sess.span_err(attr.span(), "path must end in a filename");
                            Err(())
                        }
                    }
                })
            } else if attr.has_name(sym::borrowck_graphviz_format) {
                Self::set_field(&mut ret.formatter, tcx, &attr, |s| match s {
                    sym::gen_kill | sym::two_phase => Ok(s),
                    _ => {
                        tcx.sess.span_err(attr.span(), "unknown formatter");
                        Err(())
                    }
                })
            } else {
                Ok(())
            };

            result = result.and(attr_result);
        }

        result.map(|()| ret)
    }

    fn set_field<T>(
        field: &mut Option<T>,
        tcx: TyCtxt<'tcx>,
        attr: &ast::NestedMetaItem,
        mapper: impl FnOnce(Symbol) -> Result<T, ()>,
    ) -> Result<(), ()> {
        if field.is_some() {
            tcx.sess
                .span_err(attr.span(), &format!("duplicate values for `{}`", attr.name_or_empty()));

            return Err(());
        }

        if let Some(s) = attr.value_str() {
            *field = Some(mapper(s)?);
            Ok(())
        } else {
            tcx.sess
                .span_err(attr.span(), &format!("`{}` requires an argument", attr.name_or_empty()));
            Err(())
        }
    }

    /// Returns the path where dataflow results should be written, or `None`
    /// `borrowck_graphviz_postflow` was not specified.
    ///
    /// This performs the following transformation to the argument of `borrowck_graphviz_postflow`:
    ///
    /// "path/suffix.dot" -> "path/analysis_name_suffix.dot"
    fn output_path(&self, analysis_name: &str) -> Option<PathBuf> {
        let mut ret = self.basename_and_suffix.as_ref().cloned()?;
        let suffix = ret.file_name().unwrap(); // Checked when parsing attrs

        let mut file_name: OsString = analysis_name.into();
        file_name.push("_");
        file_name.push(suffix);
        ret.set_file_name(file_name);

        Some(ret)
    }
}
<commit_msg>Add `FIXME` for faster cached block transfer functions<commit_after>//! A solver for dataflow problems.

use std::borrow::BorrowMut;
use std::ffi::OsString;
use std::fs;
use std::path::PathBuf;

use rustc_ast as ast;
use rustc_data_structures::work_queue::WorkQueue;
use rustc_graphviz as dot;
use rustc_hir::def_id::DefId;
use rustc_index::bit_set::BitSet;
use rustc_index::vec::{Idx, IndexVec};
use rustc_middle::mir::{self, traversal, BasicBlock};
use rustc_middle::ty::{self, TyCtxt};
use rustc_span::symbol::{sym, Symbol};

use super::fmt::DebugWithContext;
use super::graphviz;
use super::{
    visit_results, Analysis, Direction, GenKill, GenKillAnalysis, GenKillSet, JoinSemiLattice,
    ResultsCursor, ResultsVisitor,
};
use crate::util::pretty::dump_enabled;

/// A dataflow analysis that has converged to fixpoint.
pub struct Results<'tcx, A>
where
    A: Analysis<'tcx>,
{
    pub analysis: A,
    pub(super) entry_sets: IndexVec<BasicBlock, A::Domain>,
}

impl<A> Results<'tcx, A>
where
    A: Analysis<'tcx>,
{
    /// Creates a `ResultsCursor` that can inspect these `Results`.
    pub fn into_results_cursor(self, body: &'mir mir::Body<'tcx>) -> ResultsCursor<'mir, 'tcx, A> {
        ResultsCursor::new(body, self)
    }

    /// Gets the dataflow state for the given block.
    pub fn entry_set_for_block(&self, block: BasicBlock) -> &A::Domain {
        &self.entry_sets[block]
    }

    pub fn visit_with(
        &self,
        body: &'mir mir::Body<'tcx>,
        blocks: impl IntoIterator<Item = BasicBlock>,
        vis: &mut impl ResultsVisitor<'mir, 'tcx, FlowState = A::Domain>,
    ) {
        visit_results(body, blocks, self, vis)
    }

    pub fn visit_reachable_with(
        &self,
        body: &'mir mir::Body<'tcx>,
        vis: &mut impl ResultsVisitor<'mir, 'tcx, FlowState = A::Domain>,
    ) {
        let blocks = mir::traversal::reachable(body);
        visit_results(body, blocks.map(|(bb, _)| bb), self, vis)
    }

    pub fn visit_in_rpo_with(
        &self,
        body: &'mir mir::Body<'tcx>,
        vis: &mut impl ResultsVisitor<'mir, 'tcx, FlowState = A::Domain>,
    ) {
        let blocks = mir::traversal::reverse_postorder(body);
        visit_results(body, blocks.map(|(bb, _)| bb), self, vis)
    }
}

/// A solver for dataflow problems.
pub struct Engine<'a, 'tcx, A>
where
    A: Analysis<'tcx>,
{
    tcx: TyCtxt<'tcx>,
    body: &'a mir::Body<'tcx>,
    def_id: DefId,
    dead_unwinds: Option<&'a BitSet<BasicBlock>>,
    entry_sets: IndexVec<BasicBlock, A::Domain>,
    analysis: A,

    /// Cached, cumulative transfer functions for each block.
    //
    // FIXME(ecstaticmorse): This boxed `Fn` trait object is invoked inside a tight loop for
    // gen/kill problems on cyclic CFGs. This is not ideal, but it doesn't seem to degrade
    // performance in practice. I've tried a few ways to avoid this, but they have downsides. See
    // the message for the commit that added this FIXME for more information.
    apply_trans_for_block: Option<Box<dyn Fn(BasicBlock, &mut A::Domain)>>,
}

impl<A, D, T> Engine<'a, 'tcx, A>
where
    A: GenKillAnalysis<'tcx, Idx = T, Domain = D>,
    D: Clone + JoinSemiLattice + GenKill<T> + BorrowMut<BitSet<T>>,
    T: Idx,
{
    /// Creates a new `Engine` to solve a gen-kill dataflow problem.
    pub fn new_gen_kill(
        tcx: TyCtxt<'tcx>,
        body: &'a mir::Body<'tcx>,
        def_id: DefId,
        analysis: A,
    ) -> Self {
        // If there are no back-edges in the control-flow graph, we only ever need to apply the
        // transfer function for each block exactly once (assuming that we process blocks in RPO).
        //
        // In this case, there's no need to compute the block transfer functions ahead of time.
        if !body.is_cfg_cyclic() {
            return Self::new(tcx, body, def_id, analysis, None);
        }

        // Otherwise, compute and store the cumulative transfer function for each block.

        let identity = GenKillSet::identity(analysis.bottom_value(body).borrow().domain_size());
        let mut trans_for_block = IndexVec::from_elem(identity, body.basic_blocks());

        for (block, block_data) in body.basic_blocks().iter_enumerated() {
            let trans = &mut trans_for_block[block];
            A::Direction::gen_kill_effects_in_block(&analysis, trans, block, block_data);
        }

        let apply_trans = Box::new(move |bb: BasicBlock, state: &mut A::Domain| {
            trans_for_block[bb].apply(state.borrow_mut());
        });

        Self::new(tcx, body, def_id, analysis, Some(apply_trans as Box<_>))
    }
}

impl<A, D> Engine<'a, 'tcx, A>
where
    A: Analysis<'tcx, Domain = D>,
    D: Clone + JoinSemiLattice,
{
    /// Creates a new `Engine` to solve a dataflow problem with an arbitrary transfer
    /// function.
    ///
    /// Gen-kill problems should use `new_gen_kill`, which will coalesce transfer functions for
    /// better performance.
    pub fn new_generic(
        tcx: TyCtxt<'tcx>,
        body: &'a mir::Body<'tcx>,
        def_id: DefId,
        analysis: A,
    ) -> Self {
        Self::new(tcx, body, def_id, analysis, None)
    }

    fn new(
        tcx: TyCtxt<'tcx>,
        body: &'a mir::Body<'tcx>,
        def_id: DefId,
        analysis: A,
        apply_trans_for_block: Option<Box<dyn Fn(BasicBlock, &mut A::Domain)>>,
    ) -> Self {
        let bottom_value = analysis.bottom_value(body);
        let mut entry_sets = IndexVec::from_elem(bottom_value.clone(), body.basic_blocks());
        analysis.initialize_start_block(body, &mut entry_sets[mir::START_BLOCK]);

        if A::Direction::is_backward() && entry_sets[mir::START_BLOCK] != bottom_value {
            bug!("`initialize_start_block` is not yet supported for backward dataflow analyses");
        }

        Engine {
            analysis,
            tcx,
            body,
            def_id,
            dead_unwinds: None,
            entry_sets,
            apply_trans_for_block,
        }
    }

    /// Signals that we do not want dataflow state to propagate across unwind edges for these
    /// `BasicBlock`s.
    ///
    /// You must take care that `dead_unwinds` does not contain a `BasicBlock` that *can* actually
    /// unwind during execution. Otherwise, your dataflow results will not be correct.
    pub fn dead_unwinds(mut self, dead_unwinds: &'a BitSet<BasicBlock>) -> Self {
        self.dead_unwinds = Some(dead_unwinds);
        self
    }

    /// Computes the fixpoint for this dataflow problem and returns it.
    pub fn iterate_to_fixpoint(self) -> Results<'tcx, A>
    where
        A::Domain: DebugWithContext<A>,
    {
        let Engine {
            analysis,
            body,
            dead_unwinds,
            def_id,
            mut entry_sets,
            tcx,
            apply_trans_for_block,
            ..
        } = self;

        let mut dirty_queue: WorkQueue<BasicBlock> =
            WorkQueue::with_none(body.basic_blocks().len());

        if A::Direction::is_forward() {
            for (bb, _) in traversal::reverse_postorder(body) {
                dirty_queue.insert(bb);
            }
        } else {
            // Reverse post-order on the reverse CFG may generate a better iteration order for
            // backward dataflow analyses, but probably not enough to matter.
            for (bb, _) in traversal::postorder(body) {
                dirty_queue.insert(bb);
            }
        }

        let mut state = analysis.bottom_value(body);
        while let Some(bb) = dirty_queue.pop() {
            let bb_data = &body[bb];

            // Apply the block transfer function, using the cached one if it exists.
            state.clone_from(&entry_sets[bb]);
            match &apply_trans_for_block {
                Some(apply) => apply(bb, &mut state),
                None => A::Direction::apply_effects_in_block(&analysis, &mut state, bb, bb_data),
            }

            A::Direction::join_state_into_successors_of(
                &analysis,
                tcx,
                body,
                dead_unwinds,
                &mut state,
                (bb, bb_data),
                |target: BasicBlock, state: &A::Domain| {
                    let set_changed = entry_sets[target].join(state);
                    if set_changed {
                        dirty_queue.insert(target);
                    }
                },
            );
        }

        let results = Results { analysis, entry_sets };

        let res = write_graphviz_results(tcx, def_id, &body, &results);
        if let Err(e) = res {
            warn!("Failed to write graphviz dataflow results: {}", e);
        }

        results
    }
}

// Graphviz

/// Writes a DOT file containing the results of a dataflow analysis if the user requested it via
/// `rustc_mir` attributes.
fn write_graphviz_results<A>(
    tcx: TyCtxt<'tcx>,
    def_id: DefId,
    body: &mir::Body<'tcx>,
    results: &Results<'tcx, A>,
) -> std::io::Result<()>
where
    A: Analysis<'tcx>,
    A::Domain: DebugWithContext<A>,
{
    let attrs = match RustcMirAttrs::parse(tcx, def_id) {
        Ok(attrs) => attrs,

        // Invalid `rustc_mir` attrs are reported in `RustcMirAttrs::parse`
        Err(()) => return Ok(()),
    };

    let path = match attrs.output_path(A::NAME) {
        Some(path) => path,

        None if tcx.sess.opts.debugging_opts.dump_mir_dataflow
            && dump_enabled(tcx, A::NAME, def_id) =>
        {
            let mut path = PathBuf::from(&tcx.sess.opts.debugging_opts.dump_mir_dir);

            let item_name = ty::print::with_forced_impl_filename_line(|| {
                tcx.def_path(def_id).to_filename_friendly_no_crate()
            });
            path.push(format!("rustc.{}.{}.dot", item_name, A::NAME));
            path
        }

        None => return Ok(()),
    };

    let style = match attrs.formatter {
        Some(sym::two_phase) => graphviz::OutputStyle::BeforeAndAfter,
        _ => graphviz::OutputStyle::AfterOnly,
    };

    debug!("printing dataflow results for {:?} to {}", def_id, path.display());
    let mut buf = Vec::new();

    let graphviz = graphviz::Formatter::new(body, def_id, results, style);
    dot::render_opts(&graphviz, &mut buf, &[dot::RenderOption::Monospace])?;

    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent)?;
    }
    fs::write(&path, buf)?;

    Ok(())
}

#[derive(Default)]
struct RustcMirAttrs {
    basename_and_suffix: Option<PathBuf>,
    formatter: Option<Symbol>,
}

impl RustcMirAttrs {
    fn parse(tcx: TyCtxt<'tcx>, def_id: DefId) -> Result<Self, ()> {
        let attrs = tcx.get_attrs(def_id);

        let mut result = Ok(());
        let mut ret = RustcMirAttrs::default();

        let rustc_mir_attrs = attrs
            .iter()
            .filter(|attr| tcx.sess.check_name(attr, sym::rustc_mir))
            .flat_map(|attr| attr.meta_item_list().into_iter().flat_map(|v| v.into_iter()));

        for attr in rustc_mir_attrs {
            let attr_result = if attr.has_name(sym::borrowck_graphviz_postflow) {
                Self::set_field(&mut ret.basename_and_suffix, tcx, &attr, |s| {
                    let path = PathBuf::from(s.to_string());
                    match path.file_name() {
                        Some(_) => Ok(path),
                        None => {
                            tcx.sess.span_err(attr.span(), "path must end in a filename");
                            Err(())
                        }
                    }
                })
            } else if attr.has_name(sym::borrowck_graphviz_format) {
                Self::set_field(&mut ret.formatter, tcx, &attr, |s| match s {
                    sym::gen_kill | sym::two_phase => Ok(s),
                    _ => {
                        tcx.sess.span_err(attr.span(), "unknown formatter");
                        Err(())
                    }
                })
            } else {
                Ok(())
            };

            result = result.and(attr_result);
        }

        result.map(|()| ret)
    }

    fn set_field<T>(
        field: &mut Option<T>,
        tcx: TyCtxt<'tcx>,
        attr: &ast::NestedMetaItem,
        mapper: impl FnOnce(Symbol) -> Result<T, ()>,
    ) -> Result<(), ()> {
        if field.is_some() {
            tcx.sess
                .span_err(attr.span(), &format!("duplicate values for `{}`", attr.name_or_empty()));

            return Err(());
        }

        if let Some(s) = attr.value_str() {
            *field = Some(mapper(s)?);
            Ok(())
        } else {
            tcx.sess
                .span_err(attr.span(), &format!("`{}` requires an argument", attr.name_or_empty()));
            Err(())
        }
    }

    /// Returns the path where dataflow results should be written, or `None`
    /// `borrowck_graphviz_postflow` was not specified.
    ///
    /// This performs the following transformation to the argument of `borrowck_graphviz_postflow`:
    ///
    /// "path/suffix.dot" -> "path/analysis_name_suffix.dot"
    fn output_path(&self, analysis_name: &str) -> Option<PathBuf> {
        let mut ret = self.basename_and_suffix.as_ref().cloned()?;
        let suffix = ret.file_name().unwrap(); // Checked when parsing attrs

        let mut file_name: OsString = analysis_name.into();
        file_name.push("_");
        file_name.push(suffix);
        ret.set_file_name(file_name);

        Some(ret)
    }
}
<|endoftext|>


<commit_before>use crate::sync::mpsc;

use futures::future::poll_fn;
use loom::future::block_on;
use loom::thread;

#[test]
fn closing_tx() {
    loom::model(|| {
        let (mut tx, mut rx) = mpsc::channel(16);

        thread::spawn(move || {
            tx.try_send(()).unwrap();
            drop(tx);
        });

        let v = block_on(poll_fn(|cx| rx.poll_recv(cx)));
        assert!(v.is_some());

        let v = block_on(poll_fn(|cx| rx.poll_recv(cx)));
        assert!(v.is_none());
    });
}
<commit_msg>Test some more mpsc behavior with loom (#2246)<commit_after>use crate::sync::mpsc;

use futures::future::poll_fn;
use loom::future::block_on;
use loom::thread;

#[test]
fn closing_tx() {
    loom::model(|| {
        let (mut tx, mut rx) = mpsc::channel(16);

        thread::spawn(move || {
            tx.try_send(()).unwrap();
            drop(tx);
        });

        let v = block_on(poll_fn(|cx| rx.poll_recv(cx)));
        assert!(v.is_some());

        let v = block_on(poll_fn(|cx| rx.poll_recv(cx)));
        assert!(v.is_none());
    });
}

#[test]
fn closing_unbounded_tx() {
    loom::model(|| {
        let (tx, mut rx) = mpsc::unbounded_channel();

        thread::spawn(move || {
            tx.send(()).unwrap();
            drop(tx);
        });

        let v = block_on(poll_fn(|cx| rx.poll_recv(cx)));
        assert!(v.is_some());

        let v = block_on(poll_fn(|cx| rx.poll_recv(cx)));
        assert!(v.is_none());
    });
}

#[test]
fn dropping_tx() {
    loom::model(|| {
        let (tx, mut rx) = mpsc::channel::<()>(16);

        for _ in 0..2 {
            let tx = tx.clone();
            thread::spawn(move || {
                drop(tx);
            });
        }
        drop(tx);

        let v = block_on(poll_fn(|cx| rx.poll_recv(cx)));
        assert!(v.is_none());
    });
}

#[test]
fn dropping_unbounded_tx() {
    loom::model(|| {
        let (tx, mut rx) = mpsc::unbounded_channel::<()>();

        for _ in 0..2 {
            let tx = tx.clone();
            thread::spawn(move || {
                drop(tx);
            });
        }
        drop(tx);

        let v = block_on(poll_fn(|cx| rx.poll_recv(cx)));
        assert!(v.is_none());
    });
}
<|endoftext|>