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

use cmp;
use mem;
use ptr;

/// Rotation is much faster if it has access to a little bit of memory. This
/// union provides a RawVec-like interface, but to a fixed-size stack buffer.
#[allow(unions_with_drop_fields)]
union RawArray<T> {
    /// Ensure this is appropriately aligned for T, and is big
    /// enough for two elements even if T is enormous.
    typed: [T; 2],
    /// For normally-sized types, especially things like u8, having more
    /// than 2 in the buffer is necessary for usefulness, so pad it out
    /// enough to be helpful, but not so big as to risk overflow.
    _extra: [usize; 32],
}

impl<T> RawArray<T> {
    fn new() -> Self {
        unsafe { mem::uninitialized() }
    }
    fn ptr(&self) -> *mut T {
        unsafe { &self.typed as *const T as *mut T }
    }
    fn cap() -> usize {
        if mem::size_of::<T>() == 0 {
            usize::max_value()
        } else {
            mem::size_of::<Self>() / mem::size_of::<T>()
        }
    }
}

/// Rotates the range `[mid-left, mid+right)` such that the element at `mid`
/// becomes the first element.  Equivalently, rotates the range `left`
/// elements to the left or `right` elements to the right.
///
/// # Safety
///
/// The specified range must be valid for reading and writing.
/// The type `T` must have non-zero size.
///
/// # Algorithm
///
/// For longer rotations, swap the left-most `delta = min(left, right)`
/// elements with the right-most `delta` elements.  LLVM vectorizes this,
/// which is profitable as we only reach this step for a "large enough"
/// rotation.  Doing this puts `delta` elements on the larger side into the
/// correct position, leaving a smaller rotate problem.  Demonstration:
///
/// ```text
/// [ 6 7 8 9 10 11 12 13 . 1 2 3 4 5 ]
/// 1 2 3 4 5 [ 11 12 13 . 6 7 8 9 10 ]
/// 1 2 3 4 5 [ 8 9 10 . 6 7 ] 11 12 13
/// 1 2 3 4 5 6 7 [ 10 . 8 9 ] 11 12 13
/// 1 2 3 4 5 6 7 [ 9 . 8 ] 10 11 12 13
/// 1 2 3 4 5 6 7 8 [ . ] 9 10 11 12 13
/// ```
///
/// Once the rotation is small enough, copy some elements into a stack
/// buffer, `memmove` the others, and move the ones back from the buffer.
pub unsafe fn ptr_rotate<T>(mut left: usize, mid: *mut T, mut right: usize) {
    loop {
        let delta = cmp::min(left, right);
        if delta <= RawArray::<T>::cap() {
            break;
        }

        ptr_swap_n(
            mid.offset(-(left as isize)),
            mid.offset((right-delta) as isize),
            delta);

        if left <= right {
            right -= delta;
        } else {
            left -= delta;
        }
    }

    let rawarray = RawArray::new();
    let buf = rawarray.ptr();

    let dim = mid.offset(-(left as isize)).offset(right as isize);
    if left <= right {
        ptr::copy_nonoverlapping(mid.offset(-(left as isize)), buf, left);
        ptr::copy(mid, mid.offset(-(left as isize)), right);
        ptr::copy_nonoverlapping(buf, dim, left);
    }
    else {
        ptr::copy_nonoverlapping(mid, buf, right);
        ptr::copy(mid.offset(-(left as isize)), dim, left);
        ptr::copy_nonoverlapping(buf, mid.offset(-(left as isize)), right);
    }
}

unsafe fn ptr_swap_u8(a: *mut u8, b: *mut u8, n: usize) {
    for i in 0..n {
        ptr::swap(a.offset(i as isize), b.offset(i as isize));
    }
}
unsafe fn ptr_swap_u16(a: *mut u16, b: *mut u16, n: usize) {
    for i in 0..n {
        ptr::swap(a.offset(i as isize), b.offset(i as isize));
    }
}
unsafe fn ptr_swap_u32(a: *mut u32, b: *mut u32, n: usize) {
    for i in 0..n {
        ptr::swap(a.offset(i as isize), b.offset(i as isize));
    }
}
unsafe fn ptr_swap_u64(a: *mut u64, b: *mut u64, n: usize) {
    for i in 0..n {
        ptr::swap(a.offset(i as isize), b.offset(i as isize));
    }
}

unsafe fn ptr_swap_n<T>(a: *mut T, b: *mut T, n: usize) {
    // Doing this as a generic is 16% & 40% slower in two of the `String`
    // benchmarks, as (based on the block names) LLVM doesn't vectorize it.
    // Since this is just operating on raw memory, dispatch to a version
    // with appropriate alignment.  Helps with code size as well, by
    // avoiding monomorphizing different unrolled loops for `i32`,
    // `u32`, `f32`, `[u32; 1]`, etc.
    let size_of_t = mem::size_of::<T>();
    let align_of_t = mem::align_of::<T>();

    let a64 = mem::align_of::<u64>();
    if a64 == 8 && align_of_t % a64 == 0 {
        return ptr_swap_u64(a as *mut u64, b as *mut u64, n * (size_of_t / 8));
    }

    let a32 = mem::align_of::<u32>();
    if a32 == 4 && align_of_t % a32 == 0 {
        return ptr_swap_u32(a as *mut u32, b as *mut u32, n * (size_of_t / 4));
    }

    let a16 = mem::align_of::<u16>();
    if a16 == 2 && align_of_t % a16 == 0 {
        return ptr_swap_u16(a as *mut u16, b as *mut u16, n * (size_of_t / 2));
    }

    ptr_swap_u8(a as *mut u8, b as *mut u8, n * size_of_t);
}
<commit_msg>Remove the optimization in ptr_swap_n<commit_after>// Copyright 2012-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.

use cmp;
use mem;
use ptr;

/// Rotation is much faster if it has access to a little bit of memory. This
/// union provides a RawVec-like interface, but to a fixed-size stack buffer.
#[allow(unions_with_drop_fields)]
union RawArray<T> {
    /// Ensure this is appropriately aligned for T, and is big
    /// enough for two elements even if T is enormous.
    typed: [T; 2],
    /// For normally-sized types, especially things like u8, having more
    /// than 2 in the buffer is necessary for usefulness, so pad it out
    /// enough to be helpful, but not so big as to risk overflow.
    _extra: [usize; 32],
}

impl<T> RawArray<T> {
    fn new() -> Self {
        unsafe { mem::uninitialized() }
    }
    fn ptr(&self) -> *mut T {
        unsafe { &self.typed as *const T as *mut T }
    }
    fn cap() -> usize {
        if mem::size_of::<T>() == 0 {
            usize::max_value()
        } else {
            mem::size_of::<Self>() / mem::size_of::<T>()
        }
    }
}

/// Rotates the range `[mid-left, mid+right)` such that the element at `mid`
/// becomes the first element.  Equivalently, rotates the range `left`
/// elements to the left or `right` elements to the right.
///
/// # Safety
///
/// The specified range must be valid for reading and writing.
/// The type `T` must have non-zero size.
///
/// # Algorithm
///
/// For longer rotations, swap the left-most `delta = min(left, right)`
/// elements with the right-most `delta` elements.  LLVM vectorizes this,
/// which is profitable as we only reach this step for a "large enough"
/// rotation.  Doing this puts `delta` elements on the larger side into the
/// correct position, leaving a smaller rotate problem.  Demonstration:
///
/// ```text
/// [ 6 7 8 9 10 11 12 13 . 1 2 3 4 5 ]
/// 1 2 3 4 5 [ 11 12 13 . 6 7 8 9 10 ]
/// 1 2 3 4 5 [ 8 9 10 . 6 7 ] 11 12 13
/// 1 2 3 4 5 6 7 [ 10 . 8 9 ] 11 12 13
/// 1 2 3 4 5 6 7 [ 9 . 8 ] 10 11 12 13
/// 1 2 3 4 5 6 7 8 [ . ] 9 10 11 12 13
/// ```
///
/// Once the rotation is small enough, copy some elements into a stack
/// buffer, `memmove` the others, and move the ones back from the buffer.
pub unsafe fn ptr_rotate<T>(mut left: usize, mid: *mut T, mut right: usize) {
    loop {
        let delta = cmp::min(left, right);
        if delta <= RawArray::<T>::cap() {
            break;
        }

        ptr_swap_n(
            mid.offset(-(left as isize)),
            mid.offset((right-delta) as isize),
            delta);

        if left <= right {
            right -= delta;
        } else {
            left -= delta;
        }
    }

    let rawarray = RawArray::new();
    let buf = rawarray.ptr();

    let dim = mid.offset(-(left as isize)).offset(right as isize);
    if left <= right {
        ptr::copy_nonoverlapping(mid.offset(-(left as isize)), buf, left);
        ptr::copy(mid, mid.offset(-(left as isize)), right);
        ptr::copy_nonoverlapping(buf, dim, left);
    }
    else {
        ptr::copy_nonoverlapping(mid, buf, right);
        ptr::copy(mid.offset(-(left as isize)), dim, left);
        ptr::copy_nonoverlapping(buf, mid.offset(-(left as isize)), right);
    }
}

unsafe fn ptr_swap_n<T>(a: *mut T, b: *mut T, n: usize) {
    for i in 0..n {
        // These are nonoverlapping, so use mem::swap instead of ptr::swap
        mem::swap(&mut *a.offset(i as isize), &mut *b.offset(i as isize));
    }
}
<|endoftext|>


<commit_before><commit_msg>Fix clippy warning: next() instead of nth(0)<commit_after><|endoftext|>


<commit_before><commit_msg>adding buikd.rs<commit_after>extern crate gcc;

fn main () {
    println!("cargo:rustc-flags=-l dylib=stdc++");
}<|endoftext|>


<commit_before><commit_msg>rust 2<commit_after>fn get_next_fibs((n, m): (u32, u32)) -> (u32, u32) {
  (m, n + m)
}

fn main () {
  let mut current_fibs = (1, 1);
  let mut sum: u32 = 0;

  loop {
    current_fibs = get_next_fibs(current_fibs);

    let (n, _) = current_fibs;

    if n > 4000000 {
      break;
    }

    if n % 2 == 0 {
      sum += n;
    }
  }

  println!("{}", sum);
}
<|endoftext|>


