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

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OSRng;

#[cfg(unix)]
mod imp {
    use Rng;
    use reader::ReaderRng;
    use std::io::{IoResult, File};

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    ///
    /// This does not block.
    #[cfg(unix)]
    pub struct OSRng {
        inner: ReaderRng<File>
    }

    impl OSRng {
        /// Create a new `OSRng`.
        pub fn new() -> IoResult<OSRng> {
            let reader = try!(File::open(&Path::new("/dev/urandom")));
            let reader_rng = ReaderRng::new(reader);

            Ok(OSRng { inner: reader_rng })
        }
    }

    impl Rng for OSRng {
        fn next_u32(&mut self) -> u32 {
            self.inner.next_u32()
        }
        fn next_u64(&mut self) -> u64 {
            self.inner.next_u64()
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            self.inner.fill_bytes(v)
        }
    }
}

#[cfg(windows)]
mod imp {
    use Rng;
    use std::cast;
    use std::io::{IoResult, IoError};
    use std::libc::{c_ulong, DWORD, BYTE, LPCSTR, BOOL};
    use std::os;
    use std::rt::stack;

    type HCRYPTPROV = c_ulong;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    ///
    /// This does not block.
    pub struct OSRng {
        hcryptprov: HCRYPTPROV
    }

    static PROV_RSA_FULL: DWORD = 1;
    static CRYPT_SILENT: DWORD = 64;
    static CRYPT_VERIFYCONTEXT: DWORD = 0xF0000000;
    static NTE_BAD_SIGNATURE: DWORD = 0x80090006;

    extern "system" {
        fn CryptAcquireContextA(phProv: *mut HCRYPTPROV,
                                pszContainer: LPCSTR,
                                pszProvider: LPCSTR,
                                dwProvType: DWORD,
                                dwFlags: DWORD) -> BOOL;
        fn CryptGenRandom(hProv: HCRYPTPROV,
                          dwLen: DWORD,
                          pbBuffer: *mut BYTE) -> BOOL;
        fn CryptReleaseContext(hProv: HCRYPTPROV, dwFlags: DWORD) -> BOOL;
    }

    impl OSRng {
        /// Create a new `OSRng`.
        pub fn new() -> IoResult<OSRng> {
            let mut hcp = 0;
            let mut ret = unsafe {
                CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            // It turns out that if we can't acquire a context with the
            // NTE_BAD_SIGNATURE error code, the documentation states:
            //
            //     The provider DLL signature could not be verified. Either the
            //     DLL or the digital signature has been tampered with.
            //
            // Sounds fishy, no? As it turns out, our signature can be bad
            // because our Thread Information Block (TIB) isn't exactly what it
            // expects. As to why, I have no idea. The only data we store in the
            // TIB is the stack limit for each thread, but apparently that's
            // enough to make the signature valid.
            //
            // Furthermore, this error only happens the *first* time we call
            // CryptAcquireContext, so we don't have to worry about future
            // calls.
            //
            // Anyway, the fix employed here is that if we see this error, we
            // pray that we're not close to the end of the stack, temporarily
            // set the stack limit to 0 (what the TIB originally was), acquire a
            // context, and then reset the stack limit.
            //
            // Again, I'm not sure why this is the fix, nor why we're getting
            // this error. All I can say is that this seems to allow libnative
            // to progress where it otherwise would be hindered. Who knew?
            if ret == 0 && os::errno() as DWORD == NTE_BAD_SIGNATURE {
                unsafe {
                    let limit = stack::get_sp_limit();
                    stack::record_sp_limit(0);
                    ret = CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                               PROV_RSA_FULL,
                                               CRYPT_VERIFYCONTEXT | CRYPT_SILENT);
                    stack::record_sp_limit(limit);
                }
            }

            if ret == 0 {
                Err(IoError::last_error())
            } else {
                Ok(OSRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OSRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(v);
            unsafe { cast::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(v);
            unsafe { cast::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                fail!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }

    impl Drop for OSRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                fail!("couldn't release context: {}", os::last_os_error());
            }
        }
    }
}

#[cfg(test)]
mod test {
    use super::OSRng;
    use Rng;
    use std::task;

    #[test]
    fn test_os_rng() {
        let mut r = OSRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8, .. 1000];
        r.fill_bytes(v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in range(0, 20) {
            let (tx, rx) = channel();
            txs.push(tx);
            task::spawn(proc() {
                // wait until all the tasks are ready to go.
                rx.recv();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OSRng::new().unwrap();
                task::deschedule();
                let mut v = [0u8, .. 1000];

                for _ in range(0, 100) {
                    r.next_u32();
                    task::deschedule();
                    r.next_u64();
                    task::deschedule();
                    r.fill_bytes(v);
                    task::deschedule();
                }
            })
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(())
        }
    }
}
<commit_msg>auto merge of #13315 : alexcrichton/rust/libc, r=alexcrichton,me<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OSRng;

#[cfg(unix)]
mod imp {
    use Rng;
    use reader::ReaderRng;
    use std::io::{IoResult, File};

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    ///
    /// This does not block.
    #[cfg(unix)]
    pub struct OSRng {
        inner: ReaderRng<File>
    }

    impl OSRng {
        /// Create a new `OSRng`.
        pub fn new() -> IoResult<OSRng> {
            let reader = try!(File::open(&Path::new("/dev/urandom")));
            let reader_rng = ReaderRng::new(reader);

            Ok(OSRng { inner: reader_rng })
        }
    }

    impl Rng for OSRng {
        fn next_u32(&mut self) -> u32 {
            self.inner.next_u32()
        }
        fn next_u64(&mut self) -> u64 {
            self.inner.next_u64()
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            self.inner.fill_bytes(v)
        }
    }
}

#[cfg(windows)]
mod imp {
    extern crate libc;

    use Rng;
    use std::cast;
    use std::io::{IoResult, IoError};
    use std::os;
    use std::rt::stack;
    use self::libc::{c_ulong, DWORD, BYTE, LPCSTR, BOOL};

    type HCRYPTPROV = c_ulong;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    ///
    /// This does not block.
    pub struct OSRng {
        hcryptprov: HCRYPTPROV
    }

    static PROV_RSA_FULL: DWORD = 1;
    static CRYPT_SILENT: DWORD = 64;
    static CRYPT_VERIFYCONTEXT: DWORD = 0xF0000000;
    static NTE_BAD_SIGNATURE: DWORD = 0x80090006;

    extern "system" {
        fn CryptAcquireContextA(phProv: *mut HCRYPTPROV,
                                pszContainer: LPCSTR,
                                pszProvider: LPCSTR,
                                dwProvType: DWORD,
                                dwFlags: DWORD) -> BOOL;
        fn CryptGenRandom(hProv: HCRYPTPROV,
                          dwLen: DWORD,
                          pbBuffer: *mut BYTE) -> BOOL;
        fn CryptReleaseContext(hProv: HCRYPTPROV, dwFlags: DWORD) -> BOOL;
    }

    impl OSRng {
        /// Create a new `OSRng`.
        pub fn new() -> IoResult<OSRng> {
            let mut hcp = 0;
            let mut ret = unsafe {
                CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            // It turns out that if we can't acquire a context with the
            // NTE_BAD_SIGNATURE error code, the documentation states:
            //
            //     The provider DLL signature could not be verified. Either the
            //     DLL or the digital signature has been tampered with.
            //
            // Sounds fishy, no? As it turns out, our signature can be bad
            // because our Thread Information Block (TIB) isn't exactly what it
            // expects. As to why, I have no idea. The only data we store in the
            // TIB is the stack limit for each thread, but apparently that's
            // enough to make the signature valid.
            //
            // Furthermore, this error only happens the *first* time we call
            // CryptAcquireContext, so we don't have to worry about future
            // calls.
            //
            // Anyway, the fix employed here is that if we see this error, we
            // pray that we're not close to the end of the stack, temporarily
            // set the stack limit to 0 (what the TIB originally was), acquire a
            // context, and then reset the stack limit.
            //
            // Again, I'm not sure why this is the fix, nor why we're getting
            // this error. All I can say is that this seems to allow libnative
            // to progress where it otherwise would be hindered. Who knew?
            if ret == 0 && os::errno() as DWORD == NTE_BAD_SIGNATURE {
                unsafe {
                    let limit = stack::get_sp_limit();
                    stack::record_sp_limit(0);
                    ret = CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                               PROV_RSA_FULL,
                                               CRYPT_VERIFYCONTEXT | CRYPT_SILENT);
                    stack::record_sp_limit(limit);
                }
            }

            if ret == 0 {
                Err(IoError::last_error())
            } else {
                Ok(OSRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OSRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(v);
            unsafe { cast::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(v);
            unsafe { cast::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                fail!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }

    impl Drop for OSRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                fail!("couldn't release context: {}", os::last_os_error());
            }
        }
    }
}

#[cfg(test)]
mod test {
    use super::OSRng;
    use Rng;
    use std::task;

    #[test]
    fn test_os_rng() {
        let mut r = OSRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8, .. 1000];
        r.fill_bytes(v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in range(0, 20) {
            let (tx, rx) = channel();
            txs.push(tx);
            task::spawn(proc() {
                // wait until all the tasks are ready to go.
                rx.recv();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OSRng::new().unwrap();
                task::deschedule();
                let mut v = [0u8, .. 1000];

                for _ in range(0, 100) {
                    r.next_u32();
                    task::deschedule();
                    r.next_u64();
                    task::deschedule();
                    r.fill_bytes(v);
                    task::deschedule();
                }
            })
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(())
        }
    }
}
<|endoftext|>


<commit_before>use libc::{gid_t, pid_t, uid_t};

/// Credentials of a process
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
pub struct UCred {
    /// PID (process ID) of the process
    pid: Option<pid_t>,
    /// UID (user ID) of the process
    uid: uid_t,
    /// GID (group ID) of the process
    gid: gid_t,
}

impl UCred {
    /// Gets UID (user ID) of the process.
    pub fn uid(&self) -> uid_t {
        self.uid
    }

    /// Gets GID (group ID) of the process.
    pub fn gid(&self) -> gid_t {
        self.gid
    }

    /// Gets PID (process ID) of the process.
    ///
    /// This is only implemented under Linux, Android, iOS, macOS, Solaris and
    /// Illumos. On other platforms this will always return `None`.
    pub fn pid(&self) -> Option<pid_t> {
        self.pid
    }
}

#[cfg(any(target_os = "linux", target_os = "android"))]
pub(crate) use self::impl_linux::get_peer_cred;

#[cfg(any(
    target_os = "dragonfly",
    target_os = "freebsd",
    target_os = "netbsd",
    target_os = "openbsd"
))]
pub(crate) use self::impl_bsd::get_peer_cred;

#[cfg(any(target_os = "macos", target_os = "ios"))]
pub(crate) use self::impl_macos::get_peer_cred;

#[cfg(any(target_os = "solaris", target_os = "illumos"))]
pub(crate) use self::impl_solaris::get_peer_cred;

#[cfg(any(target_os = "linux", target_os = "android"))]
pub(crate) mod impl_linux {
    use crate::net::unix::UnixStream;

