<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::{Checksum, ChecksumCapabilities, 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;

        // Disable Checksum verification for improving performance.
        capabilities.checksum = ChecksumCapabilities::default();
        capabilities.checksum.ipv4 = Checksum::Tx;
        capabilities.checksum.tcp = Checksum::Tx;
        capabilities.checksum.udp = Checksum::Tx;
        capabilities.checksum.icmpv4 = Checksum::Tx;
        capabilities.checksum.icmpv6 = Checksum::Tx;

        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));
            // FIXME: This should follows system's setting. 7200 is Linux's default.
            socket.set_timeout(Some(Duration::from_secs(7200)));

            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>Re-enable rx checksum<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));
            // FIXME: This should follows system's setting. 7200 is Linux's default.
            socket.set_timeout(Some(Duration::from_secs(7200)));

            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>// pest. Smart PEGs in Rust
// Copyright (C) 2016  Dragoș Tiselice
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

/// A `macro` that defines each rule as a method on a `Parser`.
///
/// # Examples
///
/// ```
/// # #[macro_use] extern crate pest;
/// # use pest::Input;
/// # use pest::StringInput;
/// # use pest::Rdp;
///
/// # fn main() {
/// impl_rdp!(MyRdp);
///
/// impl MyRdp {
///     grammar! {
///         exp = { paren ~ exp | [""] }
///         paren = { ["("] ~ exp ~ [")"] }
///     }
/// }
///
/// let mut parser = MyRdp::new(Box::new(StringInput::new("(())((())())()")));
///
/// assert!(parser.exp());
/// assert!(parser.end());
/// # }
/// ```
#[macro_export]
macro_rules! grammar {
    // handle parens
    ( @conv $slf:ident [ ( $( $head:tt )* ) $( $tail:tt )* ] [ $( $optail:tt )* ]
      [ $( $output:tt )* ] ) => {
        grammar!(@conv $slf [ $( $head )* @rp $( $tail )* ] [ @lp $( $optail )* ]
                 [ $( $output )* ])
    };
    ( @conv $slf:ident [ @rp $( $tail:tt )* ] [ @lp $( $optail:tt )* ] [ $( $output:tt )* ] ) => {
        grammar!(@conv $slf [ $( $tail )* ] [ $( $optail )* ] [ $( $output )* ])
    };
    ( @conv $slf:ident [ @rp $( $tail:tt )* ] [ $op:tt $( $optail:tt )* ]
      [ $( $output:tt )* ] ) => {
        grammar!(@conv $slf [ @rp $( $tail )* ] [ $( $optail )* ] [ $( $output )* $op ])
    };

    // handle operands
    ( @conv $slf:ident [ ~ $( $tail:tt )* ] [ ~ $( $optail:tt )* ] [ $( $output:tt )* ] ) => {
        grammar!(@conv $slf [ ~ $( $tail )* ] [ $( $optail )* ] [ $( $output )* ~ ])
    };
    ( @conv $slf:ident [ ~ $( $tail:tt )* ] [ $( $optail:tt )* ] $output:tt) => {
        grammar!(@conv $slf [ $( $tail )* ] [ ~ $( $optail )* ] $output)
    };
    ( @conv $slf:ident [ | $( $tail:tt )* ] [ ~ $( $optail:tt )* ] [ $( $output:tt )* ] ) => {
        grammar!(@conv $slf [ | $( $tail )* ] [ $( $optail )* ] [ $( $output )* ~ ])
    };
    ( @conv $slf:ident [ | $( $tail:tt )* ] [ | $( $optail:tt )* ] [ $( $output:tt )* ] ) => {
        grammar!(@conv $slf [ | $( $tail )* ] [ $( $optail )* ] [ $( $output )* | ])
    };
    ( @conv $slf:ident [ | $( $tail:tt )* ] [ $( $optail:tt )* ] $output:tt) => {
        grammar!(@conv $slf [ $( $tail )* ] [ | $( $optail )* ] $output)
    };

    // handle everything else
    ( @conv $slf:ident [ $head:tt $( $tail:tt )* ] $ops:tt [ $( $output:tt )* ] ) => {
        grammar!(@conv $slf [ $( $tail )* ] $ops [ $( $output )* $head ])
    };

    // output remaining operators
    ( @conv $slf:ident [] [] [ $( $output:tt )* ] ) => {
        grammar!(@process $slf [] [ $( $output )* ])
    };
    ( @conv $slf:ident [] [ $op:tt $( $optail:tt )* ] [ $( $output:tt )* ] ) => {
        grammar!(@conv $slf [] [ $( $optail )* ] [ $( $output )* $op ])
    };

    // match
    ( @mtc $slf:ident (( $exp:expr )) ) => (($exp));
    ( @mtc $slf:ident [ $str:expr ]) => ($slf.matches($str));
    ( @mtc $slf:ident $rule:ident) => ($slf.$rule());

    // process postfix
    ( @process $_slf:ident [( $result:expr )] [] ) => ($result);
    ( @process $slf:ident [ $b:tt $a:tt $( $tail:tt )* ] [ ~ $( $optail:tt )* ] ) => {
        {
            let result = $slf.try(Box::new(move |parser| {
                grammar!(@mtc parser $a) && grammar!(@mtc parser $b)
            }));

            grammar!(@process $slf [(( result )) $( $tail )* ] [ $( $optail )* ])
        }
    };
    ( @process $slf:ident [ $b:tt $a:tt $( $tail:tt )* ] [ | $( $optail:tt )* ] ) => {
        {
            let result = grammar!(@mtc $slf $a) || grammar!(@mtc $slf $b);

            grammar!(@process $slf [(( result )) $( $tail )* ] [ $( $optail )* ])
        }
    };
    ( @process $slf:ident [] [ $single:tt ] ) => {
        grammar!(@mtc $slf $single)
    };
    ( @process $slf:ident [ $( $optail:tt )* ] [ $head:tt $( $tail:tt )* ] ) => {
        grammar!(@process $slf [ $head $( $optail )* ] [ $( $tail )* ])
    };

    ( $( $name:ident = { $( $ts:tt )* } )* ) => {
        $(
            #[allow(unused_parens)]
            #[allow(unused_variables)]
            pub fn $name(&mut self) -> bool {
                grammar!(@conv self [ $( $ts )* ] [] [])
            }
        )*
    };
}

#[cfg(test)]
mod tests {
    use super::super::Rdp;
    use super::super::Parser;
    use super::super::Input;
    use super::super::StringInput;

    impl_rdp!(MyRdp);

    impl MyRdp {
        grammar! {
            exp = { paren ~ exp | [""] }
            paren = { ["("] ~ exp ~ [")"] }
        }
    }

