<commit_before>use vec2;
use vec3;
use num::Num;

#[inline(always)]
pub fn transform_mat3<'a, 'b, T: Num>(out: &'a mut [T; 4], a: &'b [T; 4], m: [T; 9]) -> &'a mut [T; 4] {
    out[0] = a[0] * m[0] + a[1] * m[3] + a[2] * m[6];
    out[1] = a[0] * m[1] + a[1] * m[4] + a[2] * m[7];
    out[2] = a[0] * m[2] + a[1] * m[5] + a[2] * m[8];
    out
}

#[inline(always)]
pub fn transform_mat4<'a, 'b, T: Num>(out: &'a mut [T; 4], a: &'b [T; 4], m: [T; 16]) -> &'a mut [T; 4] {
    out[0] = a[0] * m[0] + a[1] * m[4] + a[2] * m[8] + a[3] * m[12];
    out[1] = a[0] * m[1] + a[1] * m[5] + a[2] * m[9] + a[3] * m[13];
    out[2] = a[0] * m[2] + a[1] * m[6] + a[2] * m[10] + a[3] * m[14];
    out[3] = a[0] * m[3] + a[1] * m[7] + a[2] * m[11] + a[3] * m[15];
    out
}

#[inline(always)]
pub fn transform_mat4_rotation<'a, 'b, T: Num>(out: &'a mut [T; 4], a: &'b [T; 4], m: [T; 16]) -> &'a mut [T; 4] {
    out[0] = a[0] * m[0] + a[1] * m[4] + a[2] * m[8];
    out[1] = a[0] * m[1] + a[1] * m[5] + a[2] * m[9];
    out[2] = a[0] * m[2] + a[1] * m[6] + a[2] * m[10];
    out[3] = a[0] * m[3] + a[1] * m[7] + a[2] * m[11] + a[3] * m[15];
    out
}

#[inline(always)]
pub fn transform_projection<'a, 'b, T: Num>(out: &'a mut [T; 4], a: &'b [T; 4], m: [T; 16]) -> &'a mut [T; 4] {
    let mut d = a[0] * m[3] + a[1] * m[7] + a[2] * m[11] + a[3] * m[15];
    d = if d != T::zero() {T::one() / d} else {d};

    out[0] = (a[0] * m[0] + a[1] * m[4] + a[2] * m[8] + a[3] * m[12]) * d;
    out[1] = (a[0] * m[1] + a[1] * m[5] + a[2] * m[9] + a[3] * m[13]) * d;
    out[2] = (a[0] * m[2] + a[1] * m[6] + a[2] * m[10] + a[3] * m[14]) * d;
    out[3] = (a[0] * m[3] + a[1] * m[7] + a[2] * m[11] + a[3] * m[15]) * d;
    out
}

#[inline(always)]
pub fn transform_projection_no_position<'a, 'b, T: Num>(out: &'a mut [T; 4], a: &'b [T; 4], m: [T; 16]) -> &'a mut [T; 4] {
    let mut d = a[0] * m[3] + a[1] * m[7] + a[2] * m[11] + a[3] * m[15];
    d = if d != T::zero() {T::one() / d} else {d};

    out[0] = (a[0] * m[0] + a[1] * m[4] + a[2] * m[8] + a[3]) * d;
    out[1] = (a[0] * m[1] + a[1] * m[5] + a[2] * m[9] + a[3]) * d;
    out[2] = (a[0] * m[2] + a[1] * m[6] + a[2] * m[10] + a[3]) * d;
    out[3] = (a[0] * m[3] + a[1] * m[7] + a[2] * m[11] + a[3] * m[15]) * d;
    out
}