    use libc::{c_void, getsockopt, socklen_t, SOL_SOCKET, SO_PEERCRED};
    use std::{io, mem};

    use libc::ucred;

    pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result<super::UCred> {
        use std::os::unix::io::AsRawFd;

        unsafe {
            let raw_fd = sock.as_raw_fd();

            let mut ucred = ucred {
                pid: 0,
                uid: 0,
                gid: 0,
            };

            let ucred_size = mem::size_of::<ucred>();

            // These paranoid checks should be optimized-out
            assert!(mem::size_of::<u32>() <= mem::size_of::<usize>());
            assert!(ucred_size <= u32::MAX as usize);

            let mut ucred_size = ucred_size as socklen_t;

            let ret = getsockopt(
                raw_fd,
                SOL_SOCKET,
                SO_PEERCRED,
                &mut ucred as *mut ucred as *mut c_void,
                &mut ucred_size,
            );
            if ret == 0 && ucred_size as usize == mem::size_of::<ucred>() {
                Ok(super::UCred {
                    uid: ucred.uid,
                    gid: ucred.gid,
                    pid: Some(ucred.pid),
                })
            } else {
                Err(io::Error::last_os_error())
            }
        }
    }
}

#[cfg(any(
    target_os = "dragonfly",
    target_os = "freebsd",
    target_os = "netbsd",
    target_os = "openbsd"
))]
pub(crate) mod impl_bsd {
    use crate::net::unix::UnixStream;

    use libc::getpeereid;
    use std::io;
    use std::mem::MaybeUninit;
    use std::os::unix::io::AsRawFd;

    pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result<super::UCred> {
        unsafe {
            let raw_fd = sock.as_raw_fd();

            let mut uid = MaybeUninit::uninit();
            let mut gid = MaybeUninit::uninit();

            let ret = getpeereid(raw_fd, uid.as_mut_ptr(), gid.as_mut_ptr());

            if ret == 0 {
                Ok(super::UCred {
                    uid: uid.assume_init(),
                    gid: gid.assume_init(),
                    pid: None,
                })
            } else {
                Err(io::Error::last_os_error())
            }
        }
    }
}

#[cfg(any(target_os = "macos", target_os = "ios"))]
pub(crate) mod impl_macos {
    use crate::net::unix::UnixStream;

    use libc::{c_void, getpeereid, getsockopt, pid_t, LOCAL_PEEREPID, SOL_LOCAL};
    use std::io;
    use std::mem::size_of;
    use std::mem::MaybeUninit;
    use std::os::unix::io::AsRawFd;

    pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result<super::UCred> {
        unsafe {
            let raw_fd = sock.as_raw_fd();

            let mut uid = MaybeUninit::uninit();
            let mut gid = MaybeUninit::uninit();
            let mut pid: MaybeUninit<pid_t> = MaybeUninit::uninit();
            let mut pid_size: MaybeUninit<u32> = MaybeUninit::new(size_of::<pid_t>() as u32);

            if getsockopt(
                raw_fd,
                SOL_LOCAL,
                LOCAL_PEEREPID,
                pid.as_mut_ptr() as *mut c_void,
                pid_size.as_mut_ptr(),
            ) != 0
            {
                return Err(io::Error::last_os_error());
            }

            assert!(pid_size.assume_init() == (size_of::<pid_t>() as u32));

            let ret = getpeereid(raw_fd, uid.as_mut_ptr(), gid.as_mut_ptr());

            if ret == 0 {
                Ok(super::UCred {
                    uid: uid.assume_init(),
                    gid: gid.assume_init(),
                    pid: Some(pid.assume_init()),
                })
            } else {
                Err(io::Error::last_os_error())
            }
        }
    }
}

#[cfg(any(target_os = "solaris", target_os = "illumos"))]
pub(crate) mod impl_solaris {
    use crate::net::unix::UnixStream;
    use std::io;
    use std::os::unix::io::AsRawFd;
    use std::ptr;

    pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result<super::UCred> {
        unsafe {
            let raw_fd = sock.as_raw_fd();

            let mut cred = ptr::null_mut();
            let ret = libc::getpeerucred(raw_fd, &mut cred);

            if ret == 0 {
                let uid = libc::ucred_geteuid(cred);
                let gid = libc::ucred_getegid(cred);
                let pid = libc::ucred_getpid(cred);

                libc::ucred_free(cred);

                Ok(super::UCred {
                    uid,
                    gid,
                    pid: Some(pid),
                })
            } else {
                Err(io::Error::last_os_error())
            }
        }
    }
}
<commit_msg>net: make ucred return pid field for OpenBSD (#3919)<commit_after>use libc::{gid_t, pid_t, uid_t};

/// Credentials of a process
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
pub struct UCred {
    /// PID (process ID) of the process
    pid: Option<pid_t>,
    /// UID (user ID) of the process
    uid: uid_t,
    /// GID (group ID) of the process
    gid: gid_t,
}

impl UCred {
    /// Gets UID (user ID) of the process.
    pub fn uid(&self) -> uid_t {
        self.uid
    }

    /// Gets GID (group ID) of the process.
    pub fn gid(&self) -> gid_t {
        self.gid
    }

    /// Gets PID (process ID) of the process.
    ///
    /// This is only implemented under Linux, Android, iOS, macOS, Solaris and
    /// Illumos. On other platforms this will always return `None`.
    pub fn pid(&self) -> Option<pid_t> {
        self.pid
    }
}

#[cfg(any(target_os = "linux", target_os = "android", target_os = "openbsd"))]
pub(crate) use self::impl_linux::get_peer_cred;

#[cfg(any(target_os = "netbsd"))]
pub(crate) use self::impl_netbsd::get_peer_cred;

#[cfg(any(target_os = "dragonfly", target_os = "freebsd"))]
pub(crate) use self::impl_bsd::get_peer_cred;

#[cfg(any(target_os = "macos", target_os = "ios"))]
pub(crate) use self::impl_macos::get_peer_cred;

#[cfg(any(target_os = "solaris", target_os = "illumos"))]
pub(crate) use self::impl_solaris::get_peer_cred;

#[cfg(any(target_os = "linux", target_os = "android", target_os = "openbsd"))]
pub(crate) mod impl_linux {
    use crate::net::unix::UnixStream;

    use libc::{c_void, getsockopt, socklen_t, SOL_SOCKET, SO_PEERCRED};
    use std::{io, mem};

    #[cfg(target_os = "openbsd")]
    use libc::sockpeercred as ucred;
    #[cfg(any(target_os = "linux", target_os = "android"))]
    use libc::ucred;

    pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result<super::UCred> {
        use std::os::unix::io::AsRawFd;

        unsafe {
            let raw_fd = sock.as_raw_fd();

            let mut ucred = ucred {
                pid: 0,
                uid: 0,
                gid: 0,
            };

            let ucred_size = mem::size_of::<ucred>();

            // These paranoid checks should be optimized-out
            assert!(mem::size_of::<u32>() <= mem::size_of::<usize>());
            assert!(ucred_size <= u32::MAX as usize);

            let mut ucred_size = ucred_size as socklen_t;

            let ret = getsockopt(
                raw_fd,
                SOL_SOCKET,
                SO_PEERCRED,
                &mut ucred as *mut ucred as *mut c_void,
                &mut ucred_size,
            );
            if ret == 0 && ucred_size as usize == mem::size_of::<ucred>() {
                Ok(super::UCred {
                    uid: ucred.uid,
                    gid: ucred.gid,
                    pid: Some(ucred.pid),
                })
            } else {
                Err(io::Error::last_os_error())
            }
        }
    }
}

#[cfg(any(target_os = "netbsd"))]
pub(crate) mod impl_netbsd {
    use crate::net::unix::UnixStream;

    use libc::{c_void, getsockopt, socklen_t, unpcbid, LOCAL_PEEREID, SOL_SOCKET};
    use std::io;
    use std::mem::size_of;
    use std::os::unix::io::AsRawFd;

    pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result<super::UCred> {
        unsafe {
            let raw_fd = sock.as_raw_fd();

            let mut unpcbid = unpcbid {
                unp_pid: 0,
                unp_euid: 0,
                unp_egid: 0,
            };

            let unpcbid_size = size_of::<unpcbid>();
            let mut unpcbid_size = unpcbid_size as socklen_t;

            let ret = getsockopt(
                raw_fd,
                SOL_SOCKET,
                LOCAL_PEEREID,
                &mut unpcbid as *mut unpcbid as *mut c_void,
                &mut unpcbid_size,
            );
            if ret == 0 && unpcbid_size as usize == size_of::<unpcbid>() {
                Ok(super::UCred {
                    uid: unpcbid.unp_euid,
                    gid: unpcbid.unp_egid,
                    pid: Some(unpcbid.unp_pid),
                })
            } else {
                Err(io::Error::last_os_error())
            }
        }
    }
}

#[cfg(any(target_os = "dragonfly", target_os = "freebsd"))]
pub(crate) mod impl_bsd {
    use crate::net::unix::UnixStream;

    use libc::getpeereid;
    use std::io;
    use std::mem::MaybeUninit;
    use std::os::unix::io::AsRawFd;

    pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result<super::UCred> {
        unsafe {
            let raw_fd = sock.as_raw_fd();

            let mut uid = MaybeUninit::uninit();
            let mut gid = MaybeUninit::uninit();

            let ret = getpeereid(raw_fd, uid.as_mut_ptr(), gid.as_mut_ptr());

            if ret == 0 {
                Ok(super::UCred {
                    uid: uid.assume_init(),
                    gid: gid.assume_init(),
                    pid: None,
                })
            } else {
                Err(io::Error::last_os_error())
            }
        }
    }
}

#[cfg(any(target_os = "macos", target_os = "ios"))]
pub(crate) mod impl_macos {
    use crate::net::unix::UnixStream;

    use libc::{c_void, getpeereid, getsockopt, pid_t, LOCAL_PEEREPID, SOL_LOCAL};
    use std::io;
    use std::mem::size_of;
    use std::mem::MaybeUninit;
    use std::os::unix::io::AsRawFd;

    pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result<super::UCred> {
        unsafe {
            let raw_fd = sock.as_raw_fd();

            let mut uid = MaybeUninit::uninit();
            let mut gid = MaybeUninit::uninit();
            let mut pid: MaybeUninit<pid_t> = MaybeUninit::uninit();
            let mut pid_size: MaybeUninit<u32> = MaybeUninit::new(size_of::<pid_t>() as u32);

            if getsockopt(
                raw_fd,
                SOL_LOCAL,
                LOCAL_PEEREPID,
                pid.as_mut_ptr() as *mut c_void,
                pid_size.as_mut_ptr(),
            ) != 0
            {
                return Err(io::Error::last_os_error());
            }

            assert!(pid_size.assume_init() == (size_of::<pid_t>() as u32));

            let ret = getpeereid(raw_fd, uid.as_mut_ptr(), gid.as_mut_ptr());

            if ret == 0 {
                Ok(super::UCred {
                    uid: uid.assume_init(),
                    gid: gid.assume_init(),
                    pid: Some(pid.assume_init()),
                })
            } else {
                Err(io::Error::last_os_error())
            }
        }
    }
}

#[cfg(any(target_os = "solaris", target_os = "illumos"))]
pub(crate) mod impl_solaris {
    use crate::net::unix::UnixStream;
    use std::io;
    use std::os::unix::io::AsRawFd;
    use std::ptr;

    pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result<super::UCred> {
        unsafe {
            let raw_fd = sock.as_raw_fd();

            let mut cred = ptr::null_mut();
            let ret = libc::getpeerucred(raw_fd, &mut cred);

            if ret == 0 {
                let uid = libc::ucred_geteuid(cred);
                let gid = libc::ucred_getegid(cred);
                let pid = libc::ucred_getpid(cred);

                libc::ucred_free(cred);

                Ok(super::UCred {
                    uid,
                    gid,
                    pid: Some(pid),
                })
            } else {
                Err(io::Error::last_os_error())
            }
        }
    }
}
<|endoftext|>


<commit_before>extern crate hyper;
extern crate env_logger;
extern crate tools;

use std::sync::Mutex;
use std::fs::File;
use std::io::BufReader;
use std::io::Read;
use std::path::Path;
use std::env;

use hyper::status::StatusCode;
use hyper::method::Method;
use hyper::server::{Request, Response};
use hyper::header::ContentType;
use hyper::mime::Mime;
use hyper::uri;

use tools::config;
use tools::blame;
use tools::format;

struct WebRequest {
    path: String,
}

struct WebResponse {
    status: StatusCode,
    content_type: String,
    output: String,
}

fn not_found() -> WebResponse {
    WebResponse {
        status: StatusCode::NotFound,
        content_type: "text/plain".to_owned(),
        output: "Not found".to_owned()
    }
}

fn handle_static(path: String, content_type: Option<&str>) -> WebResponse {
    let source_file = match File::open(&path) {
        Ok(f) => f,
        Err(_) => {
            return not_found();
        },
    };
    let mut reader = BufReader::new(&source_file);
    let mut input = String::new();
    match reader.read_to_string(&mut input) {
        Ok(_) => {},
        Err(_) => {
            return not_found();
        }
    }

    let inferred_content_type = match Path::new(&path).extension() {
        Some(ext) =>
            match ext.to_str().unwrap() {
                "css" => "text/css",
                "js" => "text/javascript",
                _ => "text/html",
            },
        None => "text/html"
    };
    let content_type = match content_type {
        Some(ct) => ct,
        None => inferred_content_type,
    };

    WebResponse { status: StatusCode::Ok, content_type: content_type.to_owned(), output: input }
}

fn handle(cfg: &config::Config, req: WebRequest) -> WebResponse {
    let path = req.path.clone();
    let path = path[1..].split('/').collect::<Vec<_>>();

    if path.len() > 0 && path[0] == "static" {
        let path = cfg.mozsearch_path.clone() + &req.path;
        return handle_static(path, None);
    }

    if path.len() < 2 {
        return not_found();
    }

    let tree_name = &path[0];
    let kind = &path[1];

    println!("DBG {:?} {} {}", path, tree_name, kind);

    match &kind[..] {
        "rev" => {
            if path.len() < 3 {
                return not_found();
            }

            let rev = &path[2];
            let path = path.clone().split_off(3);
            let path = path.join("/");

            let mut writer = Vec::new();
            match format::format_path(cfg, &tree_name, &rev, &path, &mut writer) {
                Ok(()) => {
                    let output = String::from_utf8(writer).unwrap();
                    WebResponse { status: StatusCode::Ok, content_type: "text/html".to_owned(), output: output }
                },
                Err(err) =>
                    WebResponse {
                        status: StatusCode::InternalServerError,
                        content_type: "text/plain".to_owned(),
                        output: err.to_owned(),
                    }
            }
        },

        "source" => {
            let path = path.clone().split_off(2);
            let path = path.join("/");

            let tree_config = cfg.trees.get(*tree_name).unwrap();

            let path = format!("{}/file/{}", tree_config.paths.index_path, path);
            return handle_static(path, Some("text/html"));
        },

        "diff" => {
            if path.len() < 3 {
                return not_found();
            }

            let rev = &path[2];
            let path = path.clone().split_off(3);
            let path = path.join("/");

            let mut writer = Vec::new();
            match format::format_diff(cfg, &tree_name, &rev, &path, &mut writer) {
                Ok(()) => {
                    let output = String::from_utf8(writer).unwrap();
                    WebResponse { status: StatusCode::Ok, content_type: "text/html".to_owned(), output: output }
                },
                Err(err) =>
                    WebResponse {
                        status: StatusCode::InternalServerError,
                        content_type: "text/plain".to_owned(),
                        output: err.to_owned(),
                    }
            }
        },

        "commit" => {
            if path.len() < 3 {
                return not_found();
            }

            let rev = &path[2];

            let mut writer = Vec::new();
            match format::format_commit(cfg, &tree_name, &rev, &mut writer) {
                Ok(()) => {
                    let output = String::from_utf8(writer).unwrap();
                    WebResponse { status: StatusCode::Ok, content_type: "text/html".to_owned(), output: output }
                },
                Err(err) =>
                    WebResponse {
                        status: StatusCode::InternalServerError,
                        content_type: "text/plain".to_owned(),
                        output: err.to_owned(),
                    }
            }
        },

        "commit-info" => {
            if path.len() < 3 {
                return not_found();
            }

            let rev = &path[2];
            match blame::get_commit_info(&cfg, tree_name, rev) {
                Ok(json) =>
                    WebResponse {
                        status: StatusCode::Ok,
                        content_type: "application/json".to_owned(),
                        output: json
                    },
                Err(err) =>
                    WebResponse {
                        status: StatusCode::InternalServerError,
                        content_type: "text/plain".to_owned(),
                        output: err.to_owned(),
                    }
            }
        },

        _ => {
            not_found()
        }
    }
}

fn main() {
    env_logger::init().unwrap();

    let cfg = config::load(&env::args().nth(1).unwrap(), true);
    let cfg = Mutex::new(cfg);

    let handler = move |req: Request, mut res: Response| {
        if req.method != Method::Get {
            *res.status_mut() = StatusCode::MethodNotAllowed;
            let resp = format!("Invalid method").into_bytes();
            res.send(&resp).unwrap();
            return;
        }

        let path = match req.uri {
            uri::RequestUri::AbsolutePath(path) => path,
            uri::RequestUri::AbsoluteUri(url) => url.path().to_owned(),
            _ => panic!("Unexpected URI"),
        };

        let guard = match cfg.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        };
        let ref cfg = *guard;

        let response = handle(cfg, WebRequest { path: path });

        *res.status_mut() = response.status;
        let output = response.output.into_bytes();
        let mime: Mime = response.content_type.parse().unwrap();
        res.headers_mut().set(ContentType(mime));
        res.send(&output).unwrap();
    };

    println!("On 8001");
    let _listening = hyper::Server::http("127.0.0.1:8001").unwrap().handle(handler);
}
<commit_msg>Remove extra whitespace<commit_after>extern crate hyper;
extern crate env_logger;
extern crate tools;

use std::sync::Mutex;
use std::fs::File;
use std::io::BufReader;
use std::io::Read;
use std::path::Path;
use std::env;

use hyper::status::StatusCode;
use hyper::method::Method;
use hyper::server::{Request, Response};
use hyper::header::ContentType;
use hyper::mime::Mime;
use hyper::uri;

use tools::config;
use tools::blame;
use tools::format;

struct WebRequest {
    path: String,
}

struct WebResponse {
    status: StatusCode,
    content_type: String,
    output: String,
}

fn not_found() -> WebResponse {
    WebResponse {
        status: StatusCode::NotFound,
        content_type: "text/plain".to_owned(),
        output: "Not found".to_owned()
    }
}

fn handle_static(path: String, content_type: Option<&str>) -> WebResponse {
    let source_file = match File::open(&path) {
        Ok(f) => f,
        Err(_) => {
            return not_found();
        },
    };
    let mut reader = BufReader::new(&source_file);
    let mut input = String::new();
    match reader.read_to_string(&mut input) {
        Ok(_) => {},
        Err(_) => {
            return not_found();
        }
    }

    let inferred_content_type = match Path::new(&path).extension() {
        Some(ext) =>
            match ext.to_str().unwrap() {
                "css" => "text/css",
                "js" => "text/javascript",
                _ => "text/html",
            },
        None => "text/html"
    };
    let content_type = match content_type {
        Some(ct) => ct,
        None => inferred_content_type,
    };

    WebResponse { status: StatusCode::Ok, content_type: content_type.to_owned(), output: input }
}

