<commit_before>//! SMTP client
//!
//! `SmtpConnection` allows manually sending SMTP commands.
//!
//! ```rust,no_run
//! # use std::error::Error;
//!
//! # #[cfg(feature = "smtp-transport")]
//! # fn main() -> Result<(), Box<dyn Error>> {
//! use lettre::transport::smtp::{SMTP_PORT, extension::ClientId, commands::*, client::SmtpConnection};
//!
//! let hello = ClientId::Domain("my_hostname".to_string());
//! let mut client = SmtpConnection::connect(&("localhost", SMTP_PORT), None, &hello, None)?;
//! client.command(
//!         Mail::new(Some("user@example.com".parse()?), vec![])
//!     )?;
//! client.command(
//!         Rcpt::new("user@example.org".parse()?, vec![])
//!       )?;
//! client.command(Data)?;
//! client.message("Test email".as_bytes())?;
//! client.command(Quit)?;
//! # Ok(())
//! # }
//! ```

#[cfg(feature = "serde")]
use std::fmt::Debug;

#[cfg(any(feature = "tokio02", feature = "tokio1", feature = "async-std1"))]
pub(crate) use self::async_connection::AsyncSmtpConnection;
#[cfg(any(feature = "tokio02", feature = "tokio1", feature = "async-std1"))]
pub(crate) use self::async_net::AsyncNetworkStream;
use self::net::NetworkStream;
#[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
pub(super) use self::tls::InnerTlsParameters;
pub use self::{
    connection::SmtpConnection,
    mock::MockStream,
    tls::{Certificate, Tls, TlsParameters, TlsParametersBuilder},
};

#[cfg(any(feature = "tokio02", feature = "tokio1", feature = "async-std1"))]
mod async_connection;
#[cfg(any(feature = "tokio02", feature = "tokio1", feature = "async-std1"))]
mod async_net;
mod connection;
mod mock;
mod net;
mod tls;

/// The codec used for transparency
#[derive(Default, Clone, Copy, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct ClientCodec {
    escape_count: u8,
}

impl ClientCodec {
    /// Creates a new client codec
    pub fn new() -> Self {
        ClientCodec::default()
    }

    /// Adds transparency
    fn encode(&mut self, frame: &[u8], buf: &mut Vec<u8>) {
        match frame.len() {
            0 => {
                match self.escape_count {
                    0 => buf.extend_from_slice(b"\r\n.\r\n"),
                    1 => buf.extend_from_slice(b"\n.\r\n"),
                    2 => buf.extend_from_slice(b".\r\n"),
                    _ => unreachable!(),
                }
                self.escape_count = 0;
            }
            _ => {
                let mut start = 0;
                for (idx, byte) in frame.iter().enumerate() {
                    match self.escape_count {
                        0 => self.escape_count = if *byte == b'\r' { 1 } else { 0 },
                        1 => self.escape_count = if *byte == b'\n' { 2 } else { 0 },
                        2 => self.escape_count = if *byte == b'.' { 3 } else { 0 },
                        _ => unreachable!(),
                    }
                    if self.escape_count == 3 {
                        self.escape_count = 0;
                        buf.extend_from_slice(&frame[start..idx]);
                        buf.extend_from_slice(b".");
                        start = idx;
                    }
                }
                buf.extend_from_slice(&frame[start..]);
            }
        }
    }
}