<commit_before>use std::fmt::Display;
use thiserror::Error;

#[derive(Error, Debug)]
#[error("braced error: {msg}")]
struct BracedError {
    msg: String,
}

#[derive(Error, Debug)]
#[error("braced error")]
struct BracedUnused {
    extra: usize,
}

#[derive(Error, Debug)]
#[error("tuple error: {0}")]
struct TupleError(usize);

#[derive(Error, Debug)]
#[error("unit error")]
struct UnitError;

#[derive(Error, Debug)]
enum EnumError {
    #[error("braced error: {id}")]
    Braced { id: usize },
    #[error("tuple error: {0}")]
    Tuple(usize),
    #[error("unit error")]
    Unit,
}

#[derive(Error, Debug)]
#[error("{MSG}: {id:?} (code {CODE:?})")]
struct WithConstant {
    id: &'static str,
}

const MSG: &str = "failed to do";
const CODE: usize = 9;

#[derive(Error, Debug)]
#[error("{0}")]
enum Inherit {
    Unit(UnitError),
    #[error("other error")]
    Other(UnitError),
}

#[derive(Error, Debug)]
#[error("fn main() {{}}")]
struct Braces;

#[derive(Error, Debug)]
#[error("1 + 1 = {}", 1 + 1)]
struct Arithmetic;

#[derive(Error, Debug)]
#[error("!bool = {}", not(.0))]
struct NestedShorthand(bool);

#[derive(Error, Debug)]
#[error("...")]
pub enum Void {}

fn not(bool: &bool) -> bool {
    !*bool
}

fn assert<T: Display>(expected: &str, value: T) {
    assert_eq!(expected, value.to_string());
}

#[test]
fn test_display() {
    let msg = "T".to_owned();
    assert("braced error: T", BracedError { msg });
    assert("braced error", BracedUnused { extra: 0 });
    assert("tuple error: 0", TupleError(0));
    assert("unit error", UnitError);
    assert("braced error: 0", EnumError::Braced { id: 0 });
    assert("tuple error: 0", EnumError::Tuple(0));
    assert("unit error", EnumError::Unit);
    assert("failed to do: \"\" (code 9)", WithConstant { id: "" });
    assert("unit error", Inherit::Unit(UnitError));
    assert("other error", Inherit::Other(UnitError));
    assert("fn main() {}", Braces);
    assert("1 + 1 = 2", Arithmetic);
    assert("!bool = false", NestedShorthand(true));
}
<commit_msg>Break up test_display into individual test cases<commit_after>use std::fmt::Display;
use thiserror::Error;

fn assert<T: Display>(expected: &str, value: T) {
    assert_eq!(expected, value.to_string());
}

#[test]
fn test_braced() {
    #[derive(Error, Debug)]
    #[error("braced error: {msg}")]
    struct Error {
        msg: String,
    }

    let msg = "T".to_owned();
    assert("braced error: T", Error { msg });
}

#[test]
fn test_braced_unused() {
    #[derive(Error, Debug)]
    #[error("braced error")]
    struct Error {
        extra: usize,
    }

    assert("braced error", Error { extra: 0 });
}

#[test]
fn test_tuple() {
    #[derive(Error, Debug)]
    #[error("tuple error: {0}")]
    struct Error(usize);

    assert("tuple error: 0", Error(0));
}

#[test]
fn test_unit() {
    #[derive(Error, Debug)]
    #[error("unit error")]
    struct Error;

    assert("unit error", Error);
}

#[test]
fn test_enum() {
    #[derive(Error, Debug)]
    enum Error {
        #[error("braced error: {id}")]
        Braced { id: usize },
        #[error("tuple error: {0}")]
        Tuple(usize),
        #[error("unit error")]
        Unit,
    }

    assert("braced error: 0", Error::Braced { id: 0 });
    assert("tuple error: 0", Error::Tuple(0));
    assert("unit error", Error::Unit);
}

#[test]
fn test_constants() {
    #[derive(Error, Debug)]
    #[error("{MSG}: {id:?} (code {CODE:?})")]
    struct Error {
        id: &'static str,
    }

    const MSG: &str = "failed to do";
    const CODE: usize = 9;

    assert("failed to do: \"\" (code 9)", Error { id: "" });
}

