<commit_before>// Copyright 2018, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

use futures_channel::{mpsc, oneshot};
use futures_core::future::Future;
use futures_core::stream::Stream;
use futures_core::task;
use futures_core::task::Poll;
use futures_util::future::FutureExt;
use futures_util::stream::StreamExt;
use std::marker::Unpin;
use std::pin;
use std::pin::Pin;

use Continue;
use MainContext;
use Priority;
use Source;

/// Represents a `Future` around a `glib::Source`. The future will
/// be resolved once the source has provided a value
pub struct SourceFuture<F, T> {
    create_source: Option<F>,
    source: Option<(Source, oneshot::Receiver<T>)>,
}

impl<F, T: 'static> SourceFuture<F, T>
where
    F: FnOnce(oneshot::Sender<T>) -> Source + 'static,
{
    /// Create a new `SourceFuture`
    ///
    /// The provided closure should return a newly created `glib::Source` when called
    /// and pass the value provided by the source to the oneshot sender that is passed
    /// to the closure.
    pub fn new(create_source: F) -> SourceFuture<F, T> {
        SourceFuture {
            create_source: Some(create_source),
            source: None,
        }
    }
}

impl<F, T> Unpin for SourceFuture<F, T> {}

impl<F, T> Future for SourceFuture<F, T>
where
    F: FnOnce(oneshot::Sender<T>) -> Source + 'static,
{
    type Output = T;

    fn poll(mut self: pin::Pin<&mut Self>, ctx: &mut task::Context) -> Poll<T> {
        let SourceFuture {
            ref mut create_source,
            ref mut source,
            ..
        } = *self;

        if let Some(create_source) = create_source.take() {
            let main_context = MainContext::ref_thread_default();
            assert!(
                main_context.is_owner(),
                "Spawning futures only allowed if the thread is owning the MainContext"
            );

            // Channel for sending back the Source result to our future here.
            //
            // In theory we could directly continue polling the
            // corresponding task from the Source callback,
            // however this would break at the very least
            // the g_main_current_source() API.
            let (send, recv) = oneshot::channel();

            let s = create_source(send);

            s.attach(Some(&main_context));
            *source = Some((s, recv));
        }

        // At this point we must have a receiver
        let res = {
            let &mut (_, ref mut receiver) = source.as_mut().unwrap();
            receiver.poll_unpin(ctx)
        };
        #[allow(clippy::match_wild_err_arm)]
        match res {
            Poll::Ready(Err(_)) => panic!("Source sender was unexpectedly closed"),
            Poll::Ready(Ok(v)) => {
                // Get rid of the reference to the source, it triggered
                let _ = source.take();
                Poll::Ready(v)
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

impl<T, F> Drop for SourceFuture<T, F> {
    fn drop(&mut self) {
        // Get rid of the source, we don't care anymore if it still triggers
        if let Some((source, _)) = self.source.take() {
            source.destroy();
        }
    }
}

/// Create a `Future` that will resolve after the given number of milliseconds.
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn timeout_future(value: u32) -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> {
    timeout_future_with_priority(::PRIORITY_DEFAULT, value)
}

/// Create a `Future` that will resolve after the given number of milliseconds.
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn timeout_future_with_priority(
    priority: Priority,
    value: u32,
) -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> {
    Box::pin(SourceFuture::new(move |send| {
        let mut send = Some(send);
        ::timeout_source_new(value, None, priority, move || {
            let _ = send.take().unwrap().send(());
            Continue(false)
        })
    }))
}

/// Create a `Future` that will resolve after the given number of seconds.
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn timeout_future_seconds(value: u32) -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> {
    timeout_future_seconds_with_priority(::PRIORITY_DEFAULT, value)
}

/// Create a `Future` that will resolve after the given number of seconds.
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn timeout_future_seconds_with_priority(
    priority: Priority,
    value: u32,
) -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> {
    Box::pin(SourceFuture::new(move |send| {
        let mut send = Some(send);
        ::timeout_source_new_seconds(value, None, priority, move || {
            let _ = send.take().unwrap().send(());
            Continue(false)
        })
    }))
}

/// Create a `Future` that will resolve once the child process with the given pid exits
///
/// The `Future` will resolve to the pid of the child process and the exit code.
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn child_watch_future(
    pid: ::Pid,
) -> Pin<Box<dyn Future<Output = (::Pid, i32)> + Send + 'static>> {
    child_watch_future_with_priority(::PRIORITY_DEFAULT, pid)
}

/// Create a `Future` that will resolve once the child process with the given pid exits
///
/// The `Future` will resolve to the pid of the child process and the exit code.
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn child_watch_future_with_priority(
    priority: Priority,
    pid: ::Pid,
) -> Pin<Box<dyn Future<Output = (::Pid, i32)> + Send + 'static>> {
    Box::pin(SourceFuture::new(move |send| {
        let mut send = Some(send);
        ::child_watch_source_new(pid, None, priority, move |pid, code| {
            let _ = send.take().unwrap().send((pid, code));
        })
    }))
}

#[cfg(any(unix, feature = "dox"))]
/// Create a `Future` that will resolve once the given UNIX signal is raised
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn unix_signal_future(signum: i32) -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> {
    unix_signal_future_with_priority(::PRIORITY_DEFAULT, signum)
}

#[cfg(any(unix, feature = "dox"))]
/// Create a `Future` that will resolve once the given UNIX signal is raised
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn unix_signal_future_with_priority(
    priority: Priority,
    signum: i32,
) -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> {
    Box::pin(SourceFuture::new(move |send| {
        let mut send = Some(send);
        ::unix_signal_source_new(signum, None, priority, move || {
            let _ = send.take().unwrap().send(());
            Continue(false)
        })
    }))
}

/// Represents a `Stream` around a `glib::Source`. The stream will
/// be provide all values that are provided by the source
pub struct SourceStream<F, T> {
    create_source: Option<F>,
    source: Option<(Source, mpsc::UnboundedReceiver<T>)>,
}

impl<F, T> Unpin for SourceStream<F, T> {}

impl<F, T: 'static> SourceStream<F, T>
where
    F: FnOnce(mpsc::UnboundedSender<T>) -> Source + 'static,
{
    /// Create a new `SourceStream`
    ///
    /// The provided closure should return a newly created `glib::Source` when called
    /// and pass the values provided by the source to the sender that is passed
    /// to the closure.
    pub fn new(create_source: F) -> SourceStream<F, T> {
        SourceStream {
            create_source: Some(create_source),
            source: None,
        }
    }
}

impl<F, T> Stream for SourceStream<F, T>
where
    F: FnOnce(mpsc::UnboundedSender<T>) -> Source + 'static,
{
    type Item = T;

    fn poll_next(mut self: pin::Pin<&mut Self>, ctx: &mut task::Context) -> Poll<Option<T>> {
        let SourceStream {
            ref mut create_source,
            ref mut source,
            ..
        } = *self;

        if let Some(create_source) = create_source.take() {
            let main_context = MainContext::ref_thread_default();
            assert!(
                main_context.is_owner(),
                "Spawning futures only allowed if the thread is owning the MainContext"
            );

            // Channel for sending back the Source result to our future here.
            //
            // In theory we could directly continue polling the
            // corresponding task from the Source callback,
            // however this would break at the very least
            // the g_main_current_source() API.
            let (send, recv) = mpsc::unbounded();

            let s = create_source(send);

            s.attach(Some(&main_context));
            *source = Some((s, recv));
        }

        // At this point we must have a receiver
        let res = {
            let &mut (_, ref mut receiver) = source.as_mut().unwrap();
            receiver.poll_next_unpin(ctx)
        };
        #[allow(clippy::match_wild_err_arm)]
        match res {
            Poll::Ready(v) => {
                if v.is_none() {
                    // Get rid of the reference to the source, it triggered
                    let _ = source.take();
                }
                Poll::Ready(v)
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

impl<T, F> Drop for SourceStream<T, F> {
    fn drop(&mut self) {
        // Get rid of the source, we don't care anymore if it still triggers
        if let Some((source, _)) = self.source.take() {
            source.destroy();
        }
    }
}

/// Create a `Stream` that will provide a value every given number of milliseconds.
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn interval_stream(value: u32) -> Pin<Box<dyn Stream<Item = ()> + Send + 'static>> {
    interval_stream_with_priority(::PRIORITY_DEFAULT, value)
}

/// Create a `Stream` that will provide a value every given number of milliseconds.
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn interval_stream_with_priority(
    priority: Priority,
    value: u32,
) -> Pin<Box<dyn Stream<Item = ()> + Send + 'static>> {
    Box::pin(SourceStream::new(move |send| {
        ::timeout_source_new(value, None, priority, move || {
            if send.unbounded_send(()).is_err() {
                Continue(false)
            } else {
                Continue(true)
            }
        })
    }))
}

/// Create a `Stream` that will provide a value every given number of seconds.
///
/// The `Stream` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn interval_stream_seconds(value: u32) -> Pin<Box<dyn Stream<Item = ()> + Send + 'static>> {
    interval_stream_seconds_with_priority(::PRIORITY_DEFAULT, value)
}

/// Create a `Stream` that will provide a value every given number of seconds.
///
/// The `Stream` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn interval_stream_seconds_with_priority(
    priority: Priority,
    value: u32,
) -> Pin<Box<dyn Stream<Item = ()> + Send + 'static>> {
    Box::pin(SourceStream::new(move |send| {
        ::timeout_source_new_seconds(value, None, priority, move || {
            if send.unbounded_send(()).is_err() {
                Continue(false)
            } else {
                Continue(true)
            }
        })
    }))
}

#[cfg(any(unix, feature = "dox"))]
/// Create a `Stream` that will provide a value whenever the given UNIX signal is raised
///
/// The `Stream` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn unix_signal_stream(signum: i32) -> Pin<Box<dyn Stream<Item = ()> + Send + 'static>> {
    unix_signal_stream_with_priority(::PRIORITY_DEFAULT, signum)
}