#[inline(always)]
pub fn position_mat32<'a, 'b, T: Num>(out: &'a mut [T; 4], m: [T; 6]) -> &'a mut [T; 4] {
    out[0] = m[4];
    out[1] = m[5];
    out
}

#[inline(always)]
pub fn position_mat4<'a, 'b, T: Num>(out: &'a mut [T; 4], m: [T; 16]) -> &'a mut [T; 4] {
    out[0] = m[12];
    out[1] = m[13];
    out[2] = m[14];
    out[3] = m[15];
    out
}

#[inline(always)]
pub fn scale_mat2<'a, 'b, T: Num>(out: &'a mut [T; 4], m: &'b [T; 4]) -> &'a mut [T; 4] {
    out[0] = vec2::length_values(m[0], m[2]);
    out[1] = vec2::length_values(m[1], m[3]);
    out[2] = T::one();
    out[3] = T::one();
    out
}

#[inline(always)]
pub fn scale_mat32<'a, 'b, T: Num>(out: &'a mut [T; 4], m: [T; 6]) -> &'a mut [T; 4] {
    out[0] = vec3::length_values(m[0], m[2], T::zero());
    out[1] = vec3::length_values(m[1], m[3], T::zero());
    out[2] = T::one();
    out[3] = T::one();
    out
}

#[inline(always)]
pub fn scale_mat3<'a, 'b, T: Num>(out: &'a mut [T; 4], m: [T; 9]) -> &'a mut [T; 4] {
    out[0] = vec3::length_values(m[0], m[3], m[6]);
    out[1] = vec3::length_values(m[1], m[4], m[7]);
    out[2] = vec3::length_values(m[2], m[5], m[8]);
    out[3] = T::one();
    out
}

#[inline(always)]
pub fn scale_mat4<'a, 'b, T: Num>(out: &'a mut [T; 4], m: [T; 16]) -> &'a mut [T; 4] {
    out[0] = vec3::length_values(m[0], m[4], m[8]);
    out[1] = vec3::length_values(m[1], m[5], m[9]);
    out[2] = vec3::length_values(m[2], m[6], m[10]);
    out[3] = T::one();
    out
}
<commit_msg>updates<commit_after>use vec2;
use vec3;
use num::Num;

#[inline(always)]
pub fn transform_mat3<'a, 'b, T: Num>(out: &'a mut [T; 4], a: &'b [T; 4], m: &'b [T; 9]) -> &'a mut [T; 4] {
    out[0] = a[0] * m[0] + a[1] * m[3] + a[2] * m[6];
    out[1] = a[0] * m[1] + a[1] * m[4] + a[2] * m[7];
    out[2] = a[0] * m[2] + a[1] * m[5] + a[2] * m[8];
    out
}

#[inline(always)]
pub fn transform_mat4<'a, 'b, T: Num>(out: &'a mut [T; 4], a: &'b [T; 4], m: &'b [T; 16]) -> &'a mut [T; 4] {
    out[0] = a[0] * m[0] + a[1] * m[4] + a[2] * m[8] + a[3] * m[12];
    out[1] = a[0] * m[1] + a[1] * m[5] + a[2] * m[9] + a[3] * m[13];
    out[2] = a[0] * m[2] + a[1] * m[6] + a[2] * m[10] + a[3] * m[14];
    out[3] = a[0] * m[3] + a[1] * m[7] + a[2] * m[11] + a[3] * m[15];
    out
}

#[inline(always)]
pub fn transform_mat4_rotation<'a, 'b, T: Num>(out: &'a mut [T; 4], a: &'b [T; 4], m: &'b [T; 16]) -> &'a mut [T; 4] {
    out[0] = a[0] * m[0] + a[1] * m[4] + a[2] * m[8];
    out[1] = a[0] * m[1] + a[1] * m[5] + a[2] * m[9];
    out[2] = a[0] * m[2] + a[1] * m[6] + a[2] * m[10];
    out[3] = a[0] * m[3] + a[1] * m[7] + a[2] * m[11] + a[3] * m[15];
    out
}

#[inline(always)]
pub fn transform_projection<'a, 'b, T: Num>(out: &'a mut [T; 4], a: &'b [T; 4], m: &'b [T; 16]) -> &'a mut [T; 4] {
    let mut d = a[0] * m[3] + a[1] * m[7] + a[2] * m[11] + a[3] * m[15];
    d = if d != T::zero() {T::one() / d} else {d};

    out[0] = (a[0] * m[0] + a[1] * m[4] + a[2] * m[8] + a[3] * m[12]) * d;
    out[1] = (a[0] * m[1] + a[1] * m[5] + a[2] * m[9] + a[3] * m[13]) * d;
    out[2] = (a[0] * m[2] + a[1] * m[6] + a[2] * m[10] + a[3] * m[14]) * d;
    out[3] = (a[0] * m[3] + a[1] * m[7] + a[2] * m[11] + a[3] * m[15]) * d;
    out
}

#[inline(always)]
pub fn transform_projection_no_position<'a, 'b, T: Num>(out: &'a mut [T; 4], a: &'b [T; 4], m: &'b [T; 16]) -> &'a mut [T; 4] {
    let mut d = a[0] * m[3] + a[1] * m[7] + a[2] * m[11] + a[3] * m[15];
    d = if d != T::zero() {T::one() / d} else {d};

    out[0] = (a[0] * m[0] + a[1] * m[4] + a[2] * m[8] + a[3]) * d;
    out[1] = (a[0] * m[1] + a[1] * m[5] + a[2] * m[9] + a[3]) * d;
    out[2] = (a[0] * m[2] + a[1] * m[6] + a[2] * m[10] + a[3]) * d;
    out[3] = (a[0] * m[3] + a[1] * m[7] + a[2] * m[11] + a[3] * m[15]) * d;
    out
}

#[inline(always)]
pub fn position_mat32<'a, 'b, T: Num>(out: &'a mut [T; 4], m: &'b [T; 6]) -> &'a mut [T; 4] {
    out[0] = m[4];
    out[1] = m[5];
    out
}

#[inline(always)]
pub fn position_mat4<'a, 'b, T: Num>(out: &'a mut [T; 4], m: &'b [T; 16]) -> &'a mut [T; 4] {
    out[0] = m[12];
    out[1] = m[13];
    out[2] = m[14];
    out[3] = m[15];
    out
}

#[inline(always)]
pub fn scale_mat2<'a, 'b, T: Num>(out: &'a mut [T; 4], m: &'b [T; 4]) -> &'a mut [T; 4] {
    out[0] = vec2::length_values(m[0], m[2]);
    out[1] = vec2::length_values(m[1], m[3]);
    out[2] = T::one();
    out[3] = T::one();
    out
}

#[inline(always)]
pub fn scale_mat32<'a, 'b, T: Num>(out: &'a mut [T; 4], m: &'b [T; 6]) -> &'a mut [T; 4] {
    out[0] = vec3::length_values(m[0], m[2], T::zero());
    out[1] = vec3::length_values(m[1], m[3], T::zero());
    out[2] = T::one();
    out[3] = T::one();
    out
}

#[inline(always)]
pub fn scale_mat3<'a, 'b, T: Num>(out: &'a mut [T; 4], m: &'b [T; 9]) -> &'a mut [T; 4] {
    out[0] = vec3::length_values(m[0], m[3], m[6]);
    out[1] = vec3::length_values(m[1], m[4], m[7]);
    out[2] = vec3::length_values(m[2], m[5], m[8]);
    out[3] = T::one();
    out
}

#[inline(always)]
pub fn scale_mat4<'a, 'b, T: Num>(out: &'a mut [T; 4], m: &'b [T; 16]) -> &'a mut [T; 4] {
    out[0] = vec3::length_values(m[0], m[4], m[8]);
    out[1] = vec3::length_values(m[1], m[5], m[9]);
    out[2] = vec3::length_values(m[2], m[6], m[10]);
    out[3] = T::one();
    out
}
<|endoftext|>


<commit_before>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for implementing transport layer protocols
//!
//! The transport module provides the ability to send and receive packets at
//! the transport layer using IPv4 or IPv6. It also enables layer 3 networking
//! for specific transport protocols, using IPv4 only.
//!
//! Note that this is limited by operating system support - for example, on OS
//! X and FreeBSD, it is impossible to implement protocols which are already
//! implemented in the kernel such as TCP and UDP.

#![macro_use]

extern crate libc;

use self::TransportProtocol::{Ipv4, Ipv6};
use self::TransportChannelType::{Layer3, Layer4};

use std::io;
use std::io::Error;
use std::iter::repeat;
use std::net::{self, IpAddr};
use std::mem;
use std::sync::Arc;

use packet::Packet;
use packet::ip::IpNextHeaderProtocol;
use packet::ipv4::Ipv4Packet;
use packet::udp::UdpPacket;

use internal;
use sockets;

/// Represents a transport layer protocol
#[derive(Clone, Copy)]
pub enum TransportProtocol {
    /// Represents a transport protocol built on top of IPv4
    Ipv4(IpNextHeaderProtocol),
    /// Represents a transport protocol built on top of IPv6
    Ipv6(IpNextHeaderProtocol),
}

/// Type of transport channel to present
#[derive(Clone, Copy)]
pub enum TransportChannelType {
    /// The application will send and receive transport layer packets
    Layer4(TransportProtocol),
    /// The application will send and receive IPv4 packets, with the specified transport protocol
    Layer3(IpNextHeaderProtocol),
}

/// Structure used for sending at the transport layer. Should be created with transport_channel()
pub struct TransportSender {
    socket: Arc<internal::FileDesc>,
    _channel_type: TransportChannelType,
}

/// Structure used for sending at the transport layer. Should be created with transport_channel()
pub struct TransportReceiver {
    socket: Arc<internal::FileDesc>,
    buffer: Vec<u8>,
    channel_type: TransportChannelType,
}

/// Create a new (TransportSender, TransportReceiver) pair
///
/// This allows for sending and receiving packets at the transport layer. The buffer size should be
/// large enough to handle the largest packet you wish to receive.
///
/// The channel type specifies what layer to send and receive packets at, and the transport
/// protocol you wish to implement. For example, `Layer4(Ipv4(IpNextHeaderProtocols::Udp))` would
/// allow sending and receiving UDP packets using IPv4; whereas Layer3(IpNextHeaderProtocols::Udp)
/// would include the IPv4 Header in received values, and require manual construction of an IP
/// header when sending.
pub fn transport_channel(buffer_size: usize,
                         channel_type: TransportChannelType)
    -> io::Result<(TransportSender, TransportReceiver)> {
    use std::net;

    // This hack makes sure that winsock is initialised
    let _ = {
        let ip = net::Ipv4Addr::new(255, 255, 255, 255);
        let sockaddr = net::SocketAddr::V4(net::SocketAddrV4::new(ip, 0));

        net::UdpSocket::bind(sockaddr)
    };

    let socket = unsafe {
        match channel_type {
            Layer4(Ipv4(IpNextHeaderProtocol(proto))) | Layer3(IpNextHeaderProtocol(proto)) =>
                sockets::socket(sockets::AF_INET, sockets::SOCK_RAW, proto as libc::c_int),
            Layer4(Ipv6(IpNextHeaderProtocol(proto))) =>
                sockets::socket(sockets::AF_INET6, sockets::SOCK_RAW, proto as libc::c_int),
        }
    };
    if socket == sockets::INVALID_SOCKET {
        return Err(Error::last_os_error());
    }

    if match channel_type {
        Layer3(_) | Layer4(Ipv4(_)) => true,
        _ => false,
    } {
        let hincl: libc::c_int = match channel_type {
            Layer4(..) => 0,
            _ => 1,
        };
        let res = unsafe {
            sockets::setsockopt(socket,
                                sockets::IPPROTO_IP,
                                sockets::IP_HDRINCL,
                                (&hincl as *const libc::c_int) as sockets::Buf,
                                mem::size_of::<libc::c_int>() as sockets::SockLen)
        };
        if res == -1 {
            let err = Error::last_os_error();
            unsafe {
                sockets::close(socket);
            }
            return Err(err);
        }
    }

    let sock = Arc::new(internal::FileDesc { fd: socket });
    let sender = TransportSender {
        socket: sock.clone(),
        _channel_type: channel_type,
    };
    let receiver = TransportReceiver {
        socket: sock,
        buffer: repeat(0u8).take(buffer_size).collect(),
        channel_type: channel_type,
    };

    Ok((sender, receiver))
}

impl TransportSender {
    fn send<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        let mut caddr = unsafe { mem::zeroed() };
        let sockaddr = match dst {
            IpAddr::V4(ip_addr) => net::SocketAddr::V4(net::SocketAddrV4::new(ip_addr, 0)),
            IpAddr::V6(ip_addr) => net::SocketAddr::V6(net::SocketAddrV6::new(ip_addr, 0, 0, 0)),
        };
        let slen = internal::addr_to_sockaddr(sockaddr, &mut caddr);
        let caddr_ptr = (&caddr as *const sockets::SockAddrStorage) as *const sockets::SockAddr;

        internal::send_to(self.socket.fd, packet.packet(), caddr_ptr, slen)
    }

    /// Send a packet to the provided desination
    #[inline]
    pub fn send_to<T: Packet>(&mut self, packet: T, destination: IpAddr) -> io::Result<usize> {
        self.send_to_impl(packet, destination)
    }

    #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        self.send(packet, dst)
    }

    #[cfg(any(target_os = "freebsd", target_os = "macos"))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        use packet::MutablePacket;
        use packet::ipv4::MutableIpv4Packet;

        // FreeBSD and OS X expect total length and fragment offset fields of IPv4
        // packets to be in host byte order rather than network byte order. Fragment offset is the
        // ip_off field in the ip struct and contains both the offset and the three flag bits.
        // See `man 4 ip`/Raw IP Sockets)
        if let Layer3(_) = self._channel_type {
            let mut mut_slice: Vec<u8> = repeat(0u8).take(packet.packet().len()).collect();

            let mut new_packet = MutableIpv4Packet::new(&mut mut_slice[..]).unwrap();
            new_packet.clone_from(&packet);
            let length = new_packet.get_total_length().to_be();
            new_packet.set_total_length(length);
            {
                // Turn fragment offset into host order
                let d = new_packet.packet_mut();
                let host_order = u16::from_be((d[6] as u16) << 8 | d[7] as u16);
                d[6] = (host_order >> 8) as u8;
                d[7] = host_order as u8;
            }
            return self.send(new_packet, dst);
        }

        self.send(packet, dst)
    }
}

/// Create an iterator for some packet type.
///
/// Usage:
/// ```
/// transport_channel_iterator!(Ipv4Packet, // Type to iterate over
///                             Ipv4TransportChannelIterator, // Name for iterator struct
///                             ipv4_packet_iter) // Name of function to create iterator
/// ```
#[macro_export]
macro_rules! transport_channel_iterator {
    ($ty:ident, $iter:ident, $func:ident) => (
        /// An iterator over packets of type $ty
        pub struct $iter<'a> {
            tr: &'a mut TransportReceiver
        }
        /// Return a packet iterator with packets of type $ty for some transport receiver
        pub fn $func(tr: &mut TransportReceiver) -> $iter {
            $iter {
                tr: tr
            }
        }
        impl<'a> $iter<'a> {
            /// Get the next ($ty, IpAddr) pair for the given channel
            pub fn next(&mut self) -> io::Result<($ty, IpAddr)> {
                let mut caddr: sockets::SockAddrStorage = unsafe { mem::zeroed() };
                let res = internal::recv_from(self.tr.socket.fd,
                                              &mut self.tr.buffer[..],
                                              &mut caddr);

                let offset = match self.tr.channel_type {
                    Layer4(Ipv4(_)) => {
                        let ip_header = Ipv4Packet::new(&self.tr.buffer[..]).unwrap();

                        ip_header.get_header_length() as usize * 4usize
                    },
                    Layer3(_) => {
                        fixup_packet(&mut self.tr.buffer[..]);

                        0
                    },
                    _ => 0
                };
                return match res {
                    Ok(len) => {
                        let packet = $ty::new(&self.tr.buffer[offset..len]).unwrap();
                        let addr = internal::sockaddr_to_addr(
                                        &caddr,
                                        mem::size_of::<sockets::SockAddrStorage>()
                                   );
                        let ip = match addr.unwrap() {
                            net::SocketAddr::V4(sa) => IpAddr::V4(*sa.ip()),
                            net::SocketAddr::V6(sa) => IpAddr::V6(*sa.ip()),
                        };
                        Ok((packet, ip))
                    },
                    Err(e) => Err(e),
                };

                #[cfg(any(target_os = "freebsd", target_os = "macos"))]
                fn fixup_packet(buffer: &mut [u8]) {
                    use packet::ipv4::MutableIpv4Packet;

                    let buflen = buffer.len();
                    let mut new_packet = MutableIpv4Packet::new(buffer).unwrap();

                    let length = u16::from_be(new_packet.get_total_length());
                    new_packet.set_total_length(length);

                    // OS X does this awesome thing where it removes the header length
                    // from the total length sometimes.
                    let length = new_packet.get_total_length() as usize +
                                 (new_packet.get_header_length() as usize * 4usize);
                    if length == buflen {
                        new_packet.set_total_length(length as u16)
                    }

                    let offset = u16::from_be(new_packet.get_fragment_offset());
                    new_packet.set_fragment_offset(offset);
                }

                #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
                fn fixup_packet(_buffer: &mut [u8]) {}
            }
        }
    )
}

transport_channel_iterator!(Ipv4Packet,
                            Ipv4TransportChannelIterator,
                            ipv4_packet_iter);

transport_channel_iterator!(UdpPacket,
                            UdpTransportChannelIterator,
                            udp_packet_iter);
<commit_msg>fix minor typos in src/transport.rs<commit_after>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Support for implementing transport layer protocols
//!
//! The transport module provides the ability to send and receive packets at
//! the transport layer using IPv4 or IPv6. It also enables layer 3 networking
//! for specific transport protocols, using IPv4 only.
//!
//! Note that this is limited by operating system support - for example, on OS
//! X and FreeBSD, it is impossible to implement protocols which are already
//! implemented in the kernel such as TCP and UDP.

