//===----------------------------------------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#ifndef USES_ALLOC_TYPES_H
#define USES_ALLOC_TYPES_H

#include <cassert>
#include <cstdlib>
#include <memory>

#include "test_macros.h"
#include "test_workarounds.h"
#include "type_id.h"

// There are two forms of uses-allocator construction:
//   (1) UA_AllocArg: 'T(allocator_arg_t, Alloc const&, Args&&...)'
//   (2) UA_AllocLast: 'T(Args&&..., Alloc const&)'
// 'UA_None' represents non-uses allocator construction.
enum class UsesAllocatorType
{
  UA_None      = 0,
  UA_AllocArg  = 2,
  UA_AllocLast = 4
};
constexpr UsesAllocatorType UA_None      = UsesAllocatorType::UA_None;
constexpr UsesAllocatorType UA_AllocArg  = UsesAllocatorType::UA_AllocArg;
constexpr UsesAllocatorType UA_AllocLast = UsesAllocatorType::UA_AllocLast;

inline const char* toString(UsesAllocatorType UA)
{
  switch (UA)
  {
    case UA_None:
      return "UA_None";
    case UA_AllocArg:
      return "UA_AllocArg";
    case UA_AllocLast:
      return "UA_AllocLast";
    default:
      std::abort();
  }
}

#define COMPARE_ALLOC_TYPE(LHS, RHS) CompareVerbose(#LHS, LHS, #RHS, RHS)

inline bool CompareVerbose(const char* LHSString, UsesAllocatorType LHS, const char* RHSString, UsesAllocatorType RHS)
{
  if (LHS == RHS)
  {
    return true;
  }
  std::printf(
    "UsesAllocatorType's don't match:\n%s %s\n----------\n%s %s\n", LHSString, toString(LHS), RHSString, toString(RHS));
  return false;
}

template <class Alloc, std::size_t N>
class UsesAllocatorV1;
// Implements form (1) of uses-allocator construction from the specified
// 'Alloc' type and exactly 'N' additional arguments. It also provides
// non-uses allocator construction from 'N' arguments. This test type
// blows up when form (2) of uses-allocator is even considered.

template <class Alloc, std::size_t N>
class UsesAllocatorV2;
// Implements form (2) of uses-allocator construction from the specified
// 'Alloc' type and exactly 'N' additional arguments. It also provides
// non-uses allocator construction from 'N' arguments.

template <class Alloc, std::size_t N>
class UsesAllocatorV3;
// Implements both form (1) and (2) of uses-allocator construction from
// the specified 'Alloc' type and exactly 'N' additional arguments. It also
// provides non-uses allocator construction from 'N' arguments.

template <class Alloc, std::size_t>
class NotUsesAllocator;
// Implements both form (1) and (2) of uses-allocator construction from
// the specified 'Alloc' type and exactly 'N' additional arguments. It also
// provides non-uses allocator construction from 'N' arguments. However
// 'NotUsesAllocator' never provides a 'allocator_type' typedef so it is
// never automatically uses-allocator constructed.

template <class... ArgTypes, class TestType>
bool checkConstruct(TestType& value, UsesAllocatorType form, typename TestType::CtorAlloc const& alloc)
// Check that 'value' was constructed using the specified 'form' of
// construction and with the specified 'ArgTypes...'. Additionally
// check that 'value' was constructed using the specified 'alloc'.
{
  if (form == UA_None)
  {
    return value.template checkConstruct<ArgTypes&&...>(form);
  }
  else
  {
    return value.template checkConstruct<ArgTypes&&...>(form, alloc);
  }
}

template <class... ArgTypes, class TestType>
bool checkConstruct(TestType& value, UsesAllocatorType form)
{
  return value.template checkConstruct<ArgTypes&&...>(form);
}

template <class TestType>
bool checkConstructionEquiv(TestType& T, TestType& U)
// check that 'T' and 'U' where initialized in the exact same manner.
{
  return T.checkConstructEquiv(U);
}

