<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
                            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
}
<commit_msg>TCP tun poll_* may wake the same task when replacing waker<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()) {
                if !old_waker.will_wake(cx.waker()) {
                    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()) {
                if !old_waker.will_wake(cx.waker()) {
                    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()) {
            if !old_waker.will_wake(cx.waker()) {
                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>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]
use std::{cmp::min, env};

use crate::{
    cluster::Cluster,
    experiments::{
        CompatiblityTestParams, CpuFlamegraphParams, Experiment, ExperimentParam,
        PerformanceBenchmarkParams, PerformanceBenchmarkThreeRegionSimulationParams,
        RebootRandomValidatorsParams, RecoveryTimeParams,
    },
};
use anyhow::{format_err, Result};

pub struct ExperimentSuite {
    pub experiments: Vec<Box<dyn Experiment>>,
}

impl ExperimentSuite {
    fn new_pre_release(cluster: &Cluster) -> Self {
        let mut experiments: Vec<Box<dyn Experiment>> = vec![];
        if env::var("RECOVERY_EXP").is_ok() {
            experiments.push(Box::new(
                RecoveryTimeParams {
                    num_accounts_to_mint: 100_000,
                }
                .build(cluster),
            ));
        }
        let count = min(3, cluster.validator_instances().len() / 3);
        // Reboot different sets of 3 validators *100 times
        for _ in 0..10 {
            let b = Box::new(RebootRandomValidatorsParams::new(count, 0).build(cluster));
            experiments.push(b);
        }
        experiments.push(Box::new(
            PerformanceBenchmarkParams::new_nodes_down(0).build(cluster),
        ));
        experiments.push(Box::new(
            PerformanceBenchmarkParams::new_nodes_down(10).build(cluster),
        ));
        experiments.push(Box::new(
            PerformanceBenchmarkThreeRegionSimulationParams {}.build(cluster),
        ));
        experiments.push(Box::new(
            PerformanceBenchmarkParams::new_fixed_tps(0, 10).build(cluster),
        ));
        // TODO, comment it out to unblock CI during investigation
        // experiments.push(Box::new(TwinValidatorsParams { pair: 1 }.build(cluster)));
        experiments.push(Box::new(
            CpuFlamegraphParams { duration_secs: 60 }.build(cluster),
        ));
        Self { experiments }
    }

    fn new_perf_suite(cluster: &Cluster) -> Self {
        let mut experiments: Vec<Box<dyn Experiment>> = vec![];
        experiments.push(Box::new(
            PerformanceBenchmarkParams::new_nodes_down(0).build(cluster),
        ));
        experiments.push(Box::new(
            PerformanceBenchmarkParams::new_nodes_down(10).build(cluster),
        ));
        experiments.push(Box::new(
            PerformanceBenchmarkThreeRegionSimulationParams {}.build(cluster),
        ));
        experiments.push(Box::new(
            PerformanceBenchmarkParams::new_fixed_tps(0, 10).build(cluster),
        ));
        Self { experiments }
    }

    fn new_land_blocking_suite(cluster: &Cluster) -> Self {
        let mut experiments: Vec<Box<dyn Experiment>> = vec![];
        experiments.push(Box::new(
            PerformanceBenchmarkParams::new_nodes_down(0).build(cluster),
        ));
        Self { experiments }
    }

    fn new_land_blocking_compat_suite(cluster: &Cluster) -> Result<Self> {
        let count: usize = match env::var("BATCH_SIZE") {
            Ok(val) => val
                .parse()
                .map_err(|e| format_err!("Failed to parse BATCH_SIZE {}: {}", val, e))?,
            Err(_) => cluster.validator_instances().len() / 2,
        };
        let updated_image_tag = env::var("UPDATE_TO_TAG")
            .map_err(|_| format_err!("Expected environment variable UPDATE_TO_TAG"))?;
        let mut experiments: Vec<Box<dyn Experiment>> = vec![];
        experiments.push(Box::new(
            CompatiblityTestParams {
                count,
                updated_image_tag,
            }
            .build(cluster),
        ));
        experiments.extend(Self::new_land_blocking_suite(cluster).experiments);
        Ok(Self { experiments })
    }

    pub fn new_by_name(cluster: &Cluster, name: &str) -> Result<Self> {
        match name {
            "perf" => Ok(Self::new_perf_suite(cluster)),
            "pre_release" => Ok(Self::new_pre_release(cluster)),
            "land_blocking" => Ok(Self::new_land_blocking_suite(cluster)),
            "land_blocking_compat" => Self::new_land_blocking_compat_suite(cluster),
            other => Err(format_err!("Unknown suite: {}", other)),
        }
    }
}
<commit_msg>[cluster-test] Added twin test back to pre-release suite<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]
use std::{cmp::min, env};

use crate::{
    cluster::Cluster,
    experiments::{
        CompatiblityTestParams, CpuFlamegraphParams, Experiment, ExperimentParam,
        PerformanceBenchmarkParams, PerformanceBenchmarkThreeRegionSimulationParams,
        RebootRandomValidatorsParams, RecoveryTimeParams, TwinValidatorsParams,
    },
};
use anyhow::{format_err, Result};

pub struct ExperimentSuite {
    pub experiments: Vec<Box<dyn Experiment>>,
}

impl ExperimentSuite {
    fn new_pre_release(cluster: &Cluster) -> Self {
        let mut experiments: Vec<Box<dyn Experiment>> = vec![];
        if env::var("RECOVERY_EXP").is_ok() {
            experiments.push(Box::new(
                RecoveryTimeParams {
                    num_accounts_to_mint: 100_000,
                }
                .build(cluster),
            ));
        }
        let count = min(3, cluster.validator_instances().len() / 3);
        // Reboot different sets of 3 validators *100 times
        for _ in 0..10 {
            let b = Box::new(RebootRandomValidatorsParams::new(count, 0).build(cluster));
            experiments.push(b);
        }
        experiments.push(Box::new(
            PerformanceBenchmarkParams::new_nodes_down(0).build(cluster),
        ));
        experiments.push(Box::new(
            PerformanceBenchmarkParams::new_nodes_down(10).build(cluster),
        ));
        experiments.push(Box::new(
            PerformanceBenchmarkThreeRegionSimulationParams {}.build(cluster),
        ));
        experiments.push(Box::new(
            PerformanceBenchmarkParams::new_fixed_tps(0, 10).build(cluster),
        ));
        experiments.push(Box::new(TwinValidatorsParams { pair: 1 }.build(cluster)));
        experiments.push(Box::new(
            CpuFlamegraphParams { duration_secs: 60 }.build(cluster),
        ));
        Self { experiments }
    }

    fn new_perf_suite(cluster: &Cluster) -> Self {
        let mut experiments: Vec<Box<dyn Experiment>> = vec![];
        experiments.push(Box::new(
            PerformanceBenchmarkParams::new_nodes_down(0).build(cluster),
        ));
        experiments.push(Box::new(
            PerformanceBenchmarkParams::new_nodes_down(10).build(cluster),
        ));
        experiments.push(Box::new(
            PerformanceBenchmarkThreeRegionSimulationParams {}.build(cluster),
        ));
        experiments.push(Box::new(
            PerformanceBenchmarkParams::new_fixed_tps(0, 10).build(cluster),
        ));
        Self { experiments }
    }

    fn new_land_blocking_suite(cluster: &Cluster) -> Self {
        let mut experiments: Vec<Box<dyn Experiment>> = vec![];
        experiments.push(Box::new(
            PerformanceBenchmarkParams::new_nodes_down(0).build(cluster),
        ));
        Self { experiments }
    }

    fn new_land_blocking_compat_suite(cluster: &Cluster) -> Result<Self> {
        let count: usize = match env::var("BATCH_SIZE") {
            Ok(val) => val
                .parse()
                .map_err(|e| format_err!("Failed to parse BATCH_SIZE {}: {}", val, e))?,
            Err(_) => cluster.validator_instances().len() / 2,
        };
        let updated_image_tag = env::var("UPDATE_TO_TAG")
            .map_err(|_| format_err!("Expected environment variable UPDATE_TO_TAG"))?;
        let mut experiments: Vec<Box<dyn Experiment>> = vec![];
        experiments.push(Box::new(
            CompatiblityTestParams {
                count,
                updated_image_tag,
            }
            .build(cluster),
        ));
        experiments.extend(Self::new_land_blocking_suite(cluster).experiments);
        Ok(Self { experiments })
    }

    pub fn new_by_name(cluster: &Cluster, name: &str) -> Result<Self> {
        match name {
            "perf" => Ok(Self::new_perf_suite(cluster)),
            "pre_release" => Ok(Self::new_pre_release(cluster)),
            "land_blocking" => Ok(Self::new_land_blocking_suite(cluster)),
            "land_blocking_compat" => Self::new_land_blocking_compat_suite(cluster),
            other => Err(format_err!("Unknown suite: {}", other)),
        }
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::result::Result as StdResult;

use base::ast;
use base::fnv::FnvMap;
use base::kind::{self, ArcKind, Kind, KindCache, KindEnv};
use base::merge;
use base::symbol::Symbol;
use base::types::{self, AppVec, ArcType, BuiltinType, Field, Generic, Type, Walker};

use substitution::{Constraints, Substitutable, Substitution};
use unify::{self, Error as UnifyError, Unifiable, Unifier, UnifierState};

pub type Error<I> = UnifyError<ArcKind, KindError<I>>;

pub type Result<T> = StdResult<T, Error<Symbol>>;

/// Struct containing methods for kindchecking types
pub struct KindCheck<'a> {
    variables: Vec<Generic<Symbol>>,
    /// Type bindings local to the current kindcheck invocation
    locals: Vec<(Symbol, ArcKind)>,
    info: &'a (KindEnv + 'a),
    idents: &'a (ast::IdentEnv<Ident = Symbol> + 'a),
    pub subs: Substitution<ArcKind>,
    kind_cache: KindCache,
    /// A cached one argument kind function, `Type -> Type`
    function1_kind: ArcKind,
    /// A cached two argument kind function, `Type -> Type -> Type`
    function2_kind: ArcKind,
}

fn walk_move_kind<F>(kind: ArcKind, f: &mut F) -> ArcKind
where
    F: FnMut(&Kind) -> Option<ArcKind>,
{
    match walk_move_kind2(&kind, f) {
        Some(kind) => kind,
        None => kind,
    }
}
fn walk_move_kind2<F>(kind: &ArcKind, f: &mut F) -> Option<ArcKind>
where
    F: FnMut(&Kind) -> Option<ArcKind>,
{
    let new = f(kind);
    let new2 = {
        let kind = new.as_ref().unwrap_or(kind);
        match **kind {
            Kind::Function(ref arg, ref ret) => {
                let arg_new = walk_move_kind2(arg, f);
                let ret_new = walk_move_kind2(ret, f);
                merge::merge(arg, arg_new, ret, ret_new, Kind::function)
            }
            Kind::Hole | Kind::Type | Kind::Variable(_) | Kind::Row => None,
        }
    };
    new2.or(new)
}

impl<'a> KindCheck<'a> {
    pub fn new(
        info: &'a (KindEnv + 'a),
        idents: &'a (ast::IdentEnv<Ident = Symbol> + 'a),
        kind_cache: KindCache,
    ) -> KindCheck<'a> {
        let typ = kind_cache.typ();
        let function1_kind = Kind::function(typ.clone(), typ.clone());
        KindCheck {
            variables: Vec::new(),
            locals: Vec::new(),
            info: info,
            idents: idents,
            subs: Substitution::new(()),
            function1_kind: function1_kind.clone(),
            function2_kind: Kind::function(typ, function1_kind),
            kind_cache: kind_cache,
        }
    }

    pub fn add_local(&mut self, name: Symbol, kind: ArcKind) {
        self.locals.push((name, kind));
    }

    pub fn set_variables(&mut self, variables: &[Generic<Symbol>]) {
        self.variables.clear();
        self.variables.extend(variables.iter().cloned());
    }

    pub fn type_kind(&self) -> ArcKind {
        self.kind_cache.typ()
    }

    pub fn function1_kind(&self) -> ArcKind {
        self.function1_kind.clone()
    }

    pub fn function2_kind(&self) -> ArcKind {
        self.function2_kind.clone()
    }

    pub fn row_kind(&self) -> ArcKind {
        self.kind_cache.row()
    }

    pub fn instantiate_kinds(&mut self, kind: &mut ArcKind) {
        match *ArcKind::make_mut(kind) {
            // Can't assign a new var to `kind` here because it is borrowed mutably...
            // We'll rely on fall-through instead.
            Kind::Hole => (),
            Kind::Variable(_) => ice!("Unexpected kind variable while instantiating"),
            Kind::Function(ref mut lhs, ref mut rhs) => {
                self.instantiate_kinds(lhs);
                self.instantiate_kinds(rhs);
                return;
            }
            Kind::Row | Kind::Type => return,
        }
        *kind = self.subs.new_var();
    }

    fn find(&mut self, id: &Symbol) -> Result<ArcKind> {
        let kind = self.variables
            .iter()
            .find(|var| var.id == *id)
            .map(|t| t.kind.clone())
            .or_else(|| {
                self.locals.iter().find(|t| t.0 == *id).map(|t| t.1.clone())
            })
            .or_else(|| self.info.find_kind(id))
            .map_or_else(
                || {
                    let id_str = self.idents.string(id);
                    if id_str.starts_with(char::is_uppercase) {
                        Err(UnifyError::Other(KindError::UndefinedType(id.clone())))
                    } else {
                        // Create a new variable
                        self.locals.push((id.clone(), self.subs.new_var()));
                        Ok(self.locals.last().unwrap().1.clone())
                    }
                },
                Ok,
            );

        if let Ok(ref kind) = kind {
            debug!("Find kind: {} => {}", self.idents.string(&id), kind);
        }

        kind
    }

    // Kindhecks `typ`, infering it to be of kind `Type`
    pub fn kindcheck_type(&mut self, typ: &mut ArcType) -> Result<ArcKind> {
        let type_kind = self.type_kind();
        self.kindcheck_expected(typ, &type_kind)
    }

    pub fn kindcheck_expected(&mut self, typ: &mut ArcType, expected: &ArcKind) -> Result<ArcKind> {
        debug!("Kindcheck {:?}", typ);
        let (kind, t) = self.kindcheck(typ)?;
        let kind = self.unify(expected, kind)?;
        *typ = self.finalize_type(t);
        debug!("Done {:?}", typ);
        Ok(kind)
    }

    fn builtin_kind(&self, typ: BuiltinType) -> ArcKind {
        match typ {
            BuiltinType::String |
            BuiltinType::Byte |
            BuiltinType::Char |
            BuiltinType::Int |
            BuiltinType::Float => self.type_kind(),
            BuiltinType::Array => self.function1_kind(),
            BuiltinType::Function => self.function2_kind(),
        }
    }

    fn kindcheck(&mut self, typ: &ArcType) -> Result<(ArcKind, ArcType)> {
        match **typ {
            Type::Hole | Type::Opaque | Type::Variable(_) => Ok((self.subs.new_var(), typ.clone())),
            Type::Generic(ref gen) => {
                let mut gen = gen.clone();
                gen.kind = self.find(&gen.id)?;
                Ok((gen.kind.clone(), Type::generic(gen)))
            }
            Type::Builtin(builtin_typ) => Ok((self.builtin_kind(builtin_typ), typ.clone())),
            Type::Forall(ref params, ref typ) => {
                let (ret_kind, ret_typ) = self.kindcheck(typ)?;
                let kind = params.into_iter().rev().fold(ret_kind, |body, param| {
                    Kind::function(param.kind.clone(), body)
                });

                Ok((kind, ret_typ.clone()))
            }
            Type::App(ref ctor, ref args) => {
                let (mut kind, ctor) = self.kindcheck(ctor)?;
                let mut new_args = AppVec::new();
                for arg in args {
                    let f = Kind::function(self.subs.new_var(), self.subs.new_var());
                    kind = self.unify(&f, kind)?;
                    kind = match *kind {
                        Kind::Function(ref arg_kind, ref ret) => {
                            let (actual, new_arg) = self.kindcheck(arg)?;
                            new_args.push(new_arg);
                            self.unify(arg_kind, actual)?;
                            ret.clone()
                        }
                        _ => {
                            return Err(UnifyError::TypeMismatch(
                                self.function1_kind(),
                                kind.clone(),
                            ))
                        }
                    };
                }
                Ok((kind, Type::app(ctor, new_args)))
            }
            Type::Variant(ref row) => {
                let row: Result<_> = row.row_iter()
                    .map(|field| {
                        let (kind, typ) = self.kindcheck(&field.typ)?;
                        let type_kind = self.type_kind();
                        self.unify(&type_kind, kind)?;
                        Ok(Field::new(field.name.clone(), typ))
                    })
                    .collect();

                Ok((self.type_kind(), Type::variant(row?)))
            }
            Type::Record(ref row) => {
                let (kind, row) = self.kindcheck(row)?;
                let row_kind = self.row_kind();
                self.unify(&row_kind, kind)?;
                Ok((self.type_kind(), ArcType::from(Type::Record(row))))
            }
            Type::ExtendRow {
                ref types,
                ref fields,
                ref rest,
            } => {
                let fields: Result<_> = fields
                    .iter()
                    .map(|field| {
                        let (kind, typ) = self.kindcheck(&field.typ)?;
                        let type_kind = self.type_kind();
                        self.unify(&type_kind, kind)?;
                        Ok(Field::new(field.name.clone(), typ))
                    })
                    .collect();

                let (kind, rest) = self.kindcheck(rest)?;
                let row_kind = self.row_kind();
                self.unify(&row_kind, kind)?;

                Ok((row_kind, Type::extend_row(types.clone(), fields?, rest)))
            }
            Type::EmptyRow => Ok((self.row_kind(), typ.clone())),
            Type::Ident(ref id) => self.find(id).map(|kind| (kind, typ.clone())),
            Type::Alias(ref alias) => self.find(&alias.name).map(|kind| (kind, typ.clone())),
        }
    }

    fn unify(&mut self, expected: &ArcKind, mut actual: ArcKind) -> Result<ArcKind> {
        debug!("Unify {:?} <=> {:?}", expected, actual);
        let result = unify::unify(&self.subs, (), expected, &actual);
        match result {
            Ok(k) => Ok(k),
            Err(_errors) => {
                let mut expected = expected.clone();
                expected = update_kind(&self.subs, expected, None);
                actual = update_kind(&self.subs, actual, None);
                Err(UnifyError::TypeMismatch(expected, actual))
            }
        }
    }

    pub fn finalize_type(&self, typ: ArcType) -> ArcType {
        let default = Some(&self.kind_cache.typ);
        types::walk_move_type(typ, &mut |typ| match **typ {
            Type::Variable(ref var) => {
                let mut kind = var.kind.clone();
                kind = update_kind(&self.subs, kind, default);
                Some(Type::variable(types::TypeVariable {
                    id: var.id,
                    kind: kind,
                }))
            }
            Type::Generic(ref var) => Some(Type::generic(self.finalize_generic(var))),
            _ => None,
        })
    }
    pub fn finalize_generic(&self, var: &Generic<Symbol>) -> Generic<Symbol> {
        let mut kind = var.kind.clone();
        kind = update_kind(&self.subs, kind, Some(&self.kind_cache.typ));
        Generic::new(var.id.clone(), kind)
    }
}