/// Returns the string replacing all the CRLF with "\<CRLF\>"
/// Used for debug displays
#[cfg(feature = "tracing")]
pub(super) fn escape_crlf(string: &str) -> String {
    string.replace("\r\n", "<CRLF>")
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_codec() {
        let mut codec = ClientCodec::new();
        let mut buf: Vec<u8> = vec![];

        codec.encode(b"test\r\n", &mut buf);
        codec.encode(b".\r\n", &mut buf);
        codec.encode(b"\r\ntest", &mut buf);
        codec.encode(b"te\r\n.\r\nst", &mut buf);
        codec.encode(b"test", &mut buf);
        codec.encode(b"test.", &mut buf);
        codec.encode(b"test\n", &mut buf);
        codec.encode(b".test\n", &mut buf);
        codec.encode(b"test", &mut buf);
        assert_eq!(
            String::from_utf8(buf).unwrap(),
            "test\r\n..\r\n\r\ntestte\r\n..\r\nsttesttest.test\n.test\ntest"
        );
    }

    #[test]
    #[cfg(feature = "log")]
    fn test_escape_crlf() {
        assert_eq!(escape_crlf("\r\n"), "<CRLF>");
        assert_eq!(escape_crlf("EHLO my_name\r\n"), "EHLO my_name<CRLF>");
        assert_eq!(
            escape_crlf("EHLO my_name\r\nSIZE 42\r\n"),
            "EHLO my_name<CRLF>SIZE 42<CRLF>"
        );
    }
}
<commit_msg>Make ClientCodec private (#541)<commit_after>//! SMTP client
//!
//! `SmtpConnection` allows manually sending SMTP commands.
//!
//! ```rust,no_run
//! # use std::error::Error;
//!
//! # #[cfg(feature = "smtp-transport")]
//! # fn main() -> Result<(), Box<dyn Error>> {
//! use lettre::transport::smtp::{SMTP_PORT, extension::ClientId, commands::*, client::SmtpConnection};
//!
//! let hello = ClientId::Domain("my_hostname".to_string());
//! let mut client = SmtpConnection::connect(&("localhost", SMTP_PORT), None, &hello, None)?;
//! client.command(
//!         Mail::new(Some("user@example.com".parse()?), vec![])
//!     )?;
//! client.command(
//!         Rcpt::new("user@example.org".parse()?, vec![])
//!       )?;
//! client.command(Data)?;
//! client.message("Test email".as_bytes())?;
//! client.command(Quit)?;
//! # Ok(())
//! # }
//! ```

#[cfg(feature = "serde")]
use std::fmt::Debug;

#[cfg(any(feature = "tokio02", feature = "tokio1", feature = "async-std1"))]
pub(crate) use self::async_connection::AsyncSmtpConnection;
#[cfg(any(feature = "tokio02", feature = "tokio1", feature = "async-std1"))]
pub(crate) use self::async_net::AsyncNetworkStream;
use self::net::NetworkStream;
#[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
pub(super) use self::tls::InnerTlsParameters;
pub use self::{
    connection::SmtpConnection,
    mock::MockStream,
    tls::{Certificate, Tls, TlsParameters, TlsParametersBuilder},
};

#[cfg(any(feature = "tokio02", feature = "tokio1", feature = "async-std1"))]
mod async_connection;
#[cfg(any(feature = "tokio02", feature = "tokio1", feature = "async-std1"))]
mod async_net;
mod connection;
mod mock;
mod net;
mod tls;

/// The codec used for transparency
#[derive(Default, Clone, Copy, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
struct ClientCodec {
    escape_count: u8,
}

impl ClientCodec {
    /// Creates a new client codec
    pub fn new() -> Self {
        ClientCodec::default()
    }

    /// Adds transparency
    fn encode(&mut self, frame: &[u8], buf: &mut Vec<u8>) {
        match frame.len() {
            0 => {
                match self.escape_count {
                    0 => buf.extend_from_slice(b"\r\n.\r\n"),
                    1 => buf.extend_from_slice(b"\n.\r\n"),
                    2 => buf.extend_from_slice(b".\r\n"),
                    _ => unreachable!(),
                }
                self.escape_count = 0;
            }
            _ => {
                let mut start = 0;
                for (idx, byte) in frame.iter().enumerate() {
                    match self.escape_count {
                        0 => self.escape_count = if *byte == b'\r' { 1 } else { 0 },
                        1 => self.escape_count = if *byte == b'\n' { 2 } else { 0 },
                        2 => self.escape_count = if *byte == b'.' { 3 } else { 0 },
                        _ => unreachable!(),
                    }
                    if self.escape_count == 3 {
                        self.escape_count = 0;
                        buf.extend_from_slice(&frame[start..idx]);
                        buf.extend_from_slice(b".");
                        start = idx;
                    }
                }
                buf.extend_from_slice(&frame[start..]);
            }
        }
    }
}

/// Returns the string replacing all the CRLF with "\<CRLF\>"
/// Used for debug displays
#[cfg(feature = "tracing")]
pub(super) fn escape_crlf(string: &str) -> String {
    string.replace("\r\n", "<CRLF>")
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_codec() {
        let mut codec = ClientCodec::new();
        let mut buf: Vec<u8> = vec![];

        codec.encode(b"test\r\n", &mut buf);
        codec.encode(b".\r\n", &mut buf);
        codec.encode(b"\r\ntest", &mut buf);
        codec.encode(b"te\r\n.\r\nst", &mut buf);
        codec.encode(b"test", &mut buf);
        codec.encode(b"test.", &mut buf);
        codec.encode(b"test\n", &mut buf);
        codec.encode(b".test\n", &mut buf);
        codec.encode(b"test", &mut buf);
        assert_eq!(
            String::from_utf8(buf).unwrap(),
            "test\r\n..\r\n\r\ntestte\r\n..\r\nsttesttest.test\n.test\ntest"
        );
    }

    #[test]
    #[cfg(feature = "log")]
    fn test_escape_crlf() {
        assert_eq!(escape_crlf("\r\n"), "<CRLF>");
        assert_eq!(escape_crlf("EHLO my_name\r\n"), "EHLO my_name<CRLF>");
        assert_eq!(
            escape_crlf("EHLO my_name\r\nSIZE 42\r\n"),
            "EHLO my_name<CRLF>SIZE 42<CRLF>"
        );
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Portions Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the THIRD-PARTY file.

use super::SyscallReturnCode;
use libc::{
    c_int, c_void, pthread_kill, pthread_t, sigaction, sigfillset, siginfo_t, sigset_t, EINVAL,
    SIGHUP, SIGSYS,
};
use std::io;
use std::mem;
use std::os::unix::thread::JoinHandleExt;
use std::ptr::null_mut;
use std::thread::JoinHandle;

/// Type that represents a signal handler function.
pub type SignalHandler =
    extern "C" fn(num: c_int, info: *mut siginfo_t, _unused: *mut c_void) -> ();

extern "C" {
    fn __libc_current_sigrtmin() -> c_int;
    fn __libc_current_sigrtmax() -> c_int;
}

/// Returns the minimum (inclusive) real-time signal number.
#[allow(non_snake_case)]
fn SIGRTMIN() -> c_int {
    unsafe { __libc_current_sigrtmin() }
}

/// Returns the maximum (inclusive) real-time signal number.
#[allow(non_snake_case)]
fn SIGRTMAX() -> c_int {
    unsafe { __libc_current_sigrtmax() }
}

/// Verifies that a signal number is valid when sent to a vCPU.
///
/// VCPU signals need to have values enclosed within the OS limits for realtime signals.
/// Returns either `Ok(num)` or `Err(EINVAL)`.
///
/// # Arguments
///
/// * `signum`: signal number.
///
fn validate_vcpu_signal_num(signum: c_int) -> io::Result<c_int> {
    let actual_num = signum + SIGRTMIN();
    if actual_num <= SIGRTMAX() {
        Ok(actual_num)
    } else {
        Err(io::Error::from_raw_os_error(EINVAL))
    }
}

/// Verifies that a signal number is valid when sent to the process.
///
/// Signals can take values between `SIGHUB` and `SIGSYS`.
/// Returns either `Ok(num)` or `Err(EINVAL)`.
///
/// # Arguments
///
/// * `signum`: signal number.
///
fn validate_signal_num(num: c_int) -> io::Result<c_int> {
    if num >= SIGHUP && num <= SIGSYS {
        Ok(num)
    } else {
        Err(io::Error::from_raw_os_error(EINVAL))
    }
}

/// Registers `handler` as the vCPU's signal handler of `signum`.
///
/// This is considered unsafe because the given handler will be called asynchronously, interrupting
/// whatever the thread was doing and therefore must only do async-signal-safe operations.
///
/// # Arguments
///
/// * `signum`: signal number.
/// * `handler`: signal handler functor.
///
/// # Example
///
/// ```
/// extern crate libc;
/// extern crate sys_util;
///
/// use libc::{c_int, c_void, raise, siginfo_t};
/// use sys_util::register_vcpu_signal_handler;
///
/// extern "C" fn handle_signal(_: c_int, _: *mut siginfo_t, _: *mut c_void) {}
/// extern "C" { fn __libc_current_sigrtmin() -> c_int; }
///
/// fn main() {
///     // Register dummy signal handler for `SIGRTMIN`.
///     assert!(unsafe { register_vcpu_signal_handler(0, handle_signal).is_ok() });
///     // Raise `SIGRTMIN`.
///     unsafe { raise(__libc_current_sigrtmin()); }
///     // Assert that the process is still alive.
///     assert!(true);
/// }
/// ```
///
pub unsafe fn register_vcpu_signal_handler(
    signum: c_int,
    handler: SignalHandler,
) -> io::Result<()> {
    let num = validate_vcpu_signal_num(signum)?;
    // Safe, because this is a POD struct.
    let mut sigact: sigaction = mem::zeroed();
    sigact.sa_flags = libc::SA_SIGINFO;
    sigact.sa_sigaction = handler as usize;
    SyscallReturnCode(sigaction(num, &sigact, null_mut())).into_empty_result()
}

/// Registers `handler` as the process' signal handler of `signum`.
///
/// # Arguments
///
/// * `signum`: signal number.
/// * `handler`: signal handler functor.
///
/// # Example
///
/// ```
/// extern crate libc;
/// extern crate sys_util;
///
/// use std::sync::atomic::{AtomicBool, Ordering, ATOMIC_BOOL_INIT};
/// use libc::{c_int, c_void, raise, siginfo_t, SIGUSR1};
/// use sys_util::register_signal_handler;
///
/// static HANDLER_CALLED: AtomicBool = ATOMIC_BOOL_INIT;
/// extern "C" fn handle_signal(_: c_int, _: *mut siginfo_t, _: *mut c_void) {
///     HANDLER_CALLED.store(true, Ordering::SeqCst);
/// }
///
/// fn main() {
///     assert!(unsafe { register_signal_handler(SIGUSR1, handle_signal).is_ok() });
///     unsafe { raise(SIGUSR1); }
///     assert!(HANDLER_CALLED.load(Ordering::SeqCst));
/// }
/// ```
///
pub fn register_signal_handler(signum: c_int, handler: SignalHandler) -> Result<(), io::Error> {
    let num = validate_signal_num(signum)?;
    // Safe, because this is a POD struct.
    let mut sigact: sigaction = unsafe { mem::zeroed() };
    sigact.sa_flags = libc::SA_SIGINFO;
    sigact.sa_sigaction = handler as usize;

    // We set all the bits of sa_mask, so all signals are blocked on the current thread while the
    // SIGSYS handler is executing. Safe because the parameter is valid and we check the return
    // value.
    if unsafe { sigfillset(&mut sigact.sa_mask as *mut sigset_t) } < 0 {
        return Err(io::Error::last_os_error());
    }

    // Safe because the parameters are valid and we check the return value.
    unsafe { SyscallReturnCode(sigaction(num, &sigact, null_mut())).into_empty_result() }
}

/// Trait for threads that can be signalled via `pthread_kill`.
///
/// Note that this is only useful for signals between SIGRTMIN and SIGRTMAX because these are
/// guaranteed to not be used by the C runtime.
///
/// This is marked unsafe because the implementation of this trait must guarantee that the returned
/// pthread_t is valid and has a lifetime at least that of the trait object.
pub unsafe trait Killable {
    fn pthread_handle(&self) -> pthread_t;

    /// Sends the signal `num + SIGRTMIN` to this killable thread.
    ///
    /// The value of `num + SIGRTMIN` must not exceed `SIGRTMAX`.
    fn kill(&self, num: i32) -> io::Result<()> {
        let num = validate_vcpu_signal_num(num)?;

        // Safe because we ensure we are using a valid pthread handle, a valid signal number, and
        // check the return result.
        SyscallReturnCode(unsafe { pthread_kill(self.pthread_handle() as _, num) })
            .into_empty_result()
    }
}

// Safe because we fulfill our contract of returning a genuine pthread handle.
unsafe impl<T> Killable for JoinHandle<T> {
    fn pthread_handle(&self) -> pthread_t {
        self.as_pthread_t() as _
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;
    use std::time::Duration;

    use libc::SIGSYS;

    static mut SIGNAL_HANDLER_CALLED: bool = false;

    extern "C" fn handle_signal(_: c_int, _: *mut siginfo_t, _: *mut c_void) {
        unsafe {
            SIGNAL_HANDLER_CALLED = true;
        }
    }

    #[test]
    fn test_register_signal_handler() {
        assert!(register_signal_handler(SIGSYS, handle_signal).is_ok());
        assert!(register_signal_handler(SIGSYS + 1, handle_signal).is_err());
    }

    #[test]
    #[allow(clippy::empty_loop)]
    fn test_register_vcpu_handler() {
        // Error case: invalid vCPU signal.
        unsafe {
            assert!(register_vcpu_signal_handler(SIGRTMAX(), handle_signal).is_err());
        }

        let killable = thread::spawn(|| thread::current().id());
        let killable_id = killable.join().unwrap();
        assert_ne!(killable_id, thread::current().id());

        // We install a signal handler for the specified signal; otherwise the whole process will
        // be brought down when the signal is received, as part of the default behaviour. Signal
        // handlers are global, so we install this before starting the thread.
        unsafe {
            register_vcpu_signal_handler(0, handle_signal)
                .expect("failed to register vcpu signal handler");
        }

        let killable = thread::spawn(|| loop {});

        let res = killable.kill(SIGRTMAX());
        assert!(res.is_err());
        format!("{:?}", res);

        unsafe {
            assert!(!SIGNAL_HANDLER_CALLED);
        }

        assert!(killable.kill(0).is_ok());

        // We're waiting to detect that the signal handler has been called.
        const MAX_WAIT_ITERS: u32 = 20;
        let mut iter_count = 0;
        loop {
            thread::sleep(Duration::from_millis(100));

            if unsafe { SIGNAL_HANDLER_CALLED } {
                break;
            }

            iter_count += 1;
            // timeout if we wait too long
            assert!(iter_count <= MAX_WAIT_ITERS);
        }
        // Our signal handler doesn't do anything which influences the killable thread, so the
        // previous signal is effectively ignored. If we were to join killable here, we would block
        // forever as the loop keeps running. Since we don't join, the thread will become detached
        // as the handle is dropped, and will be killed when the process/main thread exits.
    }
}
<commit_msg>Remove empty loops.<commit_after>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Portions Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the THIRD-PARTY file.

use super::SyscallReturnCode;
use libc::{
    c_int, c_void, pthread_kill, pthread_t, sigaction, sigfillset, siginfo_t, sigset_t, EINVAL,
    SIGHUP, SIGSYS,
};
use std::io;
use std::mem;
use std::os::unix::thread::JoinHandleExt;
use std::ptr::null_mut;
use std::thread::JoinHandle;

/// Type that represents a signal handler function.
pub type SignalHandler =
    extern "C" fn(num: c_int, info: *mut siginfo_t, _unused: *mut c_void) -> ();

extern "C" {
    fn __libc_current_sigrtmin() -> c_int;
    fn __libc_current_sigrtmax() -> c_int;
}

/// Returns the minimum (inclusive) real-time signal number.
#[allow(non_snake_case)]
fn SIGRTMIN() -> c_int {
    unsafe { __libc_current_sigrtmin() }
}

/// Returns the maximum (inclusive) real-time signal number.
#[allow(non_snake_case)]
fn SIGRTMAX() -> c_int {
    unsafe { __libc_current_sigrtmax() }
}

/// Verifies that a signal number is valid when sent to a vCPU.
///
/// VCPU signals need to have values enclosed within the OS limits for realtime signals.
/// Returns either `Ok(num)` or `Err(EINVAL)`.
///
/// # Arguments
///
/// * `signum`: signal number.
///
fn validate_vcpu_signal_num(signum: c_int) -> io::Result<c_int> {
    let actual_num = signum + SIGRTMIN();
    if actual_num <= SIGRTMAX() {
        Ok(actual_num)
    } else {
        Err(io::Error::from_raw_os_error(EINVAL))
    }
}

/// Verifies that a signal number is valid when sent to the process.
///
/// Signals can take values between `SIGHUB` and `SIGSYS`.
/// Returns either `Ok(num)` or `Err(EINVAL)`.
///
/// # Arguments
///
/// * `signum`: signal number.
///
fn validate_signal_num(num: c_int) -> io::Result<c_int> {
    if num >= SIGHUP && num <= SIGSYS {
        Ok(num)
    } else {
        Err(io::Error::from_raw_os_error(EINVAL))
    }
}

/// Registers `handler` as the vCPU's signal handler of `signum`.
///
/// This is considered unsafe because the given handler will be called asynchronously, interrupting
/// whatever the thread was doing and therefore must only do async-signal-safe operations.
///
/// # Arguments
///
/// * `signum`: signal number.
/// * `handler`: signal handler functor.
///
/// # Example
///
/// ```
/// extern crate libc;
/// extern crate sys_util;
///
/// use libc::{c_int, c_void, raise, siginfo_t};
/// use sys_util::register_vcpu_signal_handler;
///
/// extern "C" fn handle_signal(_: c_int, _: *mut siginfo_t, _: *mut c_void) {}
/// extern "C" { fn __libc_current_sigrtmin() -> c_int; }
///
/// fn main() {
///     // Register dummy signal handler for `SIGRTMIN`.
///     assert!(unsafe { register_vcpu_signal_handler(0, handle_signal).is_ok() });
///     // Raise `SIGRTMIN`.
///     unsafe { raise(__libc_current_sigrtmin()); }
///     // Assert that the process is still alive.
///     assert!(true);
/// }
/// ```
///
pub unsafe fn register_vcpu_signal_handler(
    signum: c_int,
    handler: SignalHandler,
) -> io::Result<()> {
    let num = validate_vcpu_signal_num(signum)?;
    // Safe, because this is a POD struct.
    let mut sigact: sigaction = mem::zeroed();
    sigact.sa_flags = libc::SA_SIGINFO;
    sigact.sa_sigaction = handler as usize;
    SyscallReturnCode(sigaction(num, &sigact, null_mut())).into_empty_result()
}

/// Registers `handler` as the process' signal handler of `signum`.
///
/// # Arguments
///
/// * `signum`: signal number.
/// * `handler`: signal handler functor.
///
/// # Example
///
/// ```
/// extern crate libc;
/// extern crate sys_util;
///
/// use std::sync::atomic::{AtomicBool, Ordering, ATOMIC_BOOL_INIT};
/// use libc::{c_int, c_void, raise, siginfo_t, SIGUSR1};
/// use sys_util::register_signal_handler;
///
/// static HANDLER_CALLED: AtomicBool = ATOMIC_BOOL_INIT;
/// extern "C" fn handle_signal(_: c_int, _: *mut siginfo_t, _: *mut c_void) {
///     HANDLER_CALLED.store(true, Ordering::SeqCst);
/// }
///
/// fn main() {
///     assert!(unsafe { register_signal_handler(SIGUSR1, handle_signal).is_ok() });
///     unsafe { raise(SIGUSR1); }
///     assert!(HANDLER_CALLED.load(Ordering::SeqCst));
/// }
/// ```
///
pub fn register_signal_handler(signum: c_int, handler: SignalHandler) -> Result<(), io::Error> {
    let num = validate_signal_num(signum)?;
    // Safe, because this is a POD struct.
    let mut sigact: sigaction = unsafe { mem::zeroed() };
    sigact.sa_flags = libc::SA_SIGINFO;
    sigact.sa_sigaction = handler as usize;

    // We set all the bits of sa_mask, so all signals are blocked on the current thread while the
    // SIGSYS handler is executing. Safe because the parameter is valid and we check the return
    // value.
    if unsafe { sigfillset(&mut sigact.sa_mask as *mut sigset_t) } < 0 {
        return Err(io::Error::last_os_error());
    }

    // Safe because the parameters are valid and we check the return value.
    unsafe { SyscallReturnCode(sigaction(num, &sigact, null_mut())).into_empty_result() }
}

/// Trait for threads that can be signalled via `pthread_kill`.
///
/// Note that this is only useful for signals between SIGRTMIN and SIGRTMAX because these are
/// guaranteed to not be used by the C runtime.
///
/// This is marked unsafe because the implementation of this trait must guarantee that the returned
/// pthread_t is valid and has a lifetime at least that of the trait object.
pub unsafe trait Killable {
    fn pthread_handle(&self) -> pthread_t;

    /// Sends the signal `num + SIGRTMIN` to this killable thread.
    ///
    /// The value of `num + SIGRTMIN` must not exceed `SIGRTMAX`.
    fn kill(&self, num: i32) -> io::Result<()> {
        let num = validate_vcpu_signal_num(num)?;

        // Safe because we ensure we are using a valid pthread handle, a valid signal number, and
        // check the return result.
        SyscallReturnCode(unsafe { pthread_kill(self.pthread_handle() as _, num) })
            .into_empty_result()
    }
}

// Safe because we fulfill our contract of returning a genuine pthread handle.
unsafe impl<T> Killable for JoinHandle<T> {
    fn pthread_handle(&self) -> pthread_t {
        self.as_pthread_t() as _
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;
    use std::time::Duration;

    use libc::SIGSYS;

    static mut SIGNAL_HANDLER_CALLED: bool = false;

    extern "C" fn handle_signal(_: c_int, _: *mut siginfo_t, _: *mut c_void) {
        unsafe {
            SIGNAL_HANDLER_CALLED = true;
        }
    }

    #[test]
    fn test_register_signal_handler() {
        assert!(register_signal_handler(SIGSYS, handle_signal).is_ok());
        assert!(register_signal_handler(SIGSYS + 1, handle_signal).is_err());
    }

    #[test]
    fn test_register_vcpu_handler() {
        // Error case: invalid vCPU signal.
        unsafe {
            assert!(register_vcpu_signal_handler(SIGRTMAX(), handle_signal).is_err());
        }

        let killable = thread::spawn(|| thread::current().id());
        let killable_id = killable.join().unwrap();
        assert_ne!(killable_id, thread::current().id());

        // We install a signal handler for the specified signal; otherwise the whole process will
        // be brought down when the signal is received, as part of the default behaviour. Signal
        // handlers are global, so we install this before starting the thread.
        unsafe {
            register_vcpu_signal_handler(0, handle_signal)
                .expect("failed to register vcpu signal handler");
        }

        let killable = thread::spawn(thread::park);

        let res = killable.kill(SIGRTMAX());
        assert!(res.is_err());
        format!("{:?}", res);

        unsafe {
            assert!(!SIGNAL_HANDLER_CALLED);
        }

        assert!(killable.kill(0).is_ok());

        // We're waiting to detect that the signal handler has been called.
        const MAX_WAIT_ITERS: u32 = 20;
        let mut iter_count = 0;
        loop {
            thread::sleep(Duration::from_millis(100));

            if unsafe { SIGNAL_HANDLER_CALLED } {
                break;
            }

            iter_count += 1;
            // timeout if we wait too long
            assert!(iter_count <= MAX_WAIT_ITERS);
        }
        // Our signal handler doesn't do anything which influences the killable thread, so the
        // previous signal is effectively ignored. If we were to join killable here, we would block
        // forever as the loop keeps running. Since we don't join, the thread will become detached
        // as the handle is dropped, and will be killed when the process/main thread exits.
    }
}
<|endoftext|>


<commit_before>//! SIMD and vendor intrinsics support library.
//!
//! This documentation is for the `coresimd` crate, but you probably want to
//! use the [`stdsimd` crate][stdsimd] which should have more complete
//! documentation.
//!
//! [stdsimd]: https://rust-lang-nursery.github.io/stdsimd/x86_64/stdsimd/

#![cfg_attr(feature = "strict", deny(warnings))]
#![allow(dead_code)]
#![allow(unused_features)]
#![feature(const_fn, link_llvm_intrinsics, platform_intrinsics, repr_simd,
           simd_ffi, target_feature, cfg_target_feature, i128_type, asm,
           integer_atomics, stmt_expr_attributes, core_intrinsics,
           crate_in_paths, no_core, attr_literals, rustc_attrs, stdsimd,
           staged_api, fn_must_use, core_float, core_slice_ext, align_offset,
           doc_cfg)]
#![cfg_attr(test,
            feature(proc_macro, test, attr_literals, abi_vectorcall,
                    untagged_unions))]
#![cfg_attr(feature = "cargo-clippy",
            allow(inline_always, too_many_arguments, cast_sign_loss,
                  cast_lossless, cast_possible_wrap,
                  cast_possible_truncation, cast_precision_loss,
                  shadow_reuse, cyclomatic_complexity, similar_names,
                  many_single_char_names))]
