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