fn update_kind(subs: &Substitution<ArcKind>, kind: ArcKind, default: Option<&ArcKind>) -> ArcKind {
    walk_move_kind(kind, &mut |kind| match *kind {
        Kind::Variable(id) => subs.find_type_for_var(id)
            .cloned()
            .or_else(|| default.cloned()),
        _ => None,
    })
}

/// Enumeration possible errors other than mismatch and occurs when kindchecking
#[derive(Debug, PartialEq)]
pub enum KindError<I> {
    /// The type is not defined in the current scope
    UndefinedType(I),
}

impl<I> fmt::Display for KindError<I>
where
    I: fmt::Display,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            KindError::UndefinedType(ref name) => write!(f, "Type '{}' is not defined", name),
        }
    }
}

pub fn fmt_kind_error<I>(error: &Error<I>, f: &mut fmt::Formatter) -> fmt::Result
where
    I: fmt::Display,
{
    use unify::Error::*;
    match *error {
        TypeMismatch(ref expected, ref actual) => write!(
            f,
            "Kind mismatch\nExpected: {}\nFound: {}",
            expected,
            actual
        ),
        Substitution(ref err) => write!(f, "{}", err),
        Other(ref err) => write!(f, "{}", err),
    }
}

impl Substitutable for ArcKind {
    type Variable = u32;
    type Factory = ();