////////////////////////////////////////////////////////////////////////////////
namespace detail
{

template <bool IsZero, size_t N, class ArgList, class... Args>
struct TakeNImp;

template <class ArgList, class... Args>
struct TakeNImp<true, 0, ArgList, Args...>
{
  typedef ArgList type;
};

template <size_t N, class... A1, class F, class... R>
struct TakeNImp<false, N, ArgumentListID<A1...>, F, R...> : TakeNImp<N - 1 == 0, N - 1, ArgumentListID<A1..., F>, R...>
{};

template <size_t N, class... Args>
struct TakeNArgs : TakeNImp<N == 0, N, ArgumentListID<>, Args...>
{};

template <class T>
struct Identity
{
  typedef T type;
};

template <class T>
using IdentityT = typename Identity<T>::type;

template <bool Value>
using EnableIfB = typename std::enable_if<Value, bool>::type;

} // end namespace detail

// FIXME: UsesAllocatorTestBase needs some special logic to deal with
// polymorphic allocators. However we don't want to include
// <experimental/memory_resource> in this header. Therefore in order
// to inject this behavior later we use a trait.
// See test_memory_resource.h for more info.
template <class Alloc>
struct TransformErasedTypeAlloc
{
  using type = Alloc;
};

using detail::EnableIfB;

struct AllocLastTag
{};

template <class Alloc, bool = std::is_default_constructible<Alloc>::value>
struct UsesAllocatorTestBaseStorage
{
  Alloc allocator;
  UsesAllocatorTestBaseStorage() = default;
  UsesAllocatorTestBaseStorage(Alloc const& a)
      : allocator(a)
  {}
  const Alloc* get_allocator() const
  {
    return &allocator;
  }
};

template <class Alloc>
struct UsesAllocatorTestBaseStorage<Alloc, false>
{
  union
  {
    char dummy;
    Alloc alloc;
  };
  bool has_alloc = false;

  UsesAllocatorTestBaseStorage()
      : dummy()
      , has_alloc(false)
  {}
  UsesAllocatorTestBaseStorage(Alloc const& a)
      : alloc(a)
      , has_alloc(true)
  {}
  ~UsesAllocatorTestBaseStorage()
  {
    if (has_alloc)
    {
      alloc.~Alloc();
    }
  }

  Alloc const* get_allocator() const
  {
    if (!has_alloc)
    {
      return nullptr;
    }
    return &alloc;
  }
};

template <class Self, class Alloc>
struct UsesAllocatorTestBase
{
public:
  using CtorAlloc = typename TransformErasedTypeAlloc<Alloc>::type;

  template <class... ArgTypes>
  bool checkConstruct(UsesAllocatorType expectType) const
  {
    auto expectArgs = &makeArgumentID<ArgTypes...>();
    return COMPARE_ALLOC_TYPE(expectType, constructor_called) && COMPARE_TYPEID(args_id, expectArgs);
  }

  template <class... ArgTypes>
  bool checkConstruct(UsesAllocatorType expectType, CtorAlloc const& expectAlloc) const
  {
    auto ExpectID = &makeArgumentID<ArgTypes...>();
    return COMPARE_ALLOC_TYPE(expectType, constructor_called) && COMPARE_TYPEID(args_id, ExpectID) && has_alloc()
        && expectAlloc == *get_alloc();
  }

  bool checkConstructEquiv(UsesAllocatorTestBase& O) const
  {
    if (has_alloc() != O.has_alloc())
    {
      return false;
    }
    return COMPARE_ALLOC_TYPE(constructor_called, O.constructor_called) && COMPARE_TYPEID(args_id, O.args_id)
        && (!has_alloc() || *get_alloc() == *O.get_alloc());
  }

protected:
  explicit UsesAllocatorTestBase(const TypeID* aid)
      : args_id(aid)
      , constructor_called(UA_None)
      , alloc_store()
  {}

