macros: add pin! macro (#2163)

Used for stack pinning and based on `pin_mut!` from the pin-util crate.

Pinning is used often when working with stream operators and the select!
macro. Given the small size of `pin!` it makes more sense to include a
version than re-export one from a separate crate or require the user to
depend on `pin-util` themselves.

4 files changed, 134 insertions, 7 deletions

diff --git a/tests-integration/tests/macros_pin.rs b/tests-integration/tests/macros_pin.rs
new file mode 100644
index 00000000..37d8f70f
--- /dev/null
+++ b/tests-integration/tests/macros_pin.rs
@@ -0,0 +1,12 @@
+use futures::executor::block_on;
+
+async fn my_async_fn() {}
+
+#[test]
+fn pin() {
+    block_on(async {
+        let future = my_async_fn();
+        tokio::pin!(future);
+        (&mut future).await
+    });
+}
diff --git a/tokio/src/macros/mod.rs b/tokio/src/macros/mod.rs
index d35a6411..b5b53ccb 100644
--- a/tokio/src/macros/mod.rs
+++ b/tokio/src/macros/mod.rs
@@ -14,6 +14,9 @@ mod join;
 mod loom;
 
 #[macro_use]
+mod pin;
+
+#[macro_use]
 mod ready;
 
 cfg_macros! {
@@ -24,8 +27,6 @@ cfg_macros! {
 #[macro_use]
 mod thread_local;
 
-cfg_macros! {
-    // Includes re-exports needed to implement macros
-    #[doc(hidden)]
-    pub mod support;
-}
+// Includes re-exports needed to implement macros
+#[doc(hidden)]
+pub mod support;
diff --git a/tokio/src/macros/pin.rs b/tokio/src/macros/pin.rs
new file mode 100644
index 00000000..e8511d3f
--- /dev/null
+++ b/tokio/src/macros/pin.rs
@@ -0,0 +1,112 @@
+/// Pins a value on the stack.
+///
+/// Calls to `async fn` return anonymous [`Future`] values that are `!Unpin`.
+/// These values must be pinned before they can be polled. Calling `.await` will
+/// handle this, but consumes the future. If it is required to call `.await` on
+/// a `&mut _` reference, the caller is responsible for pinning the future.
+///
+/// Pinning may be done by allocating with [`Box::pin`] or by using the stack
+/// with the `pin!` macro.
+///
+/// The following will **fail to compile**:
+///
+/// ```compile_fail
+/// async fn my_async_fn() {
+///     // async logic here
+/// }
+///
+/// #[tokio::main]
+/// async fn main() {
+///     let mut future = my_async_fn();
+///     (&mut future).await;
+/// }
+/// ```
+///
+/// To make this work requires pinning:
+///
+/// ```
+/// use tokio::pin;
+///
+/// async fn my_async_fn() {
+///     // async logic here
+/// }
+///
+/// #[tokio::main]
+/// async fn main() {
+///     let future = my_async_fn();
+///     pin!(future);
+///
+///     (&mut future).await;
+/// }
+/// ```
+///
+/// Pinning is useful when using `select!` and stream operators that require `T:
+/// Stream + Unpin`.
+///
+/// [`Future`]: https://doc.rust-lang.org/std/future/trait.Future.html
+/// [`Box::pin`]: #
+///
+/// # Usage
+///
+/// The `pin!` macro takes **identifiers** as arguments. It does **not** work
+/// with expressions.
+///
+/// The following does not compile as an expression is passed to `pin!`.
+///
+/// ```compile_fail
+/// async fn my_async_fn() {
+///     // async logic here
+/// }
+///
+/// #[tokio::main]
+/// async fn main() {
+///     let mut future = pin!(my_async_fn());
+///     (&mut future).await;
+/// }
+/// ```
+///
+/// # Examples
+///
+/// Using with select:
+///
+/// ```
+/// use tokio::{pin, select};
+/// use tokio::stream::{self, StreamExt};
+///
+/// async fn my_async_fn() {
+///     // async logic here
+/// }
+///
+/// #[tokio::main]
+/// async fn main() {
+///     let mut stream = stream::iter(vec![1, 2, 3, 4]);
+///
+///     let future = my_async_fn();
+///     pin!(future);
+///
+///     loop {
+///         select! {
+///             _ = &mut future => {
+///                 // Stop looping `future` will be polled after completion
+///                 break;
+///             }
+///             Some(val) = stream.next() => {
+///                 println!("got value = {}", val);
+///             }
+///         }
+///     }
+/// }
+/// ```
+#[macro_export]
+macro_rules! pin {
+    ($($x:ident),*) => { $(
+        // Move the value to ensure that it is owned
+        let mut $x = $x;
+        // Shadow the original binding so that it can't be directly accessed
+        // ever again.
+        #[allow(unused_mut)]
+        let mut $x = unsafe {
+            $crate::macros::support::Pin::new_unchecked(&mut $x)
+        };
+    )* }
+}
diff --git a/tokio/src/macros/support.rs b/tokio/src/macros/support.rs
index 1f6fcd3d..fc1cdfcf 100644
--- a/tokio/src/macros/support.rs
+++ b/tokio/src/macros/support.rs
@@ -1,5 +1,7 @@
-pub use crate::future::{maybe_done, poll_fn};
-pub use crate::util::thread_rng_n;
+cfg_macros! {
+    pub use crate::future::{maybe_done, poll_fn};
+    pub use crate::util::thread_rng_n;
+}
 
 pub use std::future::Future;
 pub use std::pin::Pin;