#![cfg_attr(test, allow(unused_imports))]
#![no_core]
#![unstable(feature = "stdsimd", issue = "0")]
#![doc(test(attr(deny(warnings))),
       test(attr(allow(dead_code, deprecated, unused_variables,
                       unused_mut))))]

#[cfg_attr(not(test), macro_use)]
extern crate core as _core;
#[cfg(test)]
#[macro_use]
extern crate std;
#[cfg(test)]
#[macro_use]
extern crate stdsimd;
#[cfg(test)]
extern crate stdsimd_test;
#[cfg(test)]
extern crate test;

macro_rules! test_v16 {
    ($item: item) => {};
}
macro_rules! test_v32 {
    ($item: item) => {};
}
macro_rules! test_v64 {
    ($item: item) => {};
}
macro_rules! test_v128 {
    ($item: item) => {};
}
macro_rules! test_v256 {
    ($item: item) => {};
}
macro_rules! test_v512 {
    ($item: item) => {};
}
macro_rules! vector_impl {
    ($([$f:ident, $($args:tt)*]),*) => { $($f!($($args)*);)* }
}

#[path = "../../../coresimd/mod.rs"]
mod coresimd;

pub use coresimd::arch;
pub use coresimd::simd;

#[allow(unused_imports)]
use _core::clone;
#[allow(unused_imports)]
use _core::cmp;
#[allow(unused_imports)]
use _core::convert;
#[allow(unused_imports)]
use _core::default;
#[allow(unused_imports)]
use _core::fmt;
#[allow(unused_imports)]
use _core::hash;
#[allow(unused_imports)]
use _core::intrinsics;
#[allow(unused_imports)]
use _core::iter;
#[allow(unused_imports)]
use _core::marker;
#[allow(unused_imports)]
use _core::mem;
#[allow(unused_imports)]
use _core::num;
#[allow(unused_imports)]
use _core::ops;
#[allow(unused_imports)]
use _core::option;
#[allow(unused_imports)]
use _core::prelude;
#[allow(unused_imports)]
use _core::ptr;
#[allow(unused_imports)]
use _core::result;
#[allow(unused_imports)]
use _core::slice;
<commit_msg>i128 is now stable.<commit_after>//! SIMD and vendor intrinsics support library.
//!
//! This documentation is for the `coresimd` crate, but you probably want to
//! use the [`stdsimd` crate][stdsimd] which should have more complete
//! documentation.
//!
//! [stdsimd]: https://rust-lang-nursery.github.io/stdsimd/x86_64/stdsimd/

#![cfg_attr(feature = "strict", deny(warnings))]
#![allow(dead_code)]
#![allow(unused_features)]
#![feature(const_fn, link_llvm_intrinsics, platform_intrinsics, repr_simd,
           simd_ffi, target_feature, cfg_target_feature, asm,
           integer_atomics, stmt_expr_attributes, core_intrinsics,
           crate_in_paths, no_core, attr_literals, rustc_attrs, stdsimd,
           staged_api, fn_must_use, core_float, core_slice_ext, align_offset,
           doc_cfg)]
#![cfg_attr(test,
            feature(proc_macro, test, attr_literals, abi_vectorcall,
                    untagged_unions))]
#![cfg_attr(feature = "cargo-clippy",
            allow(inline_always, too_many_arguments, cast_sign_loss,
                  cast_lossless, cast_possible_wrap,
                  cast_possible_truncation, cast_precision_loss,
                  shadow_reuse, cyclomatic_complexity, similar_names,
                  many_single_char_names))]
#![cfg_attr(test, allow(unused_imports))]
#![no_core]
#![unstable(feature = "stdsimd", issue = "0")]
#![doc(test(attr(deny(warnings))),
       test(attr(allow(dead_code, deprecated, unused_variables,
                       unused_mut))))]

#[cfg_attr(not(test), macro_use)]
extern crate core as _core;
#[cfg(test)]
#[macro_use]
extern crate std;
#[cfg(test)]
#[macro_use]
extern crate stdsimd;
#[cfg(test)]
extern crate stdsimd_test;
#[cfg(test)]
extern crate test;

macro_rules! test_v16 {
    ($item: item) => {};
}
macro_rules! test_v32 {
    ($item: item) => {};
}
macro_rules! test_v64 {
    ($item: item) => {};
}
macro_rules! test_v128 {
    ($item: item) => {};
}
macro_rules! test_v256 {
    ($item: item) => {};
}
macro_rules! test_v512 {
    ($item: item) => {};
}
macro_rules! vector_impl {
    ($([$f:ident, $($args:tt)*]),*) => { $($f!($($args)*);)* }
}

#[path = "../../../coresimd/mod.rs"]
mod coresimd;

pub use coresimd::arch;
pub use coresimd::simd;

