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

use log::{error, trace};
use parking_lot::Mutex as ParkingMutex;
use shadowsocks::relay::socks5::Address;
use smoltcp::{
    iface::{Interface, InterfaceBuilder, Routes, SocketHandle},
    phy::{DeviceCapabilities, Medium},
    socket::{TcpSocket, TcpSocketBuffer, TcpState},
    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;

struct TcpSocketManager {
    iface: Interface<'static, VirtTunDevice>,
    manager_notify: Arc<Notify>,
}

impl TcpSocketManager {
    fn notify(&self) {
        self.manager_notify.notify_waiters();
    }
}

type SharedTcpSocketManager = Arc<ParkingMutex<TcpSocketManager>>;

struct TcpConnection {
    socket_handle: SocketHandle,
    manager: SharedTcpSocketManager,
}

impl Drop for TcpConnection {
    fn drop(&mut self) {
        let mut manager = self.manager.lock();
        manager.iface.remove_socket(self.socket_handle);
    }
}

impl TcpConnection {
    fn new(socket: TcpSocket<'static>, manager: SharedTcpSocketManager) -> TcpConnection {
        let socket_handle = {
            let mut manager = manager.lock();
            manager.iface.add_socket(socket)
        };

        TcpConnection { socket_handle, manager }
    }
}

impl AsyncRead for TcpConnection {
    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        let mut manager = self.manager.lock();
        {
            let socket = manager.iface.get_socket::<TcpSocket>(self.socket_handle);
            if !socket.is_open() {
                return Ok(()).into();
            }

            if socket.can_recv() {
                let recv_buf = unsafe { mem::transmute::<_, &mut [u8]>(buf.unfilled_mut()) };
                let n = match socket.recv_slice(recv_buf) {
                    Ok(n) => n,
                    Err(err) => return Err(io::Error::new(ErrorKind::Other, err)).into(),
                };
                buf.advance(n);
            } else {
                socket.register_recv_waker(cx.waker());
                return Poll::Pending;
            }
        }

        manager.notify();
        Ok(()).into()
    }
}

impl AsyncWrite for TcpConnection {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        let mut manager = self.manager.lock();
        let n = {
            let socket = manager.iface.get_socket::<TcpSocket>(self.socket_handle);
            if !socket.is_open() {
                return Err(ErrorKind::BrokenPipe.into()).into();
            }
            if socket.can_send() {
                match socket.send_slice(buf) {
                    Ok(n) => n,
                    Err(err) => return Err(io::Error::new(ErrorKind::Other, err)).into(),
                }
            } else {
                socket.register_send_waker(cx.waker());
                return Poll::Pending;
            }
        };

        manager.notify();
        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 manager = self.manager.lock();
        {
            let socket = manager.iface.get_socket::<TcpSocket>(self.socket_handle);
            // close the transmission half.
            if socket.is_open() {
                socket.close();
            }

            if socket.state() != TcpState::Closed {
                socket.register_send_waker(cx.waker());
                return Poll::Pending;
            }
        }
        manager.notify();
        Ok(()).into()
    }
}

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(),
        }));

        let manager_handle = {
            let manager = manager.clone();
            let manager_notify = manager_notify.clone();
            tokio::spawn(async move {
                loop {
                    let next_duration = {
                        let mut manager = manager.lock();

                        let before_poll = Instant::now();
                        let updated_sockets = match manager.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);
                        }

                        let next_duration = manager
                            .iface
                            .poll_delay(after_poll)
                            .unwrap_or(Duration::from_millis(50));

                        next_duration
                    };

                    tokio::task::yield_now().await;

                    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();
            // 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
            let send_buffer_size = accept_opts.tcp.send_buffer_size.unwrap_or(16384);
            let recv_buffer_size = accept_opts.tcp.recv_buffer_size.unwrap_or(87380);

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

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

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

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

            // 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_waiters();
    }

    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>remove unnecessary yield in TCP manager's poll task<commit_after>use std::{
    collections::BTreeMap,
    io::{self, ErrorKind},
    mem,
    net::{IpAddr, SocketAddr},
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
    time::Duration as StdDuration,
};

use log::{error, trace};
use parking_lot::Mutex as ParkingMutex;
use shadowsocks::relay::socks5::Address;
use smoltcp::{
    iface::{Interface, InterfaceBuilder, Routes, SocketHandle},
    phy::{DeviceCapabilities, Medium},
    socket::{TcpSocket, TcpSocketBuffer, TcpState},
    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;

struct TcpSocketManager {
    iface: Interface<'static, VirtTunDevice>,
    manager_notify: Arc<Notify>,
}

impl TcpSocketManager {
    fn notify(&self) {
        self.manager_notify.notify_waiters();
    }
}

type SharedTcpSocketManager = Arc<ParkingMutex<TcpSocketManager>>;

struct TcpConnection {
    socket_handle: SocketHandle,
    manager: SharedTcpSocketManager,
}

impl Drop for TcpConnection {
    fn drop(&mut self) {
        let mut manager = self.manager.lock();
        manager.iface.remove_socket(self.socket_handle);
    }
}

impl TcpConnection {
    fn new(socket: TcpSocket<'static>, manager: SharedTcpSocketManager) -> TcpConnection {
        let socket_handle = {
            let mut manager = manager.lock();
            manager.iface.add_socket(socket)
        };

        TcpConnection { socket_handle, manager }
    }
}

impl AsyncRead for TcpConnection {
    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        let mut manager = self.manager.lock();
        {
            let socket = manager.iface.get_socket::<TcpSocket>(self.socket_handle);
            if !socket.is_open() {
                return Ok(()).into();
            }

            if socket.can_recv() {
                let recv_buf = unsafe { mem::transmute::<_, &mut [u8]>(buf.unfilled_mut()) };
                let n = match socket.recv_slice(recv_buf) {
                    Ok(n) => n,
                    Err(err) => return Err(io::Error::new(ErrorKind::Other, err)).into(),
                };
                buf.advance(n);
            } else {
                socket.register_recv_waker(cx.waker());
                return Poll::Pending;
            }
        }

        manager.notify();
        Ok(()).into()
    }
}

impl AsyncWrite for TcpConnection {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        let mut manager = self.manager.lock();
        let n = {
            let socket = manager.iface.get_socket::<TcpSocket>(self.socket_handle);
            if !socket.is_open() {
                return Err(ErrorKind::BrokenPipe.into()).into();
            }
            if socket.can_send() {
                match socket.send_slice(buf) {
                    Ok(n) => n,
                    Err(err) => return Err(io::Error::new(ErrorKind::Other, err)).into(),
                }
            } else {
                socket.register_send_waker(cx.waker());
                return Poll::Pending;
            }
        };

