#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <linux/membarrier.h>
#include <sys/syscall.h>
#include <sys/utsname.h>

#if __mips__
#  include <asm/cachectl.h>
#endif

#include "log.h"

namespace wvoec3 {

bool supports_membarrier_syscall() {
  // Check kernel version supports membarrier(2); version 4.16 is required.
  static constexpr int kRequiredMajor = 4;
  static constexpr int kRequiredMinor = 16;
  struct utsname uts;
  int major = 0, minor = 0;
  if (uname(&uts) != 0 || strcmp(uts.sysname, "Linux") != 0 ||
      sscanf(uts.release, "%d.%d", &major, &minor) != 2) {
    LOGE("Failed to get kernel version");
    return false;
  }
  if (major < kRequiredMajor ||
      (major == kRequiredMajor && minor < kRequiredMinor)) {
    LOGW(
        "Kernel version %d, %d < required %d, %d. Some membarrier_function "
        "calls may not be available. Although this is not fatal, missing "
        "membarrier support may cause run-time errors on some platforms.",
        major, minor, kRequiredMajor, kRequiredMinor);
    return false;
  }
  return true;
}

#if defined(__NR_membarrier)

// membarrier function is used to complete flush of CPU instruction pipelines.
// Reference:
// * membarrier(2)
// * art::jit::JitMemoryRegion::CommitCode
int membarrier_function(int command) {
  static bool supports_membarrier = supports_membarrier_syscall();
  if (supports_membarrier) {
    return syscall(__NR_membarrier, command, 0);
  }
  errno = ENOSYS;
  return -1;
}

#else  // __NR_membarrier
#  pragma message \
      "(info): __NR_membarrier unavailable; membarrier_function will be no-op."

int membarrier_function(int) {
  errno = ENOSYS;
  return -1;
}

#endif  // __NR_membarrier

// The Cache Flush function is very processor dependent, and is used by the
// Level3 code using platform-specific flags.
// Linked in separately from the rest of the build, since this code isn't
// obfuscated and relies on system flags to function properly.
void clear_cache_function(void *page, size_t len) {
  // Note on cross platform support.  If __has_builtin is not defined as a
  // preprocessor function, we cannot use
  // "#if defined(__has_builtin) && __has_builtin(..)".
  // So, instead, we will define USED_BUILTIN_CLEAR_CACHE if both conditions
  // are true, and use "#ifndef  USED_BUILTIN_CLEAR_CACHE" instead of #else.
#ifdef __has_builtin
#  if __has_builtin(__builtin___clear_cache)
#    pragma message "(info): clear_cache_function is using __builtin___clear_cache."
#      define USED_BUILTIN_CLEAR_CACHE
  char *begin = static_cast<char *>(page);
  char *end = begin + len;
  __builtin___clear_cache(begin, end);
#  endif
#endif

#ifndef USED_BUILTIN_CLEAR_CACHE
#  if defined(__arm__)
#    pragma message "(info): clear_cache_function is using arm asm."
#      define USED_ARM_ASM_CLEAR_CACHE
  // ARM Cache Flush System Call:
  char *begin = static_cast<char *>(page);
  char *end = begin + len;
  const int syscall = 0xf0002;
  __asm __volatile(
      "push      {r0, r1, r2, r7}\n"
      "mov       r0, %0\n"
      "mov       r1, %1\n"
      "mov       r7, %2\n"
      "mov       r2, #0x0\n"
      "svc       0x00000000\n"
      "pop       {r0, r1, r2, r7}\n"
      :
      : "r"(begin), "r"(end), "r"(syscall)
      : "r0", "r1", "r7");
#  elif defined(__aarch64__)
#    pragma message "(info): clear_cache_function is using arm64 asm."
#      define USED_ARM_ASM_CLEAR_CACHE
  uint64_t begin = static_cast<uint64_t>(reinterpret_cast<uintptr_t>(page));
  uint64_t end = begin + len;
  register uint64_t ctr_el0;
  __asm __volatile("mrs %0, ctr_el0" : "=r"(ctr_el0));
  // register CTR_EL0 [19:16] contains dcache line size
  const size_t dcache_line_size = 4 << ((ctr_el0 >> 16) & 0xF);
  register uint64_t addr;
  for (addr = begin; addr < end; addr += dcache_line_size) {
    __asm __volatile("dc cvau, %0" ::"r"(addr));
  }
  __asm __volatile("dsb ish");
  //  register CTR_EL0 [3:0] contains icache line size
  const size_t icache_line_size = 4 << (ctr_el0 & 0xF);
  for (addr = begin; addr < end; addr += icache_line_size) {
    __asm __volatile("ic ivau, %0" ::"r"(addr));
  }
  __asm __volatile("dsb ish");
  __asm __volatile("isb sy");
#  elif __mips__
#    pragma message "(info): clear_cache_function is using mips asm."
  int result = syscall(__NR_cacheflush, page, len, ICACHE);
  if (result) {
    fprintf(stderr, "cacheflush failed!: errno=%d %s\n", errno,
            strerror(errno));
    exit(-1);  // TODO(fredgc): figure out more graceful error handling.
  }
#  else
#    pragma message "(info): clear_cache_function is not doing anything."
  // TODO(fredgc): silence warning about unused variables.
#  endif
#endif

#if defined(__arm__) || defined(__aarch64__)
#  if !defined(USED_BUILTIN_CLEAR_CACHE) && !defined(USED_ARM_ASM_CLEAR_CACHE)
#    error "clear cache function unavailable, which is required for ARM."
#  endif
#  pragma message "(info): inserting membarrier_function calls."
  membarrier_function(MEMBARRIER_CMD_REGISTER_PRIVATE_EXPEDITED_SYNC_CORE);
  membarrier_function(MEMBARRIER_CMD_PRIVATE_EXPEDITED_SYNC_CORE);
#endif  // defined(__arm__) || defined(__aarch64__)
}

}  // namespace wvoec3