<commit_before>//! A simple crate for executing work on a thread pool, and getting back a
//! future.
//!
//! This crate provides a simple thread pool abstraction for running work
//! externally from the current thread that's running. An instance of `Future`
//! is handed back to represent that the work may be done later, and further
//! computations can be chained along with it as well.
//!
//! ```rust
//! extern crate futures;
//! extern crate futures_cpupool;
//!
//! use futures::Future;
//! use futures_cpupool::CpuPool;
//!
//! # fn long_running_future(a: u32) -> futures::future::BoxFuture<u32, ()> {
//! #     futures::future::result(Ok(a)).boxed()
//! # }
//! # fn main() {
//!
//! // Create a worker thread pool with four threads
//! let pool = CpuPool::new(4);
//!
//! // Execute some work on the thread pool, optionally closing over data.
//! let a = pool.spawn(long_running_future(2));
//! let b = pool.spawn(long_running_future(100));
//!
//! // Express some further computation once the work is completed on the thread
//! // pool.
//! let c = a.join(b).map(|(a, b)| a + b).wait().unwrap();
//!
//! // Print out the result
//! println!("{:?}", c);
//! # }
//! ```

#![deny(missing_docs)]
#![deny(missing_debug_implementations)]

extern crate futures;
extern crate num_cpus;

use std::panic::{self, AssertUnwindSafe};
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc;
use std::thread;
use std::fmt;

use futures::{IntoFuture, Future, Poll, Async};
use futures::future::{lazy, Executor, ExecuteError};
use futures::sync::oneshot::{channel, Sender, Receiver};
use futures::executor::{self, Run, Executor as OldExecutor};

/// A thread pool intended to run CPU intensive work.
///
/// This thread pool will hand out futures representing the completed work
/// that happens on the thread pool itself, and the futures can then be later
/// composed with other work as part of an overall computation.
///
/// The worker threads associated with a thread pool are kept alive so long as
/// there is an open handle to the `CpuPool` or there is work running on them. Once
/// all work has been drained and all references have gone away the worker
/// threads will be shut down.
///
/// Currently `CpuPool` implements `Clone` which just clones a new reference to
/// the underlying thread pool.
///
/// **Note:** if you use CpuPool inside a library it's better accept a
/// `Builder` object for thread configuration rather than configuring just
/// pool size.  This not only future proof for other settings but also allows
/// user to attach monitoring tools to lifecycle hooks.
pub struct CpuPool {
    inner: Arc<Inner>,
}

/// Thread pool configuration object
///
/// Builder starts with a number of workers equal to the number
/// of CPUs on the host. But you can change it until you call `create()`.
pub struct Builder {
    pool_size: usize,
    name_prefix: Option<String>,
    after_start: Option<Arc<Fn() + Send + Sync>>,
    before_stop: Option<Arc<Fn() + Send + Sync>>,
}

struct MySender<F, T> {
    fut: F,
    tx: Option<Sender<T>>,
    keep_running_flag: Arc<AtomicBool>,
}

trait AssertSendSync: Send + Sync {}
impl AssertSendSync for CpuPool {}

struct Inner {
    tx: Mutex<mpsc::Sender<Message>>,
    rx: Mutex<mpsc::Receiver<Message>>,
    cnt: AtomicUsize,
    size: usize,
}

impl fmt::Debug for CpuPool {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("CpuPool")
            .field("size", &self.inner.size)
            .finish()
    }
}

impl fmt::Debug for Builder {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Builder")
            .field("pool_size", &self.pool_size)
            .field("name_prefix", &self.name_prefix)
            .finish()
    }
}

/// The type of future returned from the `CpuPool::spawn` function, which
/// proxies the futures running on the thread pool.
///
/// This future will resolve in the same way as the underlying future, and it
/// will propagate panics.
#[must_use]
#[derive(Debug)]
pub struct CpuFuture<T, E> {
    inner: Receiver<thread::Result<Result<T, E>>>,
    keep_running_flag: Arc<AtomicBool>,
}

enum Message {
    Run(Run),
    Close,
}

impl CpuPool {
    /// Creates a new thread pool with `size` worker threads associated with it.
    ///
    /// The returned handle can use `execute` to run work on this thread pool,
    /// and clones can be made of it to get multiple references to the same
    /// thread pool.
    ///
    /// This is a shortcut for:
    ///
    /// ```rust
    /// # use futures_cpupool::{Builder, CpuPool};
    /// #
    /// # fn new(size: usize) -> CpuPool {
    /// Builder::new().pool_size(size).create()
    /// # }
    /// ```
    ///
    /// # Panics
    ///
    /// Panics if `size == 0`.
    pub fn new(size: usize) -> CpuPool {
        Builder::new().pool_size(size).create()
    }

    /// Creates a new thread pool with a number of workers equal to the number
    /// of CPUs on the host.
    ///
    /// This is a shortcut for:
    ///
    /// ```rust
    /// # use futures_cpupool::{Builder, CpuPool};
    /// #
    /// # fn new_num_cpus() -> CpuPool {
    /// Builder::new().create()
    /// # }
    /// ```
    pub fn new_num_cpus() -> CpuPool {
        Builder::new().create()
    }

    /// Spawns a future to run on this thread pool, returning a future
    /// representing the produced value.
    ///
    /// This function will execute the future `f` on the associated thread
    /// pool, and return a future representing the finished computation. The
    /// returned future serves as a proxy to the computation that `F` is
    /// running.
    ///
    /// To simply run an arbitrary closure on a thread pool and extract the
    /// result, you can use the `future::lazy` combinator to defer work to
    /// executing on the thread pool itself.
    ///
    /// Note that if the future `f` panics it will be caught by default and the
    /// returned future will propagate the panic. That is, panics will not tear
    /// down the thread pool and will be propagated to the returned future's
    /// `poll` method if queried.
    ///
    /// If the returned future is dropped then this `CpuPool` will attempt to
    /// cancel the computation, if possible. That is, if the computation is in
    /// the middle of working, it will be interrupted when possible.
    pub fn spawn<F>(&self, f: F) -> CpuFuture<F::Item, F::Error>
        where F: Future + Send + 'static,
              F::Item: Send + 'static,
              F::Error: Send + 'static,
    {
        let (tx, rx) = channel();
        let keep_running_flag = Arc::new(AtomicBool::new(false));
        // AssertUnwindSafe is used here because `Send + 'static` is basically
        // an alias for an implementation of the `UnwindSafe` trait but we can't
        // express that in the standard library right now.
        let sender = MySender {
            fut: AssertUnwindSafe(f).catch_unwind(),
            tx: Some(tx),
            keep_running_flag: keep_running_flag.clone(),
        };
        executor::spawn(sender).execute(self.inner.clone());
        CpuFuture { inner: rx , keep_running_flag: keep_running_flag.clone() }
    }

    /// Spawns a closure on this thread pool.
    ///
    /// This function is a convenience wrapper around the `spawn` function above
    /// for running a closure wrapped in `future::lazy`. It will spawn the
    /// function `f` provided onto the thread pool, and continue to run the
    /// future returned by `f` on the thread pool as well.
    ///
    /// The returned future will be a handle to the result produced by the
    /// future that `f` returns.
    pub fn spawn_fn<F, R>(&self, f: F) -> CpuFuture<R::Item, R::Error>
        where F: FnOnce() -> R + Send + 'static,
              R: IntoFuture + 'static,
              R::Future: Send + 'static,
              R::Item: Send + 'static,
              R::Error: Send + 'static,
    {
        self.spawn(lazy(f))
    }
}

impl<F> Executor<F> for CpuPool
    where F: Future<Item = (), Error = ()> + Send + 'static,
{
    fn execute(&self, future: F) -> Result<(), ExecuteError<F>> {
        executor::spawn(future).execute(self.inner.clone());
        Ok(())
    }
}

impl Inner {
    fn send(&self, msg: Message) {
        self.tx.lock().unwrap().send(msg).unwrap();
    }

    fn work(&self, after_start: Option<Arc<Fn() + Send + Sync>>, before_stop: Option<Arc<Fn() + Send + Sync>>) {
        after_start.map(|fun| fun());
        loop {
            let msg = self.rx.lock().unwrap().recv().unwrap();
            match msg {
                Message::Run(r) => r.run(),
                Message::Close => break,
            }
        }
        before_stop.map(|fun| fun());
    }
}

impl Clone for CpuPool {
    fn clone(&self) -> CpuPool {
        self.inner.cnt.fetch_add(1, Ordering::Relaxed);
        CpuPool { inner: self.inner.clone() }
    }
}

impl Drop for CpuPool {
    fn drop(&mut self) {
        if self.inner.cnt.fetch_sub(1, Ordering::Relaxed) == 1 {
            for _ in 0..self.inner.size {
                self.inner.send(Message::Close);
            }
        }
    }
}

impl OldExecutor for Inner {
    fn execute(&self, run: Run) {
        self.send(Message::Run(run))
    }
}

impl<T, E> CpuFuture<T, E> {
    /// Drop this future without canceling the underlying future.
    ///
    /// When `CpuFuture` is dropped, `CpuPool` will try to abort the underlying
    /// future. This function can be used when user wants to drop but keep
    /// executing the underlying future.
    pub fn forget(self) {
        self.keep_running_flag.store(true, Ordering::SeqCst);
    }
}

impl<T: Send + 'static, E: Send + 'static> Future for CpuFuture<T, E> {
    type Item = T;
    type Error = E;

    fn poll(&mut self) -> Poll<T, E> {
        match self.inner.poll().expect("shouldn't be canceled") {
            Async::Ready(Ok(Ok(e))) => Ok(e.into()),
            Async::Ready(Ok(Err(e))) => Err(e),
            Async::Ready(Err(e)) => panic::resume_unwind(e),
            Async::NotReady => Ok(Async::NotReady),
        }
    }
}

impl<F: Future> Future for MySender<F, Result<F::Item, F::Error>> {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<(), ()> {
        if let Ok(Async::Ready(_)) = self.tx.as_mut().unwrap().poll_cancel() {
            if !self.keep_running_flag.load(Ordering::SeqCst) {
                // Cancelled, bail out
                return Ok(().into())
            }
        }

        let res = match self.fut.poll() {
            Ok(Async::Ready(e)) => Ok(e),
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Err(e) => Err(e),
        };

        // if the receiving end has gone away then that's ok, we just ignore the
        // send error here.
        drop(self.tx.take().unwrap().send(res));
        Ok(Async::Ready(()))
    }
}

impl Builder {
    /// Create a builder a number of workers equal to the number
    /// of CPUs on the host.
    pub fn new() -> Builder {
        Builder {
            pool_size: num_cpus::get(),
            name_prefix: None,
            after_start: None,
            before_stop: None,
        }
    }

    /// Set size of a future CpuPool
    ///
    /// The size of a thread pool is the number of worker threads spawned
    pub fn pool_size(&mut self, size: usize) -> &mut Self {
        self.pool_size = size;
        self
    }