fn handle(cfg: &config::Config, req: WebRequest) -> WebResponse {
    let path = req.path.clone();
    let path = path[1..].split('/').collect::<Vec<_>>();

    if path.len() > 0 && path[0] == "static" {
        let path = cfg.mozsearch_path.clone() + &req.path;
        return handle_static(path, None);
    }

    if path.len() < 2 {
        return not_found();
    }

    let tree_name = &path[0];
    let kind = &path[1];

    println!("DBG {:?} {} {}", path, tree_name, kind);

    match &kind[..] {
        "rev" => {
            if path.len() < 3 {
                return not_found();
            }

            let rev = &path[2];
            let path = path.clone().split_off(3);
            let path = path.join("/");

            let mut writer = Vec::new();
            match format::format_path(cfg, &tree_name, &rev, &path, &mut writer) {
                Ok(()) => {
                    let output = String::from_utf8(writer).unwrap();
                    WebResponse { status: StatusCode::Ok, content_type: "text/html".to_owned(), output: output }
                },
                Err(err) =>
                    WebResponse {
                        status: StatusCode::InternalServerError,
                        content_type: "text/plain".to_owned(),
                        output: err.to_owned(),
                    }
            }
        },

        "source" => {
            let path = path.clone().split_off(2);
            let path = path.join("/");

            let tree_config = cfg.trees.get(*tree_name).unwrap();

            let path = format!("{}/file/{}", tree_config.paths.index_path, path);
            return handle_static(path, Some("text/html"));
        },

        "diff" => {
            if path.len() < 3 {
                return not_found();
            }

            let rev = &path[2];
            let path = path.clone().split_off(3);
            let path = path.join("/");

            let mut writer = Vec::new();
            match format::format_diff(cfg, &tree_name, &rev, &path, &mut writer) {
                Ok(()) => {
                    let output = String::from_utf8(writer).unwrap();
                    WebResponse { status: StatusCode::Ok, content_type: "text/html".to_owned(), output: output }
                },
                Err(err) =>
                    WebResponse {
                        status: StatusCode::InternalServerError,
                        content_type: "text/plain".to_owned(),
                        output: err.to_owned(),
                    }
            }
        },

        "commit" => {
            if path.len() < 3 {
                return not_found();
            }

            let rev = &path[2];

            let mut writer = Vec::new();
            match format::format_commit(cfg, &tree_name, &rev, &mut writer) {
                Ok(()) => {
                    let output = String::from_utf8(writer).unwrap();
                    WebResponse { status: StatusCode::Ok, content_type: "text/html".to_owned(), output: output }
                },
                Err(err) =>
                    WebResponse {
                        status: StatusCode::InternalServerError,
                        content_type: "text/plain".to_owned(),
                        output: err.to_owned(),
                    }
            }
        },

        "commit-info" => {
            if path.len() < 3 {
                return not_found();
            }

            let rev = &path[2];
            match blame::get_commit_info(&cfg, tree_name, rev) {
                Ok(json) =>
                    WebResponse {
                        status: StatusCode::Ok,
                        content_type: "application/json".to_owned(),
                        output: json
                    },
                Err(err) =>
                    WebResponse {
                        status: StatusCode::InternalServerError,
                        content_type: "text/plain".to_owned(),
                        output: err.to_owned(),
                    }
            }
        },

        _ => {
            not_found()
        }
    }
}

fn main() {
    env_logger::init().unwrap();

    let cfg = config::load(&env::args().nth(1).unwrap(), true);
    let cfg = Mutex::new(cfg);

    let handler = move |req: Request, mut res: Response| {
        if req.method != Method::Get {
            *res.status_mut() = StatusCode::MethodNotAllowed;
            let resp = format!("Invalid method").into_bytes();
            res.send(&resp).unwrap();
            return;
        }

        let path = match req.uri {
            uri::RequestUri::AbsolutePath(path) => path,
            uri::RequestUri::AbsoluteUri(url) => url.path().to_owned(),
            _ => panic!("Unexpected URI"),
        };

        let guard = match cfg.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        };
        let ref cfg = *guard;

        let response = handle(cfg, WebRequest { path: path });

        *res.status_mut() = response.status;
        let output = response.output.into_bytes();
        let mime: Mime = response.content_type.parse().unwrap();
        res.headers_mut().set(ContentType(mime));
        res.send(&output).unwrap();
    };

    println!("On 8001");
    let _listening = hyper::Server::http("127.0.0.1:8001").unwrap().handle(handler);
}
<|endoftext|>


<commit_before>//! Name resolution for expressions.
use std::sync::Arc;

use hir_expand::name::Name;
use ra_arena::{Arena, Idx};
use rustc_hash::FxHashMap;

use crate::{
    body::Body,
    db::DefDatabase,
    expr::{Expr, ExprId, Pat, PatId, Statement},
    DefWithBodyId,
};

pub type ScopeId = Idx<ScopeData>;

#[derive(Debug, PartialEq, Eq)]
pub struct ExprScopes {
    scopes: Arena<ScopeData>,
    scope_by_expr: FxHashMap<ExprId, ScopeId>,
}

#[derive(Debug, PartialEq, Eq)]
pub struct ScopeEntry {
    name: Name,
    pat: PatId,
}

impl ScopeEntry {
    pub fn name(&self) -> &Name {
        &self.name
    }

    pub fn pat(&self) -> PatId {
        self.pat
    }
}

#[derive(Debug, PartialEq, Eq)]
pub struct ScopeData {
    parent: Option<ScopeId>,
    entries: Vec<ScopeEntry>,
}

impl ExprScopes {
    pub(crate) fn expr_scopes_query(db: &dyn DefDatabase, def: DefWithBodyId) -> Arc<ExprScopes> {
        let body = db.body(def);
        Arc::new(ExprScopes::new(&*body))
    }

    fn new(body: &Body) -> ExprScopes {
        let mut scopes =
            ExprScopes { scopes: Arena::default(), scope_by_expr: FxHashMap::default() };
        let root = scopes.root_scope();
        scopes.add_params_bindings(body, root, &body.params);
        compute_expr_scopes(body.body_expr, body, &mut scopes, root);
        scopes
    }

    pub fn entries(&self, scope: ScopeId) -> &[ScopeEntry] {
        &self.scopes[scope].entries
    }

    pub fn scope_chain(&self, scope: Option<ScopeId>) -> impl Iterator<Item = ScopeId> + '_ {
        std::iter::successors(scope, move |&scope| self.scopes[scope].parent)
    }

    pub fn resolve_name_in_scope(&self, scope: ScopeId, name: &Name) -> Option<&ScopeEntry> {
        self.scope_chain(Some(scope))
            .find_map(|scope| self.entries(scope).iter().find(|it| it.name == *name))
    }

    pub fn scope_for(&self, expr: ExprId) -> Option<ScopeId> {
        self.scope_by_expr.get(&expr).copied()
    }

    pub fn scope_by_expr(&self) -> &FxHashMap<ExprId, ScopeId> {
        &self.scope_by_expr
    }

    fn root_scope(&mut self) -> ScopeId {
        self.scopes.alloc(ScopeData { parent: None, entries: vec![] })
    }

    fn new_scope(&mut self, parent: ScopeId) -> ScopeId {
        self.scopes.alloc(ScopeData { parent: Some(parent), entries: vec![] })
    }

    fn add_bindings(&mut self, body: &Body, scope: ScopeId, pat: PatId) {
        match &body[pat] {
            Pat::Bind { name, .. } => {
                // bind can have a sub pattern, but it's actually not allowed
                // to bind to things in there
                let entry = ScopeEntry { name: name.clone(), pat };
                self.scopes[scope].entries.push(entry)
            }
            p => p.walk_child_pats(|pat| self.add_bindings(body, scope, pat)),
        }
    }

    fn add_params_bindings(&mut self, body: &Body, scope: ScopeId, params: &[PatId]) {
        params.iter().for_each(|pat| self.add_bindings(body, scope, *pat));
    }

    fn set_scope(&mut self, node: ExprId, scope: ScopeId) {
        self.scope_by_expr.insert(node, scope);
    }
}

fn compute_block_scopes(
    statements: &[Statement],
    tail: Option<ExprId>,
    body: &Body,
    scopes: &mut ExprScopes,
    mut scope: ScopeId,
) {
    for stmt in statements {
        match stmt {
            Statement::Let { pat, initializer, .. } => {
                if let Some(expr) = initializer {
                    scopes.set_scope(*expr, scope);
                    compute_expr_scopes(*expr, body, scopes, scope);
                }
                scope = scopes.new_scope(scope);
                scopes.add_bindings(body, scope, *pat);
            }
            Statement::Expr(expr) => {
                scopes.set_scope(*expr, scope);
                compute_expr_scopes(*expr, body, scopes, scope);
            }
        }
    }
    if let Some(expr) = tail {
        compute_expr_scopes(expr, body, scopes, scope);
    }
}

fn compute_expr_scopes(expr: ExprId, body: &Body, scopes: &mut ExprScopes, scope: ScopeId) {
    scopes.set_scope(expr, scope);
    match &body[expr] {
        Expr::Block { statements, tail, .. } => {
            compute_block_scopes(&statements, *tail, body, scopes, scope);
        }
        Expr::For { iterable, pat, body: body_expr, .. } => {
            compute_expr_scopes(*iterable, body, scopes, scope);
            let scope = scopes.new_scope(scope);
            scopes.add_bindings(body, scope, *pat);
            compute_expr_scopes(*body_expr, body, scopes, scope);
        }
        Expr::Lambda { args, body: body_expr, .. } => {
            let scope = scopes.new_scope(scope);
            scopes.add_params_bindings(body, scope, &args);
            compute_expr_scopes(*body_expr, body, scopes, scope);
        }
        Expr::Match { expr, arms } => {
            compute_expr_scopes(*expr, body, scopes, scope);
            for arm in arms {
                let scope = scopes.new_scope(scope);
                scopes.add_bindings(body, scope, arm.pat);
                if let Some(guard) = arm.guard {
                    scopes.set_scope(guard, scope);
                    compute_expr_scopes(guard, body, scopes, scope);
                }
                scopes.set_scope(arm.expr, scope);
                compute_expr_scopes(arm.expr, body, scopes, scope);
            }
        }
        e => e.walk_child_exprs(|e| compute_expr_scopes(e, body, scopes, scope)),
    };
}

#[cfg(test)]
mod tests {
    use hir_expand::{name::AsName, InFile};
    use ra_db::{fixture::WithFixture, FileId, SourceDatabase};
    use ra_syntax::{algo::find_node_at_offset, ast, AstNode};
    use test_utils::{assert_eq_text, extract_offset, mark};

    use crate::{db::DefDatabase, test_db::TestDB, FunctionId, ModuleDefId};

    fn find_function(db: &TestDB, file_id: FileId) -> FunctionId {
        let krate = db.test_crate();
        let crate_def_map = db.crate_def_map(krate);

        let module = crate_def_map.modules_for_file(file_id).next().unwrap();
        let (_, def) = crate_def_map[module].scope.entries().next().unwrap();
        match def.take_values().unwrap() {
            ModuleDefId::FunctionId(it) => it,
            _ => panic!(),
        }
    }