    #[test]
    fn basic() {
        let mut parser = MyRdp::new(Box::new(StringInput::new("(())((())())()")));

        assert!(parser.exp());
        assert!(parser.end());
    }

    #[test]
    fn fail() {
        let mut parser = MyRdp::new(Box::new(StringInput::new("(())((())())(")));

        assert!(parser.exp());
        assert!(!parser.end());
    }
}
<commit_msg>Added repetition rules. Fixes #2.<commit_after>// pest. Smart PEGs in Rust
// Copyright (C) 2016  Dragoș Tiselice
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

/// A `macro` that defines each rule as a method on a `Parser`.
///
/// # Examples
///
/// ```
/// # #[macro_use] extern crate pest;
/// # use pest::Input;
/// # use pest::StringInput;
/// # use pest::Rdp;
///
/// # fn main() {
/// impl_rdp!(MyRdp);
///
/// impl MyRdp {
///     grammar! {
///         exp = { paren ~ exp | [""] }
///         paren = { ["("] ~ exp ~ [")"] }
///     }
/// }
///
/// let mut parser = MyRdp::new(Box::new(StringInput::new("(())((())())()")));
///
/// assert!(parser.exp());
/// assert!(parser.end());
/// # }
/// ```
#[macro_export]
macro_rules! grammar {
    // handle parens
    ( @conv $slf:ident [ ( $( $head:tt )* ) $( $tail:tt )* ] [ $( $optail:tt )* ]
      [ $( $output:tt )* ] ) => {
        grammar!(@conv $slf [ $( $head )* @rp $( $tail )* ] [ @lp $( $optail )* ]
                 [ $( $output )* ])
    };
    ( @conv $slf:ident [ @rp $( $tail:tt )* ] [ @lp $( $optail:tt )* ] [ $( $output:tt )* ] ) => {
        grammar!(@conv $slf [ $( $tail )* ] [ $( $optail )* ] [ $( $output )* ])
    };
    ( @conv $slf:ident [ @rp $( $tail:tt )* ] [ $op:tt $( $optail:tt )* ]
      [ $( $output:tt )* ] ) => {
        grammar!(@conv $slf [ @rp $( $tail )* ] [ $( $optail )* ] [ $( $output )* $op ])
    };

    // handle operands
    ( @conv $slf:ident [ ~ $( $tail:tt )* ] [ ~ $( $optail:tt )* ] [ $( $output:tt )* ] ) => {
        grammar!(@conv $slf [ ~ $( $tail )* ] [ $( $optail )* ] [ $( $output )* ~ ])
    };
    ( @conv $slf:ident [ ~ $( $tail:tt )* ] [ $( $optail:tt )* ] $output:tt) => {
        grammar!(@conv $slf [ $( $tail )* ] [ ~ $( $optail )* ] $output)
    };
    ( @conv $slf:ident [ | $( $tail:tt )* ] [ ~ $( $optail:tt )* ] [ $( $output:tt )* ] ) => {
        grammar!(@conv $slf [ | $( $tail )* ] [ $( $optail )* ] [ $( $output )* ~ ])
    };
    ( @conv $slf:ident [ | $( $tail:tt )* ] [ | $( $optail:tt )* ] [ $( $output:tt )* ] ) => {
        grammar!(@conv $slf [ | $( $tail )* ] [ $( $optail )* ] [ $( $output )* | ])
    };
    ( @conv $slf:ident [ | $( $tail:tt )* ] [ $( $optail:tt )* ] $output:tt) => {
        grammar!(@conv $slf [ $( $tail )* ] [ | $( $optail )* ] $output)
    };

    // handle everything else
    ( @conv $slf:ident [ $head:tt $( $tail:tt )* ] $ops:tt [ $( $output:tt )* ] ) => {
        grammar!(@conv $slf [ $( $tail )* ] $ops [ $( $output )* $head ])
    };

    // output remaining operators
    ( @conv $slf:ident [] [] [ $( $output:tt )* ] ) => {
        grammar!(@process $slf [] [ $( $output )* ])
    };
    ( @conv $slf:ident [] [ $op:tt $( $optail:tt )* ] [ $( $output:tt )* ] ) => {
        grammar!(@conv $slf [] [ $( $optail )* ] [ $( $output )* $op ])
    };

    // match
    ( @mtc $slf:ident (( $exp:expr )) ) => (($exp));
    ( @mtc $slf:ident [ $str:expr ]) => ($slf.matches($str));
    ( @mtc $slf:ident $rule:ident) => ($slf.$rule());

    // process postfix
    ( @process $_slf:ident [( $result:expr )] [] ) => ($result);
    ( @process $slf:ident [ $b:tt $a:tt $( $tail:tt )* ] [ ~ $( $optail:tt )* ] ) => {
        {
            let result = $slf.try(Box::new(move |parser| {
                grammar!(@mtc parser $a) && grammar!(@mtc parser $b)
            }));

            grammar!(@process $slf [(( result )) $( $tail )* ] [ $( $optail )* ])
        }
    };
    ( @process $slf:ident [ $b:tt $a:tt $( $tail:tt )* ] [ | $( $optail:tt )* ] ) => {
        {
            let result = grammar!(@mtc $slf $a) || grammar!(@mtc $slf $b);

            grammar!(@process $slf [(( result )) $( $tail )* ] [ $( $optail )* ])
        }
    };
    ( @process $slf:ident [ $a:tt $( $tail:tt )* ] [ * $( $optail:tt )* ] ) => {
        {
            let result = {
                loop {
                    if !grammar!(@mtc $slf $a) {
                        break
                    }
                }

                true
            };

            grammar!(@process $slf [(( result )) $( $tail )* ] [ $( $optail )* ])
        }
    };
    ( @process $slf:ident [ $a:tt $( $tail:tt )* ] [ + $( $optail:tt )* ] ) => {
        {
            let result = if grammar!(@mtc $slf $a) {
                loop {
                    if !grammar!(@mtc $slf $a) {
                        break
                    }
                }

                true
            } else {
                false
            };

            grammar!(@process $slf [(( result )) $( $tail )* ] [ $( $optail )* ])
        }
    };
    ( @process $slf:ident [] [ $single:tt ] ) => {
        grammar!(@mtc $slf $single)
    };
    ( @process $slf:ident [ $( $optail:tt )* ] [ $head:tt $( $tail:tt )* ] ) => {
        grammar!(@process $slf [ $head $( $optail )* ] [ $( $tail )* ])
    };

    ( $( $name:ident = { $( $ts:tt )* } )* ) => {
        $(
            #[allow(unused_parens)]
            #[allow(unused_variables)]
            pub fn $name(&mut self) -> bool {
                grammar!(@conv self [ $( $ts )* ] [] [])
            }
        )*
    };
}

