// Compiler for PHP (aka KPHP)
// Copyright (c) 2020 LLC «V Kontakte»
// Distributed under the GPL v3 License, see LICENSE.notice.txt

#include "runtime/allocator.h"

#include <cstdlib>
#include <cstring>
#include <dlfcn.h>
#include <unistd.h>

#include "common/algorithms/find.h"
#include "common/containers/final_action.h"
#include "common/fast-backtrace.h"
#include "common/macos-ports.h"
#include "common/wrappers/likely.h"

#include "runtime/critical_section.h"
#include "runtime/kphp-backtrace.h"
#include "runtime/memory_resource_impl//dealer.h"
#include "runtime/php_assert.h"
#include "server/server-log.h"

namespace dl {
namespace {

memory_resource::Dealer &get_memory_dealer() noexcept {
  static memory_resource::Dealer dealer;
  return dealer;
}

} // namespace

bool script_allocator_enabled = false;
long long query_num = 0;

memory_resource::unsynchronized_pool_resource &get_default_script_allocator() noexcept {
  php_assert(script_allocator_enabled);
  php_assert(get_memory_dealer().is_default_allocator_used());
  return get_memory_dealer().current_script_resource();
}

void set_current_script_allocator(memory_resource::unsynchronized_pool_resource &resource, bool force_enable) noexcept {
  get_memory_dealer().set_current_script_resource(resource);
  if (force_enable) {
    php_assert(!script_allocator_enabled);
    script_allocator_enabled = true;
  }
}

void restore_default_script_allocator(bool force_disable) noexcept {
  get_memory_dealer().restore_default_script_resource();
  if (force_disable) {
    php_assert(script_allocator_enabled);
    script_allocator_enabled = false;
  }
}

const memory_resource::MemoryStats &get_script_memory_stats() noexcept {
  return get_memory_dealer().current_script_resource().get_memory_stats();
}

size_t get_heap_memory_used() noexcept {
  return get_memory_dealer().get_heap_resource().memory_used();
}

void global_init_script_allocator() noexcept {
  auto &dealer = get_memory_dealer();
  php_assert(dealer.heap_script_resource_replacer());
  php_assert(!is_malloc_replaced());
  php_assert(!script_allocator_enabled);
  php_assert(!query_num);

  CriticalSectionGuard lock;
  dealer.drop_replacer();
  query_num++;
}

void init_script_allocator(void *buffer, size_t script_mem_size, size_t oom_handling_mem_size) noexcept {
  auto &dealer = get_memory_dealer();
  php_assert(!dealer.heap_script_resource_replacer());
  php_assert(dealer.is_default_allocator_used());
  php_assert(!is_malloc_replaced());

  CriticalSectionGuard lock;
  dealer.current_script_resource().init(buffer, script_mem_size, oom_handling_mem_size);
  script_allocator_enabled = true;
  query_num++;
}

void free_script_allocator() noexcept {
  auto &dealer = get_memory_dealer();
  php_assert(!dealer.heap_script_resource_replacer());
  php_assert(dealer.is_default_allocator_used());
  php_assert(!is_malloc_replaced());

  script_allocator_enabled = false;
}

void *allocate(size_t size) noexcept {
  php_assert(size);
  auto &dealer = get_memory_dealer();
  if (auto *heap_replacer = dealer.heap_script_resource_replacer()) {
    return heap_replacer->allocate(size);
  }
  if (unlikely(!script_allocator_enabled)) {
    php_critical_error("Trying to call allocate for non runned script, n = %zu", size);
    return nullptr;
  }

  return dealer.current_script_resource().allocate(size);
}

void *allocate0(size_t size) noexcept {
  php_assert(size);
  auto &dealer = get_memory_dealer();
  if (auto *heap_replacer = dealer.heap_script_resource_replacer()) {
    return heap_replacer->allocate0(size);
  }
  if (unlikely(!script_allocator_enabled)) {
    php_critical_error("Trying to call allocate0 for non runned script, n = %zu", size);
    return nullptr;
  }

  return dealer.current_script_resource().allocate0(size);
}

void *reallocate(void *mem, size_t new_size, size_t old_size) noexcept {
  php_assert(new_size > old_size);
  auto &dealer = get_memory_dealer();
  if (auto *heap_replacer = dealer.heap_script_resource_replacer()) {
    return heap_replacer->reallocate(mem, new_size, old_size);
  }
  if (unlikely(!script_allocator_enabled)) {
    php_critical_error("Trying to call reallocate for non runned script, p = %p, new_size = %zu, old_size = %zu", mem, new_size, old_size);
    return mem;
  }

  return dealer.current_script_resource().reallocate(mem, new_size, old_size);
}

void deallocate(void *mem, size_t size) noexcept {
  php_assert(size);
  auto &dealer = get_memory_dealer();
  if (auto *heap_replacer = dealer.heap_script_resource_replacer()) {
    return heap_replacer->deallocate(mem, size);
  }

  if (script_allocator_enabled) {
    dealer.current_script_resource().deallocate(mem, size);
  }
}

void *heap_allocate(size_t size) noexcept {
  php_assert(!query_num || !is_malloc_replaced());
  return get_memory_dealer().get_heap_resource().allocate(size);
}

void *heap_reallocate(void *mem, size_t new_size, size_t old_size) noexcept {
  return get_memory_dealer().get_heap_resource().reallocate(mem, new_size, old_size);
}

void heap_deallocate(void *mem, size_t size) noexcept {
  return get_memory_dealer().get_heap_resource().deallocate(mem, size);
}

namespace {

// guaranteed alignment of dl::allocate
constexpr size_t MALLOC_REPLACER_SIZE_OFFSET = 8;
constexpr size_t MALLOC_REPLACER_MAX_ALLOC = 0xFFFFFF00;

} // namespace

void *script_allocator_malloc(size_t size) noexcept {
  static_assert(sizeof(size_t) <= MALLOC_REPLACER_SIZE_OFFSET, "small size offset");
  if (size > MALLOC_REPLACER_MAX_ALLOC - MALLOC_REPLACER_SIZE_OFFSET) {
    php_warning("attempt to allocate too much memory: %lu", size);
    return nullptr;
  }
  const size_t real_allocate = size + MALLOC_REPLACER_SIZE_OFFSET;
  void *mem = allocate(real_allocate);

  if (unlikely(!mem)) {
    php_warning("not enough memory to continue: %lu", size);
    return mem;
  }
  *static_cast<size_t *>(mem) = real_allocate;
  return static_cast<char *>(mem) + MALLOC_REPLACER_SIZE_OFFSET;
}

void *script_allocator_calloc(size_t nmemb, size_t size) noexcept {
  void *res = script_allocator_malloc(nmemb * size);
  if (unlikely(res == nullptr)) {
    return nullptr;
  }
  return memset(res, 0, nmemb * size);
}

void *script_allocator_realloc(void *p, size_t new_size) noexcept {
  if (p == nullptr) {
    return script_allocator_malloc(new_size);
  }

  if (new_size == 0) {
    script_allocator_free(p);
    return nullptr;
  }

  void *real_p = static_cast<char *>(p) - sizeof(size_t);
  const size_t old_size = *static_cast<size_t *>(real_p);

  // TODO reuse reallocate method
  void *new_p = script_allocator_malloc(new_size);
  if (likely(new_p != nullptr)) {
    memcpy(new_p, p, std::min(new_size, old_size));
    dl::deallocate(real_p, old_size);
  }
  return new_p;
}

char *script_allocator_strdup(const char *str) noexcept {
  const size_t len = strlen(str) + 1;
  char *res = static_cast<char *>(script_allocator_malloc(len));
  if (unlikely(res == nullptr)) {
    return nullptr;
  }
  memcpy(res, str, len);
  return res;
}

void script_allocator_free(void *mem) noexcept {
  if (mem) {
    mem = static_cast<char *>(mem) - MALLOC_REPLACER_SIZE_OFFSET;
    deallocate(mem, *static_cast<size_t *>(mem));
  }
}

namespace {

class MallocStateHolder : vk::not_copyable {
public:
  static MallocStateHolder &get() noexcept {
    static MallocStateHolder state;
    return state;
  }