#[test]
fn test_inherit() {
    #[derive(Error, Debug)]
    #[error("{0}")]
    enum Error {
        Some(&'static str),
        #[error("other error")]
        Other(&'static str),
    }

    assert("some error", Error::Some("some error"));
    assert("other error", Error::Other("..."));
}

#[test]
fn test_brace_escape() {
    #[derive(Error, Debug)]
    #[error("fn main() {{}}")]
    struct Error;

    assert("fn main() {}", Error);
}

#[test]
fn test_expr() {
    #[derive(Error, Debug)]
    #[error("1 + 1 = {}", 1 + 1)]
    struct Error;
    assert("1 + 1 = 2", Error);
}

#[test]
fn test_nested() {
    #[derive(Error, Debug)]
    #[error("!bool = {}", not(.0))]
    struct Error(bool);

    fn not(bool: &bool) -> bool {
        !*bool
    }

    assert("!bool = false", Error(true));
}

#[test]
fn test_void() {
    #[derive(Error, Debug)]
    #[error("...")]
    pub enum Error {}
}
<|endoftext|>


<commit_before><commit_msg>Add At as a DepFn and implement for tuples<commit_after><|endoftext|>


<commit_before>//! UDP relay proxy server

use std::{
    io::{self, Cursor},
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use bytes::BytesMut;
use futures::{future, future::AbortHandle, stream::FuturesUnordered, StreamExt};
use log::{debug, error, info, trace, warn};
use lru_time_cache::{Entry, LruCache};
use tokio::{
    self,
    net::udp::{RecvHalf, SendHalf},
    sync::{mpsc, Mutex},
};

use crate::{
    config::ServerConfig,
    context::{Context, SharedContext},
    relay::{
        flow::{SharedMultiServerFlowStatistic, SharedServerFlowStatistic},
        socks5::Address,
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

use super::{
    crypto_io::{decrypt_payload, encrypt_payload},
    DEFAULT_TIMEOUT,
    MAXIMUM_UDP_PAYLOAD_SIZE,
};

// Represent a UDP association
struct UdpAssociation {
    // local -> remote Queue
    // Drops tx, will close local -> remote task
    tx: mpsc::Sender<Vec<u8>>,

    // local <- remote task life watcher
    watcher: AbortHandle,
}

impl Drop for UdpAssociation {
    fn drop(&mut self) {
        self.watcher.abort();
    }
}

impl UdpAssociation {
    /// Create an association with addr
    async fn associate(
        context: SharedContext,
        svr_idx: usize,
        src_addr: SocketAddr,
        mut response_tx: mpsc::Sender<(SocketAddr, BytesMut)>,
    ) -> io::Result<UdpAssociation> {
        // Create a socket for receiving packets
        let local_addr = match context.config().local_addr {
            None => {
                // Let system allocate an address for us
                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)
            }
            Some(ref addr) => {
                // Uses configured local address
                addr.bind_addr(&context).await?
            }
        };
        let remote_udp = create_udp_socket(&local_addr).await?;

        let local_addr = remote_udp.local_addr().expect("could not determine port bound to");
        debug!("created UDP Association for {} from {}", src_addr, local_addr);

        // Create a channel for sending packets to remote
        // FIXME: Channel size 1024?
        let (tx, mut rx) = mpsc::channel::<Vec<u8>>(1024);

        // Splits socket into sender and receiver
        let (mut receiver, mut sender) = remote_udp.split();

        let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT);

        // local -> remote
        {
            let context = context.clone();
            tokio::spawn(async move {
                let svr_cfg = context.server_config(svr_idx);

                while let Some(pkt) = rx.recv().await {
                    // pkt is already a raw packet, so just send it
                    if let Err(err) =
                        UdpAssociation::relay_l2r(&context, src_addr, &mut sender, &pkt[..], timeout, svr_cfg).await
                    {
                        error!("failed to relay packet, {} -> ..., error: {}", src_addr, err);

                        // FIXME: Ignore? Or how to deal with it?
                    }
                }

                debug!("UDP ASSOCIATE {} -> .. finished", src_addr);
            });
        }

        let (r2l_task, close_flag) = future::abortable(async move {
            let svr_cfg = context.server_config(svr_idx);

            loop {
                // Read and send back to source
                match UdpAssociation::relay_r2l(&context, src_addr, &mut receiver, &mut response_tx, svr_cfg).await {
                    Ok(..) => {}
                    Err(err) => {
                        error!("failed to receive packet, {} <- .., error: {}", src_addr, err);

                        // FIXME: Don't break, or if you can find a way to drop the UdpAssociation
                        // break;
                    }
                }
            }
        });

        // local <- remote
        tokio::spawn(async move {
            let _ = r2l_task.await;

            debug!("UDP ASSOCIATE {} <- .. finished", src_addr);
        });

        Ok(UdpAssociation {
            tx,
            watcher: close_flag,
        })
    }

    /// Relay packets from local to remote
    async fn relay_l2r(
        context: &Context,
        src: SocketAddr,
        remote_udp: &mut SendHalf,
        pkt: &[u8],
        timeout: Duration,
        svr_cfg: &ServerConfig,
    ) -> io::Result<()> {
        // First of all, decrypt payload CLIENT -> SERVER
        let decrypted_pkt = match decrypt_payload(context, svr_cfg.method(), svr_cfg.key(), pkt) {
            Ok(Some(pkt)) => pkt,
            Ok(None) => {
                error!("failed to decrypt pkt in UDP relay, packet too short");
                let err = io::Error::new(io::ErrorKind::InvalidData, "packet too short");
                return Err(err);
            }
            Err(err) => {
                error!("failed to decrypt pkt in UDP relay: {}", err);
                let err = io::Error::new(io::ErrorKind::InvalidData, "decrypt failed");
                return Err(err);
            }
        };

        // CLIENT -> SERVER protocol: ADDRESS + PAYLOAD
        let mut cur = Cursor::new(decrypted_pkt);

        let addr = Address::read_from(&mut cur).await?;

        debug!("UDP ASSOCIATE {} <-> {} establishing", src, addr);

        if context.check_outbound_blocked(&addr) {
            warn!("outbound {} is blocked by ACL rules", addr);
            return Ok(());
        }

        // Take out internal buffer for optimizing one byte copy
        let header_len = cur.position() as usize;
        let decrypted_pkt = cur.into_inner();
        let body = &decrypted_pkt[header_len..];

        let send_len = match addr {
            Address::SocketAddress(ref remote_addr) => {
                debug!(
                    "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                    src,
                    addr,
                    remote_addr,
                    body.len()
                );
                try_timeout(remote_udp.send_to(body, remote_addr), Some(timeout)).await?
            }
            Address::DomainNameAddress(ref dname, port) => lookup_outbound_then!(context, dname, port, |remote_addr| {
                match try_timeout(remote_udp.send_to(body, &remote_addr), Some(timeout)).await {
                    Ok(l) => {
                        debug!(
                            "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                            src,
                            addr,
                            remote_addr,
                            body.len()
                        );
                        Ok(l)
                    }
                    Err(err) => {
                        error!(
                            "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                            src,
                            addr,
                            remote_addr,
                            body.len()
                        );
                        Err(err)
                    }
                }
            })
            .map(|(_, l)| l)?,
        };

        assert_eq!(body.len(), send_len);

        Ok(())
    }

    /// Relay packets from remote to local
    async fn relay_r2l(
        context: &Context,
        src_addr: SocketAddr,
        remote_udp: &mut RecvHalf,
        response_tx: &mut mpsc::Sender<(SocketAddr, BytesMut)>,
        svr_cfg: &ServerConfig,
    ) -> io::Result<()> {
        // Waiting for response from server SERVER -> CLIENT
        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        let mut remote_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
        let (remote_recv_len, remote_addr) = remote_udp.recv_from(&mut remote_buf).await?;

        debug!(
            "UDP ASSOCIATE {} <- {}, payload length {} bytes",
            src_addr, remote_addr, remote_recv_len
        );

        // FIXME: The Address should be the Address that client sent
        let addr = Address::SocketAddress(remote_addr);

        // CLIENT <- SERVER protocol: ADDRESS + PAYLOAD
        let mut send_buf = Vec::new();
        addr.write_to_buf(&mut send_buf);
        send_buf.extend_from_slice(&remote_buf[..remote_recv_len]);

        let mut encrypt_buf = BytesMut::new();
        encrypt_payload(context, svr_cfg.method(), svr_cfg.key(), &send_buf, &mut encrypt_buf)?;

        // Send back to src_addr
        if let Err(err) = response_tx.send((src_addr, encrypt_buf)).await {
            error!("failed to send packet into response channel, error: {}", err);

            // FIXME: What to do? Ignore?
        }

        Ok(())
    }

    // Send packet to remote
    //
    // Return `Err` if receiver have been closed
    async fn send(&mut self, pkt: Vec<u8>) {
        if let Err(..) = self.tx.send(pkt).await {
            // SHOULDn't HAPPEN
            unreachable!("UDP Association local -> remote Queue closed unexpectly");
        }
    }
}

fn serialize(saddr: &SocketAddr) -> [u8; 18] {
    let mut result = [0; 18];
    result[..16].copy_from_slice(&match saddr.ip() {
        IpAddr::V4(ref ip) => ip.to_ipv6_mapped(),
        IpAddr::V6(ref ip) => ip.clone(),
    }.octets());
    result[16..].copy_from_slice(&saddr.port().to_ne_bytes());
    result
}

async fn listen(context: SharedContext, flow_stat: SharedServerFlowStatistic, svr_idx: usize) -> io::Result<()> {
    let svr_cfg = context.server_config(svr_idx);
    let listen_addr = svr_cfg.addr().bind_addr(&context).await?;

    let listener = create_udp_socket(&listen_addr).await?;
    let local_addr = listener.local_addr().expect("determine port bound to");
    info!("shadowsocks UDP listening on {}", local_addr);

    let (mut r, mut w) = listener.split();

    // NOTE: Associations are only eliminated by expire time by default
    // So it may exhaust all available file descriptors
    let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT);
    let assoc_map = if let Some(max_assoc) = context.config().udp_max_associations {
        LruCache::with_expiry_duration_and_capacity(timeout, max_assoc)
    } else {
        LruCache::with_expiry_duration(timeout)
    };
    let assoc_map = Arc::new(Mutex::new(assoc_map));

    // FIXME: Channel size 1024?
    let (tx, mut rx) = mpsc::channel::<(SocketAddr, BytesMut)>(1024);

    {
        // Tokio task for sending data back to clients

        let assoc_map = assoc_map.clone();
        let flow_stat = flow_stat.clone();

        tokio::spawn(async move {
            while let Some((src, pkt)) = rx.recv().await {
                let cache_key = serialize(&src);
                {
                    let mut amap = assoc_map.lock().await;

                    // Check or update expire time
                    if amap.get(&cache_key).is_none() {
                        debug!(
                            "UDP association {} <-> ... is already expired, throwing away packet {} bytes",
                            src,
                            pkt.len()
                        );
                        continue;
                    }
                }

                if let Err(err) = w.send_to(&pkt, &src).await {
                    error!("UDP packet send failed, err: {:?}", err);
                    break;
                }

                flow_stat.udp().incr_tx(pkt.len() as u64);
            }

            // FIXME: How to stop the outer listener Future?
        });
    }

    let mut pkt_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];

    loop {
        let (recv_len, src) = r.recv_from(&mut pkt_buf).await?;

        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        let pkt = &pkt_buf[..recv_len];

        trace!("received UDP packet from {}, length {} bytes", src, recv_len);
        flow_stat.udp().incr_rx(pkt.len() as u64);

        if recv_len == 0 {
            // For windows, it will generate a ICMP Port Unreachable Message
            // https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-recvfrom
            // Which will result in recv_from return 0.
            //
            // It cannot be solved here, because `WSAGetLastError` is already set.
            //
            // See `relay::udprelay::utils::create_socket` for more detail.
            continue;
        }

        // Check ACL
        if context.check_client_blocked(&src) {
            warn!("client {} is blocked by ACL rules", src);
            continue;
        }

        // Check or (re)create an association
        {
            // Locks the whole association map
            let mut assoc_map = assoc_map.lock().await;

            // Get or create an association
            let assoc = match assoc_map.entry(serialize(&src)) {
                Entry::Occupied(oc) => oc.into_mut(),
                Entry::Vacant(vc) => vc.insert(
                    UdpAssociation::associate(context.clone(), svr_idx, src, tx.clone())
                        .await
                        .expect("create udp association"),
                ),
            };

            // FIXME: Lock is still kept for a mutable reference
            // Send to local -> remote task
            assoc.send(pkt.to_vec()).await;
        }
    }
}

/// Starts a UDP relay server
pub async fn run(context: SharedContext, flow_stat: SharedMultiServerFlowStatistic) -> io::Result<()> {
    let vec_fut = FuturesUnordered::new();

    for (svr_idx, svr_cfg) in context.config().server.iter().enumerate() {
        let context = context.clone();
        let flow_stat = flow_stat
            .get(svr_cfg.addr().port())
            .expect("port not existed in multi-server flow statistic")
            .clone();

        let svr_fut = listen(context, flow_stat, svr_idx);
        vec_fut.push(svr_fut);
    }

    match vec_fut.into_future().await.0 {
        Some(res) => {
            error!("one of UDP servers exited unexpectly, result: {:?}", res);
            let err = io::Error::new(io::ErrorKind::Other, "server exited unexpectly");
            Err(err)
        }
        None => unreachable!(),
    }
}
<commit_msg>Reduce unnecessary cloning<commit_after>//! UDP relay proxy server

use std::{
    io::{self, Cursor},
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use bytes::BytesMut;
use futures::{future, future::AbortHandle, stream::FuturesUnordered, StreamExt};
use log::{debug, error, info, trace, warn};
use lru_time_cache::{Entry, LruCache};
use tokio::{
    self,
    net::udp::{RecvHalf, SendHalf},
    sync::{mpsc, Mutex},
};

use crate::{
    config::ServerConfig,
    context::{Context, SharedContext},
    relay::{
        flow::{SharedMultiServerFlowStatistic, SharedServerFlowStatistic},
        socks5::Address,
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

use super::{
    crypto_io::{decrypt_payload, encrypt_payload},
    DEFAULT_TIMEOUT,
    MAXIMUM_UDP_PAYLOAD_SIZE,
};

// Represent a UDP association
struct UdpAssociation {
    // local -> remote Queue
    // Drops tx, will close local -> remote task
    tx: mpsc::Sender<Vec<u8>>,

    // local <- remote task life watcher
    watcher: AbortHandle,
}

impl Drop for UdpAssociation {
    fn drop(&mut self) {
        self.watcher.abort();
    }
}

impl UdpAssociation {
    /// Create an association with addr
    async fn associate(
        context: SharedContext,
        svr_idx: usize,
        src_addr: SocketAddr,
        mut response_tx: mpsc::Sender<(SocketAddr, BytesMut)>,
    ) -> io::Result<UdpAssociation> {
        // Create a socket for receiving packets
        let local_addr = match context.config().local_addr {
            None => {
                // Let system allocate an address for us
                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)
            }
            Some(ref addr) => {
                // Uses configured local address
                addr.bind_addr(&context).await?
            }
        };
        let remote_udp = create_udp_socket(&local_addr).await?;

        let local_addr = remote_udp.local_addr().expect("could not determine port bound to");
        debug!("created UDP Association for {} from {}", src_addr, local_addr);

        // Create a channel for sending packets to remote
        // FIXME: Channel size 1024?
        let (tx, mut rx) = mpsc::channel::<Vec<u8>>(1024);

        // Splits socket into sender and receiver
        let (mut receiver, mut sender) = remote_udp.split();

        let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT);

        // local -> remote
        {
            let context = context.clone();
            tokio::spawn(async move {
                let svr_cfg = context.server_config(svr_idx);

                while let Some(pkt) = rx.recv().await {
                    // pkt is already a raw packet, so just send it
                    if let Err(err) =
                        UdpAssociation::relay_l2r(&context, src_addr, &mut sender, &pkt[..], timeout, svr_cfg).await
                    {
                        error!("failed to relay packet, {} -> ..., error: {}", src_addr, err);

                        // FIXME: Ignore? Or how to deal with it?
                    }
                }

                debug!("UDP ASSOCIATE {} -> .. finished", src_addr);
            });
        }

        let (r2l_task, close_flag) = future::abortable(async move {
            let svr_cfg = context.server_config(svr_idx);

            loop {
                // Read and send back to source
                match UdpAssociation::relay_r2l(&context, src_addr, &mut receiver, &mut response_tx, svr_cfg).await {
                    Ok(..) => {}
                    Err(err) => {
                        error!("failed to receive packet, {} <- .., error: {}", src_addr, err);

                        // FIXME: Don't break, or if you can find a way to drop the UdpAssociation
                        // break;
                    }
                }
            }
        });

        // local <- remote
        tokio::spawn(async move {
            let _ = r2l_task.await;

            debug!("UDP ASSOCIATE {} <- .. finished", src_addr);
        });

        Ok(UdpAssociation {
            tx,
            watcher: close_flag,
        })
    }

    /// Relay packets from local to remote
    async fn relay_l2r(
        context: &Context,
        src: SocketAddr,
        remote_udp: &mut SendHalf,
        pkt: &[u8],
        timeout: Duration,
        svr_cfg: &ServerConfig,
    ) -> io::Result<()> {
        // First of all, decrypt payload CLIENT -> SERVER
        let decrypted_pkt = match decrypt_payload(context, svr_cfg.method(), svr_cfg.key(), pkt) {
            Ok(Some(pkt)) => pkt,
            Ok(None) => {
                error!("failed to decrypt pkt in UDP relay, packet too short");
                let err = io::Error::new(io::ErrorKind::InvalidData, "packet too short");
                return Err(err);
            }
            Err(err) => {
                error!("failed to decrypt pkt in UDP relay: {}", err);
                let err = io::Error::new(io::ErrorKind::InvalidData, "decrypt failed");
                return Err(err);
            }
        };

        // CLIENT -> SERVER protocol: ADDRESS + PAYLOAD
        let mut cur = Cursor::new(decrypted_pkt);

        let addr = Address::read_from(&mut cur).await?;

        debug!("UDP ASSOCIATE {} <-> {} establishing", src, addr);

        if context.check_outbound_blocked(&addr) {
            warn!("outbound {} is blocked by ACL rules", addr);
            return Ok(());
        }

        // Take out internal buffer for optimizing one byte copy
        let header_len = cur.position() as usize;
        let decrypted_pkt = cur.into_inner();
        let body = &decrypted_pkt[header_len..];

        let send_len = match addr {
            Address::SocketAddress(ref remote_addr) => {
                debug!(
                    "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                    src,
                    addr,
                    remote_addr,
                    body.len()
                );
                try_timeout(remote_udp.send_to(body, remote_addr), Some(timeout)).await?
            }
            Address::DomainNameAddress(ref dname, port) => lookup_outbound_then!(context, dname, port, |remote_addr| {
                match try_timeout(remote_udp.send_to(body, &remote_addr), Some(timeout)).await {
                    Ok(l) => {
                        debug!(
                            "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                            src,
                            addr,
                            remote_addr,
                            body.len()
                        );
                        Ok(l)
                    }
                    Err(err) => {
                        error!(
                            "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                            src,
                            addr,
                            remote_addr,
                            body.len()
                        );
                        Err(err)
                    }
                }
            })
            .map(|(_, l)| l)?,
        };

        assert_eq!(body.len(), send_len);

        Ok(())
    }

    /// Relay packets from remote to local
    async fn relay_r2l(
        context: &Context,
        src_addr: SocketAddr,
        remote_udp: &mut RecvHalf,
        response_tx: &mut mpsc::Sender<(SocketAddr, BytesMut)>,
        svr_cfg: &ServerConfig,
    ) -> io::Result<()> {
        // Waiting for response from server SERVER -> CLIENT
        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        let mut remote_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
        let (remote_recv_len, remote_addr) = remote_udp.recv_from(&mut remote_buf).await?;

        debug!(
            "UDP ASSOCIATE {} <- {}, payload length {} bytes",
            src_addr, remote_addr, remote_recv_len
        );

        // FIXME: The Address should be the Address that client sent
        let addr = Address::SocketAddress(remote_addr);

        // CLIENT <- SERVER protocol: ADDRESS + PAYLOAD
        let mut send_buf = Vec::new();
        addr.write_to_buf(&mut send_buf);
        send_buf.extend_from_slice(&remote_buf[..remote_recv_len]);

        let mut encrypt_buf = BytesMut::new();
        encrypt_payload(context, svr_cfg.method(), svr_cfg.key(), &send_buf, &mut encrypt_buf)?;

        // Send back to src_addr
        if let Err(err) = response_tx.send((src_addr, encrypt_buf)).await {
            error!("failed to send packet into response channel, error: {}", err);

            // FIXME: What to do? Ignore?
        }

        Ok(())
    }

    // Send packet to remote
    //
    // Return `Err` if receiver have been closed
    async fn send(&mut self, pkt: Vec<u8>) {
        if let Err(..) = self.tx.send(pkt).await {
            // SHOULDn't HAPPEN
            unreachable!("UDP Association local -> remote Queue closed unexpectly");
        }
    }
}

fn serialize(saddr: &SocketAddr) -> [u8; 18] {
    let mut result = [0; 18];
    result[..16].copy_from_slice(&match saddr.ip() {
        IpAddr::V4(ref ip) => &ip.to_ipv6_mapped(),
        IpAddr::V6(ref ip) => ip,
    }.octets());
    result[16..].copy_from_slice(&saddr.port().to_ne_bytes());
    result
}

async fn listen(context: SharedContext, flow_stat: SharedServerFlowStatistic, svr_idx: usize) -> io::Result<()> {
    let svr_cfg = context.server_config(svr_idx);
    let listen_addr = svr_cfg.addr().bind_addr(&context).await?;

    let listener = create_udp_socket(&listen_addr).await?;
    let local_addr = listener.local_addr().expect("determine port bound to");
    info!("shadowsocks UDP listening on {}", local_addr);

    let (mut r, mut w) = listener.split();

    // NOTE: Associations are only eliminated by expire time by default
    // So it may exhaust all available file descriptors
    let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT);
    let assoc_map = if let Some(max_assoc) = context.config().udp_max_associations {
        LruCache::with_expiry_duration_and_capacity(timeout, max_assoc)
    } else {
        LruCache::with_expiry_duration(timeout)
    };
    let assoc_map = Arc::new(Mutex::new(assoc_map));

    // FIXME: Channel size 1024?
    let (tx, mut rx) = mpsc::channel::<(SocketAddr, BytesMut)>(1024);

    {
        // Tokio task for sending data back to clients

        let assoc_map = assoc_map.clone();
        let flow_stat = flow_stat.clone();

        tokio::spawn(async move {
            while let Some((src, pkt)) = rx.recv().await {
                let cache_key = serialize(&src);
                {
                    let mut amap = assoc_map.lock().await;

                    // Check or update expire time
                    if amap.get(&cache_key).is_none() {
                        debug!(
                            "UDP association {} <-> ... is already expired, throwing away packet {} bytes",
                            src,
                            pkt.len()
                        );
                        continue;
                    }
                }

                if let Err(err) = w.send_to(&pkt, &src).await {
                    error!("UDP packet send failed, err: {:?}", err);
                    break;
                }

                flow_stat.udp().incr_tx(pkt.len() as u64);
            }

            // FIXME: How to stop the outer listener Future?
        });
    }

    let mut pkt_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];

    loop {
        let (recv_len, src) = r.recv_from(&mut pkt_buf).await?;

        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        let pkt = &pkt_buf[..recv_len];

        trace!("received UDP packet from {}, length {} bytes", src, recv_len);
        flow_stat.udp().incr_rx(pkt.len() as u64);

        if recv_len == 0 {
            // For windows, it will generate a ICMP Port Unreachable Message
            // https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-recvfrom
            // Which will result in recv_from return 0.
            //
            // It cannot be solved here, because `WSAGetLastError` is already set.
            //
            // See `relay::udprelay::utils::create_socket` for more detail.
            continue;
        }

        // Check ACL
        if context.check_client_blocked(&src) {
            warn!("client {} is blocked by ACL rules", src);
            continue;
        }

        // Check or (re)create an association
        {
            // Locks the whole association map
            let mut assoc_map = assoc_map.lock().await;

            // Get or create an association
            let assoc = match assoc_map.entry(serialize(&src)) {
                Entry::Occupied(oc) => oc.into_mut(),
                Entry::Vacant(vc) => vc.insert(
                    UdpAssociation::associate(context.clone(), svr_idx, src, tx.clone())
                        .await
                        .expect("create udp association"),
                ),
            };

            // FIXME: Lock is still kept for a mutable reference
            // Send to local -> remote task
            assoc.send(pkt.to_vec()).await;
        }
    }
}