#[cfg(test)]
mod tests {
    use super::super::Rdp;
    use super::super::Parser;
    use super::super::Input;
    use super::super::StringInput;

    impl_rdp!(MyRdp);

    impl MyRdp {
        grammar! {
            exp = { paren ~ exp | [""] }
            paren = { ["("] ~ exp ~ [")"] }
            rep_zero = { ["a"]* }
            rep_one = { ["a"]+ }
        }
    }

    #[test]
    fn basic() {
        let mut parser = MyRdp::new(Box::new(StringInput::new("(())((())())()")));

        assert!(parser.exp());
        assert!(parser.end());
    }

    #[test]
    fn fail() {
        let mut parser = MyRdp::new(Box::new(StringInput::new("(())((())())(")));

        assert!(parser.exp());
        assert!(!parser.end());
    }

    #[test]
    fn rep_zero_empty() {
        let mut parser = MyRdp::new(Box::new(StringInput::new("")));

        assert!(parser.rep_zero());
        assert!(parser.end());
    }

    #[test]
    fn rep_zero_long() {
        let mut parser = MyRdp::new(Box::new(StringInput::new("aaaa")));

        assert!(parser.rep_zero());
        assert!(parser.end());
    }

    #[test]
    fn rep_zero_wrong() {
        let mut parser = MyRdp::new(Box::new(StringInput::new("aaaab")));

        assert!(parser.rep_zero());
        assert!(!parser.end());
    }

    #[test]
    fn rep_one_empty() {
        let mut parser = MyRdp::new(Box::new(StringInput::new("")));

        assert!(!parser.rep_one());
    }

    #[test]
    fn rep_one_long() {
        let mut parser = MyRdp::new(Box::new(StringInput::new("aaaa")));

        assert!(parser.rep_one());
        assert!(parser.end());
    }

    #[test]
    fn rep_one_wrong() {
        let mut parser = MyRdp::new(Box::new(StringInput::new("b")));

        assert!(!parser.rep_one());
        assert!(!parser.end());
    }
}
<|endoftext|>


<commit_before>use super::{Core, Result};
use std::num::Wrapping;

pub fn absolute_word(core: &mut Core) -> Result<u32> {
    core.read_imm_i16().map(|res| res as u32)
}
pub fn absolute_long(core: &mut Core) -> Result<u32> {
    core.read_imm_u32()
}
pub fn predecrement_ay_8(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ay!(core);
    Ok(predecrement_8(core, reg_ndx))
}
pub fn postincrement_ay_8(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ay!(core);
    Ok(postincrement_8(core, reg_ndx))
}
pub fn predecrement_ay_16(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ay!(core);
    Ok(predecrement_16(core, reg_ndx))
}
pub fn postincrement_ay_16(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ay!(core);
    Ok(postincrement_16(core, reg_ndx))
}
pub fn predecrement_ay_32(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ay!(core);
    Ok(predecrement_32(core, reg_ndx))
}
pub fn postincrement_ay_32(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ay!(core);
    Ok(postincrement_32(core, reg_ndx))
}
pub fn address_indirect_ay(core: &mut Core) -> Result<u32> {
    Ok(ay!(core))
}
pub fn displacement_ay(core: &mut Core) -> Result<u32> {
    let reg_val = core.dar[ir_ay!(core)];
    displacement(core, reg_val)
}
pub fn displacement_pc(core: &mut Core) -> Result<u32> {
    let old_pc = core.pc;
    displacement(core, old_pc)
}
pub fn index_ay(core: &mut Core) -> Result<u32> {
    let reg_val = core.dar[ir_ay!(core)];
    index(core, reg_val)
}
pub fn index_pc(core: &mut Core) -> Result<u32> {
    let pc = core.pc;
    index(core, pc)
}
pub fn predecrement_ax_8(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ax!(core);
    Ok(predecrement_8(core, reg_ndx))
}
pub fn predecrement_ax_16(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ax!(core);
    Ok(predecrement_16(core, reg_ndx))
}
pub fn predecrement_ax_32(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ax!(core);
    Ok(predecrement_32(core, reg_ndx))
}
pub fn postincrement_ax_8(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ax!(core);
    Ok(postincrement_8(core, reg_ndx))
}
pub fn postincrement_ax_16(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ax!(core);
    Ok(postincrement_16(core, reg_ndx))
}
pub fn postincrement_ax_32(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ax!(core);
    Ok(postincrement_32(core, reg_ndx))
}

fn predecrement_8(core: &mut Core, reg_ndx: usize) -> u32 {
    // pre-decrement
    core.dar[reg_ndx] = (Wrapping(core.dar[reg_ndx]) - match reg_ndx {
        15 => Wrapping(2), // A7 is kept even
         _ => Wrapping(1)
    }).0;
    core.dar[reg_ndx]
}
fn postincrement_8(core: &mut Core, reg_ndx: usize) -> u32 {
    // post-increment
    let ea = core.dar[reg_ndx];
    core.dar[reg_ndx] = (Wrapping(core.dar[reg_ndx]) + match reg_ndx {
        15 => Wrapping(2), // A7 is kept even
         _ => Wrapping(1)
    }).0;
    ea
}
fn predecrement_16(core: &mut Core, reg_ndx: usize) -> u32 {
    // pre-decrement
    core.dar[reg_ndx] = (Wrapping(core.dar[reg_ndx]) - Wrapping(2)).0;
    core.dar[reg_ndx]
}
fn postincrement_16(core: &mut Core, reg_ndx: usize) -> u32 {
    // post-increment
    let ea = core.dar[reg_ndx];
    core.dar[reg_ndx] = (Wrapping(core.dar[reg_ndx]) + Wrapping(2)).0;
    ea
}
fn predecrement_32(core: &mut Core, reg_ndx: usize) -> u32 {
    // pre-decrement
    core.dar[reg_ndx] = (Wrapping(core.dar[reg_ndx]) - Wrapping(4)).0;
    core.dar[reg_ndx]
}
fn postincrement_32(core: &mut Core, reg_ndx: usize) -> u32 {
    // post-increment
    let ea = core.dar[reg_ndx];
    core.dar[reg_ndx] = (Wrapping(core.dar[reg_ndx]) + Wrapping(4)).0;
    ea
}
pub fn displacement(core: &mut Core, reg_val: u32) -> Result<u32> {
    let displacement = try!(core.read_imm_i16());
    let ea = (Wrapping(reg_val) + Wrapping(displacement as u32)).0;
    Ok(ea)
}
// Brief Extension Word format (see M68000 PRM section 2.1)
const LONG_INDEX_MASK: u16 = 0x0800;
fn index(core: &mut Core, reg_val: u32) -> Result<u32> {
    let extension = try!(core.read_imm_u16());
    // top four bits = (D/A RRR) matches our register array layout
    let xreg_ndx = (extension>>12) as usize;
    let xn = core.dar[xreg_ndx];
    let xn = if (extension & LONG_INDEX_MASK) > 0 {xn} else {(xn as i16) as u32};

      let index = extension as i8;
    let ea = (Wrapping(reg_val) + Wrapping(xn) + Wrapping(index as u32)).0;
    Ok(ea)
}