#[allow(unused_imports)]
use _core::clone;
#[allow(unused_imports)]
use _core::cmp;
#[allow(unused_imports)]
use _core::convert;
#[allow(unused_imports)]
use _core::default;
#[allow(unused_imports)]
use _core::fmt;
#[allow(unused_imports)]
use _core::hash;
#[allow(unused_imports)]
use _core::intrinsics;
#[allow(unused_imports)]
use _core::iter;
#[allow(unused_imports)]
use _core::marker;
#[allow(unused_imports)]
use _core::mem;
#[allow(unused_imports)]
use _core::num;
#[allow(unused_imports)]
use _core::ops;
#[allow(unused_imports)]
use _core::option;
#[allow(unused_imports)]
use _core::prelude;
#[allow(unused_imports)]
use _core::ptr;
#[allow(unused_imports)]
use _core::result;
#[allow(unused_imports)]
use _core::slice;
<|endoftext|>


<commit_before>//! Flycheck provides the functionality needed to run `cargo check` or
//! another compatible command (f.x. clippy) in a background thread and provide
//! LSP diagnostics based on the output of the command.

#![warn(rust_2018_idioms, unused_lifetimes, semicolon_in_expressions_from_macros)]

use std::{
    fmt, io,
    process::{ChildStderr, ChildStdout, Command, Stdio},
    time::Duration,
};

use crossbeam_channel::{never, select, unbounded, Receiver, Sender};
use paths::AbsPathBuf;
use rustc_hash::FxHashMap;
use serde::Deserialize;
use stdx::{process::streaming_output, JodChild};

pub use cargo_metadata::diagnostic::{
    Applicability, Diagnostic, DiagnosticCode, DiagnosticLevel, DiagnosticSpan,
    DiagnosticSpanMacroExpansion,
};

#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub enum InvocationStrategy {
    Once,
    #[default]
    PerWorkspace,
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub enum InvocationLocation {
    Root(AbsPathBuf),
    #[default]
    Workspace,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum FlycheckConfig {
    CargoCommand {
        command: String,
        target_triple: Option<String>,
        all_targets: bool,
        no_default_features: bool,
        all_features: bool,
        features: Vec<String>,
        extra_args: Vec<String>,
        extra_env: FxHashMap<String, String>,
    },
    CustomCommand {
        command: String,
        args: Vec<String>,
        extra_env: FxHashMap<String, String>,
        invocation_strategy: InvocationStrategy,
        invocation_location: InvocationLocation,
    },
}

impl fmt::Display for FlycheckConfig {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            FlycheckConfig::CargoCommand { command, .. } => write!(f, "cargo {}", command),
            FlycheckConfig::CustomCommand { command, args, .. } => {
                write!(f, "{} {}", command, args.join(" "))
            }
        }
    }
}

/// Flycheck wraps the shared state and communication machinery used for
/// running `cargo check` (or other compatible command) and providing
/// diagnostics based on the output.
/// The spawned thread is shut down when this struct is dropped.
#[derive(Debug)]
pub struct FlycheckHandle {
    // XXX: drop order is significant
    sender: Sender<Restart>,
    _thread: jod_thread::JoinHandle,
    id: usize,
}

impl FlycheckHandle {
    pub fn spawn(
        id: usize,
        sender: Box<dyn Fn(Message) + Send>,
        config: FlycheckConfig,
        workspace_root: AbsPathBuf,
    ) -> FlycheckHandle {
        let actor = FlycheckActor::new(id, sender, config, workspace_root);
        let (sender, receiver) = unbounded::<Restart>();
        let thread = jod_thread::Builder::new()
            .name("Flycheck".to_owned())
            .spawn(move || actor.run(receiver))
            .expect("failed to spawn thread");
        FlycheckHandle { id, sender, _thread: thread }
    }

    /// Schedule a re-start of the cargo check worker.
    pub fn restart(&self) {
        self.sender.send(Restart::Yes).unwrap();
    }

    /// Stop this cargo check worker.
    pub fn cancel(&self) {
        self.sender.send(Restart::No).unwrap();
    }

    pub fn id(&self) -> usize {
        self.id
    }
}

pub enum Message {
    /// Request adding a diagnostic with fixes included to a file
    AddDiagnostic { id: usize, workspace_root: AbsPathBuf, diagnostic: Diagnostic },

    /// Request check progress notification to client
    Progress {
        /// Flycheck instance ID
        id: usize,
        progress: Progress,
    },
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Message::AddDiagnostic { id, workspace_root, diagnostic } => f
                .debug_struct("AddDiagnostic")
                .field("id", id)
                .field("workspace_root", workspace_root)
                .field("diagnostic_code", &diagnostic.code.as_ref().map(|it| &it.code))
                .finish(),
            Message::Progress { id, progress } => {
                f.debug_struct("Progress").field("id", id).field("progress", progress).finish()
            }
        }
    }
}

#[derive(Debug)]
pub enum Progress {
    DidStart,
    DidCheckCrate(String),
    DidFinish(io::Result<()>),
    DidCancel,
    DidFailToRestart(String),
}

enum Restart {
    Yes,
    No,
}

/// A [`FlycheckActor`] is a single check instance of a workspace.
struct FlycheckActor {
    /// The workspace id of this flycheck instance.
    id: usize,
    sender: Box<dyn Fn(Message) + Send>,
    config: FlycheckConfig,
    /// Either the workspace root of the workspace we are flychecking,
    /// or the project root of the project.
    root: AbsPathBuf,
    /// CargoHandle exists to wrap around the communication needed to be able to
    /// run `cargo check` without blocking. Currently the Rust standard library
    /// doesn't provide a way to read sub-process output without blocking, so we
    /// have to wrap sub-processes output handling in a thread and pass messages
    /// back over a channel.
    cargo_handle: Option<CargoHandle>,
}

enum Event {
    Restart(Restart),
    CheckEvent(Option<CargoMessage>),
}

impl FlycheckActor {
    fn new(
        id: usize,
        sender: Box<dyn Fn(Message) + Send>,
        config: FlycheckConfig,
        workspace_root: AbsPathBuf,
    ) -> FlycheckActor {
        tracing::info!(%id, ?workspace_root, "Spawning flycheck");
        FlycheckActor { id, sender, config, root: workspace_root, cargo_handle: None }
    }

    fn report_progress(&self, progress: Progress) {
        self.send(Message::Progress { id: self.id, progress });
    }

    fn next_event(&self, inbox: &Receiver<Restart>) -> Option<Event> {
        let check_chan = self.cargo_handle.as_ref().map(|cargo| &cargo.receiver);
        if let Ok(msg) = inbox.try_recv() {
            // give restarts a preference so check outputs don't block a restart or stop
            return Some(Event::Restart(msg));
        }
        select! {
            recv(inbox) -> msg => msg.ok().map(Event::Restart),
            recv(check_chan.unwrap_or(&never())) -> msg => Some(Event::CheckEvent(msg.ok())),
        }
    }

    fn run(mut self, inbox: Receiver<Restart>) {
        'event: while let Some(event) = self.next_event(&inbox) {
            match event {
                Event::Restart(Restart::No) => {
                    self.cancel_check_process();
                }
                Event::Restart(Restart::Yes) => {
                    // Cancel the previously spawned process
                    self.cancel_check_process();
                    while let Ok(restart) = inbox.recv_timeout(Duration::from_millis(50)) {
                        // restart chained with a stop, so just cancel
                        if let Restart::No = restart {
                            continue 'event;
                        }
                    }

                    let command = self.check_command();
                    tracing::debug!(?command, "will restart flycheck");
                    match CargoHandle::spawn(command) {
                        Ok(cargo_handle) => {
                            tracing::debug!(
                                command = ?self.check_command(),
                                "did  restart flycheck"
                            );
                            self.cargo_handle = Some(cargo_handle);
                            self.report_progress(Progress::DidStart);
                        }
                        Err(error) => {
                            self.report_progress(Progress::DidFailToRestart(format!(
                                "Failed to run the following command: {:?} error={}",
                                self.check_command(),
                                error
                            )));
                        }
                    }
                }
                Event::CheckEvent(None) => {
                    tracing::debug!(flycheck_id = self.id, "flycheck finished");

                    // Watcher finished
                    let cargo_handle = self.cargo_handle.take().unwrap();
                    let res = cargo_handle.join();
                    if res.is_err() {
                        tracing::error!(
                            "Flycheck failed to run the following command: {:?}",
                            self.check_command()
                        );
                    }
                    self.report_progress(Progress::DidFinish(res));
                }
                Event::CheckEvent(Some(message)) => match message {
                    CargoMessage::CompilerArtifact(msg) => {
                        self.report_progress(Progress::DidCheckCrate(msg.target.name));
                    }

                    CargoMessage::Diagnostic(msg) => {
                        self.send(Message::AddDiagnostic {
                            id: self.id,
                            workspace_root: self.root.clone(),
                            diagnostic: msg,
                        });
                    }
                },
            }
        }
        // If we rerun the thread, we need to discard the previous check results first
        self.cancel_check_process();
    }

    fn cancel_check_process(&mut self) {
        if let Some(cargo_handle) = self.cargo_handle.take() {
            tracing::debug!(
                command = ?self.check_command(),
                "did  cancel flycheck"
            );
            cargo_handle.cancel();
            self.report_progress(Progress::DidCancel);
        }
    }

    fn check_command(&self) -> Command {
        let (mut cmd, args) = match &self.config {
            FlycheckConfig::CargoCommand {
                command,
                target_triple,
                no_default_features,
                all_targets,
                all_features,
                extra_args,
                features,
                extra_env,
            } => {
                let mut cmd = Command::new(toolchain::cargo());
                cmd.arg(command);
                cmd.args(&["--workspace", "--message-format=json"]);

                if let Some(target) = target_triple {
                    cmd.args(&["--target", target.as_str()]);
                }
                if *all_targets {
                    cmd.arg("--all-targets");
                }
                if *all_features {
                    cmd.arg("--all-features");
                } else {
                    if *no_default_features {
                        cmd.arg("--no-default-features");
                    }
                    if !features.is_empty() {
                        cmd.arg("--features");
                        cmd.arg(features.join(" "));
                    }
                }
                cmd.envs(extra_env);
                (cmd, extra_args)
            }
            FlycheckConfig::CustomCommand {
                command,
                args,
                extra_env,
                invocation_strategy,
                invocation_location,
            } => {
                let mut cmd = Command::new(command);
                cmd.envs(extra_env);

                match invocation_location {
                    InvocationLocation::Workspace => {
                        match invocation_strategy {
                            InvocationStrategy::Once => {
                                cmd.current_dir(&self.root);
                            }
                            InvocationStrategy::PerWorkspace => {
                                // FIXME: cmd.current_dir(&affected_workspace);
                                cmd.current_dir(&self.root);
                            }
                        }
                    }
                    InvocationLocation::Root(root) => {
                        cmd.current_dir(root);
                    }
                }

                (cmd, args)
            }
        };

        cmd.args(args);
        cmd
    }

    fn send(&self, check_task: Message) {
        (self.sender)(check_task);
    }
}

/// A handle to a cargo process used for fly-checking.
struct CargoHandle {
    /// The handle to the actual cargo process. As we cannot cancel directly from with
    /// a read syscall dropping and therefor terminating the process is our best option.
    child: JodChild,
    thread: jod_thread::JoinHandle<io::Result<(bool, String)>>,
    receiver: Receiver<CargoMessage>,
}

