//===----------------------------------------------------------------------===//
//
// 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 ALLOCATORS_H
#define ALLOCATORS_H

#include <cuda/std/type_traits>
#include <cuda/std/utility>

#include "test_macros.h"

template <class T>
class A1
{
  int id_;

public:
  TEST_FUNC explicit A1(int id = 0) noexcept
      : id_(id)
  {}

  using value_type = T;

  TEST_FUNC int id() const
  {
    return id_;
  }

  STATIC_MEMBER_VAR(copy_called, bool)
  STATIC_MEMBER_VAR(move_called, bool)
  STATIC_MEMBER_VAR(allocate_called, bool)

  TEST_FUNC static cuda::std::pair<T*, cuda::std::size_t>& deallocate_called()
  {
    NV_IF_ELSE_TARGET(NV_IS_DEVICE,
                      (__shared__ cuda::std::pair<T*, cuda::std::size_t> v; return v;),
                      (static cuda::std::pair<T*, cuda::std::size_t> v = 0; return v;))
  }

  TEST_FUNC A1(const A1& a) noexcept
      : id_(a.id())
  {
    copy_called() = true;
  }
  TEST_FUNC A1(A1&& a) noexcept
      : id_(a.id())
  {
    move_called() = true;
  }
  TEST_FUNC A1& operator=(const A1& a) noexcept
  {
    id_           = a.id();
    copy_called() = true;
    return *this;
  }
  TEST_FUNC A1& operator=(A1&& a) noexcept
  {
    id_           = a.id();
    move_called() = true;
    return *this;
  }

  template <class U>
  TEST_FUNC A1(const A1<U>& a) noexcept
      : id_(a.id())
  {
    copy_called() = true;
  }
  template <class U>
  TEST_FUNC A1(A1<U>&& a) noexcept
      : id_(a.id())
  {
    move_called() = true;
  }

  TEST_FUNC T* allocate(cuda::std::size_t n)
  {
    allocate_called() = true;
    return (T*) n;
  }

  TEST_FUNC void deallocate(T* p, cuda::std::size_t n) noexcept
  {
    deallocate_called() = cuda::std::pair<T*, cuda::std::size_t>(p, n);
  }

  TEST_FUNC cuda::std::size_t max_size() const
  {
    return id_;
  }
};

template <class T, class U>
inline TEST_FUNC bool operator==(const A1<T>& x, const A1<U>& y)
{
  return x.id() == y.id();
}

template <class T, class U>
inline TEST_FUNC bool operator!=(const A1<T>& x, const A1<U>& y)
{
  return !(x == y);
}

template <class T>
class A2
{
  int id_;

public:
  TEST_FUNC explicit A2(int id = 0) noexcept
      : id_(id)
  {}

  using value_type = T;

  using size_type       = unsigned;
  using difference_type = int;

  using propagate_on_container_move_assignment = cuda::std::true_type;

  TEST_FUNC int id() const
  {
    return id_;
  }

  STATIC_MEMBER_VAR(copy_called, bool)
  STATIC_MEMBER_VAR(move_called, bool)
  STATIC_MEMBER_VAR(allocate_called, bool)

  TEST_FUNC A2(const A2& a) noexcept
      : id_(a.id())
  {
    copy_called() = true;
  }
  TEST_FUNC A2(A2&& a) noexcept
      : id_(a.id())
  {
    move_called() = true;
  }
  TEST_FUNC A2& operator=(const A2& a) noexcept
  {
    id_           = a.id();
    copy_called() = true;
    return *this;
  }
  TEST_FUNC A2& operator=(A2&& a) noexcept
  {
    id_           = a.id();
    move_called() = true;
    return *this;
  }

  TEST_FUNC T* allocate(cuda::std::size_t, const void* hint)
  {
    allocate_called() = true;
    return (T*) const_cast<void*>(hint);
  }
};

template <class T, class U>
inline TEST_FUNC bool operator==(const A2<T>& x, const A2<U>& y)
{
  return x.id() == y.id();
}

template <class T, class U>
inline TEST_FUNC bool operator!=(const A2<T>& x, const A2<U>& y)
{
  return !(x == y);
}

template <class T>
class A3
{
  int id_;

public:
  TEST_FUNC explicit A3(int id = 0) noexcept
      : id_(id)
  {}

  using value_type = T;

  using propagate_on_container_copy_assignment = cuda::std::true_type;
  using propagate_on_container_swap            = cuda::std::true_type;

  TEST_FUNC int id() const
  {
    return id_;
  }

  STATIC_MEMBER_VAR(copy_called, bool)
  STATIC_MEMBER_VAR(move_called, bool)
  STATIC_MEMBER_VAR(constructed, bool)
  STATIC_MEMBER_VAR(destroy_called, bool)

  TEST_FUNC A3(const A3& a) noexcept
      : id_(a.id())
  {
    copy_called() = true;
  }
  TEST_FUNC A3(A3&& a) noexcept
      : id_(a.id())
  {
    move_called() = true;
  }
  TEST_FUNC A3& operator=(const A3& a) noexcept
  {
    id_           = a.id();
    copy_called() = true;
    return *this;
  }
  TEST_FUNC A3& operator=(A3&& a) noexcept
  {
    id_           = a.id();
    move_called() = true;
    return *this;
  }

  template <class U, class... Args>
  TEST_FUNC void construct(U* p, Args&&... args)
  {
    ::new (p) U(cuda::std::forward<Args>(args)...);
    constructed() = true;
  }

  template <class U>
  TEST_FUNC void destroy(U* p) noexcept
  {
    p->~U();
    destroy_called() = true;
  }

  TEST_FUNC A3 select_on_container_copy_construction() const
  {
    return A3(-1);
  }
};

template <class T, class U>
inline TEST_FUNC bool operator==(const A3<T>& x, const A3<U>& y)
{
  return x.id() == y.id();
}

template <class T, class U>
inline TEST_FUNC bool operator!=(const A3<T>& x, const A3<U>& y)
{
  return !(x == y);
}

#endif // ALLOCATORS_H