#[cfg(test)]
mod tests {
    use super::super::Core;
    use super::super::effective_address::{predecrement_8, postincrement_8};

    #[test]
    fn predecrement_wraps() {
        let mut core = Core::new(0x40);
        for i in 0..8 {
            // pre-decrement should wrap to 0xFFFFFFFF
            core.dar[8+i as usize] = 0;
        }
        let ea = predecrement_8(&mut core, 8+0);
        assert_eq!(0xFFFFFFFF, ea);
    }
    #[test]
    fn predecrement_8_wraps_a7_by_two() {
        let mut core = Core::new(0x40);
        for i in 0..8 {
            // pre-decrement should wrap to 0xFFFFFFFF
            core.dar[8+i as usize] = 0;
        }
        let ea = predecrement_8(&mut core, 8+7);
        // a7 is kept even
        assert_eq!(0xFFFFFFFE, ea);
    }
    #[test]
    fn postincrement_wraps() {
        let mut core = Core::new(0x40);
        for i in 0..8 {
            // pre-decrement should wrap to 0xFFFFFFFF
            core.dar[8+i as usize] = 0xFFFFFFFF;
        }
        let ea = postincrement_8(&mut core, 8+0);
        assert_eq!(0xFFFFFFFF, ea);
        assert_eq!(0x0, core.dar[8+0]);
    }
    #[test]
    fn postincrement_8_wraps_a7_by_two() {
        let mut core = Core::new(0x40);
        for i in 0..8 {
            // pre-decrement should wrap to 0xFFFFFFFF
            core.dar[8+i as usize] = 0xFFFFFFFE;
        }
        let ea = postincrement_8(&mut core, 8+7);
        // a7 is kept even
        assert_eq!(0xFFFFFFFE, ea);
        assert_eq!(0x0, core.dar[8+7]);
    }
}
<commit_msg>Add ax-based address indirect, index and displacement<commit_after>use super::{Core, Result};
use std::num::Wrapping;

pub fn absolute_word(core: &mut Core) -> Result<u32> {
    core.read_imm_i16().map(|res| res as u32)
}
pub fn absolute_long(core: &mut Core) -> Result<u32> {
    core.read_imm_u32()
}
pub fn predecrement_ay_8(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ay!(core);
    Ok(predecrement_8(core, reg_ndx))
}
pub fn postincrement_ay_8(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ay!(core);
    Ok(postincrement_8(core, reg_ndx))
}
pub fn predecrement_ay_16(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ay!(core);
    Ok(predecrement_16(core, reg_ndx))
}
pub fn postincrement_ay_16(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ay!(core);
    Ok(postincrement_16(core, reg_ndx))
}
pub fn predecrement_ay_32(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ay!(core);
    Ok(predecrement_32(core, reg_ndx))
}
pub fn postincrement_ay_32(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ay!(core);
    Ok(postincrement_32(core, reg_ndx))
}
pub fn address_indirect_ay(core: &mut Core) -> Result<u32> {
    Ok(ay!(core))
}
pub fn address_indirect_ax(core: &mut Core) -> Result<u32> {
    Ok(ax!(core))
}
pub fn displacement_ay(core: &mut Core) -> Result<u32> {
    let reg_val = core.dar[ir_ay!(core)];
    displacement(core, reg_val)
}
pub fn displacement_ax(core: &mut Core) -> Result<u32> {
    let reg_val = core.dar[ir_ax!(core)];
    displacement(core, reg_val)
}
pub fn displacement_pc(core: &mut Core) -> Result<u32> {
    let old_pc = core.pc;
    displacement(core, old_pc)
}
pub fn index_ay(core: &mut Core) -> Result<u32> {
    let reg_val = core.dar[ir_ay!(core)];
    index(core, reg_val)
}
pub fn index_ax(core: &mut Core) -> Result<u32> {
    let reg_val = core.dar[ir_ax!(core)];
    index(core, reg_val)
}
pub fn index_pc(core: &mut Core) -> Result<u32> {
    let pc = core.pc;
    index(core, pc)
}
pub fn predecrement_ax_8(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ax!(core);
    Ok(predecrement_8(core, reg_ndx))
}
pub fn predecrement_ax_16(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ax!(core);
    Ok(predecrement_16(core, reg_ndx))
}
pub fn predecrement_ax_32(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ax!(core);
    Ok(predecrement_32(core, reg_ndx))
}
pub fn postincrement_ax_8(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ax!(core);
    Ok(postincrement_8(core, reg_ndx))
}
pub fn postincrement_ax_16(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ax!(core);
    Ok(postincrement_16(core, reg_ndx))
}
pub fn postincrement_ax_32(core: &mut Core) -> Result<u32> {
    let reg_ndx = ir_ax!(core);
    Ok(postincrement_32(core, reg_ndx))
}