    /// Set thread name prefix of a future CpuPool
    ///
    /// Thread name prefix is used for generating thread names. For example, if prefix is
    /// `my-pool-`, then threads in the pool will get names like `my-pool-1` etc.
    pub fn name_prefix<S: Into<String>>(&mut self, name_prefix: S) -> &mut Self {
        self.name_prefix = Some(name_prefix.into());
        self
    }

    /// Execute function `f` right after each thread is started but before
    /// running any jobs on it
    ///
    /// This is initially intended for bookkeeping and monitoring uses
    pub fn after_start<F>(&mut self, f: F) -> &mut Self
        where F: Fn() + Send + Sync + 'static
    {
        self.after_start = Some(Arc::new(f));
        self
    }

    /// Execute function `f` before each worker thread stops
    ///
    /// This is initially intended for bookkeeping and monitoring uses
    pub fn before_stop<F>(&mut self, f: F) -> &mut Self
        where F: Fn() + Send + Sync + 'static
    {
        self.before_stop = Some(Arc::new(f));
        self
    }

    /// Create CpuPool with configured parameters
    ///
    /// # Panics
    ///
    /// Panics if `pool_size == 0`.
    pub fn create(&mut self) -> CpuPool {
        let (tx, rx) = mpsc::channel();
        let pool = CpuPool {
            inner: Arc::new(Inner {
                tx: Mutex::new(tx),
                rx: Mutex::new(rx),
                cnt: AtomicUsize::new(1),
                size: self.pool_size,
            }),
        };
        assert!(self.pool_size > 0);

        for counter in 0..self.pool_size {
            let inner = pool.inner.clone();
            let after_start = self.after_start.clone();
            let before_stop = self.before_stop.clone();
            let mut thread_builder = thread::Builder::new();
            if let Some(ref name_prefix) = self.name_prefix {
                thread_builder = thread_builder.name(format!("{}{}", name_prefix, counter));
            }
            thread_builder.spawn(move || inner.work(after_start, before_stop)).unwrap();
        }
        return pool
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::mpsc;

    #[test]
    fn test_drop_after_start() {
        let (tx, rx) = mpsc::sync_channel(2);
        let cpu_pool = Builder::new()
            .pool_size(2)
            .after_start(move || tx.send(1).unwrap()).create();

        // After Builder is deconstructed, the tx should be droped
        // so that we can use rx as an iterator.
        let count = rx.into_iter().count();
        assert_eq!(count, 2);
    }
}
<commit_msg>update comments of `after_start` and `before_stop`.<commit_after>//! A simple crate for executing work on a thread pool, and getting back a
//! future.
//!
//! This crate provides a simple thread pool abstraction for running work
//! externally from the current thread that's running. An instance of `Future`
//! is handed back to represent that the work may be done later, and further
//! computations can be chained along with it as well.
//!
//! ```rust
//! extern crate futures;
//! extern crate futures_cpupool;
//!
//! use futures::Future;
//! use futures_cpupool::CpuPool;
//!
//! # fn long_running_future(a: u32) -> futures::future::BoxFuture<u32, ()> {
//! #     futures::future::result(Ok(a)).boxed()
//! # }
//! # fn main() {
//!
//! // Create a worker thread pool with four threads
//! let pool = CpuPool::new(4);
//!
//! // Execute some work on the thread pool, optionally closing over data.
//! let a = pool.spawn(long_running_future(2));
//! let b = pool.spawn(long_running_future(100));
//!
//! // Express some further computation once the work is completed on the thread
//! // pool.
//! let c = a.join(b).map(|(a, b)| a + b).wait().unwrap();
//!
//! // Print out the result
//! println!("{:?}", c);
//! # }
//! ```

#![deny(missing_docs)]
#![deny(missing_debug_implementations)]

extern crate futures;
extern crate num_cpus;

use std::panic::{self, AssertUnwindSafe};
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc;
use std::thread;
use std::fmt;

use futures::{IntoFuture, Future, Poll, Async};
use futures::future::{lazy, Executor, ExecuteError};
use futures::sync::oneshot::{channel, Sender, Receiver};
use futures::executor::{self, Run, Executor as OldExecutor};

/// A thread pool intended to run CPU intensive work.
///
/// This thread pool will hand out futures representing the completed work
/// that happens on the thread pool itself, and the futures can then be later
/// composed with other work as part of an overall computation.
///
/// The worker threads associated with a thread pool are kept alive so long as
/// there is an open handle to the `CpuPool` or there is work running on them. Once
/// all work has been drained and all references have gone away the worker
/// threads will be shut down.
///
/// Currently `CpuPool` implements `Clone` which just clones a new reference to
/// the underlying thread pool.
///
/// **Note:** if you use CpuPool inside a library it's better accept a
/// `Builder` object for thread configuration rather than configuring just
/// pool size.  This not only future proof for other settings but also allows
/// user to attach monitoring tools to lifecycle hooks.
pub struct CpuPool {
    inner: Arc<Inner>,
}

/// Thread pool configuration object
///
/// Builder starts with a number of workers equal to the number
/// of CPUs on the host. But you can change it until you call `create()`.
pub struct Builder {
    pool_size: usize,
    name_prefix: Option<String>,
    after_start: Option<Arc<Fn() + Send + Sync>>,
    before_stop: Option<Arc<Fn() + Send + Sync>>,
}

struct MySender<F, T> {
    fut: F,
    tx: Option<Sender<T>>,
    keep_running_flag: Arc<AtomicBool>,
}

trait AssertSendSync: Send + Sync {}
impl AssertSendSync for CpuPool {}

struct Inner {
    tx: Mutex<mpsc::Sender<Message>>,
    rx: Mutex<mpsc::Receiver<Message>>,
    cnt: AtomicUsize,
    size: usize,
}

impl fmt::Debug for CpuPool {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("CpuPool")
            .field("size", &self.inner.size)
            .finish()
    }
}

impl fmt::Debug for Builder {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Builder")
            .field("pool_size", &self.pool_size)
            .field("name_prefix", &self.name_prefix)
            .finish()
    }
}

/// The type of future returned from the `CpuPool::spawn` function, which
/// proxies the futures running on the thread pool.
///
/// This future will resolve in the same way as the underlying future, and it
/// will propagate panics.
#[must_use]
#[derive(Debug)]
pub struct CpuFuture<T, E> {
    inner: Receiver<thread::Result<Result<T, E>>>,
    keep_running_flag: Arc<AtomicBool>,
}

enum Message {
    Run(Run),
    Close,
}

impl CpuPool {
    /// Creates a new thread pool with `size` worker threads associated with it.
    ///
    /// The returned handle can use `execute` to run work on this thread pool,
    /// and clones can be made of it to get multiple references to the same
    /// thread pool.
    ///
    /// This is a shortcut for:
    ///
    /// ```rust
    /// # use futures_cpupool::{Builder, CpuPool};
    /// #
    /// # fn new(size: usize) -> CpuPool {
    /// Builder::new().pool_size(size).create()
    /// # }
    /// ```
    ///
    /// # Panics
    ///
    /// Panics if `size == 0`.
    pub fn new(size: usize) -> CpuPool {
        Builder::new().pool_size(size).create()
    }

    /// Creates a new thread pool with a number of workers equal to the number
    /// of CPUs on the host.
    ///
    /// This is a shortcut for:
    ///
    /// ```rust
    /// # use futures_cpupool::{Builder, CpuPool};
    /// #
    /// # fn new_num_cpus() -> CpuPool {
    /// Builder::new().create()
    /// # }
    /// ```
    pub fn new_num_cpus() -> CpuPool {
        Builder::new().create()
    }

    /// Spawns a future to run on this thread pool, returning a future
    /// representing the produced value.
    ///
    /// This function will execute the future `f` on the associated thread
    /// pool, and return a future representing the finished computation. The
    /// returned future serves as a proxy to the computation that `F` is
    /// running.
    ///
    /// To simply run an arbitrary closure on a thread pool and extract the
    /// result, you can use the `future::lazy` combinator to defer work to
    /// executing on the thread pool itself.
    ///
    /// Note that if the future `f` panics it will be caught by default and the
    /// returned future will propagate the panic. That is, panics will not tear
    /// down the thread pool and will be propagated to the returned future's
    /// `poll` method if queried.
    ///
    /// If the returned future is dropped then this `CpuPool` will attempt to
    /// cancel the computation, if possible. That is, if the computation is in
    /// the middle of working, it will be interrupted when possible.
    pub fn spawn<F>(&self, f: F) -> CpuFuture<F::Item, F::Error>
        where F: Future + Send + 'static,
              F::Item: Send + 'static,
              F::Error: Send + 'static,
    {
        let (tx, rx) = channel();
        let keep_running_flag = Arc::new(AtomicBool::new(false));
        // AssertUnwindSafe is used here because `Send + 'static` is basically
        // an alias for an implementation of the `UnwindSafe` trait but we can't
        // express that in the standard library right now.
        let sender = MySender {
            fut: AssertUnwindSafe(f).catch_unwind(),
            tx: Some(tx),
            keep_running_flag: keep_running_flag.clone(),
        };
        executor::spawn(sender).execute(self.inner.clone());
        CpuFuture { inner: rx , keep_running_flag: keep_running_flag.clone() }
    }

    /// Spawns a closure on this thread pool.
    ///
    /// This function is a convenience wrapper around the `spawn` function above
    /// for running a closure wrapped in `future::lazy`. It will spawn the
    /// function `f` provided onto the thread pool, and continue to run the
    /// future returned by `f` on the thread pool as well.
    ///
    /// The returned future will be a handle to the result produced by the
    /// future that `f` returns.
    pub fn spawn_fn<F, R>(&self, f: F) -> CpuFuture<R::Item, R::Error>
        where F: FnOnce() -> R + Send + 'static,
              R: IntoFuture + 'static,
              R::Future: Send + 'static,
              R::Item: Send + 'static,
              R::Error: Send + 'static,
    {
        self.spawn(lazy(f))
    }
}

impl<F> Executor<F> for CpuPool
    where F: Future<Item = (), Error = ()> + Send + 'static,
{
    fn execute(&self, future: F) -> Result<(), ExecuteError<F>> {
        executor::spawn(future).execute(self.inner.clone());
        Ok(())
    }
}

impl Inner {
    fn send(&self, msg: Message) {
        self.tx.lock().unwrap().send(msg).unwrap();
    }

    fn work(&self, after_start: Option<Arc<Fn() + Send + Sync>>, before_stop: Option<Arc<Fn() + Send + Sync>>) {
        after_start.map(|fun| fun());
        loop {
            let msg = self.rx.lock().unwrap().recv().unwrap();
            match msg {
                Message::Run(r) => r.run(),
                Message::Close => break,
            }
        }
        before_stop.map(|fun| fun());
    }
}

impl Clone for CpuPool {
    fn clone(&self) -> CpuPool {
        self.inner.cnt.fetch_add(1, Ordering::Relaxed);
        CpuPool { inner: self.inner.clone() }
    }
}