impl CargoHandle {
    fn spawn(mut command: Command) -> std::io::Result<CargoHandle> {
        command.stdout(Stdio::piped()).stderr(Stdio::piped()).stdin(Stdio::null());
        let mut child = JodChild::spawn(command)?;

        let stdout = child.stdout.take().unwrap();
        let stderr = child.stderr.take().unwrap();

        let (sender, receiver) = unbounded();
        let actor = CargoActor::new(sender, stdout, stderr);
        let thread = jod_thread::Builder::new()
            .name("CargoHandle".to_owned())
            .spawn(move || actor.run())
            .expect("failed to spawn thread");
        Ok(CargoHandle { child, thread, receiver })
    }

    fn cancel(mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
    }

    fn join(mut self) -> io::Result<()> {
        let _ = self.child.kill();
        let exit_status = self.child.wait()?;
        let (read_at_least_one_message, error) = self.thread.join()?;
        if read_at_least_one_message || exit_status.success() {
            Ok(())
        } else {
            Err(io::Error::new(io::ErrorKind::Other, format!(
                "Cargo watcher failed, the command produced no valid metadata (exit code: {:?}):\n{}",
                exit_status, error
            )))
        }
    }
}

struct CargoActor {
    sender: Sender<CargoMessage>,
    stdout: ChildStdout,
    stderr: ChildStderr,
}

impl CargoActor {
    fn new(sender: Sender<CargoMessage>, stdout: ChildStdout, stderr: ChildStderr) -> CargoActor {
        CargoActor { sender, stdout, stderr }
    }

    fn run(self) -> io::Result<(bool, String)> {
        // We manually read a line at a time, instead of using serde's
        // stream deserializers, because the deserializer cannot recover
        // from an error, resulting in it getting stuck, because we try to
        // be resilient against failures.
        //
        // Because cargo only outputs one JSON object per line, we can
        // simply skip a line if it doesn't parse, which just ignores any
        // erroneous output.

        let mut error = String::new();
        let mut read_at_least_one_message = false;
        let output = streaming_output(
            self.stdout,
            self.stderr,
            &mut |line| {
                read_at_least_one_message = true;

                // Try to deserialize a message from Cargo or Rustc.
                let mut deserializer = serde_json::Deserializer::from_str(line);
                deserializer.disable_recursion_limit();
                if let Ok(message) = JsonMessage::deserialize(&mut deserializer) {
                    match message {
                        // Skip certain kinds of messages to only spend time on what's useful
                        JsonMessage::Cargo(message) => match message {
                            cargo_metadata::Message::CompilerArtifact(artifact)
                                if !artifact.fresh =>
                            {
                                self.sender.send(CargoMessage::CompilerArtifact(artifact)).unwrap();
                            }
                            cargo_metadata::Message::CompilerMessage(msg) => {
                                self.sender.send(CargoMessage::Diagnostic(msg.message)).unwrap();
                            }
                            _ => (),
                        },
                        JsonMessage::Rustc(message) => {
                            self.sender.send(CargoMessage::Diagnostic(message)).unwrap();
                        }
                    }
                }
            },
            &mut |line| {
                error.push_str(line);
                error.push('\n');
            },
        );
        match output {
            Ok(_) => Ok((read_at_least_one_message, error)),
            Err(e) => Err(io::Error::new(e.kind(), format!("{:?}: {}", e, error))),
        }
    }
}

enum CargoMessage {
    CompilerArtifact(cargo_metadata::Artifact),
    Diagnostic(Diagnostic),
}

#[derive(Deserialize)]
#[serde(untagged)]
enum JsonMessage {
    Cargo(cargo_metadata::Message),
    Rustc(Diagnostic),
}
<commit_msg>fix: Fix standard flycheck command not being executed in the workspace it is being invoked for<commit_after>//! Flycheck provides the functionality needed to run `cargo check` or
//! another compatible command (f.x. clippy) in a background thread and provide
//! LSP diagnostics based on the output of the command.

#![warn(rust_2018_idioms, unused_lifetimes, semicolon_in_expressions_from_macros)]

use std::{
    fmt, io,
    process::{ChildStderr, ChildStdout, Command, Stdio},
    time::Duration,
};

use crossbeam_channel::{never, select, unbounded, Receiver, Sender};
use paths::AbsPathBuf;
use rustc_hash::FxHashMap;
use serde::Deserialize;
use stdx::{process::streaming_output, JodChild};

pub use cargo_metadata::diagnostic::{
    Applicability, Diagnostic, DiagnosticCode, DiagnosticLevel, DiagnosticSpan,
    DiagnosticSpanMacroExpansion,
};

#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub enum InvocationStrategy {
    Once,
    #[default]
    PerWorkspace,
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub enum InvocationLocation {
    Root(AbsPathBuf),
    #[default]
    Workspace,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum FlycheckConfig {
    CargoCommand {
        command: String,
        target_triple: Option<String>,
        all_targets: bool,
        no_default_features: bool,
        all_features: bool,
        features: Vec<String>,
        extra_args: Vec<String>,
        extra_env: FxHashMap<String, String>,
    },
    CustomCommand {
        command: String,
        args: Vec<String>,
        extra_env: FxHashMap<String, String>,
        invocation_strategy: InvocationStrategy,
        invocation_location: InvocationLocation,
    },
}

impl fmt::Display for FlycheckConfig {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            FlycheckConfig::CargoCommand { command, .. } => write!(f, "cargo {}", command),
            FlycheckConfig::CustomCommand { command, args, .. } => {
                write!(f, "{} {}", command, args.join(" "))
            }
        }
    }
}

/// Flycheck wraps the shared state and communication machinery used for
/// running `cargo check` (or other compatible command) and providing
/// diagnostics based on the output.
/// The spawned thread is shut down when this struct is dropped.
#[derive(Debug)]
pub struct FlycheckHandle {
    // XXX: drop order is significant
    sender: Sender<Restart>,
    _thread: jod_thread::JoinHandle,
    id: usize,
}

impl FlycheckHandle {
    pub fn spawn(
        id: usize,
        sender: Box<dyn Fn(Message) + Send>,
        config: FlycheckConfig,
        workspace_root: AbsPathBuf,
    ) -> FlycheckHandle {
        let actor = FlycheckActor::new(id, sender, config, workspace_root);
        let (sender, receiver) = unbounded::<Restart>();
        let thread = jod_thread::Builder::new()
            .name("Flycheck".to_owned())
            .spawn(move || actor.run(receiver))
            .expect("failed to spawn thread");
        FlycheckHandle { id, sender, _thread: thread }
    }

    /// Schedule a re-start of the cargo check worker.
    pub fn restart(&self) {
        self.sender.send(Restart::Yes).unwrap();
    }

    /// Stop this cargo check worker.
    pub fn cancel(&self) {
        self.sender.send(Restart::No).unwrap();
    }

    pub fn id(&self) -> usize {
        self.id
    }
}

pub enum Message {
    /// Request adding a diagnostic with fixes included to a file
    AddDiagnostic { id: usize, workspace_root: AbsPathBuf, diagnostic: Diagnostic },

    /// Request check progress notification to client
    Progress {
        /// Flycheck instance ID
        id: usize,
        progress: Progress,
    },
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Message::AddDiagnostic { id, workspace_root, diagnostic } => f
                .debug_struct("AddDiagnostic")
                .field("id", id)
                .field("workspace_root", workspace_root)
                .field("diagnostic_code", &diagnostic.code.as_ref().map(|it| &it.code))
                .finish(),
            Message::Progress { id, progress } => {
                f.debug_struct("Progress").field("id", id).field("progress", progress).finish()
            }
        }
    }
}

#[derive(Debug)]
pub enum Progress {
    DidStart,
    DidCheckCrate(String),
    DidFinish(io::Result<()>),
    DidCancel,
    DidFailToRestart(String),
}

enum Restart {
    Yes,
    No,
}

/// A [`FlycheckActor`] is a single check instance of a workspace.
struct FlycheckActor {
    /// The workspace id of this flycheck instance.
    id: usize,
    sender: Box<dyn Fn(Message) + Send>,
    config: FlycheckConfig,
    /// Either the workspace root of the workspace we are flychecking,
    /// or the project root of the project.
    root: AbsPathBuf,
    /// CargoHandle exists to wrap around the communication needed to be able to
    /// run `cargo check` without blocking. Currently the Rust standard library
    /// doesn't provide a way to read sub-process output without blocking, so we
    /// have to wrap sub-processes output handling in a thread and pass messages
    /// back over a channel.
    cargo_handle: Option<CargoHandle>,
}

enum Event {
    Restart(Restart),
    CheckEvent(Option<CargoMessage>),
}

impl FlycheckActor {
    fn new(
        id: usize,
        sender: Box<dyn Fn(Message) + Send>,
        config: FlycheckConfig,
        workspace_root: AbsPathBuf,
    ) -> FlycheckActor {
        tracing::info!(%id, ?workspace_root, "Spawning flycheck");
        FlycheckActor { id, sender, config, root: workspace_root, cargo_handle: None }
    }

    fn report_progress(&self, progress: Progress) {
        self.send(Message::Progress { id: self.id, progress });
    }

    fn next_event(&self, inbox: &Receiver<Restart>) -> Option<Event> {
        let check_chan = self.cargo_handle.as_ref().map(|cargo| &cargo.receiver);
        if let Ok(msg) = inbox.try_recv() {
            // give restarts a preference so check outputs don't block a restart or stop
            return Some(Event::Restart(msg));
        }
        select! {
            recv(inbox) -> msg => msg.ok().map(Event::Restart),
            recv(check_chan.unwrap_or(&never())) -> msg => Some(Event::CheckEvent(msg.ok())),
        }
    }