  UsesAllocatorTestBase(UsesAllocatorTestBase const&)
      : args_id(&makeArgumentID<Self const&>())
      , constructor_called(UA_None)
      , alloc_store()
  {}

  UsesAllocatorTestBase(UsesAllocatorTestBase&&)
      : args_id(&makeArgumentID<Self&&>())
      , constructor_called(UA_None)
      , alloc_store()
  {}

  template <class... Args>
  UsesAllocatorTestBase(std::allocator_arg_t, CtorAlloc const& a, Args&&...)
      : args_id(&makeArgumentID<Args&&...>())
      , constructor_called(UA_AllocArg)
      , alloc_store(a)
  {}

  template <class... Args, class ArgsIDL = detail::TakeNArgs<sizeof...(Args) - 1, Args&&...>>
  UsesAllocatorTestBase(AllocLastTag, Args&&... args)
      : args_id(&makeTypeIDImp<typename ArgsIDL::type>())
      , constructor_called(UA_AllocLast)
      , alloc_store(UsesAllocatorTestBase::getAllocatorFromPack(typename ArgsIDL::type{}, std::forward<Args>(args)...))
  {}

private:
  template <class... LArgs, class... Args>
  static CtorAlloc getAllocatorFromPack(ArgumentListID<LArgs...>, Args&&... args)
  {
    return UsesAllocatorTestBase::getAllocatorFromPackImp<LArgs const&...>(args...);
  }

  template <class... LArgs>
  static CtorAlloc getAllocatorFromPackImp(typename detail::Identity<LArgs>::type..., CtorAlloc const& alloc)
  {
    return alloc;
  }

  bool has_alloc() const
  {
    return alloc_store.get_allocator() != nullptr;
  }
  const CtorAlloc* get_alloc() const
  {
    return alloc_store.get_allocator();
  }

public:
  const TypeID* args_id;
  UsesAllocatorType constructor_called = UA_None;
  UsesAllocatorTestBaseStorage<CtorAlloc> alloc_store;
};

template <class Alloc, size_t Arity>
class UsesAllocatorV1 : public UsesAllocatorTestBase<UsesAllocatorV1<Alloc, Arity>, Alloc>
{
public:
  typedef Alloc allocator_type;

  using Base      = UsesAllocatorTestBase<UsesAllocatorV1, Alloc>;
  using CtorAlloc = typename Base::CtorAlloc;

  UsesAllocatorV1()
      : Base(&makeArgumentID<>())
  {}

  UsesAllocatorV1(UsesAllocatorV1 const&)
      : Base(&makeArgumentID<UsesAllocatorV1 const&>())
  {}
  UsesAllocatorV1(UsesAllocatorV1&&)
      : Base(&makeArgumentID<UsesAllocatorV1&&>())
  {}
  // Non-Uses Allocator Ctor
  template <class... Args, EnableIfB<sizeof...(Args) == Arity> = false>
  UsesAllocatorV1(Args&&...)
      : Base(&makeArgumentID<Args&&...>())
  {}

  // Uses Allocator Arg Ctor
  template <class... Args>
  UsesAllocatorV1(std::allocator_arg_t tag, CtorAlloc const& a, Args&&... args)
      : Base(tag, a, std::forward<Args>(args)...)
  {}

  // BLOWS UP: Uses Allocator Last Ctor
  template <class First, class... Args, EnableIfB<sizeof...(Args) == Arity> Dummy = false>
  constexpr UsesAllocatorV1(First&&, Args&&...)
  {
    static_assert(!std::is_same<First, First>::value, "");
  }
};

template <class Alloc, size_t Arity>
class UsesAllocatorV2 : public UsesAllocatorTestBase<UsesAllocatorV2<Alloc, Arity>, Alloc>
{
public:
  typedef Alloc allocator_type;

  using Base      = UsesAllocatorTestBase<UsesAllocatorV2, Alloc>;
  using CtorAlloc = typename Base::CtorAlloc;