#[cfg(any(unix, feature = "dox"))]
/// Create a `Stream` that will provide a value whenever the given UNIX signal is raised
///
/// The `Stream` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn unix_signal_stream_with_priority(
    priority: Priority,
    signum: i32,
) -> Pin<Box<dyn Stream<Item = ()> + Send + 'static>> {
    Box::pin(SourceStream::new(move |send| {
        ::unix_signal_source_new(signum, None, priority, move || {
            if send.unbounded_send(()).is_err() {
                Continue(false)
            } else {
                Continue(true)
            }
        })
    }))
}

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

    #[test]
    fn test_timeout() {
        let c = MainContext::new();

        let res = c.block_on(timeout_future(20));

        assert_eq!(res, ());
    }

    #[test]
    fn test_timeout_send() {
        let c = MainContext::new();
        let l = ::MainLoop::new(Some(&c), false);

        let l_clone = l.clone();
        c.spawn(timeout_future(20).then(move |()| {
            l_clone.quit();
            futures_util::future::ready(())
        }));

        l.run();
    }

    #[test]
    fn test_interval() {
        let c = MainContext::new();

        let mut count = 0;

        {
            let count = &mut count;
            let res = c.block_on(
                interval_stream(20)
                    .take(2)
                    .for_each(|()| {
                        *count = *count + 1;

                        futures_util::future::ready(())
                    })
                    .map(|_| ()),
            );

            assert_eq!(res, ());
        }

        assert_eq!(count, 2);
    }

    #[test]
    fn test_timeout_and_channel() {
        let c = MainContext::default();

        let res = c.block_on(timeout_future(20).then(|()| {
            let (sender, receiver) = oneshot::channel();

            thread::spawn(move || {
                sender.send(1).unwrap();
            });

            receiver.then(|i| futures_util::future::ready(i.unwrap()))
        }));

        assert_eq!(res, 1);
    }
}
<commit_msg>clippy: Clean up issues in source_futures tests<commit_after>// Copyright 2018, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

use futures_channel::{mpsc, oneshot};
use futures_core::future::Future;
use futures_core::stream::Stream;
use futures_core::task;
use futures_core::task::Poll;
use futures_util::future::FutureExt;
use futures_util::stream::StreamExt;
use std::marker::Unpin;
use std::pin;
use std::pin::Pin;

use Continue;
use MainContext;
use Priority;
use Source;

/// Represents a `Future` around a `glib::Source`. The future will
/// be resolved once the source has provided a value
pub struct SourceFuture<F, T> {
    create_source: Option<F>,
    source: Option<(Source, oneshot::Receiver<T>)>,
}

impl<F, T: 'static> SourceFuture<F, T>
where
    F: FnOnce(oneshot::Sender<T>) -> Source + 'static,
{
    /// Create a new `SourceFuture`
    ///
    /// The provided closure should return a newly created `glib::Source` when called
    /// and pass the value provided by the source to the oneshot sender that is passed
    /// to the closure.
    pub fn new(create_source: F) -> SourceFuture<F, T> {
        SourceFuture {
            create_source: Some(create_source),
            source: None,
        }
    }
}

impl<F, T> Unpin for SourceFuture<F, T> {}

impl<F, T> Future for SourceFuture<F, T>
where
    F: FnOnce(oneshot::Sender<T>) -> Source + 'static,
{
    type Output = T;

    fn poll(mut self: pin::Pin<&mut Self>, ctx: &mut task::Context) -> Poll<T> {
        let SourceFuture {
            ref mut create_source,
            ref mut source,
            ..
        } = *self;

        if let Some(create_source) = create_source.take() {
            let main_context = MainContext::ref_thread_default();
            assert!(
                main_context.is_owner(),
                "Spawning futures only allowed if the thread is owning the MainContext"
            );

            // Channel for sending back the Source result to our future here.
            //
            // In theory we could directly continue polling the
            // corresponding task from the Source callback,
            // however this would break at the very least
            // the g_main_current_source() API.
            let (send, recv) = oneshot::channel();

            let s = create_source(send);

            s.attach(Some(&main_context));
            *source = Some((s, recv));
        }

        // At this point we must have a receiver
        let res = {
            let &mut (_, ref mut receiver) = source.as_mut().unwrap();
            receiver.poll_unpin(ctx)
        };
        #[allow(clippy::match_wild_err_arm)]
        match res {
            Poll::Ready(Err(_)) => panic!("Source sender was unexpectedly closed"),
            Poll::Ready(Ok(v)) => {
                // Get rid of the reference to the source, it triggered
                let _ = source.take();
                Poll::Ready(v)
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

impl<T, F> Drop for SourceFuture<T, F> {
    fn drop(&mut self) {
        // Get rid of the source, we don't care anymore if it still triggers
        if let Some((source, _)) = self.source.take() {
            source.destroy();
        }
    }
}

/// Create a `Future` that will resolve after the given number of milliseconds.
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn timeout_future(value: u32) -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> {
    timeout_future_with_priority(::PRIORITY_DEFAULT, value)
}

/// Create a `Future` that will resolve after the given number of milliseconds.
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn timeout_future_with_priority(
    priority: Priority,
    value: u32,
) -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> {
    Box::pin(SourceFuture::new(move |send| {
        let mut send = Some(send);
        ::timeout_source_new(value, None, priority, move || {
            let _ = send.take().unwrap().send(());
            Continue(false)
        })
    }))
}

/// Create a `Future` that will resolve after the given number of seconds.
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn timeout_future_seconds(value: u32) -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> {
    timeout_future_seconds_with_priority(::PRIORITY_DEFAULT, value)
}

/// Create a `Future` that will resolve after the given number of seconds.
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn timeout_future_seconds_with_priority(
    priority: Priority,
    value: u32,
) -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> {
    Box::pin(SourceFuture::new(move |send| {
        let mut send = Some(send);
        ::timeout_source_new_seconds(value, None, priority, move || {
            let _ = send.take().unwrap().send(());
            Continue(false)
        })
    }))
}

/// Create a `Future` that will resolve once the child process with the given pid exits
///
/// The `Future` will resolve to the pid of the child process and the exit code.
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn child_watch_future(
    pid: ::Pid,
) -> Pin<Box<dyn Future<Output = (::Pid, i32)> + Send + 'static>> {
    child_watch_future_with_priority(::PRIORITY_DEFAULT, pid)
}

/// Create a `Future` that will resolve once the child process with the given pid exits
///
/// The `Future` will resolve to the pid of the child process and the exit code.
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn child_watch_future_with_priority(
    priority: Priority,
    pid: ::Pid,
) -> Pin<Box<dyn Future<Output = (::Pid, i32)> + Send + 'static>> {
    Box::pin(SourceFuture::new(move |send| {
        let mut send = Some(send);
        ::child_watch_source_new(pid, None, priority, move |pid, code| {
            let _ = send.take().unwrap().send((pid, code));
        })
    }))
}

#[cfg(any(unix, feature = "dox"))]
/// Create a `Future` that will resolve once the given UNIX signal is raised
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn unix_signal_future(signum: i32) -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> {
    unix_signal_future_with_priority(::PRIORITY_DEFAULT, signum)
}

#[cfg(any(unix, feature = "dox"))]
/// Create a `Future` that will resolve once the given UNIX signal is raised
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn unix_signal_future_with_priority(
    priority: Priority,
    signum: i32,
) -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> {
    Box::pin(SourceFuture::new(move |send| {
        let mut send = Some(send);
        ::unix_signal_source_new(signum, None, priority, move || {
            let _ = send.take().unwrap().send(());
            Continue(false)
        })
    }))
}

/// Represents a `Stream` around a `glib::Source`. The stream will
/// be provide all values that are provided by the source
pub struct SourceStream<F, T> {
    create_source: Option<F>,
    source: Option<(Source, mpsc::UnboundedReceiver<T>)>,
}

impl<F, T> Unpin for SourceStream<F, T> {}

impl<F, T: 'static> SourceStream<F, T>
where
    F: FnOnce(mpsc::UnboundedSender<T>) -> Source + 'static,
{
    /// Create a new `SourceStream`
    ///
    /// The provided closure should return a newly created `glib::Source` when called
    /// and pass the values provided by the source to the sender that is passed
    /// to the closure.
    pub fn new(create_source: F) -> SourceStream<F, T> {
        SourceStream {
            create_source: Some(create_source),
            source: None,
        }
    }
}

impl<F, T> Stream for SourceStream<F, T>
where
    F: FnOnce(mpsc::UnboundedSender<T>) -> Source + 'static,
{
    type Item = T;

    fn poll_next(mut self: pin::Pin<&mut Self>, ctx: &mut task::Context) -> Poll<Option<T>> {
        let SourceStream {
            ref mut create_source,
            ref mut source,
            ..
        } = *self;

        if let Some(create_source) = create_source.take() {
            let main_context = MainContext::ref_thread_default();
            assert!(
                main_context.is_owner(),
                "Spawning futures only allowed if the thread is owning the MainContext"
            );

            // Channel for sending back the Source result to our future here.
            //
            // In theory we could directly continue polling the
            // corresponding task from the Source callback,
            // however this would break at the very least
            // the g_main_current_source() API.
            let (send, recv) = mpsc::unbounded();

            let s = create_source(send);

            s.attach(Some(&main_context));
            *source = Some((s, recv));
        }

        // At this point we must have a receiver
        let res = {
            let &mut (_, ref mut receiver) = source.as_mut().unwrap();
            receiver.poll_next_unpin(ctx)
        };
        #[allow(clippy::match_wild_err_arm)]
        match res {
            Poll::Ready(v) => {
                if v.is_none() {
                    // Get rid of the reference to the source, it triggered
                    let _ = source.take();
                }
                Poll::Ready(v)
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

impl<T, F> Drop for SourceStream<T, F> {
    fn drop(&mut self) {
        // Get rid of the source, we don't care anymore if it still triggers
        if let Some((source, _)) = self.source.take() {
            source.destroy();
        }
    }
}

/// Create a `Stream` that will provide a value every given number of milliseconds.
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn interval_stream(value: u32) -> Pin<Box<dyn Stream<Item = ()> + Send + 'static>> {
    interval_stream_with_priority(::PRIORITY_DEFAULT, value)
}

/// Create a `Stream` that will provide a value every given number of milliseconds.
///
/// The `Future` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn interval_stream_with_priority(
    priority: Priority,
    value: u32,
) -> Pin<Box<dyn Stream<Item = ()> + Send + 'static>> {
    Box::pin(SourceStream::new(move |send| {
        ::timeout_source_new(value, None, priority, move || {
            if send.unbounded_send(()).is_err() {
                Continue(false)
            } else {
                Continue(true)
            }
        })
    }))
}

/// Create a `Stream` that will provide a value every given number of seconds.
///
/// The `Stream` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn interval_stream_seconds(value: u32) -> Pin<Box<dyn Stream<Item = ()> + Send + 'static>> {
    interval_stream_seconds_with_priority(::PRIORITY_DEFAULT, value)
}

/// Create a `Stream` that will provide a value every given number of seconds.
///
/// The `Stream` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn interval_stream_seconds_with_priority(
    priority: Priority,
    value: u32,
) -> Pin<Box<dyn Stream<Item = ()> + Send + 'static>> {
    Box::pin(SourceStream::new(move |send| {
        ::timeout_source_new_seconds(value, None, priority, move || {
            if send.unbounded_send(()).is_err() {
                Continue(false)
            } else {
                Continue(true)
            }
        })
    }))
}

#[cfg(any(unix, feature = "dox"))]
/// Create a `Stream` that will provide a value whenever the given UNIX signal is raised
///
/// The `Stream` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn unix_signal_stream(signum: i32) -> Pin<Box<dyn Stream<Item = ()> + Send + 'static>> {
    unix_signal_stream_with_priority(::PRIORITY_DEFAULT, signum)
}