  void replace_malloc(bool is_malloc_replaced_before) noexcept {
    php_assert(is_malloc_replaced_before == is_malloc_replaced_);

    if (!is_malloc_replaced_before) {
      last_malloc_replacement_backtrace_size_ = fast_backtrace(last_malloc_replacement_backtrace_.data(), last_malloc_replacement_backtrace_.size());
    }

    is_malloc_replaced_ = !is_malloc_replaced_;
  }

  bool is_malloc_replaced() const noexcept {
    return is_malloc_replaced_;
  }

  std::pair<void *const *, int> get_last_malloc_replacement_backtrace() const noexcept {
    return {last_malloc_replacement_backtrace_.data(), last_malloc_replacement_backtrace_size_};
  }

private:
  MallocStateHolder() = default;
  bool is_malloc_replaced_{false};
  std::array<void *, 128> last_malloc_replacement_backtrace_{};
  int last_malloc_replacement_backtrace_size_{0};
};

} // namespace

bool is_malloc_replaced() noexcept {
  return MallocStateHolder::get().is_malloc_replaced();
}

void replace_malloc_with_script_allocator() noexcept {
  MallocStateHolder::get().replace_malloc(false);
}

void rollback_malloc_replacement() noexcept {
  MallocStateHolder::get().replace_malloc(true);
}

void write_last_malloc_replacement_stacktrace(char *buf, size_t buf_size) noexcept {
  auto malloc_replacement_rollback = temporary_rollback_malloc_replacement();
  auto [raw_backtrace, backtrace_size] = MallocStateHolder::get().get_last_malloc_replacement_backtrace();
  parse_kphp_backtrace(buf, buf_size, raw_backtrace, backtrace_size);
}

MemoryReplacementGuard::MemoryReplacementGuard(memory_resource::unsynchronized_pool_resource &memory_resource, bool force_enable_disable) : force_enable_disable_(force_enable_disable) {
  dl::enter_critical_section();
  dl::set_current_script_allocator(memory_resource, force_enable_disable_);
}

MemoryReplacementGuard::~MemoryReplacementGuard() {
  dl::restore_default_script_allocator(force_enable_disable_);
  dl::leave_critical_section();
}

} // namespace dl