impl Drop for CpuPool {
    fn drop(&mut self) {
        if self.inner.cnt.fetch_sub(1, Ordering::Relaxed) == 1 {
            for _ in 0..self.inner.size {
                self.inner.send(Message::Close);
            }
        }
    }
}

impl OldExecutor for Inner {
    fn execute(&self, run: Run) {
        self.send(Message::Run(run))
    }
}

impl<T, E> CpuFuture<T, E> {
    /// Drop this future without canceling the underlying future.
    ///
    /// When `CpuFuture` is dropped, `CpuPool` will try to abort the underlying
    /// future. This function can be used when user wants to drop but keep
    /// executing the underlying future.
    pub fn forget(self) {
        self.keep_running_flag.store(true, Ordering::SeqCst);
    }
}

impl<T: Send + 'static, E: Send + 'static> Future for CpuFuture<T, E> {
    type Item = T;
    type Error = E;

    fn poll(&mut self) -> Poll<T, E> {
        match self.inner.poll().expect("shouldn't be canceled") {
            Async::Ready(Ok(Ok(e))) => Ok(e.into()),
            Async::Ready(Ok(Err(e))) => Err(e),
            Async::Ready(Err(e)) => panic::resume_unwind(e),
            Async::NotReady => Ok(Async::NotReady),
        }
    }
}

impl<F: Future> Future for MySender<F, Result<F::Item, F::Error>> {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<(), ()> {
        if let Ok(Async::Ready(_)) = self.tx.as_mut().unwrap().poll_cancel() {
            if !self.keep_running_flag.load(Ordering::SeqCst) {
                // Cancelled, bail out
                return Ok(().into())
            }
        }

        let res = match self.fut.poll() {
            Ok(Async::Ready(e)) => Ok(e),
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Err(e) => Err(e),
        };

        // if the receiving end has gone away then that's ok, we just ignore the
        // send error here.
        drop(self.tx.take().unwrap().send(res));
        Ok(Async::Ready(()))
    }
}

impl Builder {
    /// Create a builder a number of workers equal to the number
    /// of CPUs on the host.
    pub fn new() -> Builder {
        Builder {
            pool_size: num_cpus::get(),
            name_prefix: None,
            after_start: None,
            before_stop: None,
        }
    }

    /// Set size of a future CpuPool
    ///
    /// The size of a thread pool is the number of worker threads spawned
    pub fn pool_size(&mut self, size: usize) -> &mut Self {
        self.pool_size = size;
        self
    }

    /// Set thread name prefix of a future CpuPool
    ///
    /// Thread name prefix is used for generating thread names. For example, if prefix is
    /// `my-pool-`, then threads in the pool will get names like `my-pool-1` etc.
    pub fn name_prefix<S: Into<String>>(&mut self, name_prefix: S) -> &mut Self {
        self.name_prefix = Some(name_prefix.into());
        self
    }

    /// Execute function `f` right after each thread is started but before
    /// running any jobs on it.
    ///
    /// This is initially intended for bookkeeping and monitoring uses.
    /// The `f` will be deconstructed after the `builder` is deconstructed
    /// and all threads in the pool has executed it.
    pub fn after_start<F>(&mut self, f: F) -> &mut Self
        where F: Fn() + Send + Sync + 'static
    {
        self.after_start = Some(Arc::new(f));
        self
    }

    /// Execute function `f` before each worker thread stops.
    ///
    /// This is initially intended for bookkeeping and monitoring uses.
    /// The `f` will be deconstructed after the `builder` is deconstructed
    /// and all threads in the pool has executed it.
    pub fn before_stop<F>(&mut self, f: F) -> &mut Self
        where F: Fn() + Send + Sync + 'static
    {
        self.before_stop = Some(Arc::new(f));
        self
    }

    /// Create CpuPool with configured parameters
    ///
    /// # Panics
    ///
    /// Panics if `pool_size == 0`.
    pub fn create(&mut self) -> CpuPool {
        let (tx, rx) = mpsc::channel();
        let pool = CpuPool {
            inner: Arc::new(Inner {
                tx: Mutex::new(tx),
                rx: Mutex::new(rx),
                cnt: AtomicUsize::new(1),
                size: self.pool_size,
            }),
        };
        assert!(self.pool_size > 0);

        for counter in 0..self.pool_size {
            let inner = pool.inner.clone();
            let after_start = self.after_start.clone();
            let before_stop = self.before_stop.clone();
            let mut thread_builder = thread::Builder::new();
            if let Some(ref name_prefix) = self.name_prefix {
                thread_builder = thread_builder.name(format!("{}{}", name_prefix, counter));
            }
            thread_builder.spawn(move || inner.work(after_start, before_stop)).unwrap();
        }
        return pool
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::mpsc;

    #[test]
    fn test_drop_after_start() {
        let (tx, rx) = mpsc::sync_channel(2);
        let _cpu_pool = Builder::new()
            .pool_size(2)
            .after_start(move || tx.send(1).unwrap()).create();

        // After Builder is deconstructed, the tx should be droped
        // so that we can use rx as an iterator.
        let count = rx.into_iter().count();
        assert_eq!(count, 2);
    }
}
<|endoftext|>


<commit_before>use std::fs::{OpenOptions, File};
use std::io::Read;
use std::result::Result::{self, Ok, Err};

use libmodbus_rs::{Modbus, ModbusClient, ModbusRTU, SerialMode, Timeout};

use retry::{Retry, RetryError};

use hex_slice::AsHex;

use ::sunsaver::*;

pub trait SunSaverConnection {
    fn read_registers(&mut self) ->  [u16; 44];

    fn read_logged_data(&mut self) -> [u16; 32 * 16];

    fn read_response(&mut self) -> SunSaverResponse {
        SunSaverResponse::from_raw_bits(self.read_registers())
    }
}

pub struct ModbusSunSaverConnection {
    connection: Modbus,
}

impl ModbusSunSaverConnection {
    pub fn open(device: &str) -> ModbusSunSaverConnection {
        /* A Meterbus to Serial Converter (MSC) is required to adapt the Meter interface to an isolated RS-232 interface**.
           The SunSaver MPPT supports RTU mode only.
           16bit MODBUS® addresses (per the modbus.org spec)
           The serial communication parameters are:
             * BPS: 9600 baud
             * Parity: None
             * Data bits: 8
             * Stop bits: 2
             * Flow control: None
            All addresses listed are for the request PDU.
            The SunSaver MPPT default server address: 0x01. */
        debug!("Configuring device {}", device);
        let mut connection = Modbus::new_rtu(device, 9600, 'N', 8, 2).unwrap();
        assert!(connection.set_slave(0x01).is_ok());
        assert!(connection.rtu_set_serial_mode(SerialMode::MODBUS_RTU_RS232).is_ok());
        assert!(connection.set_response_timeout( Timeout { sec: 1, usec: 0 } ).is_ok());
        connection.set_debug(false).unwrap();

        let timeout = connection.get_response_timeout();
        info!("Timout {:?}", timeout);

        assert!(connection.connect().is_ok());
        debug!("Connected");

        ModbusSunSaverConnection {
            connection: connection,
        }
    }

    fn read_registers_retry(&self, address: i32, num_bit: i32, dest: &mut [u16]) -> Result<usize, RetryError> {
        match Retry::new(
            &mut || self.connection.read_registers(address, num_bit, dest),
            &mut |response| response.is_ok()
            ).try(3).wait(100).execute() {
            Ok(response) => Ok(response.unwrap() as usize),
            Err(error) => Err(error),
        }
    }
}

impl SunSaverConnection for ModbusSunSaverConnection {
    fn read_registers(&mut self) -> [u16; 44] {
        let mut response_register = [0u16; 44 as usize];
        let mut num_read_bytes = 0;
        num_read_bytes += self.read_registers_retry(0x08, 22, &mut response_register[0..22]).unwrap();
        num_read_bytes += self.read_registers_retry(0x1E, 22, &mut response_register[23..44]).unwrap();
        //if num_read_bytes != 44 {
        //    panic!("Failed to read all registers! Required 44 got {}", num_read_bytes);
        //}
        debug!("Read {} bytes", num_read_bytes);
        debug!("read reg 0x08 + 44: {:#x}", response_register.as_hex());

        response_register
    }

    fn read_logged_data(&mut self) -> [u16; 32 * 16] {
        let mut logged_data = [0u16; (32 * 16) as usize];

        for i in 0..32 {
            let offset: usize = i * 16;
            self.read_registers_retry((0x8000 + offset) as i32, 16, &mut logged_data[offset..(offset + 16)]).unwrap();
        }

        debug!("logged_data_start");
        for i in (0 as usize)..32 {
            debug!("{:#x}", logged_data[i..(i+16)].as_hex());
        }
        debug!("logged_data_end");

        logged_data
    }
}

#[derive(Debug)]
pub struct FileSunSaverConnection {
    file: File,
}

impl FileSunSaverConnection {
    pub fn open(filename: &str) -> FileSunSaverConnection {
        let file = OpenOptions::new().read(true).write(false).open(filename).unwrap();

        FileSunSaverConnection {
            file: file,
        }
    }
}

impl SunSaverConnection for FileSunSaverConnection {
    fn read_registers(&mut self) -> [u16; 44] {
        let mut response_register_u8 = [0u8; 88 as usize];
        assert!(self.file.read_exact(&mut response_register_u8).is_ok());

        let response_register_vec_u16: Vec<u16> = response_register_u8.chunks(2)
        .map(|items| {
            u16::from_le(((items[0] as u16) << 8) + (items[1] as u16))
        })
        .collect();

        let mut response_register_u16 = [0u16; 44 as usize];
        response_register_u16.clone_from_slice(&response_register_vec_u16);
        response_register_u16
    }

    fn read_logged_data(&mut self) -> [u16; 32 * 16] {
        unimplemented!();
    }
}<commit_msg>Read in single go<commit_after>use std::fs::{OpenOptions, File};
use std::io::Read;
use std::result::Result::{self, Ok, Err};

use libmodbus_rs::{Modbus, ModbusClient, ModbusRTU, SerialMode, Timeout};

use retry::{Retry, RetryError};

use hex_slice::AsHex;

use ::sunsaver::*;

pub trait SunSaverConnection {
    fn read_registers(&mut self) ->  [u16; 44];

    fn read_logged_data(&mut self) -> [u16; 32 * 16];