/// Starts a UDP relay server
pub async fn run(context: SharedContext, flow_stat: SharedMultiServerFlowStatistic) -> io::Result<()> {
    let vec_fut = FuturesUnordered::new();

    for (svr_idx, svr_cfg) in context.config().server.iter().enumerate() {
        let context = context.clone();
        let flow_stat = flow_stat
            .get(svr_cfg.addr().port())
            .expect("port not existed in multi-server flow statistic")
            .clone();

        let svr_fut = listen(context, flow_stat, svr_idx);
        vec_fut.push(svr_fut);
    }

    match vec_fut.into_future().await.0 {
        Some(res) => {
            error!("one of UDP servers exited unexpectly, result: {:?}", res);
            let err = io::Error::new(io::ErrorKind::Other, "server exited unexpectly");
            Err(err)
        }
        None => unreachable!(),
    }
}
<|endoftext|>


<commit_before><commit_msg>Fix some error cases in lexer<commit_after><|endoftext|>


<commit_before><commit_msg>Formatting<commit_after><|endoftext|>


<commit_before><commit_msg>Remove unnecessary allocations<commit_after><|endoftext|>


<commit_before><commit_msg>watch app nodes<commit_after><|endoftext|>


<commit_before><commit_msg>Add rust code for lucians-luscious-lasagna<commit_after>pub fn expected_minutes_in_oven() -> i32 {
    40
}