    fn from_variable(x: u32) -> ArcKind {
        Kind::variable(x)
    }

    fn get_var(&self) -> Option<&u32> {
        match **self {
            Kind::Variable(ref var) => Some(var),
            _ => None,
        }
    }

    fn traverse<F>(&self, f: &mut F)
    where
        F: Walker<ArcKind>,
    {
        kind::walk_kind(self, f);
    }
    fn instantiate(
        &self,
        _subs: &Substitution<Self>,
        _constraints: &FnvMap<Symbol, Constraints<Self>>,
    ) -> Self {
        self.clone()
    }
}

impl<S> Unifiable<S> for ArcKind {
    type Error = KindError<Symbol>;

    fn zip_match<U>(
        &self,
        other: &Self,
        unifier: &mut UnifierState<S, U>,
    ) -> StdResult<Option<Self>, Error<Symbol>>
    where
        U: Unifier<S, Self>,
    {
        match (&**self, &**other) {
            (&Kind::Function(ref l1, ref l2), &Kind::Function(ref r1, ref r2)) => {
                let a = unifier.try_match(l1, r1);
                let r = unifier.try_match(l2, r2);
                Ok(merge::merge(l1, a, l2, r, Kind::function))
            }
            (l, r) => if l == r {
                Ok(None)
            } else {
                Err(UnifyError::TypeMismatch(self.clone(), other.clone()))
            },
        }
    }
}
<commit_msg>Don't forget forall when kindchecking<commit_after>use std::fmt;
use std::result::Result as StdResult;

use base::ast;
use base::fnv::FnvMap;
use base::kind::{self, ArcKind, Kind, KindCache, KindEnv};
use base::merge;
use base::symbol::Symbol;
use base::types::{self, AppVec, ArcType, BuiltinType, Field, Generic, Type, Walker};

use substitution::{Constraints, Substitutable, Substitution};
use unify::{self, Error as UnifyError, Unifiable, Unifier, UnifierState};

pub type Error<I> = UnifyError<ArcKind, KindError<I>>;

pub type Result<T> = StdResult<T, Error<Symbol>>;

/// Struct containing methods for kindchecking types
pub struct KindCheck<'a> {
    variables: Vec<Generic<Symbol>>,
    /// Type bindings local to the current kindcheck invocation
    locals: Vec<(Symbol, ArcKind)>,
    info: &'a (KindEnv + 'a),
    idents: &'a (ast::IdentEnv<Ident = Symbol> + 'a),
    pub subs: Substitution<ArcKind>,
    kind_cache: KindCache,
    /// A cached one argument kind function, `Type -> Type`
    function1_kind: ArcKind,
    /// A cached two argument kind function, `Type -> Type -> Type`
    function2_kind: ArcKind,
}

fn walk_move_kind<F>(kind: ArcKind, f: &mut F) -> ArcKind
where
    F: FnMut(&Kind) -> Option<ArcKind>,
{
    match walk_move_kind2(&kind, f) {
        Some(kind) => kind,
        None => kind,
    }
}
fn walk_move_kind2<F>(kind: &ArcKind, f: &mut F) -> Option<ArcKind>
where
    F: FnMut(&Kind) -> Option<ArcKind>,
{
    let new = f(kind);
    let new2 = {
        let kind = new.as_ref().unwrap_or(kind);
        match **kind {
            Kind::Function(ref arg, ref ret) => {
                let arg_new = walk_move_kind2(arg, f);
                let ret_new = walk_move_kind2(ret, f);
                merge::merge(arg, arg_new, ret, ret_new, Kind::function)
            }
            Kind::Hole | Kind::Type | Kind::Variable(_) | Kind::Row => None,
        }
    };
    new2.or(new)
}

