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

// Modifications Copyright (c) 2025 Advanced Micro Devices, Inc.
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

#ifndef COUNT_NEW_H
#define COUNT_NEW_H

#include <cuda/std/cassert>
#include <cuda/std/cstdlib>

#if !_CCCL_COMPILER(NVRTC) && !defined(_CCCL_COMPILER_HIPRTC)
#  include <new>
#endif // !_CCCL_COMPILER(NVRTC)

#include "test_macros.h"

// TODO(HIP/AMD): When we use  __HIP_DEVICE_COMPILE__ this code is failing due to a runtime assertion.
// Maybe switch back from __HIP_PLATFORM_AMD__ to __HIP_DEVICE_COMPILE__ once a fix for the failing tests during host compile is found
#if defined(TEST_HAS_SANITIZERS) || defined(__CUDA_ARCH__) || defined(__HIP_PLATFORM_AMD__) // we cannot overload operator {new, delete} on device
#  define DISABLE_NEW_COUNT
#endif

// All checks return true when disable_checking is enabled.
#ifdef DISABLE_NEW_COUNT
TEST_GLOBAL_VARIABLE const bool MemCounter_disable_checking = true;
#else
TEST_GLOBAL_VARIABLE const bool MemCounter_disable_checking = false;
#endif

// number of allocations to throw after. Default (unsigned)-1. If
// throw_after has the default value it will never be decremented.
TEST_GLOBAL_VARIABLE const unsigned MemCounter_never_throw_value = static_cast<unsigned>(-1);

class MemCounter
{
public:
  MemCounter() = default;

  // Make MemCounter super hard to accidentally construct or copy.
  class MemCounterCtorArg_
  {};
  __host__ __device__ explicit MemCounter(MemCounterCtorArg_) {}

private:
  __host__ __device__ MemCounter(MemCounter const&)            = delete;
  __host__ __device__ MemCounter& operator=(MemCounter const&) = delete;

public:
  // Disallow any allocations from occurring. Useful for testing that
  // code doesn't perform any allocations.
  bool disable_allocations = false;

  unsigned throw_after = MemCounter_never_throw_value;

  int outstanding_new                 = 0;
  int new_called                      = 0;
  int delete_called                   = 0;
  int aligned_new_called              = 0;
  int aligned_delete_called           = 0;
  cuda::std::size_t last_new_size     = 0;
  cuda::std::size_t last_new_align    = 0;
  cuda::std::size_t last_delete_align = 0;

  int outstanding_array_new                 = 0;
  int new_array_called                      = 0;
  int delete_array_called                   = 0;
  int aligned_new_array_called              = 0;
  int aligned_delete_array_called           = 0;
  cuda::std::size_t last_new_array_size     = 0;
  cuda::std::size_t last_new_array_align    = 0;
  cuda::std::size_t last_delete_array_align = 0;

public:
  __host__ __device__ void newCalled(cuda::std::size_t s)
  {
    assert(disable_allocations == false);
    assert(s);
    if (throw_after == 0)
    {
      throw_after = MemCounter_never_throw_value;
      cuda::std::__throw_bad_alloc();
    }
    else if (throw_after != MemCounter_never_throw_value)
    {
      --throw_after;
    }
    ++new_called;
    ++outstanding_new;
    last_new_size = s;
  }

  __host__ __device__ void alignedNewCalled(cuda::std::size_t s, cuda::std::size_t a)
  {
    newCalled(s);
    ++aligned_new_called;
    last_new_align = a;
  }

  __host__ __device__ void deleteCalled(void* p)
  {
    assert(p);
    --outstanding_new;
    ++delete_called;
  }

  __host__ __device__ void alignedDeleteCalled(void* p, cuda::std::size_t a)
  {
    deleteCalled(p);
    ++aligned_delete_called;
    last_delete_align = a;
  }

  __host__ __device__ void newArrayCalled(cuda::std::size_t s)
  {
    assert(disable_allocations == false);
    assert(s);
    if (throw_after == 0)
    {
      throw_after = MemCounter_never_throw_value;
      cuda::std::__throw_bad_alloc();
    }
    else
    {
      // don't decrement throw_after here. newCalled will end up doing that.
    }
    ++outstanding_array_new;
    ++new_array_called;
    last_new_array_size = s;
  }