pub fn remaining_minutes_in_oven(actual_minutes_in_oven: i32) -> i32 {
    40 - actual_minutes_in_oven
}

pub fn preparation_time_in_minutes(number_of_layers: i32) -> i32 {
    number_of_layers * 2
}

pub fn elapsed_time_in_minutes(number_of_layers: i32, actual_minutes_in_oven: i32) -> i32 {
    preparation_time_in_minutes(number_of_layers) + actual_minutes_in_oven
}
<|endoftext|>


<commit_before><commit_msg>Add help command to print usage<commit_after><|endoftext|>


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

use rustc::dep_graph::DepNode;
use rustc::hir::def_id::DefId;
use rustc::hir::svh::Svh;
use rustc::session::Session;
use rustc::ty::TyCtxt;
use rustc_data_structures::fx::FxHashMap;
use rustc_serialize::Encodable as RustcEncodable;
use rustc_serialize::opaque::Encoder;
use std::hash::Hash;
use std::io::{self, Cursor, Write};
use std::fs::{self, File};
use std::path::PathBuf;

use IncrementalHashesMap;
use ich::Fingerprint;
use super::data::*;
use super::directory::*;
use super::hash::*;
use super::preds::*;
use super::fs::*;
use super::dirty_clean;
use super::file_format;
use calculate_svh::IchHasher;

pub fn save_dep_graph<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
                                incremental_hashes_map: &IncrementalHashesMap,
                                svh: Svh) {
    debug!("save_dep_graph()");
    let _ignore = tcx.dep_graph.in_ignore();
    let sess = tcx.sess;
    if sess.opts.incremental.is_none() {
        return;
    }

    let mut builder = DefIdDirectoryBuilder::new(tcx);
    let query = tcx.dep_graph.query();
    let mut hcx = HashContext::new(tcx, incremental_hashes_map);
    let preds = Predecessors::new(&query, &mut hcx);
    let mut current_metadata_hashes = FxHashMap();

    // IMPORTANT: We are saving the metadata hashes *before* the dep-graph,
    //            since metadata-encoding might add new entries to the
    //            DefIdDirectory (which is saved in the dep-graph file).
    save_in(sess,
            metadata_hash_export_path(sess),
            |e| encode_metadata_hashes(tcx,
                                       svh,
                                       &preds,
                                       &mut builder,
                                       &mut current_metadata_hashes,
                                       e));
    save_in(sess,
            dep_graph_path(sess),
            |e| encode_dep_graph(&preds, &mut builder, e));

    let prev_metadata_hashes = incremental_hashes_map.prev_metadata_hashes.borrow();
    dirty_clean::check_dirty_clean_metadata(tcx,
                                            &*prev_metadata_hashes,
                                            &current_metadata_hashes);
}

pub fn save_work_products(sess: &Session) {
    if sess.opts.incremental.is_none() {
        return;
    }

    debug!("save_work_products()");
    let _ignore = sess.dep_graph.in_ignore();
    let path = work_products_path(sess);
    save_in(sess, path, |e| encode_work_products(sess, e));
}

fn save_in<F>(sess: &Session, path_buf: PathBuf, encode: F)
    where F: FnOnce(&mut Encoder) -> io::Result<()>
{
    debug!("save: storing data in {}", path_buf.display());

    // delete the old dep-graph, if any
    // Note: It's important that we actually delete the old file and not just
    // truncate and overwrite it, since it might be a shared hard-link, the
    // underlying data of which we don't want to modify
    if path_buf.exists() {
        match fs::remove_file(&path_buf) {
            Ok(()) => {
                debug!("save: remove old file");
            }
            Err(err) => {
                sess.err(&format!("unable to delete old dep-graph at `{}`: {}",
                                  path_buf.display(),
                                  err));
                return;
            }
        }
    }

    // generate the data in a memory buffer
    let mut wr = Cursor::new(Vec::new());
    file_format::write_file_header(&mut wr).unwrap();
    match encode(&mut Encoder::new(&mut wr)) {
        Ok(()) => {}
        Err(err) => {
            sess.err(&format!("could not encode dep-graph to `{}`: {}",
                              path_buf.display(),
                              err));
            return;
        }
    }

    // write the data out
    let data = wr.into_inner();
    match File::create(&path_buf).and_then(|mut file| file.write_all(&data)) {
        Ok(_) => {
            debug!("save: data written to disk successfully");
        }
        Err(err) => {
            sess.err(&format!("failed to write dep-graph to `{}`: {}",
                              path_buf.display(),
                              err));
            return;
        }
    }
}

pub fn encode_dep_graph(preds: &Predecessors,
                        builder: &mut DefIdDirectoryBuilder,
                        encoder: &mut Encoder)
                        -> io::Result<()> {
    // First encode the commandline arguments hash
    let tcx = builder.tcx();
    tcx.sess.opts.dep_tracking_hash().encode(encoder)?;

    // Create a flat list of (Input, WorkProduct) edges for
    // serialization.
    let mut edges = vec![];
    for (&target, sources) in &preds.inputs {
        match *target {
            DepNode::MetaData(ref def_id) => {
                // Metadata *targets* are always local metadata nodes. We have
                // already handled those in `encode_metadata_hashes`.
                assert!(def_id.is_local());
                continue;
            }
            _ => (),
        }
        let target = builder.map(target);
        for &source in sources {
            let source = builder.map(source);
            edges.push((source, target.clone()));
        }
    }

    if tcx.sess.opts.debugging_opts.incremental_dump_hash {
        for (dep_node, hash) in &preds.hashes {
            println!("HIR hash for {:?} is {}", dep_node, hash);
        }
    }

    // Create the serialized dep-graph.
    let graph = SerializedDepGraph {
        edges: edges,
        hashes: preds.hashes
            .iter()
            .map(|(&dep_node, &hash)| {
                SerializedHash {
                    dep_node: builder.map(dep_node),
                    hash: hash,
                }
            })
            .collect(),
    };

    debug!("graph = {:#?}", graph);

    // Encode the directory and then the graph data.
    builder.directory().encode(encoder)?;
    graph.encode(encoder)?;

    Ok(())
}

pub fn encode_metadata_hashes(tcx: TyCtxt,
                              svh: Svh,
                              preds: &Predecessors,
                              builder: &mut DefIdDirectoryBuilder,
                              current_metadata_hashes: &mut FxHashMap<DefId, Fingerprint>,
                              encoder: &mut Encoder)
                              -> io::Result<()> {
    // For each `MetaData(X)` node where `X` is local, accumulate a
    // hash.  These are the metadata items we export. Downstream
    // crates will want to see a hash that tells them whether we might
    // have changed the metadata for a given item since they last
    // compiled.
    //
    // (I initially wrote this with an iterator, but it seemed harder to read.)
    let mut serialized_hashes = SerializedMetadataHashes {
        hashes: vec![],
        index_map: FxHashMap()
    };

    let mut def_id_hashes = FxHashMap();

    for (&target, sources) in &preds.inputs {
        let def_id = match *target {
            DepNode::MetaData(def_id) => {
                assert!(def_id.is_local());
                def_id
            }
            _ => continue,
        };

        let mut def_id_hash = |def_id: DefId| -> u64 {
            *def_id_hashes.entry(def_id)
                .or_insert_with(|| {
                    let index = builder.add(def_id);
                    let path = builder.lookup_def_path(index);
                    path.deterministic_hash(tcx)
                })
        };

        // To create the hash for each item `X`, we don't hash the raw
        // bytes of the metadata (though in principle we
        // could). Instead, we walk the predecessors of `MetaData(X)`
        // from the dep-graph. This corresponds to all the inputs that
        // were read to construct the metadata. To create the hash for
        // the metadata, we hash (the hash of) all of those inputs.
        debug!("save: computing metadata hash for {:?}", def_id);

        // Create a vector containing a pair of (source-id, hash).
        // The source-id is stored as a `DepNode<u64>`, where the u64
        // is the det. hash of the def-path. This is convenient
        // because we can sort this to get a stable ordering across
        // compilations, even if the def-ids themselves have changed.
        let mut hashes: Vec<(DepNode<u64>, Fingerprint)> = sources.iter()
            .map(|dep_node| {
                let hash_dep_node = dep_node.map_def(|&def_id| Some(def_id_hash(def_id))).unwrap();
                let hash = preds.hashes[dep_node];
                (hash_dep_node, hash)
            })
            .collect();

        hashes.sort();
        let mut state = IchHasher::new();
        hashes.hash(&mut state);
        let hash = state.finish();

        debug!("save: metadata hash for {:?} is {}", def_id, hash);

        if tcx.sess.opts.debugging_opts.incremental_dump_hash {
            println!("metadata hash for {:?} is {}", def_id, hash);
            for dep_node in sources {
                println!("metadata hash for {:?} depends on {:?} with hash {}",
                         def_id, dep_node, preds.hashes[dep_node]);
            }
        }

        serialized_hashes.hashes.push(SerializedMetadataHash {
            def_index: def_id.index,
            hash: hash,
        });
    }

    if tcx.sess.opts.debugging_opts.query_dep_graph {
        for serialized_hash in &serialized_hashes.hashes {
            let def_id = DefId::local(serialized_hash.def_index);

            // Store entry in the index_map
            let def_path_index = builder.add(def_id);
            serialized_hashes.index_map.insert(def_id.index, def_path_index);

            // Record hash in current_metadata_hashes
            current_metadata_hashes.insert(def_id, serialized_hash.hash);
        }

        debug!("save: stored index_map (len={}) for serialized hashes",
               serialized_hashes.index_map.len());
    }

    // Encode everything.
    svh.encode(encoder)?;
    serialized_hashes.encode(encoder)?;

    Ok(())
}

pub fn encode_work_products(sess: &Session, encoder: &mut Encoder) -> io::Result<()> {
    let work_products: Vec<_> = sess.dep_graph
        .work_products()
        .iter()
        .map(|(id, work_product)| {
            SerializedWorkProduct {
                id: id.clone(),
                work_product: work_product.clone(),
            }
        })
        .collect();

    work_products.encode(encoder)
}
<commit_msg>Add some more info to -Zincremental-info<commit_after>// Copyright 2014 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.