#![macro_use]

extern crate libc;

use self::TransportProtocol::{Ipv4, Ipv6};
use self::TransportChannelType::{Layer3, Layer4};

use std::io;
use std::io::Error;
use std::iter::repeat;
use std::net::{self, IpAddr};
use std::mem;
use std::sync::Arc;

use packet::Packet;
use packet::ip::IpNextHeaderProtocol;
use packet::ipv4::Ipv4Packet;
use packet::udp::UdpPacket;

use internal;
use sockets;

/// Represents a transport layer protocol
#[derive(Clone, Copy)]
pub enum TransportProtocol {
    /// Represents a transport protocol built on top of IPv4
    Ipv4(IpNextHeaderProtocol),
    /// Represents a transport protocol built on top of IPv6
    Ipv6(IpNextHeaderProtocol),
}

/// Type of transport channel to present
#[derive(Clone, Copy)]
pub enum TransportChannelType {
    /// The application will send and receive transport layer packets
    Layer4(TransportProtocol),
    /// The application will send and receive IPv4 packets, with the specified transport protocol
    Layer3(IpNextHeaderProtocol),
}

/// Structure used for sending at the transport layer. Should be created with transport_channel()
pub struct TransportSender {
    socket: Arc<internal::FileDesc>,
    _channel_type: TransportChannelType,
}

/// Structure used for receiving at the transport layer. Should be created with transport_channel()
pub struct TransportReceiver {
    socket: Arc<internal::FileDesc>,
    buffer: Vec<u8>,
    channel_type: TransportChannelType,
}

/// Create a new (TransportSender, TransportReceiver) pair
///
/// This allows for sending and receiving packets at the transport layer. The buffer size should be
/// large enough to handle the largest packet you wish to receive.
///
/// The channel type specifies what layer to send and receive packets at, and the transport
/// protocol you wish to implement. For example, `Layer4(Ipv4(IpNextHeaderProtocols::Udp))` would
/// allow sending and receiving UDP packets using IPv4; whereas Layer3(IpNextHeaderProtocols::Udp)
/// would include the IPv4 Header in received values, and require manual construction of an IP
/// header when sending.
pub fn transport_channel(buffer_size: usize,
                         channel_type: TransportChannelType)
    -> io::Result<(TransportSender, TransportReceiver)> {
    use std::net;

    // This hack makes sure that winsock is initialised
    let _ = {
        let ip = net::Ipv4Addr::new(255, 255, 255, 255);
        let sockaddr = net::SocketAddr::V4(net::SocketAddrV4::new(ip, 0));

        net::UdpSocket::bind(sockaddr)
    };

    let socket = unsafe {
        match channel_type {
            Layer4(Ipv4(IpNextHeaderProtocol(proto))) | Layer3(IpNextHeaderProtocol(proto)) =>
                sockets::socket(sockets::AF_INET, sockets::SOCK_RAW, proto as libc::c_int),
            Layer4(Ipv6(IpNextHeaderProtocol(proto))) =>
                sockets::socket(sockets::AF_INET6, sockets::SOCK_RAW, proto as libc::c_int),
        }
    };
    if socket == sockets::INVALID_SOCKET {
        return Err(Error::last_os_error());
    }

    if match channel_type {
        Layer3(_) | Layer4(Ipv4(_)) => true,
        _ => false,
    } {
        let hincl: libc::c_int = match channel_type {
            Layer4(..) => 0,
            _ => 1,
        };
        let res = unsafe {
            sockets::setsockopt(socket,
                                sockets::IPPROTO_IP,
                                sockets::IP_HDRINCL,
                                (&hincl as *const libc::c_int) as sockets::Buf,
                                mem::size_of::<libc::c_int>() as sockets::SockLen)
        };
        if res == -1 {
            let err = Error::last_os_error();
            unsafe {
                sockets::close(socket);
            }
            return Err(err);
        }
    }

    let sock = Arc::new(internal::FileDesc { fd: socket });
    let sender = TransportSender {
        socket: sock.clone(),
        _channel_type: channel_type,
    };
    let receiver = TransportReceiver {
        socket: sock,
        buffer: repeat(0u8).take(buffer_size).collect(),
        channel_type: channel_type,
    };

    Ok((sender, receiver))
}

impl TransportSender {
    fn send<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        let mut caddr = unsafe { mem::zeroed() };
        let sockaddr = match dst {
            IpAddr::V4(ip_addr) => net::SocketAddr::V4(net::SocketAddrV4::new(ip_addr, 0)),
            IpAddr::V6(ip_addr) => net::SocketAddr::V6(net::SocketAddrV6::new(ip_addr, 0, 0, 0)),
        };
        let slen = internal::addr_to_sockaddr(sockaddr, &mut caddr);
        let caddr_ptr = (&caddr as *const sockets::SockAddrStorage) as *const sockets::SockAddr;

        internal::send_to(self.socket.fd, packet.packet(), caddr_ptr, slen)
    }

    /// Send a packet to the provided destination
    #[inline]
    pub fn send_to<T: Packet>(&mut self, packet: T, destination: IpAddr) -> io::Result<usize> {
        self.send_to_impl(packet, destination)
    }

    #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        self.send(packet, dst)
    }

    #[cfg(any(target_os = "freebsd", target_os = "macos"))]
    fn send_to_impl<T: Packet>(&mut self, packet: T, dst: IpAddr) -> io::Result<usize> {
        use packet::MutablePacket;
        use packet::ipv4::MutableIpv4Packet;

        // FreeBSD and OS X expect total length and fragment offset fields of IPv4
        // packets to be in host byte order rather than network byte order. Fragment offset is the
        // ip_off field in the ip struct and contains both the offset and the three flag bits.
        // See `man 4 ip`/Raw IP Sockets)
        if let Layer3(_) = self._channel_type {
            let mut mut_slice: Vec<u8> = repeat(0u8).take(packet.packet().len()).collect();

            let mut new_packet = MutableIpv4Packet::new(&mut mut_slice[..]).unwrap();
            new_packet.clone_from(&packet);
            let length = new_packet.get_total_length().to_be();
            new_packet.set_total_length(length);
            {
                // Turn fragment offset into host order
                let d = new_packet.packet_mut();
                let host_order = u16::from_be((d[6] as u16) << 8 | d[7] as u16);
                d[6] = (host_order >> 8) as u8;
                d[7] = host_order as u8;
            }
            return self.send(new_packet, dst);
        }

        self.send(packet, dst)
    }
}

/// Create an iterator for some packet type.
///
/// Usage:
/// ```
/// transport_channel_iterator!(Ipv4Packet, // Type to iterate over
///                             Ipv4TransportChannelIterator, // Name for iterator struct
///                             ipv4_packet_iter) // Name of function to create iterator
/// ```
#[macro_export]
macro_rules! transport_channel_iterator {
    ($ty:ident, $iter:ident, $func:ident) => (
        /// An iterator over packets of type $ty
        pub struct $iter<'a> {
            tr: &'a mut TransportReceiver
        }
        /// Return a packet iterator with packets of type $ty for some transport receiver
        pub fn $func(tr: &mut TransportReceiver) -> $iter {
            $iter {
                tr: tr
            }
        }
        impl<'a> $iter<'a> {
            /// Get the next ($ty, IpAddr) pair for the given channel
            pub fn next(&mut self) -> io::Result<($ty, IpAddr)> {
                let mut caddr: sockets::SockAddrStorage = unsafe { mem::zeroed() };
                let res = internal::recv_from(self.tr.socket.fd,
                                              &mut self.tr.buffer[..],
                                              &mut caddr);

                let offset = match self.tr.channel_type {
                    Layer4(Ipv4(_)) => {
                        let ip_header = Ipv4Packet::new(&self.tr.buffer[..]).unwrap();

                        ip_header.get_header_length() as usize * 4usize
                    },
                    Layer3(_) => {
                        fixup_packet(&mut self.tr.buffer[..]);

                        0
                    },
                    _ => 0
                };
                return match res {
                    Ok(len) => {
                        let packet = $ty::new(&self.tr.buffer[offset..len]).unwrap();
                        let addr = internal::sockaddr_to_addr(
                                        &caddr,
                                        mem::size_of::<sockets::SockAddrStorage>()
                                   );
                        let ip = match addr.unwrap() {
                            net::SocketAddr::V4(sa) => IpAddr::V4(*sa.ip()),
                            net::SocketAddr::V6(sa) => IpAddr::V6(*sa.ip()),
                        };
                        Ok((packet, ip))
                    },
                    Err(e) => Err(e),
                };

                #[cfg(any(target_os = "freebsd", target_os = "macos"))]
                fn fixup_packet(buffer: &mut [u8]) {
                    use packet::ipv4::MutableIpv4Packet;

                    let buflen = buffer.len();
                    let mut new_packet = MutableIpv4Packet::new(buffer).unwrap();

                    let length = u16::from_be(new_packet.get_total_length());
                    new_packet.set_total_length(length);

                    // OS X does this awesome thing where it removes the header length
                    // from the total length sometimes.
                    let length = new_packet.get_total_length() as usize +
                                 (new_packet.get_header_length() as usize * 4usize);
                    if length == buflen {
                        new_packet.set_total_length(length as u16)
                    }

                    let offset = u16::from_be(new_packet.get_fragment_offset());
                    new_packet.set_fragment_offset(offset);
                }

                #[cfg(all(not(target_os = "freebsd"), not(target_os = "macos")))]
                fn fixup_packet(_buffer: &mut [u8]) {}
            }
        }
    )
}

transport_channel_iterator!(Ipv4Packet,
                            Ipv4TransportChannelIterator,
                            ipv4_packet_iter);

transport_channel_iterator!(UdpPacket,
                            UdpTransportChannelIterator,
                            udp_packet_iter);
<|endoftext|>


<commit_before>use std::collections::HashMap;

use parser::{Program, Decl, Stmt, Expr};
use types::Type;
use error::Error;

type Symtable = HashMap<String, Type>;

pub fn tc_program(p: &Program) -> Result<(), Error> {
    let mut init_symtable = try!(tc_decls(&p.decls));
    tc_stmts(&p.stmts, &mut init_symtable)
}

fn tc_decls(decls: &Vec<Decl>) -> Result<Symtable, Error> {
    let mut symtable: Symtable = HashMap::new();
    for decl in decls {
        try!(tc_decl(&decl, &mut symtable));
    }
    println!("{:?}", symtable);
    Ok(symtable)
}

fn tc_decl(decl: &Decl, symtable: &mut Symtable) -> Result<(), Error> {
    if symtable.contains_key(&decl.id) {
        return Err(Error::GenericError);
    } else {
        symtable.insert(decl.id.clone(), decl.ty);
        Ok(())
    }
}