#[cfg(any(unix, feature = "dox"))]
/// Create a `Stream` that will provide a value whenever the given UNIX signal is raised
///
/// The `Stream` must be spawned on an `Executor` backed by a `glib::MainContext`.
pub fn unix_signal_stream_with_priority(
    priority: Priority,
    signum: i32,
) -> Pin<Box<dyn Stream<Item = ()> + Send + 'static>> {
    Box::pin(SourceStream::new(move |send| {
        ::unix_signal_source_new(signum, None, priority, move || {
            if send.unbounded_send(()).is_err() {
                Continue(false)
            } else {
                Continue(true)
            }
        })
    }))
}

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

    #[test]
    fn test_timeout() {
        let c = MainContext::new();

        c.block_on(timeout_future(20));
    }

    #[test]
    fn test_timeout_send() {
        let c = MainContext::new();
        let l = ::MainLoop::new(Some(&c), false);

        let l_clone = l.clone();
        c.spawn(timeout_future(20).then(move |()| {
            l_clone.quit();
            futures_util::future::ready(())
        }));

        l.run();
    }

    #[test]
    fn test_interval() {
        let c = MainContext::new();

        let mut count = 0;

        {
            let count = &mut count;
            c.block_on(
                interval_stream(20)
                    .take(2)
                    .for_each(|()| {
                        *count += 1;

                        futures_util::future::ready(())
                    })
                    .map(|_| ()),
            );
        }

        assert_eq!(count, 2);
    }

    #[test]
    fn test_timeout_and_channel() {
        let c = MainContext::default();

        let res = c.block_on(timeout_future(20).then(|()| {
            let (sender, receiver) = oneshot::channel();

            thread::spawn(move || {
                sender.send(1).unwrap();
            });

            receiver.then(|i| futures_util::future::ready(i.unwrap()))
        }));

        assert_eq!(res, 1);
    }
}
<|endoftext|>


<commit_before>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use std::collections::BTreeMap;
use std::io::BufRead;
use std::marker::Sized;
use std::result::Result as RResult;

use toml::Value;
use uuid::Uuid;

use task_hookrs::task::Task as TTask;
use task_hookrs::import::{import_task, import_tasks};

use libimagstore::store::{FileLockEntry, Store};
use libimagstore::storeid::{IntoStoreId, StoreIdIterator};
use module_path::ModuleEntryPath;

use error::{TodoErrorKind as TEK, MapErrInto};
use result::Result;