impl<'a> KindCheck<'a> {
    pub fn new(
        info: &'a (KindEnv + 'a),
        idents: &'a (ast::IdentEnv<Ident = Symbol> + 'a),
        kind_cache: KindCache,
    ) -> KindCheck<'a> {
        let typ = kind_cache.typ();
        let function1_kind = Kind::function(typ.clone(), typ.clone());
        KindCheck {
            variables: Vec::new(),
            locals: Vec::new(),
            info: info,
            idents: idents,
            subs: Substitution::new(()),
            function1_kind: function1_kind.clone(),
            function2_kind: Kind::function(typ, function1_kind),
            kind_cache: kind_cache,
        }
    }

    pub fn add_local(&mut self, name: Symbol, kind: ArcKind) {
        self.locals.push((name, kind));
    }

    pub fn set_variables(&mut self, variables: &[Generic<Symbol>]) {
        self.variables.clear();
        self.variables.extend(variables.iter().cloned());
    }

    pub fn type_kind(&self) -> ArcKind {
        self.kind_cache.typ()
    }

    pub fn function1_kind(&self) -> ArcKind {
        self.function1_kind.clone()
    }

    pub fn function2_kind(&self) -> ArcKind {
        self.function2_kind.clone()
    }

    pub fn row_kind(&self) -> ArcKind {
        self.kind_cache.row()
    }

    pub fn instantiate_kinds(&mut self, kind: &mut ArcKind) {
        match *ArcKind::make_mut(kind) {
            // Can't assign a new var to `kind` here because it is borrowed mutably...
            // We'll rely on fall-through instead.
            Kind::Hole => (),
            Kind::Variable(_) => ice!("Unexpected kind variable while instantiating"),
            Kind::Function(ref mut lhs, ref mut rhs) => {
                self.instantiate_kinds(lhs);
                self.instantiate_kinds(rhs);
                return;
            }
            Kind::Row | Kind::Type => return,
        }
        *kind = self.subs.new_var();
    }

    fn find(&mut self, id: &Symbol) -> Result<ArcKind> {
        let kind = self.variables
            .iter()
            .find(|var| var.id == *id)
            .map(|t| t.kind.clone())
            .or_else(|| {
                self.locals.iter().find(|t| t.0 == *id).map(|t| t.1.clone())
            })
            .or_else(|| self.info.find_kind(id))
            .map_or_else(
                || {
                    let id_str = self.idents.string(id);
                    if id_str.starts_with(char::is_uppercase) {
                        Err(UnifyError::Other(KindError::UndefinedType(id.clone())))
                    } else {
                        // Create a new variable
                        self.locals.push((id.clone(), self.subs.new_var()));
                        Ok(self.locals.last().unwrap().1.clone())
                    }
                },
                Ok,
            );

        if let Ok(ref kind) = kind {
            debug!("Find kind: {} => {}", self.idents.string(&id), kind);
        }

        kind
    }

    // Kindhecks `typ`, infering it to be of kind `Type`
    pub fn kindcheck_type(&mut self, typ: &mut ArcType) -> Result<ArcKind> {
        let type_kind = self.type_kind();
        self.kindcheck_expected(typ, &type_kind)
    }

    pub fn kindcheck_expected(&mut self, typ: &mut ArcType, expected: &ArcKind) -> Result<ArcKind> {
        debug!("Kindcheck {:?}", typ);
        let (kind, t) = self.kindcheck(typ)?;
        let kind = self.unify(expected, kind)?;
        *typ = self.finalize_type(t);
        debug!("Done {:?}", typ);
        Ok(kind)
    }

    fn builtin_kind(&self, typ: BuiltinType) -> ArcKind {
        match typ {
            BuiltinType::String |
            BuiltinType::Byte |
            BuiltinType::Char |
            BuiltinType::Int |
            BuiltinType::Float => self.type_kind(),
            BuiltinType::Array => self.function1_kind(),
            BuiltinType::Function => self.function2_kind(),
        }
    }

    fn kindcheck(&mut self, typ: &ArcType) -> Result<(ArcKind, ArcType)> {
        match **typ {
            Type::Hole | Type::Opaque | Type::Variable(_) => Ok((self.subs.new_var(), typ.clone())),
            Type::Generic(ref gen) => {
                let mut gen = gen.clone();
                gen.kind = self.find(&gen.id)?;
                Ok((gen.kind.clone(), Type::generic(gen)))
            }
            Type::Builtin(builtin_typ) => Ok((self.builtin_kind(builtin_typ), typ.clone())),
            Type::Forall(ref params, ref typ) => {
                let mut params = params.clone();
                for param in &mut params {
                    param.kind = self.subs.new_var();
                    self.locals.push((param.id.clone(), param.kind.clone()));
                }
                let (ret_kind, ret_typ) = self.kindcheck(typ)?;

                let offset = self.locals.len() - params.len();
                self.locals.drain(offset..);

                Ok((ret_kind, Type::forall(params, ret_typ.clone())))
            }
            Type::App(ref ctor, ref args) => {
                let (mut kind, ctor) = self.kindcheck(ctor)?;
                let mut new_args = AppVec::new();
                for arg in args {
                    let f = Kind::function(self.subs.new_var(), self.subs.new_var());
                    kind = self.unify(&f, kind)?;
                    kind = match *kind {
                        Kind::Function(ref arg_kind, ref ret) => {
                            let (actual, new_arg) = self.kindcheck(arg)?;
                            new_args.push(new_arg);
                            self.unify(arg_kind, actual)?;
                            ret.clone()
                        }
                        _ => {
                            return Err(UnifyError::TypeMismatch(
                                self.function1_kind(),
                                kind.clone(),
                            ))
                        }
                    };
                }
                Ok((kind, Type::app(ctor, new_args)))
            }
            Type::Variant(ref row) => {
                let row: Result<_> = row.row_iter()
                    .map(|field| {
                        let (kind, typ) = self.kindcheck(&field.typ)?;
                        let type_kind = self.type_kind();
                        self.unify(&type_kind, kind)?;
                        Ok(Field::new(field.name.clone(), typ))
                    })
                    .collect();

                Ok((self.type_kind(), Type::variant(row?)))
            }
            Type::Record(ref row) => {
                let (kind, row) = self.kindcheck(row)?;
                let row_kind = self.row_kind();
                self.unify(&row_kind, kind)?;
                Ok((self.type_kind(), ArcType::from(Type::Record(row))))
            }
            Type::ExtendRow {
                ref types,
                ref fields,
                ref rest,
            } => {
                let fields: Result<_> = fields
                    .iter()
                    .map(|field| {
                        let (kind, typ) = self.kindcheck(&field.typ)?;
                        let type_kind = self.type_kind();
                        self.unify(&type_kind, kind)?;
                        Ok(Field::new(field.name.clone(), typ))
                    })
                    .collect();

                let (kind, rest) = self.kindcheck(rest)?;
                let row_kind = self.row_kind();
                self.unify(&row_kind, kind)?;

                Ok((row_kind, Type::extend_row(types.clone(), fields?, rest)))
            }
            Type::EmptyRow => Ok((self.row_kind(), typ.clone())),
            Type::Ident(ref id) => self.find(id).map(|kind| (kind, typ.clone())),
            Type::Alias(ref alias) => self.find(&alias.name).map(|kind| (kind, typ.clone())),
        }
    }

    fn unify(&mut self, expected: &ArcKind, mut actual: ArcKind) -> Result<ArcKind> {
        debug!("Unify {:?} <=> {:?}", expected, actual);
        let result = unify::unify(&self.subs, (), expected, &actual);
        match result {
            Ok(k) => Ok(k),
            Err(_errors) => {
                let mut expected = expected.clone();
                expected = update_kind(&self.subs, expected, None);
                actual = update_kind(&self.subs, actual, None);
                Err(UnifyError::TypeMismatch(expected, actual))
            }
        }
    }

    pub fn finalize_type(&self, typ: ArcType) -> ArcType {
        let default = Some(&self.kind_cache.typ);
        types::walk_move_type(typ, &mut |typ| match **typ {
            Type::Variable(ref var) => {
                let mut kind = var.kind.clone();
                kind = update_kind(&self.subs, kind, default);
                Some(Type::variable(types::TypeVariable {
                    id: var.id,
                    kind: kind,
                }))
            }
            Type::Generic(ref var) => Some(Type::generic(self.finalize_generic(var))),
            _ => None,
        })
    }
    pub fn finalize_generic(&self, var: &Generic<Symbol>) -> Generic<Symbol> {
        let mut kind = var.kind.clone();
        kind = update_kind(&self.subs, kind, Some(&self.kind_cache.typ));
        Generic::new(var.id.clone(), kind)
    }
}