    fn run(mut self, inbox: Receiver<Restart>) {
        'event: while let Some(event) = self.next_event(&inbox) {
            match event {
                Event::Restart(Restart::No) => {
                    self.cancel_check_process();
                }
                Event::Restart(Restart::Yes) => {
                    // Cancel the previously spawned process
                    self.cancel_check_process();
                    while let Ok(restart) = inbox.recv_timeout(Duration::from_millis(50)) {
                        // restart chained with a stop, so just cancel
                        if let Restart::No = restart {
                            continue 'event;
                        }
                    }

                    let command = self.check_command();
                    tracing::debug!(?command, "will restart flycheck");
                    match CargoHandle::spawn(command) {
                        Ok(cargo_handle) => {
                            tracing::debug!(
                                command = ?self.check_command(),
                                "did  restart flycheck"
                            );
                            self.cargo_handle = Some(cargo_handle);
                            self.report_progress(Progress::DidStart);
                        }
                        Err(error) => {
                            self.report_progress(Progress::DidFailToRestart(format!(
                                "Failed to run the following command: {:?} error={}",
                                self.check_command(),
                                error
                            )));
                        }
                    }
                }
                Event::CheckEvent(None) => {
                    tracing::debug!(flycheck_id = self.id, "flycheck finished");

                    // Watcher finished
                    let cargo_handle = self.cargo_handle.take().unwrap();
                    let res = cargo_handle.join();
                    if res.is_err() {
                        tracing::error!(
                            "Flycheck failed to run the following command: {:?}",
                            self.check_command()
                        );
                    }
                    self.report_progress(Progress::DidFinish(res));
                }
                Event::CheckEvent(Some(message)) => match message {
                    CargoMessage::CompilerArtifact(msg) => {
                        self.report_progress(Progress::DidCheckCrate(msg.target.name));
                    }

                    CargoMessage::Diagnostic(msg) => {
                        self.send(Message::AddDiagnostic {
                            id: self.id,
                            workspace_root: self.root.clone(),
                            diagnostic: msg,
                        });
                    }
                },
            }
        }
        // If we rerun the thread, we need to discard the previous check results first
        self.cancel_check_process();
    }

    fn cancel_check_process(&mut self) {
        if let Some(cargo_handle) = self.cargo_handle.take() {
            tracing::debug!(
                command = ?self.check_command(),
                "did  cancel flycheck"
            );
            cargo_handle.cancel();
            self.report_progress(Progress::DidCancel);
        }
    }

    fn check_command(&self) -> Command {
        let (mut cmd, args) = match &self.config {
            FlycheckConfig::CargoCommand {
                command,
                target_triple,
                no_default_features,
                all_targets,
                all_features,
                extra_args,
                features,
                extra_env,
            } => {
                let mut cmd = Command::new(toolchain::cargo());
                cmd.arg(command);
                cmd.current_dir(&self.root);
                cmd.args(&["--workspace", "--message-format=json", "--manifest-path"])
                    .arg(self.root.join("Cargo.toml").as_os_str());

                if let Some(target) = target_triple {
                    cmd.args(&["--target", target.as_str()]);
                }
                if *all_targets {
                    cmd.arg("--all-targets");
                }
                if *all_features {
                    cmd.arg("--all-features");
                } else {
                    if *no_default_features {
                        cmd.arg("--no-default-features");
                    }
                    if !features.is_empty() {
                        cmd.arg("--features");
                        cmd.arg(features.join(" "));
                    }
                }
                cmd.envs(extra_env);
                (cmd, extra_args)
            }
            FlycheckConfig::CustomCommand {
                command,
                args,
                extra_env,
                invocation_strategy,
                invocation_location,
            } => {
                let mut cmd = Command::new(command);
                cmd.envs(extra_env);

                match invocation_location {
                    InvocationLocation::Workspace => {
                        match invocation_strategy {
                            InvocationStrategy::Once => {
                                cmd.current_dir(&self.root);
                            }
                            InvocationStrategy::PerWorkspace => {
                                // FIXME: cmd.current_dir(&affected_workspace);
                                cmd.current_dir(&self.root);
                            }
                        }
                    }
                    InvocationLocation::Root(root) => {
                        cmd.current_dir(root);
                    }
                }

                (cmd, args)
            }
        };

        cmd.args(args);
        cmd
    }

    fn send(&self, check_task: Message) {
        (self.sender)(check_task);
    }
}

/// A handle to a cargo process used for fly-checking.
struct CargoHandle {
    /// The handle to the actual cargo process. As we cannot cancel directly from with
    /// a read syscall dropping and therefor terminating the process is our best option.
    child: JodChild,
    thread: jod_thread::JoinHandle<io::Result<(bool, String)>>,
    receiver: Receiver<CargoMessage>,
}

impl CargoHandle {
    fn spawn(mut command: Command) -> std::io::Result<CargoHandle> {
        command.stdout(Stdio::piped()).stderr(Stdio::piped()).stdin(Stdio::null());
        let mut child = JodChild::spawn(command)?;

        let stdout = child.stdout.take().unwrap();
        let stderr = child.stderr.take().unwrap();

        let (sender, receiver) = unbounded();
        let actor = CargoActor::new(sender, stdout, stderr);
        let thread = jod_thread::Builder::new()
            .name("CargoHandle".to_owned())
            .spawn(move || actor.run())
            .expect("failed to spawn thread");
        Ok(CargoHandle { child, thread, receiver })
    }

    fn cancel(mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
    }

    fn join(mut self) -> io::Result<()> {
        let _ = self.child.kill();
        let exit_status = self.child.wait()?;
        let (read_at_least_one_message, error) = self.thread.join()?;
        if read_at_least_one_message || exit_status.success() {
            Ok(())
        } else {
            Err(io::Error::new(io::ErrorKind::Other, format!(
                "Cargo watcher failed, the command produced no valid metadata (exit code: {:?}):\n{}",
                exit_status, error
            )))
        }
    }
}

struct CargoActor {
    sender: Sender<CargoMessage>,
    stdout: ChildStdout,
    stderr: ChildStderr,
}

impl CargoActor {
    fn new(sender: Sender<CargoMessage>, stdout: ChildStdout, stderr: ChildStderr) -> CargoActor {
        CargoActor { sender, stdout, stderr }
    }

    fn run(self) -> io::Result<(bool, String)> {
        // We manually read a line at a time, instead of using serde's
        // stream deserializers, because the deserializer cannot recover
        // from an error, resulting in it getting stuck, because we try to
        // be resilient against failures.
        //
        // Because cargo only outputs one JSON object per line, we can
        // simply skip a line if it doesn't parse, which just ignores any
        // erroneous output.

        let mut error = String::new();
        let mut read_at_least_one_message = false;
        let output = streaming_output(
            self.stdout,
            self.stderr,
            &mut |line| {
                read_at_least_one_message = true;

                // Try to deserialize a message from Cargo or Rustc.
                let mut deserializer = serde_json::Deserializer::from_str(line);
                deserializer.disable_recursion_limit();
                if let Ok(message) = JsonMessage::deserialize(&mut deserializer) {
                    match message {
                        // Skip certain kinds of messages to only spend time on what's useful
                        JsonMessage::Cargo(message) => match message {
                            cargo_metadata::Message::CompilerArtifact(artifact)
                                if !artifact.fresh =>
                            {
                                self.sender.send(CargoMessage::CompilerArtifact(artifact)).unwrap();
                            }
                            cargo_metadata::Message::CompilerMessage(msg) => {
                                self.sender.send(CargoMessage::Diagnostic(msg.message)).unwrap();
                            }
                            _ => (),
                        },
                        JsonMessage::Rustc(message) => {
                            self.sender.send(CargoMessage::Diagnostic(message)).unwrap();
                        }
                    }
                }
            },
            &mut |line| {
                error.push_str(line);
                error.push('\n');
            },
        );
        match output {
            Ok(_) => Ok((read_at_least_one_message, error)),
            Err(e) => Err(io::Error::new(e.kind(), format!("{:?}: {}", e, error))),
        }
    }
}

enum CargoMessage {
    CompilerArtifact(cargo_metadata::Artifact),
    Diagnostic(Diagnostic),
}

#[derive(Deserialize)]
#[serde(untagged)]
enum JsonMessage {
    Cargo(cargo_metadata::Message),
    Rustc(Diagnostic),
}
<|endoftext|>


<commit_before>//! A higher level attributes based on TokenTree, with also some shortcuts.

use std::{ops, sync::Arc};

use cfg::{CfgExpr, CfgOptions};
use either::Either;
use hir_expand::{hygiene::Hygiene, AstId, InFile};
use itertools::Itertools;
use mbe::ast_to_token_tree;
use syntax::{
    ast::{self, AstNode, AttrsOwner},
    SmolStr,
};
use tt::Subtree;

use crate::{
    db::DefDatabase,
    item_tree::{ItemTreeId, ItemTreeNode},
    nameres::ModuleSource,
    path::ModPath,
    src::HasChildSource,
    AdtId, AttrDefId, Lookup,
};

/// Holds documentation
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Documentation(Arc<str>);

impl Documentation {
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl Into<String> for Documentation {
    fn into(self) -> String {
        self.as_str().to_owned()
    }
}

#[derive(Default, Debug, Clone, PartialEq, Eq)]
pub struct Attrs {
    entries: Option<Arc<[Attr]>>,
}

impl ops::Deref for Attrs {
    type Target = [Attr];

    fn deref(&self) -> &[Attr] {
        match &self.entries {
            Some(it) => &*it,
            None => &[],
        }
    }
}

impl Attrs {
    pub const EMPTY: Attrs = Attrs { entries: None };

    pub(crate) fn attrs_query(db: &dyn DefDatabase, def: AttrDefId) -> Attrs {
        match def {
            AttrDefId::ModuleId(module) => {
                let def_map = db.crate_def_map(module.krate);
                let mod_data = &def_map[module.local_id];
                match mod_data.declaration_source(db) {
                    Some(it) => {
                        Attrs::from_attrs_owner(db, it.as_ref().map(|it| it as &dyn AttrsOwner))
                    }
                    None => Attrs::from_attrs_owner(
                        db,
                        mod_data.definition_source(db).as_ref().map(|src| match src {
                            ModuleSource::SourceFile(file) => file as &dyn AttrsOwner,
                            ModuleSource::Module(module) => module as &dyn AttrsOwner,
                        }),
                    ),
                }
            }
            AttrDefId::FieldId(it) => {
                let src = it.parent.child_source(db);
                match &src.value[it.local_id] {
                    Either::Left(_tuple) => Attrs::default(),
                    Either::Right(record) => Attrs::from_attrs_owner(db, src.with_value(record)),
                }
            }
            AttrDefId::EnumVariantId(var_id) => {
                let src = var_id.parent.child_source(db);
                let src = src.as_ref().map(|it| &it[var_id.local_id]);
                Attrs::from_attrs_owner(db, src.map(|it| it as &dyn AttrsOwner))
            }
            AttrDefId::AdtId(it) => match it {
                AdtId::StructId(it) => attrs_from_item_tree(it.lookup(db).id, db),
                AdtId::EnumId(it) => attrs_from_item_tree(it.lookup(db).id, db),
                AdtId::UnionId(it) => attrs_from_item_tree(it.lookup(db).id, db),
            },
            AttrDefId::TraitId(it) => attrs_from_item_tree(it.lookup(db).id, db),
            AttrDefId::MacroDefId(it) => {
                it.ast_id.map_or_else(Default::default, |ast_id| attrs_from_ast(ast_id, db))
            }
            AttrDefId::ImplId(it) => attrs_from_item_tree(it.lookup(db).id, db),
            AttrDefId::ConstId(it) => attrs_from_item_tree(it.lookup(db).id, db),
            AttrDefId::StaticId(it) => attrs_from_item_tree(it.lookup(db).id, db),
            AttrDefId::FunctionId(it) => attrs_from_item_tree(it.lookup(db).id, db),
            AttrDefId::TypeAliasId(it) => attrs_from_item_tree(it.lookup(db).id, db),
        }
    }

    pub fn from_attrs_owner(db: &dyn DefDatabase, owner: InFile<&dyn AttrsOwner>) -> Attrs {
        let hygiene = Hygiene::new(db.upcast(), owner.file_id);
        Attrs::new(owner.value, &hygiene)
    }

    pub(crate) fn new(owner: &dyn AttrsOwner, hygiene: &Hygiene) -> Attrs {
        let docs = ast::CommentIter::from_syntax_node(owner.syntax()).doc_comment_text().map(
            |docs_text| Attr {
                input: Some(AttrInput::Literal(SmolStr::new(docs_text))),
                path: ModPath::from(hir_expand::name!(doc)),
            },
        );
        let mut attrs = owner.attrs().peekable();
        let entries = if attrs.peek().is_none() && docs.is_none() {
            // Avoid heap allocation
            None
        } else {
            Some(attrs.flat_map(|ast| Attr::from_src(ast, hygiene)).chain(docs).collect())
        };
        Attrs { entries }
    }

