tokio 0.2.10

use pin_project_lite::pin_project;
use std::cell::RefCell;
use std::error::Error;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::{fmt, thread};

/// Declare a new task local storage key of type [`tokio::task::LocalKey`].
///
/// # Syntax
///
/// The macro wraps any number of static declarations and makes them task locals.
/// Publicity and attributes for each static are allowed.
///
/// # Examples
///
/// ```
/// # use tokio::task_local;
/// task_local! {
///     pub static FOO: u32;
///
///     #[allow(unused)]
///     static BAR: f32;
/// }
/// # fn main() {}
/// ```
///
/// See [LocalKey documentation][`tokio::task::LocalKey`] for more
/// information.
///
/// [`tokio::task::LocalKey`]: ../tokio/task/struct.LocalKey.html
#[macro_export]
macro_rules! task_local {
     // empty (base case for the recursion)
    () => {};

    ($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty; $($rest:tt)*) => {
        $crate::__task_local_inner!($(#[$attr])* $vis $name, $t);
        $crate::task_local!($($rest)*);
    };

    ($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty) => {
        $crate::__task_local_inner!($(#[$attr])* $vis $name, $t);
    }
}

#[doc(hidden)]
#[macro_export]
macro_rules! __task_local_inner {
    ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty) => {
        static $name: $crate::task::LocalKey<$t> = {
            std::thread_local! {
                static __KEY: std::cell::RefCell<Option<$t>> = std::cell::RefCell::new(None);
            }

            $crate::task::LocalKey { inner: __KEY }
        };
    };
}

/// A key for task-local data.
///
/// This type is generated by `task_local!` macro and unlike `thread_local!` it has
/// no concept of lazily initialization. Instead, it is designed to provide task local
/// storage the future that is passed to `set`.
///
/// # Initialization and Destruction
///
/// Initialization is done via `set` which is an `async fn` that wraps another
/// [`std::future::Future`] and will set the value on each `Future::poll` call.
/// Once the `set` future is dropped the corresponding task local value is also
/// dropped.
///
/// # Examples
///
/// ```
/// # async fn dox() {
/// tokio::task_local! {
///     static FOO: u32;
/// }
///
/// FOO.scope(1, async move {
///     assert_eq!(FOO.get(), 1);
/// }).await;
///
/// FOO.scope(2, async move {
///     assert_eq!(FOO.get(), 2);
///
///     FOO.scope(3, async move {
///         assert_eq!(FOO.get(), 3);
///     }).await;
/// }).await;
/// # }
/// ```
pub struct LocalKey<T: 'static> {
    #[doc(hidden)]
    pub inner: thread::LocalKey<RefCell<Option<T>>>,
}

impl<T: 'static> LocalKey<T> {
    /// Sets a value `T` as the task local value for the future `F`.
    ///
    /// This will run the provided future to completion and set the
    /// provided value as the task local under this key. Once the returned
    /// future is dropped so will the value passed be dropped.
    ///
    /// # async fn dox() {
    /// tokio::task_local! {
    ///     static FOO: u32;
    /// }
    ///
    /// FOO.scope(1, async move {
    ///     println!("task local value: {}", FOO.get());
    /// }).await;
    /// # }
    pub async fn scope<F>(&'static self, value: T, f: F) -> F::Output
    where
        F: Future,
    {
        TaskLocalFuture {
            local: &self,
            slot: Some(value),
            future: f,
        }
        .await
    }

    /// Access this task-local key, running the provided closure with a reference
    /// passed to the value.
    ///
    /// # Panics
    ///
    /// This function will panic if not called within a future that has not been
    /// set via `LocalKey::set`.
    pub fn with<F, R>(&'static self, f: F) -> R
    where
        F: FnOnce(&T) -> R,
    {
        self.try_with(f).expect(
            "cannot access a Task Local Storage value \
             without setting it via `LocalKey::set`",
        )
    }

    /// Access this task-local key, running the provided closure with a reference
    /// passed to the value. Unlike `with` this function will return a `Result<R, AccessError>`
    /// instead of panicking.
    pub fn try_with<F, R>(&'static self, f: F) -> Result<R, AccessError>
    where
        F: FnOnce(&T) -> R,
    {
        self.inner.with(|v| {
            if let Some(val) = v.borrow().as_ref() {
                Ok(f(val))
            } else {
                Err(AccessError { _private: () })
            }
        })
    }
}

impl<T: Copy + 'static> LocalKey<T> {
    /// Get a copy of the task-local value if it implements
    /// the `Copy` trait.
    pub fn get(&'static self) -> T {
        self.with(|v| *v)
    }
}

impl<T: 'static> fmt::Debug for LocalKey<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.pad("LocalKey { .. }")
    }
}

pin_project! {
    struct TaskLocalFuture<T: StaticLifetime, F> {
        local: &'static LocalKey<T>,
        slot: Option<T>,
        #[pin]
        future: F,
    }
}

impl<T: 'static, F: Future> Future for TaskLocalFuture<T, F> {
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        struct Guard<'a, T: 'static> {
            local: &'static LocalKey<T>,
            slot: &'a mut Option<T>,
            prev: Option<T>,
        }

        impl<T> Drop for Guard<'_, T> {
            fn drop(&mut self) {
                let value = self.local.inner.with(|c| c.replace(self.prev.take()));
                *self.slot = value;
            }
        }

        let mut project = self.project();
        let val = project.slot.take();

        let prev = project.local.inner.with(|c| c.replace(val));

        let _guard = Guard {
            prev,
            slot: &mut project.slot,
            local: *project.local,
        };

        project.future.poll(cx)
    }
}

// Required to make `pin_project` happy.
trait StaticLifetime: 'static {}
impl<T: 'static> StaticLifetime for T {}

/// An error returned by [`LocalKey::try_with`](struct.LocalKey.html#method.try_with).
#[derive(Clone, Copy, Eq, PartialEq)]
pub struct AccessError {
    _private: (),
}

impl fmt::Debug for AccessError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("AccessError").finish()
    }
}

impl fmt::Display for AccessError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt("task-local value not set", f)
    }
}

impl Error for AccessError {}