fn tc_stmts(stmts: &Vec<Stmt>, symtable: &mut Symtable) -> Result<(), Error> {
    for stmt in stmts {
        try!(tc_stmt(&stmt, symtable));
    }
    Ok(())
}

fn tc_stmt(stmt: &Stmt, symtable: &mut Symtable) -> Result<(), Error> {
    match *stmt {
        Stmt::Assign { .. } => tc_stmt_assign(stmt, symtable),
        Stmt::Read { .. } => tc_stmt_read(stmt, symtable),
        Stmt::Print { .. } => tc_stmt_print(stmt, symtable),
        Stmt::If { .. } => tc_stmt_if(stmt, symtable),
        Stmt::While { .. } => tc_stmt_while(stmt, symtable),
    }
}

fn tc_stmt_assign(stmt: &Stmt, symtable: &mut Symtable) -> Result<(), Error> {
    match *stmt {
        Stmt::Assign { pos, ref id, ref expr } => {
            let expr_ty = try!(tc_expr(expr, symtable));
            match symtable.get(id) {
                Some(&id_ty) => {
                    match (id_ty, expr_ty) {
                        (Type::Int, Type::Int) => Ok(()),
                        (Type::Int, Type::Float) => Err(Error::GenericError),
                        (Type::Float, Type::Int) => Ok(()),
                        (Type::Float, Type::Float) => Ok(()),
                    }
                }
                None => Err(Error::GenericError)
            }
        }
        _ => { Err(Error::GenericError) }
    }
}

fn tc_stmt_read(stmt: &Stmt, symtable: &Symtable) -> Result<(), Error> {
    match *stmt {
        Stmt::Read { pos, ref id } => {
            if symtable.contains_key(id) {
                Ok(())
            } else {
                Err(Error::GenericError)
            }
        }
        _ => { Err(Error::GenericError) }
    }
}

fn tc_stmt_print(stmt: &Stmt, symtable: &mut Symtable) -> Result<(), Error> {
    match *stmt {
        Stmt::Print { pos, ref expr } => {
            try!(tc_expr(expr, symtable));
            Ok(())
        }
        _ => Err(Error::GenericError)
    }
}

fn tc_stmt_if(stmt: &Stmt, symtable: &mut Symtable) -> Result<(), Error> {
    match *stmt {
        Stmt::If { pos, ref expr, ref then_stmts, ref else_stmts } => {
            let t = try!(tc_expr(expr, symtable));
            match t {
                Type::Int => {
                    try!(tc_stmts(then_stmts, symtable));
                    try!(tc_stmts(else_stmts, symtable));
                    Ok(())
                }
                _ => { Err(Error::GenericError) }
            }
        }
        _ => { Err(Error::GenericError) }
    }
}

fn tc_stmt_while(stmt: &Stmt, symtable: &mut Symtable) -> Result<(), Error> {
    match *stmt {
        Stmt::While { pos, ref expr, ref stmts } => {
            let t = try!(tc_expr(expr, symtable));
            match t {
                Type::Int => {
                    tc_stmts(stmts, symtable)
                }
                _ => {
                    Err(Error::GenericError)
                }
            }
        }
        _ => { Err(Error::GenericError) }
    }
}

fn tc_expr(expr: &Expr, symtable: &mut Symtable) -> Result<Type, Error> {
    let ty = try!(match *expr {
        Expr::Int { .. } => Ok(Type::Int),
        Expr::Float { .. } => Ok(Type::Float),
        Expr::Id { .. } => tc_expr_id(expr, symtable),
        Expr::Negate { .. } => tc_expr_negate(expr, symtable),
        Expr::Add { .. } |
        Expr::Sub { .. } |
        Expr::Mul { .. } |
        Expr::Div { .. } => tc_expr_binop(expr, symtable),
    });
    // TODO(vfoley): insert expr -> ty into symtable
    Ok(ty)
}

fn tc_expr_id(expr: &Expr, symtable: &mut Symtable) -> Result<Type, Error> {
    match *expr {
        Expr::Id { pos, ref id } => {
            match symtable.get(id) {
                Some(ty) => Ok(*ty),
                None => Err(Error::GenericError),
            }
        }
        _ => { Err(Error::GenericError) }
    }
}

fn tc_expr_negate(e: &Expr, symtable: &mut Symtable) -> Result<Type, Error> {
    tc_expr(e, symtable)
}

fn tc_expr_binop(expr: &Expr, symtable: &mut Symtable) -> Result<Type, Error> {
    match *expr {
        Expr::Add { pos, ref expr1, ref expr2 } |
        Expr::Sub { pos, ref expr1, ref expr2 } |
        Expr::Mul { pos, ref expr1, ref expr2 } |
        Expr::Div { pos, ref expr1, ref expr2 } => {
            let t1 = try!(tc_expr(expr1, symtable));
            let t2 = try!(tc_expr(expr2, symtable));

            match (t1, t2) {
                (Type::Int   , Type::Int)   => Ok(Type::Int),
                (Type::Int   , Type::Float) => Ok(Type::Float),
                (Type::Float , Type::Int)   => Ok(Type::Float),
                (Type::Float , Type::Float) => Ok(Type::Float),
            }
        }
        _ => Err(Error::GenericError)
    }
}
<commit_msg>Switch from match to if-let to reduce indenting<commit_after>use std::collections::HashMap;

use parser::{Program, Decl, Stmt, Expr};
use types::Type;
use error::Error;

type Symtable = HashMap<String, Type>;

pub fn tc_program(p: &Program) -> Result<(), Error> {
    let mut init_symtable = try!(tc_decls(&p.decls));
    tc_stmts(&p.stmts, &mut init_symtable)
}

fn tc_decls(decls: &Vec<Decl>) -> Result<Symtable, Error> {
    let mut symtable: Symtable = HashMap::new();
    for decl in decls {
        try!(tc_decl(&decl, &mut symtable));
    }
    println!("{:?}", symtable);
    Ok(symtable)
}

fn tc_decl(decl: &Decl, symtable: &mut Symtable) -> Result<(), Error> {
    if symtable.contains_key(&decl.id) {
        return Err(Error::GenericError);
    } else {
        symtable.insert(decl.id.clone(), decl.ty);
        Ok(())
    }
}

fn tc_stmts(stmts: &Vec<Stmt>, symtable: &mut Symtable) -> Result<(), Error> {
    for stmt in stmts {
        try!(tc_stmt(&stmt, symtable));
    }
    Ok(())
}

fn tc_stmt(stmt: &Stmt, symtable: &mut Symtable) -> Result<(), Error> {
    match *stmt {
        Stmt::Assign { .. } => tc_stmt_assign(stmt, symtable),
        Stmt::Read { .. } => tc_stmt_read(stmt, symtable),
        Stmt::Print { .. } => tc_stmt_print(stmt, symtable),
        Stmt::If { .. } => tc_stmt_if(stmt, symtable),
        Stmt::While { .. } => tc_stmt_while(stmt, symtable),
    }
}

fn tc_stmt_assign(stmt: &Stmt, symtable: &mut Symtable) -> Result<(), Error> {
    if let Stmt::Assign { pos, ref id, ref expr } = *stmt {
        let expr_ty = try!(tc_expr(expr, symtable));
        match symtable.get(id) {
            Some(&id_ty) => {
                match (id_ty, expr_ty) {
                    (Type::Int, Type::Int) => Ok(()),
                    (Type::Int, Type::Float) => Err(Error::GenericError),
                    (Type::Float, Type::Int) => Ok(()),
                    (Type::Float, Type::Float) => Ok(()),
                }
            }
            None => Err(Error::GenericError)
        }
    } else {
        Err(Error::GenericError)
    }
}

fn tc_stmt_read(stmt: &Stmt, symtable: &Symtable) -> Result<(), Error> {
    if let Stmt::Read { pos, ref id } = *stmt {
        if symtable.contains_key(id) {
            Ok(())
        } else {
            Err(Error::GenericError)
        }
    } else {
        Err(Error::GenericError)
    }
}

fn tc_stmt_print(stmt: &Stmt, symtable: &mut Symtable) -> Result<(), Error> {
    if let Stmt::Print { pos, ref expr } = *stmt {
        try!(tc_expr(expr, symtable));
        Ok(())
    } else {
        Err(Error::GenericError)
    }
}

fn tc_stmt_if(stmt: &Stmt, symtable: &mut Symtable) -> Result<(), Error> {
    if let Stmt::If { pos, ref expr, ref then_stmts, ref else_stmts } = *stmt {
        let t = try!(tc_expr(expr, symtable));
        match t {
            Type::Int => {
                try!(tc_stmts(then_stmts, symtable));
                try!(tc_stmts(else_stmts, symtable));
                Ok(())
            }
            _ => { Err(Error::GenericError) }
        }
    } else {
        Err(Error::GenericError)
    }
}

fn tc_stmt_while(stmt: &Stmt, symtable: &mut Symtable) -> Result<(), Error> {
    if let Stmt::While { pos, ref expr, ref stmts } = *stmt {
        let t = try!(tc_expr(expr, symtable));
        match t {
            Type::Int => {
                tc_stmts(stmts, symtable)
            }
            _ => {
                Err(Error::GenericError)
            }
        }
    } else {
        Err(Error::GenericError)
    }
}

fn tc_expr(expr: &Expr, symtable: &mut Symtable) -> Result<Type, Error> {
    let ty = try!(match *expr {
        Expr::Int { .. } => Ok(Type::Int),
        Expr::Float { .. } => Ok(Type::Float),
        Expr::Id { .. } => tc_expr_id(expr, symtable),
        Expr::Negate { .. } => tc_expr_negate(expr, symtable),
        Expr::Add { .. } |
        Expr::Sub { .. } |
        Expr::Mul { .. } |
        Expr::Div { .. } => tc_expr_binop(expr, symtable),
    });
    // TODO(vfoley): insert expr -> ty into symtable
    Ok(ty)
}

fn tc_expr_id(expr: &Expr, symtable: &mut Symtable) -> Result<Type, Error> {
    if let Expr::Id { pos, ref id } = *expr {
        match symtable.get(id) {
            Some(ty) => Ok(*ty),
            None => Err(Error::GenericError),
        }
    } else {
        Err(Error::GenericError)
    }
}

fn tc_expr_negate(e: &Expr, symtable: &mut Symtable) -> Result<Type, Error> {
    tc_expr(e, symtable)
}