        manager.notify();
        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 manager = self.manager.lock();
        {
            let socket = manager.iface.get_socket::<TcpSocket>(self.socket_handle);
            // close the transmission half.
            if socket.is_open() {
                socket.close();
            }

            if socket.state() != TcpState::Closed {
                socket.register_send_waker(cx.waker());
                return Poll::Pending;
            }
        }
        manager.notify();
        Ok(()).into()
    }
}

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(),
        }));

        let manager_handle = {
            let manager = manager.clone();
            let manager_notify = manager_notify.clone();
            tokio::spawn(async move {
                loop {
                    let next_duration = {
                        let mut manager = manager.lock();

                        let before_poll = Instant::now();
                        let updated_sockets = match manager.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);
                        }

                        let next_duration = manager
                            .iface
                            .poll_delay(after_poll)
                            .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();
            // 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
            let send_buffer_size = accept_opts.tcp.send_buffer_size.unwrap_or(16384);
            let recv_buffer_size = accept_opts.tcp.recv_buffer_size.unwrap_or(87380);

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

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

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

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

            // 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_waiters();
    }

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

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

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

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

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


<commit_before>// Copyright 2012-2015 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.

//! An implementation of SipHash 2-4.

use prelude::v1::*;

use ptr;
use super::Hasher;

/// An implementation of SipHash 2-4.
///
/// See: http://131002.net/siphash/
///
/// This is the default hashing function used by standard library (eg.
/// `collections::HashMap` uses it by default).
///
/// SipHash is a general-purpose hashing function: it runs at a good
/// speed (competitive with Spooky and City) and permits strong _keyed_
/// hashing. This lets you key your hashtables from a strong RNG, such
/// as [`rand::Rng`](https://doc.rust-lang.org/rand/rand/trait.Rng.html).
///
/// Although the SipHash algorithm is considered to be generally strong,
/// it is not intended for cryptographic purposes. As such, all
/// cryptographic uses of this implementation are _strongly discouraged_.
#[stable(feature = "rust1", since = "1.0.0")]
pub struct SipHasher {
    k0: u64,
    k1: u64,
    length: usize, // how many bytes we've processed
    // v0, v2 and v1, v3 show up in pairs in the algorithm,
    // and simd implementations of SipHash will use vectors
    // of v02 and v13. By placing them in this order in the struct,
    // the compiler can pick up on just a few simd optimizations by itself.
    v0: u64, // hash state
    v2: u64,
    v1: u64,
    v3: u64,
    tail: u64, // unprocessed bytes le
    ntail: usize, // how many bytes in tail are valid
}

// sadly, these macro definitions can't appear later,
// because they're needed in the following defs;
// this design could be improved.

macro_rules! u8to64_le {
    ($buf:expr, $i:expr) =>
    ($buf[0+$i] as u64 |
     ($buf[1+$i] as u64) << 8 |
     ($buf[2+$i] as u64) << 16 |
     ($buf[3+$i] as u64) << 24 |
     ($buf[4+$i] as u64) << 32 |
     ($buf[5+$i] as u64) << 40 |
     ($buf[6+$i] as u64) << 48 |
     ($buf[7+$i] as u64) << 56);
    ($buf:expr, $i:expr, $len:expr) =>
    ({
        let mut t = 0;
        let mut out = 0;
        while t < $len {
            out |= ($buf[t+$i] as u64) << t*8;
            t += 1;
        }
        out
    });
}

/// Load a full u64 word from a byte stream, in LE order. Use
/// `copy_nonoverlapping` to let the compiler generate the most efficient way
/// to load u64 from a possibly unaligned address.
///
/// Unsafe because: unchecked indexing at i..i+8
#[inline]
unsafe fn load_u64_le(buf: &[u8], i: usize) -> u64 {
    debug_assert!(i + 8 <= buf.len());
    let mut data = 0u64;
    ptr::copy_nonoverlapping(buf.get_unchecked(i), &mut data as *mut _ as *mut u8, 8);
    data.to_le()
}

macro_rules! rotl {
    ($x:expr, $b:expr) =>
    (($x << $b) | ($x >> (64_i32.wrapping_sub($b))))
}

macro_rules! compress {
    ($v0:expr, $v1:expr, $v2:expr, $v3:expr) =>
    ({
        $v0 = $v0.wrapping_add($v1); $v1 = rotl!($v1, 13); $v1 ^= $v0;
        $v0 = rotl!($v0, 32);
        $v2 = $v2.wrapping_add($v3); $v3 = rotl!($v3, 16); $v3 ^= $v2;
        $v0 = $v0.wrapping_add($v3); $v3 = rotl!($v3, 21); $v3 ^= $v0;
        $v2 = $v2.wrapping_add($v1); $v1 = rotl!($v1, 17); $v1 ^= $v2;
        $v2 = rotl!($v2, 32);
    })
}

impl SipHasher {
    /// Creates a new `SipHasher` with the two initial keys set to 0.
    #[inline]
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn new() -> SipHasher {
        SipHasher::new_with_keys(0, 0)
    }

    /// Creates a `SipHasher` that is keyed off the provided keys.
    #[inline]
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn new_with_keys(key0: u64, key1: u64) -> SipHasher {
        let mut state = SipHasher {
            k0: key0,
            k1: key1,
            length: 0,
            v0: 0,
            v1: 0,
            v2: 0,
            v3: 0,
            tail: 0,
            ntail: 0,
        };
        state.reset();
        state
    }

    #[inline]
    fn reset(&mut self) {
        self.length = 0;
        self.v0 = self.k0 ^ 0x736f6d6570736575;
        self.v1 = self.k1 ^ 0x646f72616e646f6d;
        self.v2 = self.k0 ^ 0x6c7967656e657261;
        self.v3 = self.k1 ^ 0x7465646279746573;
        self.ntail = 0;
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Hasher for SipHasher {
    #[inline]
    fn write(&mut self, msg: &[u8]) {
        let length = msg.len();
        self.length += length;

        let mut needed = 0;

        if self.ntail != 0 {
            needed = 8 - self.ntail;
            if length < needed {
                self.tail |= u8to64_le!(msg, 0, length) << 8 * self.ntail;
                self.ntail += length;
                return
            }

            let m = self.tail | u8to64_le!(msg, 0, needed) << 8 * self.ntail;

            self.v3 ^= m;
            compress!(self.v0, self.v1, self.v2, self.v3);
            compress!(self.v0, self.v1, self.v2, self.v3);
            self.v0 ^= m;

            self.ntail = 0;
        }

        // Buffered tail is now flushed, process new input.
        let len = length - needed;
        let left = len & 0x7;

        let mut i = needed;
        while i < len - left {
            let mi = unsafe { load_u64_le(msg, i) };

            self.v3 ^= mi;
            compress!(self.v0, self.v1, self.v2, self.v3);
            compress!(self.v0, self.v1, self.v2, self.v3);
            self.v0 ^= mi;

            i += 8;
        }

        self.tail = u8to64_le!(msg, i, left);
        self.ntail = left;
    }

    #[inline]
    fn finish(&self) -> u64 {
        let mut v0 = self.v0;
        let mut v1 = self.v1;
        let mut v2 = self.v2;
        let mut v3 = self.v3;

        let b: u64 = ((self.length as u64 & 0xff) << 56) | self.tail;

        v3 ^= b;
        compress!(v0, v1, v2, v3);
        compress!(v0, v1, v2, v3);
        v0 ^= b;

        v2 ^= 0xff;
        compress!(v0, v1, v2, v3);
        compress!(v0, v1, v2, v3);
        compress!(v0, v1, v2, v3);
        compress!(v0, v1, v2, v3);

        v0 ^ v1 ^ v2 ^ v3
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Clone for SipHasher {
    #[inline]
    fn clone(&self) -> SipHasher {
        SipHasher {
            k0: self.k0,
            k1: self.k1,
            length: self.length,
            v0: self.v0,
            v1: self.v1,
            v2: self.v2,
            v3: self.v3,
            tail: self.tail,
            ntail: self.ntail,
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Default for SipHasher {
    fn default() -> SipHasher {
        SipHasher::new()
    }
}
<commit_msg>Link to actual CSPRNG in `hash::SipHasher` documentation<commit_after>// Copyright 2012-2015 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.

//! An implementation of SipHash 2-4.

use prelude::v1::*;

use ptr;
use super::Hasher;

/// An implementation of SipHash 2-4.
///
/// See: http://131002.net/siphash/
///
/// This is currently the default hashing function used by standard library
/// (eg. `collections::HashMap` uses it by default).
///
/// SipHash is a general-purpose hashing function: it runs at a good
/// speed (competitive with Spooky and City) and permits strong _keyed_
/// hashing. This lets you key your hashtables from a strong RNG, such as
/// [`rand::os::OsRng`](https://doc.rust-lang.org/rand/rand/os/struct.OsRng.html).
///
/// Although the SipHash algorithm is considered to be generally strong,
/// it is not intended for cryptographic purposes. As such, all
/// cryptographic uses of this implementation are _strongly discouraged_.
#[stable(feature = "rust1", since = "1.0.0")]
pub struct SipHasher {
    k0: u64,
    k1: u64,
    length: usize, // how many bytes we've processed
    // v0, v2 and v1, v3 show up in pairs in the algorithm,
    // and simd implementations of SipHash will use vectors
    // of v02 and v13. By placing them in this order in the struct,
    // the compiler can pick up on just a few simd optimizations by itself.
    v0: u64, // hash state
    v2: u64,
    v1: u64,
    v3: u64,
    tail: u64, // unprocessed bytes le
    ntail: usize, // how many bytes in tail are valid
}

// sadly, these macro definitions can't appear later,
// because they're needed in the following defs;
// this design could be improved.

macro_rules! u8to64_le {
    ($buf:expr, $i:expr) =>
    ($buf[0+$i] as u64 |
     ($buf[1+$i] as u64) << 8 |
     ($buf[2+$i] as u64) << 16 |
     ($buf[3+$i] as u64) << 24 |
     ($buf[4+$i] as u64) << 32 |
     ($buf[5+$i] as u64) << 40 |
     ($buf[6+$i] as u64) << 48 |
     ($buf[7+$i] as u64) << 56);
    ($buf:expr, $i:expr, $len:expr) =>
    ({
        let mut t = 0;
        let mut out = 0;
        while t < $len {
            out |= ($buf[t+$i] as u64) << t*8;
            t += 1;
        }
        out
    });
}

/// Load a full u64 word from a byte stream, in LE order. Use
/// `copy_nonoverlapping` to let the compiler generate the most efficient way
/// to load u64 from a possibly unaligned address.
///
/// Unsafe because: unchecked indexing at i..i+8
#[inline]
unsafe fn load_u64_le(buf: &[u8], i: usize) -> u64 {
    debug_assert!(i + 8 <= buf.len());
    let mut data = 0u64;
    ptr::copy_nonoverlapping(buf.get_unchecked(i), &mut data as *mut _ as *mut u8, 8);
    data.to_le()
}

macro_rules! rotl {
    ($x:expr, $b:expr) =>
    (($x << $b) | ($x >> (64_i32.wrapping_sub($b))))
}

macro_rules! compress {
    ($v0:expr, $v1:expr, $v2:expr, $v3:expr) =>
    ({
        $v0 = $v0.wrapping_add($v1); $v1 = rotl!($v1, 13); $v1 ^= $v0;
        $v0 = rotl!($v0, 32);
        $v2 = $v2.wrapping_add($v3); $v3 = rotl!($v3, 16); $v3 ^= $v2;
        $v0 = $v0.wrapping_add($v3); $v3 = rotl!($v3, 21); $v3 ^= $v0;
        $v2 = $v2.wrapping_add($v1); $v1 = rotl!($v1, 17); $v1 ^= $v2;
        $v2 = rotl!($v2, 32);
    })
}

impl SipHasher {
    /// Creates a new `SipHasher` with the two initial keys set to 0.
    #[inline]
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn new() -> SipHasher {
        SipHasher::new_with_keys(0, 0)
    }

    /// Creates a `SipHasher` that is keyed off the provided keys.
    #[inline]
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn new_with_keys(key0: u64, key1: u64) -> SipHasher {
        let mut state = SipHasher {
            k0: key0,
            k1: key1,
            length: 0,
            v0: 0,
            v1: 0,
            v2: 0,
            v3: 0,
            tail: 0,
            ntail: 0,
        };
        state.reset();
        state
    }

    #[inline]
    fn reset(&mut self) {
        self.length = 0;
        self.v0 = self.k0 ^ 0x736f6d6570736575;
        self.v1 = self.k1 ^ 0x646f72616e646f6d;
        self.v2 = self.k0 ^ 0x6c7967656e657261;
        self.v3 = self.k1 ^ 0x7465646279746573;
        self.ntail = 0;
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Hasher for SipHasher {
    #[inline]
    fn write(&mut self, msg: &[u8]) {
        let length = msg.len();
        self.length += length;

        let mut needed = 0;

        if self.ntail != 0 {
            needed = 8 - self.ntail;
            if length < needed {
                self.tail |= u8to64_le!(msg, 0, length) << 8 * self.ntail;
                self.ntail += length;
                return
            }

            let m = self.tail | u8to64_le!(msg, 0, needed) << 8 * self.ntail;

            self.v3 ^= m;
            compress!(self.v0, self.v1, self.v2, self.v3);
            compress!(self.v0, self.v1, self.v2, self.v3);
            self.v0 ^= m;

            self.ntail = 0;
        }

        // Buffered tail is now flushed, process new input.
        let len = length - needed;
        let left = len & 0x7;

        let mut i = needed;
        while i < len - left {
            let mi = unsafe { load_u64_le(msg, i) };

            self.v3 ^= mi;
            compress!(self.v0, self.v1, self.v2, self.v3);
            compress!(self.v0, self.v1, self.v2, self.v3);
            self.v0 ^= mi;

            i += 8;
        }

        self.tail = u8to64_le!(msg, i, left);
        self.ntail = left;
    }

    #[inline]
    fn finish(&self) -> u64 {
        let mut v0 = self.v0;
        let mut v1 = self.v1;
        let mut v2 = self.v2;
        let mut v3 = self.v3;

        let b: u64 = ((self.length as u64 & 0xff) << 56) | self.tail;

        v3 ^= b;
        compress!(v0, v1, v2, v3);
        compress!(v0, v1, v2, v3);
        v0 ^= b;

        v2 ^= 0xff;
        compress!(v0, v1, v2, v3);
        compress!(v0, v1, v2, v3);
        compress!(v0, v1, v2, v3);
        compress!(v0, v1, v2, v3);

        v0 ^ v1 ^ v2 ^ v3
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Clone for SipHasher {
    #[inline]
    fn clone(&self) -> SipHasher {
        SipHasher {
            k0: self.k0,
            k1: self.k1,
            length: self.length,
            v0: self.v0,
            v1: self.v1,
            v2: self.v2,
            v3: self.v3,
            tail: self.tail,
            ntail: self.ntail,
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Default for SipHasher {
    fn default() -> SipHasher {
        SipHasher::new()
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 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.

//! This module defines the `__rust_probestack` intrinsic which is used in the
//! implementation of "stack probes" on certain platforms.
//!
//! The purpose of a stack probe is to provide a static guarantee that if a
//! thread has a guard page then a stack overflow is guaranteed to hit that
//! guard page. If a function did not have a stack probe then there's a risk of
//! having a stack frame *larger* than the guard page, so a function call could
//! skip over the guard page entirely and then later hit maybe the heap or
//! another thread, possibly leading to security vulnerabilities such as [The
//! Stack Clash], for example.
//!
//! [The Stack Clash]: https://blog.qualys.com/securitylabs/2017/06/19/the-stack-clash
//!
//! The `__rust_probestack` is called in the prologue of functions whose stack
//! size is larger than the guard page, for example larger than 4096 bytes on
//! x86. This function is then responsible for "touching" all pages relevant to
//! the stack to ensure that that if any of them are the guard page we'll hit
//! them guaranteed.
//!
//! The precise ABI for how this function operates is defined by LLVM. There's
//! no real documentation as to what this is, so you'd basically need to read
//! the LLVM source code for reference. Often though the test cases can be
//! illuminating as to the ABI that's generated, or just looking at the output
//! of `llc`.
//!
//! Note that `#[naked]` is typically used here for the stack probe because the
//! ABI corresponds to no actual ABI.
//!
//! Finally it's worth noting that at the time of this writing LLVM only has
//! support for stack probes on x86 and x86_64. There's no support for stack
//! probes on any other architecture like ARM or PowerPC64. LLVM I'm sure would
//! be more than welcome to accept such a change!

#![cfg(not(windows))] // Windows already has builtins to do this

#[naked]
#[no_mangle]
#[cfg(all(target_arch = "x86_64", not(feature = "mangled-names")))]
pub unsafe extern fn __rust_probestack() {
    // Our goal here is to touch each page between %rsp+8 and %rsp+8-%rax,
    // ensuring that if any pages are unmapped we'll make a page fault.
    //
    // The ABI here is that the stack frame size is located in `%eax`. Upon
    // return we're not supposed to modify `%esp` or `%eax`.
    asm!("
        mov    %rax,%r11        // duplicate %rax as we're clobbering %r11

        // Main loop, taken in one page increments. We're decrementing rsp by
        // a page each time until there's less than a page remaining. We're
        // guaranteed that this function isn't called unless there's more than a
        // page needed.
        //
        // Note that we're also testing against `8(%rsp)` to account for the 8
        // bytes pushed on the stack orginally with our return address. Using
        // `8(%rsp)` simulates us testing the stack pointer in the caller's
        // context.
        cmp    $$0x1000,%r11
        jna    3f
    2:
        sub    $$0x1000,%rsp
        test   %rsp,8(%rsp)
        sub    $$0x1000,%r11
        cmp    $$0x1000,%r11
        ja     2b

    3:
        // Finish up the last remaining stack space requested, getting the last
        // bits out of r11
        sub    %r11,%rsp
        test   %rsp,8(%rsp)

        // Restore the stack pointer to what it previously was when entering
        // this function. The caller will readjust the stack pointer after we
        // return.
        add    %rax,%rsp

        ret
    " ::: "memory" : "volatile");
    ::core::intrinsics::unreachable();
}

#[naked]
#[no_mangle]
#[cfg(all(target_arch = "x86", not(feature = "mangled-names")))]
pub unsafe extern fn __rust_probestack() {
    // This is the same as x86_64 above, only translated for 32-bit sizes. Note
    // that on Unix we're expected to restore everything as it was, this
    // function basically can't tamper with anything.
    //
    // The ABI here is the same as x86_64, except everything is 32-bits large.
    asm!("
        push   %ecx
        mov    %eax,%ecx

        cmp    $$0x1000,%ecx
        jna    3f
    2:
        sub    $$0x1000,%esp
        test   %esp,8(%esp)
        sub    $$0x1000,%ecx
        cmp    $$0x1000,%ecx
        ja     2b

    3:
        sub    %ecx,%esp
        test   %esp,8(%esp)

        add    %eax,%esp
        pop    %ecx
        ret
    " ::: "memory" : "volatile");
    ::core::intrinsics::unreachable();
}
<commit_msg>Explain why we need extra check in probestack.<commit_after>// Copyright 2017 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.

//! This module defines the `__rust_probestack` intrinsic which is used in the
//! implementation of "stack probes" on certain platforms.
//!
//! The purpose of a stack probe is to provide a static guarantee that if a
//! thread has a guard page then a stack overflow is guaranteed to hit that
//! guard page. If a function did not have a stack probe then there's a risk of
//! having a stack frame *larger* than the guard page, so a function call could
//! skip over the guard page entirely and then later hit maybe the heap or
//! another thread, possibly leading to security vulnerabilities such as [The
//! Stack Clash], for example.
//!
//! [The Stack Clash]: https://blog.qualys.com/securitylabs/2017/06/19/the-stack-clash
//!
//! The `__rust_probestack` is called in the prologue of functions whose stack
//! size is larger than the guard page, for example larger than 4096 bytes on
//! x86. This function is then responsible for "touching" all pages relevant to
//! the stack to ensure that that if any of them are the guard page we'll hit
//! them guaranteed.
//!
//! The precise ABI for how this function operates is defined by LLVM. There's
//! no real documentation as to what this is, so you'd basically need to read
//! the LLVM source code for reference. Often though the test cases can be
//! illuminating as to the ABI that's generated, or just looking at the output
//! of `llc`.
//!
//! Note that `#[naked]` is typically used here for the stack probe because the
//! ABI corresponds to no actual ABI.
//!
//! Finally it's worth noting that at the time of this writing LLVM only has
//! support for stack probes on x86 and x86_64. There's no support for stack
//! probes on any other architecture like ARM or PowerPC64. LLVM I'm sure would
//! be more than welcome to accept such a change!

#![cfg(not(windows))] // Windows already has builtins to do this

#[naked]
#[no_mangle]
#[cfg(all(target_arch = "x86_64", not(feature = "mangled-names")))]
pub unsafe extern fn __rust_probestack() {
    // Our goal here is to touch each page between %rsp+8 and %rsp+8-%rax,
    // ensuring that if any pages are unmapped we'll make a page fault.
    //
    // The ABI here is that the stack frame size is located in `%eax`. Upon
    // return we're not supposed to modify `%esp` or `%eax`.
    asm!("
        mov    %rax,%r11        // duplicate %rax as we're clobbering %r11

        // Main loop, taken in one page increments. We're decrementing rsp by
        // a page each time until there's less than a page remaining. We're
        // guaranteed that this function isn't called unless there's more than a
        // page needed.
        //
        // Note that we're also testing against `8(%rsp)` to account for the 8
        // bytes pushed on the stack orginally with our return address. Using
        // `8(%rsp)` simulates us testing the stack pointer in the caller's
        // context.

        // It's usually called when %rax >= 0x1000, but that's not always true.
        // Dynamic stack allocation, which is needed to implement unsized
        // rvalues, triggers stackprobe even if %rax < 0x1000.
        // Thus we have to check %r11 first to avoid segfault.
        cmp    $$0x1000,%r11
        jna    3f
    2:
        sub    $$0x1000,%rsp
        test   %rsp,8(%rsp)
        sub    $$0x1000,%r11
        cmp    $$0x1000,%r11
        ja     2b

    3:
        // Finish up the last remaining stack space requested, getting the last
        // bits out of r11
        sub    %r11,%rsp
        test   %rsp,8(%rsp)

        // Restore the stack pointer to what it previously was when entering
        // this function. The caller will readjust the stack pointer after we
        // return.
        add    %rax,%rsp

        ret
    " ::: "memory" : "volatile");
    ::core::intrinsics::unreachable();
}

#[naked]
#[no_mangle]
#[cfg(all(target_arch = "x86", not(feature = "mangled-names")))]
pub unsafe extern fn __rust_probestack() {
    // This is the same as x86_64 above, only translated for 32-bit sizes. Note
    // that on Unix we're expected to restore everything as it was, this
    // function basically can't tamper with anything.
    //
    // The ABI here is the same as x86_64, except everything is 32-bits large.
    asm!("
        push   %ecx
        mov    %eax,%ecx

        cmp    $$0x1000,%ecx
        jna    3f
    2:
        sub    $$0x1000,%esp
        test   %esp,8(%esp)
        sub    $$0x1000,%ecx
        cmp    $$0x1000,%ecx
        ja     2b

    3:
        sub    %ecx,%esp
        test   %esp,8(%esp)

        add    %eax,%esp
        pop    %ecx
        ret
    " ::: "memory" : "volatile");
    ::core::intrinsics::unreachable();
}
<|endoftext|>


<commit_before>extern crate time;

use std::fs::{File, read_dir};
use std::io::{BufReader, BufRead, Read, Write};
use std::io;
use std::path::Path;
use std::collections::HashMap;
use std::fmt;
use models::LanguageModelBuilder;
use error::Error;

pub fn find_most_common_words(corpus: &Path, num: usize) -> Vec<(String, u32)> {
    let words = files(corpus).flat_map(|file| read_words(file));

    let mut word_counts = HashMap::new();
    for word in words {
        *word_counts.entry(word).or_insert(0) += 1;
    }

    let mut counts: Vec<_> = word_counts.into_iter().collect();
    counts.sort_by(|a, b| b.1.cmp(&a.1));
    counts.truncate(num);

    counts
}

pub fn save_words(path: &Path, words: &Vec<(String, u32)>) -> io::Result<()> {
    let mut out = try!(File::create(path));
    for &(ref word, count) in words {
        try!(out.write_all(format!("{}, {}\n", &word, count).as_bytes()));
    }

    Ok(())
}

pub fn load_most_common_words(filename: &str, num: usize) -> Result<Vec<(String, u32)>, Error> {
    let file = try!(File::open(&Path::new(filename)));
    let reader = BufReader::new(file);
    let mut words = vec![];
    for line in reader.lines().take(num) {
        if let Ok(line) = line {
            let mut columns = line.split(',');
            let word = columns.next().map(|w| w.trim().to_string());
            let num = columns.next().map(|n| n.trim().parse());

            match (word, num) {
                (Some(word), Some(num)) => words.push((word, try!(num))),
                _ => return Err(Error::FormatError),
            }
        } else {
            return Err(Error::FormatError);
        }
    }

    Ok(words)
}

fn sentences<T: Read + 'static>(reader: BufReader<T>) -> Box<Iterator<Item = String>> {
    Box::new(reader.split('.' as u8)
                   .filter_map(|v| {
                       let lowercase = v.unwrap();
                       String::from_utf8(lowercase).ok().map(|s| s.to_lowercase())
                   }))
}

fn read_words<R: BufRead + 'static>(reader: R) -> Box<Iterator<Item = String>> {
    Box::new(reader.lines()
                   .filter_map(|line| line.ok())
                   .flat_map(|line| {
                       line.split(|c| {
                               match c {
                                   'a'...'z' | 'A'...'Z' => false,
                                   _ => true,
                               }
                           })
                           .filter(|word| !word.is_empty())
                           .map(|word| word.to_lowercase())
                           .collect::<Vec<_>>()
                           .into_iter()
                   }))
}

fn files(path: &Path) -> Box<Iterator<Item = BufReader<File>>> {
    Box::new(read_dir(path)
                 .unwrap()
                 .map(|path| path.unwrap().path())
                 .filter(|path| path.to_str().unwrap().ends_with(".txt"))
                 .map(|path| {
                     let file = File::open(path).unwrap();
                     BufReader::new(file)
                 }))
}

pub fn create_model(corpus: &Path, words: Vec<String>) -> LanguageModelBuilder {
    let mut builder = LanguageModelBuilder::new(10, words);

    for sentence in files(corpus).flat_map(sentences) {
        let mut acc = builder.new_sentence();

        let words = sentence.split(|c| {
                                match c {
                                    'a'...'z' => false,
                                    _ => true,
                                }
                            })
                            .filter(|w| !w.is_empty());
        for word in words {
            acc.add_word(word);
        }
    }

    builder
}
<commit_msg>use u64 rather than u32 for number of words - in theory a corpus could have > 4 billion 'the's<commit_after>extern crate time;

use std::fs::{File, read_dir};
use std::io::{BufReader, BufRead, Read, Write};
use std::io;
use std::path::Path;
use std::collections::HashMap;
use models::LanguageModelBuilder;
use error::Error;

pub fn find_most_common_words(corpus: &Path, num: usize) -> Vec<(String, u64)> {
    let words = files(corpus).flat_map(|file| read_words(file));

    let mut word_counts = HashMap::new();
    for word in words {
        *word_counts.entry(word).or_insert(0) += 1;
    }

    let mut counts: Vec<_> = word_counts.into_iter().collect();
    counts.sort_by(|a, b| b.1.cmp(&a.1));
    counts.truncate(num);

    counts
}

pub fn save_words(path: &Path, words: &Vec<(String, u64)>) -> io::Result<()> {
    let mut out = try!(File::create(path));
    for &(ref word, count) in words {
        try!(out.write_all(format!("{}, {}\n", &word, count).as_bytes()));
    }

    Ok(())
}

pub fn load_most_common_words(filename: &str, num: usize) -> Result<Vec<(String, u64)>, Error> {
    let file = try!(File::open(&Path::new(filename)));
    let reader = BufReader::new(file);
    let mut words = vec![];
    for line in reader.lines().take(num) {
        if let Ok(line) = line {
            let mut columns = line.split(',');
            let word = columns.next().map(|w| w.trim().to_string());
            let num = columns.next().map(|n| n.trim().parse());

            match (word, num) {
                (Some(word), Some(num)) => words.push((word, try!(num))),
                _ => return Err(Error::FormatError),
            }
        } else {
            return Err(Error::FormatError);
        }
    }

    Ok(words)
}

fn sentences<T: Read + 'static>(reader: BufReader<T>) -> Box<Iterator<Item = String>> {
    Box::new(reader.split('.' as u8)
                   .filter_map(|v| {
                       let lowercase = v.unwrap();
                       String::from_utf8(lowercase).ok().map(|s| s.to_lowercase())
                   }))
}

fn read_words<R: BufRead + 'static>(reader: R) -> Box<Iterator<Item = String>> {
    Box::new(reader.lines()
                   .filter_map(|line| line.ok())
                   .flat_map(|line| {
                       line.split(|c| {
                               match c {
                                   'a'...'z' | 'A'...'Z' => false,
                                   _ => true,
                               }
                           })
                           .filter(|word| !word.is_empty())
                           .map(|word| word.to_lowercase())
                           .collect::<Vec<_>>()
                           .into_iter()
                   }))
}

fn files(path: &Path) -> Box<Iterator<Item = BufReader<File>>> {
    Box::new(read_dir(path)
                 .unwrap()
                 .map(|path| path.unwrap().path())
                 .filter(|path| path.to_str().unwrap().ends_with(".txt"))
                 .map(|path| {
                     let file = File::open(path).unwrap();
                     BufReader::new(file)
                 }))
}

pub fn create_model(corpus: &Path, words: Vec<String>) -> LanguageModelBuilder {
    let mut builder = LanguageModelBuilder::new(10, words);

    for sentence in files(corpus).flat_map(sentences) {
        let mut acc = builder.new_sentence();

        let words = sentence.split(|c| {
                                match c {
                                    'a'...'z' => false,
                                    _ => true,
                                }
                            })
                            .filter(|w| !w.is_empty());
        for word in words {
            acc.add_word(word);
        }
    }

    builder
}
<|endoftext|>


<commit_before>//! Implements a command for uploading react-native projects.
use std::env;
use std::fs;
use std::path::PathBuf;
use std::process;

use chrono::Duration;
use clap::{App, Arg, ArgMatches};
use failure::{bail, Error};
use if_chain::if_chain;
use log::info;
use serde::{Deserialize, Serialize};

use crate::api::{Api, NewRelease};
use crate::config::Config;
use crate::utils::args::ArgExt;
use crate::utils::file_search::ReleaseFileSearch;
use crate::utils::file_upload::UploadContext;
use crate::utils::fs::TempFile;
use crate::utils::sourcemaps::SourceMapProcessor;
use crate::utils::system::propagate_exit_status;
use crate::utils::xcode::{InfoPlist, MayDetach};

#[derive(Serialize, Deserialize, Default, Debug)]
struct SourceMapReport {
    bundle_path: Option<PathBuf>,
    sourcemap_path: Option<PathBuf>,
}

pub fn make_app<'a, 'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
    app.about("Upload react-native projects in a Xcode build step.")
        .org_project_args()
        // legacy parameter
        .arg(
            Arg::with_name("verbose")
                .long("verbose")
                .short("v")
                .hidden(true),
        )
        .arg(Arg::with_name("force").long("force").short("f").help(
            "Force the script to run, even in debug configuration.{n}This rarely \
             does what you want because the default build script does not actually \
             produce any information that the sentry build tool could pick up on.",
        ))
        .arg(Arg::with_name("allow_fetch").long("allow-fetch").help(
            "Enable sourcemap fetching from the packager.{n}If this is enabled \
             the react native packager needs to run and sourcemaps are downloade \
             from it if the simulator platform is detected.",
        ))
        .arg(
            Arg::with_name("fetch_from")
                .long("fetch-from")
                .value_name("URL")
                .help(
                    "Set the URL to fetch sourcemaps from.{n}\
                     The default is http://127.0.0.1:8081/, where the react-native \
                     packager runs by default.",
                ),
        )
        .arg(
            Arg::with_name("force_foreground")
                .long("force-foreground")
                .help(
                    "Wait for the process to finish.{n}\
                     By default part of the build process will when triggered from Xcode \
                     detach and continue in the background.  When an error happens, \
                     a dialog is shown.  If this parameter is passed, Xcode will wait \
                     for the process to finish before the build finishes and output \
                     will be shown in the Xcode build output.",
                ),
        )
        .arg(
            Arg::with_name("build_script")
                .value_name("BUILD_SCRIPT")
                .index(1)
                .help(
                    "Optional path to the build script.{n}\
                     This is the path to the `react-native-xcode.sh` script you want \
                     to use.  By default the bundled build script is used.",
                ),
        )
        .arg(
            Arg::with_name("dist")
                .long("dist")
                .value_name("DISTRIBUTION")
                .multiple(true)
                .number_of_values(1)
                .help("The names of the distributions to publish. Can be supplied multiple times."),
        )
        .arg(
            Arg::with_name("args")
                .value_name("ARGS")
                .multiple(true)
                .last(true)
                .help("Optional arguments to pass to the build script."),
        )
        .arg(
            Arg::with_name("wait")
                .long("wait")
                .help("Wait for the server to fully process uploaded files."),
        )
}

fn find_node() -> String {
    if let Ok(path) = env::var("NODE_BINARY") {
        if !path.is_empty() {
            return path;
        }
    }
    "node".into()
}

pub fn execute(matches: &ArgMatches<'_>) -> Result<(), Error> {
    let config = Config::current();
    let (org, project) = config.get_org_and_project(matches)?;
    let api = Api::current();
    let should_wrap = matches.is_present("force")
        || match env::var("CONFIGURATION") {
            Ok(config) => !&config.contains("Debug"),
            Err(_) => bail!("Need to run this from Xcode"),
        };
    let base = env::current_dir()?;
    let script = if let Some(path) = matches.value_of("build_script") {
        base.join(path)
    } else {
        base.join("../node_modules/react-native/scripts/react-native-xcode.sh")
    }
    .canonicalize()?;

    info!(
        "Issuing a command for Organization: {} Project: {}",
        org, project
    );

    // if we allow fetching and we detect a simulator run, then we need to switch
    // to simulator mode.
    let fetch_url;
    if_chain! {
        if matches.is_present("allow_fetch");
        if let Ok(val) = env::var("PLATFORM_NAME");
        if val.ends_with("simulator");
        then {
            let url = matches.value_of("fetch_from").unwrap_or("http://127.0.0.1:8081/");
            info!("Fetching sourcemaps from {}", url);
            fetch_url = Some(url);
        } else {
            info!("Using react-native build script at {}", base.display());
            fetch_url = None;
        }
    }

    // in case we are in debug mode we directly dispatch to the script
    // and exit out early.
    if !should_wrap && fetch_url.is_none() {
        info!("Running in debug mode, skipping script wrapping.");
        let rv = process::Command::new(&script).spawn()?.wait()?;
        propagate_exit_status(rv);
        return Ok(());
    }

    info!("Parsing Info.plist");
    let plist = match InfoPlist::discover_from_env()? {
        Some(plist) => plist,
        None => bail!("Could not find info.plist"),
    };
    info!("Parse result from Info.plist: {:?}", &plist);
    let report_file = TempFile::create()?;
    let node = find_node();
    info!("Using node interpreter '{}'", &node);

    MayDetach::wrap("React native symbol handling", |md| {
        let bundle_path;
        let sourcemap_path;
        let bundle_url;
        let sourcemap_url;
        let bundle_file;
        let sourcemap_file;

        // If we have a fetch URL we need to fetch them from there now.  In that
        // case we do indeed fetch it right from the running packager and then
        // store it in temporary files for later consumption.
        if let Some(url) = fetch_url {
            if !matches.is_present("force_foreground") {
                md.may_detach()?;
            }
            let url = url.trim_end_matches('/');
            bundle_file = TempFile::create()?;
            bundle_path = bundle_file.path().to_path_buf();
            bundle_url = "~/index.ios.bundle".to_string();
            sourcemap_file = TempFile::create()?;
            sourcemap_path = sourcemap_file.path().to_path_buf();
            sourcemap_url = "~/index.ios.map".to_string();

            // wait up to 10 seconds for the server to be up.
            if !api.wait_until_available(url, Duration::seconds(10))? {
                bail!("Error: react-native packager did not respond in time");
            }

            api.download(
                &format!("{}/index.ios.bundle?platform=ios&dev=true", url),
                &mut bundle_file.open()?,
            )?;
            api.download(
                &format!("{}/index.ios.map?platform=ios&dev=true", url),
                &mut sourcemap_file.open()?,
            )?;

        // This is the case where we need to hook into the release process to
        // collect sourcemaps when they are generated.
        //
        // this invokes via an indirection of sentry-cli our wrap_call() below.
        // What is happening behind the scenes is that we switch out NODE_BINARY
        // for ourselves which is what the react-native build script normally
        // invokes.  Because we export __SENTRY_RN_WRAP_XCODE_CALL=1, the main
        // sentry-cli script will invoke our wrap_call() function below.
        //
        // That will then attempt to figure out that a react-native bundle is
        // happening to the build script, parse out the arguments, add additional
        // arguments if needed and then report the parsed arguments to a temporary
        // JSON file we load back below.
        //
        // With that we we then have all the information we need to invoke the
        // upload process.
        } else {
            let rv = process::Command::new(&script)
                .env("NODE_BINARY", env::current_exe()?.to_str().unwrap())
                .env("SENTRY_RN_REAL_NODE_BINARY", &node)
                .env(
                    "SENTRY_RN_SOURCEMAP_REPORT",
                    report_file.path().to_str().unwrap(),
                )
                .env("__SENTRY_RN_WRAP_XCODE_CALL", "1")
                .spawn()?
                .wait()?;
            propagate_exit_status(rv);

            if !matches.is_present("force_foreground") {
                md.may_detach()?;
            }
            let mut f = fs::File::open(report_file.path())?;
            let report: SourceMapReport = serde_json::from_reader(&mut f).unwrap_or_else(|_| {
                let err_msg = format!(
                    "File {} doesn't contain a valid JSON data.",
                    report_file.path().display()
                );
                panic!("{}", err_msg);
            });
            if report.bundle_path.is_none() || report.sourcemap_path.is_none() {
                println!("Warning: build produced no sourcemaps.");
                return Ok(());
            }

            bundle_path = report.bundle_path.unwrap();
            bundle_url = format!("~/{}", bundle_path.file_name().unwrap().to_string_lossy());
            sourcemap_path = report.sourcemap_path.unwrap();
            sourcemap_url = format!(
                "~/{}",
                sourcemap_path.file_name().unwrap().to_string_lossy()
            );
        }

        // now that we have all the data, we can now process and upload the
        // sourcemaps.
        println!("Processing react-native sourcemaps for Sentry upload.");
        info!("  bundle path: {}", bundle_path.display());
        info!("  sourcemap path: {}", sourcemap_path.display());

        let mut processor = SourceMapProcessor::new();
        processor.add(&bundle_url, ReleaseFileSearch::collect_file(bundle_path)?)?;
        processor.add(
            &sourcemap_url,
            ReleaseFileSearch::collect_file(sourcemap_path)?,
        )?;
        processor.rewrite(&[base.parent().unwrap().to_str().unwrap()])?;
        processor.add_sourcemap_references()?;

        let release = api.new_release(
            &org,
            &NewRelease {
                version: format!(
                    "{}@{}+{}",
                    plist.bundle_id(),
                    plist.version(),
                    plist.build()
                ),
                projects: vec![project.to_string()],
                ..Default::default()
            },
        )?;

        match matches.values_of("dist") {
            None => {
                processor.upload(&UploadContext {
                    org: &org,
                    project: Some(&project),
                    release: &release.version,
                    dist: Some(&plist.build()),
                    wait: matches.is_present("wait"),
                })?;
            }
            Some(dists) => {
                for dist in dists {
                    processor.upload(&UploadContext {
                        org: &org,
                        project: Some(&project),
                        release: &release.version,
                        dist: Some(dist),
                        wait: matches.is_present("wait"),
                    })?;
                }
            }
        }

        Ok(())
    })
}

pub fn wrap_call() -> Result<(), Error> {
    let mut args: Vec<_> = env::args().skip(1).collect();
    let mut bundle_path = None;
    let mut sourcemap_path = None;

    if args.len() > 1 && (args[1] == "bundle" || args[1] == "ram-bundle") {
        let mut iter = args.iter().fuse();
        while let Some(item) = iter.next() {
            if item == "--sourcemap-output" {
                sourcemap_path = iter.next().cloned();
            } else if let Some(rest) = item.strip_prefix("--sourcemap-output=") {
                sourcemap_path = Some(rest.to_string());
            } else if item == "--bundle-output" {
                bundle_path = iter.next().cloned();
            } else if let Some(rest) = item.strip_prefix("--bundle-output=") {
                bundle_path = Some(rest.to_string());
            }
        }
    }

    let mut sourcemap_report = SourceMapReport::default();

    if sourcemap_path.is_none() && bundle_path.is_some() {
        let mut path = env::temp_dir();
        let mut map_path = PathBuf::from(bundle_path.clone().unwrap());
        map_path.set_extension("jsbundle.map");
        path.push(map_path.file_name().unwrap());
        sourcemap_report.sourcemap_path = Some(PathBuf::from(&path));
        args.push("--sourcemap-output".into());
        args.push(path.into_os_string().into_string().unwrap());
    } else if let Some(path) = sourcemap_path {
        sourcemap_report.sourcemap_path = Some(PathBuf::from(path));
    }

    sourcemap_report.bundle_path = bundle_path.map(PathBuf::from);

    let rv = process::Command::new(env::var("SENTRY_RN_REAL_NODE_BINARY").unwrap())
        .args(&args)
        .spawn()?
        .wait()?;
    propagate_exit_status(rv);

    let mut f = fs::File::create(env::var("SENTRY_RN_SOURCEMAP_REPORT").unwrap())?;
    serde_json::to_writer(&mut f, &sourcemap_report)?;

    Ok(())
}
<commit_msg>[Fix] Allow to define a custom release name for source-maps upload (#967)<commit_after>//! Implements a command for uploading react-native projects.
use std::env;
use std::fs;
use std::path::PathBuf;
use std::process;

use chrono::Duration;
use clap::{App, Arg, ArgMatches};
use failure::{bail, Error};
use if_chain::if_chain;
use log::info;
use serde::{Deserialize, Serialize};

use crate::api::{Api, NewRelease};
use crate::config::Config;
use crate::utils::args::ArgExt;
use crate::utils::file_search::ReleaseFileSearch;
use crate::utils::file_upload::UploadContext;
use crate::utils::fs::TempFile;
use crate::utils::sourcemaps::SourceMapProcessor;
use crate::utils::system::propagate_exit_status;
use crate::utils::xcode::{InfoPlist, MayDetach};

#[derive(Serialize, Deserialize, Default, Debug)]
struct SourceMapReport {
    bundle_path: Option<PathBuf>,
    sourcemap_path: Option<PathBuf>,
}

pub fn make_app<'a, 'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
    app.about("Upload react-native projects in a Xcode build step.")
        .org_project_args()
        // legacy parameter
        .arg(
            Arg::with_name("verbose")
                .long("verbose")
                .short("v")
                .hidden(true),
        )
        .arg(Arg::with_name("force").long("force").short("f").help(
            "Force the script to run, even in debug configuration.{n}This rarely \
             does what you want because the default build script does not actually \
             produce any information that the sentry build tool could pick up on.",
        ))
        .arg(Arg::with_name("allow_fetch").long("allow-fetch").help(
            "Enable sourcemap fetching from the packager.{n}If this is enabled \
             the react native packager needs to run and sourcemaps are downloade \
             from it if the simulator platform is detected.",
        ))
        .arg(
            Arg::with_name("fetch_from")
                .long("fetch-from")
                .value_name("URL")
                .help(
                    "Set the URL to fetch sourcemaps from.{n}\
                     The default is http://127.0.0.1:8081/, where the react-native \
                     packager runs by default.",
                ),
        )
        .arg(
            Arg::with_name("force_foreground")
                .long("force-foreground")
                .help(
                    "Wait for the process to finish.{n}\
                     By default part of the build process will when triggered from Xcode \
                     detach and continue in the background.  When an error happens, \
                     a dialog is shown.  If this parameter is passed, Xcode will wait \
                     for the process to finish before the build finishes and output \
                     will be shown in the Xcode build output.",
                ),
        )
        .arg(
            Arg::with_name("build_script")
                .value_name("BUILD_SCRIPT")
                .index(1)
                .help(
                    "Optional path to the build script.{n}\
                     This is the path to the `react-native-xcode.sh` script you want \
                     to use.  By default the bundled build script is used.",
                ),
        )
        .arg(
            Arg::with_name("dist")
                .long("dist")
                .value_name("DISTRIBUTION")
                .multiple(true)
                .number_of_values(1)
                .help("The names of the distributions to publish. Can be supplied multiple times."),
        )
        .arg(
            Arg::with_name("args")
                .value_name("ARGS")
                .multiple(true)
                .last(true)
                .help("Optional arguments to pass to the build script."),
        )
        .arg(
            Arg::with_name("wait")
                .long("wait")
                .help("Wait for the server to fully process uploaded files."),
        )
}

fn find_node() -> String {
    if let Ok(path) = env::var("NODE_BINARY") {
        if !path.is_empty() {
            return path;
        }
    }
    "node".into()
}

pub fn execute(matches: &ArgMatches<'_>) -> Result<(), Error> {
    let config = Config::current();
    let (org, project) = config.get_org_and_project(matches)?;
    let api = Api::current();
    let should_wrap = matches.is_present("force")
        || match env::var("CONFIGURATION") {
            Ok(config) => !&config.contains("Debug"),
            Err(_) => bail!("Need to run this from Xcode"),
        };
    let base = env::current_dir()?;
    let script = if let Some(path) = matches.value_of("build_script") {
        base.join(path)
    } else {
        base.join("../node_modules/react-native/scripts/react-native-xcode.sh")
    }
    .canonicalize()?;

    info!(
        "Issuing a command for Organization: {} Project: {}",
        org, project
    );

    // if we allow fetching and we detect a simulator run, then we need to switch
    // to simulator mode.
    let fetch_url;
    if_chain! {
        if matches.is_present("allow_fetch");
        if let Ok(val) = env::var("PLATFORM_NAME");
        if val.ends_with("simulator");
        then {
            let url = matches.value_of("fetch_from").unwrap_or("http://127.0.0.1:8081/");
            info!("Fetching sourcemaps from {}", url);
            fetch_url = Some(url);
        } else {
            info!("Using react-native build script at {}", base.display());
            fetch_url = None;
        }
    }

    // in case we are in debug mode we directly dispatch to the script
    // and exit out early.
    if !should_wrap && fetch_url.is_none() {
        info!("Running in debug mode, skipping script wrapping.");
        let rv = process::Command::new(&script).spawn()?.wait()?;
        propagate_exit_status(rv);
        return Ok(());
    }

    info!("Parsing Info.plist");
    let plist = match InfoPlist::discover_from_env()? {
        Some(plist) => plist,
        None => bail!("Could not find info.plist"),
    };
    info!("Parse result from Info.plist: {:?}", &plist);
    let report_file = TempFile::create()?;
    let node = find_node();
    info!("Using node interpreter '{}'", &node);

    MayDetach::wrap("React native symbol handling", |md| {
        let bundle_path;
        let sourcemap_path;
        let bundle_url;
        let sourcemap_url;
        let bundle_file;
        let sourcemap_file;

        // If we have a fetch URL we need to fetch them from there now.  In that
        // case we do indeed fetch it right from the running packager and then
        // store it in temporary files for later consumption.
        if let Some(url) = fetch_url {
            if !matches.is_present("force_foreground") {
                md.may_detach()?;
            }
            let url = url.trim_end_matches('/');
            bundle_file = TempFile::create()?;
            bundle_path = bundle_file.path().to_path_buf();
            bundle_url = "~/index.ios.bundle".to_string();
            sourcemap_file = TempFile::create()?;
            sourcemap_path = sourcemap_file.path().to_path_buf();
            sourcemap_url = "~/index.ios.map".to_string();

            // wait up to 10 seconds for the server to be up.
            if !api.wait_until_available(url, Duration::seconds(10))? {
                bail!("Error: react-native packager did not respond in time");
            }

            api.download(
                &format!("{}/index.ios.bundle?platform=ios&dev=true", url),
                &mut bundle_file.open()?,
            )?;
            api.download(
                &format!("{}/index.ios.map?platform=ios&dev=true", url),
                &mut sourcemap_file.open()?,
            )?;

        // This is the case where we need to hook into the release process to
        // collect sourcemaps when they are generated.
        //
        // this invokes via an indirection of sentry-cli our wrap_call() below.
        // What is happening behind the scenes is that we switch out NODE_BINARY
        // for ourselves which is what the react-native build script normally
        // invokes.  Because we export __SENTRY_RN_WRAP_XCODE_CALL=1, the main
        // sentry-cli script will invoke our wrap_call() function below.
        //
        // That will then attempt to figure out that a react-native bundle is
        // happening to the build script, parse out the arguments, add additional
        // arguments if needed and then report the parsed arguments to a temporary
        // JSON file we load back below.
        //
        // With that we we then have all the information we need to invoke the
        // upload process.
        } else {
            let rv = process::Command::new(&script)
                .env("NODE_BINARY", env::current_exe()?.to_str().unwrap())
                .env("SENTRY_RN_REAL_NODE_BINARY", &node)
                .env(
                    "SENTRY_RN_SOURCEMAP_REPORT",
                    report_file.path().to_str().unwrap(),
                )
                .env("__SENTRY_RN_WRAP_XCODE_CALL", "1")
                .spawn()?
                .wait()?;
            propagate_exit_status(rv);

            if !matches.is_present("force_foreground") {
                md.may_detach()?;
            }
            let mut f = fs::File::open(report_file.path())?;
            let report: SourceMapReport = serde_json::from_reader(&mut f).unwrap_or_else(|_| {
                let err_msg = format!(
                    "File {} doesn't contain a valid JSON data.",
                    report_file.path().display()
                );
                panic!("{}", err_msg);
            });
            if report.bundle_path.is_none() || report.sourcemap_path.is_none() {
                println!("Warning: build produced no sourcemaps.");
                return Ok(());
            }

            bundle_path = report.bundle_path.unwrap();
            bundle_url = format!("~/{}", bundle_path.file_name().unwrap().to_string_lossy());
            sourcemap_path = report.sourcemap_path.unwrap();
            sourcemap_url = format!(
                "~/{}",
                sourcemap_path.file_name().unwrap().to_string_lossy()
            );
        }

        // now that we have all the data, we can now process and upload the
        // sourcemaps.
        println!("Processing react-native sourcemaps for Sentry upload.");
        info!("  bundle path: {}", bundle_path.display());
        info!("  sourcemap path: {}", sourcemap_path.display());

        let mut processor = SourceMapProcessor::new();
        processor.add(&bundle_url, ReleaseFileSearch::collect_file(bundle_path)?)?;
        processor.add(
            &sourcemap_url,
            ReleaseFileSearch::collect_file(sourcemap_path)?,
        )?;
        processor.rewrite(&[base.parent().unwrap().to_str().unwrap()])?;
        processor.add_sourcemap_references()?;

        let dist = env::var("SENTRY_DIST").unwrap_or(plist.build().to_string());
        let release_name = env::var("SENTRY_RELEASE").unwrap_or(format!(
            "{}@{}+{}",
            plist.bundle_id(),
            plist.version(),
            dist
        ));

        let release = api.new_release(
            &org,
            &NewRelease {
                version: release_name,
                projects: vec![project.to_string()],
                ..Default::default()
            },
        )?;

        match matches.values_of("dist") {
            None => {
                processor.upload(&UploadContext {
                    org: &org,
                    project: Some(&project),
                    release: &release.version,
                    dist: Some(&dist),
                    wait: matches.is_present("wait"),
                })?;
            }
            Some(dists) => {
                for dist in dists {
                    processor.upload(&UploadContext {
                        org: &org,
                        project: Some(&project),
                        release: &release.version,
                        dist: Some(dist),
                        wait: matches.is_present("wait"),
                    })?;
                }
            }
        }

        Ok(())
    })
}

pub fn wrap_call() -> Result<(), Error> {
    let mut args: Vec<_> = env::args().skip(1).collect();
    let mut bundle_path = None;
    let mut sourcemap_path = None;

    if args.len() > 1 && (args[1] == "bundle" || args[1] == "ram-bundle") {
        let mut iter = args.iter().fuse();
        while let Some(item) = iter.next() {
            if item == "--sourcemap-output" {
                sourcemap_path = iter.next().cloned();
            } else if let Some(rest) = item.strip_prefix("--sourcemap-output=") {
                sourcemap_path = Some(rest.to_string());
            } else if item == "--bundle-output" {
                bundle_path = iter.next().cloned();
            } else if let Some(rest) = item.strip_prefix("--bundle-output=") {
                bundle_path = Some(rest.to_string());
            }
        }
    }

    let mut sourcemap_report = SourceMapReport::default();

    if sourcemap_path.is_none() && bundle_path.is_some() {
        let mut path = env::temp_dir();
        let mut map_path = PathBuf::from(bundle_path.clone().unwrap());
        map_path.set_extension("jsbundle.map");
        path.push(map_path.file_name().unwrap());
        sourcemap_report.sourcemap_path = Some(PathBuf::from(&path));
        args.push("--sourcemap-output".into());
        args.push(path.into_os_string().into_string().unwrap());
    } else if let Some(path) = sourcemap_path {
        sourcemap_report.sourcemap_path = Some(PathBuf::from(path));
    }

    sourcemap_report.bundle_path = bundle_path.map(PathBuf::from);

    let rv = process::Command::new(env::var("SENTRY_RN_REAL_NODE_BINARY").unwrap())
        .args(&args)
        .spawn()?
        .wait()?;
    propagate_exit_status(rv);

    let mut f = fs::File::create(env::var("SENTRY_RN_SOURCEMAP_REPORT").unwrap())?;
    serde_json::to_writer(&mut f, &sourcemap_report)?;

    Ok(())
}
<|endoftext|>


<commit_before>use std::mem;
use std::result::Result;
use math::variables::{AbstVar, new_const};
use math::relationships::Relationship;

pub struct Expression {
    left_hand_side: Vec<AbstVar>,
    relationship: Relationship,
    right_hand_side: Vec<AbstVar>,
}

impl Expression {
    pub fn new(l_h_s: Vec<AbstVar>, r: Relationship, r_h_s: Vec<AbstVar>) -> Expression {
        Expression {
            left_hand_side: l_h_s,
            relationship: r,
            right_hand_side: r_h_s,
        }
    }

    pub fn lhs(&self) -> &Vec<AbstVar> {
        &self.left_hand_side
    }

    pub fn rel(&self) -> &Relationship {
        &self.relationship
    }

    pub fn rhs(&self) -> &Vec<AbstVar> {
        &self.right_hand_side
    }

    pub fn set_rel(&mut self, new_rel: Relationship) {
        self.relationship = new_rel;
    }

    pub fn add_lhs(&mut self, to_add: AbstVar) {
        self.left_hand_side.push(to_add);
    }

    pub fn move_from_lhs_side(&mut self, index: usize, insert_at_start: bool) {
        let mut to_move = self.left_hand_side.remove(index);
        to_move.change_sign();
        if self.left_hand_side.is_empty() {
            self.left_hand_side.push(new_const("zero", 0.0));
        }
        insert_side(&mut self.right_hand_side, to_move, insert_at_start);
    }

    pub fn mul_both_sides(&mut self, by: f32) {
        mul_side(&mut self.left_hand_side, by);
        mul_side(&mut self.right_hand_side, by);

        // Change sign if required
        if by.is_sign_negative() && self.relationship != Relationship::EQ {
            if self.relationship == Relationship::GEQ {
                self.relationship = Relationship::LEQ;
            } else {
                self.relationship = Relationship::GEQ;
            }
        }
    }

    pub fn swap_sides(&mut self) -> Result<&str, &str> {
        if self.relationship == Relationship::EQ {
            mem::swap(&mut self.left_hand_side, &mut self.right_hand_side);
            Ok("Swapped!")
        } else {
            Err("Invalid state: relationship must be \"EQ\".")
        }
    }
}

fn insert_side(side: &mut Vec<AbstVar>, var: AbstVar, start: bool) {
    // Preserve order of basic followed by non basic
    if start {
        side.insert(0, var);
    } else {
        let mut insert_at = 0;
        while insert_at < side.len() {
            match &mut side[insert_at] {
                &mut AbstVar::SlackVar { .. } => break,
                _ => insert_at += 1,
            };
        }
        side.insert(insert_at, var);
    }
}

fn mul_side(side: &mut Vec<AbstVar>, by: f32) {
    for var in side.iter_mut() {
        match var {
            &mut AbstVar::Variable { ref mut coefficient, .. } => {
                if *coefficient != 0.0 {
                    *coefficient = *coefficient * by;
                }
            }
            &mut AbstVar::Constant { ref mut value, .. } => {
                if *value != 0.0 {
                    *value = *value * by;
                }
            }
            _ => panic!("Unexpected variant in this program logic."),
        };
    }
}
<commit_msg>Using less precise floats seems to remove the need for ensuring no negative zeros.<commit_after>use std::mem;
use std::result::Result;
use math::variables::{AbstVar, new_const};
use math::relationships::Relationship;

pub struct Expression {
    left_hand_side: Vec<AbstVar>,
    relationship: Relationship,
    right_hand_side: Vec<AbstVar>,
}

impl Expression {
    pub fn new(l_h_s: Vec<AbstVar>, r: Relationship, r_h_s: Vec<AbstVar>) -> Expression {
        Expression {
            left_hand_side: l_h_s,
            relationship: r,
            right_hand_side: r_h_s,
        }
    }

    pub fn lhs(&self) -> &Vec<AbstVar> {
        &self.left_hand_side
    }

    pub fn rel(&self) -> &Relationship {
        &self.relationship
    }

    pub fn rhs(&self) -> &Vec<AbstVar> {
        &self.right_hand_side
    }

    pub fn set_rel(&mut self, new_rel: Relationship) {
        self.relationship = new_rel;
    }

    pub fn add_lhs(&mut self, to_add: AbstVar) {
        self.left_hand_side.push(to_add);
    }

    pub fn move_from_lhs_side(&mut self, index: usize, insert_at_start: bool) {
        let mut to_move = self.left_hand_side.remove(index);
        to_move.change_sign();
        if self.left_hand_side.is_empty() {
            self.left_hand_side.push(new_const("zero", 0.0));
        }
        insert_side(&mut self.right_hand_side, to_move, insert_at_start);
    }

    pub fn mul_both_sides(&mut self, by: f32) {
        mul_side(&mut self.left_hand_side, by);
        mul_side(&mut self.right_hand_side, by);

        // Change sign if required
        if by.is_sign_negative() && self.relationship != Relationship::EQ {
            if self.relationship == Relationship::GEQ {
                self.relationship = Relationship::LEQ;
            } else {
                self.relationship = Relationship::GEQ;
            }
        }
    }

    pub fn swap_sides(&mut self) -> Result<&str, &str> {
        if self.relationship == Relationship::EQ {
            mem::swap(&mut self.left_hand_side, &mut self.right_hand_side);
            Ok("Swapped!")
        } else {
            Err("Invalid state: relationship must be \"EQ\".")
        }
    }
}

fn insert_side(side: &mut Vec<AbstVar>, var: AbstVar, start: bool) {
    // Preserve order of basic followed by non basic
    if start {
        side.insert(0, var);
    } else {
        let mut insert_at = 0;
        while insert_at < side.len() {
            match &mut side[insert_at] {
                &mut AbstVar::SlackVar { .. } => break,
                _ => insert_at += 1,
            };
        }
        side.insert(insert_at, var);
    }
}

fn mul_side(side: &mut Vec<AbstVar>, by: f32) {
    for var in side.iter_mut() {
        match var {
            &mut AbstVar::Variable { ref mut coefficient, .. } => *coefficient = *coefficient * by,
            &mut AbstVar::Constant { ref mut value, .. } => *value = *value * by,
            _ => panic!("Unexpected variant in this program logic."),
        };
    }
}
<|endoftext|>