[coroutines] Add std::experimental::task<T> type

Summary:
Adds the coroutine `std::experimental::task<T>` type described in proposal P1056R0.
See https://wg21.link/P1056R0.

This implementation allows customization of the allocator used to allocate the
coroutine frame by passing std::allocator_arg as the first argument, followed by
the allocator to use.

This supports co_awaiting the same task multiple times. The second and
subsequent times it returns a reference to the already-computed value.

This diff also adds some implementations of other utilities that have potential for
standardization as helpers within the test/... area:
- `sync_wait(awaitable)` - See P1171R0
- `manual_reset_event`

Move the definition of the __aligned_allocation_size helper function
from <experimental/memory_resource> to <experimental/__memory>
so it can be more widely used without pulling in memory_resource.

Outstanding work:
- Use C++14 keywords directly rather than macro versions
  eg. use `noexcept` instead of `_NOEXCEPT`).
- Add support for overaligned coroutine frames.
  This may need wording in the Coroutines TS to support passing the extra `std::align_val_t`.
- Eliminate use of `if constexpr` if we want it to compile under C++14.

Patch by @lewissbaker (Lewis Baker).

llvm-svn: 357010

3 files changed, 510 insertions, 8 deletions

diff --git a/libcxx/include/experimental/__memory b/libcxx/include/experimental/__memory
index 4cf8978468c..6da6bef663a 100644
--- a/libcxx/include/experimental/__memory
+++ b/libcxx/include/experimental/__memory
@@ -73,6 +73,13 @@ struct __lfts_uses_alloc_ctor
     >
 {};
 