    fn read_response(&mut self) -> SunSaverResponse {
        SunSaverResponse::from_raw_bits(self.read_registers())
    }
}

pub struct ModbusSunSaverConnection {
    connection: Modbus,
}

impl ModbusSunSaverConnection {
    pub fn open(device: &str) -> ModbusSunSaverConnection {
        /* A Meterbus to Serial Converter (MSC) is required to adapt the Meter interface to an isolated RS-232 interface**.
           The SunSaver MPPT supports RTU mode only.
           16bit MODBUS® addresses (per the modbus.org spec)
           The serial communication parameters are:
             * BPS: 9600 baud
             * Parity: None
             * Data bits: 8
             * Stop bits: 2
             * Flow control: None
            All addresses listed are for the request PDU.
            The SunSaver MPPT default server address: 0x01. */
        debug!("Configuring device {}", device);
        let mut connection = Modbus::new_rtu(device, 9600, 'N', 8, 2).unwrap();
        assert!(connection.set_slave(0x01).is_ok());
        assert!(connection.rtu_set_serial_mode(SerialMode::MODBUS_RTU_RS232).is_ok());
        assert!(connection.set_response_timeout( Timeout { sec: 1, usec: 0 } ).is_ok());
        connection.set_debug(false).unwrap();

        let timeout = connection.get_response_timeout();
        info!("Timout {:?}", timeout);

        assert!(connection.connect().is_ok());
        debug!("Connected");

        ModbusSunSaverConnection {
            connection: connection,
        }
    }

    fn read_registers_retry(&self, address: i32, num_bit: i32, dest: &mut [u16]) -> Result<usize, RetryError> {
        match Retry::new(
            &mut || self.connection.read_registers(address, num_bit, dest),
            &mut |response| response.is_ok()
            ).try(3).wait(100).execute() {
            Ok(response) => Ok(response.unwrap() as usize),
            Err(error) => Err(error),
        }
    }
}

impl SunSaverConnection for ModbusSunSaverConnection {
    fn read_registers(&mut self) -> [u16; 44] {
        let mut response_register = [0u16; 44 as usize];
        let mut num_read_bytes = 0;
        num_read_bytes += self.read_registers_retry(0x08, 22, &mut response_register[0..22]).unwrap();
        num_read_bytes += self.read_registers_retry(0x1E, 22, &mut response_register[23..44]).unwrap();
        //if num_read_bytes != 44 {
        //    panic!("Failed to read all registers! Required 44 got {}", num_read_bytes);
        //}
        debug!("Read {} bytes", num_read_bytes);
        debug!("read reg 0x08 + 44: {:#x}", response_register.as_hex());

        response_register
    }

    fn read_logged_data(&mut self) -> [u16; 32 * 16] {
        let mut logged_data = [0u16; (32 * 16) as usize];

        self.read_registers_retry(0x8000, 32 * 16, &mut logged_data).unwrap();

        debug!("logged_data_start");
        for i in (0 as usize)..32 {
            debug!("{:#x}", logged_data[i..(i+16)].as_hex());
        }
        debug!("logged_data_end");

        logged_data
    }
}

#[derive(Debug)]
pub struct FileSunSaverConnection {
    file: File,
}

impl FileSunSaverConnection {
    pub fn open(filename: &str) -> FileSunSaverConnection {
        let file = OpenOptions::new().read(true).write(false).open(filename).unwrap();

        FileSunSaverConnection {
            file: file,
        }
    }
}

impl SunSaverConnection for FileSunSaverConnection {
    fn read_registers(&mut self) -> [u16; 44] {
        let mut response_register_u8 = [0u8; 88 as usize];
        assert!(self.file.read_exact(&mut response_register_u8).is_ok());

        let response_register_vec_u16: Vec<u16> = response_register_u8.chunks(2)
        .map(|items| {
            u16::from_le(((items[0] as u16) << 8) + (items[1] as u16))
        })
        .collect();

        let mut response_register_u16 = [0u16; 44 as usize];
        response_register_u16.clone_from_slice(&response_register_vec_u16);
        response_register_u16
    }