use rustc::dep_graph::DepNode;
use rustc::hir::def_id::DefId;
use rustc::hir::svh::Svh;
use rustc::session::Session;
use rustc::ty::TyCtxt;
use rustc_data_structures::fx::FxHashMap;
use rustc_serialize::Encodable as RustcEncodable;
use rustc_serialize::opaque::Encoder;
use std::hash::Hash;
use std::io::{self, Cursor, Write};
use std::fs::{self, File};
use std::path::PathBuf;

use IncrementalHashesMap;
use ich::Fingerprint;
use super::data::*;
use super::directory::*;
use super::hash::*;
use super::preds::*;
use super::fs::*;
use super::dirty_clean;
use super::file_format;
use calculate_svh::IchHasher;

pub fn save_dep_graph<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
                                incremental_hashes_map: &IncrementalHashesMap,
                                svh: Svh) {
    debug!("save_dep_graph()");
    let _ignore = tcx.dep_graph.in_ignore();
    let sess = tcx.sess;
    if sess.opts.incremental.is_none() {
        return;
    }

    let mut builder = DefIdDirectoryBuilder::new(tcx);
    let query = tcx.dep_graph.query();

    if tcx.sess.opts.debugging_opts.incremental_info {
        println!("incremental: {} nodes in dep-graph", query.graph.len_nodes());
        println!("incremental: {} edges in dep-graph", query.graph.len_edges());
    }

    let mut hcx = HashContext::new(tcx, incremental_hashes_map);
    let preds = Predecessors::new(&query, &mut hcx);
    let mut current_metadata_hashes = FxHashMap();

    // IMPORTANT: We are saving the metadata hashes *before* the dep-graph,
    //            since metadata-encoding might add new entries to the
    //            DefIdDirectory (which is saved in the dep-graph file).
    save_in(sess,
            metadata_hash_export_path(sess),
            |e| encode_metadata_hashes(tcx,
                                       svh,
                                       &preds,
                                       &mut builder,
                                       &mut current_metadata_hashes,
                                       e));
    save_in(sess,
            dep_graph_path(sess),
            |e| encode_dep_graph(&preds, &mut builder, e));

    let prev_metadata_hashes = incremental_hashes_map.prev_metadata_hashes.borrow();
    dirty_clean::check_dirty_clean_metadata(tcx,
                                            &*prev_metadata_hashes,
                                            &current_metadata_hashes);
}

pub fn save_work_products(sess: &Session) {
    if sess.opts.incremental.is_none() {
        return;
    }

    debug!("save_work_products()");
    let _ignore = sess.dep_graph.in_ignore();
    let path = work_products_path(sess);
    save_in(sess, path, |e| encode_work_products(sess, e));
}

fn save_in<F>(sess: &Session, path_buf: PathBuf, encode: F)
    where F: FnOnce(&mut Encoder) -> io::Result<()>
{
    debug!("save: storing data in {}", path_buf.display());

    // delete the old dep-graph, if any
    // Note: It's important that we actually delete the old file and not just
    // truncate and overwrite it, since it might be a shared hard-link, the
    // underlying data of which we don't want to modify
    if path_buf.exists() {
        match fs::remove_file(&path_buf) {
            Ok(()) => {
                debug!("save: remove old file");
            }
            Err(err) => {
                sess.err(&format!("unable to delete old dep-graph at `{}`: {}",
                                  path_buf.display(),
                                  err));
                return;
            }
        }
    }

    // generate the data in a memory buffer
    let mut wr = Cursor::new(Vec::new());
    file_format::write_file_header(&mut wr).unwrap();
    match encode(&mut Encoder::new(&mut wr)) {
        Ok(()) => {}
        Err(err) => {
            sess.err(&format!("could not encode dep-graph to `{}`: {}",
                              path_buf.display(),
                              err));
            return;
        }
    }

    // write the data out
    let data = wr.into_inner();
    match File::create(&path_buf).and_then(|mut file| file.write_all(&data)) {
        Ok(_) => {
            debug!("save: data written to disk successfully");
        }
        Err(err) => {
            sess.err(&format!("failed to write dep-graph to `{}`: {}",
                              path_buf.display(),
                              err));
            return;
        }
    }
}

pub fn encode_dep_graph(preds: &Predecessors,
                        builder: &mut DefIdDirectoryBuilder,
                        encoder: &mut Encoder)
                        -> io::Result<()> {
    // First encode the commandline arguments hash
    let tcx = builder.tcx();
    tcx.sess.opts.dep_tracking_hash().encode(encoder)?;

    // Create a flat list of (Input, WorkProduct) edges for
    // serialization.
    let mut edges = vec![];
    for (&target, sources) in &preds.inputs {
        match *target {
            DepNode::MetaData(ref def_id) => {
                // Metadata *targets* are always local metadata nodes. We have
                // already handled those in `encode_metadata_hashes`.
                assert!(def_id.is_local());
                continue;
            }
            _ => (),
        }
        let target = builder.map(target);
        for &source in sources {
            let source = builder.map(source);
            edges.push((source, target.clone()));
        }
    }

    if tcx.sess.opts.debugging_opts.incremental_dump_hash {
        for (dep_node, hash) in &preds.hashes {
            println!("HIR hash for {:?} is {}", dep_node, hash);
        }
    }

    // Create the serialized dep-graph.
    let graph = SerializedDepGraph {
        edges: edges,
        hashes: preds.hashes
            .iter()
            .map(|(&dep_node, &hash)| {
                SerializedHash {
                    dep_node: builder.map(dep_node),
                    hash: hash,
                }
            })
            .collect(),
    };

    if tcx.sess.opts.debugging_opts.incremental_info {
        println!("incremental: {} edges in serialized dep-graph", graph.edges.len());
        println!("incremental: {} hashes in serialized dep-graph", graph.hashes.len());
    }

    debug!("graph = {:#?}", graph);

    // Encode the directory and then the graph data.
    builder.directory().encode(encoder)?;
    graph.encode(encoder)?;

    Ok(())
}

pub fn encode_metadata_hashes(tcx: TyCtxt,
                              svh: Svh,
                              preds: &Predecessors,
                              builder: &mut DefIdDirectoryBuilder,
                              current_metadata_hashes: &mut FxHashMap<DefId, Fingerprint>,
                              encoder: &mut Encoder)
                              -> io::Result<()> {
    // For each `MetaData(X)` node where `X` is local, accumulate a
    // hash.  These are the metadata items we export. Downstream
    // crates will want to see a hash that tells them whether we might
    // have changed the metadata for a given item since they last
    // compiled.
    //
    // (I initially wrote this with an iterator, but it seemed harder to read.)
    let mut serialized_hashes = SerializedMetadataHashes {
        hashes: vec![],
        index_map: FxHashMap()
    };

    let mut def_id_hashes = FxHashMap();

    for (&target, sources) in &preds.inputs {
        let def_id = match *target {
            DepNode::MetaData(def_id) => {
                assert!(def_id.is_local());
                def_id
            }
            _ => continue,
        };

        let mut def_id_hash = |def_id: DefId| -> u64 {
            *def_id_hashes.entry(def_id)
                .or_insert_with(|| {
                    let index = builder.add(def_id);
                    let path = builder.lookup_def_path(index);
                    path.deterministic_hash(tcx)
                })
        };

        // To create the hash for each item `X`, we don't hash the raw
        // bytes of the metadata (though in principle we
        // could). Instead, we walk the predecessors of `MetaData(X)`
        // from the dep-graph. This corresponds to all the inputs that
        // were read to construct the metadata. To create the hash for
        // the metadata, we hash (the hash of) all of those inputs.
        debug!("save: computing metadata hash for {:?}", def_id);

        // Create a vector containing a pair of (source-id, hash).
        // The source-id is stored as a `DepNode<u64>`, where the u64
        // is the det. hash of the def-path. This is convenient
        // because we can sort this to get a stable ordering across
        // compilations, even if the def-ids themselves have changed.
        let mut hashes: Vec<(DepNode<u64>, Fingerprint)> = sources.iter()
            .map(|dep_node| {
                let hash_dep_node = dep_node.map_def(|&def_id| Some(def_id_hash(def_id))).unwrap();
                let hash = preds.hashes[dep_node];
                (hash_dep_node, hash)
            })
            .collect();

        hashes.sort();
        let mut state = IchHasher::new();
        hashes.hash(&mut state);
        let hash = state.finish();

        debug!("save: metadata hash for {:?} is {}", def_id, hash);

        if tcx.sess.opts.debugging_opts.incremental_dump_hash {
            println!("metadata hash for {:?} is {}", def_id, hash);
            for dep_node in sources {
                println!("metadata hash for {:?} depends on {:?} with hash {}",
                         def_id, dep_node, preds.hashes[dep_node]);
            }
        }

        serialized_hashes.hashes.push(SerializedMetadataHash {
            def_index: def_id.index,
            hash: hash,
        });
    }

    if tcx.sess.opts.debugging_opts.query_dep_graph {
        for serialized_hash in &serialized_hashes.hashes {
            let def_id = DefId::local(serialized_hash.def_index);

            // Store entry in the index_map
            let def_path_index = builder.add(def_id);
            serialized_hashes.index_map.insert(def_id.index, def_path_index);

            // Record hash in current_metadata_hashes
            current_metadata_hashes.insert(def_id, serialized_hash.hash);
        }

        debug!("save: stored index_map (len={}) for serialized hashes",
               serialized_hashes.index_map.len());
    }

    // Encode everything.
    svh.encode(encoder)?;
    serialized_hashes.encode(encoder)?;

    Ok(())
}

pub fn encode_work_products(sess: &Session, encoder: &mut Encoder) -> io::Result<()> {
    let work_products: Vec<_> = sess.dep_graph
        .work_products()
        .iter()
        .map(|(id, work_product)| {
            SerializedWorkProduct {
                id: id.clone(),
                work_product: work_product.clone(),
            }
        })
        .collect();

    work_products.encode(encoder)
}
<|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.

use borrow_check::borrow_set::BorrowSet;
use rustc::hir::def_id::DefId;
use rustc::mir::{ClosureRegionRequirements, ClosureOutlivesSubject, Mir};
use rustc::infer::InferCtxt;
use rustc::ty::{self, RegionKind, RegionVid};
use rustc::util::nodemap::FxHashMap;
use std::collections::BTreeSet;
use std::fmt::Debug;
use std::io;
use transform::MirSource;
use util::liveness::{LivenessResults, LocalSet};
use dataflow::FlowAtLocation;
use dataflow::MaybeInitializedPlaces;
use dataflow::move_paths::MoveData;