    fn do_check(code: &str, expected: &[&str]) {
        let (off, code) = extract_offset(code);
        let code = {
            let mut buf = String::new();
            let off: usize = off.into();
            buf.push_str(&code[..off]);
            buf.push_str("marker");
            buf.push_str(&code[off..]);
            buf
        };

        let (db, file_id) = TestDB::with_single_file(&code);

        let file_syntax = db.parse(file_id).syntax_node();
        let marker: ast::PathExpr = find_node_at_offset(&file_syntax, off).unwrap();
        let function = find_function(&db, file_id);

        let scopes = db.expr_scopes(function.into());
        let (_body, source_map) = db.body_with_source_map(function.into());

        let expr_id = source_map
            .node_expr(InFile { file_id: file_id.into(), value: &marker.into() })
            .unwrap();
        let scope = scopes.scope_for(expr_id);

        let actual = scopes
            .scope_chain(scope)
            .flat_map(|scope| scopes.entries(scope))
            .map(|it| it.name().to_string())
            .collect::<Vec<_>>()
            .join("\n");
        let expected = expected.join("\n");
        assert_eq_text!(&expected, &actual);
    }

    #[test]
    fn test_lambda_scope() {
        do_check(
            r"
            fn quux(foo: i32) {
                let f = |bar, baz: i32| {
                    <|>
                };
            }",
            &["bar", "baz", "foo"],
        );
    }

    #[test]
    fn test_call_scope() {
        do_check(
            r"
            fn quux() {
                f(|x| <|> );
            }",
            &["x"],
        );
    }

    #[test]
    fn test_method_call_scope() {
        do_check(
            r"
            fn quux() {
                z.f(|x| <|> );
            }",
            &["x"],
        );
    }

    #[test]
    fn test_loop_scope() {
        do_check(
            r"
            fn quux() {
                loop {
                    let x = ();
                    <|>
                };
            }",
            &["x"],
        );
    }

    #[test]
    fn test_match() {
        do_check(
            r"
            fn quux() {
                match () {
                    Some(x) => {
                        <|>
                    }
                };
            }",
            &["x"],
        );
    }

    #[test]
    fn test_shadow_variable() {
        do_check(
            r"
            fn foo(x: String) {
                let x : &str = &x<|>;
            }",
            &["x"],
        );
    }

    fn do_check_local_name(code: &str, expected_offset: u32) {
        let (off, code) = extract_offset(code);

        let (db, file_id) = TestDB::with_single_file(&code);

        let file = db.parse(file_id).ok().unwrap();
        let expected_name = find_node_at_offset::<ast::Name>(file.syntax(), expected_offset.into())
            .expect("failed to find a name at the target offset");
        let name_ref: ast::NameRef = find_node_at_offset(file.syntax(), off).unwrap();

        let function = find_function(&db, file_id);

        let scopes = db.expr_scopes(function.into());
        let (_body, source_map) = db.body_with_source_map(function.into());

        let expr_scope = {
            let expr_ast = name_ref.syntax().ancestors().find_map(ast::Expr::cast).unwrap();
            let expr_id =
                source_map.node_expr(InFile { file_id: file_id.into(), value: &expr_ast }).unwrap();
            scopes.scope_for(expr_id).unwrap()
        };

        let resolved = scopes.resolve_name_in_scope(expr_scope, &name_ref.as_name()).unwrap();
        let pat_src = source_map.pat_syntax(resolved.pat()).unwrap();

        let local_name = pat_src.value.either(
            |it| it.syntax_node_ptr().to_node(file.syntax()),
            |it| it.syntax_node_ptr().to_node(file.syntax()),
        );
        assert_eq!(local_name.text_range(), expected_name.syntax().text_range());
    }

    #[test]
    fn test_resolve_local_name() {
        do_check_local_name(
            r#"
            fn foo(x: i32, y: u32) {
                {
                    let z = x * 2;
                }
                {
                    let t = x<|> * 3;
                }
            }"#,
            21,
        );
    }

    #[test]
    fn test_resolve_local_name_declaration() {
        do_check_local_name(
            r#"
            fn foo(x: String) {
                let x : &str = &x<|>;
            }"#,
            21,
        );
    }

    #[test]
    fn test_resolve_local_name_shadow() {
        do_check_local_name(
            r"
            fn foo(x: String) {
                let x : &str = &x;
                x<|>
            }
            ",
            53,
        );
    }

    #[test]
    fn ref_patterns_contribute_bindings() {
        do_check_local_name(
            r"
            fn foo() {
                if let Some(&from) = bar() {
                    from<|>;
                }
            }
            ",
            53,
        );
    }

    #[test]
    fn while_let_desugaring() {
        mark::check!(infer_resolve_while_let);
        do_check_local_name(
            r#"
fn test() {
    let foo: Option<f32> = None;
    while let Option::Some(spam) = foo {
        spam<|>
    }
}
"#,
            75,
        );
    }
}
<commit_msg>Handle bindings after @ in patterns<commit_after>//! Name resolution for expressions.
use std::sync::Arc;

use hir_expand::name::Name;
use ra_arena::{Arena, Idx};
use rustc_hash::FxHashMap;

use crate::{
    body::Body,
    db::DefDatabase,
    expr::{Expr, ExprId, Pat, PatId, Statement},
    DefWithBodyId,
};

pub type ScopeId = Idx<ScopeData>;

#[derive(Debug, PartialEq, Eq)]
pub struct ExprScopes {
    scopes: Arena<ScopeData>,
    scope_by_expr: FxHashMap<ExprId, ScopeId>,
}

#[derive(Debug, PartialEq, Eq)]
pub struct ScopeEntry {
    name: Name,
    pat: PatId,
}

impl ScopeEntry {
    pub fn name(&self) -> &Name {
        &self.name
    }

    pub fn pat(&self) -> PatId {
        self.pat
    }
}

#[derive(Debug, PartialEq, Eq)]
pub struct ScopeData {
    parent: Option<ScopeId>,
    entries: Vec<ScopeEntry>,
}

impl ExprScopes {
    pub(crate) fn expr_scopes_query(db: &dyn DefDatabase, def: DefWithBodyId) -> Arc<ExprScopes> {
        let body = db.body(def);
        Arc::new(ExprScopes::new(&*body))
    }

    fn new(body: &Body) -> ExprScopes {
        let mut scopes =
            ExprScopes { scopes: Arena::default(), scope_by_expr: FxHashMap::default() };
        let root = scopes.root_scope();
        scopes.add_params_bindings(body, root, &body.params);
        compute_expr_scopes(body.body_expr, body, &mut scopes, root);
        scopes
    }

    pub fn entries(&self, scope: ScopeId) -> &[ScopeEntry] {
        &self.scopes[scope].entries
    }

    pub fn scope_chain(&self, scope: Option<ScopeId>) -> impl Iterator<Item = ScopeId> + '_ {
        std::iter::successors(scope, move |&scope| self.scopes[scope].parent)
    }

    pub fn resolve_name_in_scope(&self, scope: ScopeId, name: &Name) -> Option<&ScopeEntry> {
        self.scope_chain(Some(scope))
            .find_map(|scope| self.entries(scope).iter().find(|it| it.name == *name))
    }

    pub fn scope_for(&self, expr: ExprId) -> Option<ScopeId> {
        self.scope_by_expr.get(&expr).copied()
    }

    pub fn scope_by_expr(&self) -> &FxHashMap<ExprId, ScopeId> {
        &self.scope_by_expr
    }

    fn root_scope(&mut self) -> ScopeId {
        self.scopes.alloc(ScopeData { parent: None, entries: vec![] })
    }

    fn new_scope(&mut self, parent: ScopeId) -> ScopeId {
        self.scopes.alloc(ScopeData { parent: Some(parent), entries: vec![] })
    }

    fn add_bindings(&mut self, body: &Body, scope: ScopeId, pat: PatId) {
        let pattern = &body[pat];
        if let Pat::Bind { name, .. } = pattern {
            let entry = ScopeEntry { name: name.clone(), pat };
            self.scopes[scope].entries.push(entry);
        }

        pattern.walk_child_pats(|pat| self.add_bindings(body, scope, pat));
    }

    fn add_params_bindings(&mut self, body: &Body, scope: ScopeId, params: &[PatId]) {
        params.iter().for_each(|pat| self.add_bindings(body, scope, *pat));
    }

    fn set_scope(&mut self, node: ExprId, scope: ScopeId) {
        self.scope_by_expr.insert(node, scope);
    }
}

fn compute_block_scopes(
    statements: &[Statement],
    tail: Option<ExprId>,
    body: &Body,
    scopes: &mut ExprScopes,
    mut scope: ScopeId,
) {
    for stmt in statements {
        match stmt {
            Statement::Let { pat, initializer, .. } => {
                if let Some(expr) = initializer {
                    scopes.set_scope(*expr, scope);
                    compute_expr_scopes(*expr, body, scopes, scope);
                }
                scope = scopes.new_scope(scope);
                scopes.add_bindings(body, scope, *pat);
            }
            Statement::Expr(expr) => {
                scopes.set_scope(*expr, scope);
                compute_expr_scopes(*expr, body, scopes, scope);
            }
        }
    }
    if let Some(expr) = tail {
        compute_expr_scopes(expr, body, scopes, scope);
    }
}

fn compute_expr_scopes(expr: ExprId, body: &Body, scopes: &mut ExprScopes, scope: ScopeId) {
    scopes.set_scope(expr, scope);
    match &body[expr] {
        Expr::Block { statements, tail, .. } => {
            compute_block_scopes(&statements, *tail, body, scopes, scope);
        }
        Expr::For { iterable, pat, body: body_expr, .. } => {
            compute_expr_scopes(*iterable, body, scopes, scope);
            let scope = scopes.new_scope(scope);
            scopes.add_bindings(body, scope, *pat);
            compute_expr_scopes(*body_expr, body, scopes, scope);
        }
        Expr::Lambda { args, body: body_expr, .. } => {
            let scope = scopes.new_scope(scope);
            scopes.add_params_bindings(body, scope, &args);
            compute_expr_scopes(*body_expr, body, scopes, scope);
        }
        Expr::Match { expr, arms } => {
            compute_expr_scopes(*expr, body, scopes, scope);
            for arm in arms {
                let scope = scopes.new_scope(scope);
                scopes.add_bindings(body, scope, arm.pat);
                if let Some(guard) = arm.guard {
                    scopes.set_scope(guard, scope);
                    compute_expr_scopes(guard, body, scopes, scope);
                }
                scopes.set_scope(arm.expr, scope);
                compute_expr_scopes(arm.expr, body, scopes, scope);
            }
        }
        e => e.walk_child_exprs(|e| compute_expr_scopes(e, body, scopes, scope)),
    };
}

#[cfg(test)]
mod tests {
    use hir_expand::{name::AsName, InFile};
    use ra_db::{fixture::WithFixture, FileId, SourceDatabase};
    use ra_syntax::{algo::find_node_at_offset, ast, AstNode};
    use test_utils::{assert_eq_text, extract_offset, mark};

    use crate::{db::DefDatabase, test_db::TestDB, FunctionId, ModuleDefId};