  __host__ __device__ void alignedNewArrayCalled(cuda::std::size_t s, cuda::std::size_t a)
  {
    newArrayCalled(s);
    ++aligned_new_array_called;
    last_new_array_align = a;
  }

  __host__ __device__ void deleteArrayCalled(void* p)
  {
    assert(p);
    --outstanding_array_new;
    ++delete_array_called;
  }

  __host__ __device__ void alignedDeleteArrayCalled(void* p, cuda::std::size_t a)
  {
    deleteArrayCalled(p);
    ++aligned_delete_array_called;
    last_delete_array_align = a;
  }

  __host__ __device__ void disableAllocations()
  {
    disable_allocations = true;
  }

  __host__ __device__ void enableAllocations()
  {
    disable_allocations = false;
  }

  __host__ __device__ void reset()
  {
    disable_allocations = false;
    throw_after         = MemCounter_never_throw_value;

    outstanding_new       = 0;
    new_called            = 0;
    delete_called         = 0;
    aligned_new_called    = 0;
    aligned_delete_called = 0;
    last_new_size         = 0;
    last_new_align        = 0;

    outstanding_array_new       = 0;
    new_array_called            = 0;
    delete_array_called         = 0;
    aligned_new_array_called    = 0;
    aligned_delete_array_called = 0;
    last_new_array_size         = 0;
    last_new_array_align        = 0;
  }

public:
  __host__ __device__ bool checkOutstandingNewEq(int n) const
  {
    return MemCounter_disable_checking || n == outstanding_new;
  }

  __host__ __device__ bool checkOutstandingNewNotEq(int n) const
  {
    return MemCounter_disable_checking || n != outstanding_new;
  }

  __host__ __device__ bool checkNewCalledEq(int n) const
  {
    return MemCounter_disable_checking || n == new_called;
  }

  __host__ __device__ bool checkNewCalledNotEq(int n) const
  {
    return MemCounter_disable_checking || n != new_called;
  }

  __host__ __device__ bool checkNewCalledGreaterThan(int n) const
  {
    return MemCounter_disable_checking || new_called > n;
  }

  __host__ __device__ bool checkDeleteCalledEq(int n) const
  {
    return MemCounter_disable_checking || n == delete_called;
  }

  __host__ __device__ bool checkDeleteCalledNotEq(int n) const
  {
    return MemCounter_disable_checking || n != delete_called;
  }

  __host__ __device__ bool checkAlignedNewCalledEq(int n) const
  {
    return MemCounter_disable_checking || n == aligned_new_called;
  }

  __host__ __device__ bool checkAlignedNewCalledNotEq(int n) const
  {
    return MemCounter_disable_checking || n != aligned_new_called;
  }

  __host__ __device__ bool checkAlignedNewCalledGreaterThan(int n) const
  {
    return MemCounter_disable_checking || aligned_new_called > n;
  }

  __host__ __device__ bool checkAlignedDeleteCalledEq(int n) const
  {
    return MemCounter_disable_checking || n == aligned_delete_called;
  }

  __host__ __device__ bool checkAlignedDeleteCalledNotEq(int n) const
  {
    return MemCounter_disable_checking || n != aligned_delete_called;
  }

  __host__ __device__ bool checkLastNewSizeEq(cuda::std::size_t n) const
  {
    return MemCounter_disable_checking || n == last_new_size;
  }

  __host__ __device__ bool checkLastNewSizeNotEq(cuda::std::size_t n) const
  {
    return MemCounter_disable_checking || n != last_new_size;
  }

  __host__ __device__ bool checkLastNewAlignEq(cuda::std::size_t n) const
  {
    return MemCounter_disable_checking || n == last_new_align;
  }

  __host__ __device__ bool checkLastNewAlignNotEq(cuda::std::size_t n) const
  {
    return MemCounter_disable_checking || n != last_new_align;
  }

  __host__ __device__ bool checkLastDeleteAlignEq(cuda::std::size_t n) const
  {
    return MemCounter_disable_checking || n == last_delete_align;
  }

