[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, 396 insertions, 0 deletions

diff --git a/libcxx/test/std/experimental/task/task.basic/task_custom_allocator.pass.cpp b/libcxx/test/std/experimental/task/task.basic/task_custom_allocator.pass.cpp
new file mode 100644
index 00000000000..bc7e5989861
--- /dev/null
+++ b/libcxx/test/std/experimental/task/task.basic/task_custom_allocator.pass.cpp
@@ -0,0 +1,230 @@
+// -*- C++ -*-
+//===----------------------------------------------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is dual licensed under the MIT and the University of Illinois Open
+// Source Licenses. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+// UNSUPPORTED: c++98, c++03, c++11, c++14
+
+#include <experimental/task>
+#include <cstdlib>
+#include <cassert>
+#include <vector>
+#include <memory>
+#include <experimental/memory_resource>
+
+#include "../sync_wait.hpp"
+
+namespace coro = std::experimental::coroutines_v1;
+
+namespace
+{
+  static size_t allocator_instance_count = 0;
+
+  // A custom allocator that tracks the number of allocator instances that
+  // have been constructed/destructed as well as the number of bytes that
+  // have been allocated/deallocated using the allocator.
+  template<typename T>
+  class my_allocator {
+  public:
+    using value_type = T;
+    using size_type = std::size_t;
+    using difference_type = std::ptrdiff_t;
+    using is_always_equal = std::false_type;
+
+    explicit my_allocator(
+      std::shared_ptr<size_type> totalAllocated) noexcept
+      : totalAllocated_(std::move(totalAllocated))
+    {
+      ++allocator_instance_count;
+      assert(totalAllocated_);
+    }
+
+    my_allocator(const my_allocator& other)
+    : totalAllocated_(other.totalAllocated_)
+    {
+      ++allocator_instance_count;
+    }
+
+    my_allocator(my_allocator&& other)
+    : totalAllocated_(std::move(other.totalAllocated_))
+    {
+      ++allocator_instance_count;
+    }
+
+    template<typename U>
+    my_allocator(const my_allocator<U>& other)
+    : totalAllocated_(other.totalAllocated_)
+    {
+      ++allocator_instance_count;
+    }
+
+    template<typename U>
+    my_allocator(my_allocator<U>&& other)
+    : totalAllocated_(std::move(other.totalAllocated_))
+    {
+      ++allocator_instance_count;
+    }
+
+    ~my_allocator()
+    {
+      --allocator_instance_count;
+    }
+
+    char* allocate(size_t n) {
+      const auto byteCount = n * sizeof(T);
+      void* p = std::malloc(byteCount);
+      if (!p) {
+        throw std::bad_alloc{};
+      }
+      *totalAllocated_ += byteCount;
+      return static_cast<char*>(p);
+    }
+
+    void deallocate(char* p, size_t n) {
+      const auto byteCount = n * sizeof(T);
+      *totalAllocated_ -= byteCount;
+      std::free(p);
+    }
+  private:
+    template<typename U>
+    friend class my_allocator;
+
+    std::shared_ptr<size_type> totalAllocated_;
+  };
+}
+
+template<typename Allocator>
+coro::task<void> f(std::allocator_arg_t, [[maybe_unused]] Allocator alloc)
+{
+  co_return;
+}
+
+void test_custom_allocator_is_destructed()
+{
+  auto totalAllocated = std::make_shared<size_t>(0);
+
+  assert(allocator_instance_count == 0);
+
+  {
+    std::vector<coro::task<>> tasks;
+    tasks.push_back(
+      f(std::allocator_arg, my_allocator<char>{ totalAllocated }));
+    tasks.push_back(
+      f(std::allocator_arg, my_allocator<char>{ totalAllocated }));
+
+    assert(allocator_instance_count == 4);
+    assert(*totalAllocated > 0);
+  }
+
+  assert(allocator_instance_count == 0);
+  assert(*totalAllocated == 0);
+}
+
+void test_custom_allocator_type_rebinding()
+{
+  auto totalAllocated = std::make_shared<size_t>(0);
+  {
+    std::vector<coro::task<>> tasks;
+    tasks.emplace_back(
+      f(std::allocator_arg, my_allocator<int>{ totalAllocated }));
+    coro::sync_wait(tasks[0]);
+  }
+  assert(*totalAllocated == 0);
+  assert(allocator_instance_count == 0);
+}
+
+void test_mixed_custom_allocator_type_erasure()
+{
+  assert(allocator_instance_count == 0);
+
+  // Show that different allocators can be used within a vector of tasks
+  // of the same type. ie. that the allocator is type-erased inside the
+  // coroutine.
+  std::vector<coro::task<>> tasks;
+  tasks.push_back(f(
+    std::allocator_arg, std::allocator<char>{}));
+  tasks.push_back(f(
+    std::allocator_arg,
+    std::experimental::pmr::polymorphic_allocator<char>{
+       std::experimental::pmr::new_delete_resource() }));
+  tasks.push_back(f(
+    std::allocator_arg,
+    my_allocator<char>{ std::make_shared<size_t>(0) }));
+
+  assert(allocator_instance_count > 0);
+
+  for (auto& t : tasks)
+  {
+    coro::sync_wait(t);
+  }
+
+  tasks.clear();
+
+  assert(allocator_instance_count == 0);
+}
+
+template<typename Allocator>
+coro::task<int> add_async(std::allocator_arg_t, [[maybe_unused]] Allocator alloc, int a, int b)
+{
+  co_return a + b;
+}
+
+void test_task_custom_allocator_with_extra_args()
+{
+  std::vector<coro::task<int>> tasks;
+
+  for (int i = 0; i < 5; ++i) {
+    tasks.push_back(add_async(
+      std::allocator_arg,
+      std::allocator<char>{},
+      i, 2 * i));
+  }
+
+  for (int i = 0; i < 5; ++i)
+  {
+    assert(sync_wait(std::move(tasks[i])) == 3 * i);
+  }
+}
+
+struct some_type {
+  template<typename Allocator>
+  coro::task<int> get_async(std::allocator_arg_t, [[maybe_unused]] Allocator alloc) {
+    co_return 42;
+  }
+
+  template<typename Allocator>
+  coro::task<int> add_async(std::allocator_arg_t, [[maybe_unused]] Allocator alloc, int a, int b) {
+    co_return a + b;
+  }
+};
+
+void test_task_custom_allocator_on_member_function()
+{
+  assert(allocator_instance_count == 0);
+
+  auto totalAllocated = std::make_shared<size_t>(0);
+  some_type obj;
+  assert(sync_wait(obj.get_async(std::allocator_arg, std::allocator<char>{})) == 42);
+  assert(sync_wait(obj.get_async(std::allocator_arg, my_allocator<char>{totalAllocated})) == 42);
+  assert(sync_wait(obj.add_async(std::allocator_arg, std::allocator<char>{}, 2, 3)) == 5);
+  assert(sync_wait(obj.add_async(std::allocator_arg, my_allocator<char>{totalAllocated}, 2, 3)) == 5);
+
+  assert(allocator_instance_count == 0);
+  assert(*totalAllocated == 0);
+}
+
+int main()
+{
+  test_custom_allocator_is_destructed();
+  test_custom_allocator_type_rebinding();
+  test_mixed_custom_allocator_type_erasure();
+  test_task_custom_allocator_with_extra_args();
+  test_task_custom_allocator_on_member_function();
+
+  return 0;
+}
diff --git a/libcxx/test/std/experimental/task/task.basic/task_of_value.pass.cpp b/libcxx/test/std/experimental/task/task.basic/task_of_value.pass.cpp
new file mode 100644
index 00000000000..0d5a8c9f619
--- /dev/null
+++ b/libcxx/test/std/experimental/task/task.basic/task_of_value.pass.cpp
@@ -0,0 +1,70 @@
+// -*- C++ -*-
+//===----------------------------------------------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is dual licensed under the MIT and the University of Illinois Open
+// Source Licenses. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+// UNSUPPORTED: c++98, c++03, c++11, c++14
+
+#include <experimental/task>
+#include <string>
+#include <vector>
+#include <memory>
+#include "../sync_wait.hpp"
+
+void test_returning_move_only_type()
+{
+  auto move_only_async =
+   [](bool x) -> std::experimental::task<std::unique_ptr<int>> {
+     if (x) {
+       auto p = std::make_unique<int>(123);
+       co_return p; // Should be implicit std::move(p) here.
+     }
+
+     co_return std::make_unique<int>(456);
+   };
+
+  assert(*sync_wait(move_only_async(true)) == 123);
+  assert(*sync_wait(move_only_async(false)) == 456);
+}
+
+void test_co_return_with_curly_braces()
+{
+  auto t = []() -> std::experimental::task<std::tuple<int, std::string>>
+  {
+    co_return { 123, "test" };
+  }();
+
+  auto result = sync_wait(std::move(t));
+
+  assert(std::get<0>(result) == 123);
+  assert(std::get<1>(result) == "test");
+}
+
+void test_co_return_by_initialiser_list()
+{
+  auto t = []() -> std::experimental::task<std::vector<int>>
+  {
+    co_return { 2, 10, -1 };
+  }();
+
+  auto result = sync_wait(std::move(t));
+
+  assert(result.size() == 3);
+  assert(result[0] == 2);
+  assert(result[1] == 10);
+  assert(result[2] == -1);
+}
+
+int main()
+{
+  test_returning_move_only_type();
+  test_co_return_with_curly_braces();
+  test_co_return_by_initialiser_list();
+
+  return 0;
+}
diff --git a/libcxx/test/std/experimental/task/task.basic/task_of_void.pass.cpp b/libcxx/test/std/experimental/task/task.basic/task_of_void.pass.cpp
new file mode 100644
index 00000000000..de860b5db6a
--- /dev/null
+++ b/libcxx/test/std/experimental/task/task.basic/task_of_void.pass.cpp
@@ -0,0 +1,96 @@
+// -*- C++ -*-
+//===----------------------------------------------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is dual licensed under the MIT and the University of Illinois Open
+// Source Licenses. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+// UNSUPPORTED: c++98, c++03, c++11, c++14
+
+#include <experimental/task>
+#include "../manual_reset_event.hpp"
+#include "../sync_wait.hpp"
+
+#include <optional>
+#include <thread>
+
+namespace coro = std::experimental::coroutines_v1;
+
+static bool has_f_executed = false;
+
+static coro::task<void> f()
+{
+  has_f_executed = true;
+  co_return;
+}
+
+static void test_coroutine_executes_lazily()
+{
+  coro::task<void> t = f();
+  assert(!has_f_executed);
+  coro::sync_wait(t);
+  assert(has_f_executed);
+}
+
+static std::optional<int> last_value_passed_to_g;
+
+static coro::task<void> g(int a)
+{
+  last_value_passed_to_g = a;
+  co_return;
+}
+
+void test_coroutine_accepts_arguments()
+{
+  auto t = g(123);
+  assert(!last_value_passed_to_g);
+  coro::sync_wait(t);
+  assert(last_value_passed_to_g);
+  assert(*last_value_passed_to_g == 123);
+}
+
+int shared_value = 0;
+int read_value = 0;
+
+coro::task<void> consume_async(manual_reset_event& event)
+{
+  co_await event;
+  read_value = shared_value;
+}
+
+void produce(manual_reset_event& event)
+{
+  shared_value = 101;
+  event.set();
+}
+
+void test_async_completion()
+{
+  manual_reset_event e;
+  std::thread t1{ [&e]
+  {
+    sync_wait(consume_async(e));
+  }};
+
+  assert(read_value == 0);
+
+  std::thread t2{ [&e] { produce(e); }};
+
+  t1.join();
+
+  assert(read_value == 101);
+
+  t2.join();
+}
+
+int main()
+{
+  test_coroutine_executes_lazily();
+  test_coroutine_accepts_arguments();
+  test_async_completion();
+
+  return 0;
+}