fn tc_expr_binop(expr: &Expr, symtable: &mut Symtable) -> Result<Type, Error> {
    match *expr {
        Expr::Add { pos, ref expr1, ref expr2 } |
        Expr::Sub { pos, ref expr1, ref expr2 } |
        Expr::Mul { pos, ref expr1, ref expr2 } |
        Expr::Div { pos, ref expr1, ref expr2 } => {
            let t1 = try!(tc_expr(expr1, symtable));
            let t2 = try!(tc_expr(expr2, symtable));

            match (t1, t2) {
                (Type::Int   , Type::Int)   => Ok(Type::Int),
                (Type::Int   , Type::Float) => Ok(Type::Float),
                (Type::Float , Type::Int)   => Ok(Type::Float),
                (Type::Float , Type::Float) => Ok(Type::Float),
            }
        }
        _ => Err(Error::GenericError)
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::path::PathBuf;

use failure::Error;
use git2;
use regex::Regex;

use api::{Ref, Repo};

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum VcsType {
    Git,
    Unknown,
}

#[derive(Copy, Clone)]
pub enum GitReference<'a> {
    Commit(git2::Oid),
    Symbolic(&'a str),
}

impl<'a> fmt::Display for GitReference<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            GitReference::Commit(ref c) => write!(f, "{}", c),
            GitReference::Symbolic(ref s) => write!(f, "{}", s),
        }
    }
}

#[derive(Debug)]
pub struct CommitSpec {
    pub repo: String,
    pub path: Option<PathBuf>,
    pub rev: String,
    pub prev_rev: Option<String>,
}

impl fmt::Display for CommitSpec {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}@{}", &self.repo, &self.rev)
    }
}

#[derive(Debug, PartialEq, Eq)]
enum VcsProvider {
    Generic,
    Git,
    GitHub,
    Bitbucket,
    Vsts,
}

impl fmt::Display for VcsProvider {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            VcsProvider::Generic => write!(f, "generic"),
            VcsProvider::Git => write!(f, "git"),
            VcsProvider::GitHub => write!(f, "github"),
            VcsProvider::Bitbucket => write!(f, "bitbucket"),
            VcsProvider::Vsts => write!(f, "vsts"),
        }
    }
}

#[derive(Debug, PartialEq, Eq)]
struct VcsUrl {
    pub provider: VcsProvider,
    pub id: String,
    pub ty: VcsType,
}

fn parse_rev_range(rng: &str) -> (Option<String>, String) {
    if rng == "" {
        return (None, "HEAD".into());
    }
    let mut iter = rng.rsplitn(2, "..");
    let rev = iter.next().unwrap_or("HEAD");
    (iter.next().map(|x| x.to_string()), rev.to_string())
}

impl CommitSpec {
    pub fn parse(s: &str) -> Result<CommitSpec, Error> {
        lazy_static! {
            static ref SPEC_RE: Regex = Regex::new(r"^([^@#]+)(?:#([^@]+))?(?:@(.+))?$").unwrap();
        }
        if let Some(caps) = SPEC_RE.captures(s) {
            let (prev_rev, rev) = parse_rev_range(caps.get(3).map(|x| x.as_str()).unwrap_or(""));
            Ok(CommitSpec {
                repo: caps[1].to_string(),
                path: caps.get(2).map(|x| PathBuf::from(x.as_str())),
                rev,
                prev_rev,
            })
        } else {
            bail!("Could not parse commit spec '{}'", s)
        }
    }

    pub fn reference(&self) -> GitReference {
        if let Ok(oid) = git2::Oid::from_str(&self.rev) {
            GitReference::Commit(oid)
        } else {
            GitReference::Symbolic(&self.rev)
        }
    }

    pub fn prev_reference(&self) -> Option<GitReference> {
        self.prev_rev.as_ref().map(|rev| {
            if let Ok(oid) = git2::Oid::from_str(rev) {
                GitReference::Commit(oid)
            } else {
                GitReference::Symbolic(rev)
            }
        })
    }
}

fn strip_git_suffix(s: &str) -> &str {
    s.trim_right_matches(".git")
}

impl VcsUrl {
    pub fn parse(url: &str) -> VcsUrl {
        lazy_static! {
            static ref GIT_URL_RE: Regex =
                Regex::new(r"^(?:ssh|https?)://(?:[^@]+@)?([^/]+)/(.+)$").unwrap();
            static ref GIT_SSH_RE: Regex = Regex::new(r"^(?:[^@]+@)?([^/]+):(.+)$").unwrap();
        }

        if let Some(caps) = GIT_URL_RE.captures(url) {
            if let Some(rv) = VcsUrl::from_git_parts(&caps[1], &caps[2]) {
                return rv;
            } else {
                return VcsUrl {
                    provider: VcsProvider::Generic,
                    id: format!("{}/{}", &caps[1], strip_git_suffix(&caps[2])),
                    ty: VcsType::Unknown,
                };
            }
        }

        if let Some(caps) = GIT_SSH_RE.captures(url) {
            if let Some(rv) = VcsUrl::from_git_parts(&caps[1], &caps[2]) {
                return rv;
            } else {
                return VcsUrl {
                    provider: VcsProvider::Git,
                    id: url.into(),
                    ty: VcsType::Git,
                };
            }
        }

        VcsUrl {
            provider: VcsProvider::Generic,
            id: url.into(),
            ty: VcsType::Unknown,
        }
    }

    fn from_git_parts(host: &str, path: &str) -> Option<VcsUrl> {
        lazy_static! {
            static ref VS_DOMAIN_RE: Regex = Regex::new(r"^([^.]+)\.visualstudio.com$").unwrap();
            static ref VS_GIT_PATH_RE: Regex = Regex::new(r"^_git/(.+?)(?:\.git)?$").unwrap();
        }
        if host == "github.com" {
            return Some(VcsUrl {
                provider: VcsProvider::GitHub,
                id: strip_git_suffix(path).into(),
                ty: VcsType::Git,
            });
        } else if host == "bitbucket.org" {
            return Some(VcsUrl {
                provider: VcsProvider::Bitbucket,
                id: strip_git_suffix(path).into(),
                ty: VcsType::Git,
            });
        } else if let Some(caps) = VS_DOMAIN_RE.captures(host) {
            let username = &caps[1];
            if let Some(caps) = VS_GIT_PATH_RE.captures(path) {
                return Some(VcsUrl {
                    provider: VcsProvider::Vsts,
                    id: format!("{}/{}", username, &caps[1]),
                    ty: VcsType::Git,
                });
            }
        }
        None
    }
}

fn is_matching_url(a: &str, b: &str) -> bool {
    VcsUrl::parse(a) == VcsUrl::parse(b)
}

fn find_reference_url(repo: &str, repos: &[Repo]) -> Result<String, Error> {
    let mut found_non_git = false;
    for configured_repo in repos {
        if configured_repo.name != repo {
            continue;
        }

        match configured_repo.provider.id.as_str() {
            "git"
            | "github"
            | "bitbucket"
            | "visualstudio"
            | "integrations:github"
            | "integrations:github_enterprise"
            | "integrations:bitbucket"
            | "integrations:vsts" => {
                if let Some(ref url) = configured_repo.url {
                    debug!("  Got reference URL for repo {}: {}", repo, url);
                    return Ok(url.clone());
                }
            }
            _ => {
                debug!("  unknown repository {} skipped", configured_repo);
                found_non_git = true;
            }
        }
    }

    if found_non_git {
        bail!("For non git repositories explicit revisions are required");
    } else {
        bail!("Could not find matching repository for {}", repo);
    }
}

fn find_matching_rev(
    reference: GitReference,
    spec: &CommitSpec,
    repos: &[Repo],
    disable_discovery: bool,
) -> Result<Option<String>, Error> {
    macro_rules! log_match {
        ($ex:expr) => {{
            let val = $ex;
            info!("  -> found matching revision {}", val);
            val
        }};
    }

    info!("Resolving {} ({})", &reference, spec);

    let r = match reference {
        GitReference::Commit(commit) => {
            return Ok(Some(log_match!(commit.to_string())));
        }
        GitReference::Symbolic(r) => r,
    };

    let (repo, discovery) = if let Some(ref path) = spec.path {
        (git2::Repository::open(path)?, false)
    } else {
        (git2::Repository::open_from_env()?, !disable_discovery)
    };

    let reference_url = find_reference_url(&spec.repo, repos)?;
    debug!("  Looking for reference URL {}", &reference_url);

    // direct reference in root repository found.  If we are in discovery
    // mode we want to also check for matching URLs.
    if_chain! {
        if let Ok(remote) = repo.find_remote("origin");
        if let Some(url) = remote.url();
        then {
            if !discovery || is_matching_url(url, &reference_url) {
                debug!("  found match: {} == {}", url, &reference_url);
                let head = repo.revparse_single(r)?;
                return Ok(Some(log_match!(head.id().to_string())));
            } else {
                debug!("  not a match: {} != {}", url, &reference_url);
            }
        }
    }

    // in discovery mode we want to find that repo in associated submodules.
    for submodule in repo.submodules()? {
        if let Some(submodule_url) = submodule.url() {
            debug!("  found submodule with URL {}", submodule_url);
            if is_matching_url(submodule_url, &reference_url) {
                debug!(
                    "  found submodule match: {} == {}",
                    submodule_url, &reference_url
                );
                // heads on submodules is easier so let's start with that
                // because that does not require the submodule to be
                // checked out.
                if r == "HEAD" {
                    if let Some(head_oid) = submodule.head_id() {
                        return Ok(Some(log_match!(head_oid.to_string())));
                    }
                }

                // otherwise we need to open the submodule which requires
                // it to be checked out.
                if let Ok(subrepo) = submodule.open() {
                    let head = subrepo.revparse_single(r)?;
                    return Ok(Some(log_match!(head.id().to_string())));
                }
            } else {
                debug!(
                    "  not a submodule match: {} != {}",
                    submodule_url, &reference_url
                );
            }
        }
    }

    info!("  -> no matching revision found");
    Ok(None)
}

fn find_matching_revs(
    spec: &CommitSpec,
    repos: &[Repo],
    disable_discovery: bool,
) -> Result<(Option<String>, String), Error> {
    fn error(r: GitReference, repo: &str) -> Error {
        format_err!(
            "Could not find commit '{}' for '{}'. If you do not have local \
             checkouts of the repositories in question referencing tags or \
             other references will not work and you need to refer to \
             revisions explicitly.",
            r,
            repo
        )
    }

    let rev = if let Some(rev) =
        find_matching_rev(spec.reference(), &spec, &repos[..], disable_discovery)?
    {
        rev
    } else {
        return Err(error(spec.reference(), &spec.repo));
    };

    let prev_rev = if let Some(rev) = spec.prev_reference() {
        if let Some(rv) = find_matching_rev(rev, &spec, &repos[..], disable_discovery)? {
            Some(rv)
        } else {
            return Err(error(rev, &spec.repo));
        }
    } else {
        None
    };

    Ok((prev_rev, rev))
}

pub fn find_head() -> Result<String, Error> {
    let repo = git2::Repository::open_from_env()?;
    let head = repo.revparse_single("HEAD")?;
    Ok(head.id().to_string())
}

