Updated <experimental/task> WIP

[lewissbaker]

Summary

Addressed some feedback from initial version.

• factored out allocator padding calculations
• use anonymous unions instead of char-buffers and reinterpret_cast
• make co_await'ing an invalid task have undefined behaviour
• run clang-format over the file to fix some formatting

include/experimental/task

+// -*- C++ -*-
+//===----------------------------- coroutine -----------------------------===//
+//
+//                     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/coroutine>
+
+#include <exception>
+#include <utility>
+#include <type_traits>
+
+#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>
+
+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;
+
+  static _LIBCPP_CONSTEXPR size_t __pad(size_t __value,
+                                        size_t __alignment) _NOEXCEPT {
+    return (__value + (__alignment - 1u)) & ~(__alignment - 1u);
+  }
+
+  static _LIBCPP_CONSTEXPR size_t
+  __get_dealloc_func_offset(size_t __frameSize) _NOEXCEPT {
+    return __pad(__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 __pad(__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 {
+    return __get_allocator_offset<_Alloc>(__frameSize) + sizeof(_Alloc);
+  }
+
+  _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:
+  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& __allocator,
+                            _Args&...) {
+    static_assert(is_same<typename _Alloc::value_type, char>::value,
+                  "task<T> coroutine custom allocator must have a 'value_type' "
+                  "of 'char'.");
+
+    void* __pointer = __allocator.allocate(
+        __get_padded_frame_size_with_allocator<_Alloc>(__size));
+
+    _DeallocFunc*& __deallocFunc = __get_dealloc_func(__pointer, __size);
+    __deallocFunc = [](void* __pointer, size_t __size) _NOEXCEPT {
+      static_assert(is_nothrow_move_constructible<_Alloc>::value,
+                    "task<T> coroutine custom allocator requires a noexcept "
+                    "move constructor");
+
+      _Alloc& __allocatorInFrame = __get_allocator<_Alloc>(__pointer, __size);
+      _Alloc __allocatorOnStack = move(__allocatorInFrame);
+      __allocatorInFrame.~_Alloc();
+
+      size_t __paddedSize =
+          __get_padded_frame_size_with_allocator<_Alloc>(__size);
+
+      // 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);
+    };
+
+    // Copy the allocator into the heap frame.
+
+    static_assert(is_nothrow_copy_constructible<_Alloc>::value,
+                  "task<T> coroutine custom allocator requires a noexcept copy "
+                  "constructor");
+
+    new (static_cast<void*>(_VSTD::addressof(
+        __get_allocator<_Alloc>(__pointer, __size)))) _Alloc(__allocator);
+
+    return __pointer;
+  }
+
+  _LIBCPP_INLINE_VISIBILITY
+  static void operator delete(void* __pointer, size_t __size) {
+    __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 class __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 {}
+
+  ~__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
+  _Handle get_return_object() _NOEXCEPT { return _Handle::from_promise(*this); }
+
+  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<_Tp, _Value>::value)) {
+    new (static_cast<void*>(_VSTD::addressof(__value_)))
+        _Tp(static_cast<_Value&&>(__value));
+
+    // Only set __state_ after successfully constructing the value.
+    // If constructor throws then state will be updated by unhandled_exception().
+    __state_ = _State::__value;
+  }
+
+  _Tp& __lvalue_result() {
+    __throw_if_exception();
+    return __value_;
+  }
+
+  _Tp&& __rvalue_result() {
+    __throw_if_exception();
+    return static_cast<_Tp&&>(__value_);
+  }
+
+private:
+  void __throw_if_exception() {
+#ifndef _LIBCPP_NO_EXCEPTIONS
+    if (__state_ == _State::__exception) {
+      rethrow_exception(__exception_);
+    }
+#endif
+  }
+
+  enum class _State {
+    __no_value,
+    __value
+#ifndef _LIBCPP_NO_EXCEPTIONS
+    ,
+    __exception
+#endif
+  };
+
+  _State __state_ = _State::__no_value;
+  union {
+    char __empty_;
+    _Tp __value_;
+#ifndef _LIBCPP_NO_EXCEPTIONS
+    exception_ptr __exception_;
+#endif
+  };
+};
+
+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 {}
+
+  ~__task_promise() {
+#ifndef _LIBCPP_NO_EXCEPTIONS
+    if (__has_exception_) {
+      __exception_.~exception_ptr();
+    }
+#endif
+  }
+
+  _LIBCPP_INLINE_VISIBILITY
+  _Handle get_return_object() _NOEXCEPT { return _Handle::from_promise(*this); }
+
+  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_;
+#ifndef _LIBCPP_NO_EXCEPTIONS
+    exception_ptr __exception_;
+#endif
+  };
+#ifndef _LIBCPP_NO_EXCEPTIONS
+  bool __has_exception_ = false;
+#endif
+};
+
+template <>
+class __task_promise<void> final : public __task_promise_base {
+  using _Handle = coroutine_handle<__task_promise>;
+
+public:
+  _Handle get_return_object() _NOEXCEPT { return _Handle::from_promise(*this); }
+
+  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
+  }
+
+#ifndef _LIBCPP_NO_EXCEPTIONS
+  exception_ptr __exception_;
+#endif
+};
+
+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(_Handle __coro) _NOEXCEPT : __coro_(__coro) {}
+
+  _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
+  task& operator=(task __other) _NOEXCEPT {
+    swap(__other);
+    return *this;
+  }
+
+  _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:
+  _Handle __coro_;
+};
+
+_LIBCPP_END_NAMESPACE_EXPERIMENTAL_COROUTINES
+
+#endif
+
+////////////////////////
+// Tests for task<T>
+
+using namespace std::experimental;
+
+task<void> f1() { co_return; }
+
+task<void> f2() { co_await f1(); }
+
+task<void> f3() {
+  task<void> t = f1();
+  co_await t;
+}
+
+task<int> g1() { co_return 123; }
+
+task<int> g2() { co_return co_await g1() + 100; }
+
+task<int> g3() {
+  // co_await rvalue returns rvalue reference to result.
+  {
+    // BUG: This will leave x as a dangling rvalue reference.
+    // This code is not intended to be executed, it's just here
+    // to check the return type of co_await'ing a task<T>&&.
+    decltype(auto) x = co_await g1();
+    static_assert(std::is_same<decltype(x), int&&>::value);
+  }
+
+  // co_await lvalue returns lvalue reference to result
+  task<int> t = g1();
+  decltype(auto) y = co_await t;
+  static_assert(std::is_same_v<decltype(y), int&>);
+
+  co_return y * 2;
+}
+
+class MoveOnly {
+public:
+  MoveOnly();
+  MoveOnly(const MoveOnly&) = delete;
+  MoveOnly(MoveOnly&&) noexcept;
+  ~MoveOnly();
+  int get() const;
+
+private:
+  void* _data;
+};
+
+task<MoveOnly> h1(bool x) {
+  if (x) {
+    MoveOnly value;
+    co_return std::move(value);
+  }
+  co_return MoveOnly{};
+}
+
+task<int> h2() {
+  auto x = co_await h1(true);
+  co_return x.get();
+}
+
+static int x;
+static int y;
+
+task<int&> h3(bool cond) { co_return cond ? x : y; }
+
+task<void> h3consumer() {
+  int& result = co_await h3(true);
+  result = 32;
+}
+
+// custom allocator tests
+
+template <typename Allocator>
+task<void> a1(std::allocator_arg_t, Allocator allocator, bool x) {
+  co_return;
+}
+
+class my_allocator {
+public:
+  using value_type = char;
+  my_allocator();
+  my_allocator(const my_allocator&) noexcept;
+  my_allocator(my_allocator&&) noexcept;
+  char* allocate(size_t n);
+  void deallocate(char* p, size_t n);
+
+private:
+  void* state;
+};
+
+task<void> a2() {
+  my_allocator alloc;
+  a1(std::allocator_arg, alloc, true);
+  co_return;
+}
+
+task<void> a2a() {
+  my_allocator alloc;
+  return a1(std::allocator_arg, alloc, true);
+}
+
+task<void> a3() {
+  co_await a1(std::allocator_arg, std::allocator<char>{}, false);
+}
+
+task<void> a3a() {
+  return a1(std::allocator_arg, std::allocator<char>{}, false);
+}
 No newline at end of file

include/module.modulemap

module experimental {
 
     module algorithm {
       header "experimental/algorithm"
       export *
     }
      module coroutine {
       requires coroutines
       header "experimental/coroutine"
+      header "experimental/task"
       export *
     }
     module deque {
       header "experimental/deque"
       export *
     }
     module dynarray {
       header "experimental/dynarray"