// sanitizers aren't happy with custom realization of malloc-like functions
#if !ASAN_ENABLED and !defined(__APPLE__)

extern "C" void *__libc_malloc(size_t size);
extern "C" void *__libc_calloc(size_t nmemb, size_t size);
extern "C" void *__libc_realloc(void *mem, size_t new_size);
extern "C" void __libc_free(void *mem);
extern "C" void *__libc_memalign(size_t alignment, size_t size);

void *malloc(size_t size) {
  return dl::is_malloc_replaced() ? dl::script_allocator_malloc(size) : __libc_malloc(size);
}

void *calloc(size_t nmemb, size_t size) {
  return dl::is_malloc_replaced() ? dl::script_allocator_calloc(nmemb, size) : __libc_calloc(nmemb, size);
}

void *realloc(void *mem, size_t new_size) {
  return dl::is_malloc_replaced() ? dl::script_allocator_realloc(mem, new_size) : __libc_realloc(mem, new_size);
}

void free(void *mem) {
  return dl::is_malloc_replaced() ? dl::script_allocator_free(mem) : __libc_free(mem);
}

void *aligned_alloc(size_t alignment, size_t size) {
  if (dl::is_malloc_replaced()) {
    // script allocator gives addresses aligned to 8 bytes
    php_assert(alignment <= 8);
    return dl::script_allocator_malloc(size);
  }
  return __libc_memalign(alignment, size);
}

void *memalign(size_t alignment, size_t size) {
  return aligned_alloc(alignment, size);
}

// put realization posix_memalign function here if you want to use it under script allocator

#endif

// replace global operators new and delete for linked C++ code
void *operator new(size_t size) {
  auto *res = std::malloc(size);
  if (!res) {
    php_critical_error("nullptr from malloc");
  }
  return res;
}

void *operator new(size_t size, const std::nothrow_t &) noexcept {
  return std::malloc(size);
}

void *operator new[](size_t size) {
  return operator new(size);
}

void *operator new[](size_t size, const std::nothrow_t &) noexcept {
  return operator new(size, std::nothrow);
}

void operator delete(void *mem) noexcept {
  return std::free(mem);
}

void operator delete(void *mem, const std::nothrow_t &) noexcept {
  return std::free(mem);
}

void operator delete[](void *mem) noexcept {
  return std::free(mem);
}

void operator delete[](void *mem, const std::nothrow_t &) noexcept {
  return std::free(mem);
}

void operator delete(void *mem, size_t) noexcept {
  return std::free(mem);
}

void operator delete[](void *mem, size_t) noexcept {
  return std::free(mem);
}