use util as mir_util;
use util::pretty::{self, ALIGN};
use self::mir_util::PassWhere;

mod constraint_generation;
pub mod explain_borrow;
pub(crate) mod region_infer;
mod renumber;
mod subtype_constraint_generation;
pub(crate) mod type_check;
mod universal_regions;

use self::region_infer::RegionInferenceContext;
use self::universal_regions::UniversalRegions;

/// Rewrites the regions in the MIR to use NLL variables, also
/// scraping out the set of universal regions (e.g., region parameters)
/// declared on the function. That set will need to be given to
/// `compute_regions`.
pub(in borrow_check) fn replace_regions_in_mir<'cx, 'gcx, 'tcx>(
    infcx: &InferCtxt<'cx, 'gcx, 'tcx>,
    def_id: DefId,
    param_env: ty::ParamEnv<'tcx>,
    mir: &mut Mir<'tcx>,
) -> UniversalRegions<'tcx> {
    debug!("replace_regions_in_mir(def_id={:?})", def_id);

    // Compute named region information. This also renumbers the inputs/outputs.
    let universal_regions = UniversalRegions::new(infcx, def_id, param_env);

    // Replace all remaining regions with fresh inference variables.
    renumber::renumber_mir(infcx, mir);

    let source = MirSource::item(def_id);
    mir_util::dump_mir(infcx.tcx, None, "renumber", &0, source, mir, |_, _| Ok(()));

    universal_regions
}

/// Computes the (non-lexical) regions from the input MIR.
///
/// This may result in errors being reported.
pub(in borrow_check) fn compute_regions<'cx, 'gcx, 'tcx>(
    infcx: &InferCtxt<'cx, 'gcx, 'tcx>,
    def_id: DefId,
    universal_regions: UniversalRegions<'tcx>,
    mir: &Mir<'tcx>,
    param_env: ty::ParamEnv<'gcx>,
    flow_inits: &mut FlowAtLocation<MaybeInitializedPlaces<'cx, 'gcx, 'tcx>>,
    move_data: &MoveData<'tcx>,
    _borrow_set: &BorrowSet<'tcx>,
) -> (
    RegionInferenceContext<'tcx>,
    Option<ClosureRegionRequirements<'gcx>>,
) {
    // Run the MIR type-checker.
    let liveness = &LivenessResults::compute(mir);
    let constraint_sets = &type_check::type_check(
        infcx,
        param_env,
        mir,
        def_id,
        &universal_regions,
        &liveness,
        flow_inits,
        move_data,
    );

    // Create the region inference context, taking ownership of the region inference
    // data that was contained in `infcx`.
    let var_origins = infcx.take_region_var_origins();
    let mut regioncx = RegionInferenceContext::new(var_origins, universal_regions, mir);
    subtype_constraint_generation::generate(&mut regioncx, mir, constraint_sets);

    // Generate non-subtyping constraints.
    constraint_generation::generate_constraints(infcx, &mut regioncx, &mir);

    // Solve the region constraints.
    let closure_region_requirements = regioncx.solve(infcx, &mir, def_id);

    // Dump MIR results into a file, if that is enabled. This let us
    // write unit-tests, as well as helping with debugging.
    dump_mir_results(
        infcx,
        liveness,
        MirSource::item(def_id),
        &mir,
        &regioncx,
        &closure_region_requirements,
    );

    // We also have a `#[rustc_nll]` annotation that causes us to dump
    // information
    dump_annotation(infcx, &mir, def_id, &regioncx, &closure_region_requirements);

    (regioncx, closure_region_requirements)
}

fn dump_mir_results<'a, 'gcx, 'tcx>(
    infcx: &InferCtxt<'a, 'gcx, 'tcx>,
    liveness: &LivenessResults,
    source: MirSource,
    mir: &Mir<'tcx>,
    regioncx: &RegionInferenceContext,
    closure_region_requirements: &Option<ClosureRegionRequirements>,
) {
    if !mir_util::dump_enabled(infcx.tcx, "nll", source) {
        return;
    }

    let regular_liveness_per_location: FxHashMap<_, _> = mir.basic_blocks()
        .indices()
        .flat_map(|bb| {
            let mut results = vec![];
            liveness
                .regular
                .simulate_block(&mir, bb, |location, local_set| {
                    results.push((location, local_set.clone()));
                });
            results
        })
        .collect();

    let drop_liveness_per_location: FxHashMap<_, _> = mir.basic_blocks()
        .indices()
        .flat_map(|bb| {
            let mut results = vec![];
            liveness
                .drop
                .simulate_block(&mir, bb, |location, local_set| {
                    results.push((location, local_set.clone()));
                });
            results
        })
        .collect();

    mir_util::dump_mir(infcx.tcx, None, "nll", &0, source, mir, |pass_where, out| {
        match pass_where {
            // Before the CFG, dump out the values for each region variable.
            PassWhere::BeforeCFG => {
                regioncx.dump_mir(out)?;

                if let Some(closure_region_requirements) = closure_region_requirements {
                    writeln!(out, "|")?;
                    writeln!(out, "| Free Region Constraints")?;
                    for_each_region_constraint(closure_region_requirements, &mut |msg| {
                        writeln!(out, "| {}", msg)
                    })?;
                }
            }

            // Before each basic block, dump out the values
            // that are live on entry to the basic block.
            PassWhere::BeforeBlock(bb) => {
                let s = live_variable_set(&liveness.regular.ins[bb], &liveness.drop.ins[bb]);
                writeln!(out, "    | Live variables on entry to {:?}: {}", bb, s)?;
            }

            PassWhere::BeforeLocation(location) => {
                let s = live_variable_set(
                    &regular_liveness_per_location[&location],
                    &drop_liveness_per_location[&location],
                );
                writeln!(
                    out,
                    "{:ALIGN$} | Live variables on entry to {:?}: {}",
                    "",
                    location,
                    s,
                    ALIGN = ALIGN
                )?;
            }

            PassWhere::AfterLocation(_) | PassWhere::AfterCFG => {}
        }
        Ok(())
    });

    // Also dump the inference graph constraints as a graphviz file.
    let _: io::Result<()> = do_catch! {{
        let mut file =
            pretty::create_dump_file(infcx.tcx, "regioncx.dot", None, "nll", &0, source)?;
        regioncx.dump_graphviz(&mut file)?;
    }};
}

fn dump_annotation<'a, 'gcx, 'tcx>(
    infcx: &InferCtxt<'a, 'gcx, 'tcx>,
    mir: &Mir<'tcx>,
    mir_def_id: DefId,
    regioncx: &RegionInferenceContext,
    closure_region_requirements: &Option<ClosureRegionRequirements>,
) {
    let tcx = infcx.tcx;
    let base_def_id = tcx.closure_base_def_id(mir_def_id);
    if !tcx.has_attr(base_def_id, "rustc_regions") {
        return;
    }

    // When the enclosing function is tagged with `#[rustc_regions]`,
    // we dump out various bits of state as warnings. This is useful
    // for verifying that the compiler is behaving as expected.  These
    // warnings focus on the closure region requirements -- for
    // viewing the intraprocedural state, the -Zdump-mir output is
    // better.

    if let Some(closure_region_requirements) = closure_region_requirements {
        let mut err = tcx.sess
            .diagnostic()
            .span_note_diag(mir.span, "External requirements");

        regioncx.annotate(&mut err);

        err.note(&format!(
            "number of external vids: {}",
            closure_region_requirements.num_external_vids
        ));

        // Dump the region constraints we are imposing *between* those
        // newly created variables.
        for_each_region_constraint(closure_region_requirements, &mut |msg| {
            err.note(msg);
            Ok(())
        }).unwrap();

        err.emit();
    } else {
        let mut err = tcx.sess
            .diagnostic()
            .span_note_diag(mir.span, "No external requirements");
        regioncx.annotate(&mut err);
        err.emit();
    }
}

fn for_each_region_constraint(
    closure_region_requirements: &ClosureRegionRequirements,
    with_msg: &mut dyn FnMut(&str) -> io::Result<()>,
) -> io::Result<()> {
    for req in &closure_region_requirements.outlives_requirements {
        let subject: &dyn Debug = match &req.subject {
            ClosureOutlivesSubject::Region(subject) => subject,
            ClosureOutlivesSubject::Ty(ty) => ty,
        };
        with_msg(&format!(
            "where {:?}: {:?}",
            subject,
            req.outlived_free_region,
        ))?;
    }
    Ok(())
}

/// Right now, we piggy back on the `ReVar` to store our NLL inference
/// regions. These are indexed with `RegionVid`. This method will
/// assert that the region is a `ReVar` and extract its internal index.
/// This is reasonable because in our MIR we replace all universal regions
/// with inference variables.
pub trait ToRegionVid {
    fn to_region_vid(self) -> RegionVid;
}

impl<'tcx> ToRegionVid for &'tcx RegionKind {
    fn to_region_vid(self) -> RegionVid {
        if let ty::ReVar(vid) = self {
            *vid
        } else {
            bug!("region is not an ReVar: {:?}", self)
        }
    }
}

impl ToRegionVid for RegionVid {
    fn to_region_vid(self) -> RegionVid {
        self
    }
}

fn live_variable_set(regular: &LocalSet, drops: &LocalSet) -> String {
    // sort and deduplicate:
    let all_locals: BTreeSet<_> = regular.iter().chain(drops.iter()).collect();

    // construct a string with each local, including `(drop)` if it is
    // only dropped, versus a regular use.
    let mut string = String::new();
    for local in all_locals {
        string.push_str(&format!("{:?}", local));

        if !regular.contains(&local) {
            assert!(drops.contains(&local));
            string.push_str(" (drop)");
        }

        string.push_str(", ");
    }

    let len = if string.is_empty() {
        0
    } else {
        string.len() - 2
    };

    format!("[{}]", &string[..len])
}
<commit_msg>use `crate` visibility modifier<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.