fn predecrement_8(core: &mut Core, reg_ndx: usize) -> u32 {
    // pre-decrement
    core.dar[reg_ndx] = (Wrapping(core.dar[reg_ndx]) - match reg_ndx {
        15 => Wrapping(2), // A7 is kept even
         _ => Wrapping(1)
    }).0;
    core.dar[reg_ndx]
}
fn postincrement_8(core: &mut Core, reg_ndx: usize) -> u32 {
    // post-increment
    let ea = core.dar[reg_ndx];
    core.dar[reg_ndx] = (Wrapping(core.dar[reg_ndx]) + match reg_ndx {
        15 => Wrapping(2), // A7 is kept even
         _ => Wrapping(1)
    }).0;
    ea
}
fn predecrement_16(core: &mut Core, reg_ndx: usize) -> u32 {
    // pre-decrement
    core.dar[reg_ndx] = (Wrapping(core.dar[reg_ndx]) - Wrapping(2)).0;
    core.dar[reg_ndx]
}
fn postincrement_16(core: &mut Core, reg_ndx: usize) -> u32 {
    // post-increment
    let ea = core.dar[reg_ndx];
    core.dar[reg_ndx] = (Wrapping(core.dar[reg_ndx]) + Wrapping(2)).0;
    ea
}
fn predecrement_32(core: &mut Core, reg_ndx: usize) -> u32 {
    // pre-decrement
    core.dar[reg_ndx] = (Wrapping(core.dar[reg_ndx]) - Wrapping(4)).0;
    core.dar[reg_ndx]
}
fn postincrement_32(core: &mut Core, reg_ndx: usize) -> u32 {
    // post-increment
    let ea = core.dar[reg_ndx];
    core.dar[reg_ndx] = (Wrapping(core.dar[reg_ndx]) + Wrapping(4)).0;
    ea
}
pub fn displacement(core: &mut Core, reg_val: u32) -> Result<u32> {
    let displacement = try!(core.read_imm_i16());
    let ea = (Wrapping(reg_val) + Wrapping(displacement as u32)).0;
    Ok(ea)
}
// Brief Extension Word format (see M68000 PRM section 2.1)
const LONG_INDEX_MASK: u16 = 0x0800;
fn index(core: &mut Core, reg_val: u32) -> Result<u32> {
    let extension = try!(core.read_imm_u16());
    // top four bits = (D/A RRR) matches our register array layout
    let xreg_ndx = (extension>>12) as usize;
    let xn = core.dar[xreg_ndx];
    let xn = if (extension & LONG_INDEX_MASK) > 0 {xn} else {(xn as i16) as u32};

      let index = extension as i8;
    let ea = (Wrapping(reg_val) + Wrapping(xn) + Wrapping(index as u32)).0;
    Ok(ea)
}

#[cfg(test)]
mod tests {
    use super::super::Core;
    use super::super::effective_address::{predecrement_8, postincrement_8};

    #[test]
    fn predecrement_wraps() {
        let mut core = Core::new(0x40);
        for i in 0..8 {
            // pre-decrement should wrap to 0xFFFFFFFF
            core.dar[8+i as usize] = 0;
        }
        let ea = predecrement_8(&mut core, 8+0);
        assert_eq!(0xFFFFFFFF, ea);
    }
    #[test]
    fn predecrement_8_wraps_a7_by_two() {
        let mut core = Core::new(0x40);
        for i in 0..8 {
            // pre-decrement should wrap to 0xFFFFFFFF
            core.dar[8+i as usize] = 0;
        }
        let ea = predecrement_8(&mut core, 8+7);
        // a7 is kept even
        assert_eq!(0xFFFFFFFE, ea);
    }
    #[test]
    fn postincrement_wraps() {
        let mut core = Core::new(0x40);
        for i in 0..8 {
            // pre-decrement should wrap to 0xFFFFFFFF
            core.dar[8+i as usize] = 0xFFFFFFFF;
        }
        let ea = postincrement_8(&mut core, 8+0);
        assert_eq!(0xFFFFFFFF, ea);
        assert_eq!(0x0, core.dar[8+0]);
    }
    #[test]
    fn postincrement_8_wraps_a7_by_two() {
        let mut core = Core::new(0x40);
        for i in 0..8 {
            // pre-decrement should wrap to 0xFFFFFFFF
            core.dar[8+i as usize] = 0xFFFFFFFE;
        }
        let ea = postincrement_8(&mut core, 8+7);
        // a7 is kept even
        assert_eq!(0xFFFFFFFE, ea);
        assert_eq!(0x0, core.dar[8+7]);
    }
}
<|endoftext|>


<commit_before>use ffi;
use HasLua;
use CopyRead;
use ConsumeRead;
use Push;
use LuaTable;
use std::intrinsics::TypeId;

extern fn destructor_wrapper(lua: *mut ffi::lua_State) -> ::libc::c_int {
    use std::mem;

    let impl_raw = unsafe { ffi::lua_touserdata(lua, ffi::lua_upvalueindex(1)) };
    let imp: fn(*mut ffi::lua_State)->::libc::c_int = unsafe { mem::transmute(impl_raw) };

    imp(lua)
}

fn destructor_impl<T>(lua: *mut ffi::lua_State) -> ::libc::c_int {
    use std::mem;

    let obj = unsafe { ffi::lua_touserdata(lua, -1) };
    let obj: &mut T = unsafe { mem::transmute(obj) };
    mem::replace(obj, unsafe { mem::uninitialized() });

    0
}

/// Pushes an object as a user data.
///
/// In Lua, a user data is anything that is not recognized by Lua. When the script attempts to
///  copy a user data, instead only a reference to the data is copied.
///
/// The way a Lua script can use the user data depends on the content of the **metatable**, which is
///  a Lua table linked to the object.
///
/// [See this link for more infos.](http://www.lua.org/manual/5.2/manual.html#2.4)
///
/// # Arguments
///  * metatable: Function that fills the metatable of the object.
// TODO: the type must be Send because the Lua context is Send, but this conflicts with &str
#[experimental]
pub fn push_userdata<L: HasLua, T: 'static>(data: T, lua: &mut L,
    metatable: |&mut LuaTable<L>|) -> uint
{
    use std::mem;

    let typeid = format!("{}", TypeId::of::<T>());

    let luaDataRaw = unsafe { ffi::lua_newuserdata(lua.use_lua(), mem::size_of_val(&data) as ::libc::size_t) };
    let luaData: *mut T = unsafe { mem::transmute(luaDataRaw) };
    unsafe { use std::ptr; ptr::write(luaData, data) };

    let lua_raw = lua.use_lua();

    // creating a metatable
    unsafe {
        ffi::lua_newtable(lua.use_lua());

        // index "__typeid" corresponds to the hash of the TypeId of T
        "__typeid".push_to_lua(lua);
        typeid.push_to_lua(lua);
        ffi::lua_settable(lua.use_lua(), -3);

        // index "__gc" call the object's destructor
        {
            "__gc".push_to_lua(lua);

            // pushing destructor_impl as a lightuserdata
            let destructor_impl: fn(*mut ffi::lua_State)->::libc::c_int = destructor_impl::<T>;
            ffi::lua_pushlightuserdata(lua.use_lua(), mem::transmute(destructor_impl));

            // pushing destructor_wrapper as a closure
            ffi::lua_pushcclosure(lua.use_lua(), mem::transmute(destructor_wrapper), 1);

            ffi::lua_settable(lua.use_lua(), -3);
        }

        {
            let mut table = ConsumeRead::read_from_variable(::LoadedVariable{ lua: lua, size: 1 }).ok().unwrap();
            metatable(&mut table);
            mem::forget(table);
        }

        ffi::lua_setmetatable(lua_raw, -2);
    }

    1
}