    fn read_logged_data(&mut self) -> [u16; 32 * 16] {
        unimplemented!();
    }
}<|endoftext|>


<commit_before>//! This modules implements a function to resolve a path `foo::bar::baz` to a
//! def, which is used within the name resolution.
//!
//! When name resolution is finished, the result of resolving a path is either
//! `Some(def)` or `None`. However, when we are in process of resolving imports
//! or macros, there's a third possibility:
//!
//!   I can't resolve this path right now, but I might be resolve this path
//!   later, when more macros are expanded.
//!
//! `ReachedFixedPoint` signals about this.

use std::iter::successors;

use base_db::Edition;
use hir_expand::name::Name;
use hir_expand::name;
use test_utils::mark;

use crate::{
    db::DefDatabase,
    item_scope::BUILTIN_SCOPE,
    nameres::{BuiltinShadowMode, DefMap},
    path::{ModPath, PathKind},
    per_ns::PerNs,
    visibility::{RawVisibility, Visibility},
    AdtId, CrateId, EnumVariantId, LocalModuleId, ModuleDefId, ModuleId,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum ResolveMode {
    Import,
    Other,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum ReachedFixedPoint {
    Yes,
    No,
}

#[derive(Debug, Clone)]
pub(super) struct ResolvePathResult {
    pub(super) resolved_def: PerNs,
    pub(super) segment_index: Option<usize>,
    pub(super) reached_fixedpoint: ReachedFixedPoint,
    pub(super) krate: Option<CrateId>,
}

impl ResolvePathResult {
    fn empty(reached_fixedpoint: ReachedFixedPoint) -> ResolvePathResult {
        ResolvePathResult::with(PerNs::none(), reached_fixedpoint, None, None)
    }

    fn with(
        resolved_def: PerNs,
        reached_fixedpoint: ReachedFixedPoint,
        segment_index: Option<usize>,
        krate: Option<CrateId>,
    ) -> ResolvePathResult {
        ResolvePathResult { resolved_def, reached_fixedpoint, segment_index, krate }
    }
}

impl DefMap {
    pub(super) fn resolve_name_in_extern_prelude(&self, name: &Name) -> PerNs {
        if name == &name!(self) {
            return PerNs::types(ModuleId { krate: self.krate, local_id: self.root }.into(),
                Visibility::Public,
            );
        }
        self.extern_prelude
            .get(name)
            .map_or(PerNs::none(), |&it| PerNs::types(it, Visibility::Public))
    }

    pub(crate) fn resolve_visibility(
        &self,
        db: &dyn DefDatabase,
        original_module: LocalModuleId,
        visibility: &RawVisibility,
    ) -> Option<Visibility> {
        match visibility {
            RawVisibility::Module(path) => {
                let (result, remaining) =
                    self.resolve_path(db, original_module, &path, BuiltinShadowMode::Module);
                if remaining.is_some() {
                    return None;
                }
                let types = result.take_types()?;
                match types {
                    ModuleDefId::ModuleId(m) => Some(Visibility::Module(m)),
                    _ => {
                        // error: visibility needs to refer to module
                        None
                    }
                }
            }
            RawVisibility::Public => Some(Visibility::Public),
        }
    }

    // Returns Yes if we are sure that additions to `ItemMap` wouldn't change
    // the result.
    pub(super) fn resolve_path_fp_with_macro(
        &self,
        db: &dyn DefDatabase,
        mode: ResolveMode,
        original_module: LocalModuleId,
        path: &ModPath,
        shadow: BuiltinShadowMode,
    ) -> ResolvePathResult {
        let mut result = ResolvePathResult::empty(ReachedFixedPoint::No);
        result.segment_index = Some(usize::max_value());

        let mut current_map = self;
        loop {
            let new = current_map.resolve_path_fp_with_macro_single(
                db,
                mode,
                original_module,
                path,
                shadow,
            );

            // Merge `new` into `result`.
            result.resolved_def = result.resolved_def.or(new.resolved_def);
            if result.reached_fixedpoint == ReachedFixedPoint::No {
                result.reached_fixedpoint = new.reached_fixedpoint;
            }
            // FIXME: this doesn't seem right; what if the different namespace resolutions come from different crates?
            result.krate = result.krate.or(new.krate);
            result.segment_index = result.segment_index.min(new.segment_index);

            match &current_map.parent {
                Some(map) => current_map = map,
                None => return result,
            }
        }
    }

    pub(super) fn resolve_path_fp_with_macro_single(
        &self,
        db: &dyn DefDatabase,
        mode: ResolveMode,
        original_module: LocalModuleId,
        path: &ModPath,
        shadow: BuiltinShadowMode,
    ) -> ResolvePathResult {
        let mut segments = path.segments.iter().enumerate();
        let mut curr_per_ns: PerNs = match path.kind {
            PathKind::DollarCrate(krate) => {
                if krate == self.krate {
                    mark::hit!(macro_dollar_crate_self);
                    PerNs::types(
                        ModuleId { krate: self.krate, local_id: self.root }.into(),
                        Visibility::Public,
                    )
                } else {
                    let def_map = db.crate_def_map(krate);
                    let module = ModuleId { krate, local_id: def_map.root };
                    mark::hit!(macro_dollar_crate_other);
                    PerNs::types(module.into(), Visibility::Public)
                }
            }
            PathKind::Crate => PerNs::types(
                ModuleId { krate: self.krate, local_id: self.root }.into(),
                Visibility::Public,
            ),
            // plain import or absolute path in 2015: crate-relative with
            // fallback to extern prelude (with the simplification in
            // rust-lang/rust#57745)
            // FIXME there must be a nicer way to write this condition
            PathKind::Plain | PathKind::Abs
                if self.edition == Edition::Edition2015
                    && (path.kind == PathKind::Abs || mode == ResolveMode::Import) =>
            {
                let (_, segment) = match segments.next() {
                    Some((idx, segment)) => (idx, segment),
                    None => return ResolvePathResult::empty(ReachedFixedPoint::Yes),
                };
                log::debug!("resolving {:?} in crate root (+ extern prelude)", segment);
                self.resolve_name_in_crate_root_or_extern_prelude(&segment)
            }
            PathKind::Plain => {
                let (_, segment) = match segments.next() {
                    Some((idx, segment)) => (idx, segment),
                    None => return ResolvePathResult::empty(ReachedFixedPoint::Yes),
                };
                // The first segment may be a builtin type. If the path has more
                // than one segment, we first try resolving it as a module
                // anyway.
                // FIXME: If the next segment doesn't resolve in the module and
                // BuiltinShadowMode wasn't Module, then we need to try
                // resolving it as a builtin.
                let prefer_module =
                    if path.segments.len() == 1 { shadow } else { BuiltinShadowMode::Module };

                log::debug!("resolving {:?} in module", segment);
                self.resolve_name_in_module(db, original_module, &segment, prefer_module)
            }
            PathKind::Super(lvl) => {
                let m = successors(Some(original_module), |m| self.modules[*m].parent)
                    .nth(lvl as usize);
                if let Some(local_id) = m {
                    PerNs::types(
                        ModuleId { krate: self.krate, local_id }.into(),
                        Visibility::Public,
                    )
                } else {
                    log::debug!("super path in root module");
                    return ResolvePathResult::empty(ReachedFixedPoint::Yes);
                }
            }
            PathKind::Abs => {
                // 2018-style absolute path -- only extern prelude
                let segment = match segments.next() {
                    Some((_, segment)) => segment,
                    None => return ResolvePathResult::empty(ReachedFixedPoint::Yes),
                };
                if let Some(def) = self.extern_prelude.get(&segment) {
                    log::debug!("absolute path {:?} resolved to crate {:?}", path, def);
                    PerNs::types(*def, Visibility::Public)
                } else {
                    return ResolvePathResult::empty(ReachedFixedPoint::No); // extern crate declarations can add to the extern prelude
                }
            }
        };

        for (i, segment) in segments {
            let (curr, vis) = match curr_per_ns.take_types_vis() {
                Some(r) => r,
                None => {
                    // we still have path segments left, but the path so far
                    // didn't resolve in the types namespace => no resolution
                    // (don't break here because `curr_per_ns` might contain
                    // something in the value namespace, and it would be wrong
                    // to return that)
                    return ResolvePathResult::empty(ReachedFixedPoint::No);
                }
            };
            // resolve segment in curr

            curr_per_ns = match curr {
                ModuleDefId::ModuleId(module) => {
                    if module.krate != self.krate {
                        let path = ModPath {
                            segments: path.segments[i..].to_vec(),
                            kind: PathKind::Super(0),
                        };
                        log::debug!("resolving {:?} in other crate", path);
                        let defp_map = db.crate_def_map(module.krate);
                        let (def, s) = defp_map.resolve_path(db, module.local_id, &path, shadow);
                        return ResolvePathResult::with(
                            def,
                            ReachedFixedPoint::Yes,
                            s.map(|s| s + i),
                            Some(module.krate),
                        );
                    }

                    // Since it is a qualified path here, it should not contains legacy macros
                    self[module.local_id].scope.get(&segment)
                }
                ModuleDefId::AdtId(AdtId::EnumId(e)) => {
                    // enum variant
                    mark::hit!(can_import_enum_variant);
                    let enum_data = db.enum_data(e);
                    match enum_data.variant(&segment) {
                        Some(local_id) => {
                            let variant = EnumVariantId { parent: e, local_id };
                            match &*enum_data.variants[local_id].variant_data {
                                crate::adt::VariantData::Record(_) => {
                                    PerNs::types(variant.into(), Visibility::Public)
                                }
                                crate::adt::VariantData::Tuple(_)
                                | crate::adt::VariantData::Unit => {
                                    PerNs::both(variant.into(), variant.into(), Visibility::Public)
                                }
                            }
                        }
                        None => {
                            return ResolvePathResult::with(
                                PerNs::types(e.into(), vis),
                                ReachedFixedPoint::Yes,
                                Some(i),
                                Some(self.krate),
                            );
                        }
                    }
                }
                s => {
                    // could be an inherent method call in UFCS form
                    // (`Struct::method`), or some other kind of associated item
                    log::debug!(
                        "path segment {:?} resolved to non-module {:?}, but is not last",
                        segment,
                        curr,
                    );

                    return ResolvePathResult::with(
                        PerNs::types(s, vis),
                        ReachedFixedPoint::Yes,
                        Some(i),
                        Some(self.krate),
                    );
                }
            };
        }

        ResolvePathResult::with(curr_per_ns, ReachedFixedPoint::Yes, None, Some(self.krate))
    }

    fn resolve_name_in_module(
        &self,
        db: &dyn DefDatabase,
        module: LocalModuleId,
        name: &Name,
        shadow: BuiltinShadowMode,
    ) -> PerNs {
        // Resolve in:
        //  - legacy scope of macro
        //  - current module / scope
        //  - extern prelude
        //  - std prelude
        let from_legacy_macro = self[module]
            .scope
            .get_legacy_macro(name)
            .map_or_else(PerNs::none, |m| PerNs::macros(m, Visibility::Public));
        let from_scope = self[module].scope.get(name);
        let from_builtin = BUILTIN_SCOPE.get(name).copied().unwrap_or_else(PerNs::none);
        let from_scope_or_builtin = match shadow {
            BuiltinShadowMode::Module => from_scope.or(from_builtin),
            BuiltinShadowMode::Other => {
                if let Some(ModuleDefId::ModuleId(_)) = from_scope.take_types() {
                    from_builtin.or(from_scope)
                } else {
                    from_scope.or(from_builtin)
                }
            }
        };
        let from_extern_prelude = self
            .extern_prelude
            .get(name)
            .map_or(PerNs::none(), |&it| PerNs::types(it, Visibility::Public));
        let from_prelude = self.resolve_in_prelude(db, name);

        from_legacy_macro.or(from_scope_or_builtin).or(from_extern_prelude).or(from_prelude)
    }

    fn resolve_name_in_crate_root_or_extern_prelude(&self, name: &Name) -> PerNs {
        let from_crate_root = self[self.root].scope.get(name);
        let from_extern_prelude = self.resolve_name_in_extern_prelude(name);

        from_crate_root.or(from_extern_prelude)
    }

    fn resolve_in_prelude(&self, db: &dyn DefDatabase, name: &Name) -> PerNs {
        if let Some(prelude) = self.prelude {
            let keep;
            let def_map = if prelude.krate == self.krate {
                self
            } else {
                // Extend lifetime
                keep = db.crate_def_map(prelude.krate);
                &keep
            };
            def_map[prelude.local_id].scope.get(name)
        } else {
            PerNs::none()
        }
    }
}
<commit_msg>cargo fmt<commit_after>//! This modules implements a function to resolve a path `foo::bar::baz` to a
//! def, which is used within the name resolution.
//!
//! When name resolution is finished, the result of resolving a path is either
//! `Some(def)` or `None`. However, when we are in process of resolving imports
//! or macros, there's a third possibility:
//!
//!   I can't resolve this path right now, but I might be resolve this path
//!   later, when more macros are expanded.
//!
//! `ReachedFixedPoint` signals about this.

use std::iter::successors;

use base_db::Edition;
use hir_expand::name;
use hir_expand::name::Name;
use test_utils::mark;

use crate::{
    db::DefDatabase,
    item_scope::BUILTIN_SCOPE,
    nameres::{BuiltinShadowMode, DefMap},
    path::{ModPath, PathKind},
    per_ns::PerNs,
    visibility::{RawVisibility, Visibility},
    AdtId, CrateId, EnumVariantId, LocalModuleId, ModuleDefId, ModuleId,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum ResolveMode {
    Import,
    Other,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum ReachedFixedPoint {
    Yes,
    No,
}

#[derive(Debug, Clone)]
pub(super) struct ResolvePathResult {
    pub(super) resolved_def: PerNs,
    pub(super) segment_index: Option<usize>,
    pub(super) reached_fixedpoint: ReachedFixedPoint,
    pub(super) krate: Option<CrateId>,
}

impl ResolvePathResult {
    fn empty(reached_fixedpoint: ReachedFixedPoint) -> ResolvePathResult {
        ResolvePathResult::with(PerNs::none(), reached_fixedpoint, None, None)
    }

    fn with(
        resolved_def: PerNs,
        reached_fixedpoint: ReachedFixedPoint,
        segment_index: Option<usize>,
        krate: Option<CrateId>,
    ) -> ResolvePathResult {
        ResolvePathResult { resolved_def, reached_fixedpoint, segment_index, krate }
    }
}

impl DefMap {
    pub(super) fn resolve_name_in_extern_prelude(&self, name: &Name) -> PerNs {
        if name == &name!(self) {
            return PerNs::types(
                ModuleId { krate: self.krate, local_id: self.root }.into(),
                Visibility::Public,
            );
        }
        self.extern_prelude
            .get(name)
            .map_or(PerNs::none(), |&it| PerNs::types(it, Visibility::Public))
    }

    pub(crate) fn resolve_visibility(
        &self,
        db: &dyn DefDatabase,
        original_module: LocalModuleId,
        visibility: &RawVisibility,
    ) -> Option<Visibility> {
        match visibility {
            RawVisibility::Module(path) => {
                let (result, remaining) =
                    self.resolve_path(db, original_module, &path, BuiltinShadowMode::Module);
                if remaining.is_some() {
                    return None;
                }
                let types = result.take_types()?;
                match types {
                    ModuleDefId::ModuleId(m) => Some(Visibility::Module(m)),
                    _ => {
                        // error: visibility needs to refer to module
                        None
                    }
                }
            }
            RawVisibility::Public => Some(Visibility::Public),
        }
    }

    // Returns Yes if we are sure that additions to `ItemMap` wouldn't change
    // the result.
    pub(super) fn resolve_path_fp_with_macro(
        &self,
        db: &dyn DefDatabase,
        mode: ResolveMode,
        original_module: LocalModuleId,
        path: &ModPath,
        shadow: BuiltinShadowMode,
    ) -> ResolvePathResult {
        let mut result = ResolvePathResult::empty(ReachedFixedPoint::No);
        result.segment_index = Some(usize::max_value());

        let mut current_map = self;
        loop {
            let new = current_map.resolve_path_fp_with_macro_single(
                db,
                mode,
                original_module,
                path,
                shadow,
            );

            // Merge `new` into `result`.
            result.resolved_def = result.resolved_def.or(new.resolved_def);
            if result.reached_fixedpoint == ReachedFixedPoint::No {
                result.reached_fixedpoint = new.reached_fixedpoint;
            }
            // FIXME: this doesn't seem right; what if the different namespace resolutions come from different crates?
            result.krate = result.krate.or(new.krate);
            result.segment_index = result.segment_index.min(new.segment_index);

            match &current_map.parent {
                Some(map) => current_map = map,
                None => return result,
            }
        }
    }

    pub(super) fn resolve_path_fp_with_macro_single(
        &self,
        db: &dyn DefDatabase,
        mode: ResolveMode,
        original_module: LocalModuleId,
        path: &ModPath,
        shadow: BuiltinShadowMode,
    ) -> ResolvePathResult {
        let mut segments = path.segments.iter().enumerate();
        let mut curr_per_ns: PerNs = match path.kind {
            PathKind::DollarCrate(krate) => {
                if krate == self.krate {
                    mark::hit!(macro_dollar_crate_self);
                    PerNs::types(
                        ModuleId { krate: self.krate, local_id: self.root }.into(),
                        Visibility::Public,
                    )
                } else {
                    let def_map = db.crate_def_map(krate);
                    let module = ModuleId { krate, local_id: def_map.root };
                    mark::hit!(macro_dollar_crate_other);
                    PerNs::types(module.into(), Visibility::Public)
                }
            }
            PathKind::Crate => PerNs::types(
                ModuleId { krate: self.krate, local_id: self.root }.into(),
                Visibility::Public,
            ),
            // plain import or absolute path in 2015: crate-relative with
            // fallback to extern prelude (with the simplification in
            // rust-lang/rust#57745)
            // FIXME there must be a nicer way to write this condition
            PathKind::Plain | PathKind::Abs
                if self.edition == Edition::Edition2015
                    && (path.kind == PathKind::Abs || mode == ResolveMode::Import) =>
            {
                let (_, segment) = match segments.next() {
                    Some((idx, segment)) => (idx, segment),
                    None => return ResolvePathResult::empty(ReachedFixedPoint::Yes),
                };
                log::debug!("resolving {:?} in crate root (+ extern prelude)", segment);
                self.resolve_name_in_crate_root_or_extern_prelude(&segment)
            }
            PathKind::Plain => {
                let (_, segment) = match segments.next() {
                    Some((idx, segment)) => (idx, segment),
                    None => return ResolvePathResult::empty(ReachedFixedPoint::Yes),
                };
                // The first segment may be a builtin type. If the path has more
                // than one segment, we first try resolving it as a module
                // anyway.
                // FIXME: If the next segment doesn't resolve in the module and
                // BuiltinShadowMode wasn't Module, then we need to try
                // resolving it as a builtin.
                let prefer_module =
                    if path.segments.len() == 1 { shadow } else { BuiltinShadowMode::Module };

                log::debug!("resolving {:?} in module", segment);
                self.resolve_name_in_module(db, original_module, &segment, prefer_module)
            }
            PathKind::Super(lvl) => {
                let m = successors(Some(original_module), |m| self.modules[*m].parent)
                    .nth(lvl as usize);
                if let Some(local_id) = m {
                    PerNs::types(
                        ModuleId { krate: self.krate, local_id }.into(),
                        Visibility::Public,
                    )
                } else {
                    log::debug!("super path in root module");
                    return ResolvePathResult::empty(ReachedFixedPoint::Yes);
                }
            }
            PathKind::Abs => {
                // 2018-style absolute path -- only extern prelude
                let segment = match segments.next() {
                    Some((_, segment)) => segment,
                    None => return ResolvePathResult::empty(ReachedFixedPoint::Yes),
                };
                if let Some(def) = self.extern_prelude.get(&segment) {
                    log::debug!("absolute path {:?} resolved to crate {:?}", path, def);
                    PerNs::types(*def, Visibility::Public)
                } else {
                    return ResolvePathResult::empty(ReachedFixedPoint::No); // extern crate declarations can add to the extern prelude
                }
            }
        };

        for (i, segment) in segments {
            let (curr, vis) = match curr_per_ns.take_types_vis() {
                Some(r) => r,
                None => {
                    // we still have path segments left, but the path so far
                    // didn't resolve in the types namespace => no resolution
                    // (don't break here because `curr_per_ns` might contain
                    // something in the value namespace, and it would be wrong
                    // to return that)
                    return ResolvePathResult::empty(ReachedFixedPoint::No);
                }
            };
            // resolve segment in curr

            curr_per_ns = match curr {
                ModuleDefId::ModuleId(module) => {
                    if module.krate != self.krate {
                        let path = ModPath {
                            segments: path.segments[i..].to_vec(),
                            kind: PathKind::Super(0),
                        };
                        log::debug!("resolving {:?} in other crate", path);
                        let defp_map = db.crate_def_map(module.krate);
                        let (def, s) = defp_map.resolve_path(db, module.local_id, &path, shadow);
                        return ResolvePathResult::with(
                            def,
                            ReachedFixedPoint::Yes,
                            s.map(|s| s + i),
                            Some(module.krate),
                        );
                    }

                    // Since it is a qualified path here, it should not contains legacy macros
                    self[module.local_id].scope.get(&segment)
                }
                ModuleDefId::AdtId(AdtId::EnumId(e)) => {
                    // enum variant
                    mark::hit!(can_import_enum_variant);
                    let enum_data = db.enum_data(e);
                    match enum_data.variant(&segment) {
                        Some(local_id) => {
                            let variant = EnumVariantId { parent: e, local_id };
                            match &*enum_data.variants[local_id].variant_data {
                                crate::adt::VariantData::Record(_) => {
                                    PerNs::types(variant.into(), Visibility::Public)
                                }
                                crate::adt::VariantData::Tuple(_)
                                | crate::adt::VariantData::Unit => {
                                    PerNs::both(variant.into(), variant.into(), Visibility::Public)
                                }
                            }
                        }
                        None => {
                            return ResolvePathResult::with(
                                PerNs::types(e.into(), vis),
                                ReachedFixedPoint::Yes,
                                Some(i),
                                Some(self.krate),
                            );
                        }
                    }
                }
                s => {
                    // could be an inherent method call in UFCS form
                    // (`Struct::method`), or some other kind of associated item
                    log::debug!(
                        "path segment {:?} resolved to non-module {:?}, but is not last",
                        segment,
                        curr,
                    );

                    return ResolvePathResult::with(
                        PerNs::types(s, vis),
                        ReachedFixedPoint::Yes,
                        Some(i),
                        Some(self.krate),
                    );
                }
            };
        }

        ResolvePathResult::with(curr_per_ns, ReachedFixedPoint::Yes, None, Some(self.krate))
    }

    fn resolve_name_in_module(
        &self,
        db: &dyn DefDatabase,
        module: LocalModuleId,
        name: &Name,
        shadow: BuiltinShadowMode,
    ) -> PerNs {
        // Resolve in:
        //  - legacy scope of macro
        //  - current module / scope
        //  - extern prelude
        //  - std prelude
        let from_legacy_macro = self[module]
            .scope
            .get_legacy_macro(name)
            .map_or_else(PerNs::none, |m| PerNs::macros(m, Visibility::Public));
        let from_scope = self[module].scope.get(name);
        let from_builtin = BUILTIN_SCOPE.get(name).copied().unwrap_or_else(PerNs::none);
        let from_scope_or_builtin = match shadow {
            BuiltinShadowMode::Module => from_scope.or(from_builtin),
            BuiltinShadowMode::Other => {
                if let Some(ModuleDefId::ModuleId(_)) = from_scope.take_types() {
                    from_builtin.or(from_scope)
                } else {
                    from_scope.or(from_builtin)
                }
            }
        };
        let from_extern_prelude = self
            .extern_prelude
            .get(name)
            .map_or(PerNs::none(), |&it| PerNs::types(it, Visibility::Public));
        let from_prelude = self.resolve_in_prelude(db, name);

        from_legacy_macro.or(from_scope_or_builtin).or(from_extern_prelude).or(from_prelude)
    }

    fn resolve_name_in_crate_root_or_extern_prelude(&self, name: &Name) -> PerNs {
        let from_crate_root = self[self.root].scope.get(name);
        let from_extern_prelude = self.resolve_name_in_extern_prelude(name);

        from_crate_root.or(from_extern_prelude)
    }

    fn resolve_in_prelude(&self, db: &dyn DefDatabase, name: &Name) -> PerNs {
        if let Some(prelude) = self.prelude {
            let keep;
            let def_map = if prelude.krate == self.krate {
                self
            } else {
                // Extend lifetime
                keep = db.crate_def_map(prelude.krate);
                &keep
            };
            def_map[prelude.local_id].scope.get(name)
        } else {
            PerNs::none()
        }
    }
}
<|endoftext|>


<commit_before>use syntax::{
    ast::{edit::AstNodeEdit, make, AstNode, BlockExpr, ElseBranch, Expr, IfExpr, MatchArm},
    NodeOrToken,
    SyntaxKind::{COMMA, WHITESPACE},
};

use crate::{AssistContext, AssistId, AssistKind, Assists};

// Assist: move_guard_to_arm_body
//
// Moves match guard into match arm body.
//
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         Action::Move { distance } $0if distance > 10 => foo(),
//         _ => (),
//     }
// }
// ```
// ->
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         Action::Move { distance } => if distance > 10 {
//             foo()
//         },
//         _ => (),
//     }
// }
// ```
pub(crate) fn move_guard_to_arm_body(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let match_arm = ctx.find_node_at_offset::<MatchArm>()?;
    let guard = match_arm.guard()?;
    if ctx.offset() > guard.syntax().text_range().end() {
        cov_mark::hit!(move_guard_unapplicable_in_arm_body);
        return None;
    }
    let space_before_guard = guard.syntax().prev_sibling_or_token();

    // FIXME: support `if let` guards too
    if guard.let_token().is_some() {
        return None;
    }
    let guard_condition = guard.expr()?;
    let arm_expr = match_arm.expr()?;
    let if_expr = make::expr_if(
        make::condition(guard_condition, None),
        make::block_expr(None, Some(arm_expr.clone())),
        None,
    )
    .indent(arm_expr.indent_level());

    let target = guard.syntax().text_range();
    acc.add(
        AssistId("move_guard_to_arm_body", AssistKind::RefactorRewrite),
        "Move guard to arm body",
        target,
        |edit| {
            match space_before_guard {
                Some(element) if element.kind() == WHITESPACE => {
                    edit.delete(element.text_range());
                }
                _ => (),
            };

            edit.delete(guard.syntax().text_range());
            edit.replace_ast(arm_expr, if_expr);
        },
    )
}

// Assist: move_arm_cond_to_match_guard
//
// Moves if expression from match arm body into a guard.
//
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         Action::Move { distance } => $0if distance > 10 { foo() },
//         _ => (),
//     }
// }
// ```
// ->
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         Action::Move { distance } if distance > 10 => foo(),
//         _ => (),
//     }
// }
// ```
pub(crate) fn move_arm_cond_to_match_guard(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let match_arm: MatchArm = ctx.find_node_at_offset::<MatchArm>()?;
    let match_pat = match_arm.pat()?;
    let arm_body = match_arm.expr()?;

    let mut replace_node = None;
    let if_expr: IfExpr = IfExpr::cast(arm_body.syntax().clone()).or_else(|| {
        let block_expr = BlockExpr::cast(arm_body.syntax().clone())?;
        if let Expr::IfExpr(e) = block_expr.tail_expr()? {
            replace_node = Some(block_expr.syntax().clone());
            Some(e)
        } else {
            None
        }
    })?;
    let replace_node = replace_node.unwrap_or_else(|| if_expr.syntax().clone());

    let cond = if_expr.condition()?;
    let then_block = if_expr.then_branch()?;

    // Not support moving if let to arm guard
    if cond.is_pattern_cond() {
        return None;
    }

    let buf = format!(" if {}", cond.syntax().text());

    acc.add(
        AssistId("move_arm_cond_to_match_guard", AssistKind::RefactorRewrite),
        "Move condition to match guard",
        replace_node.text_range(),
        |edit| {
            let then_only_expr = then_block.statements().next().is_none();

            match &then_block.tail_expr() {
                Some(then_expr) if then_only_expr => {
                    edit.replace(replace_node.text_range(), then_expr.syntax().text())
                }
                _ if replace_node != *if_expr.syntax() => {
                    // Dedent because if_expr is in a BlockExpr
                    let replace_with = then_block.dedent(1.into()).syntax().text();
                    edit.replace(replace_node.text_range(), replace_with)
                }
                _ => edit.replace(replace_node.text_range(), then_block.syntax().text()),
            }

            edit.insert(match_pat.syntax().text_range().end(), buf);

            // If with only an else branch
            if let Some(ElseBranch::Block(else_block)) = if_expr.else_branch() {
                let then_arm_end = match_arm.syntax().text_range().end();
                if then_block.tail_expr().is_some() && then_only_expr {
                    // Insert comma for expression if there isn't one
                    match match_arm.syntax().last_child_or_token() {
                        Some(NodeOrToken::Token(t)) if t.kind() == COMMA => {}
                        _ => edit.insert(then_arm_end, ","),
                    }
                }
                let else_only_expr = else_block.statements().next().is_none();
                let indent_level = match_arm.indent_level();
                let spaces = "    ".repeat(indent_level.0 as _);
                edit.insert(then_arm_end, format!("\n{}{} => ", spaces, match_pat));
                match &else_block.tail_expr() {
                    Some(else_expr) if else_only_expr => {
                        edit.insert(then_arm_end, else_expr.syntax().text());
                        edit.insert(then_arm_end, ",");
                    }
                    _ if replace_node != *if_expr.syntax() => {
                        edit.insert(then_arm_end, else_block.dedent(1.into()).syntax().text());
                    }
                    _ => edit.insert(then_arm_end, else_block.syntax().text()),
                }
            }
        },
    )
}

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

    use crate::tests::{check_assist, check_assist_not_applicable, check_assist_target};

    #[test]
    fn move_guard_to_arm_body_range() {
        cov_mark::check!(move_guard_unapplicable_in_arm_body);
        check_assist_not_applicable(
            move_guard_to_arm_body,
            r#"
fn main() {
    match 92 {
        x if x > 10 => $0false,
        _ => true
    }
}
"#,
        );
    }
    #[test]
    fn move_guard_to_arm_body_target() {
        check_assist_target(
            move_guard_to_arm_body,
            r#"
fn main() {
    match 92 {
        x $0if x > 10 => false,
        _ => true
    }
}
"#,
            r#"if x > 10"#,
        );
    }

    #[test]
    fn move_guard_to_arm_body_works() {
        check_assist(
            move_guard_to_arm_body,
            r#"
fn main() {
    match 92 {
        x $0if x > 10 => false,
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x => if x > 10 {
            false
        },
        _ => true
    }
}
"#,
        );
    }

    #[test]
    fn move_guard_to_arm_body_works_complex_match() {
        check_assist(
            move_guard_to_arm_body,
            r#"
fn main() {
    match 92 {
        $0x @ 4 | x @ 5    if x > 5 => true,
        _ => false
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x @ 4 | x @ 5 => if x > 5 {
            true
        },
        _ => false
    }
}
"#,
        );
    }

    #[test]
    fn move_arm_cond_to_match_guard_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => if x > 10 { $0false },
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => false,
        _ => true
    }
}
"#,
        );
    }

    #[test]
    fn move_arm_cond_in_block_to_match_guard_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => {
            $0if x > 10 {
                false
            }
        },
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => false,
        _ => true
    }
}
"#,
        );
    }

    #[test]
    fn move_arm_cond_to_match_guard_if_let_not_works() {
        check_assist_not_applicable(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => if let 62 = x { $0false },
        _ => true
    }
}
"#,
        );
    }

    #[test]
    fn move_arm_cond_to_match_guard_if_empty_body_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => if x > 10 { $0 },
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => {  },
        _ => true
    }
}
"#,
        );
    }

    #[test]
    fn move_arm_cond_to_match_guard_if_multiline_body_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => if x > 10 {
            92;$0
            false
        },
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => {
            92;
            false
        },
        _ => true
    }
}
"#,
        );
    }

    #[test]
    fn move_arm_cond_in_block_to_match_guard_if_multiline_body_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => {
            if x > 10 {
                92;$0
                false
            }
        }
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => {
            92;
            false
        }
        _ => true
    }
}
"#,
        )
    }

    #[test]
    fn move_arm_cond_to_match_guard_with_else_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => if x > 10 {$0
            false
        } else {
            true
        }
        _ => true,
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => false,
        x => true,
        _ => true,
    }
}
"#,
        )
    }

    #[test]
    fn move_arm_cond_to_match_guard_with_else_block_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => {
            if x > 10 {$0
                false
            } else {
                true
            }
        }
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => false,
        x => true,
        _ => true
    }
}
"#,
        )
    }

    #[test]
    fn move_arm_cond_to_match_guard_else_if_empty_body_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => if x > 10 { $0 } else { },
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => {  },
        x => { }
        _ => true
    }
}
"#,
        );
    }

    #[test]
    fn move_arm_cond_to_match_guard_with_else_multiline_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => if x > 10 {
            92;$0
            false
        } else {
            true
        }
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => {
            92;
            false
        }
        x => true,
        _ => true
    }
}
"#,
        )
    }

    #[test]
    fn move_arm_cond_to_match_guard_with_else_multiline_else_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => {
            if x > 10 {$0
                false
            } else {
                42;
                true
            }
        }
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => false,
        x => {
            42;
            true
        }
        _ => true
    }
}
"#,
        )
    }

    #[test]
    fn move_arm_cond_to_match_guard_with_else_last_arm_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        3 => true,
        x => {
            if x > 10 {$0
                false
            } else {
                92;
                true
            }
        }
    }
}
"#,
            r#"