fn update_kind(subs: &Substitution<ArcKind>, kind: ArcKind, default: Option<&ArcKind>) -> ArcKind {
    walk_move_kind(kind, &mut |kind| match *kind {
        Kind::Variable(id) => subs.find_type_for_var(id)
            .cloned()
            .or_else(|| default.cloned()),
        _ => None,
    })
}

/// Enumeration possible errors other than mismatch and occurs when kindchecking
#[derive(Debug, PartialEq)]
pub enum KindError<I> {
    /// The type is not defined in the current scope
    UndefinedType(I),
}

impl<I> fmt::Display for KindError<I>
where
    I: fmt::Display,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            KindError::UndefinedType(ref name) => write!(f, "Type '{}' is not defined", name),
        }
    }
}

pub fn fmt_kind_error<I>(error: &Error<I>, f: &mut fmt::Formatter) -> fmt::Result
where
    I: fmt::Display,
{
    use unify::Error::*;
    match *error {
        TypeMismatch(ref expected, ref actual) => write!(
            f,
            "Kind mismatch\nExpected: {}\nFound: {}",
            expected,
            actual
        ),
        Substitution(ref err) => write!(f, "{}", err),
        Other(ref err) => write!(f, "{}", err),
    }
}

impl Substitutable for ArcKind {
    type Variable = u32;
    type Factory = ();

    fn from_variable(x: u32) -> ArcKind {
        Kind::variable(x)
    }

    fn get_var(&self) -> Option<&u32> {
        match **self {
            Kind::Variable(ref var) => Some(var),
            _ => None,
        }
    }

    fn traverse<F>(&self, f: &mut F)
    where
        F: Walker<ArcKind>,
    {
        kind::walk_kind(self, f);
    }
    fn instantiate(
        &self,
        _subs: &Substitution<Self>,
        _constraints: &FnvMap<Symbol, Constraints<Self>>,
    ) -> Self {
        self.clone()
    }
}

impl<S> Unifiable<S> for ArcKind {
    type Error = KindError<Symbol>;

    fn zip_match<U>(
        &self,
        other: &Self,
        unifier: &mut UnifierState<S, U>,
    ) -> StdResult<Option<Self>, Error<Symbol>>
    where
        U: Unifier<S, Self>,
    {
        match (&**self, &**other) {
            (&Kind::Function(ref l1, ref l2), &Kind::Function(ref r1, ref r2)) => {
                let a = unifier.try_match(l1, r1);
                let r = unifier.try_match(l2, r2);
                Ok(merge::merge(l1, a, l2, r, Kind::function))
            }
            (l, r) => if l == r {
                Ok(None)
            } else {
                Err(UnifyError::TypeMismatch(self.clone(), other.clone()))
            },
        }
    }
}
<|endoftext|>


<commit_before>use super::validate_module;
use builder::module;
use elements::{BlockType, ExportEntry, External, FunctionType, GlobalEntry, GlobalType,
				ImportEntry, InitExpr, Internal, MemoryType, Opcode, Opcodes, TableType,
				ValueType};

#[test]
fn empty_is_valid() {
	let module = module().build();
	assert!(validate_module(&module).is_ok());
}

#[test]
fn mem_limits() {
	// min > max
	let m = module()
		.memory()
			.with_min(10)
			.with_max(Some(9))
			.build()
		.build();
	assert!(validate_module(&m).is_err());

	// min = max
	let m = module()
		.memory()
			.with_min(10)
			.with_max(Some(10))
			.build()
		.build();
	assert!(validate_module(&m).is_ok());

	// mem is always valid without max
	let m = module()
		.memory()
			.with_min(10)
			.build()
		.build();
	assert!(validate_module(&m).is_ok());
}

#[test]
fn table_limits() {
	// min > max
	let m = module()
		.table()
			.with_min(10)
			.with_max(Some(9))
			.build()
		.build();
	assert!(validate_module(&m).is_err());

	// min = max
	let m = module()
		.table()
			.with_min(10)
			.with_max(Some(10))
			.build()
		.build();
	assert!(validate_module(&m).is_ok());

	// table is always valid without max
	let m = module()
		.table()
			.with_min(10)
			.build()
		.build();
	assert!(validate_module(&m).is_ok());
}

#[test]
fn global_init_const() {
	let m = module()
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, true),
				InitExpr::new(
					vec![Opcode::I32Const(42), Opcode::End]
				)
			)
		)
		.build();
	assert!(validate_module(&m).is_ok());

	// init expr type differs from declared global type
	let m = module()
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I64, true),
				InitExpr::new(vec![Opcode::I32Const(42), Opcode::End])
			)
		)
		.build();
	assert!(validate_module(&m).is_err());
}

#[test]
fn global_init_global() {
	let m = module()
		.with_import(
			ImportEntry::new(
				"env".into(),
				"ext_global".into(),
				External::Global(GlobalType::new(ValueType::I32, false))
			)
		)
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, true),
				InitExpr::new(vec![Opcode::GetGlobal(0), Opcode::End])
			)
		)
		.build();
	assert!(validate_module(&m).is_ok());

	// get_global can reference only previously defined globals
	let m = module()
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, true),
				InitExpr::new(vec![Opcode::GetGlobal(0), Opcode::End])
			)
		)
		.build();
	assert!(validate_module(&m).is_err());

	// get_global can reference only const globals
	let m = module()
		.with_import(
			ImportEntry::new(
				"env".into(),
				"ext_global".into(),
				External::Global(GlobalType::new(ValueType::I32, true))
			)
		)
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, true),
				InitExpr::new(vec![Opcode::GetGlobal(0), Opcode::End])
			)
		)
		.build();
	assert!(validate_module(&m).is_err());

	// get_global in init_expr can only refer to imported globals.
	let m = module()
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, false),
				InitExpr::new(vec![Opcode::I32Const(0), Opcode::End])
			)
		)
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, true),
				InitExpr::new(vec![Opcode::GetGlobal(0), Opcode::End])
			)
		)
		.build();
	assert!(validate_module(&m).is_err());
}

#[test]
fn global_init_misc() {
	// without delimiting End opcode
	let m = module()
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, true),
				InitExpr::new(vec![Opcode::I32Const(42)])
			)
		)
		.build();
	assert!(validate_module(&m).is_err());

	// empty init expr
	let m = module()
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, true),
				InitExpr::new(vec![Opcode::End])
			)
		)
		.build();
	assert!(validate_module(&m).is_err());

	// not an constant opcode used
	let m = module()
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, true),
				InitExpr::new(vec![Opcode::Unreachable, Opcode::End])
			)
		)
		.build();
	assert!(validate_module(&m).is_err());
}

#[test]
fn module_limits_validity() {
	// module cannot contain more than 1 memory atm.
	let m = module()
		.with_import(
			ImportEntry::new(
				"core".into(),
				"memory".into(),
				External::Memory(MemoryType::new(10, None))
			)
		)
		.memory()
			.with_min(10)
			.build()
		.build();
	assert!(validate_module(&m).is_err());

	// module cannot contain more than 1 table atm.
	let m = module()
		.with_import(
			ImportEntry::new(
				"core".into(),
				"table".into(),
				External::Table(TableType::new(10, None))
			)
		)
		.table()
			.with_min(10)
			.build()
		.build();
	assert!(validate_module(&m).is_err());
}

// TODO: pepyakin
// #[test]
// fn if_else_with_return_type_validation() {
// 	let module_instance = ModuleInstance::new(Weak::default(), "test".into(), module().build()).unwrap();
// 	let mut context = FunctionValidationContext::new(&module_instance, None, &[], 1024, 1024, FunctionSignature::Module(&FunctionType::default()));