use borrow_check::borrow_set::BorrowSet;
use rustc::hir::def_id::DefId;
use rustc::mir::{ClosureRegionRequirements, ClosureOutlivesSubject, Mir};
use rustc::infer::InferCtxt;
use rustc::ty::{self, RegionKind, RegionVid};
use rustc::util::nodemap::FxHashMap;
use std::collections::BTreeSet;
use std::fmt::Debug;
use std::io;
use transform::MirSource;
use util::liveness::{LivenessResults, LocalSet};
use dataflow::FlowAtLocation;
use dataflow::MaybeInitializedPlaces;
use dataflow::move_paths::MoveData;

use util as mir_util;
use util::pretty::{self, ALIGN};
use self::mir_util::PassWhere;

mod constraint_generation;
pub mod explain_borrow;
crate mod region_infer;
mod renumber;
mod subtype_constraint_generation;
crate mod type_check;
mod universal_regions;

use self::region_infer::RegionInferenceContext;
use self::universal_regions::UniversalRegions;

/// Rewrites the regions in the MIR to use NLL variables, also
/// scraping out the set of universal regions (e.g., region parameters)
/// declared on the function. That set will need to be given to
/// `compute_regions`.
pub(in borrow_check) fn replace_regions_in_mir<'cx, 'gcx, 'tcx>(
    infcx: &InferCtxt<'cx, 'gcx, 'tcx>,
    def_id: DefId,
    param_env: ty::ParamEnv<'tcx>,
    mir: &mut Mir<'tcx>,
) -> UniversalRegions<'tcx> {
    debug!("replace_regions_in_mir(def_id={:?})", def_id);

    // Compute named region information. This also renumbers the inputs/outputs.
    let universal_regions = UniversalRegions::new(infcx, def_id, param_env);

    // Replace all remaining regions with fresh inference variables.
    renumber::renumber_mir(infcx, mir);

    let source = MirSource::item(def_id);
    mir_util::dump_mir(infcx.tcx, None, "renumber", &0, source, mir, |_, _| Ok(()));

    universal_regions
}

/// Computes the (non-lexical) regions from the input MIR.
///
/// This may result in errors being reported.
pub(in borrow_check) fn compute_regions<'cx, 'gcx, 'tcx>(
    infcx: &InferCtxt<'cx, 'gcx, 'tcx>,
    def_id: DefId,
    universal_regions: UniversalRegions<'tcx>,
    mir: &Mir<'tcx>,
    param_env: ty::ParamEnv<'gcx>,
    flow_inits: &mut FlowAtLocation<MaybeInitializedPlaces<'cx, 'gcx, 'tcx>>,
    move_data: &MoveData<'tcx>,
    _borrow_set: &BorrowSet<'tcx>,
) -> (
    RegionInferenceContext<'tcx>,
    Option<ClosureRegionRequirements<'gcx>>,
) {
    // Run the MIR type-checker.
    let liveness = &LivenessResults::compute(mir);
    let constraint_sets = &type_check::type_check(
        infcx,
        param_env,
        mir,
        def_id,
        &universal_regions,
        &liveness,
        flow_inits,
        move_data,
    );

    // Create the region inference context, taking ownership of the region inference
    // data that was contained in `infcx`.
    let var_origins = infcx.take_region_var_origins();
    let mut regioncx = RegionInferenceContext::new(var_origins, universal_regions, mir);
    subtype_constraint_generation::generate(&mut regioncx, mir, constraint_sets);

    // Generate non-subtyping constraints.
    constraint_generation::generate_constraints(infcx, &mut regioncx, &mir);

    // Solve the region constraints.
    let closure_region_requirements = regioncx.solve(infcx, &mir, def_id);

    // Dump MIR results into a file, if that is enabled. This let us
    // write unit-tests, as well as helping with debugging.
    dump_mir_results(
        infcx,
        liveness,
        MirSource::item(def_id),
        &mir,
        &regioncx,
        &closure_region_requirements,
    );

    // We also have a `#[rustc_nll]` annotation that causes us to dump
    // information
    dump_annotation(infcx, &mir, def_id, &regioncx, &closure_region_requirements);

    (regioncx, closure_region_requirements)
}

fn dump_mir_results<'a, 'gcx, 'tcx>(
    infcx: &InferCtxt<'a, 'gcx, 'tcx>,
    liveness: &LivenessResults,
    source: MirSource,
    mir: &Mir<'tcx>,
    regioncx: &RegionInferenceContext,
    closure_region_requirements: &Option<ClosureRegionRequirements>,
) {
    if !mir_util::dump_enabled(infcx.tcx, "nll", source) {
        return;
    }

    let regular_liveness_per_location: FxHashMap<_, _> = mir.basic_blocks()
        .indices()
        .flat_map(|bb| {
            let mut results = vec![];
            liveness
                .regular
                .simulate_block(&mir, bb, |location, local_set| {
                    results.push((location, local_set.clone()));
                });
            results
        })
        .collect();

    let drop_liveness_per_location: FxHashMap<_, _> = mir.basic_blocks()
        .indices()
        .flat_map(|bb| {
            let mut results = vec![];
            liveness
                .drop
                .simulate_block(&mir, bb, |location, local_set| {
                    results.push((location, local_set.clone()));
                });
            results
        })
        .collect();

    mir_util::dump_mir(infcx.tcx, None, "nll", &0, source, mir, |pass_where, out| {
        match pass_where {
            // Before the CFG, dump out the values for each region variable.
            PassWhere::BeforeCFG => {
                regioncx.dump_mir(out)?;

                if let Some(closure_region_requirements) = closure_region_requirements {
                    writeln!(out, "|")?;
                    writeln!(out, "| Free Region Constraints")?;
                    for_each_region_constraint(closure_region_requirements, &mut |msg| {
                        writeln!(out, "| {}", msg)
                    })?;
                }
            }

            // Before each basic block, dump out the values
            // that are live on entry to the basic block.
            PassWhere::BeforeBlock(bb) => {
                let s = live_variable_set(&liveness.regular.ins[bb], &liveness.drop.ins[bb]);
                writeln!(out, "    | Live variables on entry to {:?}: {}", bb, s)?;
            }

            PassWhere::BeforeLocation(location) => {
                let s = live_variable_set(
                    &regular_liveness_per_location[&location],
                    &drop_liveness_per_location[&location],
                );
                writeln!(
                    out,
                    "{:ALIGN$} | Live variables on entry to {:?}: {}",
                    "",
                    location,
                    s,
                    ALIGN = ALIGN
                )?;
            }

            PassWhere::AfterLocation(_) | PassWhere::AfterCFG => {}
        }
        Ok(())
    });

    // Also dump the inference graph constraints as a graphviz file.
    let _: io::Result<()> = do_catch! {{
        let mut file =
            pretty::create_dump_file(infcx.tcx, "regioncx.dot", None, "nll", &0, source)?;
        regioncx.dump_graphviz(&mut file)?;
    }};
}

fn dump_annotation<'a, 'gcx, 'tcx>(
    infcx: &InferCtxt<'a, 'gcx, 'tcx>,
    mir: &Mir<'tcx>,
    mir_def_id: DefId,
    regioncx: &RegionInferenceContext,
    closure_region_requirements: &Option<ClosureRegionRequirements>,
) {
    let tcx = infcx.tcx;
    let base_def_id = tcx.closure_base_def_id(mir_def_id);
    if !tcx.has_attr(base_def_id, "rustc_regions") {
        return;
    }

    // When the enclosing function is tagged with `#[rustc_regions]`,
    // we dump out various bits of state as warnings. This is useful
    // for verifying that the compiler is behaving as expected.  These
    // warnings focus on the closure region requirements -- for
    // viewing the intraprocedural state, the -Zdump-mir output is
    // better.

    if let Some(closure_region_requirements) = closure_region_requirements {
        let mut err = tcx.sess
            .diagnostic()
            .span_note_diag(mir.span, "External requirements");

        regioncx.annotate(&mut err);

        err.note(&format!(
            "number of external vids: {}",
            closure_region_requirements.num_external_vids
        ));

        // Dump the region constraints we are imposing *between* those
        // newly created variables.
        for_each_region_constraint(closure_region_requirements, &mut |msg| {
            err.note(msg);
            Ok(())
        }).unwrap();

        err.emit();
    } else {
        let mut err = tcx.sess
            .diagnostic()
            .span_note_diag(mir.span, "No external requirements");
        regioncx.annotate(&mut err);
        err.emit();
    }
}

fn for_each_region_constraint(
    closure_region_requirements: &ClosureRegionRequirements,
    with_msg: &mut dyn FnMut(&str) -> io::Result<()>,
) -> io::Result<()> {
    for req in &closure_region_requirements.outlives_requirements {
        let subject: &dyn Debug = match &req.subject {
            ClosureOutlivesSubject::Region(subject) => subject,
            ClosureOutlivesSubject::Ty(ty) => ty,
        };
        with_msg(&format!(
            "where {:?}: {:?}",
            subject,
            req.outlived_free_region,
        ))?;
    }
    Ok(())
}

/// Right now, we piggy back on the `ReVar` to store our NLL inference
/// regions. These are indexed with `RegionVid`. This method will
/// assert that the region is a `ReVar` and extract its internal index.
/// This is reasonable because in our MIR we replace all universal regions
/// with inference variables.
pub trait ToRegionVid {
    fn to_region_vid(self) -> RegionVid;
}

impl<'tcx> ToRegionVid for &'tcx RegionKind {
    fn to_region_vid(self) -> RegionVid {
        if let ty::ReVar(vid) = self {
            *vid
        } else {
            bug!("region is not an ReVar: {:?}", self)
        }
    }
}

impl ToRegionVid for RegionVid {
    fn to_region_vid(self) -> RegionVid {
        self
    }
}

fn live_variable_set(regular: &LocalSet, drops: &LocalSet) -> String {
    // sort and deduplicate:
    let all_locals: BTreeSet<_> = regular.iter().chain(drops.iter()).collect();

    // construct a string with each local, including `(drop)` if it is
    // only dropped, versus a regular use.
    let mut string = String::new();
    for local in all_locals {
        string.push_str(&format!("{:?}", local));

        if !regular.contains(&local) {
            assert!(drops.contains(&local));
            string.push_str(" (drop)");
        }

        string.push_str(", ");
    }

    let len = if string.is_empty() {
        0
    } else {
        string.len() - 2
    };

    format!("[{}]", &string[..len])
}
<|endoftext|>