// TODO: the type must be Send because the Lua context is Send, but this conflicts with &str
#[experimental]
pub fn read_copy_userdata<L: HasLua, T: Clone + 'static>(lua: &mut L, index: ::libc::c_int) -> Option<T> {
    unsafe {
        let expectedTypeid = format!("{}", TypeId::of::<T>());

        let dataPtr = ffi::lua_touserdata(lua.use_lua(), index);
        if dataPtr.is_null() {
            return None;
        }

        if ffi::lua_getmetatable(lua.use_lua(), -1) == 0 {
            return None;
        }

        "__typeid".push_to_lua(lua);
        ffi::lua_gettable(lua.use_lua(), -2);
        if CopyRead::read_from_lua(lua, -1) != Some(expectedTypeid) {
            return None;
        }
        ffi::lua_pop(lua.use_lua(), -2);

        let data: &T = ::std::mem::transmute(dataPtr);
        Some(data.clone())
    }
}
<commit_msg>Parameters for push_userdata and read_copy_userdata must now be Send<commit_after>use ffi;
use HasLua;
use CopyRead;
use ConsumeRead;
use Push;
use LuaTable;
use std::intrinsics::TypeId;

extern fn destructor_wrapper(lua: *mut ffi::lua_State) -> ::libc::c_int {
    use std::mem;

    let impl_raw = unsafe { ffi::lua_touserdata(lua, ffi::lua_upvalueindex(1)) };
    let imp: fn(*mut ffi::lua_State)->::libc::c_int = unsafe { mem::transmute(impl_raw) };

    imp(lua)
}

fn destructor_impl<T>(lua: *mut ffi::lua_State) -> ::libc::c_int {
    use std::mem;

    let obj = unsafe { ffi::lua_touserdata(lua, -1) };
    let obj: &mut T = unsafe { mem::transmute(obj) };
    mem::replace(obj, unsafe { mem::uninitialized() });

    0
}

/// Pushes an object as a user data.
///
/// In Lua, a user data is anything that is not recognized by Lua. When the script attempts to
///  copy a user data, instead only a reference to the data is copied.
///
/// The way a Lua script can use the user data depends on the content of the **metatable**, which is
///  a Lua table linked to the object.
///
/// [See this link for more infos.](http://www.lua.org/manual/5.2/manual.html#2.4)
///
/// # Arguments
///  * metatable: Function that fills the metatable of the object.
#[experimental]
pub fn push_userdata<L: HasLua, T: 'static + Send>(data: T, lua: &mut L,
    metatable: |&mut LuaTable<L>|) -> uint
{
    use std::mem;

    let typeid = format!("{}", TypeId::of::<T>());

    let luaDataRaw = unsafe { ffi::lua_newuserdata(lua.use_lua(), mem::size_of_val(&data) as ::libc::size_t) };
    let luaData: *mut T = unsafe { mem::transmute(luaDataRaw) };
    unsafe { use std::ptr; ptr::write(luaData, data) };

    let lua_raw = lua.use_lua();

    // creating a metatable
    unsafe {
        ffi::lua_newtable(lua.use_lua());

        // index "__typeid" corresponds to the hash of the TypeId of T
        "__typeid".push_to_lua(lua);
        typeid.push_to_lua(lua);
        ffi::lua_settable(lua.use_lua(), -3);

        // index "__gc" call the object's destructor
        {
            "__gc".push_to_lua(lua);

            // pushing destructor_impl as a lightuserdata
            let destructor_impl: fn(*mut ffi::lua_State)->::libc::c_int = destructor_impl::<T>;
            ffi::lua_pushlightuserdata(lua.use_lua(), mem::transmute(destructor_impl));

            // pushing destructor_wrapper as a closure
            ffi::lua_pushcclosure(lua.use_lua(), mem::transmute(destructor_wrapper), 1);

            ffi::lua_settable(lua.use_lua(), -3);
        }

        {
            let mut table = ConsumeRead::read_from_variable(::LoadedVariable{ lua: lua, size: 1 }).ok().unwrap();
            metatable(&mut table);
            mem::forget(table);
        }

        ffi::lua_setmetatable(lua_raw, -2);
    }

    1
}

#[experimental]
pub fn read_copy_userdata<L: HasLua, T: Clone + 'static>(lua: &mut L, index: ::libc::c_int) -> Option<T> {
    unsafe {
        let expectedTypeid = format!("{}", TypeId::of::<T>());

        let dataPtr = ffi::lua_touserdata(lua.use_lua(), index);
        if dataPtr.is_null() {
            return None;
        }

        if ffi::lua_getmetatable(lua.use_lua(), -1) == 0 {
            return None;
        }

        "__typeid".push_to_lua(lua);
        ffi::lua_gettable(lua.use_lua(), -2);
        if CopyRead::read_from_lua(lua, -1) != Some(expectedTypeid) {
            return None;
        }
        ffi::lua_pop(lua.use_lua(), -2);

        let data: &T = ::std::mem::transmute(dataPtr);
        Some(data.clone())
    }
}
<|endoftext|>


<commit_before>//! Defines the BSDF which acts as a container for composing the various BRDFs
//! and BTDFs that describe the surface's properties

use std::cmp;
use enum_set::EnumSet;

use linalg::{self, Normal, Vector, Point};
use film::Colorf;
use geometry::DifferentialGeometry;
use bxdf::{BxDF, BxDFType};
use sampler::Sample;

/// The BSDF contains the various BRDFs and BTDFs that describe the surface's properties
/// at some point. It also transforms incident and outgoing light directions into
/// shading space to make the BxDFs easier to implement.
/// TODO: We really need the memory pool. Each time we get the bsdf from a
/// material we need to allocate a decent amount of stuff since they each need
/// their own tangent, bitangent and differential geometry reference.
pub struct BSDF<'a> {
    /// The hit point
    pub p: Point,
    /// Shading normal, may be perturbed by bump mapping
    pub n: Normal,
    /// The actual geometry normal
    pub ng: Normal,
    /// Tangent vector for the surface
    pub tan: Vector,
    /// Bitangent vector for the surface
    pub bitan: Vector,
    /// Refractive index of the geometry
    pub eta: f32,
    /// TODO: Currently a Vec is safe to use but once in the memory pool it
    /// will leak since it won't be dropped. This would also migrate our BxDFs
    /// from Box<BxDF> to &BxDF. When unboxed traits land we can move to unboxed
    /// BxDFs here though.
    bxdfs: &'a [Box<BxDF + Send + Sync>],
}