  UsesAllocatorV2()
      : Base(&makeArgumentID<>())
  {}
  UsesAllocatorV2(UsesAllocatorV2 const&)
      : Base(&makeArgumentID<UsesAllocatorV2 const&>())
  {}
  UsesAllocatorV2(UsesAllocatorV2&&)
      : Base(&makeArgumentID<UsesAllocatorV2&&>())
  {}

  // Non-Uses Allocator Ctor
  template <class... Args, EnableIfB<sizeof...(Args) == Arity> = false>
  UsesAllocatorV2(Args&&...)
      : Base(&makeArgumentID<Args&&...>())
  {}

  // Uses Allocator Last Ctor
  template <class... Args, EnableIfB<sizeof...(Args) == Arity + 1> = false>
  UsesAllocatorV2(Args&&... args)
      : Base(AllocLastTag{}, std::forward<Args>(args)...)
  {}
};

template <class Alloc, size_t Arity>
class UsesAllocatorV3 : public UsesAllocatorTestBase<UsesAllocatorV3<Alloc, Arity>, Alloc>
{
public:
  typedef Alloc allocator_type;

  using Base      = UsesAllocatorTestBase<UsesAllocatorV3, Alloc>;
  using CtorAlloc = typename Base::CtorAlloc;

  UsesAllocatorV3()
      : Base(&makeArgumentID<>())
  {}
  UsesAllocatorV3(UsesAllocatorV3 const&)
      : Base(&makeArgumentID<UsesAllocatorV3 const&>())
  {}
  UsesAllocatorV3(UsesAllocatorV3&&)
      : Base(&makeArgumentID<UsesAllocatorV3&&>())
  {}

  // Non-Uses Allocator Ctor
  template <class... Args, EnableIfB<sizeof...(Args) == Arity> = false>
  UsesAllocatorV3(Args&&...)
      : Base(&makeArgumentID<Args&&...>())
  {}

  // Uses Allocator Arg Ctor
  template <class... Args>
  UsesAllocatorV3(std::allocator_arg_t tag, CtorAlloc const& alloc, Args&&... args)
      : Base(tag, alloc, std::forward<Args>(args)...)
  {}

  // Uses Allocator Last Ctor
  template <class... Args, EnableIfB<sizeof...(Args) == Arity + 1> = false>
  UsesAllocatorV3(Args&&... args)
      : Base(AllocLastTag{}, std::forward<Args>(args)...)
  {}
};

template <class Alloc, size_t Arity>
class NotUsesAllocator : public UsesAllocatorTestBase<NotUsesAllocator<Alloc, Arity>, Alloc>
{
public:
  // no allocator_type typedef provided

  using Base      = UsesAllocatorTestBase<NotUsesAllocator, Alloc>;
  using CtorAlloc = typename Base::CtorAlloc;

  NotUsesAllocator()
      : Base(&makeArgumentID<>())
  {}
  NotUsesAllocator(NotUsesAllocator const&)
      : Base(&makeArgumentID<NotUsesAllocator const&>())
  {}
  NotUsesAllocator(NotUsesAllocator&&)
      : Base(&makeArgumentID<NotUsesAllocator&&>())
  {}
  // Non-Uses Allocator Ctor
  template <class... Args, EnableIfB<sizeof...(Args) == Arity> = false>
  NotUsesAllocator(Args&&...)
      : Base(&makeArgumentID<Args&&...>())
  {}

  // Uses Allocator Arg Ctor
  template <class... Args>
  NotUsesAllocator(std::allocator_arg_t tag, CtorAlloc const& alloc, Args&&... args)
      : Base(tag, alloc, std::forward<Args>(args)...)
  {}

  // Uses Allocator Last Ctor
  template <class... Args, EnableIfB<sizeof...(Args) == Arity + 1> = false>
  NotUsesAllocator(Args&&... args)
      : Base(AllocLastTag{}, std::forward<Args>(args)...)
  {}
};

#endif /* USES_ALLOC_TYPES_H */