/// Given commit specs and repos this returns a list of head commits
/// from it.
pub fn find_heads(specs: Option<Vec<CommitSpec>>, repos: &[Repo]) -> Result<Vec<Ref>, Error> {
    let mut rv = vec![];

    // if commit specs were explicitly provided find head commits with
    // limited amounts of magic.
    if let Some(specs) = specs {
        for spec in &specs {
            let (prev_rev, rev) = find_matching_revs(&spec, &repos[..], specs.len() == 1)?;
            rv.push(Ref {
                repo: spec.repo.clone(),
                rev,
                prev_rev,
            });
        }

    // otherwise apply all the magic available
    } else {
        for repo in repos {
            let spec = CommitSpec {
                repo: repo.name.to_string(),
                path: None,
                rev: "HEAD".into(),
                prev_rev: None,
            };
            if let Some(rev) = find_matching_rev(spec.reference(), &spec, &repos[..], false)? {
                rv.push(Ref {
                    repo: repo.name.to_string(),
                    rev,
                    prev_rev: None,
                });
            }
        }
    }

    Ok(rv)
}

#[test]
fn test_url_parsing() {
    assert_eq!(
        VcsUrl::parse("http://github.com/mitsuhiko/flask"),
        VcsUrl {
            provider: VcsProvider::GitHub,
            id: "mitsuhiko/flask".into(),
            ty: VcsType::Git,
        }
    );
    assert_eq!(
        VcsUrl::parse("git@github.com:mitsuhiko/flask.git"),
        VcsUrl {
            provider: VcsProvider::GitHub,
            id: "mitsuhiko/flask".into(),
            ty: VcsType::Git,
        }
    );
    assert_eq!(
        VcsUrl::parse("http://bitbucket.org/mitsuhiko/flask"),
        VcsUrl {
            provider: VcsProvider::Bitbucket,
            id: "mitsuhiko/flask".into(),
            ty: VcsType::Git,
        }
    );
    assert_eq!(
        VcsUrl::parse("git@bitbucket.org:mitsuhiko/flask.git"),
        VcsUrl {
            provider: VcsProvider::Bitbucket,
            id: "mitsuhiko/flask".into(),
            ty: VcsType::Git,
        }
    );
    assert_eq!(
        VcsUrl::parse("https://neilmanvar.visualstudio.com/_git/sentry-demo"),
        VcsUrl {
            provider: VcsProvider::Vsts,
            id: "neilmanvar/sentry-demo".into(),
            ty: VcsType::Git,
        }
    );
    assert_eq!(
        VcsUrl::parse("https://github.myenterprise.com/mitsuhiko/flask.git"),
        VcsUrl {
            provider: VcsProvider::Generic,
            id: "github.myenterprise.com/mitsuhiko/flask".into(),
            ty: VcsType::Unknown,
        }
    )
}

#[test]
fn test_url_normalization() {
    assert!(is_matching_url(
        "http://github.com/mitsuhiko/flask",
        "git@github.com:mitsuhiko/flask.git"
    ));
}
<commit_msg>Add gitlab to supported VCS providers<commit_after>use std::fmt;
use std::path::PathBuf;

use failure::Error;
use git2;
use regex::Regex;

use api::{Ref, Repo};

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum VcsType {
    Git,
    Unknown,
}

#[derive(Copy, Clone)]
pub enum GitReference<'a> {
    Commit(git2::Oid),
    Symbolic(&'a str),
}

impl<'a> fmt::Display for GitReference<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            GitReference::Commit(ref c) => write!(f, "{}", c),
            GitReference::Symbolic(ref s) => write!(f, "{}", s),
        }
    }
}

#[derive(Debug)]
pub struct CommitSpec {
    pub repo: String,
    pub path: Option<PathBuf>,
    pub rev: String,
    pub prev_rev: Option<String>,
}

impl fmt::Display for CommitSpec {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}@{}", &self.repo, &self.rev)
    }
}

#[derive(Debug, PartialEq, Eq)]
enum VcsProvider {
    Generic,
    Git,
    GitHub,
    GitLab,
    Bitbucket,
    Vsts,
}

impl fmt::Display for VcsProvider {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            VcsProvider::Generic => write!(f, "generic"),
            VcsProvider::Git => write!(f, "git"),
            VcsProvider::GitHub => write!(f, "github"),
            VcsProvider::GitLab => write!(f, "gitlab"),
            VcsProvider::Bitbucket => write!(f, "bitbucket"),
            VcsProvider::Vsts => write!(f, "vsts"),
        }
    }
}

#[derive(Debug, PartialEq, Eq)]
struct VcsUrl {
    pub provider: VcsProvider,
    pub id: String,
    pub ty: VcsType,
}

fn parse_rev_range(rng: &str) -> (Option<String>, String) {
    if rng == "" {
        return (None, "HEAD".into());
    }
    let mut iter = rng.rsplitn(2, "..");
    let rev = iter.next().unwrap_or("HEAD");
    (iter.next().map(|x| x.to_string()), rev.to_string())
}

impl CommitSpec {
    pub fn parse(s: &str) -> Result<CommitSpec, Error> {
        lazy_static! {
            static ref SPEC_RE: Regex = Regex::new(r"^([^@#]+)(?:#([^@]+))?(?:@(.+))?$").unwrap();
        }
        if let Some(caps) = SPEC_RE.captures(s) {
            let (prev_rev, rev) = parse_rev_range(caps.get(3).map(|x| x.as_str()).unwrap_or(""));
            Ok(CommitSpec {
                repo: caps[1].to_string(),
                path: caps.get(2).map(|x| PathBuf::from(x.as_str())),
                rev,
                prev_rev,
            })
        } else {
            bail!("Could not parse commit spec '{}'", s)
        }
    }

    pub fn reference(&self) -> GitReference {
        if let Ok(oid) = git2::Oid::from_str(&self.rev) {
            GitReference::Commit(oid)
        } else {
            GitReference::Symbolic(&self.rev)
        }
    }

    pub fn prev_reference(&self) -> Option<GitReference> {
        self.prev_rev.as_ref().map(|rev| {
            if let Ok(oid) = git2::Oid::from_str(rev) {
                GitReference::Commit(oid)
            } else {
                GitReference::Symbolic(rev)
            }
        })
    }
}

fn strip_git_suffix(s: &str) -> &str {
    s.trim_right_matches(".git")
}

impl VcsUrl {
    pub fn parse(url: &str) -> VcsUrl {
        lazy_static! {
            static ref GIT_URL_RE: Regex =
                Regex::new(r"^(?:ssh|https?)://(?:[^@]+@)?([^/]+)/(.+)$").unwrap();
            static ref GIT_SSH_RE: Regex = Regex::new(r"^(?:[^@]+@)?([^/]+):(.+)$").unwrap();
        }

        if let Some(caps) = GIT_URL_RE.captures(url) {
            if let Some(rv) = VcsUrl::from_git_parts(&caps[1], &caps[2]) {
                return rv;
            } else {
                return VcsUrl {
                    provider: VcsProvider::Generic,
                    id: format!("{}/{}", &caps[1], strip_git_suffix(&caps[2])),
                    ty: VcsType::Unknown,
                };
            }
        }

        if let Some(caps) = GIT_SSH_RE.captures(url) {
            if let Some(rv) = VcsUrl::from_git_parts(&caps[1], &caps[2]) {
                return rv;
            } else {
                return VcsUrl {
                    provider: VcsProvider::Git,
                    id: url.into(),
                    ty: VcsType::Git,
                };
            }
        }

        VcsUrl {
            provider: VcsProvider::Generic,
            id: url.into(),
            ty: VcsType::Unknown,
        }
    }

    fn from_git_parts(host: &str, path: &str) -> Option<VcsUrl> {
        lazy_static! {
            static ref VS_DOMAIN_RE: Regex = Regex::new(r"^([^.]+)\.visualstudio.com$").unwrap();
            static ref VS_GIT_PATH_RE: Regex = Regex::new(r"^_git/(.+?)(?:\.git)?$").unwrap();
        }
        if host == "github.com" {
            return Some(VcsUrl {
                provider: VcsProvider::GitHub,
                id: strip_git_suffix(path).into(),
                ty: VcsType::Git,
            });
        } else if host == "gitlab.com" {
            return Some(VcsUrl {
                provider: VcsProvider::GitLab,
                id: strip_git_suffix(path).into(),
                ty: VcsType::Git,
            });
        } else if host == "bitbucket.org" {
            return Some(VcsUrl {
                provider: VcsProvider::Bitbucket,
                id: strip_git_suffix(path).into(),
                ty: VcsType::Git,
            });
        } else if let Some(caps) = VS_DOMAIN_RE.captures(host) {
            let username = &caps[1];
            if let Some(caps) = VS_GIT_PATH_RE.captures(path) {
                return Some(VcsUrl {
                    provider: VcsProvider::Vsts,
                    id: format!("{}/{}", username, &caps[1]),
                    ty: VcsType::Git,
                });
            }
        }
        None
    }
}

fn is_matching_url(a: &str, b: &str) -> bool {
    VcsUrl::parse(a) == VcsUrl::parse(b)
}

fn find_reference_url(repo: &str, repos: &[Repo]) -> Result<String, Error> {
    let mut found_non_git = false;
    for configured_repo in repos {
        if configured_repo.name != repo {
            continue;
        }

        match configured_repo.provider.id.as_str() {
            "git"
            | "github"
            | "bitbucket"
            | "visualstudio"
            | "integrations:github"
            | "integrations:github_enterprise"
            | "integrations:gitlab"
            | "integrations:bitbucket"
            | "integrations:vsts" => {
                if let Some(ref url) = configured_repo.url {
                    debug!("  Got reference URL for repo {}: {}", repo, url);
                    return Ok(url.clone());
                }
            }
            _ => {
                debug!("  unknown repository {} skipped", configured_repo);
                found_non_git = true;
            }
        }
    }

    if found_non_git {
        bail!("For non git repositories explicit revisions are required");
    } else {
        bail!("Could not find matching repository for {}", repo);
    }
}