impl<'a> BSDF<'a> {
    /// Create a new BSDF using the BxDFs passed to shade the differential geometry with
    /// refractive index `eta`
    pub fn new(bxdfs: &'a [Box<BxDF + Send + Sync>], eta: f32,
               dg: &DifferentialGeometry<'a>)
               -> BSDF<'a> {
        let n = dg.n.normalized();
        let mut bitan = dg.dp_du.normalized();
        let tan = linalg::cross(&n, &bitan);
        bitan = linalg::cross(&tan, &n);
        BSDF { p: dg.p, n: n, ng: dg.ng, tan: tan, bitan: bitan, bxdfs: bxdfs, eta: eta }
    }
    /// Return the total number of BxDFs
    pub fn num_bxdfs(&self) -> usize { self.bxdfs.len() }
    /// Return the number of BxDFs matching the flags
    pub fn num_matching(&self, flags: EnumSet<BxDFType>) -> usize {
        self.bxdfs.iter().filter(|ref x| x.matches(flags)).count()
    }
    /// Transform the vector from world space to shading space
    pub fn to_shading(&self, v: &Vector) -> Vector {
        Vector::new(linalg::dot(v, &self.bitan), linalg::dot(v, &self.tan),
                    linalg::dot(v, &self.n))
    }
    /// Transform the vectro from shading space to world space
    pub fn from_shading(&self, v: &Vector) -> Vector {
        Vector::new(self.bitan.x * v.x + self.tan.x * v.y + self.n.x * v.z,
                    self.bitan.y * v.x + self.tan.y * v.y + self.n.y * v.z,
                    self.bitan.z * v.x + self.tan.z * v.y + self.n.z * v.z)
    }
    /// Evaluate the BSDF for the outgoing and incident light directions
    /// `w_o` and `w_i` in world space, sampling the desired subset of BxDFs
    /// selected by the flags passed. `wo_world` and `wi_world` should point from
    /// the hit point in the outgoing and incident light directions respectively.
    pub fn eval(&self, wo_world: &Vector, wi_world: &Vector, mut flags: EnumSet<BxDFType>) -> Colorf {
        let w_o = self.to_shading(wo_world);
        let w_i = self.to_shading(wi_world);
        // Determine if we should evaluate reflection or transmission based on the
        // geometry normal and the light directions
        if linalg::dot(wo_world, &self.ng) * linalg::dot(wi_world, &self.ng) > 0.0 {
            flags.remove(&BxDFType::Transmission);
        } else {
            flags.remove(&BxDFType::Reflection);
        }
        // Find all matching BxDFs and add their contribution to the material's color
        self.bxdfs.iter().filter_map(|ref x| if x.matches(flags) { Some(x.eval(&w_o, &w_i)) } else { None })
            .fold(Colorf::broadcast(0.0), |x, y| x + y)
    }
    /// Sample a component of the BSDF to get an incident light direction for light
    /// leaving the surface along `w_o`.
    /// `samples` are the 3 random values to use when sampling a component of the BSDF
    /// and a the chosen BSDF
    /// Returns the color, direction, pdf and the type of BxDF that was sampled.
    pub fn sample(&self, wo_world: &Vector, flags: EnumSet<BxDFType>, samples: &Sample)
        -> (Colorf, Vector, f32, EnumSet<BxDFType>) {
        let n_matching = self.num_matching(flags);
        if n_matching == 0 {
            return (Colorf::broadcast(0.0), Vector::broadcast(0.0), 0.0, EnumSet::new());
        }
        let comp = cmp::min((samples.one_d * n_matching as f32) as usize, n_matching - 1);
        let bxdf = self.matching_at(comp, flags);
        let w_o = self.to_shading(wo_world).normalized();
        let (mut f, w_i, mut pdf) = bxdf.sample(&w_o, &samples.two_d);
        if w_i.length_sqr() == 0.0 {
            return (Colorf::broadcast(0.0), Vector::broadcast(0.0), 0.0, EnumSet::new());
        }

        let wi_world = self.from_shading(&w_i).normalized();
        // TODO: We re-use our functions but actually do a lot of redundant computation. I'm not
        // sure that the compiler will eliminate it. Should just copy in the code from pdf and eval
        if !bxdf.bxdf_type().contains(&BxDFType::Specular) && n_matching > 1 {
            pdf = self.pdf(wo_world, &wi_world, flags);
        }
        if n_matching > 1 {
            pdf /= n_matching as f32;
        }

        if !bxdf.bxdf_type().contains(&BxDFType::Specular) {
            f = self.eval(wo_world, &wi_world, flags);
        }
        (f, wi_world, pdf, bxdf.bxdf_type())
    }
    /// Compute the pdf for sampling the pair of incident and outgoing light directions for
    /// the BxDFs matching the flags set
    pub fn pdf(&self, wo_world: &Vector, wi_world: &Vector, flags: EnumSet<BxDFType>) -> f32 {
        // should also normalize?
        let w_o = self.to_shading(wo_world).normalized();
        let w_i = self.to_shading(wi_world).normalized();
        let (pdf_val, n_comps) = self.bxdfs.iter()
            .filter_map(|ref x| if x.matches(flags) { Some(x.pdf(&w_o, &w_i)) } else { None })
            .fold((0.0, 0), |(p, n), y| (p + y, n + 1));
        if n_comps > 0 {
            pdf_val / n_comps as f32
        } else {
            0.0
        }
    }
    /// Get the `i`th BxDF that matches the flags passed. There should not be fewer than i
    /// BxDFs that match the flags
    fn matching_at(&self, i: usize, flags: EnumSet<BxDFType>) -> &Box<BxDF + Send + Sync> {
        let mut it = self.bxdfs.iter().filter(|ref x| x.matches(flags)).skip(i);
        match it.next() {
            Some(b) => b,
            None => panic!("Out of bounds index for BxDF type {:?}", flags)
        }
    }
}

<commit_msg>Small spacing tweak<commit_after>//! Defines the BSDF which acts as a container for composing the various BRDFs
//! and BTDFs that describe the surface's properties

use std::cmp;
use enum_set::EnumSet;

use linalg::{self, Normal, Vector, Point};
use film::Colorf;
use geometry::DifferentialGeometry;
use bxdf::{BxDF, BxDFType};
use sampler::Sample;