+// Round __s up to next multiple of __a.
+inline _LIBCPP_INLINE_VISIBILITY _LIBCPP_CONSTEXPR
+size_t __aligned_allocation_size(size_t __s, size_t __a) _NOEXCEPT
+{
+    return (__s + __a - 1) & ~(__a - 1);
+}
+
 template <class _Tp, class _Alloc, class ..._Args>
 inline _LIBCPP_INLINE_VISIBILITY
 void __lfts_user_alloc_construct(
diff --git a/libcxx/include/experimental/memory_resource b/libcxx/include/experimental/memory_resource
index f999fb9befd..897bd1f8882 100644
--- a/libcxx/include/experimental/memory_resource
+++ b/libcxx/include/experimental/memory_resource
@@ -86,14 +86,6 @@ _LIBCPP_PUSH_MACROS
 
 _LIBCPP_BEGIN_NAMESPACE_LFTS_PMR
 
-// Round __s up to next multiple of __a.
-inline _LIBCPP_INLINE_VISIBILITY
-size_t __aligned_allocation_size(size_t __s, size_t __a) _NOEXCEPT
-{
-    _LIBCPP_ASSERT(__s + __a > __s, "aligned allocation size overflows");
-    return (__s + __a - 1) & ~(__a - 1);
-}
-
 // 8.5, memory.resource
 class _LIBCPP_TYPE_VIS memory_resource
 {
diff --git a/libcxx/include/experimental/task b/libcxx/include/experimental/task
new file mode 100644
index 00000000000..2bdcaf2ba34
--- /dev/null
+++ b/libcxx/include/experimental/task
@@ -0,0 +1,503 @@
+// -*- C++ -*-
+//===------------------------------- task ---------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _LIBCPP_EXPERIMENTAL_TASK
+#define _LIBCPP_EXPERIMENTAL_TASK
+
+#include <experimental/__config>
+#include <experimental/__memory>
+#include <experimental/coroutine>
+
+#include <exception>
+#include <type_traits>
+#include <utility>
+
+#if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
+#pragma GCC system_header
+#endif
+
+#ifdef _LIBCPP_HAS_NO_COROUTINES
+#if defined(_LIBCPP_WARNING)
+_LIBCPP_WARNING("<experimental/task> cannot be used with this compiler")
+#else
+#warning <experimental/task> cannot be used with this compiler
+#endif
+#endif
+
+_LIBCPP_BEGIN_NAMESPACE_EXPERIMENTAL_COROUTINES
+
+////// task<T>
+
+template <typename _Tp = void>
+class task;
+
+struct __task_promise_final_awaitable {
+  _LIBCPP_INLINE_VISIBILITY
+  _LIBCPP_CONSTEXPR bool await_ready() const _NOEXCEPT { return false; }
+
+  template <typename _TaskPromise>
+  _LIBCPP_INLINE_VISIBILITY coroutine_handle<>
+  await_suspend(coroutine_handle<_TaskPromise> __coro) const _NOEXCEPT {
+    _LIBCPP_ASSERT(
+        __coro.promise().__continuation_,
+        "Coroutine completed without a valid continuation attached.");
+    return __coro.promise().__continuation_;
+  }
+
+  _LIBCPP_INLINE_VISIBILITY
+  void await_resume() const _NOEXCEPT {}
+};
+
+class _LIBCPP_TYPE_VIS __task_promise_base {
+  using _DeallocFunc = void(void* __ptr, size_t __size) _NOEXCEPT;
+
+  template <typename _Alloc>
+  static constexpr bool __allocator_needs_to_be_stored =
+      !allocator_traits<_Alloc>::is_always_equal::value ||
+      !is_default_constructible_v<_Alloc>;
+
+  static _LIBCPP_CONSTEXPR size_t
+  __get_dealloc_func_offset(size_t __frameSize) _NOEXCEPT {
+    return _VSTD_LFTS::__aligned_allocation_size(__frameSize,
+                                                 alignof(_DeallocFunc*));
+  }
+
+  static _LIBCPP_CONSTEXPR size_t
+  __get_padded_frame_size(size_t __frameSize) _NOEXCEPT {
+    return __get_dealloc_func_offset(__frameSize) + sizeof(_DeallocFunc*);
+  }
+
+  template <typename _Alloc>
+  static _LIBCPP_CONSTEXPR size_t
+  __get_allocator_offset(size_t __frameSize) _NOEXCEPT {
+    return _VSTD_LFTS::__aligned_allocation_size(
+        __get_padded_frame_size(__frameSize), alignof(_Alloc));
+  }
+
+  template <typename _Alloc>
+  static _LIBCPP_CONSTEXPR size_t
+  __get_padded_frame_size_with_allocator(size_t __frameSize) _NOEXCEPT {
+    if constexpr (__allocator_needs_to_be_stored<_Alloc>) {
+      return __get_allocator_offset<_Alloc>(__frameSize) + sizeof(_Alloc);
+    } else {
+      return __get_padded_frame_size(__frameSize);
+    }
+  }
+
+  _LIBCPP_INLINE_VISIBILITY
+  static _DeallocFunc*& __get_dealloc_func(void* __frameStart,
+                                           size_t __frameSize) _NOEXCEPT {
+    return *reinterpret_cast<_DeallocFunc**>(
+        static_cast<char*>(__frameStart) +
+        __get_dealloc_func_offset(__frameSize));
+  }
+
+  template <typename _Alloc>
+  _LIBCPP_INLINE_VISIBILITY static _Alloc&
+  __get_allocator(void* __frameStart, size_t __frameSize) _NOEXCEPT {
+    return *reinterpret_cast<_Alloc*>(
+        static_cast<char*>(__frameStart) +
+        __get_allocator_offset<_Alloc>(__frameSize));
+  }
+
+public:
+  __task_promise_base() _NOEXCEPT = default;
+
+  // Explicitly disable special member functions.
+  __task_promise_base(const __task_promise_base&) = delete;
+  __task_promise_base(__task_promise_base&&) = delete;
+  __task_promise_base& operator=(const __task_promise_base&) = delete;
+  __task_promise_base& operator=(__task_promise_base&&) = delete;
+
+  static void* operator new(size_t __size) {
+    // Allocate space for an extra pointer immediately after __size that holds
+    // the type-erased deallocation function.
+    void* __pointer = ::operator new(__get_padded_frame_size(__size));
+
+    _DeallocFunc*& __deallocFunc = __get_dealloc_func(__pointer, __size);
+    __deallocFunc = [](void* __pointer, size_t __size) _NOEXCEPT {
+      ::operator delete(__pointer, __get_padded_frame_size(__size));
+    };
+
+    return __pointer;
+  }
+
+  template <typename _Alloc, typename... _Args>
+  static void* operator new(size_t __size, allocator_arg_t, _Alloc& __alloc,
+                            _Args&...) {
+    using _CharAlloc =
+        typename allocator_traits<_Alloc>::template rebind_alloc<char>;
+
+    _CharAlloc __charAllocator{__alloc};
+
+    void* __pointer = __charAllocator.allocate(
+        __get_padded_frame_size_with_allocator<_CharAlloc>(__size));
+
+    _DeallocFunc*& __deallocFunc = __get_dealloc_func(__pointer, __size);
+    __deallocFunc = [](void* __pointer, size_t __size) _NOEXCEPT {
+      // Allocators are required to not throw from their move constructors
+      // however they aren't required to be declared noexcept so we can't
+      // actually check this with a static_assert.
+      //
+      // static_assert(is_nothrow_move_constructible<_Alloc>::value,
+      //              "task<T> coroutine custom allocator requires a noexcept "
+      //              "move constructor");
+
+      size_t __paddedSize =
+          __get_padded_frame_size_with_allocator<_CharAlloc>(__size);
+
+      if constexpr (__allocator_needs_to_be_stored<_CharAlloc>) {
+        _CharAlloc& __allocatorInFrame =
+            __get_allocator<_CharAlloc>(__pointer, __size);
+        _CharAlloc __allocatorOnStack = _VSTD::move(__allocatorInFrame);
+        __allocatorInFrame.~_CharAlloc();
+        // Allocator requirements state that deallocate() must not throw.
+        // See [allocator.requirements] from C++ standard.
+        // We are relying on that here.
+        __allocatorOnStack.deallocate(static_cast<char*>(__pointer),
+                                      __paddedSize);
+      } else {
+        _CharAlloc __alloc;
+        __alloc.deallocate(static_cast<char*>(__pointer), __paddedSize);
+      }
+    };
+
+    // Copy the allocator into the heap frame (if required)
+    if constexpr (__allocator_needs_to_be_stored<_CharAlloc>) {
+      // task<T> coroutine custom allocation requires the copy constructor to
+      // not throw but we can't rely on it being declared noexcept.
+      // If it did throw we'd leak the allocation here.
+      ::new (static_cast<void*>(
+          _VSTD::addressof(__get_allocator<_CharAlloc>(__pointer, __size))))
+          _CharAlloc(_VSTD::move(__charAllocator));
+    }
+
+    return __pointer;
+  }
+
+  template <typename _This, typename _Alloc, typename... _Args>
+  static void* operator new(size_t __size, _This&, allocator_arg_t __allocArg, _Alloc& __alloc,
+                            _Args&...) {
+    return __task_promise_base::operator new(__size, __allocArg, __alloc);
+  }
+
+  _LIBCPP_INLINE_VISIBILITY
+  static void operator delete(void* __pointer, size_t __size)_NOEXCEPT {
+    __get_dealloc_func(__pointer, __size)(__pointer, __size);
+  }
+
+  _LIBCPP_INLINE_VISIBILITY
+  suspend_always initial_suspend() const _NOEXCEPT { return {}; }
+
+  _LIBCPP_INLINE_VISIBILITY
+  __task_promise_final_awaitable final_suspend() _NOEXCEPT { return {}; }
+
+  _LIBCPP_INLINE_VISIBILITY
+  void __set_continuation(coroutine_handle<> __continuation) {
+    _LIBCPP_ASSERT(!__continuation_, "task already has a continuation");
+    __continuation_ = __continuation;
+  }
+
+private:
+  friend struct __task_promise_final_awaitable;
+
+  coroutine_handle<> __continuation_;
+};
+
+template <typename _Tp>
+class _LIBCPP_TEMPLATE_VIS __task_promise final : public __task_promise_base {
+  using _Handle = coroutine_handle<__task_promise>;
+
+public:
+  __task_promise() _NOEXCEPT : __state_(_State::__no_value) {}
+
+  ~__task_promise() {
+    switch (__state_) {
+    case _State::__value:
+      __value_.~_Tp();
+      break;
+#ifndef _LIBCPP_NO_EXCEPTIONS
+    case _State::__exception:
+      __exception_.~exception_ptr();
+      break;
+#endif
+    case _State::__no_value:
+      break;
+    };
+  }
+
+  _LIBCPP_INLINE_VISIBILITY
+  task<_Tp> get_return_object() _NOEXCEPT;
+
+  void unhandled_exception() _NOEXCEPT {
+#ifndef _LIBCPP_NO_EXCEPTIONS
+    ::new (static_cast<void*>(&__exception_))
+        exception_ptr(current_exception());
+    __state_ = _State::__exception;
+#else
+    _LIBCPP_ASSERT(
+        false, "task<T> coroutine unexpectedly called unhandled_exception()");
+#endif
+  }
+
+  // Only enable return_value() overload if _Tp is implicitly constructible from
+  // _Value
+  template <typename _Value,
+            enable_if_t<is_convertible<_Value, _Tp>::value, int> = 0>
+  void return_value(_Value&& __value)
+      _NOEXCEPT_((is_nothrow_constructible_v<_Tp, _Value>)) {
+    __construct_value(static_cast<_Value&&>(__value));
+  }
+
+  template <typename _Value>
+  auto return_value(std::initializer_list<_Value> __initializer) _NOEXCEPT_(
+      (is_nothrow_constructible_v<_Tp, std::initializer_list<_Value>>))
+      -> std::enable_if_t<
+          std::is_constructible_v<_Tp, std::initializer_list<_Value>>> {
+    __construct_value(_VSTD::move(__initializer));
+  }
+
+  auto return_value(_Tp&& __value)
+      _NOEXCEPT_((is_nothrow_move_constructible_v<_Tp>))
+          -> std::enable_if_t<std::is_move_constructible_v<_Tp>> {
+    __construct_value(static_cast<_Tp&&>(__value));
+  }
+
+  _Tp& __lvalue_result() {
+    __throw_if_exception();
+    return __value_;
+  }
+
+  _Tp __rvalue_result() {
+    __throw_if_exception();
+    return static_cast<_Tp&&>(__value_);
+  }
+
+private:
+  template <typename... _Args>
+  void __construct_value(_Args&&... __args) {
+    ::new (static_cast<void*>(_VSTD::addressof(__value_)))
+        _Tp(static_cast<_Args&&>(__args)...);
+
+    // Only set __state_ after successfully constructing the value.
+    // If constructor throws then state will be updated by
+    // unhandled_exception().
+    __state_ = _State::__value;
+  }
+
+  void __throw_if_exception() {
+#ifndef _LIBCPP_NO_EXCEPTIONS
+    if (__state_ == _State::__exception) {
+      rethrow_exception(__exception_);
+    }
+#endif
+  }
+
+  enum class _State { __no_value, __value, __exception };
+
+  _State __state_ = _State::__no_value;
+  union {
+    char __empty_;
+    _Tp __value_;
+    exception_ptr __exception_;
+  };
+};
+
+template <typename _Tp>
+class __task_promise<_Tp&> final : public __task_promise_base {
+  using _Ptr = _Tp*;
+  using _Handle = coroutine_handle<__task_promise>;
+
+public:
+  __task_promise() _NOEXCEPT = default;
+
+  ~__task_promise() {
+#ifndef _LIBCPP_NO_EXCEPTIONS
+    if (__has_exception_) {
+      __exception_.~exception_ptr();
+    }
+#endif
+  }
+
+  _LIBCPP_INLINE_VISIBILITY
+  task<_Tp&> get_return_object() _NOEXCEPT;
+
+  void unhandled_exception() _NOEXCEPT {
+#ifndef _LIBCPP_NO_EXCEPTIONS
+    ::new (static_cast<void*>(&__exception_))
+        exception_ptr(current_exception());
+    __has_exception_ = true;
+#else
+    _LIBCPP_ASSERT(
+        false, "task<T> coroutine unexpectedly called unhandled_exception()");
+#endif
+  }
+
+  void return_value(_Tp& __value) _NOEXCEPT {
+    ::new (static_cast<void*>(&__pointer_)) _Ptr(_VSTD::addressof(__value));
+  }
+
+  _Tp& __lvalue_result() {
+    __throw_if_exception();
+    return *__pointer_;
+  }
+
+  _Tp& __rvalue_result() { return __lvalue_result(); }
+
+private:
+  void __throw_if_exception() {
+#ifndef _LIBCPP_NO_EXCEPTIONS
+    if (__has_exception_) {
+      rethrow_exception(__exception_);
+    }
+#endif
+  }
+
+  union {
+    char __empty_;
+    _Ptr __pointer_;
+    exception_ptr __exception_;
+  };
+  bool __has_exception_ = false;
+};
+
+template <>
+class __task_promise<void> final : public __task_promise_base {
+  using _Handle = coroutine_handle<__task_promise>;
+
+public:
+  task<void> get_return_object() _NOEXCEPT;
+
+  void return_void() _NOEXCEPT {}
+
+  void unhandled_exception() _NOEXCEPT {
+#ifndef _LIBCPP_NO_EXCEPTIONS
+    __exception_ = current_exception();
+#endif
+  }
+
+  void __lvalue_result() { __throw_if_exception(); }
+
+  void __rvalue_result() { __throw_if_exception(); }
+
+private:
+  void __throw_if_exception() {
+#ifndef _LIBCPP_NO_EXCEPTIONS
+    if (__exception_) {
+      rethrow_exception(__exception_);
+    }
+#endif
+  }
+
+  exception_ptr __exception_;
+};
+
+template <typename _Tp>
+class _LIBCPP_TEMPLATE_VIS _LIBCPP_NODISCARD_AFTER_CXX17 task {
+public:
+  using promise_type = __task_promise<_Tp>;
+
+private:
+  using _Handle = coroutine_handle<__task_promise<_Tp>>;
+
+  class _AwaiterBase {
+  public:
+    _AwaiterBase(_Handle __coro) _NOEXCEPT : __coro_(__coro) {}
+
+    _LIBCPP_INLINE_VISIBILITY
+    bool await_ready() const { return __coro_.done(); }
+
+    _LIBCPP_INLINE_VISIBILITY
+    _Handle await_suspend(coroutine_handle<> __continuation) const {
+      __coro_.promise().__set_continuation(__continuation);
+      return __coro_;
+    }
+
+  protected:
+    _Handle __coro_;
+  };
+
+public:
+  _LIBCPP_INLINE_VISIBILITY
+  task(task&& __other) _NOEXCEPT
+      : __coro_(_VSTD::exchange(__other.__coro_, {})) {}
+
+  task(const task&) = delete;
+  task& operator=(const task&) = delete;
+
+  _LIBCPP_INLINE_VISIBILITY
+  ~task() {
+    if (__coro_)
+      __coro_.destroy();
+  }
+
+  _LIBCPP_INLINE_VISIBILITY
+  void swap(task& __other) _NOEXCEPT { _VSTD::swap(__coro_, __other.__coro_); }
+
+  _LIBCPP_INLINE_VISIBILITY
+  auto operator co_await() & {
+    class _Awaiter : public _AwaiterBase {
+    public:
+      using _AwaiterBase::_AwaiterBase;
+
+      _LIBCPP_INLINE_VISIBILITY
+      decltype(auto) await_resume() {
+        return this->__coro_.promise().__lvalue_result();
+      }
+    };
+
+    _LIBCPP_ASSERT(__coro_,
+                   "Undefined behaviour to co_await an invalid task<T>");
+    return _Awaiter{__coro_};
+  }
+
+  _LIBCPP_INLINE_VISIBILITY
+  auto operator co_await() && {
+    class _Awaiter : public _AwaiterBase {
+    public:
+      using _AwaiterBase::_AwaiterBase;
+
+      _LIBCPP_INLINE_VISIBILITY
+      decltype(auto) await_resume() {
+        return this->__coro_.promise().__rvalue_result();
+      }
+    };
+
+    _LIBCPP_ASSERT(__coro_,
+                   "Undefined behaviour to co_await an invalid task<T>");
+    return _Awaiter{__coro_};
+  }
+
+private:
+  friend class __task_promise<_Tp>;
+
+  _LIBCPP_INLINE_VISIBILITY
+  task(_Handle __coro) _NOEXCEPT : __coro_(__coro) {}
+
+  _Handle __coro_;
+};
+
+template <typename _Tp>
+task<_Tp> __task_promise<_Tp>::get_return_object() _NOEXCEPT {
+  return task<_Tp>{_Handle::from_promise(*this)};
+}
+
+template <typename _Tp>
+task<_Tp&> __task_promise<_Tp&>::get_return_object() _NOEXCEPT {
+  return task<_Tp&>{_Handle::from_promise(*this)};
+}
+
+task<void> __task_promise<void>::get_return_object() _NOEXCEPT {
+  return task<void>{_Handle::from_promise(*this)};
+}
+
+_LIBCPP_END_NAMESPACE_EXPERIMENTAL_COROUTINES
+
+#endif