fn find_matching_rev(
    reference: GitReference,
    spec: &CommitSpec,
    repos: &[Repo],
    disable_discovery: bool,
) -> Result<Option<String>, Error> {
    macro_rules! log_match {
        ($ex:expr) => {{
            let val = $ex;
            info!("  -> found matching revision {}", val);
            val
        }};
    }

    info!("Resolving {} ({})", &reference, spec);

    let r = match reference {
        GitReference::Commit(commit) => {
            return Ok(Some(log_match!(commit.to_string())));
        }
        GitReference::Symbolic(r) => r,
    };

    let (repo, discovery) = if let Some(ref path) = spec.path {
        (git2::Repository::open(path)?, false)
    } else {
        (git2::Repository::open_from_env()?, !disable_discovery)
    };

    let reference_url = find_reference_url(&spec.repo, repos)?;
    debug!("  Looking for reference URL {}", &reference_url);

    // direct reference in root repository found.  If we are in discovery
    // mode we want to also check for matching URLs.
    if_chain! {
        if let Ok(remote) = repo.find_remote("origin");
        if let Some(url) = remote.url();
        then {
            if !discovery || is_matching_url(url, &reference_url) {
                debug!("  found match: {} == {}", url, &reference_url);
                let head = repo.revparse_single(r)?;
                return Ok(Some(log_match!(head.id().to_string())));
            } else {
                debug!("  not a match: {} != {}", url, &reference_url);
            }
        }
    }

    // in discovery mode we want to find that repo in associated submodules.
    for submodule in repo.submodules()? {
        if let Some(submodule_url) = submodule.url() {
            debug!("  found submodule with URL {}", submodule_url);
            if is_matching_url(submodule_url, &reference_url) {
                debug!(
                    "  found submodule match: {} == {}",
                    submodule_url, &reference_url
                );
                // heads on submodules is easier so let's start with that
                // because that does not require the submodule to be
                // checked out.
                if r == "HEAD" {
                    if let Some(head_oid) = submodule.head_id() {
                        return Ok(Some(log_match!(head_oid.to_string())));
                    }
                }

                // otherwise we need to open the submodule which requires
                // it to be checked out.
                if let Ok(subrepo) = submodule.open() {
                    let head = subrepo.revparse_single(r)?;
                    return Ok(Some(log_match!(head.id().to_string())));
                }
            } else {
                debug!(
                    "  not a submodule match: {} != {}",
                    submodule_url, &reference_url
                );
            }
        }
    }

    info!("  -> no matching revision found");
    Ok(None)
}

fn find_matching_revs(
    spec: &CommitSpec,
    repos: &[Repo],
    disable_discovery: bool,
) -> Result<(Option<String>, String), Error> {
    fn error(r: GitReference, repo: &str) -> Error {
        format_err!(
            "Could not find commit '{}' for '{}'. If you do not have local \
             checkouts of the repositories in question referencing tags or \
             other references will not work and you need to refer to \
             revisions explicitly.",
            r,
            repo
        )
    }

    let rev = if let Some(rev) =
        find_matching_rev(spec.reference(), &spec, &repos[..], disable_discovery)?
    {
        rev
    } else {
        return Err(error(spec.reference(), &spec.repo));
    };

    let prev_rev = if let Some(rev) = spec.prev_reference() {
        if let Some(rv) = find_matching_rev(rev, &spec, &repos[..], disable_discovery)? {
            Some(rv)
        } else {
            return Err(error(rev, &spec.repo));
        }
    } else {
        None
    };

    Ok((prev_rev, rev))
}

pub fn find_head() -> Result<String, Error> {
    let repo = git2::Repository::open_from_env()?;
    let head = repo.revparse_single("HEAD")?;
    Ok(head.id().to_string())
}

/// Given commit specs and repos this returns a list of head commits
/// from it.
pub fn find_heads(specs: Option<Vec<CommitSpec>>, repos: &[Repo]) -> Result<Vec<Ref>, Error> {
    let mut rv = vec![];

    // if commit specs were explicitly provided find head commits with
    // limited amounts of magic.
    if let Some(specs) = specs {
        for spec in &specs {
            let (prev_rev, rev) = find_matching_revs(&spec, &repos[..], specs.len() == 1)?;
            rv.push(Ref {
                repo: spec.repo.clone(),
                rev,
                prev_rev,
            });
        }

    // otherwise apply all the magic available
    } else {
        for repo in repos {
            let spec = CommitSpec {
                repo: repo.name.to_string(),
                path: None,
                rev: "HEAD".into(),
                prev_rev: None,
            };
            if let Some(rev) = find_matching_rev(spec.reference(), &spec, &repos[..], false)? {
                rv.push(Ref {
                    repo: repo.name.to_string(),
                    rev,
                    prev_rev: None,
                });
            }
        }
    }

    Ok(rv)
}

#[test]
fn test_url_parsing() {
    assert_eq!(
        VcsUrl::parse("http://github.com/mitsuhiko/flask"),
        VcsUrl {
            provider: VcsProvider::GitHub,
            id: "mitsuhiko/flask".into(),
            ty: VcsType::Git,
        }
    );
    assert_eq!(
        VcsUrl::parse("git@github.com:mitsuhiko/flask.git"),
        VcsUrl {
            provider: VcsProvider::GitHub,
            id: "mitsuhiko/flask".into(),
            ty: VcsType::Git,
        }
    );
    assert_eq!(
        VcsUrl::parse("http://bitbucket.org/mitsuhiko/flask"),
        VcsUrl {
            provider: VcsProvider::Bitbucket,
            id: "mitsuhiko/flask".into(),
            ty: VcsType::Git,
        }
    );
    assert_eq!(
        VcsUrl::parse("git@bitbucket.org:mitsuhiko/flask.git"),
        VcsUrl {
            provider: VcsProvider::Bitbucket,
            id: "mitsuhiko/flask".into(),
            ty: VcsType::Git,
        }
    );
    assert_eq!(
        VcsUrl::parse("https://neilmanvar.visualstudio.com/_git/sentry-demo"),
        VcsUrl {
            provider: VcsProvider::Vsts,
            id: "neilmanvar/sentry-demo".into(),
            ty: VcsType::Git,
        }
    );
    assert_eq!(
        VcsUrl::parse("https://github.myenterprise.com/mitsuhiko/flask.git"),
        VcsUrl {
            provider: VcsProvider::Generic,
            id: "github.myenterprise.com/mitsuhiko/flask".into(),
            ty: VcsType::Unknown,
        }
    );
    assert_eq!(
        VcsUrl::parse("https://gitlab.com/gitlab-org/gitlab-ce"),
        VcsUrl {
            provider: VcsProvider::GitLab,
            id: "gitlab-org/gitlab-ce".into(),
            ty: VcsType::Git,
        }
    )
}

#[test]
fn test_url_normalization() {
    assert!(is_matching_url(
        "http://github.com/mitsuhiko/flask",
        "git@github.com:mitsuhiko/flask.git"
    ));
}
<|endoftext|>


<commit_before>// Copyright 2012-2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! write docs here

use std::comm::channel;

type WorkUnit<ARG, RV> = (ARG, Sender<RV>);

enum MessageToWorker<ARG, RV> {
    Work(WorkUnit<ARG, RV>),
    Halt,
}

enum MessageToDispatcher<ARG, RV> {
    Dispatch(WorkUnit<ARG, RV>),
    HaltAll,
    RegisterWorker(Sender<Sender<Sender<MessageToWorker<ARG, RV>>>>),
}

/// A queue that distributes work items to worker tasks.
pub struct WorkQueue<ARG, RV> {
    priv dispatcher: Sender<MessageToDispatcher<ARG, RV>>,
}

/// A proxy to a `WorkQueue`. It can be freely cloned to use from multiple tasks.
pub struct Dispatcher<ARG, RV> {
    priv inbox: Receiver<MessageToDispatcher<ARG, RV>>,
}

/// A proxy to a `WorkQueue`. It can be freely cloned to use from multiple tasks.
pub struct WorkQueueProxy<ARG, RV> {
    priv dispatcher: Sender<MessageToDispatcher<ARG, RV>>,
}

/// A worker that executes tasks from its parent queue.
pub struct Worker<ARG, RV> {
    priv ask_for_work: Sender<Sender<MessageToWorker<ARG, RV>>>,
}

/// Create a new work queue.
pub fn WorkQueue<ARG:Send, RV:Send>() -> (WorkQueue<ARG, RV>, Dispatcher<ARG, RV>) {
    let (dispatcher, dispatcher_inbox) = channel::<MessageToDispatcher<ARG, RV>>();
    return (
        WorkQueue{dispatcher: dispatcher},
        Dispatcher{inbox: dispatcher_inbox},
    );
}

impl<ARG:Send, RV:Send> WorkQueue<ARG, RV> {
    /// Create a copyable proxy that can be used to push work units.
    pub fn proxy(&self) -> WorkQueueProxy<ARG, RV> {
        return WorkQueueProxy{dispatcher: self.dispatcher.clone()};
    }

    /// Create a new worker.
    pub fn worker(&self) -> Worker<ARG, RV> {
        let (reg_s, reg_r) = channel::<Sender<Sender<MessageToWorker<ARG, RV>>>>();
        self.dispatcher.send(RegisterWorker(reg_s));
        return Worker{ask_for_work: reg_r.recv()};
    }

    /// Push a work item to this queue.
    pub fn push(&self, arg: ARG) -> Receiver<RV> {
        let (rv, wait_for_rv) = channel::<RV>();
        self.dispatcher.send(Dispatch((arg, rv)));
        return wait_for_rv;
    }
}

// rustc complais "cannot implement a destructor on a structure with
// type parameters", but our destruction is safe, we only send
// a simple message on a channel.
#[unsafe_destructor]
impl<ARG:Send, RV:Send> Drop for WorkQueue<ARG, RV> {
    fn drop(&mut self) {
        self.dispatcher.send(HaltAll);
    }
}

impl<ARG:Send, RV:Send> Dispatcher<ARG, RV> {
    /// Run the dispatcher loop. It will stop when the parent WorkQueue
    /// object is dropped.
    pub fn run(&self) {
        let (want_work, idle_worker) = channel::<Sender<MessageToWorker<ARG, RV>>>();
        let mut worker_count = 0;
        let idle_worker = idle_worker;
        loop {
            match self.inbox.recv() {
                Dispatch(work_item) => {
                    idle_worker.recv().send(Work(work_item));
                },
                RegisterWorker(want_idle_sender) => {
                    worker_count += 1;
                    want_idle_sender.send(want_work.clone());
                }
                HaltAll => {
                    while worker_count > 0 {
                        idle_worker.recv().send(Halt);
                        worker_count -= 1;
                    }
                    return;
                },
            };
        }
    }
}

impl<ARG:Send, RV:Send> WorkQueueProxy<ARG, RV> {
    /// Push a work item to this queue.
    pub fn push(&self, arg: ARG) -> Receiver<RV> {
        let (rv, wait_for_rv) = channel::<RV>();
        self.dispatcher.send(Dispatch((arg, rv)));
        return wait_for_rv;
    }
}

impl<ARG:Send, RV:Send> Clone for WorkQueueProxy<ARG, RV> {
    fn clone(&self) -> WorkQueueProxy<ARG, RV> {
        return WorkQueueProxy{dispatcher: self.dispatcher.clone()};
    }
}

impl<ARG:Send, RV:Send> Worker<ARG, RV> {
    pub fn run(&self, fun: |arg: ARG| -> RV) {
        loop {
            let (idle, work_unit) = channel::<MessageToWorker<ARG, RV>>();
            self.ask_for_work.send(idle);
            match work_unit.recv() {
                Work((arg, rv)) => rv.send((fun)(arg)),
                Halt            => return
            };
        }
    }
}