fn main() {
    match 92 {
        3 => true,
        x if x > 10 => false,
        x => {
            92;
            true
        }
    }
}
"#,
        )
    }

    #[test]
    fn move_arm_cond_to_match_guard_with_else_comma_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        3 => true,
        x => if x > 10 {$0
            false
        } else {
            92;
            true
        },
    }
}
"#,
            r#"
fn main() {
    match 92 {
        3 => true,
        x if x > 10 => false,
        x => {
            92;
            true
        }
    }
}
"#,
        )
    }
}
<commit_msg>Add coverage marks<commit_after>use syntax::{
    ast::{edit::AstNodeEdit, make, AstNode, BlockExpr, ElseBranch, Expr, IfExpr, MatchArm},
    NodeOrToken,
    SyntaxKind::{COMMA, WHITESPACE},
};

use crate::{AssistContext, AssistId, AssistKind, Assists};

// Assist: move_guard_to_arm_body
//
// Moves match guard into match arm body.
//
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         Action::Move { distance } $0if distance > 10 => foo(),
//         _ => (),
//     }
// }
// ```
// ->
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         Action::Move { distance } => if distance > 10 {
//             foo()
//         },
//         _ => (),
//     }
// }
// ```
pub(crate) fn move_guard_to_arm_body(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let match_arm = ctx.find_node_at_offset::<MatchArm>()?;
    let guard = match_arm.guard()?;
    if ctx.offset() > guard.syntax().text_range().end() {
        cov_mark::hit!(move_guard_unapplicable_in_arm_body);
        return None;
    }
    let space_before_guard = guard.syntax().prev_sibling_or_token();

    // FIXME: support `if let` guards too
    if guard.let_token().is_some() {
        return None;
    }
    let guard_condition = guard.expr()?;
    let arm_expr = match_arm.expr()?;
    let if_expr = make::expr_if(
        make::condition(guard_condition, None),
        make::block_expr(None, Some(arm_expr.clone())),
        None,
    )
    .indent(arm_expr.indent_level());

    let target = guard.syntax().text_range();
    acc.add(
        AssistId("move_guard_to_arm_body", AssistKind::RefactorRewrite),
        "Move guard to arm body",
        target,
        |edit| {
            match space_before_guard {
                Some(element) if element.kind() == WHITESPACE => {
                    edit.delete(element.text_range());
                }
                _ => (),
            };

            edit.delete(guard.syntax().text_range());
            edit.replace_ast(arm_expr, if_expr);
        },
    )
}

// Assist: move_arm_cond_to_match_guard
//
// Moves if expression from match arm body into a guard.
//
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         Action::Move { distance } => $0if distance > 10 { foo() },
//         _ => (),
//     }
// }
// ```
// ->
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         Action::Move { distance } if distance > 10 => foo(),
//         _ => (),
//     }
// }
// ```
pub(crate) fn move_arm_cond_to_match_guard(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let match_arm: MatchArm = ctx.find_node_at_offset::<MatchArm>()?;
    let match_pat = match_arm.pat()?;
    let arm_body = match_arm.expr()?;

    let mut replace_node = None;
    let if_expr: IfExpr = IfExpr::cast(arm_body.syntax().clone()).or_else(|| {
        let block_expr = BlockExpr::cast(arm_body.syntax().clone())?;
        if let Expr::IfExpr(e) = block_expr.tail_expr()? {
            replace_node = Some(block_expr.syntax().clone());
            Some(e)
        } else {
            None
        }
    })?;
    let replace_node = replace_node.unwrap_or_else(|| if_expr.syntax().clone());

    let cond = if_expr.condition()?;
    let then_block = if_expr.then_branch()?;

    // Not support moving if let to arm guard
    if cond.is_pattern_cond() {
        return None;
    }

    let buf = format!(" if {}", cond.syntax().text());

    acc.add(
        AssistId("move_arm_cond_to_match_guard", AssistKind::RefactorRewrite),
        "Move condition to match guard",
        replace_node.text_range(),
        |edit| {
            let then_only_expr = then_block.statements().next().is_none();

            match &then_block.tail_expr() {
                Some(then_expr) if then_only_expr => {
                    edit.replace(replace_node.text_range(), then_expr.syntax().text())
                }
                _ if replace_node != *if_expr.syntax() => {
                    // Dedent because if_expr is in a BlockExpr
                    let replace_with = then_block.dedent(1.into()).syntax().text();
                    edit.replace(replace_node.text_range(), replace_with)
                }
                _ => edit.replace(replace_node.text_range(), then_block.syntax().text()),
            }

            edit.insert(match_pat.syntax().text_range().end(), buf);

            // If with only an else branch
            if let Some(ElseBranch::Block(else_block)) = if_expr.else_branch() {
                let then_arm_end = match_arm.syntax().text_range().end();
                if then_block.tail_expr().is_some() && then_only_expr {
                    // Insert comma for expression if there isn't one
                    match match_arm.syntax().last_child_or_token() {
                        Some(NodeOrToken::Token(t)) if t.kind() == COMMA => {}
                        _ => edit.insert(then_arm_end, ","),
                    }
                }
                let else_only_expr = else_block.statements().next().is_none();
                let indent_level = match_arm.indent_level();
                let spaces = "    ".repeat(indent_level.0 as _);
                edit.insert(then_arm_end, format!("\n{}{} => ", spaces, match_pat));
                match &else_block.tail_expr() {
                    Some(else_expr) if else_only_expr => {
                        cov_mark::hit!(move_guard_ifelse_expr_only);
                        edit.insert(then_arm_end, else_expr.syntax().text());
                        edit.insert(then_arm_end, ",");
                    }
                    _ if replace_node != *if_expr.syntax() => {
                        cov_mark::hit!(move_guard_ifelse_in_block);
                        edit.insert(then_arm_end, else_block.dedent(1.into()).syntax().text());
                    }
                    _ => {
                        cov_mark::hit!(move_guard_ifelse_else_block);
                        edit.insert(then_arm_end, else_block.syntax().text());
                    }
                }
            }
        },
    )
}

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

    use crate::tests::{check_assist, check_assist_not_applicable, check_assist_target};

    #[test]
    fn move_guard_to_arm_body_range() {
        cov_mark::check!(move_guard_unapplicable_in_arm_body);
        check_assist_not_applicable(
            move_guard_to_arm_body,
            r#"
fn main() {
    match 92 {
        x if x > 10 => $0false,
        _ => true
    }
}
"#,
        );
    }
    #[test]
    fn move_guard_to_arm_body_target() {
        check_assist_target(
            move_guard_to_arm_body,
            r#"
fn main() {
    match 92 {
        x $0if x > 10 => false,
        _ => true
    }
}
"#,
            r#"if x > 10"#,
        );
    }

    #[test]
    fn move_guard_to_arm_body_works() {
        check_assist(
            move_guard_to_arm_body,
            r#"
fn main() {
    match 92 {
        x $0if x > 10 => false,
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x => if x > 10 {
            false
        },
        _ => true
    }
}
"#,
        );
    }

    #[test]
    fn move_guard_to_arm_body_works_complex_match() {
        check_assist(
            move_guard_to_arm_body,
            r#"
fn main() {
    match 92 {
        $0x @ 4 | x @ 5    if x > 5 => true,
        _ => false
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x @ 4 | x @ 5 => if x > 5 {
            true
        },
        _ => false
    }
}
"#,
        );
    }

    #[test]
    fn move_arm_cond_to_match_guard_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => if x > 10 { $0false },
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => false,
        _ => true
    }
}
"#,
        );
    }

    #[test]
    fn move_arm_cond_in_block_to_match_guard_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => {
            $0if x > 10 {
                false
            }
        },
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => false,
        _ => true
    }
}
"#,
        );
    }

    #[test]
    fn move_arm_cond_to_match_guard_if_let_not_works() {
        check_assist_not_applicable(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => if let 62 = x { $0false },
        _ => true
    }
}
"#,
        );
    }

    #[test]
    fn move_arm_cond_to_match_guard_if_empty_body_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => if x > 10 { $0 },
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => {  },
        _ => true
    }
}
"#,
        );
    }

    #[test]
    fn move_arm_cond_to_match_guard_if_multiline_body_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => if x > 10 {
            92;$0
            false
        },
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => {
            92;
            false
        },
        _ => true
    }
}
"#,
        );
    }

    #[test]
    fn move_arm_cond_in_block_to_match_guard_if_multiline_body_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => {
            if x > 10 {
                92;$0
                false
            }
        }
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => {
            92;
            false
        }
        _ => true
    }
}
"#,
        )
    }

    #[test]
    fn move_arm_cond_to_match_guard_with_else_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => if x > 10 {$0
            false
        } else {
            true
        }
        _ => true,
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => false,
        x => true,
        _ => true,
    }
}
"#,
        )
    }

    #[test]
    fn move_arm_cond_to_match_guard_with_else_block_works() {
        cov_mark::check!(move_guard_ifelse_expr_only);
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => {
            if x > 10 {$0
                false
            } else {
                true
            }
        }
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => false,
        x => true,
        _ => true
    }
}
"#,
        )
    }

    #[test]
    fn move_arm_cond_to_match_guard_else_if_empty_body_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => if x > 10 { $0 } else { },
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => {  },
        x => { }
        _ => true
    }
}
"#,
        );
    }

    #[test]
    fn move_arm_cond_to_match_guard_with_else_multiline_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => if x > 10 {
            92;$0
            false
        } else {
            true
        }
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => {
            92;
            false
        }
        x => true,
        _ => true
    }
}
"#,
        )
    }

    #[test]
    fn move_arm_cond_to_match_guard_with_else_multiline_else_works() {
        cov_mark::check!(move_guard_ifelse_else_block);
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => if x > 10 {$0
            false
        } else {
            42;
            true
        }
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => false,
        x => {
            42;
            true
        }
        _ => true
    }
}
"#,
        )
    }

    #[test]
    fn move_arm_cond_to_match_guard_with_else_multiline_else_block_works() {
        cov_mark::check!(move_guard_ifelse_in_block);
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        x => {
            if x > 10 {$0
                false
            } else {
                42;
                true
            }
        }
        _ => true
    }
}
"#,
            r#"
