3 files changed, 152 insertions, 16 deletions
diff --git a/src/libstd/future.rs b/src/libstd/future.rs
new file mode 100644
index 00000000000..2da775fdc94
--- /dev/null
+++ b/src/libstd/future.rs
@@ -0,0 +1,116 @@
+// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+//! Asynchronous values.
+
+use core::cell::Cell;
+use core::marker::Unpin;
+use core::mem::PinMut;
+use core::option::Option;
+use core::ptr::NonNull;
+use core::task::{self, Poll};
+use core::ops::{Drop, Generator, GeneratorState};
+
+#[doc(inline)]
+pub use core::future::*;
+
+/// Wrap a future in a generator.
+///
+/// This function returns a `GenFuture` underneath, but hides it in `impl Trait` to give
+/// better error messages (`impl Future` rather than `GenFuture<[closure.....]>`).
+#[unstable(feature = "gen_future", issue = "50547")]
+pub fn from_generator<T: Generator<Yield = ()>>(x: T) -> impl Future<Output = T::Return> {
+    GenFuture(x)
+}
+
+/// A wrapper around generators used to implement `Future` for `async`/`await` code.
+#[unstable(feature = "gen_future", issue = "50547")]
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
+struct GenFuture<T: Generator<Yield = ()>>(T);
+
+// We rely on the fact that async/await futures are immovable in order to create
+// self-referential borrows in the underlying generator.
+impl<T: Generator<Yield = ()>> !Unpin for GenFuture<T> {}
+
+#[unstable(feature = "gen_future", issue = "50547")]
+impl<T: Generator<Yield = ()>> Future for GenFuture<T> {
+    type Output = T::Return;
+    fn poll(self: PinMut<Self>, cx: &mut task::Context) -> Poll<Self::Output> {
+        set_task_cx(cx, || match unsafe { PinMut::get_mut(self).0.resume() } {
+            GeneratorState::Yielded(()) => Poll::Pending,
+            GeneratorState::Complete(x) => Poll::Ready(x),
+        })
+    }
+}
+
+thread_local! {
+    static TLS_CX: Cell<Option<NonNull<task::Context<'static>>>> = Cell::new(None);
+}
+
+struct SetOnDrop(Option<NonNull<task::Context<'static>>>);
+
+impl Drop for SetOnDrop {
+    fn drop(&mut self) {
+        TLS_CX.with(|tls_cx| {
+            tls_cx.set(self.0.take());
+        });
+    }
+}
+
+#[unstable(feature = "gen_future", issue = "50547")]
+/// Sets the thread-local task context used by async/await futures.
+pub fn set_task_cx<F, R>(cx: &mut task::Context, f: F) -> R
+where
+    F: FnOnce() -> R
+{
+    let old_cx = TLS_CX.with(|tls_cx| {
+        tls_cx.replace(NonNull::new(
+            cx
+                as *mut task::Context
+                as *mut ()
+                as *mut task::Context<'static>
+        ))
+    });
+    let _reset_cx = SetOnDrop(old_cx);
+    f()
+}
+
+#[unstable(feature = "gen_future", issue = "50547")]
+/// Retrieves the thread-local task context used by async/await futures.
+///
+/// This function acquires exclusive access to the task context.
+///
+/// Panics if no task has been set or if the task context has already been
+/// retrived by a surrounding call to get_task_cx.
+pub fn get_task_cx<F, R>(f: F) -> R
+where
+    F: FnOnce(&mut task::Context) -> R
+{
+    let cx_ptr = TLS_CX.with(|tls_cx| {
+        // Clear the entry so that nested `with_get_cx` calls
+        // will fail or set their own value.
+        tls_cx.replace(None)
+    });
+    let _reset_cx = SetOnDrop(cx_ptr);
+
+    let mut cx_ptr = cx_ptr.expect(
+        "TLS task::Context not set. This is a rustc bug. \
+        Please file an issue on https://github.com/rust-lang/rust.");
+    unsafe { f(cx_ptr.as_mut()) }
+}
+
+#[unstable(feature = "gen_future", issue = "50547")]
+/// Polls a future in the current thread-local task context.
+pub fn poll_in_task_cx<F>(f: &mut PinMut<F>) -> Poll<F::Output>
+where
+    F: Future
+{
+    get_task_cx(|cx| f.reborrow().poll(cx))
+}
diff --git a/src/libstd/lib.rs b/src/libstd/lib.rs
index a6061e96ae5..caad924ea5b 100644
--- a/src/libstd/lib.rs
+++ b/src/libstd/lib.rs
@@ -263,6 +263,7 @@
 #![feature(fn_traits)]
 #![feature(fnbox)]
 #![feature(futures_api)]
+#![feature(generator_trait)]
 #![feature(hashmap_internals)]
 #![feature(int_error_internals)]
 #![feature(integer_atomics)]
@@ -462,22 +463,6 @@ pub use core::u128;
 #[stable(feature = "core_hint", since = "1.27.0")]
 pub use core::hint;
 
-#[unstable(feature = "futures_api",
-           reason = "futures in libcore are unstable",
-           issue = "50547")]
-pub mod task {
-    //! Types and Traits for working with asynchronous tasks.
-    #[doc(inline)]
-    pub use core::task::*;
-    #[doc(inline)]
-    pub use alloc_crate::task::*;
-}
-
-#[unstable(feature = "futures_api",
-           reason = "futures in libcore are unstable",
-           issue = "50547")]
-pub use core::future;
-
 pub mod f32;
 pub mod f64;
 
@@ -499,6 +484,22 @@ pub mod process;
 pub mod sync;
 pub mod time;
 
+#[unstable(feature = "futures_api",
+           reason = "futures in libcore are unstable",
+           issue = "50547")]
+pub mod task {
+    //! Types and Traits for working with asynchronous tasks.
+    #[doc(inline)]
+    pub use core::task::*;
+    #[doc(inline)]
+    pub use alloc_crate::task::*;
+}
+
+#[unstable(feature = "futures_api",
+           reason = "futures in libcore are unstable",
+           issue = "50547")]
+pub mod future;
+
 // Platform-abstraction modules
 #[macro_use]
 mod sys_common;
diff --git a/src/libstd/macros.rs b/src/libstd/macros.rs
index 8da70f5717e..a856e7736fb 100644
--- a/src/libstd/macros.rs
+++ b/src/libstd/macros.rs
@@ -213,6 +213,25 @@ macro_rules! eprintln {
     ($fmt:expr, $($arg:tt)*) => (eprint!(concat!($fmt, "\n"), $($arg)*));
 }
 
+#[macro_export]
+#[unstable(feature = "await_macro", issue = "50547")]
+#[allow_internal_unstable]
+#[allow_internal_unsafe]
+macro_rules! await {
+    ($e:expr) => { {
+        let mut pinned = $e;
+        let mut pinned = unsafe { $crate::mem::PinMut::new_unchecked(&mut pinned) };
+        loop {
+            match $crate::future::poll_in_task_cx(&mut pinned) {
+                // FIXME(cramertj) prior to stabilizing await, we have to ensure that this
+                // can't be used to create a generator on stable via `|| await!()`.
+                $crate::task::Poll::Pending => yield,
+                $crate::task::Poll::Ready(x) => break x,
+            }
+        }
+    } }
+}
+
 /// A macro to select an event from a number of receivers.
 ///
 /// This macro is used to wait for the first event to occur on a number of