    fn find_function(db: &TestDB, file_id: FileId) -> FunctionId {
        let krate = db.test_crate();
        let crate_def_map = db.crate_def_map(krate);

        let module = crate_def_map.modules_for_file(file_id).next().unwrap();
        let (_, def) = crate_def_map[module].scope.entries().next().unwrap();
        match def.take_values().unwrap() {
            ModuleDefId::FunctionId(it) => it,
            _ => panic!(),
        }
    }

    fn do_check(ra_fixture: &str, expected: &[&str]) {
        let (off, code) = extract_offset(ra_fixture);
        let code = {
            let mut buf = String::new();
            let off: usize = off.into();
            buf.push_str(&code[..off]);
            buf.push_str("marker");
            buf.push_str(&code[off..]);
            buf
        };

        let (db, file_id) = TestDB::with_single_file(&code);

        let file_syntax = db.parse(file_id).syntax_node();
        let marker: ast::PathExpr = find_node_at_offset(&file_syntax, off).unwrap();
        let function = find_function(&db, file_id);

        let scopes = db.expr_scopes(function.into());
        let (_body, source_map) = db.body_with_source_map(function.into());

        let expr_id = source_map
            .node_expr(InFile { file_id: file_id.into(), value: &marker.into() })
            .unwrap();
        let scope = scopes.scope_for(expr_id);

        let actual = scopes
            .scope_chain(scope)
            .flat_map(|scope| scopes.entries(scope))
            .map(|it| it.name().to_string())
            .collect::<Vec<_>>()
            .join("\n");
        let expected = expected.join("\n");
        assert_eq_text!(&expected, &actual);
    }

    #[test]
    fn test_lambda_scope() {
        do_check(
            r"
            fn quux(foo: i32) {
                let f = |bar, baz: i32| {
                    <|>
                };
            }",
            &["bar", "baz", "foo"],
        );
    }

    #[test]
    fn test_call_scope() {
        do_check(
            r"
            fn quux() {
                f(|x| <|> );
            }",
            &["x"],
        );
    }

    #[test]
    fn test_method_call_scope() {
        do_check(
            r"
            fn quux() {
                z.f(|x| <|> );
            }",
            &["x"],
        );
    }

    #[test]
    fn test_loop_scope() {
        do_check(
            r"
            fn quux() {
                loop {
                    let x = ();
                    <|>
                };
            }",
            &["x"],
        );
    }

    #[test]
    fn test_match() {
        do_check(
            r"
            fn quux() {
                match () {
                    Some(x) => {
                        <|>
                    }
                };
            }",
            &["x"],
        );
    }

    #[test]
    fn test_shadow_variable() {
        do_check(
            r"
            fn foo(x: String) {
                let x : &str = &x<|>;
            }",
            &["x"],
        );
    }

    #[test]
    fn test_bindings_after_at() {
        do_check(
            r"
            fn foo() {
                match Some(()) {
                    opt @ Some(unit) => {
                        <|>
                    }
                    _ => {}
                }
            }",
            &["opt", "unit"],
        );
    }

    fn do_check_local_name(code: &str, expected_offset: u32) {
        let (off, code) = extract_offset(code);

        let (db, file_id) = TestDB::with_single_file(&code);

        let file = db.parse(file_id).ok().unwrap();
        let expected_name = find_node_at_offset::<ast::Name>(file.syntax(), expected_offset.into())
            .expect("failed to find a name at the target offset");
        let name_ref: ast::NameRef = find_node_at_offset(file.syntax(), off).unwrap();

        let function = find_function(&db, file_id);

        let scopes = db.expr_scopes(function.into());
        let (_body, source_map) = db.body_with_source_map(function.into());

        let expr_scope = {
            let expr_ast = name_ref.syntax().ancestors().find_map(ast::Expr::cast).unwrap();
            let expr_id =
                source_map.node_expr(InFile { file_id: file_id.into(), value: &expr_ast }).unwrap();
            scopes.scope_for(expr_id).unwrap()
        };

        let resolved = scopes.resolve_name_in_scope(expr_scope, &name_ref.as_name()).unwrap();
        let pat_src = source_map.pat_syntax(resolved.pat()).unwrap();

        let local_name = pat_src.value.either(
            |it| it.syntax_node_ptr().to_node(file.syntax()),
            |it| it.syntax_node_ptr().to_node(file.syntax()),
        );
        assert_eq!(local_name.text_range(), expected_name.syntax().text_range());
    }

    #[test]
    fn test_resolve_local_name() {
        do_check_local_name(
            r#"
            fn foo(x: i32, y: u32) {
                {
                    let z = x * 2;
                }
                {
                    let t = x<|> * 3;
                }
            }"#,
            21,
        );
    }

    #[test]
    fn test_resolve_local_name_declaration() {
        do_check_local_name(
            r#"
            fn foo(x: String) {
                let x : &str = &x<|>;
            }"#,
            21,
        );
    }

    #[test]
    fn test_resolve_local_name_shadow() {
        do_check_local_name(
            r"
            fn foo(x: String) {
                let x : &str = &x;
                x<|>
            }
            ",
            53,
        );
    }

    #[test]
    fn ref_patterns_contribute_bindings() {
        do_check_local_name(
            r"
            fn foo() {
                if let Some(&from) = bar() {
                    from<|>;
                }
            }
            ",
            53,
        );
    }

    #[test]
    fn while_let_desugaring() {
        mark::check!(infer_resolve_while_let);
        do_check_local_name(
            r#"
fn test() {
    let foo: Option<f32> = None;
    while let Option::Some(spam) = foo {
        spam<|>
    }
}
"#,
            75,
        );
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 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 session::config::Options;

use std::fs;
use std::io::{self, StdoutLock, Write};
use std::time::Instant;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum ProfileCategory {
    Parsing,
    Expansion,
    TypeChecking,
    BorrowChecking,
    Codegen,
    Linking,
    Other,
}

struct Categories<T> {
    parsing: T,
    expansion: T,
    type_checking: T,
    borrow_checking: T,
    codegen: T,
    linking: T,
    other: T,
}

impl<T: Default> Categories<T> {
    fn new() -> Categories<T> {
        Categories {
            parsing: T::default(),
            expansion: T::default(),
            type_checking: T::default(),
            borrow_checking: T::default(),
            codegen: T::default(),
            linking: T::default(),
            other: T::default(),
        }
    }
}

impl<T> Categories<T> {
    fn get(&self, category: ProfileCategory) -> &T {
        match category {
            ProfileCategory::Parsing => &self.parsing,
            ProfileCategory::Expansion => &self.expansion,
            ProfileCategory::TypeChecking => &self.type_checking,
            ProfileCategory::BorrowChecking => &self.borrow_checking,
            ProfileCategory::Codegen => &self.codegen,
            ProfileCategory::Linking => &self.linking,
            ProfileCategory::Other => &self.other,
        }
    }

    fn set(&mut self, category: ProfileCategory, value: T) {
        match category {
            ProfileCategory::Parsing => self.parsing = value,
            ProfileCategory::Expansion => self.expansion = value,
            ProfileCategory::TypeChecking => self.type_checking = value,
            ProfileCategory::BorrowChecking => self.borrow_checking = value,
            ProfileCategory::Codegen => self.codegen = value,
            ProfileCategory::Linking => self.linking = value,
            ProfileCategory::Other => self.other = value,
        }
    }
}

struct CategoryData {
    times: Categories<u64>,
    query_counts: Categories<(u64, u64)>,
}

impl CategoryData {
    fn new() -> CategoryData {
        CategoryData {
            times: Categories::new(),
            query_counts: Categories::new(),
        }
    }

    fn print(&self, lock: &mut StdoutLock) {
        macro_rules! p {
            ($name:tt, $rustic_name:ident) => {
                let (hits, total) = self.query_counts.$rustic_name;
                let (hits, total) = if total > 0 {
                    (format!("{:.2}", (((hits as f32) / (total as f32)) * 100.0)), total.to_string())
                } else {
                    ("".into(), "".into())
                };

                writeln!(
                   lock,
                   "| {0: <16} | {1: <14} | {2: <14} | {3: <8} |",
                   $name,
                   self.times.$rustic_name / 1_000_000,
                   total,
                   hits
                ).unwrap();
            };
        }

        writeln!(lock, "| Phase            | Time (ms)      | Queries        | Hits (%) |").unwrap();
        writeln!(lock, "| ---------------- | -------------- | -------------- | -------- |").unwrap();

        p!("Parsing", parsing);
        p!("Expansion", expansion);
        p!("TypeChecking", type_checking);
        p!("BorrowChecking", borrow_checking);
        p!("Codegen", codegen);
        p!("Linking", linking);
        p!("Other", other);
    }

    fn json(&self) -> String {
        macro_rules! j {
            ($category:tt, $rustic_name:ident) => {
                format!("{{ \"category\": {}, \"time_ms\": {} }}", stringify!($category), self.times.$rustic_name / 1_000_000)
            }
        }

        format!("[
            {},
            {},
            {},
            {},
            {},
            {},
            {}
        ]",
            j!("Parsing", parsing),
            j!("Expansion", expansion),
            j!("TypeChecking", type_checking),
            j!("BorrowChecking", borrow_checking),
            j!("Codegen", codegen),
            j!("Linking", linking),
            j!("Other", other)
        )
    }
}

pub struct SelfProfiler {
    timer_stack: Vec<ProfileCategory>,
    data: CategoryData,
    current_timer: Instant,
}

pub struct ProfilerActivity<'a>(ProfileCategory, &'a mut SelfProfiler);

impl<'a> Drop for ProfilerActivity<'a> {
    fn drop(&mut self) {
        let ProfilerActivity (category, profiler) = self;

        profiler.end_activity(*category);
    }
}

impl SelfProfiler {
    pub fn new() -> SelfProfiler {
        let mut profiler = SelfProfiler {
            timer_stack: Vec::new(),
            data: CategoryData::new(),
            current_timer: Instant::now(),
        };

        profiler.start_activity(ProfileCategory::Other);

        profiler
    }

    pub fn start_activity(&mut self, category: ProfileCategory) {
        match self.timer_stack.last().cloned() {
            None => {
                self.current_timer = Instant::now();
            },
            Some(current_category) if current_category == category => {
                //since the current category is the same as the new activity's category,
                //we don't need to do anything with the timer, we just need to push it on the stack
            }
            Some(current_category) => {
                let elapsed = self.stop_timer();

                //record the current category's time
                let new_time = self.data.times.get(current_category) + elapsed;
                self.data.times.set(current_category, new_time);
            }
        }

        //push the new category
        self.timer_stack.push(category);
    }