fn main() {
    match 92 {
        x if x > 10 => false,
        x => {
            42;
            true
        }
        _ => true
    }
}
"#,
        )
    }

    #[test]
    fn move_arm_cond_to_match_guard_with_else_last_arm_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        3 => true,
        x => {
            if x > 10 {$0
                false
            } else {
                92;
                true
            }
        }
    }
}
"#,
            r#"
fn main() {
    match 92 {
        3 => true,
        x if x > 10 => false,
        x => {
            92;
            true
        }
    }
}
"#,
        )
    }

    #[test]
    fn move_arm_cond_to_match_guard_with_else_comma_works() {
        check_assist(
            move_arm_cond_to_match_guard,
            r#"
fn main() {
    match 92 {
        3 => true,
        x => if x > 10 {$0
            false
        } else {
            92;
            true
        },
    }
}
"#,
            r#"
fn main() {
    match 92 {
        3 => true,
        x if x > 10 => false,
        x => {
            92;
            true
        }
    }
}
"#,
        )
    }
}
<|endoftext|>


<commit_before>//! FIXME: write short doc here

use ra_db::SourceDatabase;
use ra_syntax::{ast, AstNode};

use crate::{db::RootDatabase, display::ToNav, FilePosition, NavigationTarget, RangeInfo};