// 	Validator::validate_function(&mut context, BlockType::NoResult, &[
// 		Opcode::I32Const(1),
// 		Opcode::If(BlockType::NoResult),
// 			Opcode::I32Const(1),
// 			Opcode::If(BlockType::Value(ValueType::I32)),
// 				Opcode::I32Const(1),
// 			Opcode::Else,
// 				Opcode::I32Const(2),
// 			Opcode::End,
// 		Opcode::Drop,
// 		Opcode::End,
// 		Opcode::End,
// 	]).unwrap();
// }
<commit_msg>Test recursive calls<commit_after>use super::validate_module;
use builder::module;
use elements::{BlockType, ExportEntry, External, FunctionType, GlobalEntry, GlobalType,
				ImportEntry, InitExpr, Internal, MemoryType, Opcode, Opcodes, TableType,
				ValueType};

#[test]
fn empty_is_valid() {
	let module = module().build();
	assert!(validate_module(&module).is_ok());
}

#[test]
fn mem_limits() {
	// min > max
	let m = module()
		.memory()
			.with_min(10)
			.with_max(Some(9))
			.build()
		.build();
	assert!(validate_module(&m).is_err());

	// min = max
	let m = module()
		.memory()
			.with_min(10)
			.with_max(Some(10))
			.build()
		.build();
	assert!(validate_module(&m).is_ok());

	// mem is always valid without max
	let m = module()
		.memory()
			.with_min(10)
			.build()
		.build();
	assert!(validate_module(&m).is_ok());
}

#[test]
fn table_limits() {
	// min > max
	let m = module()
		.table()
			.with_min(10)
			.with_max(Some(9))
			.build()
		.build();
	assert!(validate_module(&m).is_err());

	// min = max
	let m = module()
		.table()
			.with_min(10)
			.with_max(Some(10))
			.build()
		.build();
	assert!(validate_module(&m).is_ok());

	// table is always valid without max
	let m = module()
		.table()
			.with_min(10)
			.build()
		.build();
	assert!(validate_module(&m).is_ok());
}

#[test]
fn global_init_const() {
	let m = module()
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, true),
				InitExpr::new(
					vec![Opcode::I32Const(42), Opcode::End]
				)
			)
		)
		.build();
	assert!(validate_module(&m).is_ok());

	// init expr type differs from declared global type
	let m = module()
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I64, true),
				InitExpr::new(vec![Opcode::I32Const(42), Opcode::End])
			)
		)
		.build();
	assert!(validate_module(&m).is_err());
}

#[test]
fn global_init_global() {
	let m = module()
		.with_import(
			ImportEntry::new(
				"env".into(),
				"ext_global".into(),
				External::Global(GlobalType::new(ValueType::I32, false))
			)
		)
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, true),
				InitExpr::new(vec![Opcode::GetGlobal(0), Opcode::End])
			)
		)
		.build();
	assert!(validate_module(&m).is_ok());

	// get_global can reference only previously defined globals
	let m = module()
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, true),
				InitExpr::new(vec![Opcode::GetGlobal(0), Opcode::End])
			)
		)
		.build();
	assert!(validate_module(&m).is_err());

	// get_global can reference only const globals
	let m = module()
		.with_import(
			ImportEntry::new(
				"env".into(),
				"ext_global".into(),
				External::Global(GlobalType::new(ValueType::I32, true))
			)
		)
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, true),
				InitExpr::new(vec![Opcode::GetGlobal(0), Opcode::End])
			)
		)
		.build();
	assert!(validate_module(&m).is_err());

	// get_global in init_expr can only refer to imported globals.
	let m = module()
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, false),
				InitExpr::new(vec![Opcode::I32Const(0), Opcode::End])
			)
		)
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, true),
				InitExpr::new(vec![Opcode::GetGlobal(0), Opcode::End])
			)
		)
		.build();
	assert!(validate_module(&m).is_err());
}

#[test]
fn global_init_misc() {
	// without delimiting End opcode
	let m = module()
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, true),
				InitExpr::new(vec![Opcode::I32Const(42)])
			)
		)
		.build();
	assert!(validate_module(&m).is_err());

	// empty init expr
	let m = module()
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, true),
				InitExpr::new(vec![Opcode::End])
			)
		)
		.build();
	assert!(validate_module(&m).is_err());

	// not an constant opcode used
	let m = module()
		.with_global(
			GlobalEntry::new(
				GlobalType::new(ValueType::I32, true),
				InitExpr::new(vec![Opcode::Unreachable, Opcode::End])
			)
		)
		.build();
	assert!(validate_module(&m).is_err());
}

#[test]
fn module_limits_validity() {
	// module cannot contain more than 1 memory atm.
	let m = module()
		.with_import(
			ImportEntry::new(
				"core".into(),
				"memory".into(),
				External::Memory(MemoryType::new(10, None))
			)
		)
		.memory()
			.with_min(10)
			.build()
		.build();
	assert!(validate_module(&m).is_err());

	// module cannot contain more than 1 table atm.
	let m = module()
		.with_import(
			ImportEntry::new(
				"core".into(),
				"table".into(),
				External::Table(TableType::new(10, None))
			)
		)
		.table()
			.with_min(10)
			.build()
		.build();
	assert!(validate_module(&m).is_err());
}

#[test]
fn funcs() {
	// recursive function calls is legal.
	let m = module()
		.function()
			.signature().return_type().i32().build()
			.body().with_opcodes(Opcodes::new(vec![
				Opcode::Call(1),
				Opcode::End,
			])).build()
			.build()
		.function()
			.signature().return_type().i32().build()
			.body().with_opcodes(Opcodes::new(vec![
				Opcode::Call(0),
				Opcode::End,
			])).build()
			.build()
		.build();
	assert!(validate_module(&m).is_ok());
}

// TODO: pepyakin
// #[test]
// fn if_else_with_return_type_validation() {
// 	let module_instance = ModuleInstance::new(Weak::default(), "test".into(), module().build()).unwrap();
// 	let mut context = FunctionValidationContext::new(&module_instance, None, &[], 1024, 1024, FunctionSignature::Module(&FunctionType::default()));

// 	Validator::validate_function(&mut context, BlockType::NoResult, &[
// 		Opcode::I32Const(1),
// 		Opcode::If(BlockType::NoResult),
// 			Opcode::I32Const(1),
// 			Opcode::If(BlockType::Value(ValueType::I32)),
// 				Opcode::I32Const(1),
// 			Opcode::Else,
// 				Opcode::I32Const(2),
// 			Opcode::End,
// 		Opcode::Drop,
// 		Opcode::End,
// 		Opcode::End,
// 	]).unwrap();
// }
<|endoftext|>


<commit_before>//! An enum mapping type.
//!
//! It is implemented using an array type, so using it is as fast as using Rust
//! arrays.
//!
//! # Examples
//!
//! ```
//! #[macro_use]
//! extern crate enum_map;
//!
//! use enum_map::EnumMap;
//!
//! #[derive(Debug, EnumMap)]
//! enum Example {
//!     A,
//!     B,
//!     C,
//! }
//!
//! fn main() {
//!     let mut map = enum_map! {
//!         Example::A => 1,
//!         Example::B => 2,
//!         Example::C => 3,
//!     };
//!     map[Example::C] = 4;
//!
//!     assert_eq!(map[Example::A], 1);
//!
//!     for (key, &value) in &map {
//!         println!("{:?} has {} as value.", key, value);
//!     }
//! }
//! ```

#![no_std]
#![deny(missing_docs)]

#[macro_use]
extern crate array_macro;
// This allows to quiet "proc macro crates and `#[no_link]` crates have no
// effect without `#[macro_use]`" warning. Just using extern DOES have
// an effect of letting me export the macro.
#[allow(unused_imports)]
#[macro_use]
extern crate enum_map_derive;

mod enummap_impls;
mod internal;
mod iter;
mod serde;

pub use internal::Internal;
pub use iter::{Iter, IterMut};
// `*` here means re-exporting a derive procedural macro.
pub use enum_map_derive::*;

use core::slice;