  __host__ __device__ bool checkLastDeleteAlignNotEq(cuda::std::size_t n) const
  {
    return MemCounter_disable_checking || n != last_delete_align;
  }

  __host__ __device__ bool checkOutstandingArrayNewEq(int n) const
  {
    return MemCounter_disable_checking || n == outstanding_array_new;
  }

  __host__ __device__ bool checkOutstandingArrayNewNotEq(int n) const
  {
    return MemCounter_disable_checking || n != outstanding_array_new;
  }

  __host__ __device__ bool checkNewArrayCalledEq(int n) const
  {
    return MemCounter_disable_checking || n == new_array_called;
  }

  __host__ __device__ bool checkNewArrayCalledNotEq(int n) const
  {
    return MemCounter_disable_checking || n != new_array_called;
  }

  __host__ __device__ bool checkDeleteArrayCalledEq(int n) const
  {
    return MemCounter_disable_checking || n == delete_array_called;
  }

  __host__ __device__ bool checkDeleteArrayCalledNotEq(int n) const
  {
    return MemCounter_disable_checking || n != delete_array_called;
  }

  __host__ __device__ bool checkAlignedNewArrayCalledEq(int n) const
  {
    return MemCounter_disable_checking || n == aligned_new_array_called;
  }

  __host__ __device__ bool checkAlignedNewArrayCalledNotEq(int n) const
  {
    return MemCounter_disable_checking || n != aligned_new_array_called;
  }

  __host__ __device__ bool checkAlignedNewArrayCalledGreaterThan(int n) const
  {
    return MemCounter_disable_checking || aligned_new_array_called > n;
  }

  __host__ __device__ bool checkAlignedDeleteArrayCalledEq(int n) const
  {
    return MemCounter_disable_checking || n == aligned_delete_array_called;
  }

  __host__ __device__ bool checkAlignedDeleteArrayCalledNotEq(int n) const
  {
    return MemCounter_disable_checking || n != aligned_delete_array_called;
  }

  __host__ __device__ bool checkLastNewArraySizeEq(cuda::std::size_t n) const
  {
    return MemCounter_disable_checking || n == last_new_array_size;
  }

  __host__ __device__ bool checkLastNewArraySizeNotEq(cuda::std::size_t n) const
  {
    return MemCounter_disable_checking || n != last_new_array_size;
  }

  __host__ __device__ bool checkLastNewArrayAlignEq(cuda::std::size_t n) const
  {
    return MemCounter_disable_checking || n == last_new_array_align;
  }

  __host__ __device__ bool checkLastNewArrayAlignNotEq(cuda::std::size_t n) const
  {
    return MemCounter_disable_checking || n != last_new_array_align;
  }
};

TEST_GLOBAL_VARIABLE MemCounter counter{};

__host__ __device__ inline constexpr MemCounter* getGlobalMemCounter()
{
  return &counter;
}

TEST_GLOBAL_VARIABLE MemCounter& globalMemCounter = *getGlobalMemCounter();

#ifndef DISABLE_NEW_COUNT
void* operator new(cuda::std::size_t s)
{
  getGlobalMemCounter()->newCalled(s);
  void* ret = malloc(s);
  if (ret == nullptr)
  {
    cuda::std::__throw_bad_alloc();
  }
  return ret;
}

void operator delete(void* p) noexcept
{
  getGlobalMemCounter()->deleteCalled(p);
  free(p);
}

#  if TEST_COMPILER(GCC)
void operator delete(void* p, cuda::std::size_t) noexcept
{
  getGlobalMemCounter()->deleteCalled(p);
  free(p);
}
#  endif // TEST_COMPILER(GCC)

void* operator new[](cuda::std::size_t s)
{
  getGlobalMemCounter()->newArrayCalled(s);
  return operator new(s);
}

void operator delete[](void* p) noexcept
{
  getGlobalMemCounter()->deleteArrayCalled(p);
  operator delete(p);
}

#  if TEST_COMPILER(GCC)
void operator delete[](void* p, cuda::std::size_t) noexcept
{
  getGlobalMemCounter()->deleteArrayCalled(p);
  operator delete(p);
}
#  endif // TEST_COMPILER(GCC)