pub(crate) fn goto_type_definition(
    db: &RootDatabase,
    position: FilePosition,
) -> Option<RangeInfo<Vec<NavigationTarget>>> {
    let parse = db.parse(position.file_id);

    let node = parse.tree().syntax().token_at_offset(position.offset).find_map(|token| {
        token
            .parent()
            .ancestors()
            .find(|n| ast::Expr::cast(n.clone()).is_some() || ast::Pat::cast(n.clone()).is_some())
    })?;

    let analyzer =
        hir::SourceAnalyzer::new(db, hir::Source::new(position.file_id.into(), &node), None);

    let ty: hir::Ty = if let Some(ty) =
        ast::Expr::cast(node.clone()).and_then(|e| analyzer.type_of(db, &e))
    {
        ty
    } else if let Some(ty) = ast::Pat::cast(node.clone()).and_then(|p| analyzer.type_of_pat(db, &p))
    {
        ty
    } else {
        return None;
    };

    let adt_def = analyzer.autoderef(db, ty).find_map(|ty| ty.as_adt().map(|adt| adt.0))?;

    let nav = adt_def.to_nav(db);
    Some(RangeInfo::new(node.text_range(), vec![nav]))
}

#[cfg(test)]
mod tests {
    use crate::mock_analysis::analysis_and_position;

    fn check_goto(fixture: &str, expected: &str) {
        let (analysis, pos) = analysis_and_position(fixture);

        let mut navs = analysis.goto_type_definition(pos).unwrap().unwrap().info;
        assert_eq!(navs.len(), 1);
        let nav = navs.pop().unwrap();
        nav.assert_match(expected);
    }

    #[test]
    fn goto_type_definition_works_simple() {
        check_goto(
            "
            //- /lib.rs
            struct Foo;
            fn foo() {
                let f: Foo;
                f<|>
            }
            ",
            "Foo STRUCT_DEF FileId(1) [0; 11) [7; 10)",
        );
    }

    #[test]
    fn goto_type_definition_works_simple_ref() {
        check_goto(
            "
            //- /lib.rs
            struct Foo;
            fn foo() {
                let f: &Foo;
                f<|>
            }
            ",
            "Foo STRUCT_DEF FileId(1) [0; 11) [7; 10)",
        );
    }
}
<commit_msg>Goto type def through macros<commit_after>//! FIXME: write short doc here

use hir::db::AstDatabase;
use ra_syntax::{ast, AstNode};

use crate::{
    db::RootDatabase, display::ToNav, expand::descend_into_macros, FilePosition, NavigationTarget,
    RangeInfo,
};

pub(crate) fn goto_type_definition(
    db: &RootDatabase,
    position: FilePosition,
) -> Option<RangeInfo<Vec<NavigationTarget>>> {
    let file = db.parse_or_expand(position.file_id.into())?;
    let token = file.token_at_offset(position.offset).filter(|it| !it.kind().is_trivia()).next()?;
    let token = descend_into_macros(db, position.file_id, token);

    let node = token.ast.ancestors().find_map(|token| {
        token
            .ancestors()
            .find(|n| ast::Expr::cast(n.clone()).is_some() || ast::Pat::cast(n.clone()).is_some())
    })?;

    let analyzer = hir::SourceAnalyzer::new(db, token.with_ast(&node), None);

    let ty: hir::Ty = if let Some(ty) =
        ast::Expr::cast(node.clone()).and_then(|e| analyzer.type_of(db, &e))
    {
        ty
    } else if let Some(ty) = ast::Pat::cast(node.clone()).and_then(|p| analyzer.type_of_pat(db, &p))
    {
        ty
    } else {
        return None;
    };

    let adt_def = analyzer.autoderef(db, ty).find_map(|ty| ty.as_adt().map(|adt| adt.0))?;

    let nav = adt_def.to_nav(db);
    Some(RangeInfo::new(node.text_range(), vec![nav]))
}

#[cfg(test)]
mod tests {
    use crate::mock_analysis::analysis_and_position;

    fn check_goto(fixture: &str, expected: &str) {
        let (analysis, pos) = analysis_and_position(fixture);

        let mut navs = analysis.goto_type_definition(pos).unwrap().unwrap().info;
        assert_eq!(navs.len(), 1);
        let nav = navs.pop().unwrap();
        nav.assert_match(expected);
    }

    #[test]
    fn goto_type_definition_works_simple() {
        check_goto(
            "
            //- /lib.rs
            struct Foo;
            fn foo() {
                let f: Foo;
                f<|>
            }
            ",
            "Foo STRUCT_DEF FileId(1) [0; 11) [7; 10)",
        );
    }

    #[test]
    fn goto_type_definition_works_simple_ref() {
        check_goto(
            "
            //- /lib.rs
            struct Foo;
            fn foo() {
                let f: &Foo;
                f<|>
            }
            ",
            "Foo STRUCT_DEF FileId(1) [0; 11) [7; 10)",
        );
    }

    #[test]
    fn goto_type_definition_works_through_macro() {
        check_goto(
            "
            //- /lib.rs
            macro_rules! id {
                ($($tt:tt)*) => { $($tt)* }
            }
            struct Foo {}
            id! {
                fn bar() {
                    let f<|> = Foo {};
                }
            }
            ",
            "Foo STRUCT_DEF FileId(1) [52; 65) [59; 62)",
        );
    }
}
<|endoftext|>