// -*- 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;
}