#[cfg(test)]
mod test {
    use std::task::spawn;
    use super::WorkQueue;

    #[test]
    fn test_queue() {
        let (queue, dispatcher) = WorkQueue::<int, int>();
        spawn(proc() { dispatcher.run() });
        for _ in range(0, 3) {
            let worker = queue.worker();
            spawn(proc() { worker.run(|arg| arg * 2); });
        }

        let return_list: ~[int] =
            range(0, 10)
            .map(|c| queue.push(c))
            .map(|rv| rv.recv())
            .collect();

        assert_eq!(return_list, ~[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]);
    }

    #[test]
    fn test_enqueue_from_tasks() {
        let (queue, dispatcher) = WorkQueue::<int, int>();
        spawn(proc() { dispatcher.run() });
        for _ in range(0, 3) {
            let worker = queue.worker();
            spawn(proc() { worker.run(|arg| arg * 2); });
        }
        let mut promise_list: ~[Receiver<int>] = ~[];
        let queue_proxy = queue.proxy();
        for c in range(0, 10) {
            let queue_proxy_clone = queue_proxy.clone();
            let (done, promise) = channel::<int>();
            promise_list.push(promise);
            spawn(proc() {
                let done = done;
                let queue = queue_proxy_clone;
                let rv = queue.push(c);
                done.send(rv.recv());
            });
        }

        let return_list: ~[int] =
            promise_list
            .iter()
            .map(|promise| promise.recv())
            .collect();
        assert_eq!(return_list, ~[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]);
    }
}

#[cfg(test)]
mod bench {
    extern crate test;

    use self::test::BenchHarness;
    use super::WorkQueue;

    #[bench]
    fn bench_50_tasks_4_threads(b: &mut BenchHarness) {
        let (queue, dispatcher) = WorkQueue::<int, int>();
        spawn(proc() { dispatcher.run() });
        for _ in range(0, 4) {
            let worker = queue.worker();
            spawn(proc() { worker.run(|arg| arg * 2); });
        }
        b.iter(|| {
            let _: ~[int] =
                range(0, 50)
                .map(|_| queue.push(1))
                .map(|rv| rv.recv())
                .collect();
        });
    }

    #[bench]
    fn bench_spawn_5_workers(b: &mut BenchHarness) {
        b.iter(|| {
            let (queue, dispatcher) = WorkQueue::<int, int>();
            spawn(proc() { dispatcher.run() });
            for _ in range(0, 5) {
                let worker = queue.worker();
                spawn(proc() { worker.run(|arg| arg * 2); });
            }
        });
    }
}
<commit_msg>module documentation<commit_after>// Copyright 2012-2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/*!

A simple worker queue.

# Example

```rust
use std::task;

// Create the queue
let (queue, dispatcher) = WorkQueue::<int, int>();

// Spawn a dispatcher thread
task::spawn(proc() { dispatcher.run(); });

// Spawn 10 workers
for _ in range(0, 10) {
    let worker = queue.worker();
    task::spawn(proc() { worker.run(|arg| arg + 2); });
}

// Execute jobs on the queue
let answers: ~[int] =
    range(0, 1000)
    .map(|n| queue.push(n))
    .map(|rv| rv.recv())
    .collect();
assert_eq!(answers, range(2, 1002).collect());

// The dispatcher and workers are shut down when `queue` goes out ouf scope
```

Jobs can be queued from other tasks:

```rust
let proxy = queue.proxy();
assert!(task::try(proc() {
    assert_eq!(proxy.push(5).recv(), 7);
}).is_ok());
```

*/

use std::comm::channel;

type WorkUnit<ARG, RV> = (ARG, Sender<RV>);

enum MessageToWorker<ARG, RV> {
    Work(WorkUnit<ARG, RV>),
    Halt,
}

enum MessageToDispatcher<ARG, RV> {
    Dispatch(WorkUnit<ARG, RV>),
    HaltAll,
    RegisterWorker(Sender<Sender<Sender<MessageToWorker<ARG, RV>>>>),
}

/// A queue that distributes work items to worker tasks.
pub struct WorkQueue<ARG, RV> {
    priv dispatcher: Sender<MessageToDispatcher<ARG, RV>>,
}

/// A proxy to a `WorkQueue`. It can be freely cloned to use from multiple tasks.
pub struct Dispatcher<ARG, RV> {
    priv inbox: Receiver<MessageToDispatcher<ARG, RV>>,
}

/// A proxy to a `WorkQueue`. It can be freely cloned to use from multiple tasks.
pub struct WorkQueueProxy<ARG, RV> {
    priv dispatcher: Sender<MessageToDispatcher<ARG, RV>>,
}

/// A worker that executes tasks from its parent queue.
pub struct Worker<ARG, RV> {
    priv ask_for_work: Sender<Sender<MessageToWorker<ARG, RV>>>,
}

/// Create a new work queue.
pub fn WorkQueue<ARG:Send, RV:Send>() -> (WorkQueue<ARG, RV>, Dispatcher<ARG, RV>) {
    let (dispatcher, dispatcher_inbox) = channel::<MessageToDispatcher<ARG, RV>>();
    return (
        WorkQueue{dispatcher: dispatcher},
        Dispatcher{inbox: dispatcher_inbox},
    );
}

impl<ARG:Send, RV:Send> WorkQueue<ARG, RV> {
    /// Create a copyable proxy that can be used to push work units.
    pub fn proxy(&self) -> WorkQueueProxy<ARG, RV> {
        return WorkQueueProxy{dispatcher: self.dispatcher.clone()};
    }

    /// Create a new worker.
    pub fn worker(&self) -> Worker<ARG, RV> {
        let (reg_s, reg_r) = channel::<Sender<Sender<MessageToWorker<ARG, RV>>>>();
        self.dispatcher.send(RegisterWorker(reg_s));
        return Worker{ask_for_work: reg_r.recv()};
    }

    /// Push a work item to this queue.
    pub fn push(&self, arg: ARG) -> Receiver<RV> {
        let (rv, wait_for_rv) = channel::<RV>();
        self.dispatcher.send(Dispatch((arg, rv)));
        return wait_for_rv;
    }
}

// rustc complais "cannot implement a destructor on a structure with
// type parameters", but our destruction is safe, we only send
// a simple message on a channel.
#[unsafe_destructor]
impl<ARG:Send, RV:Send> Drop for WorkQueue<ARG, RV> {
    fn drop(&mut self) {
        self.dispatcher.send(HaltAll);
    }
}

impl<ARG:Send, RV:Send> Dispatcher<ARG, RV> {
    /// Run the dispatcher loop. It will stop when the parent WorkQueue
    /// object is dropped.
    pub fn run(&self) {
        let (want_work, idle_worker) = channel::<Sender<MessageToWorker<ARG, RV>>>();
        let mut worker_count = 0;
        let idle_worker = idle_worker;
        loop {
            match self.inbox.recv() {
                Dispatch(work_item) => {
                    idle_worker.recv().send(Work(work_item));
                },
                RegisterWorker(want_idle_sender) => {
                    worker_count += 1;
                    want_idle_sender.send(want_work.clone());
                }
                HaltAll => {
                    while worker_count > 0 {
                        idle_worker.recv().send(Halt);
                        worker_count -= 1;
                    }
                    return;
                },
            };
        }
    }
}

impl<ARG:Send, RV:Send> WorkQueueProxy<ARG, RV> {
    /// Push a work item to this queue.
    pub fn push(&self, arg: ARG) -> Receiver<RV> {
        let (rv, wait_for_rv) = channel::<RV>();
        self.dispatcher.send(Dispatch((arg, rv)));
        return wait_for_rv;
    }
}

impl<ARG:Send, RV:Send> Clone for WorkQueueProxy<ARG, RV> {
    fn clone(&self) -> WorkQueueProxy<ARG, RV> {
        return WorkQueueProxy{dispatcher: self.dispatcher.clone()};
    }
}

impl<ARG:Send, RV:Send> Worker<ARG, RV> {
    pub fn run(&self, fun: |arg: ARG| -> RV) {
        loop {
            let (idle, work_unit) = channel::<MessageToWorker<ARG, RV>>();
            self.ask_for_work.send(idle);
            match work_unit.recv() {
                Work((arg, rv)) => rv.send((fun)(arg)),
                Halt            => return
            };
        }
    }
}

#[cfg(test)]
mod test {
    use std::task::spawn;
    use super::WorkQueue;

    #[test]
    fn test_queue() {
        let (queue, dispatcher) = WorkQueue::<int, int>();
        spawn(proc() { dispatcher.run() });
        for _ in range(0, 3) {
            let worker = queue.worker();
            spawn(proc() { worker.run(|arg| arg * 2); });
        }

        let return_list: ~[int] =
            range(0, 10)
            .map(|c| queue.push(c))
            .map(|rv| rv.recv())
            .collect();

        assert_eq!(return_list, ~[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]);
    }

    #[test]
    fn test_enqueue_from_tasks() {
        let (queue, dispatcher) = WorkQueue::<int, int>();
        spawn(proc() { dispatcher.run() });
        for _ in range(0, 3) {
            let worker = queue.worker();
            spawn(proc() { worker.run(|arg| arg * 2); });
        }
        let mut promise_list: ~[Receiver<int>] = ~[];
        let queue_proxy = queue.proxy();
        for c in range(0, 10) {
            let queue_proxy_clone = queue_proxy.clone();
            let (done, promise) = channel::<int>();
            promise_list.push(promise);
            spawn(proc() {
                let done = done;
                let queue = queue_proxy_clone;
                let rv = queue.push(c);
                done.send(rv.recv());
            });
        }

        let return_list: ~[int] =
            promise_list
            .iter()
            .map(|promise| promise.recv())
            .collect();
        assert_eq!(return_list, ~[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]);
    }
}

#[cfg(test)]
mod bench {
    extern crate test;

    use self::test::BenchHarness;
    use super::WorkQueue;

    #[bench]
    fn bench_50_tasks_4_threads(b: &mut BenchHarness) {
        let (queue, dispatcher) = WorkQueue::<int, int>();
        spawn(proc() { dispatcher.run() });
        for _ in range(0, 4) {
            let worker = queue.worker();
            spawn(proc() { worker.run(|arg| arg * 2); });
        }
        b.iter(|| {
            let _: ~[int] =
                range(0, 50)
                .map(|_| queue.push(1))
                .map(|rv| rv.recv())
                .collect();
        });
    }

    #[bench]
    fn bench_spawn_5_workers(b: &mut BenchHarness) {
        b.iter(|| {
            let (queue, dispatcher) = WorkQueue::<int, int>();
            spawn(proc() { dispatcher.run() });
            for _ in range(0, 5) {
                let worker = queue.worker();
                spawn(proc() { worker.run(|arg| arg * 2); });
            }
        });
    }
}
<|endoftext|>