/// Task struct containing a `FileLockEntry`
pub trait Task<'a> {
    fn import_task_from_reader<R: BufRead>(store: &'a Store, r: R) -> Result<(Self, String, Uuid)>
        where Self: Sized;
    fn get_task_from_import<R: BufRead>(store: &'a Store, r: R) -> Result<RResult<Self, String>>
        where Self: Sized;
    fn get_task_from_string(store: &'a Store, s: String) -> Result<RResult<Self, String>>
        where Self: Sized;
    fn get_task_from_uuid(store: &'a Store, uuid: Uuid) -> Result<Option<Self>>
        where Self: Sized;
    fn retrieve_task_from_import<R: BufRead>(store: &'a Store, r: R) -> Result<Self>
        where Self: Sized;
    fn retrieve_task_from_string(store: &'a Store, s: String) -> Result<Self>
        where Self: Sized;
    fn delete_tasks_by_imports<R: BufRead>(store: &Store, r: R) -> Result<()>;
    fn delete_task_by_uuid(store: &Store, uuid: Uuid) -> Result<()>;
    fn all_tasks(store: &Store) -> Result<StoreIdIterator>;
    fn new_from_twtask(store: &'a Store, task: TTask) -> Result<Self>
        where Self: Sized;
}

impl<'a> Task<'a> for FileLockEntry<'a> {

    fn import_task_from_reader<R: BufRead>(store: &'a Store, mut r: R) -> Result<(FileLockEntry<'a>, String, Uuid)> {
        let mut line = String::new();
        try!(r.read_line(&mut line).map_err_into(TEK::UTF8Error));
        import_task(&line.as_str())
            .map_err_into(TEK::ImportError)
            .and_then(|t| {
                let uuid = t.uuid().clone();
                Self::new_from_twtask(store, t).map(|t| (t, line, uuid))
            })
    }

    /// Get a task from an import string. That is: read the imported string, get the UUID from it
    /// and try to load this UUID from store.
    ///
    /// Possible return values are:
    ///
    /// * Ok(Ok(Task))
    /// * Ok(Err(String)) - where the String is the String read from the `r` parameter
    /// * Err(_)          - where the error is an error that happened during evaluation
    ///
    fn get_task_from_import<R: BufRead>(store: &'a Store, mut r: R) -> Result<RResult<FileLockEntry<'a>, String>> {
        let mut line = String::new();
        try!(r.read_line(&mut line).map_err_into(TEK::UTF8Error));
        Self::get_task_from_string(store, line)
    }

    /// Get a task from a String. The String is expected to contain the JSON-representation of the
    /// Task to get from the store (only the UUID really matters in this case)
    ///
    /// For an explanation on the return values see `Task::get_from_import()`.
    fn get_task_from_string(store: &'a Store, s: String) -> Result<RResult<FileLockEntry<'a>, String>> {
        import_task(s.as_str())
            .map_err_into(TEK::ImportError)
            .map(|t| t.uuid().clone())
            .and_then(|uuid| Self::get_task_from_uuid(store, uuid))
            .and_then(|o| match o {
                None    => Ok(Err(s)),
                Some(t) => Ok(Ok(t)),
            })
    }

    /// Get a task from an UUID.
    ///
    /// If there is no task with this UUID, this returns `Ok(None)`.
    fn get_task_from_uuid(store: &'a Store, uuid: Uuid) -> Result<Option<FileLockEntry<'a>>> {
        ModuleEntryPath::new(format!("taskwarrior/{}", uuid))
            .into_storeid()
            .and_then(|store_id| store.get(store_id))
            .map_err_into(TEK::StoreError)
    }

    /// Same as Task::get_from_import() but uses Store::retrieve() rather than Store::get(), to
    /// implicitely create the task if it does not exist.
    fn retrieve_task_from_import<R: BufRead>(store: &'a Store, mut r: R) -> Result<FileLockEntry<'a>> {
        let mut line = String::new();
        try!(r.read_line(&mut line).map_err_into(TEK::UTF8Error));
        Self::retrieve_task_from_string(store, line)
    }

    /// Retrieve a task from a String. The String is expected to contain the JSON-representation of
    /// the Task to retrieve from the store (only the UUID really matters in this case)
    fn retrieve_task_from_string(store: &'a Store, s: String) -> Result<FileLockEntry<'a>> {
        Self::get_task_from_string(store, s)
            .and_then(|opt| match opt {
                Ok(task)    => Ok(task),
                Err(string) => import_task(string.as_str())
                    .map_err_into(TEK::ImportError)
                    .and_then(|t| Self::new_from_twtask(store, t)),
            })
    }

    fn delete_tasks_by_imports<R: BufRead>(store: &Store, r: R) -> Result<()> {
        use serde_json::ser::to_string as serde_to_string;
        use task_hookrs::status::TaskStatus;

        for (counter, res_ttask) in import_tasks(r).into_iter().enumerate() {
            match res_ttask {
                Ok(ttask) => {
                    if counter % 2 == 1 {
                        // Only every second task is needed, the first one is the
                        // task before the change, and the second one after
                        // the change. The (maybe modified) second one is
                        // expected by taskwarrior.
                        match serde_to_string(&ttask).map_err_into(TEK::ImportError) {
                            // use println!() here, as we talk with TW
                            Ok(val) => println!("{}", val),
                            Err(e)  => return Err(e),
                        }

                        // Taskwarrior does not have the concept of deleted tasks, but only modified
                        // ones.
                        //
                        // Here we check if the status of a task is deleted and if yes, we delete it
                        // from the store.
                        if *ttask.status() == TaskStatus::Deleted {
                            match Self::delete_task_by_uuid(store, *ttask.uuid()) {
                                Ok(_)  => info!("Deleted task {}", *ttask.uuid()),
                                Err(e) => return Err(e),
                            }
                        }
                    } // end if c % 2
                },
                Err(e) => return Err(e).map_err_into(TEK::ImportError),
            }
        }
        Ok(())
    }

    fn delete_task_by_uuid(store: &Store, uuid: Uuid) -> Result<()> {
        ModuleEntryPath::new(format!("taskwarrior/{}", uuid))
            .into_storeid()
            .and_then(|id| store.delete(id))
            .map_err_into(TEK::StoreError)
    }

    fn all_tasks(store: &Store) -> Result<StoreIdIterator> {
        store.retrieve_for_module("todo/taskwarrior")
            .map_err_into(TEK::StoreError)
    }

    fn new_from_twtask(store: &'a Store, task: TTask) -> Result<FileLockEntry<'a>> {
        use toml_query::read::TomlValueReadExt;
        use toml_query::set::TomlValueSetExt;

        let uuid     = task.uuid();
        ModuleEntryPath::new(format!("taskwarrior/{}", uuid))
            .into_storeid()
            .map_err_into(TEK::StoreIdError)
            .and_then(|id| {
                store.retrieve(id)
                    .map_err_into(TEK::StoreError)
                    .and_then(|mut fle| {
                        {
                            let hdr = fle.get_header_mut();
                            if try!(hdr.read("todo").map_err_into(TEK::StoreError)).is_none() {
                                try!(hdr
                                    .set("todo", Value::Table(BTreeMap::new()))
                                    .map_err_into(TEK::StoreError));
                            }

                            try!(hdr.set("todo.uuid", Value::String(format!("{}",uuid)))
                                 .map_err_into(TEK::StoreError));
                        }

                        // If none of the errors above have returned the function, everything is fine
                        Ok(fle)
                    })
            })

    }

}

<commit_msg>Changed Task trait to TaskStore trait, implemented for the Store instead of the FileLockEntry<commit_after>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use std::collections::BTreeMap;
use std::io::BufRead;
use std::result::Result as RResult;

use toml::Value;
use uuid::Uuid;

use task_hookrs::task::Task as TTask;
use task_hookrs::import::{import_task, import_tasks};

use libimagstore::store::{FileLockEntry, Store};
use libimagstore::storeid::{IntoStoreId, StoreIdIterator};
use module_path::ModuleEntryPath;

use error::{TodoErrorKind as TEK, MapErrInto};
use result::Result;

/// Task struct containing a `FileLockEntry`
pub trait TaskStore<'a> {
    fn import_task_from_reader<R: BufRead>(&'a self, r: R) -> Result<(FileLockEntry<'a>, String, Uuid)>;
    fn get_task_from_import<R: BufRead>(&'a self, r: R) -> Result<RResult<FileLockEntry<'a>, String>>;
    fn get_task_from_string(&'a self, s: String) -> Result<RResult<FileLockEntry<'a>, String>>;
    fn get_task_from_uuid(&'a self, uuid: Uuid) -> Result<Option<FileLockEntry<'a>>>;
    fn retrieve_task_from_import<R: BufRead>(&'a self, r: R) -> Result<FileLockEntry<'a>>;
    fn retrieve_task_from_string(&'a self, s: String) -> Result<FileLockEntry<'a>>;
    fn delete_tasks_by_imports<R: BufRead>(&self, r: R) -> Result<()>;
    fn delete_task_by_uuid(&self, uuid: Uuid) -> Result<()>;
    fn all_tasks(&self) -> Result<StoreIdIterator>;
    fn new_from_twtask(&'a self, task: TTask) -> Result<FileLockEntry<'a>>;
}

impl<'a> TaskStore<'a> for Store {

    fn import_task_from_reader<R: BufRead>(&'a self, mut r: R) -> Result<(FileLockEntry<'a>, String, Uuid)> {
        let mut line = String::new();
        try!(r.read_line(&mut line).map_err_into(TEK::UTF8Error));
        import_task(&line.as_str())
            .map_err_into(TEK::ImportError)
            .and_then(|t| {
                let uuid = t.uuid().clone();
                self.new_from_twtask(t).map(|t| (t, line, uuid))
            })
    }

    /// Get a task from an import string. That is: read the imported string, get the UUID from it
    /// and try to load this UUID from store.
    ///
    /// Possible return values are:
    ///
    /// * Ok(Ok(Task))
    /// * Ok(Err(String)) - where the String is the String read from the `r` parameter
    /// * Err(_)          - where the error is an error that happened during evaluation
    ///
    fn get_task_from_import<R: BufRead>(&'a self, mut r: R) -> Result<RResult<FileLockEntry<'a>, String>> {
        let mut line = String::new();
        try!(r.read_line(&mut line).map_err_into(TEK::UTF8Error));
        self.get_task_from_string(line)
    }

    /// Get a task from a String. The String is expected to contain the JSON-representation of the
    /// Task to get from the store (only the UUID really matters in this case)
    ///
    /// For an explanation on the return values see `Task::get_from_import()`.
    fn get_task_from_string(&'a self, s: String) -> Result<RResult<FileLockEntry<'a>, String>> {
        import_task(s.as_str())
            .map_err_into(TEK::ImportError)
            .map(|t| t.uuid().clone())
            .and_then(|uuid| self.get_task_from_uuid(uuid))
            .and_then(|o| match o {
                None    => Ok(Err(s)),
                Some(t) => Ok(Ok(t)),
            })
    }

    /// Get a task from an UUID.
    ///
    /// If there is no task with this UUID, this returns `Ok(None)`.
    fn get_task_from_uuid(&'a self, uuid: Uuid) -> Result<Option<FileLockEntry<'a>>> {
        ModuleEntryPath::new(format!("taskwarrior/{}", uuid))
            .into_storeid()
            .and_then(|store_id| self.get(store_id))
            .map_err_into(TEK::StoreError)
    }

    /// Same as Task::get_from_import() but uses Store::retrieve() rather than Store::get(), to
    /// implicitely create the task if it does not exist.
    fn retrieve_task_from_import<R: BufRead>(&'a self, mut r: R) -> Result<FileLockEntry<'a>> {
        let mut line = String::new();
        try!(r.read_line(&mut line).map_err_into(TEK::UTF8Error));
        self.retrieve_task_from_string(line)
    }

    /// Retrieve a task from a String. The String is expected to contain the JSON-representation of
    /// the Task to retrieve from the store (only the UUID really matters in this case)
    fn retrieve_task_from_string(&'a self, s: String) -> Result<FileLockEntry<'a>> {
        self.get_task_from_string(s)
            .and_then(|opt| match opt {
                Ok(task)    => Ok(task),
                Err(string) => import_task(string.as_str())
                    .map_err_into(TEK::ImportError)
                    .and_then(|t| self.new_from_twtask(t)),
            })
    }

    fn delete_tasks_by_imports<R: BufRead>(&self, r: R) -> Result<()> {
        use serde_json::ser::to_string as serde_to_string;
        use task_hookrs::status::TaskStatus;

        for (counter, res_ttask) in import_tasks(r).into_iter().enumerate() {
            match res_ttask {
                Ok(ttask) => {
                    if counter % 2 == 1 {
                        // Only every second task is needed, the first one is the
                        // task before the change, and the second one after
                        // the change. The (maybe modified) second one is
                        // expected by taskwarrior.
                        match serde_to_string(&ttask).map_err_into(TEK::ImportError) {
                            // use println!() here, as we talk with TW
                            Ok(val) => println!("{}", val),
                            Err(e)  => return Err(e),
                        }

                        // Taskwarrior does not have the concept of deleted tasks, but only modified
                        // ones.
                        //
                        // Here we check if the status of a task is deleted and if yes, we delete it
                        // from the store.
                        if *ttask.status() == TaskStatus::Deleted {
                            match self.delete_task_by_uuid(*ttask.uuid()) {
                                Ok(_)  => info!("Deleted task {}", *ttask.uuid()),
                                Err(e) => return Err(e),
                            }
                        }
                    } // end if c % 2
                },
                Err(e) => return Err(e).map_err_into(TEK::ImportError),
            }
        }
        Ok(())
    }

    fn delete_task_by_uuid(&self, uuid: Uuid) -> Result<()> {
        ModuleEntryPath::new(format!("taskwarrior/{}", uuid))
            .into_storeid()
            .and_then(|id| self.delete(id))
            .map_err_into(TEK::StoreError)
    }

    fn all_tasks(&self) -> Result<StoreIdIterator> {
        self.retrieve_for_module("todo/taskwarrior")
            .map_err_into(TEK::StoreError)
    }

    fn new_from_twtask(&'a self, task: TTask) -> Result<FileLockEntry<'a>> {
        use toml_query::read::TomlValueReadExt;
        use toml_query::set::TomlValueSetExt;

        let uuid     = task.uuid();
        ModuleEntryPath::new(format!("taskwarrior/{}", uuid))
            .into_storeid()
            .map_err_into(TEK::StoreIdError)
            .and_then(|id| {
                self.retrieve(id)
                    .map_err_into(TEK::StoreError)
                    .and_then(|mut fle| {
                        {
                            let hdr = fle.get_header_mut();
                            if try!(hdr.read("todo").map_err_into(TEK::StoreError)).is_none() {
                                try!(hdr
                                    .set("todo", Value::Table(BTreeMap::new()))
                                    .map_err_into(TEK::StoreError));
                            }

                            try!(hdr.set("todo.uuid", Value::String(format!("{}",uuid)))
                                 .map_err_into(TEK::StoreError));
                        }

                        // If none of the errors above have returned the function, everything is fine
                        Ok(fle)
                    })
            })

    }

}

<|endoftext|>


<commit_before>use std::cmp::min;
use std::ops::Index;

use objc::{Id, IdSlice, IntoIdVector, Owned, Ownership, ShareId};

use {INSArray, INSCopying, INSObject, NSArray, NSEnumerator};

pub trait INSDictionary : INSObject {
    type Key: INSObject;
    type Value: INSObject;
    type Own: Ownership;

    fn count(&self) -> usize {
        unsafe {
            msg_send![self, count]
        }
    }

    fn object_for(&self, key: &Self::Key) -> Option<&Self::Value> {
        unsafe {
            let obj: *mut Self::Value = msg_send![self, objectForKey:key];
            obj.as_ref()
        }
    }

    fn all_keys(&self) -> Vec<&Self::Key> {
        let keys = unsafe {
            let keys: *mut NSArray<Self::Key> = msg_send![self, allKeys];
            &*keys
        };
        keys.to_vec()
    }

    fn all_values(&self) -> Vec<&Self::Value> {
        let vals = unsafe {
            let vals: *mut NSArray<Self::Value> = msg_send![self, allValues];
            &*vals
        };
        vals.to_vec()
    }

    fn key_enumerator(&self) -> NSEnumerator<Self::Key> {
        unsafe {
            let result = msg_send![self, keyEnumerator];
            NSEnumerator::from_ptr(result)
        }
    }

    fn object_enumerator(&self) -> NSEnumerator<Self::Value> {
        unsafe {
            let result = msg_send![self, objectEnumerator];
            NSEnumerator::from_ptr(result)
        }
    }

    fn keys_and_objects(&self) -> (Vec<&Self::Key>, Vec<&Self::Value>) {
        let len = self.count();
        let mut keys: Vec<&Self::Key> = Vec::with_capacity(len);
        let mut objs: Vec<&Self::Value> = Vec::with_capacity(len);
        unsafe {
            let _: () = msg_send![self, getObjects:objs.as_ptr()
                                           andKeys:keys.as_ptr()];
            keys.set_len(len);
            objs.set_len(len);
        }
        (keys, objs)
    }

    unsafe fn from_refs<T>(keys: &[&T], vals: &[&Self::Value]) -> Id<Self>
            where T: INSCopying<Output=Self::Key> {
        let cls = <Self as INSObject>::class();
        let count = min(keys.len(), vals.len());
        let obj: *mut Self = msg_send![cls, alloc];
        let obj: *mut Self = msg_send![obj, initWithObjects:vals.as_ptr()
                                                    forKeys:keys.as_ptr()
                                                      count:count];
        Id::from_retained_ptr(obj)
    }

    fn from_keys_and_objects<T>(keys: &[&T],
            vals: Vec<Id<Self::Value, Self::Own>>) -> Id<Self>
            where T: INSCopying<Output=Self::Key> {
        let vals_refs = vals.as_refs_slice();
        unsafe {
            INSDictionary::from_refs(keys, vals_refs)
        }
    }

    fn into_keys_and_objects(dict: Id<Self>) ->
            (Vec<ShareId<Self::Key>>, Vec<Id<Self::Value, Self::Own>>) {
        let (keys, objs) = dict.keys_and_objects();
        unsafe {
            (keys.into_id_vec(), objs.into_id_vec())
        }
    }
}

object_struct!(NSDictionary<K, V>);

impl<K, V> INSDictionary for NSDictionary<K, V>
        where K: INSObject, V: INSObject {
    type Key = K;
    type Value = V;
    type Own = Owned;
}

impl<K, V> Index<K> for NSDictionary<K, V> where K: INSObject, V: INSObject {
    type Output = V;

    fn index(&self, index: &K) -> &V {
        self.object_for(index).unwrap()
    }
}

#[cfg(test)]
mod tests {
    use objc::{Id};
    use {INSObject, INSString, NSObject, NSString};
    use super::{INSDictionary, NSDictionary};

    fn sample_dict(key: &str) -> Id<NSDictionary<NSString, NSObject>> {
        let string: Id<NSString> = INSString::from_str(key);
        let obj: Id<NSObject> = INSObject::new();
        INSDictionary::from_keys_and_objects(&[&*string], vec![obj])
    }

    #[test]
    fn test_count() {
        let dict = sample_dict("abcd");
        assert!(dict.count() == 1);
    }

    #[test]
    fn test_object_for() {
        let dict = sample_dict("abcd");

        let string: Id<NSString> = INSString::from_str("abcd");
        assert!(dict.object_for(&*string).is_some());

        let string: Id<NSString> = INSString::from_str("abcde");
        assert!(dict.object_for(&*string).is_none());
    }

    #[test]
    fn test_all_keys() {
        let dict = sample_dict("abcd");
        let keys = dict.all_keys();

        assert!(keys.len() == 1);
        assert!(keys[0].as_str() == "abcd");
    }

    #[test]
    fn test_all_values() {
        let dict = sample_dict("abcd");
        let vals = dict.all_values();

        assert!(vals.len() == 1);
    }

    #[test]
    fn test_keys_and_objects() {
        let dict = sample_dict("abcd");
        let (keys, objs) = dict.keys_and_objects();

        assert!(keys.len() == 1);
        assert!(objs.len() == 1);
        assert!(keys[0].as_str() == "abcd");
        assert!(objs[0] == dict.object_for(keys[0]).unwrap());
    }

    #[test]
    fn test_key_enumerator() {
        let dict = sample_dict("abcd");
        assert!(dict.key_enumerator().count() == 1);
        assert!(dict.key_enumerator().next().unwrap().as_str() == "abcd");
    }

    #[test]
    fn test_object_enumerator() {
        let dict = sample_dict("abcd");
        assert!(dict.object_enumerator().count() == 1);
    }
}
<commit_msg>Modified some dictionary methods to return arrays.<commit_after>use std::cmp::min;
use std::ops::Index;
use std::ptr;

use objc::{Id, IdSlice, Owned, Ownership};

use {INSArray, INSCopying, INSObject, NSArray, NSSharedArray, NSEnumerator};

pub trait INSDictionary : INSObject {
    type Key: INSObject;
    type Value: INSObject;
    type Own: Ownership;

    fn count(&self) -> usize {
        unsafe {
            msg_send![self, count]
        }
    }

    fn object_for(&self, key: &Self::Key) -> Option<&Self::Value> {
        unsafe {
            let obj: *mut Self::Value = msg_send![self, objectForKey:key];
            obj.as_ref()
        }
    }

    fn keys(&self) -> Vec<&Self::Key> {
        let len = self.count();
        let mut keys: Vec<&Self::Key> = Vec::with_capacity(len);
        unsafe {
            let _: () = msg_send![self, getObjects:ptr::null_mut::<Self::Value>()
                                           andKeys:keys.as_mut_ptr()];
            keys.set_len(len);
        }
        keys
    }

    fn values(&self) -> Vec<&Self::Value> {
        let len = self.count();
        let mut vals: Vec<&Self::Value> = Vec::with_capacity(len);
        unsafe {
            let _: () = msg_send![self, getObjects:vals.as_mut_ptr()
                                           andKeys:ptr::null_mut::<Self::Key>()];
            vals.set_len(len);
        }
        vals
    }

    fn keys_and_objects(&self) -> (Vec<&Self::Key>, Vec<&Self::Value>) {
        let len = self.count();
        let mut keys: Vec<&Self::Key> = Vec::with_capacity(len);
        let mut objs: Vec<&Self::Value> = Vec::with_capacity(len);
        unsafe {
            let _: () = msg_send![self, getObjects:objs.as_mut_ptr()
                                           andKeys:keys.as_mut_ptr()];
            keys.set_len(len);
            objs.set_len(len);
        }
        (keys, objs)
    }

    fn key_enumerator(&self) -> NSEnumerator<Self::Key> {
        unsafe {
            let result = msg_send![self, keyEnumerator];
            NSEnumerator::from_ptr(result)
        }
    }

    fn object_enumerator(&self) -> NSEnumerator<Self::Value> {
        unsafe {
            let result = msg_send![self, objectEnumerator];
            NSEnumerator::from_ptr(result)
        }
    }

    fn keys_array(&self) -> Id<NSSharedArray<Self::Key>> {
        unsafe {
            let keys: *mut NSSharedArray<Self::Key> = msg_send![self, allKeys];
            Id::from_ptr(keys)
        }
    }

    unsafe fn from_refs<T>(keys: &[&T], vals: &[&Self::Value]) -> Id<Self>
            where T: INSCopying<Output=Self::Key> {
        let cls = <Self as INSObject>::class();
        let count = min(keys.len(), vals.len());
        let obj: *mut Self = msg_send![cls, alloc];
        let obj: *mut Self = msg_send![obj, initWithObjects:vals.as_ptr()
                                                    forKeys:keys.as_ptr()
                                                      count:count];
        Id::from_retained_ptr(obj)
    }

    fn from_keys_and_objects<T>(keys: &[&T],
            vals: Vec<Id<Self::Value, Self::Own>>) -> Id<Self>
            where T: INSCopying<Output=Self::Key> {
        let vals_refs = vals.as_refs_slice();
        unsafe {
            INSDictionary::from_refs(keys, vals_refs)
        }
    }

    fn into_values_array(dict: Id<Self>) -> Id<NSArray<Self::Value, Self::Own>> {
        unsafe {
            let vals = msg_send![dict, allValues];
            Id::from_ptr(vals)
        }
    }
}

object_struct!(NSDictionary<K, V>);

impl<K, V> INSDictionary for NSDictionary<K, V>
        where K: INSObject, V: INSObject {
    type Key = K;
    type Value = V;
    type Own = Owned;
}

impl<K, V> Index<K> for NSDictionary<K, V> where K: INSObject, V: INSObject {
    type Output = V;

    fn index(&self, index: &K) -> &V {
        self.object_for(index).unwrap()
    }
}

#[cfg(test)]
mod tests {
    use objc::{Id};
    use {INSArray, INSObject, INSString, NSObject, NSString};
    use super::{INSDictionary, NSDictionary};

    fn sample_dict(key: &str) -> Id<NSDictionary<NSString, NSObject>> {
        let string: Id<NSString> = INSString::from_str(key);
        let obj: Id<NSObject> = INSObject::new();
        INSDictionary::from_keys_and_objects(&[&*string], vec![obj])
    }

    #[test]
    fn test_count() {
        let dict = sample_dict("abcd");
        assert!(dict.count() == 1);
    }

    #[test]
    fn test_object_for() {
        let dict = sample_dict("abcd");

        let string: Id<NSString> = INSString::from_str("abcd");
        assert!(dict.object_for(&*string).is_some());

        let string: Id<NSString> = INSString::from_str("abcde");
        assert!(dict.object_for(&*string).is_none());
    }

    #[test]
    fn test_keys() {
        let dict = sample_dict("abcd");
        let keys = dict.keys();

        assert!(keys.len() == 1);
        assert!(keys[0].as_str() == "abcd");
    }

    #[test]
    fn test_values() {
        let dict = sample_dict("abcd");
        let vals = dict.values();

        assert!(vals.len() == 1);
    }

    #[test]
    fn test_keys_and_objects() {
        let dict = sample_dict("abcd");
        let (keys, objs) = dict.keys_and_objects();

        assert!(keys.len() == 1);
        assert!(objs.len() == 1);
        assert!(keys[0].as_str() == "abcd");
        assert!(objs[0] == dict.object_for(keys[0]).unwrap());
    }

    #[test]
    fn test_key_enumerator() {
        let dict = sample_dict("abcd");
        assert!(dict.key_enumerator().count() == 1);
        assert!(dict.key_enumerator().next().unwrap().as_str() == "abcd");
    }

    #[test]
    fn test_object_enumerator() {
        let dict = sample_dict("abcd");
        assert!(dict.object_enumerator().count() == 1);
    }

    #[test]
    fn test_arrays() {
        let dict = sample_dict("abcd");

        let keys = dict.keys_array();
        assert!(keys.count() == 1);
        assert!(keys.object_at(0).as_str() == "abcd");

        let objs = INSDictionary::into_values_array(dict);
        assert!(objs.count() == 1);
    }
}
<|endoftext|>


<commit_before>use std::time::Duration;
use stdweb::web::{window, IWindowOrWorker, TimeoutHandle};

#[derive(Debug)]
pub struct Timeout {
    handle: TimeoutHandle,
}

impl Timeout {
    pub fn new<F>(f: F, duration: Duration) -> Timeout
    where
        F: 'static + FnMut(),
    {
        Timeout {
            handle: window().set_clearable_timeout(f, duration.as_millis() as u32),
        }
    }
}

impl Drop for Timeout {
    fn drop(&mut self) {
        let handle = std::mem::replace(&mut self.handle, unsafe { std::mem::uninitialized() });
        handle.clear();
    }
}
<commit_msg>fix: remove deprecated usage of `mem::uninitialized`. (#1341)<commit_after>use std::time::Duration;
use stdweb::web::{window, IWindowOrWorker, TimeoutHandle};

#[derive(Debug)]
pub struct Timeout {
    handle: Option<TimeoutHandle>,
}

impl Timeout {
    pub fn new<F>(f: F, duration: Duration) -> Timeout
    where
        F: 'static + FnMut(),
    {
        Timeout {
            handle: Some(window().set_clearable_timeout(f, duration.as_millis() as u32)),
        }
    }
}

impl Drop for Timeout {
    fn drop(&mut self) {
        let handle = self.handle.take().unwrap();
        handle.clear();
    }
}
<|endoftext|>


<commit_before>//! Definition of the Shared combinator, a future that is cloneable,
//! and can be polled in multiple threads.
//!
//! # Examples
//!
//! ```
//! # extern crate futures;
//! # extern crate futures_executor;
//! use futures::prelude::*;
//! use futures::future;
//! use futures_executor::block_on;
//!
//! # fn main() {
//! let future = future::ready(6);
//! let shared1 = future.shared();
//! let shared2 = shared1.clone();
//! assert_eq!(6, *block_on(shared1));
//! assert_eq!(6, *block_on(shared2));
//! # }
//! ```

use slab::Slab;

use std::fmt;
use std::cell::UnsafeCell;
use std::marker::Unpin;
use std::mem::PinMut;
use std::sync::{Arc, Mutex};
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::SeqCst;
use std::task::local_waker_from_nonlocal;

use futures_core::{Future, Poll};
use futures_core::task::{self, Wake, Waker};

/// A future that is cloneable and can be polled in multiple threads.
/// Use `Future::shared()` method to convert any future into a `Shared` future.
#[must_use = "futures do nothing unless polled"]
pub struct Shared<F: Future> {
    inner: Arc<Inner<F>>,
    waker_key: usize,
}

struct Inner<F: Future> {
    future_or_output: UnsafeCell<FutureOrOutput<F>>,
    notifier: Arc<Notifier>,
}

struct Notifier {
    state: AtomicUsize,
    wakers: Mutex<Slab<Option<Waker>>>,
}

// The future itself is polled behind the `Arc`, so it won't be moved
// when `Shared` is moved.
impl<F: Future> Unpin for Shared<F> {}

impl<F: Future> fmt::Debug for Shared<F> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Shared")
            .field("inner", &self.inner)
            .field("waker_key", &self.waker_key)
            .finish()
    }
}

impl<F: Future> fmt::Debug for Inner<F> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Inner")
            .finish()
    }
}

enum FutureOrOutput<F: Future> {
    Future(F),
    Output(Arc<F::Output>),
}

unsafe impl<F> Send for Inner<F>
where
    F: Future + Send,
    F::Output: Send + Sync,
{}

unsafe impl<F> Sync for Inner<F>
where
    F: Future + Send,
    F::Output: Send + Sync,
{}

const IDLE: usize = 0;
const POLLING: usize = 1;
const REPOLL: usize = 2;
const COMPLETE: usize = 3;
const POISONED: usize = 4;

const NULL_WAKER_KEY: usize = usize::max_value();

pub fn new<F: Future>(future: F) -> Shared<F> {
    Shared {
        inner: Arc::new(Inner {
            future_or_output: UnsafeCell::new(FutureOrOutput::Future(future)),
            notifier: Arc::new(Notifier {
                state: AtomicUsize::new(IDLE),
                wakers: Mutex::new(Slab::new()),
            }),
        }),
        waker_key: NULL_WAKER_KEY,
    }
}

impl<F> Shared<F> where F: Future {
    /// If any clone of this `Shared` has completed execution, returns its result immediately
    /// without blocking. Otherwise, returns None without triggering the work represented by
    /// this `Shared`.
    pub fn peek(&self) -> Option<Arc<F::Output>> {
        match self.inner.notifier.state.load(SeqCst) {
            COMPLETE => {
                Some(unsafe { self.clone_output() })
            }
            POISONED => panic!("inner future panicked during poll"),
            _ => None,
        }
    }

    /// Registers the current task to receive a wakeup when `Inner` is awoken.
    fn set_waker(&mut self, cx: &mut task::Context) {
        let mut wakers = self.inner.notifier.wakers.lock().unwrap();
        if self.waker_key == NULL_WAKER_KEY {
            self.waker_key = wakers.insert(Some(cx.waker().clone()));
        } else {
            let waker_slot = &mut wakers[self.waker_key];
            let needs_replacement = if let Some(old_waker) = waker_slot {
                // If there's still an unwoken waker in the slot, only replace
                // if the current one wouldn't wake the same task.
                !old_waker.will_wake(cx.waker())
            } else {
                true
            };
            if needs_replacement {
                *waker_slot = Some(cx.waker().clone());
            }
        }
        debug_assert!(self.waker_key != NULL_WAKER_KEY);
    }

    /// Safety: callers must first ensure that `self.inner.state`
    /// is `COMPLETE`
    unsafe fn clone_output(&self) -> Arc<F::Output> {
        if let FutureOrOutput::Output(item) = &*self.inner.future_or_output.get() {
            item.clone()
        } else {
            unreachable!()
        }
    }
}

impl<F: Future> Future for Shared<F> {
    type Output = Arc<F::Output>;

    fn poll(mut self: PinMut<Self>, cx: &mut task::Context) -> Poll<Self::Output> {
        let this = &mut *self;

        this.set_waker(cx);

        match this.inner.notifier.state.compare_and_swap(IDLE, POLLING, SeqCst) {
            IDLE => {
                // Lock acquired, fall through
            }
            POLLING | REPOLL => {
                // Another task is currently polling, at this point we just want
                // to ensure that our task handle is currently registered

                return Poll::Pending
            }
            COMPLETE => {
                return unsafe { Poll::Ready(this.clone_output()) };
            }
            POISONED => panic!("inner future panicked during poll"),
            _ => unreachable!(),
        }

        let waker = local_waker_from_nonlocal(this.inner.notifier.clone());
        let mut cx = cx.with_waker(&waker);

        loop {
            struct Reset<'a>(&'a AtomicUsize);

            impl<'a> Drop for Reset<'a> {
                fn drop(&mut self) {
                    use std::thread;

                    if thread::panicking() {
                        self.0.store(POISONED, SeqCst);
                    }
                }
            }

            let _reset = Reset(&this.inner.notifier.state);

            // Poll the future
            let res = unsafe {
                if let FutureOrOutput::Future(future) = &mut *this.inner.future_or_output.get() {
                    PinMut::new_unchecked(future).poll(&mut cx)
                } else {
                    unreachable!()
                }
            };
            match res {
                Poll::Pending => {
                    // Not ready, try to release the handle
                    match this.inner.notifier.state.compare_and_swap(POLLING, IDLE, SeqCst) {
                        POLLING => {
                            // Success
                            return Poll::Pending;
                        }
                        REPOLL => {
                            // Gotta poll again!
                            let prev = this.inner.notifier.state.swap(POLLING, SeqCst);
                            assert_eq!(prev, REPOLL);
                        }
                        _ => unreachable!(),
                    }
                }
                Poll::Ready(output) => {
                    let output = Arc::new(output);
                    unsafe {
                        *this.inner.future_or_output.get() =
                            FutureOrOutput::Output(output.clone());
                    }
                    // Complete the future
                    this.inner.notifier.state.store(COMPLETE, SeqCst);
                    Wake::wake(&this.inner.notifier);
                    return Poll::Ready(output);
                }
            }
        }
    }
}

impl<F> Clone for Shared<F> where F: Future {
    fn clone(&self) -> Self {
        Shared {
            inner: self.inner.clone(),
            waker_key: NULL_WAKER_KEY,
        }
    }
}

impl<F> Drop for Shared<F> where F: Future {
    fn drop(&mut self) {
        if self.waker_key != NULL_WAKER_KEY {
            if let Ok(mut wakers) = self.inner.notifier.wakers.lock() {
                wakers.remove(self.waker_key);
            }
        }
    }
}

impl Wake for Notifier {
    fn wake(arc_self: &Arc<Self>) {
        arc_self.state.compare_and_swap(POLLING, REPOLL, SeqCst);

        let wakers = &mut *arc_self.wakers.lock().unwrap();
        for (_key, opt_waker) in wakers {
            if let Some(waker) = opt_waker.take() {
                waker.wake();
            }
        }
    }
}
<commit_msg>FutureExt::Shared::poll optimizations<commit_after>//! Definition of the Shared combinator, a future that is cloneable,
//! and can be polled in multiple threads.
//!
//! # Examples
//!
//! ```
//! # extern crate futures;
//! # extern crate futures_executor;
//! use futures::prelude::*;
//! use futures::future;
//! use futures_executor::block_on;
//!
//! # fn main() {
//! let future = future::ready(6);
//! let shared1 = future.shared();
//! let shared2 = shared1.clone();
//! assert_eq!(6, *block_on(shared1));
//! assert_eq!(6, *block_on(shared2));
//! # }
//! ```

use slab::Slab;

use std::fmt;
use std::cell::UnsafeCell;
use std::marker::Unpin;
use std::mem::PinMut;
use std::sync::{Arc, Mutex};
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::SeqCst;
use std::task::local_waker_from_nonlocal;

use futures_core::{Future, Poll};
use futures_core::task::{self, Wake, Waker};

/// A future that is cloneable and can be polled in multiple threads.
/// Use `Future::shared()` method to convert any future into a `Shared` future.
#[must_use = "futures do nothing unless polled"]
pub struct Shared<F: Future> {
    inner: Arc<Inner<F>>,
    waker_key: usize,
}

struct Inner<F: Future> {
    future_or_output: UnsafeCell<FutureOrOutput<F>>,
    notifier: Arc<Notifier>,
}

struct Notifier {
    state: AtomicUsize,
    wakers: Mutex<Option<Slab<Option<Waker>>>>,
}

// The future itself is polled behind the `Arc`, so it won't be moved
// when `Shared` is moved.
impl<F: Future> Unpin for Shared<F> {}

impl<F: Future> fmt::Debug for Shared<F> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Shared")
            .field("inner", &self.inner)
            .field("waker_key", &self.waker_key)
            .finish()
    }
}

impl<F: Future> fmt::Debug for Inner<F> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Inner")
            .finish()
    }
}

enum FutureOrOutput<F: Future> {
    Future(F),
    Output(Arc<F::Output>),
}

unsafe impl<F> Send for Inner<F>
where
    F: Future + Send,
    F::Output: Send + Sync,
{}

unsafe impl<F> Sync for Inner<F>
where
    F: Future + Send,
    F::Output: Send + Sync,
{}

const IDLE: usize = 0;
const POLLING: usize = 1;
const REPOLL: usize = 2;
const COMPLETE: usize = 3;
const POISONED: usize = 4;

const NULL_WAKER_KEY: usize = usize::max_value();

pub fn new<F: Future>(future: F) -> Shared<F> {
    Shared {
        inner: Arc::new(Inner {
            future_or_output: UnsafeCell::new(FutureOrOutput::Future(future)),
            notifier: Arc::new(Notifier {
                state: AtomicUsize::new(IDLE),
                wakers: Mutex::new(Some(Slab::new())),
            }),
        }),
        waker_key: NULL_WAKER_KEY,
    }
}

impl<F> Shared<F> where F: Future {
    /// If any clone of this `Shared` has completed execution, returns its result immediately
    /// without blocking. Otherwise, returns None without triggering the work represented by
    /// this `Shared`.
    pub fn peek(&self) -> Option<Arc<F::Output>> {
        match self.inner.notifier.state.load(SeqCst) {
            COMPLETE => {
                Some(unsafe { self.clone_output() })
            }
            POISONED => panic!("inner future panicked during poll"),
            _ => None,
        }
    }

    /// Registers the current task to receive a wakeup when `Inner` is awoken.
    fn set_waker(&mut self, cx: &mut task::Context) {
        // Acquire the lock first before checking COMPLETE to ensure there
        // isn't a race.
        let mut wakers = self.inner.notifier.wakers.lock().unwrap();
        let wakers = if let Some(wakers) = wakers.as_mut() {
            wakers
        } else {
            // The value is already available, so there's no need to set the waker.
            return
        };
        if self.waker_key == NULL_WAKER_KEY {
            self.waker_key = wakers.insert(Some(cx.waker().clone()));
        } else {
            let waker_slot = &mut wakers[self.waker_key];
            let needs_replacement = if let Some(old_waker) = waker_slot {
                // If there's still an unwoken waker in the slot, only replace
                // if the current one wouldn't wake the same task.
                !old_waker.will_wake(cx.waker())
            } else {
                true
            };
            if needs_replacement {
                *waker_slot = Some(cx.waker().clone());
            }
        }
        debug_assert!(self.waker_key != NULL_WAKER_KEY);
    }

    /// Safety: callers must first ensure that `self.inner.state`
    /// is `COMPLETE`
    unsafe fn clone_output(&self) -> Arc<F::Output> {
        if let FutureOrOutput::Output(item) = &*self.inner.future_or_output.get() {
            item.clone()
        } else {
            unreachable!()
        }
    }
}

impl<F: Future> Future for Shared<F> {
    type Output = Arc<F::Output>;

    fn poll(mut self: PinMut<Self>, cx: &mut task::Context) -> Poll<Self::Output> {
        let this = &mut *self;

        this.set_waker(cx);

        match this.inner.notifier.state.compare_and_swap(IDLE, POLLING, SeqCst) {
            IDLE => {
                // Lock acquired, fall through
            }
            POLLING | REPOLL => {
                // Another task is currently polling, at this point we just want
                // to ensure that our task handle is currently registered

                return Poll::Pending
            }
            COMPLETE => {
                return unsafe { Poll::Ready(this.clone_output()) };
            }
            POISONED => panic!("inner future panicked during poll"),
            _ => unreachable!(),
        }

        let waker = local_waker_from_nonlocal(this.inner.notifier.clone());
        let mut cx = cx.with_waker(&waker);

        loop {
            struct Reset<'a>(&'a AtomicUsize);

            impl<'a> Drop for Reset<'a> {
                fn drop(&mut self) {
                    use std::thread;

                    if thread::panicking() {
                        self.0.store(POISONED, SeqCst);
                    }
                }
            }

            let _reset = Reset(&this.inner.notifier.state);

            // Poll the future
            let res = unsafe {
                if let FutureOrOutput::Future(future) = &mut *this.inner.future_or_output.get() {
                    PinMut::new_unchecked(future).poll(&mut cx)
                } else {
                    unreachable!()
                }
            };
            match res {
                Poll::Pending => {
                    // Not ready, try to release the handle
                    match this.inner.notifier.state.compare_and_swap(POLLING, IDLE, SeqCst) {
                        POLLING => {
                            // Success
                            return Poll::Pending;
                        }
                        REPOLL => {
                            // Gotta poll again!
                            let prev = this.inner.notifier.state.swap(POLLING, SeqCst);
                            assert_eq!(prev, REPOLL);
                        }
                        _ => unreachable!(),
                    }
                }
                Poll::Ready(output) => {
                    let output = Arc::new(output);
                    unsafe {
                        *this.inner.future_or_output.get() =
                            FutureOrOutput::Output(output.clone());
                    }

                    // Complete the future
                    let mut lock = this.inner.notifier.wakers.lock().unwrap();
                    this.inner.notifier.state.store(COMPLETE, SeqCst);
                    let wakers = &mut lock.take().unwrap();
                    for (_key, opt_waker) in wakers {
                        if let Some(waker) = opt_waker.take() {
                            waker.wake();
                        }
                    }
                    return Poll::Ready(output);
                }
            }
        }
    }
}

impl<F> Clone for Shared<F> where F: Future {
    fn clone(&self) -> Self {
        Shared {
            inner: self.inner.clone(),
            waker_key: NULL_WAKER_KEY,
        }
    }
}

impl<F> Drop for Shared<F> where F: Future {
    fn drop(&mut self) {
        if self.waker_key != NULL_WAKER_KEY {
            if let Ok(mut wakers) = self.inner.notifier.wakers.lock() {
                if let Some(wakers) = wakers.as_mut() {
                    wakers.remove(self.waker_key);
                }
            }
        }
    }
}

impl Wake for Notifier {
    fn wake(arc_self: &Arc<Self>) {
        arc_self.state.compare_and_swap(POLLING, REPOLL, SeqCst);

        let wakers = &mut *arc_self.wakers.lock().unwrap();
        if let Some(wakers) = wakers.as_mut() {
            for (_key, opt_waker) in wakers {
                if let Some(waker) = opt_waker.take() {
                    waker.wake();
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use futures_core::future::Future;
use futures_core::task::{LocalWaker, Poll, Wake, Waker};
use slab::Slab;
use std::fmt;
use std::cell::UnsafeCell;
use std::marker::Unpin;
use std::mem;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::SeqCst;
use std::task::local_waker_from_nonlocal;

/// A future that is cloneable and can be polled in multiple threads.
/// Use `Future::shared()` method to convert any future into a `Shared` future.
#[must_use = "futures do nothing unless polled"]
pub struct Shared<Fut: Future> {
    inner: Arc<Inner<Fut>>,
    waker_key: usize,
}

struct Inner<Fut: Future> {
    future_or_output: UnsafeCell<FutureOrOutput<Fut>>,
    notifier: Arc<Notifier>,
}

struct Notifier {
    state: AtomicUsize,
    wakers: Mutex<Option<Slab<Option<Waker>>>>,
}

// The future itself is polled behind the `Arc`, so it won't be moved
// when `Shared` is moved.
impl<Fut: Future> Unpin for Shared<Fut> {}

impl<Fut: Future> fmt::Debug for Shared<Fut> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Shared")
            .field("inner", &self.inner)
            .field("waker_key", &self.waker_key)
            .finish()
    }
}

impl<Fut: Future> fmt::Debug for Inner<Fut> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Inner")
            .finish()
    }
}

enum FutureOrOutput<Fut: Future> {
    Future(Fut),
    Output(Fut::Output),
    Done,
}

unsafe impl<Fut> Send for Inner<Fut>
where
    Fut: Future + Send,
    Fut::Output: Send + Sync,
{}

unsafe impl<Fut> Sync for Inner<Fut>
where
    Fut: Future + Send,
    Fut::Output: Send + Sync,
{}

const IDLE: usize = 0;
const POLLING: usize = 1;
const REPOLL: usize = 2;
const COMPLETE: usize = 3;
const POISONED: usize = 4;

const NULL_WAKER_KEY: usize = usize::max_value();

impl<Fut: Future> Shared<Fut> {
    pub(super) fn new(future: Fut) -> Shared<Fut> {
        Shared {
            inner: Arc::new(Inner {
                future_or_output: UnsafeCell::new(FutureOrOutput::Future(future)),
                notifier: Arc::new(Notifier {
                    state: AtomicUsize::new(IDLE),
                    wakers: Mutex::new(Some(Slab::new())),
                }),
            }),
            waker_key: NULL_WAKER_KEY,
        }
    }
}

impl<Fut> Shared<Fut>
where
    Fut: Future,
    Fut::Output: Clone,
{
    /// If any clone of this `Shared` has completed execution, returns its result immediately
    /// without blocking. Otherwise, returns None without triggering the work represented by
    /// this `Shared`.
    pub fn peek(&self) -> Option<Fut::Output> {
        match self.inner.notifier.state.load(SeqCst) {
            COMPLETE => Some(unsafe { self.clone_output() }),
            POISONED => panic!("inner future panicked during poll"),
            _ => None,
        }
    }

    /// Registers the current task to receive a wakeup when `Inner` is awoken.
    fn set_waker(&mut self, lw: &LocalWaker) {
        // Acquire the lock first before checking COMPLETE to ensure there
        // isn't a race.
        let mut wakers = self.inner.notifier.wakers.lock().unwrap();
        let wakers = if let Some(wakers) = wakers.as_mut() {
            wakers
        } else {
            // The value is already available, so there's no need to set the waker.
            return
        };
        if self.waker_key == NULL_WAKER_KEY {
            self.waker_key = wakers.insert(Some(lw.clone().into_waker()));
        } else {
            let waker_slot = &mut wakers[self.waker_key];
            let needs_replacement = if let Some(old_waker) = waker_slot {
                // If there's still an unwoken waker in the slot, only replace
                // if the current one wouldn't wake the same task.
                !lw.will_wake_nonlocal(old_waker)
            } else {
                true
            };
            if needs_replacement {
                *waker_slot = Some(lw.clone().into_waker());
            }
        }
        debug_assert!(self.waker_key != NULL_WAKER_KEY);
    }

    /// Safety: callers must first ensure that `self.inner.state`
    /// is `COMPLETE`
    unsafe fn take_or_clone_output(&mut self) -> Fut::Output {
        if let Some(inner) = Arc::get_mut(&mut self.inner) {
            let output = mem::replace(
                &mut inner.future_or_output,
                UnsafeCell::new(FutureOrOutput::Done)
            );
            if let FutureOrOutput::Output(item) = output.into_inner() {
                return item;
            }
        }
        self.clone_output()
    }

    /// Safety: callers must first ensure that `self.inner.state`
    /// is `COMPLETE`
    unsafe fn clone_output(&self) -> Fut::Output {
        match &*self.inner.future_or_output.get() {
            FutureOrOutput::Output(item) => item.clone(),
            FutureOrOutput::Done => panic!("Shared future polled again after completion"),
            FutureOrOutput::Future(_) => unreachable!(),
        }
    }
}

impl<Fut: Future> Future for Shared<Fut>
where
    Fut::Output: Clone,
{
    type Output = Fut::Output;

    fn poll(mut self: Pin<&mut Self>, lw: &LocalWaker) -> Poll<Self::Output> {
        let this = &mut *self;

        this.set_waker(lw);

        match this.inner.notifier.state.compare_and_swap(IDLE, POLLING, SeqCst) {
            IDLE => {
                // Lock acquired, fall through
            }
            POLLING | REPOLL => {
                // Another task is currently polling, at this point we just want
                // to ensure that our task handle is currently registered

                return Poll::Pending;
            }
            COMPLETE => {
                return unsafe { Poll::Ready(this.take_or_clone_output()) };
            }
            POISONED => panic!("inner future panicked during poll"),
            _ => unreachable!(),
        }

        let waker = local_waker_from_nonlocal(this.inner.notifier.clone());
        let lw = &waker;

        struct Reset<'a>(&'a AtomicUsize);

        impl<'a> Drop for Reset<'a> {
            fn drop(&mut self) {
                use std::thread;

                if thread::panicking() {
                    self.0.store(POISONED, SeqCst);
                }
            }
        }

        let output = loop {
            let _reset = Reset(&this.inner.notifier.state);

            // Poll the future
            let res = unsafe {
                if let FutureOrOutput::Future(future) =
                    &mut *this.inner.future_or_output.get()
                {
                    Pin::new_unchecked(future).poll(lw)
                } else {
                    unreachable!()
                }
            };
            match res {
                Poll::Pending => {
                    // Not ready, try to release the handle
                    match this
                        .inner
                        .notifier
                        .state
                        .compare_and_swap(POLLING, IDLE, SeqCst)
                    {
                        POLLING => {
                            // Success
                            return Poll::Pending;
                        }
                        REPOLL => {
                            // Gotta poll again!
                            let prev =
                                this.inner.notifier.state.swap(POLLING, SeqCst);
                            assert_eq!(prev, REPOLL);
                        }
                        _ => unreachable!(),
                    }
                }
                Poll::Ready(output) => break output,
            }
        };

        if Arc::get_mut(&mut this.inner).is_some() {
            unsafe {
                *this.inner.future_or_output.get() = FutureOrOutput::Done;
            }
            return Poll::Ready(output);
        }

        let _reset = Reset(&this.inner.notifier.state);

        unsafe {
            *this.inner.future_or_output.get() =
                FutureOrOutput::Output(output.clone());
        }

        // Complete the future
        let mut lock = this.inner.notifier.wakers.lock().unwrap();
        this.inner.notifier.state.store(COMPLETE, SeqCst);
        let wakers = &mut lock.take().unwrap();
        for (_key, opt_waker) in wakers {
            if let Some(waker) = opt_waker.take() {
                waker.wake();
            }
        }
        Poll::Ready(output)
    }
}

impl<Fut> Clone for Shared<Fut>
where
    Fut: Future,
{
    fn clone(&self) -> Self {
        Shared {
            inner: self.inner.clone(),
            waker_key: NULL_WAKER_KEY,
        }
    }
}

impl<Fut> Drop for Shared<Fut>
where
    Fut: Future,
{
    fn drop(&mut self) {
        if self.waker_key != NULL_WAKER_KEY {
            if let Ok(mut wakers) = self.inner.notifier.wakers.lock() {
                if let Some(wakers) = wakers.as_mut() {
                    wakers.remove(self.waker_key);
                }
            }
        }
    }
}

impl Wake for Notifier {
    fn wake(arc_self: &Arc<Self>) {
        arc_self.state.compare_and_swap(POLLING, REPOLL, SeqCst);

        let wakers = &mut *arc_self.wakers.lock().unwrap();
        if let Some(wakers) = wakers.as_mut() {
            for (_key, opt_waker) in wakers {
                if let Some(waker) = opt_waker.take() {
                    waker.wake();
                }
            }
        }
    }
}
<commit_msg>Avoid panicking in Shared::peek if the future is completed<commit_after>use futures_core::future::Future;
use futures_core::task::{LocalWaker, Poll, Wake, Waker};
use slab::Slab;
use std::fmt;
use std::cell::UnsafeCell;
use std::marker::Unpin;
use std::mem;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::SeqCst;
use std::task::local_waker_from_nonlocal;

/// A future that is cloneable and can be polled in multiple threads.
/// Use `Future::shared()` method to convert any future into a `Shared` future.
#[must_use = "futures do nothing unless polled"]
pub struct Shared<Fut: Future> {
    inner: Arc<Inner<Fut>>,
    waker_key: usize,
}

struct Inner<Fut: Future> {
    future_or_output: UnsafeCell<FutureOrOutput<Fut>>,
    notifier: Arc<Notifier>,
}

struct Notifier {
    state: AtomicUsize,
    wakers: Mutex<Option<Slab<Option<Waker>>>>,
}

// The future itself is polled behind the `Arc`, so it won't be moved
// when `Shared` is moved.
impl<Fut: Future> Unpin for Shared<Fut> {}

impl<Fut: Future> fmt::Debug for Shared<Fut> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Shared")
            .field("inner", &self.inner)
            .field("waker_key", &self.waker_key)
            .finish()
    }
}

impl<Fut: Future> fmt::Debug for Inner<Fut> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Inner")
            .finish()
    }
}

enum FutureOrOutput<Fut: Future> {
    Future(Fut),
    Output(Fut::Output),
    Done,
}

unsafe impl<Fut> Send for Inner<Fut>
where
    Fut: Future + Send,
    Fut::Output: Send + Sync,
{}

unsafe impl<Fut> Sync for Inner<Fut>
where
    Fut: Future + Send,
    Fut::Output: Send + Sync,
{}

const IDLE: usize = 0;
const POLLING: usize = 1;
const REPOLL: usize = 2;
const COMPLETE: usize = 3;
const POISONED: usize = 4;

const NULL_WAKER_KEY: usize = usize::max_value();

impl<Fut: Future> Shared<Fut> {
    pub(super) fn new(future: Fut) -> Shared<Fut> {
        Shared {
            inner: Arc::new(Inner {
                future_or_output: UnsafeCell::new(FutureOrOutput::Future(future)),
                notifier: Arc::new(Notifier {
                    state: AtomicUsize::new(IDLE),
                    wakers: Mutex::new(Some(Slab::new())),
                }),
            }),
            waker_key: NULL_WAKER_KEY,
        }
    }
}

impl<Fut> Shared<Fut>
where
    Fut: Future,
    Fut::Output: Clone,
{
    /// If any clone of this `Shared` has completed execution and there are still at least one
    /// clone of this `Shared` alive then, returns its result immediately
    /// without blocking. Otherwise, returns None without triggering the work represented by
    /// this `Shared`.
    pub fn peek(&self) -> Option<Fut::Output> {
        match self.inner.notifier.state.load(SeqCst) {
            COMPLETE => unsafe { self.clone_output() },
            POISONED => panic!("inner future panicked during poll"),
            _ => None,
        }
    }

    /// Registers the current task to receive a wakeup when `Inner` is awoken.
    fn set_waker(&mut self, lw: &LocalWaker) {
        // Acquire the lock first before checking COMPLETE to ensure there
        // isn't a race.
        let mut wakers = self.inner.notifier.wakers.lock().unwrap();
        let wakers = if let Some(wakers) = wakers.as_mut() {
            wakers
        } else {
            // The value is already available, so there's no need to set the waker.
            return
        };
        if self.waker_key == NULL_WAKER_KEY {
            self.waker_key = wakers.insert(Some(lw.clone().into_waker()));
        } else {
            let waker_slot = &mut wakers[self.waker_key];
            let needs_replacement = if let Some(old_waker) = waker_slot {
                // If there's still an unwoken waker in the slot, only replace
                // if the current one wouldn't wake the same task.
                !lw.will_wake_nonlocal(old_waker)
            } else {
                true
            };
            if needs_replacement {
                *waker_slot = Some(lw.clone().into_waker());
            }
        }
        debug_assert!(self.waker_key != NULL_WAKER_KEY);
    }

    /// Safety: callers must first ensure that `self.inner.state`
    /// is `COMPLETE`
    unsafe fn take_or_clone_output(&mut self) -> Fut::Output {
        if let Some(inner) = Arc::get_mut(&mut self.inner) {
            let output = mem::replace(
                &mut inner.future_or_output,
                UnsafeCell::new(FutureOrOutput::Done)
            );
            if let FutureOrOutput::Output(item) = output.into_inner() {
                return item;
            }
        }
        self.clone_output()
            .unwrap_or_else(|| panic!("Shared future polled again after completion"))
    }

    /// Safety: callers must first ensure that `self.inner.state`
    /// is `COMPLETE`
    unsafe fn clone_output(&self) -> Option<Fut::Output> {
        match &*self.inner.future_or_output.get() {
            FutureOrOutput::Output(item) => Some(item.clone()),
            FutureOrOutput::Done => None,
            FutureOrOutput::Future(_) => unreachable!(),
        }
    }
}

impl<Fut: Future> Future for Shared<Fut>
where
    Fut::Output: Clone,
{
    type Output = Fut::Output;

    fn poll(mut self: Pin<&mut Self>, lw: &LocalWaker) -> Poll<Self::Output> {
        let this = &mut *self;

        this.set_waker(lw);

        match this.inner.notifier.state.compare_and_swap(IDLE, POLLING, SeqCst) {
            IDLE => {
                // Lock acquired, fall through
            }
            POLLING | REPOLL => {
                // Another task is currently polling, at this point we just want
                // to ensure that our task handle is currently registered

                return Poll::Pending;
            }
            COMPLETE => {
                return unsafe { Poll::Ready(this.take_or_clone_output()) };
            }
            POISONED => panic!("inner future panicked during poll"),
            _ => unreachable!(),
        }

        let waker = local_waker_from_nonlocal(this.inner.notifier.clone());
        let lw = &waker;

        struct Reset<'a>(&'a AtomicUsize);

        impl<'a> Drop for Reset<'a> {
            fn drop(&mut self) {
                use std::thread;

                if thread::panicking() {
                    self.0.store(POISONED, SeqCst);
                }
            }
        }

        let output = loop {
            let _reset = Reset(&this.inner.notifier.state);

            // Poll the future
            let res = unsafe {
                if let FutureOrOutput::Future(future) =
                    &mut *this.inner.future_or_output.get()
                {
                    Pin::new_unchecked(future).poll(lw)
                } else {
                    unreachable!()
                }
            };
            match res {
                Poll::Pending => {
                    // Not ready, try to release the handle
                    match this
                        .inner
                        .notifier
                        .state
                        .compare_and_swap(POLLING, IDLE, SeqCst)
                    {
                        POLLING => {
                            // Success
                            return Poll::Pending;
                        }
                        REPOLL => {
                            // Gotta poll again!
                            let prev =
                                this.inner.notifier.state.swap(POLLING, SeqCst);
                            assert_eq!(prev, REPOLL);
                        }
                        _ => unreachable!(),
                    }
                }
                Poll::Ready(output) => break output,
            }
        };

        if Arc::get_mut(&mut this.inner).is_some() {
            unsafe {
                *this.inner.future_or_output.get() = FutureOrOutput::Done;
            }
            return Poll::Ready(output);
        }

        let _reset = Reset(&this.inner.notifier.state);

        unsafe {
            *this.inner.future_or_output.get() =
                FutureOrOutput::Output(output.clone());
        }

        // Complete the future
        let mut lock = this.inner.notifier.wakers.lock().unwrap();
        this.inner.notifier.state.store(COMPLETE, SeqCst);
        let wakers = &mut lock.take().unwrap();
        for (_key, opt_waker) in wakers {
            if let Some(waker) = opt_waker.take() {
                waker.wake();
            }
        }
        Poll::Ready(output)
    }
}

impl<Fut> Clone for Shared<Fut>
where
    Fut: Future,
{
    fn clone(&self) -> Self {
        Shared {
            inner: self.inner.clone(),
            waker_key: NULL_WAKER_KEY,
        }
    }
}

impl<Fut> Drop for Shared<Fut>
where
    Fut: Future,
{
    fn drop(&mut self) {
        if self.waker_key != NULL_WAKER_KEY {
            if let Ok(mut wakers) = self.inner.notifier.wakers.lock() {
                if let Some(wakers) = wakers.as_mut() {
                    wakers.remove(self.waker_key);
                }
            }
        }
    }
}

impl Wake for Notifier {
    fn wake(arc_self: &Arc<Self>) {
        arc_self.state.compare_and_swap(POLLING, REPOLL, SeqCst);

        let wakers = &mut *arc_self.wakers.lock().unwrap();
        if let Some(wakers) = wakers.as_mut() {
            for (_key, opt_waker) in wakers {
                if let Some(waker) = opt_waker.take() {
                    waker.wake();
                }
            }
        }
    }
}
<|endoftext|>