#  if _LIBCUDACXX_HAS_ALIGNED_ALLOCATION()
#    if defined(_WIN32)
#      define USE_ALIGNED_ALLOC
#    endif

void* operator new(cuda::std::size_t s, cuda::std::align_val_t av)
{
  const cuda::std::size_t a = static_cast<cuda::std::size_t>(av);
  getGlobalMemCounter()->alignedNewCalled(s, a);
  void* ret;
#    ifdef USE_ALIGNED_ALLOC
  ret = _aligned_malloc(s, a);
#    else
  posix_memalign(&ret, a, s);
#    endif
  if (ret == nullptr)
  {
    cuda::std::__throw_bad_alloc();
  }
  return ret;
}

void operator delete(void* p, cuda::std::align_val_t av) noexcept
{
  const cuda::std::size_t a = static_cast<cuda::std::size_t>(av);
  getGlobalMemCounter()->alignedDeleteCalled(p, a);
  if (p)
  {
#    ifdef USE_ALIGNED_ALLOC
    ::_aligned_free(p);
#    else
    ::free(p);
#    endif
  }
}

void* operator new[](cuda::std::size_t s, cuda::std::align_val_t av)
{
  const cuda::std::size_t a = static_cast<cuda::std::size_t>(av);
  getGlobalMemCounter()->alignedNewArrayCalled(s, a);
  return operator new(s, av);
}

void operator delete[](void* p, cuda::std::align_val_t av) noexcept
{
  const cuda::std::size_t a = static_cast<cuda::std::size_t>(av);
  getGlobalMemCounter()->alignedDeleteArrayCalled(p, a);
  return operator delete(p, av);
}

#  endif // _LIBCUDACXX_HAS_ALIGNED_ALLOCATION()

#endif // DISABLE_NEW_COUNT

struct DisableAllocationGuard
{
  __host__ __device__ explicit DisableAllocationGuard(bool disable = true)
      : m_disabled(disable)
  {
    // Don't re-disable if already disabled.
    if (globalMemCounter.disable_allocations == true)
    {
      m_disabled = false;
    }
    if (m_disabled)
    {
      globalMemCounter.disableAllocations();
    }
  }

  __host__ __device__ void release()
  {
    if (m_disabled)
    {
      globalMemCounter.enableAllocations();
    }
    m_disabled = false;
  }

  __host__ __device__ ~DisableAllocationGuard()
  {
    release();
  }

private:
  bool m_disabled;

  __host__ __device__ DisableAllocationGuard(DisableAllocationGuard const&)            = delete;
  __host__ __device__ DisableAllocationGuard& operator=(DisableAllocationGuard const&) = delete;
};

struct RequireAllocationGuard
{
  __host__ __device__ explicit RequireAllocationGuard(cuda::std::size_t RequireAtLeast = 1)
      : m_req_alloc(RequireAtLeast)
      , m_new_count_on_init(globalMemCounter.new_called)
      , m_outstanding_new_on_init(globalMemCounter.outstanding_new)
      , m_exactly(false)
  {}

  __host__ __device__ void requireAtLeast(cuda::std::size_t N)
  {
    m_req_alloc = N;
    m_exactly   = false;
  }
  __host__ __device__ void requireExactly(cuda::std::size_t N)
  {
    m_req_alloc = N;
    m_exactly   = true;
  }

  __host__ __device__ ~RequireAllocationGuard()
  {
    assert(globalMemCounter.checkOutstandingNewEq(static_cast<int>(m_outstanding_new_on_init)));
    cuda::std::size_t Expect = m_new_count_on_init + m_req_alloc;
    assert(globalMemCounter.checkNewCalledEq(static_cast<int>(Expect))
           || (!m_exactly && globalMemCounter.checkNewCalledGreaterThan(static_cast<int>(Expect))));
  }

private:
  cuda::std::size_t m_req_alloc;
  const cuda::std::size_t m_new_count_on_init;
  const cuda::std::size_t m_outstanding_new_on_init;
  bool m_exactly;
  __host__ __device__ RequireAllocationGuard(RequireAllocationGuard const&)            = delete;
  __host__ __device__ RequireAllocationGuard& operator=(RequireAllocationGuard const&) = delete;
};

#endif /* COUNT_NEW_H */