    pub fn merge(&self, other: Attrs) -> Attrs {
        match (&self.entries, &other.entries) {
            (None, None) => Attrs { entries: None },
            (Some(entries), None) | (None, Some(entries)) => {
                Attrs { entries: Some(entries.clone()) }
            }
            (Some(a), Some(b)) => {
                Attrs { entries: Some(a.iter().chain(b.iter()).cloned().collect()) }
            }
        }
    }

    pub fn by_key(&self, key: &'static str) -> AttrQuery<'_> {
        AttrQuery { attrs: self, key }
    }

    pub fn cfg(&self) -> Option<CfgExpr> {
        // FIXME: handle cfg_attr :-)
        let mut cfgs = self.by_key("cfg").tt_values().map(CfgExpr::parse).collect::<Vec<_>>();
        match cfgs.len() {
            0 => None,
            1 => Some(cfgs.pop().unwrap()),
            _ => Some(CfgExpr::All(cfgs)),
        }
    }
    pub(crate) fn is_cfg_enabled(&self, cfg_options: &CfgOptions) -> bool {
        match self.cfg() {
            None => true,
            Some(cfg) => cfg_options.check(&cfg) != Some(false),
        }
    }

    pub fn docs(&self) -> Option<Documentation> {
        let mut docs = String::new();
        self.by_key("doc")
            .attrs()
            .flat_map(|attr| match attr.input.as_ref()? {
                AttrInput::Literal(s) => Some(s),
                AttrInput::TokenTree(_) => None,
            })
            .intersperse(&SmolStr::new_inline("\n"))
            // No FromIterator<SmolStr> for String
            .for_each(|s| docs.push_str(s.as_str()));
        if docs.is_empty() {
            None
        } else {
            Some(Documentation(docs.into()))
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Attr {
    pub(crate) path: ModPath,
    pub(crate) input: Option<AttrInput>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AttrInput {
    /// `#[attr = "string"]`
    Literal(SmolStr),
    /// `#[attr(subtree)]`
    TokenTree(Subtree),
}

impl Attr {
    fn from_src(ast: ast::Attr, hygiene: &Hygiene) -> Option<Attr> {
        let path = ModPath::from_src(ast.path()?, hygiene)?;
        let input = if let Some(lit) = ast.literal() {
            // FIXME: escape? raw string?
            let value = lit.syntax().first_token()?.text().trim_matches('"').into();
            Some(AttrInput::Literal(value))
        } else if let Some(tt) = ast.token_tree() {
            Some(AttrInput::TokenTree(ast_to_token_tree(&tt)?.0))
        } else {
            None
        };
        Some(Attr { path, input })
    }
}

#[derive(Debug, Clone, Copy)]
pub struct AttrQuery<'a> {
    attrs: &'a Attrs,
    key: &'static str,
}

impl<'a> AttrQuery<'a> {
    pub fn tt_values(self) -> impl Iterator<Item = &'a Subtree> {
        self.attrs().filter_map(|attr| match attr.input.as_ref()? {
            AttrInput::TokenTree(it) => Some(it),
            _ => None,
        })
    }

    pub fn string_value(self) -> Option<&'a SmolStr> {
        self.attrs().find_map(|attr| match attr.input.as_ref()? {
            AttrInput::Literal(it) => Some(it),
            _ => None,
        })
    }

    pub fn exists(self) -> bool {
        self.attrs().next().is_some()
    }

    fn attrs(self) -> impl Iterator<Item = &'a Attr> {
        let key = self.key;
        self.attrs
            .iter()
            .filter(move |attr| attr.path.as_ident().map_or(false, |s| s.to_string() == key))
    }
}

fn attrs_from_ast<N>(src: AstId<N>, db: &dyn DefDatabase) -> Attrs
where
    N: ast::AttrsOwner,
{
    let src = InFile::new(src.file_id, src.to_node(db.upcast()));
    Attrs::from_attrs_owner(db, src.as_ref().map(|it| it as &dyn AttrsOwner))
}

fn attrs_from_item_tree<N: ItemTreeNode>(id: ItemTreeId<N>, db: &dyn DefDatabase) -> Attrs {
    let tree = db.item_tree(id.file_id);
    let mod_item = N::id_to_mod_item(id.value);
    tree.attrs(mod_item.into()).clone()
}
<commit_msg>Remove raw pre and suffixes from string attr literals<commit_after>//! A higher level attributes based on TokenTree, with also some shortcuts.

use std::{ops, sync::Arc};

use cfg::{CfgExpr, CfgOptions};
use either::Either;
use hir_expand::{hygiene::Hygiene, AstId, InFile};
use itertools::Itertools;
use mbe::ast_to_token_tree;
use syntax::{
    ast::{self, AstNode, AttrsOwner},
    SmolStr,
};
use tt::Subtree;

use crate::{
    db::DefDatabase,
    item_tree::{ItemTreeId, ItemTreeNode},
    nameres::ModuleSource,
    path::ModPath,
    src::HasChildSource,
    AdtId, AttrDefId, Lookup,
};

/// Holds documentation
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Documentation(Arc<str>);

impl Documentation {
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl Into<String> for Documentation {
    fn into(self) -> String {
        self.as_str().to_owned()
    }
}

#[derive(Default, Debug, Clone, PartialEq, Eq)]
pub struct Attrs {
    entries: Option<Arc<[Attr]>>,
}

impl ops::Deref for Attrs {
    type Target = [Attr];

    fn deref(&self) -> &[Attr] {
        match &self.entries {
            Some(it) => &*it,
            None => &[],
        }
    }
}

impl Attrs {
    pub const EMPTY: Attrs = Attrs { entries: None };

    pub(crate) fn attrs_query(db: &dyn DefDatabase, def: AttrDefId) -> Attrs {
        match def {
            AttrDefId::ModuleId(module) => {
                let def_map = db.crate_def_map(module.krate);
                let mod_data = &def_map[module.local_id];
                match mod_data.declaration_source(db) {
                    Some(it) => {
                        Attrs::from_attrs_owner(db, it.as_ref().map(|it| it as &dyn AttrsOwner))
                    }
                    None => Attrs::from_attrs_owner(
                        db,
                        mod_data.definition_source(db).as_ref().map(|src| match src {
                            ModuleSource::SourceFile(file) => file as &dyn AttrsOwner,
                            ModuleSource::Module(module) => module as &dyn AttrsOwner,
                        }),
                    ),
                }
            }
            AttrDefId::FieldId(it) => {
                let src = it.parent.child_source(db);
                match &src.value[it.local_id] {
                    Either::Left(_tuple) => Attrs::default(),
                    Either::Right(record) => Attrs::from_attrs_owner(db, src.with_value(record)),
                }
            }
            AttrDefId::EnumVariantId(var_id) => {
                let src = var_id.parent.child_source(db);
                let src = src.as_ref().map(|it| &it[var_id.local_id]);
                Attrs::from_attrs_owner(db, src.map(|it| it as &dyn AttrsOwner))
            }
            AttrDefId::AdtId(it) => match it {
                AdtId::StructId(it) => attrs_from_item_tree(it.lookup(db).id, db),
                AdtId::EnumId(it) => attrs_from_item_tree(it.lookup(db).id, db),
                AdtId::UnionId(it) => attrs_from_item_tree(it.lookup(db).id, db),
            },
            AttrDefId::TraitId(it) => attrs_from_item_tree(it.lookup(db).id, db),
            AttrDefId::MacroDefId(it) => {
                it.ast_id.map_or_else(Default::default, |ast_id| attrs_from_ast(ast_id, db))
            }
            AttrDefId::ImplId(it) => attrs_from_item_tree(it.lookup(db).id, db),
            AttrDefId::ConstId(it) => attrs_from_item_tree(it.lookup(db).id, db),
            AttrDefId::StaticId(it) => attrs_from_item_tree(it.lookup(db).id, db),
            AttrDefId::FunctionId(it) => attrs_from_item_tree(it.lookup(db).id, db),
            AttrDefId::TypeAliasId(it) => attrs_from_item_tree(it.lookup(db).id, db),
        }
    }

    pub fn from_attrs_owner(db: &dyn DefDatabase, owner: InFile<&dyn AttrsOwner>) -> Attrs {
        let hygiene = Hygiene::new(db.upcast(), owner.file_id);
        Attrs::new(owner.value, &hygiene)
    }

    pub(crate) fn new(owner: &dyn AttrsOwner, hygiene: &Hygiene) -> Attrs {
        let docs = ast::CommentIter::from_syntax_node(owner.syntax()).doc_comment_text().map(
            |docs_text| Attr {
                input: Some(AttrInput::Literal(SmolStr::new(docs_text))),
                path: ModPath::from(hir_expand::name!(doc)),
            },
        );
        let mut attrs = owner.attrs().peekable();
        let entries = if attrs.peek().is_none() && docs.is_none() {
            // Avoid heap allocation
            None
        } else {
            Some(attrs.flat_map(|ast| Attr::from_src(ast, hygiene)).chain(docs).collect())
        };
        Attrs { entries }
    }

    pub fn merge(&self, other: Attrs) -> Attrs {
        match (&self.entries, &other.entries) {
            (None, None) => Attrs { entries: None },
            (Some(entries), None) | (None, Some(entries)) => {
                Attrs { entries: Some(entries.clone()) }
            }
            (Some(a), Some(b)) => {
                Attrs { entries: Some(a.iter().chain(b.iter()).cloned().collect()) }
            }
        }
    }