    pub fn record_query(&mut self, category: ProfileCategory) {
        let (hits, total) = *self.data.query_counts.get(category);
        self.data.query_counts.set(category, (hits, total + 1));
    }

    pub fn record_query_hit(&mut self, category: ProfileCategory) {
        let (hits, total) = *self.data.query_counts.get(category);
        self.data.query_counts.set(category, (hits + 1, total));
    }

    pub fn end_activity(&mut self, category: ProfileCategory) {
        match self.timer_stack.pop() {
            None => bug!("end_activity() was called but there was no running activity"),
            Some(c) =>
                assert!(
                    c == category,
                    "end_activity() was called but a different activity was running"),
        }

        //check if the new running timer is in the same category as this one
        //if it is, we don't need to do anything
        if let Some(c) = self.timer_stack.last() {
            if *c == category {
                return;
            }
        }

        //the new timer is different than the previous, so record the elapsed time and start a new timer
        let elapsed = self.stop_timer();
        let new_time = self.data.times.get(category) + elapsed;
        self.data.times.set(category, new_time);
    }

    fn stop_timer(&mut self) -> u64 {
        let elapsed = self.current_timer.elapsed();

        self.current_timer = Instant::now();

        (elapsed.as_secs() * 1_000_000_000) + (elapsed.subsec_nanos() as u64)
    }

    pub fn print_results(&mut self, opts: &Options) {
        self.end_activity(ProfileCategory::Other);

        assert!(self.timer_stack.is_empty(), "there were timers running when print_results() was called");

        let out = io::stdout();
        let mut lock = out.lock();

        let crate_name = opts.crate_name.as_ref().map(|n| format!(" for {}", n)).unwrap_or_default();

        writeln!(lock, "Self profiling results{}:", crate_name).unwrap();
        writeln!(lock).unwrap();

        self.data.print(&mut lock);

        writeln!(lock).unwrap();
        writeln!(lock, "Optimization level: {:?}", opts.optimize).unwrap();

        let incremental = if opts.incremental.is_some() { "on" } else { "off" };
        writeln!(lock, "Incremental: {}", incremental).unwrap();
    }

    pub fn save_results(&self) {
        fs::write("self_profiler_results.json", self.data.json()).unwrap();
    }

    pub fn record_activity<'a>(&'a mut self, category: ProfileCategory) -> ProfilerActivity<'a> {
        self.start_activity(category);

        ProfilerActivity(category, self)
    }
}<commit_msg>Include query hits/totals<commit_after>// Copyright 2018 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 session::config::Options;

use std::fs;
use std::io::{self, StdoutLock, Write};
use std::time::Instant;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum ProfileCategory {
    Parsing,
    Expansion,
    TypeChecking,
    BorrowChecking,
    Codegen,
    Linking,
    Other,
}

struct Categories<T> {
    parsing: T,
    expansion: T,
    type_checking: T,
    borrow_checking: T,
    codegen: T,
    linking: T,
    other: T,
}

impl<T: Default> Categories<T> {
    fn new() -> Categories<T> {
        Categories {
            parsing: T::default(),
            expansion: T::default(),
            type_checking: T::default(),
            borrow_checking: T::default(),
            codegen: T::default(),
            linking: T::default(),
            other: T::default(),
        }
    }
}

impl<T> Categories<T> {
    fn get(&self, category: ProfileCategory) -> &T {
        match category {
            ProfileCategory::Parsing => &self.parsing,
            ProfileCategory::Expansion => &self.expansion,
            ProfileCategory::TypeChecking => &self.type_checking,
            ProfileCategory::BorrowChecking => &self.borrow_checking,
            ProfileCategory::Codegen => &self.codegen,
            ProfileCategory::Linking => &self.linking,
            ProfileCategory::Other => &self.other,
        }
    }

    fn set(&mut self, category: ProfileCategory, value: T) {
        match category {
            ProfileCategory::Parsing => self.parsing = value,
            ProfileCategory::Expansion => self.expansion = value,
            ProfileCategory::TypeChecking => self.type_checking = value,
            ProfileCategory::BorrowChecking => self.borrow_checking = value,
            ProfileCategory::Codegen => self.codegen = value,
            ProfileCategory::Linking => self.linking = value,
            ProfileCategory::Other => self.other = value,
        }
    }
}

struct CategoryData {
    times: Categories<u64>,
    query_counts: Categories<(u64, u64)>,
}

impl CategoryData {
    fn new() -> CategoryData {
        CategoryData {
            times: Categories::new(),
            query_counts: Categories::new(),
        }
    }

    fn print(&self, lock: &mut StdoutLock) {
        macro_rules! p {
            ($name:tt, $rustic_name:ident) => {
                let (hits, total) = self.query_counts.$rustic_name;
                let (hits, total) = if total > 0 {
                    (format!("{:.2}", (((hits as f32) / (total as f32)) * 100.0)), total.to_string())
                } else {
                    ("".into(), "".into())
                };

                writeln!(
                   lock,
                   "| {0: <16} | {1: <14} | {2: <14} | {3: <8} |",
                   $name,
                   self.times.$rustic_name / 1_000_000,
                   total,
                   hits
                ).unwrap();
            };
        }

        writeln!(lock, "| Phase            | Time (ms)      | Queries        | Hits (%) |").unwrap();
        writeln!(lock, "| ---------------- | -------------- | -------------- | -------- |").unwrap();

        p!("Parsing", parsing);
        p!("Expansion", expansion);
        p!("TypeChecking", type_checking);
        p!("BorrowChecking", borrow_checking);
        p!("Codegen", codegen);
        p!("Linking", linking);
        p!("Other", other);
    }

    fn json(&self) -> String {
        macro_rules! j {
            ($category:tt, $rustic_name:ident) => {{
                let (hits, total) = self.query_counts.$rustic_name;

                format!("{{ \"category\": {}, \"time_ms\": {}, \"query_count\": {}, \"query_hits\": {} }}",
                    stringify!($category),
                    self.times.$rustic_name / 1_000_000,
                    total,
                    format!("{:.2}", (((hits as f32) / (total as f32)) * 100.0))
                )
            }}
        }

        format!("[
            {},
            {},
            {},
            {},
            {},
            {},
            {}
        ]",
            j!("Parsing", parsing),
            j!("Expansion", expansion),
            j!("TypeChecking", type_checking),
            j!("BorrowChecking", borrow_checking),
            j!("Codegen", codegen),
            j!("Linking", linking),
            j!("Other", other)
        )
    }
}

pub struct SelfProfiler {
    timer_stack: Vec<ProfileCategory>,
    data: CategoryData,
    current_timer: Instant,
}

pub struct ProfilerActivity<'a>(ProfileCategory, &'a mut SelfProfiler);

impl<'a> Drop for ProfilerActivity<'a> {
    fn drop(&mut self) {
        let ProfilerActivity (category, profiler) = self;

        profiler.end_activity(*category);
    }
}

impl SelfProfiler {
    pub fn new() -> SelfProfiler {
        let mut profiler = SelfProfiler {
            timer_stack: Vec::new(),
            data: CategoryData::new(),
            current_timer: Instant::now(),
        };

        profiler.start_activity(ProfileCategory::Other);

        profiler
    }

    pub fn start_activity(&mut self, category: ProfileCategory) {
        match self.timer_stack.last().cloned() {
            None => {
                self.current_timer = Instant::now();
            },
            Some(current_category) if current_category == category => {
                //since the current category is the same as the new activity's category,
                //we don't need to do anything with the timer, we just need to push it on the stack
            }
            Some(current_category) => {
                let elapsed = self.stop_timer();

                //record the current category's time
                let new_time = self.data.times.get(current_category) + elapsed;
                self.data.times.set(current_category, new_time);
            }
        }

        //push the new category
        self.timer_stack.push(category);
    }

    pub fn record_query(&mut self, category: ProfileCategory) {
        let (hits, total) = *self.data.query_counts.get(category);
        self.data.query_counts.set(category, (hits, total + 1));
    }

    pub fn record_query_hit(&mut self, category: ProfileCategory) {
        let (hits, total) = *self.data.query_counts.get(category);
        self.data.query_counts.set(category, (hits + 1, total));
    }

    pub fn end_activity(&mut self, category: ProfileCategory) {
        match self.timer_stack.pop() {
            None => bug!("end_activity() was called but there was no running activity"),
            Some(c) =>
                assert!(
                    c == category,
                    "end_activity() was called but a different activity was running"),
        }

        //check if the new running timer is in the same category as this one
        //if it is, we don't need to do anything
        if let Some(c) = self.timer_stack.last() {
            if *c == category {
                return;
            }
        }

        //the new timer is different than the previous, so record the elapsed time and start a new timer
        let elapsed = self.stop_timer();
        let new_time = self.data.times.get(category) + elapsed;
        self.data.times.set(category, new_time);
    }

    fn stop_timer(&mut self) -> u64 {
        let elapsed = self.current_timer.elapsed();

        self.current_timer = Instant::now();

        (elapsed.as_secs() * 1_000_000_000) + (elapsed.subsec_nanos() as u64)
    }

    pub fn print_results(&mut self, opts: &Options) {
        self.end_activity(ProfileCategory::Other);

        assert!(self.timer_stack.is_empty(), "there were timers running when print_results() was called");

        let out = io::stdout();
        let mut lock = out.lock();

        let crate_name = opts.crate_name.as_ref().map(|n| format!(" for {}", n)).unwrap_or_default();

        writeln!(lock, "Self profiling results{}:", crate_name).unwrap();
        writeln!(lock).unwrap();

        self.data.print(&mut lock);

        writeln!(lock).unwrap();
        writeln!(lock, "Optimization level: {:?}", opts.optimize).unwrap();

        let incremental = if opts.incremental.is_some() { "on" } else { "off" };
        writeln!(lock, "Incremental: {}", incremental).unwrap();
    }

    pub fn save_results(&self) {
        fs::write("self_profiler_results.json", self.data.json()).unwrap();
    }

    pub fn record_activity<'a>(&'a mut self, category: ProfileCategory) -> ProfilerActivity<'a> {
        self.start_activity(category);

        ProfilerActivity(category, self)
    }
}<|endoftext|>


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

// chameneos

use std::option;
use std::os;
use std::strbuf::StrBuf;
use std::task;

fn print_complements() {
    let all = [Blue, Red, Yellow];
    for aa in all.iter() {
        for bb in all.iter() {
            println!("{} + {} -> {}", show_color(*aa), show_color(*bb),
                show_color(transform(*aa, *bb)));
        }
    }
}

enum color { Red, Yellow, Blue }

struct CreatureInfo {
    name: uint,
    color: color
}