/// The BSDF contains the various BRDFs and BTDFs that describe the surface's properties
/// at some point. It also transforms incident and outgoing light directions into
/// shading space to make the BxDFs easier to implement.
/// TODO: We really need the memory pool. Each time we get the bsdf from a
/// material we need to allocate a decent amount of stuff since they each need
/// their own tangent, bitangent and differential geometry reference.
pub struct BSDF<'a> {
    /// The hit point
    pub p: Point,
    /// Shading normal, may be perturbed by bump mapping
    pub n: Normal,
    /// The actual geometry normal
    pub ng: Normal,
    /// Tangent vector for the surface
    pub tan: Vector,
    /// Bitangent vector for the surface
    pub bitan: Vector,
    /// Refractive index of the geometry
    pub eta: f32,
    /// TODO: Currently a Vec is safe to use but once in the memory pool it
    /// will leak since it won't be dropped. This would also migrate our BxDFs
    /// from Box<BxDF> to &BxDF. When unboxed traits land we can move to unboxed
    /// BxDFs here though.
    bxdfs: &'a [Box<BxDF + Send + Sync>],
}

impl<'a> BSDF<'a> {
    /// Create a new BSDF using the BxDFs passed to shade the differential geometry with
    /// refractive index `eta`
    pub fn new(bxdfs: &'a [Box<BxDF + Send + Sync>], eta: f32,
               dg: &DifferentialGeometry<'a>)
               -> BSDF<'a> {
        let n = dg.n.normalized();
        let mut bitan = dg.dp_du.normalized();
        let tan = linalg::cross(&n, &bitan);
        bitan = linalg::cross(&tan, &n);
        BSDF { p: dg.p, n: n, ng: dg.ng, tan: tan, bitan: bitan, bxdfs: bxdfs, eta: eta }
    }
    /// Return the total number of BxDFs
    pub fn num_bxdfs(&self) -> usize { self.bxdfs.len() }
    /// Return the number of BxDFs matching the flags
    pub fn num_matching(&self, flags: EnumSet<BxDFType>) -> usize {
        self.bxdfs.iter().filter(|ref x| x.matches(flags)).count()
    }
    /// Transform the vector from world space to shading space
    pub fn to_shading(&self, v: &Vector) -> Vector {
        Vector::new(linalg::dot(v, &self.bitan), linalg::dot(v, &self.tan),
                    linalg::dot(v, &self.n))
    }
    /// Transform the vectro from shading space to world space
    pub fn from_shading(&self, v: &Vector) -> Vector {
        Vector::new(self.bitan.x * v.x + self.tan.x * v.y + self.n.x * v.z,
                    self.bitan.y * v.x + self.tan.y * v.y + self.n.y * v.z,
                    self.bitan.z * v.x + self.tan.z * v.y + self.n.z * v.z)
    }
    /// Evaluate the BSDF for the outgoing and incident light directions
    /// `w_o` and `w_i` in world space, sampling the desired subset of BxDFs
    /// selected by the flags passed. `wo_world` and `wi_world` should point from
    /// the hit point in the outgoing and incident light directions respectively.
    pub fn eval(&self, wo_world: &Vector, wi_world: &Vector, mut flags: EnumSet<BxDFType>) -> Colorf {
        let w_o = self.to_shading(wo_world);
        let w_i = self.to_shading(wi_world);
        // Determine if we should evaluate reflection or transmission based on the
        // geometry normal and the light directions
        if linalg::dot(wo_world, &self.ng) * linalg::dot(wi_world, &self.ng) > 0.0 {
            flags.remove(&BxDFType::Transmission);
        } else {
            flags.remove(&BxDFType::Reflection);
        }
        // Find all matching BxDFs and add their contribution to the material's color
        self.bxdfs.iter().filter_map(|ref x| if x.matches(flags) { Some(x.eval(&w_o, &w_i)) } else { None })
            .fold(Colorf::broadcast(0.0), |x, y| x + y)
    }
    /// Sample a component of the BSDF to get an incident light direction for light
    /// leaving the surface along `w_o`.
    /// `samples` are the 3 random values to use when sampling a component of the BSDF
    /// and a the chosen BSDF
    /// Returns the color, direction, pdf and the type of BxDF that was sampled.
    pub fn sample(&self, wo_world: &Vector, flags: EnumSet<BxDFType>, samples: &Sample)
        -> (Colorf, Vector, f32, EnumSet<BxDFType>) {
        let n_matching = self.num_matching(flags);
        if n_matching == 0 {
            return (Colorf::broadcast(0.0), Vector::broadcast(0.0), 0.0, EnumSet::new());
        }
        let comp = cmp::min((samples.one_d * n_matching as f32) as usize, n_matching - 1);
        let bxdf = self.matching_at(comp, flags);
        let w_o = self.to_shading(wo_world).normalized();
        let (mut f, w_i, mut pdf) = bxdf.sample(&w_o, &samples.two_d);
        if w_i.length_sqr() == 0.0 {
            return (Colorf::broadcast(0.0), Vector::broadcast(0.0), 0.0, EnumSet::new());
        }
        let wi_world = self.from_shading(&w_i).normalized();

        // TODO: We re-use our functions but actually do a lot of redundant computation. I'm not
        // sure that the compiler will eliminate it. Should just copy in the code from pdf and eval
        if !bxdf.bxdf_type().contains(&BxDFType::Specular) && n_matching > 1 {
            pdf = self.pdf(wo_world, &wi_world, flags);
        }
        if n_matching > 1 {
            pdf /= n_matching as f32;
        }

        if !bxdf.bxdf_type().contains(&BxDFType::Specular) {
            f = self.eval(wo_world, &wi_world, flags);
        }
        (f, wi_world, pdf, bxdf.bxdf_type())
    }
    /// Compute the pdf for sampling the pair of incident and outgoing light directions for
    /// the BxDFs matching the flags set
    pub fn pdf(&self, wo_world: &Vector, wi_world: &Vector, flags: EnumSet<BxDFType>) -> f32 {
        // should also normalize?
        let w_o = self.to_shading(wo_world).normalized();
        let w_i = self.to_shading(wi_world).normalized();
        let (pdf_val, n_comps) = self.bxdfs.iter()
            .filter_map(|ref x| if x.matches(flags) { Some(x.pdf(&w_o, &w_i)) } else { None })
            .fold((0.0, 0), |(p, n), y| (p + y, n + 1));
        if n_comps > 0 {
            pdf_val / n_comps as f32
        } else {
            0.0
        }
    }
    /// Get the `i`th BxDF that matches the flags passed. There should not be fewer than i
    /// BxDFs that match the flags
    fn matching_at(&self, i: usize, flags: EnumSet<BxDFType>) -> &Box<BxDF + Send + Sync> {
        let mut it = self.bxdfs.iter().filter(|ref x| x.matches(flags)).skip(i);
        match it.next() {
            Some(b) => b,
            None => panic!("Out of bounds index for BxDF type {:?}", flags)
        }
    }
}

<|endoftext|>