/// An enum mapping.
///
/// This internally uses an array which stores a value for each possible
/// enum value. To work, it requires implementation of internal (private,
/// although public due to macro limitations) trait which allows extracting
/// information about an enum, which can be automatically generated using
/// `#[derive(EnumMap)]` from `enum_map_derive` crate.
///
/// Additionally, `bool` and `u8` automatically derives from `EnumMap`. While
/// `u8` is not technically an enum, it's convenient to consider it like one.
/// In particular, [reverse-complement in benchmark game] could be using `u8`
/// as an enum.
///
/// # Examples
///
/// ```
/// #[macro_use]
/// extern crate enum_map;
///
/// use enum_map::EnumMap;
///
/// #[derive(EnumMap)]
/// enum Example {
///     A,
///     B,
///     C,
/// }
///
/// fn main() {
///     let mut map = EnumMap::new();
///     // new initializes map with default values
///     assert_eq!(map[Example::A], 0);
///     map[Example::A] = 3;
///     assert_eq!(map[Example::A], 3);
/// }
/// ```
///
/// [reverse-complement in benchmark game]:
///     http://benchmarksgame.alioth.debian.org/u64q/program.php?test=revcomp&lang=rust&id=2
#[derive(Debug)]
pub struct EnumMap<K: Internal<V>, V> {
    array: K::Array,
}

impl<K: Internal<V>, V: Default> EnumMap<K, V>
where
    K::Array: Default,
{
    /// Creates an enum map with default values.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// use enum_map::EnumMap;
    ///
    /// #[derive(EnumMap)]
    /// enum Example {
    ///     A,
    /// }
    ///
    /// fn main() {
    ///     let enum_map = EnumMap::<_, i32>::new();
    ///     assert_eq!(enum_map[Example::A], 0);
    /// }
    /// ```
    pub fn new() -> Self {
        EnumMap::default()
    }
}

impl<K: Internal<V>, V> EnumMap<K, V> {
    /// Returns number of elements in enum map.
    pub fn len(&self) -> usize {
        self.as_slice().len()
    }

    /// Returns whether the enum variant set is empty.
    ///
    /// This isn't particularly useful, as there is no real reason to use
    /// enum map for enums without variants. However, it is provided for
    /// consistency with data structures providing len method (and I will
    /// admit, to avoid clippy warnings).
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// use enum_map::EnumMap;
    ///
    /// #[derive(EnumMap)]
    /// enum Void {}
    ///
    /// #[derive(EnumMap)]
    /// enum SingleVariant {
    ///     Variant,
    /// }
    ///
    /// fn main() {
    ///     assert_eq!(EnumMap::<Void, ()>::new().is_empty(), true);
    ///     assert_eq!(EnumMap::<SingleVariant, ()>::new().is_empty(), false);
    /// }
    pub fn is_empty(&self) -> bool {
        self.as_slice().is_empty()
    }

    /// Returns an iterator over enum map.
    pub fn iter(&self) -> Iter<K, V> {
        self.into_iter()
    }

    /// Returns a mutable iterator over enum map.
    pub fn iter_mut(&mut self) -> IterMut<K, V> {
        self.into_iter()
    }

    /// Swaps two indexes.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// fn main() {
    ///     let mut map = enum_map! { false => 0, true => 1 };
    ///     map.swap(false, true);
    ///     assert_eq!(map[false], 1);
    ///     assert_eq!(map[true], 0);
    /// }
    /// ```
    pub fn swap(&mut self, a: K, b: K) {
        self.as_mut_slice().swap(a.to_usize(), b.to_usize())
    }

    /// An iterator visiting all values. The iterator type is `&V`.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// fn main() {
    ///     let map = enum_map! { false => 3, true => 4 };
    ///     let mut values = map.values();
    ///     assert_eq!(values.next(), Some(&3));
    ///     assert_eq!(values.next(), Some(&4));
    ///     assert_eq!(values.next(), None);
    /// }
    /// ```
    pub fn values(&self) -> slice::Iter<V> {
        self.as_slice().iter()
    }

    /// An iterator visiting all values mutably. The iterator type is `&mut V`.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// fn main() {
    ///     let mut map = enum_map! { _ => 2 };
    ///     for value in map.values_mut() {
    ///         *value += 2;
    ///     }
    ///     assert_eq!(map[false], 4);
    ///     assert_eq!(map[true], 4);
    /// }
    /// ```
    pub fn values_mut(&mut self) -> slice::IterMut<V> {
        self.as_mut_slice().iter_mut()
    }

    /// Converts an enum map to a slice representing values.
    pub fn as_slice(&self) -> &[V] {
        K::slice(&self.array)
    }

    /// Converts a mutable enum map to a mutable slice representing values.
    pub fn as_mut_slice(&mut self) -> &mut [V] {
        K::slice_mut(&mut self.array)
    }

    /// Returns a raw pointer to the enum map's buffer.
    ///
    /// The caller must ensure that the slice outlives the pointer this
    /// function returns, or else it will end up pointing to garbage.
    ///
    /// Modifying the container referenced by this slice may cause its buffer
    /// to be reallocated, which would also make any pointers to it invalid.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// use enum_map::EnumMap;
    ///
    /// fn main() {
    ///     let map = enum_map! { 5 => 42, _ => 0 };
    ///     assert_eq!(unsafe { *map.as_ptr().offset(5) }, 42);
    /// }
    /// ```
    pub fn as_ptr(&self) -> *const V {
        self.as_slice().as_ptr()
    }

    /// Returns an unsafe mutable pointer to the enum map's buffer.
    ///
    /// The caller must ensure that the slice outlives the pointer this
    /// function returns, or else it will end up pointing to garbage.
    ///
    /// Modifying the container referenced by this slice may cause its buffer
    /// to be reallocated, which would also make any pointers to it invalid.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// use enum_map::EnumMap;
    ///
    /// fn main() {
    ///     let mut map = enum_map! { _ => 0 };
    ///     unsafe {
    ///         *map.as_mut_ptr().offset(11) = 23
    ///     };
    ///     assert_eq!(map[11], 23);
    /// }
    /// ```
    pub fn as_mut_ptr(&mut self) -> *mut V {
        self.as_mut_slice().as_mut_ptr()
    }
}

impl<F: FnMut(K) -> V, K: Internal<V>, V> From<F> for EnumMap<K, V> {
    fn from(f: F) -> Self {
        EnumMap { array: K::from_function(f) }
    }
}

/// Enum map constructor.
///
/// This macro allows to create a new enum map in a type safe way. It takes
/// a list of `,` separated pairs separated by `=>`. Left side is `|`
/// separated list of enum keys, or `_` to match all unmatched enum keys,
/// while right side is a value.
///
/// # Examples
///
/// ```
/// #[macro_use]
/// extern crate enum_map;
///
/// #[derive(EnumMap)]
/// enum Example {
///     A,
///     B,
///     C,
///     D,
/// }
///
/// fn main() {
///     let enum_map = enum_map! {
///         Example::A | Example::B => 1,
///         Example::C => 2,
///         _ => 3,
///     };
///     assert_eq!(enum_map[Example::A], 1);
///     assert_eq!(enum_map[Example::B], 1);
///     assert_eq!(enum_map[Example::C], 2);
///     assert_eq!(enum_map[Example::D], 3);
/// }
/// ```
#[macro_export]
macro_rules! enum_map {
    {$($t:tt)*} => {
        $crate::EnumMap::from(|k| match k { $($t)* })
    };
}
<commit_msg>Remove useless whitespace in values() doc<commit_after>//! An enum mapping type.
//!
//! It is implemented using an array type, so using it is as fast as using Rust
//! arrays.
//!
//! # Examples
//!
//! ```
//! #[macro_use]
//! extern crate enum_map;
//!
//! use enum_map::EnumMap;
//!
//! #[derive(Debug, EnumMap)]
//! enum Example {
//!     A,
//!     B,
//!     C,
//! }
//!
//! fn main() {
//!     let mut map = enum_map! {
//!         Example::A => 1,
//!         Example::B => 2,
//!         Example::C => 3,
//!     };
//!     map[Example::C] = 4;
//!
//!     assert_eq!(map[Example::A], 1);
//!
//!     for (key, &value) in &map {
//!         println!("{:?} has {} as value.", key, value);
//!     }
//! }
//! ```