fn show_color(cc: color) -> &'static str {
    match cc {
        Red    => "red",
        Yellow => "yellow",
        Blue   => "blue"
    }
}

fn show_color_list(set: Vec<color>) -> StrBuf {
    let mut out = StrBuf::new();
    for col in set.iter() {
        out.push_char(' ');
        out.push_str(show_color(*col));
    }
    out
}

fn show_digit(nn: uint) -> &'static str {
    match nn {
        0 => {"zero"}
        1 => {"one"}
        2 => {"two"}
        3 => {"three"}
        4 => {"four"}
        5 => {"five"}
        6 => {"six"}
        7 => {"seven"}
        8 => {"eight"}
        9 => {"nine"}
        _ => {fail!("expected digits from 0 to 9...")}
    }
}

fn show_number(nn: uint) -> StrBuf {
    let mut out = vec![];
    let mut num = nn;
    let mut dig;
    let mut len = 0;
    if num == 0 { out.push(show_digit(0)) };

    while num != 0 {
        dig = num % 10;
        num = num / 10;
        out.push(" ");
        let s = show_digit(dig);
        out.push(s);
        len += 1 + s.len();
    }
    len += 1;
    out.push(" ");

    let mut ret = StrBuf::with_capacity(len);
    for s in out.iter().rev() {
        ret.push_str(*s);
    }
    ret
}

fn transform(aa: color, bb: color) -> color {
    match (aa, bb) {
        (Red,    Red   ) => { Red    }
        (Red,    Yellow) => { Blue   }
        (Red,    Blue  ) => { Yellow }
        (Yellow, Red   ) => { Blue   }
        (Yellow, Yellow) => { Yellow }
        (Yellow, Blue  ) => { Red    }
        (Blue,   Red   ) => { Yellow }
        (Blue,   Yellow) => { Red    }
        (Blue,   Blue  ) => { Blue   }
    }
}

fn creature(
    name: uint,
    color: color,
    from_rendezvous: Receiver<Option<CreatureInfo>>,
    to_rendezvous: Sender<CreatureInfo>,
    to_rendezvous_log: Sender<~str>
) {
    let mut color = color;
    let mut creatures_met = 0;
    let mut evil_clones_met = 0;

    loop {
        // ask for a pairing
        to_rendezvous.send(CreatureInfo {name: name, color: color});
        let resp = from_rendezvous.recv();

        // log and change, or print and quit
        match resp {
            option::Some(other_creature) => {
                color = transform(color, other_creature.color);

                // track some statistics
                creatures_met += 1;
                if other_creature.name == name {
                   evil_clones_met += 1;
                }
            }
            option::None => {
                // log creatures met and evil clones of self
                let report = format!("{} {}",
                                     creatures_met, show_number(evil_clones_met).as_slice());
                to_rendezvous_log.send(report);
                break;
            }
        }
    }
}

fn rendezvous(nn: uint, set: Vec<color>) {

    // these ports will allow us to hear from the creatures
    let (to_rendezvous, from_creatures) = channel::<CreatureInfo>();
    let (to_rendezvous_log, from_creatures_log) = channel::<~str>();

    // these channels will be passed to the creatures so they can talk to us

    // these channels will allow us to talk to each creature by 'name'/index
    let mut to_creature: Vec<Sender<Option<CreatureInfo>>> =
        set.iter().enumerate().map(|(ii, col)| {
            // create each creature as a listener with a port, and
            // give us a channel to talk to each
            let ii = ii;
            let col = *col;
            let to_rendezvous = to_rendezvous.clone();
            let to_rendezvous_log = to_rendezvous_log.clone();
            let (to_creature, from_rendezvous) = channel();
            task::spawn(proc() {
                creature(ii,
                         col,
                         from_rendezvous,
                         to_rendezvous.clone(),
                         to_rendezvous_log.clone());
            });
            to_creature
        }).collect();

    let mut creatures_met = 0;

    // set up meetings...
    for _ in range(0, nn) {
        let mut fst_creature: CreatureInfo = from_creatures.recv();
        let mut snd_creature: CreatureInfo = from_creatures.recv();

        creatures_met += 2;

        to_creature.get_mut(fst_creature.name).send(Some(snd_creature));
        to_creature.get_mut(snd_creature.name).send(Some(fst_creature));
    }

    // tell each creature to stop
    for to_one in to_creature.iter() {
        to_one.send(None);
    }

    // save each creature's meeting stats
    let mut report = Vec::new();
    for _to_one in to_creature.iter() {
        report.push(from_creatures_log.recv());
    }

    // print each color in the set
    println!("{}", show_color_list(set));

    // print each creature's stats
    for rep in report.iter() {
        println!("{}", *rep);
    }

    // print the total number of creatures met
    println!("{}", show_number(creatures_met));
}

fn main() {
    let args = os::args();
    let args = if os::getenv("RUST_BENCH").is_some() {
        vec!("".to_owned(), "200000".to_owned())
    } else if args.len() <= 1u {
        vec!("".to_owned(), "600".to_owned())
    } else {
        args.move_iter().collect()
    };

    let nn = from_str::<uint>(*args.get(1)).unwrap();

    print_complements();
    println!("");

    rendezvous(nn, vec!(Blue, Red, Yellow));
    println!("");

    rendezvous(nn,
        vec!(Blue, Red, Yellow, Red, Yellow, Blue, Red, Yellow, Red, Blue));
}
<commit_msg>auto merge of #13667 : TeXitoi/rust/shootout-chameneos-redux-fix, r=alexcrichton<commit_after>// Copyright 2012-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.

#![feature(phase)]
#[phase(syntax)] extern crate green;

use std::strbuf::StrBuf;
use std::fmt;

green_start!(main)

fn print_complements() {
    let all = [Blue, Red, Yellow];
    for aa in all.iter() {
        for bb in all.iter() {
            println!("{} + {} -> {}", *aa, *bb, transform(*aa, *bb));
        }
    }
}

enum Color { Red, Yellow, Blue }
impl fmt::Show for Color {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let str = match *self {
            Red => "red",
            Yellow => "yellow",
            Blue => "blue",
        };
        f.buf.write(str.as_bytes())
    }
}

struct CreatureInfo {
    name: uint,
    color: Color
}

fn show_color_list(set: Vec<Color>) -> StrBuf {
    let mut out = StrBuf::new();
    for col in set.iter() {
        out.push_char(' ');
        out.push_str(col.to_str());
    }
    out
}

fn show_digit(nn: uint) -> &'static str {
    match nn {
        0 => {" zero"}
        1 => {" one"}
        2 => {" two"}
        3 => {" three"}
        4 => {" four"}
        5 => {" five"}
        6 => {" six"}
        7 => {" seven"}
        8 => {" eight"}
        9 => {" nine"}
        _ => {fail!("expected digits from 0 to 9...")}
    }
}

struct Number(uint);
impl fmt::Show for Number {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut out = vec![];
        let Number(mut num) = *self;
        if num == 0 { out.push(show_digit(0)) };

        while num != 0 {
            let dig = num % 10;
            num = num / 10;
            let s = show_digit(dig);
            out.push(s);
        }

        for s in out.iter().rev() {
            try!(f.buf.write(s.as_bytes()));
        }
        Ok(())
    }
}

fn transform(aa: Color, bb: Color) -> Color {
    match (aa, bb) {
        (Red,    Red   ) => { Red    }
        (Red,    Yellow) => { Blue   }
        (Red,    Blue  ) => { Yellow }
        (Yellow, Red   ) => { Blue   }
        (Yellow, Yellow) => { Yellow }
        (Yellow, Blue  ) => { Red    }
        (Blue,   Red   ) => { Yellow }
        (Blue,   Yellow) => { Red    }
        (Blue,   Blue  ) => { Blue   }
    }
}

fn creature(
    name: uint,
    mut color: Color,
    from_rendezvous: Receiver<CreatureInfo>,
    to_rendezvous: Sender<CreatureInfo>,
    to_rendezvous_log: Sender<~str>
) {
    let mut creatures_met = 0;
    let mut evil_clones_met = 0;
    let mut rendezvous = from_rendezvous.iter();

    loop {
        // ask for a pairing
        to_rendezvous.send(CreatureInfo {name: name, color: color});

        // log and change, or quit
        match rendezvous.next() {
            Some(other_creature) => {
                color = transform(color, other_creature.color);

                // track some statistics
                creatures_met += 1;
                if other_creature.name == name {
                   evil_clones_met += 1;
                }
            }
            None => break
        }
    }
    // log creatures met and evil clones of self
    let report = format!("{}{}", creatures_met, Number(evil_clones_met));
    to_rendezvous_log.send(report);
}

fn rendezvous(nn: uint, set: Vec<Color>) {
    // these ports will allow us to hear from the creatures
    let (to_rendezvous, from_creatures) = channel::<CreatureInfo>();

    // these channels will be passed to the creatures so they can talk to us
    let (to_rendezvous_log, from_creatures_log) = channel::<~str>();

    // these channels will allow us to talk to each creature by 'name'/index
    let mut to_creature: Vec<Sender<CreatureInfo>> =
        set.iter().enumerate().map(|(ii, &col)| {
            // create each creature as a listener with a port, and
            // give us a channel to talk to each
            let to_rendezvous = to_rendezvous.clone();
            let to_rendezvous_log = to_rendezvous_log.clone();
            let (to_creature, from_rendezvous) = channel();
            spawn(proc() {
                creature(ii,
                         col,
                         from_rendezvous,
                         to_rendezvous,
                         to_rendezvous_log);
            });
            to_creature
        }).collect();

    let mut creatures_met = 0;

    // set up meetings...
    for _ in range(0, nn) {
        let fst_creature = from_creatures.recv();
        let snd_creature = from_creatures.recv();

        creatures_met += 2;

        to_creature.get_mut(fst_creature.name).send(snd_creature);
        to_creature.get_mut(snd_creature.name).send(fst_creature);
    }

    // tell each creature to stop
    drop(to_creature);

    // print each color in the set
    println!("{}", show_color_list(set));

    // print each creature's stats
    drop(to_rendezvous_log);
    for rep in from_creatures_log.iter() {
        println!("{}", rep);
    }

    // print the total number of creatures met
    println!("{}\n", Number(creatures_met));
}

fn main() {
    let nn = if std::os::getenv("RUST_BENCH").is_some() {
        200000
    } else {
        std::os::args().get(1).and_then(|arg| from_str(*arg)).unwrap_or(600)
    };

    print_complements();
    println!("");

    rendezvous(nn, vec!(Blue, Red, Yellow));

    rendezvous(nn,
        vec!(Blue, Red, Yellow, Red, Yellow, Blue, Red, Yellow, Red, Blue));
}
<|endoftext|>