// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project

#include <Kokkos_Macros.hpp>
#ifdef KOKKOS_ENABLE_EXPERIMENTAL_CXX20_MODULES
import kokkos.core;
import kokkos.core_impl;
#else
#include <Kokkos_Core.hpp>
#endif

#include <gtest/gtest.h>

template <class MemorySpace, class ExecutionSpace>
struct TestMemoryAccessViolation {
  Kokkos::Impl::SpaceAwareAccessor<MemorySpace, Kokkos::default_accessor<int>>
      acc;

  KOKKOS_FUNCTION decltype(auto) bad_access() const {
    return acc.access(nullptr, 0);
  }

  KOKKOS_FUNCTION void operator()(int) const { ++bad_access(); }

  TestMemoryAccessViolation(ExecutionSpace const& s,
                            std::string const& matcher) {
    constexpr bool accessible_from_execution_space = Kokkos::SpaceAccessibility<
        /*AccessSpace=*/ExecutionSpace,
        /*MemorySpace=*/MemorySpace>::accessible;
    EXPECT_FALSE(accessible_from_execution_space);
    EXPECT_DEATH(
        {
          Kokkos::parallel_for(Kokkos::RangePolicy<ExecutionSpace>(s, 0, 1),
                               *this);
          Kokkos::fence();
        },
        matcher);
  }
};

template <class MemorySpace, class ExecutionSpace>
void test_memory_access_violation(ExecutionSpace const& s,
                                  std::string const& m) {
  TestMemoryAccessViolation<MemorySpace, ExecutionSpace>(s, m);
}

template <class ExecutionSpace>
void test_memory_access_violations_from_host() {
  using memory_space_t = typename ExecutionSpace::memory_space;
  using exec_space_t   = Kokkos::DefaultHostExecutionSpace;
  const exec_space_t exec_space{};
  std::string const message =
      "Kokkos::SpaceAwareAccessor ERROR: attempt to access inaccessible memory "
      "space";
  test_memory_access_violation<memory_space_t, exec_space_t>(exec_space,
                                                             message);
}

template <class ExecutionSpace>
void test_memory_access_violations_from_device() {
  using memory_space_t = Kokkos::HostSpace;
  using exec_space_t   = ExecutionSpace;
  const exec_space_t exec_space{};
  // The invalid access is detected directly and we can't capture the error
  // message
#ifdef KOKKOS_ENABLE_SYCL
  std::string const message = ".*";
#else
  std::string const message =
      "Kokkos::SpaceAwareAccessor ERROR: attempt to access inaccessible memory "
      "space";
#endif
  test_memory_access_violation<memory_space_t, exec_space_t>(exec_space,
                                                             message);
}

TEST(TEST_CATEGORY_DEATH,
     mdspan_space_aware_accessor_invalid_access_from_host) {
  ::testing::FLAGS_gtest_death_test_style = "threadsafe";

  using ExecutionSpace = TEST_EXECSPACE;

  if (Kokkos::SpaceAccessibility<
          /*AccessSpace=*/Kokkos::HostSpace,
          /*MemorySpace=*/typename ExecutionSpace::memory_space>::accessible) {
    GTEST_SKIP() << "skipping since no memory access violation would occur";
  }

  test_memory_access_violations_from_host<ExecutionSpace>();
}

TEST(TEST_CATEGORY_DEATH,
     mdspan_space_aware_accessor_invalid_access_from_device) {
  ::testing::FLAGS_gtest_death_test_style = "threadsafe";

  using ExecutionSpace = TEST_EXECSPACE;

  if (Kokkos::SpaceAccessibility<
          /*AccessSpace=*/ExecutionSpace,
          /*MemorySpace=*/Kokkos::HostSpace>::accessible) {
    GTEST_SKIP() << "skipping since no memory access violation would occur";
  }

#if defined(KOKKOS_ENABLE_SYCL) && defined(NDEBUG)  // FIXME_SYCL
  if (std::is_same_v<ExecutionSpace, Kokkos::SYCL>) {
    GTEST_SKIP() << "skipping SYCL device-side abort does not work when NDEBUG "
                    "is defined";
  }
#endif
#if defined(KOKKOS_ENABLE_OPENACC)  // FIXME_OPENACC
  if (std::is_same<ExecutionSpace, Kokkos::Experimental::OpenACC>::value) {
    GTEST_SKIP() << "skipping because OpenACC backend is currently not "
                    "able to abort from the device";
  }
#endif

  test_memory_access_violations_from_device<ExecutionSpace>();
}