#![no_std]
#![deny(missing_docs)]

#[macro_use]
extern crate array_macro;
// This allows to quiet "proc macro crates and `#[no_link]` crates have no
// effect without `#[macro_use]`" warning. Just using extern DOES have
// an effect of letting me export the macro.
#[allow(unused_imports)]
#[macro_use]
extern crate enum_map_derive;

mod enummap_impls;
mod internal;
mod iter;
mod serde;

pub use internal::Internal;
pub use iter::{Iter, IterMut};
// `*` here means re-exporting a derive procedural macro.
pub use enum_map_derive::*;

use core::slice;

/// An enum mapping.
///
/// This internally uses an array which stores a value for each possible
/// enum value. To work, it requires implementation of internal (private,
/// although public due to macro limitations) trait which allows extracting
/// information about an enum, which can be automatically generated using
/// `#[derive(EnumMap)]` from `enum_map_derive` crate.
///
/// Additionally, `bool` and `u8` automatically derives from `EnumMap`. While
/// `u8` is not technically an enum, it's convenient to consider it like one.
/// In particular, [reverse-complement in benchmark game] could be using `u8`
/// as an enum.
///
/// # Examples
///
/// ```
/// #[macro_use]
/// extern crate enum_map;
///
/// use enum_map::EnumMap;
///
/// #[derive(EnumMap)]
/// enum Example {
///     A,
///     B,
///     C,
/// }
///
/// fn main() {
///     let mut map = EnumMap::new();
///     // new initializes map with default values
///     assert_eq!(map[Example::A], 0);
///     map[Example::A] = 3;
///     assert_eq!(map[Example::A], 3);
/// }
/// ```
///
/// [reverse-complement in benchmark game]:
///     http://benchmarksgame.alioth.debian.org/u64q/program.php?test=revcomp&lang=rust&id=2
#[derive(Debug)]
pub struct EnumMap<K: Internal<V>, V> {
    array: K::Array,
}

impl<K: Internal<V>, V: Default> EnumMap<K, V>
where
    K::Array: Default,
{
    /// Creates an enum map with default values.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// use enum_map::EnumMap;
    ///
    /// #[derive(EnumMap)]
    /// enum Example {
    ///     A,
    /// }
    ///
    /// fn main() {
    ///     let enum_map = EnumMap::<_, i32>::new();
    ///     assert_eq!(enum_map[Example::A], 0);
    /// }
    /// ```
    pub fn new() -> Self {
        EnumMap::default()
    }
}

impl<K: Internal<V>, V> EnumMap<K, V> {
    /// Returns number of elements in enum map.
    pub fn len(&self) -> usize {
        self.as_slice().len()
    }

    /// Returns whether the enum variant set is empty.
    ///
    /// This isn't particularly useful, as there is no real reason to use
    /// enum map for enums without variants. However, it is provided for
    /// consistency with data structures providing len method (and I will
    /// admit, to avoid clippy warnings).
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// use enum_map::EnumMap;
    ///
    /// #[derive(EnumMap)]
    /// enum Void {}
    ///
    /// #[derive(EnumMap)]
    /// enum SingleVariant {
    ///     Variant,
    /// }
    ///
    /// fn main() {
    ///     assert_eq!(EnumMap::<Void, ()>::new().is_empty(), true);
    ///     assert_eq!(EnumMap::<SingleVariant, ()>::new().is_empty(), false);
    /// }
    pub fn is_empty(&self) -> bool {
        self.as_slice().is_empty()
    }

    /// Returns an iterator over enum map.
    pub fn iter(&self) -> Iter<K, V> {
        self.into_iter()
    }

    /// Returns a mutable iterator over enum map.
    pub fn iter_mut(&mut self) -> IterMut<K, V> {
        self.into_iter()
    }

    /// Swaps two indexes.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// fn main() {
    ///     let mut map = enum_map! { false => 0, true => 1 };
    ///     map.swap(false, true);
    ///     assert_eq!(map[false], 1);
    ///     assert_eq!(map[true], 0);
    /// }
    /// ```
    pub fn swap(&mut self, a: K, b: K) {
        self.as_mut_slice().swap(a.to_usize(), b.to_usize())
    }

    /// An iterator visiting all values. The iterator type is `&V`.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// fn main() {
    ///     let map = enum_map! { false => 3, true => 4 };
    ///     let mut values = map.values();
    ///     assert_eq!(values.next(), Some(&3));
    ///     assert_eq!(values.next(), Some(&4));
    ///     assert_eq!(values.next(), None);
    /// }
    /// ```
    pub fn values(&self) -> slice::Iter<V> {
        self.as_slice().iter()
    }

    /// An iterator visiting all values mutably. The iterator type is `&mut V`.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// fn main() {
    ///     let mut map = enum_map! { _ => 2 };
    ///     for value in map.values_mut() {
    ///         *value += 2;
    ///     }
    ///     assert_eq!(map[false], 4);
    ///     assert_eq!(map[true], 4);
    /// }
    /// ```
    pub fn values_mut(&mut self) -> slice::IterMut<V> {
        self.as_mut_slice().iter_mut()
    }

    /// Converts an enum map to a slice representing values.
    pub fn as_slice(&self) -> &[V] {
        K::slice(&self.array)
    }

    /// Converts a mutable enum map to a mutable slice representing values.
    pub fn as_mut_slice(&mut self) -> &mut [V] {
        K::slice_mut(&mut self.array)
    }

    /// Returns a raw pointer to the enum map's buffer.
    ///
    /// The caller must ensure that the slice outlives the pointer this
    /// function returns, or else it will end up pointing to garbage.
    ///
    /// Modifying the container referenced by this slice may cause its buffer
    /// to be reallocated, which would also make any pointers to it invalid.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// use enum_map::EnumMap;
    ///
    /// fn main() {
    ///     let map = enum_map! { 5 => 42, _ => 0 };
    ///     assert_eq!(unsafe { *map.as_ptr().offset(5) }, 42);
    /// }
    /// ```
    pub fn as_ptr(&self) -> *const V {
        self.as_slice().as_ptr()
    }

    /// Returns an unsafe mutable pointer to the enum map's buffer.
    ///
    /// The caller must ensure that the slice outlives the pointer this
    /// function returns, or else it will end up pointing to garbage.
    ///
    /// Modifying the container referenced by this slice may cause its buffer
    /// to be reallocated, which would also make any pointers to it invalid.
    ///
    /// # Examples
    ///
    /// ```
    /// #[macro_use]
    /// extern crate enum_map;
    ///
    /// use enum_map::EnumMap;
    ///
    /// fn main() {
    ///     let mut map = enum_map! { _ => 0 };
    ///     unsafe {
    ///         *map.as_mut_ptr().offset(11) = 23
    ///     };
    ///     assert_eq!(map[11], 23);
    /// }
    /// ```
    pub fn as_mut_ptr(&mut self) -> *mut V {
        self.as_mut_slice().as_mut_ptr()
    }
}

impl<F: FnMut(K) -> V, K: Internal<V>, V> From<F> for EnumMap<K, V> {
    fn from(f: F) -> Self {
        EnumMap { array: K::from_function(f) }
    }
}

/// Enum map constructor.
///
/// This macro allows to create a new enum map in a type safe way. It takes
/// a list of `,` separated pairs separated by `=>`. Left side is `|`
/// separated list of enum keys, or `_` to match all unmatched enum keys,
/// while right side is a value.
///
/// # Examples
///
/// ```
/// #[macro_use]
/// extern crate enum_map;
///
/// #[derive(EnumMap)]
/// enum Example {
///     A,
///     B,
///     C,
///     D,
/// }
///
/// fn main() {
///     let enum_map = enum_map! {
///         Example::A | Example::B => 1,
///         Example::C => 2,
///         _ => 3,
///     };
///     assert_eq!(enum_map[Example::A], 1);
///     assert_eq!(enum_map[Example::B], 1);
///     assert_eq!(enum_map[Example::C], 2);
///     assert_eq!(enum_map[Example::D], 3);
/// }
/// ```
#[macro_export]
macro_rules! enum_map {
    {$($t:tt)*} => {
        $crate::EnumMap::from(|k| match k { $($t)* })
    };
}
<|endoftext|>