    pub fn by_key(&self, key: &'static str) -> AttrQuery<'_> {
        AttrQuery { attrs: self, key }
    }

    pub fn cfg(&self) -> Option<CfgExpr> {
        // FIXME: handle cfg_attr :-)
        let mut cfgs = self.by_key("cfg").tt_values().map(CfgExpr::parse).collect::<Vec<_>>();
        match cfgs.len() {
            0 => None,
            1 => Some(cfgs.pop().unwrap()),
            _ => Some(CfgExpr::All(cfgs)),
        }
    }
    pub(crate) fn is_cfg_enabled(&self, cfg_options: &CfgOptions) -> bool {
        match self.cfg() {
            None => true,
            Some(cfg) => cfg_options.check(&cfg) != Some(false),
        }
    }

    pub fn docs(&self) -> Option<Documentation> {
        let mut docs = String::new();
        self.by_key("doc")
            .attrs()
            .flat_map(|attr| match attr.input.as_ref()? {
                AttrInput::Literal(s) => Some(s),
                AttrInput::TokenTree(_) => None,
            })
            .intersperse(&SmolStr::new_inline("\n"))
            // No FromIterator<SmolStr> for String
            .for_each(|s| docs.push_str(s.as_str()));
        if docs.is_empty() {
            None
        } else {
            Some(Documentation(docs.into()))
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Attr {
    pub(crate) path: ModPath,
    pub(crate) input: Option<AttrInput>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AttrInput {
    /// `#[attr = "string"]`
    Literal(SmolStr),
    /// `#[attr(subtree)]`
    TokenTree(Subtree),
}

impl Attr {
    fn from_src(ast: ast::Attr, hygiene: &Hygiene) -> Option<Attr> {
        let path = ModPath::from_src(ast.path()?, hygiene)?;
        let input = if let Some(lit) = ast.literal() {
            let value = if let ast::LiteralKind::String(string) = lit.kind() {
                string.value()?.into()
            } else {
                lit.syntax().first_token()?.text().trim_matches('"').into()
            };
            Some(AttrInput::Literal(value))
        } else if let Some(tt) = ast.token_tree() {
            Some(AttrInput::TokenTree(ast_to_token_tree(&tt)?.0))
        } else {
            None
        };
        Some(Attr { path, input })
    }
}

#[derive(Debug, Clone, Copy)]
pub struct AttrQuery<'a> {
    attrs: &'a Attrs,
    key: &'static str,
}

impl<'a> AttrQuery<'a> {
    pub fn tt_values(self) -> impl Iterator<Item = &'a Subtree> {
        self.attrs().filter_map(|attr| match attr.input.as_ref()? {
            AttrInput::TokenTree(it) => Some(it),
            _ => None,
        })
    }

    pub fn string_value(self) -> Option<&'a SmolStr> {
        self.attrs().find_map(|attr| match attr.input.as_ref()? {
            AttrInput::Literal(it) => Some(it),
            _ => None,
        })
    }

    pub fn exists(self) -> bool {
        self.attrs().next().is_some()
    }

    fn attrs(self) -> impl Iterator<Item = &'a Attr> {
        let key = self.key;
        self.attrs
            .iter()
            .filter(move |attr| attr.path.as_ident().map_or(false, |s| s.to_string() == key))
    }
}

fn attrs_from_ast<N>(src: AstId<N>, db: &dyn DefDatabase) -> Attrs
where
    N: ast::AttrsOwner,
{
    let src = InFile::new(src.file_id, src.to_node(db.upcast()));
    Attrs::from_attrs_owner(db, src.as_ref().map(|it| it as &dyn AttrsOwner))
}

fn attrs_from_item_tree<N: ItemTreeNode>(id: ItemTreeId<N>, db: &dyn DefDatabase) -> Attrs {
    let tree = db.item_tree(id.file_id);
    let mod_item = N::id_to_mod_item(id.value);
    tree.attrs(mod_item.into()).clone()
}
<|endoftext|>


<commit_before>use std::sync::Arc;
use std::path::Path;
use std::collections::HashSet;

use rustc_hash::FxHashMap;

use ra_db::{
    CrateGraph, FileId, SourceRoot, SourceRootId, SourceDatabase, salsa,
};
use ra_hir::{db, HirInterner};
use ra_project_model::ProjectWorkspace;
use ra_vfs::{Vfs, VfsChange};

type Result<T> = std::result::Result<T, failure::Error>;

#[salsa::database(
    ra_db::SourceDatabaseStorage,
    db::HirDatabaseStorage,
    db::PersistentHirDatabaseStorage
)]
#[derive(Debug)]
pub struct BatchDatabase {
    runtime: salsa::Runtime<BatchDatabase>,
    interner: Arc<HirInterner>,
    // file_counter: u32,
}

impl salsa::Database for BatchDatabase {
    fn salsa_runtime(&self) -> &salsa::Runtime<BatchDatabase> {
        &self.runtime
    }
}

impl AsRef<HirInterner> for BatchDatabase {
    fn as_ref(&self) -> &HirInterner {
        &self.interner
    }
}

fn vfs_file_to_id(f: ra_vfs::VfsFile) -> FileId {
    FileId(f.0.into())
}
fn vfs_root_to_id(r: ra_vfs::VfsRoot) -> SourceRootId {
    SourceRootId(r.0.into())
}

impl BatchDatabase {
    pub fn load(crate_graph: CrateGraph, vfs: &mut Vfs) -> BatchDatabase {
        let mut db =
            BatchDatabase { runtime: salsa::Runtime::default(), interner: Default::default() };
        db.set_crate_graph(Arc::new(crate_graph));

        // wait until Vfs has loaded all roots
        let receiver = vfs.task_receiver().clone();
        let mut roots_loaded = HashSet::new();
        for task in receiver {
            vfs.handle_task(task);
            let mut done = false;
            for change in vfs.commit_changes() {
                match change {
                    VfsChange::AddRoot { root, files } => {
                        let source_root_id = vfs_root_to_id(root);
                        log::debug!(
                            "loaded source root {:?} with path {:?}",
                            source_root_id,
                            vfs.root2path(root)
                        );
                        let mut file_map = FxHashMap::default();
                        for (vfs_file, path, text) in files {
                            let file_id = vfs_file_to_id(vfs_file);
                            db.set_file_text(file_id, text);
                            db.set_file_relative_path(file_id, path.clone());
                            db.set_file_source_root(file_id, source_root_id);
                            file_map.insert(path, file_id);
                        }
                        let source_root = SourceRoot { files: file_map };
                        db.set_source_root(source_root_id, Arc::new(source_root));
                        roots_loaded.insert(source_root_id);
                        if roots_loaded.len() == vfs.num_roots() {
                            done = true;
                        }
                    }
                    VfsChange::AddFile { .. }
                    | VfsChange::RemoveFile { .. }
                    | VfsChange::ChangeFile { .. } => {
                        // log::warn!("VFS changed while loading");
                    }
                }
            }
            if done {
                break;
            }
        }

        db
    }

    pub fn load_cargo(root: impl AsRef<Path>) -> Result<(BatchDatabase, Vec<SourceRootId>)> {
        let root = root.as_ref().canonicalize()?;
        let ws = ProjectWorkspace::discover(root.as_ref())?;
        let mut roots = Vec::new();
        roots.push(root.clone());
        for pkg in ws.cargo.packages() {
            roots.push(pkg.root(&ws.cargo).to_path_buf());
        }
        for krate in ws.sysroot.crates() {
            roots.push(krate.root_dir(&ws.sysroot).to_path_buf())
        }
        let (mut vfs, roots) = Vfs::new(roots);
        let mut load = |path: &Path| {
            let vfs_file = vfs.load(path);
            log::debug!("vfs file {:?} -> {:?}", path, vfs_file);
            vfs_file.map(vfs_file_to_id)
        };
        let crate_graph = ws.to_crate_graph(&mut load);
        log::debug!("crate graph: {:?}", crate_graph);

        let local_roots = roots
            .into_iter()
            .filter(|r| vfs.root2path(*r).starts_with(&root))
            .map(vfs_root_to_id)
            .collect();

        let db = BatchDatabase::load(crate_graph, &mut vfs);
        let _ = vfs.shutdown();
        Ok((db, local_roots))
    }
}
<commit_msg>Add a smoke test for ra_batch<commit_after>use std::sync::Arc;
use std::path::Path;
use std::collections::HashSet;

use rustc_hash::FxHashMap;

use ra_db::{
    CrateGraph, FileId, SourceRoot, SourceRootId, SourceDatabase, salsa,
};
use ra_hir::{db, HirInterner};
use ra_project_model::ProjectWorkspace;
use ra_vfs::{Vfs, VfsChange};

type Result<T> = std::result::Result<T, failure::Error>;

#[salsa::database(
    ra_db::SourceDatabaseStorage,
    db::HirDatabaseStorage,
    db::PersistentHirDatabaseStorage
)]
#[derive(Debug)]
pub struct BatchDatabase {
    runtime: salsa::Runtime<BatchDatabase>,
    interner: Arc<HirInterner>,
    // file_counter: u32,
}

impl salsa::Database for BatchDatabase {
    fn salsa_runtime(&self) -> &salsa::Runtime<BatchDatabase> {
        &self.runtime
    }
}

impl AsRef<HirInterner> for BatchDatabase {
    fn as_ref(&self) -> &HirInterner {
        &self.interner
    }
}

fn vfs_file_to_id(f: ra_vfs::VfsFile) -> FileId {
    FileId(f.0.into())
}
fn vfs_root_to_id(r: ra_vfs::VfsRoot) -> SourceRootId {
    SourceRootId(r.0.into())
}

impl BatchDatabase {
    pub fn load(crate_graph: CrateGraph, vfs: &mut Vfs) -> BatchDatabase {
        let mut db =
            BatchDatabase { runtime: salsa::Runtime::default(), interner: Default::default() };
        db.set_crate_graph(Arc::new(crate_graph));

        // wait until Vfs has loaded all roots
        let receiver = vfs.task_receiver().clone();
        let mut roots_loaded = HashSet::new();
        for task in receiver {
            vfs.handle_task(task);
            let mut done = false;
            for change in vfs.commit_changes() {
                match change {
                    VfsChange::AddRoot { root, files } => {
                        let source_root_id = vfs_root_to_id(root);
                        log::debug!(
                            "loaded source root {:?} with path {:?}",
                            source_root_id,
                            vfs.root2path(root)
                        );
                        let mut file_map = FxHashMap::default();
                        for (vfs_file, path, text) in files {
                            let file_id = vfs_file_to_id(vfs_file);
                            db.set_file_text(file_id, text);
                            db.set_file_relative_path(file_id, path.clone());
                            db.set_file_source_root(file_id, source_root_id);
                            file_map.insert(path, file_id);
                        }
                        let source_root = SourceRoot { files: file_map };
                        db.set_source_root(source_root_id, Arc::new(source_root));
                        roots_loaded.insert(source_root_id);
                        if roots_loaded.len() == vfs.num_roots() {
                            done = true;
                        }
                    }
                    VfsChange::AddFile { .. }
                    | VfsChange::RemoveFile { .. }
                    | VfsChange::ChangeFile { .. } => {
                        // log::warn!("VFS changed while loading");
                    }
                }
            }
            if done {
                break;
            }
        }

        db
    }

    pub fn load_cargo(root: impl AsRef<Path>) -> Result<(BatchDatabase, Vec<SourceRootId>)> {
        let root = root.as_ref().canonicalize()?;
        let ws = ProjectWorkspace::discover(root.as_ref())?;
        let mut roots = Vec::new();
        roots.push(root.clone());
        for pkg in ws.cargo.packages() {
            roots.push(pkg.root(&ws.cargo).to_path_buf());
        }
        for krate in ws.sysroot.crates() {
            roots.push(krate.root_dir(&ws.sysroot).to_path_buf())
        }
        let (mut vfs, roots) = Vfs::new(roots);
        let mut load = |path: &Path| {
            let vfs_file = vfs.load(path);
            log::debug!("vfs file {:?} -> {:?}", path, vfs_file);
            vfs_file.map(vfs_file_to_id)
        };
        let crate_graph = ws.to_crate_graph(&mut load);
        log::debug!("crate graph: {:?}", crate_graph);

        let local_roots = roots
            .into_iter()
            .filter(|r| vfs.root2path(*r).starts_with(&root))
            .map(vfs_root_to_id)
            .collect();

        let db = BatchDatabase::load(crate_graph, &mut vfs);
        let _ = vfs.shutdown();
        Ok((db, local_roots))
    }
}

#[cfg(test)]
mod tests {
    use ra_hir::Crate;
    use super::*;

    #[test]
    fn test_loading_rust_analyzer() {
        let mut path = std::env::current_exe().unwrap();
        while !path.join("Cargo.toml").is_file() {
            path = path.parent().unwrap().to_owned();
        }
        let (db, roots) = BatchDatabase::load_cargo(path).unwrap();
        let mut num_crates = 0;
        for root in roots {
            for _krate in Crate::source_root_crates(&db, root) {
                num_crates += 1;
            }
        }

        // RA has quite a few crates, but the exact count doesn't matter
        assert!(num_crates > 20);
    }
}
<|endoftext|>