/* +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
Copyright (c) 2008-2017, Petr Kobalicek

This software is provided 'as-is', without any express or implied
warranty. In no event will the authors be held liable for any damages
arising from the use of this software.

Permission is granted to anyone to use this software for any purpose,
including commercial applications, and to alter it and redistribute it
freely, subject to the following restrictions:

1. The origin of this software must not be misrepresented; you must not
   claim that you wrote the original software. If you use this software
   in a product, an acknowledgment in the product documentation would be
   appreciated but is not required.
2. Altered source versions must be plainly marked as such, and must not be
   misrepresented as being the original software.
3. This notice may not be removed or altered from any source distribution.
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ */
#ifdef __PLUMED_HAS_ASMJIT
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wpedantic"
// [AsmJit]
// Complete x86/x64 JIT and Remote Assembler for C++.
//
// [License]
// Zlib - See LICENSE.md file in the package.

// [Export]
#define ASMJIT_EXPORTS

// [Dependencies]
#include "./assembler.h"
#include "./utils.h"
#include "./vmem.h"

// [Api-Begin]
#include "./asmjit_apibegin.h"

namespace PLMD {
namespace asmjit {

// ============================================================================
// [asmjit::ErrorHandler]
// ============================================================================

ErrorHandler::ErrorHandler() noexcept {}
ErrorHandler::~ErrorHandler() noexcept {}

// ============================================================================
// [asmjit::CodeHolder - Utilities]
// ============================================================================

static void CodeHolder_setGlobalOption(CodeHolder* self, uint32_t clear, uint32_t add) noexcept {
  // Modify global options of `CodeHolder` itself.
  self->_globalOptions = (self->_globalOptions & ~clear) | add;

  // Modify all global options of all `CodeEmitter`s attached.
  CodeEmitter* emitter = self->_emitters;
  while (emitter) {
    emitter->_globalOptions = (emitter->_globalOptions & ~clear) | add;
    emitter = emitter->_nextEmitter;
  }
}

static void CodeHolder_resetInternal(CodeHolder* self, bool releaseMemory) noexcept {
  // Detach all `CodeEmitter`s.
  while (self->_emitters)
    self->detach(self->_emitters);

  // Reset everything into its construction state.
  self->_codeInfo.reset();
  self->_globalHints = 0;
  self->_globalOptions = 0;
  self->_logger = nullptr;
  self->_errorHandler = nullptr;

  self->_unresolvedLabelsCount = 0;
  self->_trampolinesSize = 0;

  // Reset all sections.
  size_t numSections = self->_sections.getLength();
  for (size_t i = 0; i < numSections; i++) {
    SectionEntry* section = self->_sections[i];
    if (section->_buffer.hasData() && !section->_buffer.isExternal())
      Internal::releaseMemory(section->_buffer._data);
    section->_buffer._data = nullptr;
    section->_buffer._capacity = 0;
  }

  // Reset zone allocator and all containers using it.
  ZoneHeap* heap = &self->_baseHeap;

  self->_namedLabels.reset(heap);
  self->_relocations.reset();
  self->_labels.reset();
  self->_sections.reset();

  heap->reset(&self->_baseZone);
  self->_baseZone.reset(releaseMemory);
}

// ============================================================================
// [asmjit::CodeHolder - Construction / Destruction]
// ============================================================================

CodeHolder::CodeHolder() noexcept
  : _codeInfo(),
    _globalHints(0),
    _globalOptions(0),
    _emitters(nullptr),
    _cgAsm(nullptr),
    _logger(nullptr),
    _errorHandler(nullptr),
    _unresolvedLabelsCount(0),
    _trampolinesSize(0),
    _baseZone(16384 - Zone::kZoneOverhead),
    _dataZone(16384 - Zone::kZoneOverhead),
    _baseHeap(&_baseZone),
    _namedLabels(&_baseHeap) {}

CodeHolder::~CodeHolder() noexcept {
  CodeHolder_resetInternal(this, true);
}

// ============================================================================
// [asmjit::CodeHolder - Init / Reset]
// ============================================================================

Error CodeHolder::init(const CodeInfo& info) noexcept {
  // Cannot reinitialize if it's locked or there is one or more CodeEmitter
  // attached.
  if (isInitialized())
    return DebugUtils::errored(kErrorAlreadyInitialized);

  // If we are just initializing there should be no emitters attached).
  ASMJIT_ASSERT(_emitters == nullptr);

  // Create the default section and insert it to the `_sections` array.
  Error err = _sections.willGrow(&_baseHeap);
  if (err == kErrorOk) {
    SectionEntry* se = _baseZone.allocZeroedT<SectionEntry>();
    if (ASMJIT_LIKELY(se)) {
      se->_flags = SectionEntry::kFlagExec | SectionEntry::kFlagConst;
      se->_setDefaultName('.', 't', 'e', 'x', 't');
      _sections.appendUnsafe(se);
    }
    else {
      err = DebugUtils::errored(kErrorNoHeapMemory);
    }
  }

  if (ASMJIT_UNLIKELY(err)) {
    _baseZone.reset(false);
    return err;
  }
  else {
    _codeInfo = info;
    return kErrorOk;
  }
}

void CodeHolder::reset(bool releaseMemory) noexcept {
  CodeHolder_resetInternal(this, releaseMemory);
}

// ============================================================================
// [asmjit::CodeHolder - Attach / Detach]
// ============================================================================

Error CodeHolder::attach(CodeEmitter* emitter) noexcept {
  // Catch a possible misuse of the API.
  if (!emitter)
    return DebugUtils::errored(kErrorInvalidArgument);

  uint32_t type = emitter->getType();
  if (type == CodeEmitter::kTypeNone || type >= CodeEmitter::kTypeCount)
    return DebugUtils::errored(kErrorInvalidState);

  // This is suspicious, but don't fail if `emitter` matches.
  if (emitter->_code != nullptr) {
    if (emitter->_code == this) return kErrorOk;
    return DebugUtils::errored(kErrorInvalidState);
  }

  // Special case - attach `Assembler`.
  CodeEmitter** pSlot = nullptr;
  if (type == CodeEmitter::kTypeAssembler) {
    if (_cgAsm)
      return DebugUtils::errored(kErrorSlotOccupied);
    pSlot = reinterpret_cast<CodeEmitter**>(&_cgAsm);
  }

  Error err = emitter->onAttach(this);
  if (err != kErrorOk) return err;

  // Add to a single-linked list of `CodeEmitter`s.
  emitter->_nextEmitter = _emitters;
  _emitters = emitter;
  if (pSlot) *pSlot = emitter;

  // Establish the connection.
  emitter->_code = this;
  return kErrorOk;
}

Error CodeHolder::detach(CodeEmitter* emitter) noexcept {
  if (!emitter)
    return DebugUtils::errored(kErrorInvalidArgument);

  if (emitter->_code != this)
    return DebugUtils::errored(kErrorInvalidState);

  uint32_t type = emitter->getType();
  Error err = kErrorOk;

  // NOTE: We always detach if we were asked to, if error happens during
  // `emitter->onDetach()` we just propagate it, but the CodeEmitter will
  // be detached.
  if (!emitter->_destroyed) {
    if (type == CodeEmitter::kTypeAssembler)
      static_cast<Assembler*>(emitter)->sync();
    err = emitter->onDetach(this);
  }

  // Special case - detach `Assembler`.
  if (type == CodeEmitter::kTypeAssembler)
    _cgAsm = nullptr;

  // Remove from a single-linked list of `CodeEmitter`s.
  CodeEmitter** pPrev = &_emitters;
  for (;;) {
    ASMJIT_ASSERT(*pPrev != nullptr);
    CodeEmitter* cur = *pPrev;

    if (cur == emitter) {
      *pPrev = emitter->_nextEmitter;
      break;
    }

    pPrev = &cur->_nextEmitter;
  }

  emitter->_code = nullptr;
  emitter->_nextEmitter = nullptr;

  return err;
}

// ============================================================================
// [asmjit::CodeHolder - Sync]
// ============================================================================

void CodeHolder::sync() noexcept {
  if (_cgAsm) _cgAsm->sync();
}

// ============================================================================
// [asmjit::CodeHolder - Result Information]
// ============================================================================

size_t CodeHolder::getCodeSize() const noexcept {
  // Reflect all changes first.
  const_cast<CodeHolder*>(this)->sync();

  // TODO: Support sections.
  return _sections[0]->_buffer._length + getTrampolinesSize();
}

// ============================================================================
// [asmjit::CodeHolder - Logging & Error Handling]
// ============================================================================

#if !defined(ASMJIT_DISABLE_LOGGING)
void CodeHolder::setLogger(Logger* logger) noexcept {
  uint32_t opt = 0;
  if (logger) opt = CodeEmitter::kOptionLoggingEnabled;

  _logger = logger;
  CodeHolder_setGlobalOption(this, CodeEmitter::kOptionLoggingEnabled, opt);
}
#endif // !ASMJIT_DISABLE_LOGGING

Error CodeHolder::setErrorHandler(ErrorHandler* handler) noexcept {
  _errorHandler = handler;
  return kErrorOk;
}

// ============================================================================
// [asmjit::CodeHolder - Sections]
// ============================================================================

static Error CodeHolder_reserveInternal(CodeHolder* self, CodeBuffer* cb, size_t n) noexcept {
  uint8_t* oldData = cb->_data;
  uint8_t* newData;

  if (oldData && !cb->isExternal())
    newData = static_cast<uint8_t*>(Internal::reallocMemory(oldData, n));
  else
    newData = static_cast<uint8_t*>(Internal::allocMemory(n));

  if (ASMJIT_UNLIKELY(!newData))
    return DebugUtils::errored(kErrorNoHeapMemory);

  cb->_data = newData;
  cb->_capacity = n;

  // Update the `Assembler` pointers if attached. Maybe we should introduce an
  // event for this, but since only one Assembler can be attached at a time it
  // should not matter how these pointers are updated.
  Assembler* a = self->_cgAsm;
  if (a && &a->_section->_buffer == cb) {
    size_t offset = a->getOffset();

    a->_bufferData = newData;
    a->_bufferEnd  = newData + n;
    a->_bufferPtr  = newData + offset;
  }

  return kErrorOk;
}

Error CodeHolder::growBuffer(CodeBuffer* cb, size_t n) noexcept {
  // This is most likely called by `Assembler` so `sync()` shouldn't be needed,
  // however, if this is called by the user and the currently attached Assembler
  // did generate some code we could lose that, so sync now and make sure the
  // section length is updated.
  if (_cgAsm) _cgAsm->sync();

  // Now the length of the section must be valid.
  size_t length = cb->getLength();
  if (ASMJIT_UNLIKELY(n > IntTraits<uintptr_t>::maxValue() - length))
    return DebugUtils::errored(kErrorNoHeapMemory);

  // We can now check if growing the buffer is really necessary. It's unlikely
  // that this function is called while there is still room for `n` bytes.
  size_t capacity = cb->getCapacity();
  size_t required = cb->getLength() + n;
  if (ASMJIT_UNLIKELY(required <= capacity)) return kErrorOk;

  if (cb->isFixedSize())
    return DebugUtils::errored(kErrorCodeTooLarge);

  if (capacity < 8096)
    capacity = 8096;
  else
    capacity += Globals::kAllocOverhead;

  do {
    size_t old = capacity;
    if (capacity < Globals::kAllocThreshold)
      capacity *= 2;
    else
      capacity += Globals::kAllocThreshold;

    if (capacity < Globals::kAllocThreshold)
      capacity *= 2;
    else
      capacity += Globals::kAllocThreshold;

    // Overflow.
    if (ASMJIT_UNLIKELY(old > capacity))
      return DebugUtils::errored(kErrorNoHeapMemory);
  } while (capacity - Globals::kAllocOverhead < required);

  return CodeHolder_reserveInternal(this, cb, capacity - Globals::kAllocOverhead);
}

Error CodeHolder::reserveBuffer(CodeBuffer* cb, size_t n) noexcept {
  size_t capacity = cb->getCapacity();
  if (n <= capacity) return kErrorOk;

  if (cb->isFixedSize())
    return DebugUtils::errored(kErrorCodeTooLarge);

  // We must sync, as mentioned in `growBuffer()` as well.
  if (_cgAsm) _cgAsm->sync();

  return CodeHolder_reserveInternal(this, cb, n);
}

// ============================================================================
// [asmjit::CodeHolder - Labels & Symbols]
// ============================================================================

namespace {

//! \internal
//!
//! Only used to lookup a label from `_namedLabels`.
class LabelByName {
public:
  ASMJIT_INLINE LabelByName(const char* name, size_t nameLength, uint32_t hVal) noexcept
    : name(name),
      nameLength(static_cast<uint32_t>(nameLength)),
      hVal(hVal) {}

  ASMJIT_INLINE bool matches(const LabelEntry* entry) const noexcept {
    return static_cast<uint32_t>(entry->getNameLength()) == nameLength &&
           ::memcmp(entry->getName(), name, nameLength) == 0;
  }

  const char* name;
  uint32_t nameLength;
  uint32_t hVal;
};

// Returns a hash of `name` and fixes `nameLength` if it's `Globals::kInvalidIndex`.
static uint32_t CodeHolder_hashNameAndFixLen(const char* name, size_t& nameLength) noexcept {
  uint32_t hVal = 0;
  if (nameLength == Globals::kInvalidIndex) {
    size_t i = 0;
    for (;;) {
      uint8_t c = static_cast<uint8_t>(name[i]);
      if (!c) break;
      hVal = Utils::hashRound(hVal, c);
      i++;
    }
    nameLength = i;
  }
  else {
    for (size_t i = 0; i < nameLength; i++) {
      uint8_t c = static_cast<uint8_t>(name[i]);
      if (ASMJIT_UNLIKELY(!c)) return DebugUtils::errored(kErrorInvalidLabelName);
      hVal = Utils::hashRound(hVal, c);
    }
  }
  return hVal;
}

} // anonymous namespace

LabelLink* CodeHolder::newLabelLink(LabelEntry* le, uint32_t sectionId, size_t offset, intptr_t rel) noexcept {
  LabelLink* link = _baseHeap.allocT<LabelLink>();
  if (ASMJIT_UNLIKELY(!link)) return nullptr;

  link->prev = le->_links;
  le->_links = link;

  link->sectionId = sectionId;
  link->relocId = RelocEntry::kInvalidId;
  link->offset = offset;
  link->rel = rel;

  _unresolvedLabelsCount++;
  return link;
}

Error CodeHolder::newLabelId(uint32_t& idOut) noexcept {
  idOut = 0;

  size_t index = _labels.getLength();
  if (ASMJIT_LIKELY(index >= Operand::kPackedIdCount))
    return DebugUtils::errored(kErrorLabelIndexOverflow);

  ASMJIT_PROPAGATE(_labels.willGrow(&_baseHeap));
  LabelEntry* le = _baseHeap.allocZeroedT<LabelEntry>();

  if (ASMJIT_UNLIKELY(!le))
    return DebugUtils::errored(kErrorNoHeapMemory);;

  uint32_t id = Operand::packId(static_cast<uint32_t>(index));
  le->_setId(id);
  le->_parentId = 0;
  le->_sectionId = SectionEntry::kInvalidId;
  le->_offset = 0;

  _labels.appendUnsafe(le);
  idOut = id;
  return kErrorOk;
}

Error CodeHolder::newNamedLabelId(uint32_t& idOut, const char* name, size_t nameLength, uint32_t type, uint32_t parentId) noexcept {
  idOut = 0;
  uint32_t hVal = CodeHolder_hashNameAndFixLen(name, nameLength);

  if (ASMJIT_UNLIKELY(nameLength == 0))
    return DebugUtils::errored(kErrorInvalidLabelName);

  if (ASMJIT_UNLIKELY(nameLength > Globals::kMaxLabelLength))
    return DebugUtils::errored(kErrorLabelNameTooLong);

  switch (type) {
    case Label::kTypeLocal:
      if (ASMJIT_UNLIKELY(Operand::unpackId(parentId) >= _labels.getLength()))
        return DebugUtils::errored(kErrorInvalidParentLabel);

      hVal ^= parentId;
      break;

    case Label::kTypeGlobal:
      if (ASMJIT_UNLIKELY(parentId != 0))
        return DebugUtils::errored(kErrorNonLocalLabelCantHaveParent);

      break;

    default:
      return DebugUtils::errored(kErrorInvalidArgument);
  }

  // Don't allow to insert duplicates. Local labels allow duplicates that have
  // different id, this is already accomplished by having a different hashes
  // between the same label names having different parent labels.
  LabelEntry* le = _namedLabels.get(LabelByName(name, nameLength, hVal));
  if (ASMJIT_UNLIKELY(le))
    return DebugUtils::errored(kErrorLabelAlreadyDefined);

  Error err = kErrorOk;
  size_t index = _labels.getLength();

  if (ASMJIT_UNLIKELY(index >= Operand::kPackedIdCount))
    return DebugUtils::errored(kErrorLabelIndexOverflow);

  ASMJIT_PROPAGATE(_labels.willGrow(&_baseHeap));
  le = _baseHeap.allocZeroedT<LabelEntry>();

  if (ASMJIT_UNLIKELY(!le))
    return  DebugUtils::errored(kErrorNoHeapMemory);

  uint32_t id = Operand::packId(static_cast<uint32_t>(index));
  le->_hVal = hVal;
  le->_setId(id);
  le->_type = static_cast<uint8_t>(type);
  le->_parentId = 0;
  le->_sectionId = SectionEntry::kInvalidId;
  le->_offset = 0;

  if (le->_name.mustEmbed(nameLength)) {
    le->_name.setEmbedded(name, nameLength);
  }
  else {
    char* nameExternal = static_cast<char*>(_dataZone.dup(name, nameLength, true));
    if (ASMJIT_UNLIKELY(!nameExternal))
      return DebugUtils::errored(kErrorNoHeapMemory);
    le->_name.setExternal(nameExternal, nameLength);
  }

  _labels.appendUnsafe(le);
  _namedLabels.put(le);

  idOut = id;
  return err;
}

uint32_t CodeHolder::getLabelIdByName(const char* name, size_t nameLength, uint32_t parentId) noexcept {
  uint32_t hVal = CodeHolder_hashNameAndFixLen(name, nameLength);
  if (ASMJIT_UNLIKELY(!nameLength)) return 0;

  LabelEntry* le = _namedLabels.get(LabelByName(name, nameLength, hVal));
  return le ? le->getId() : static_cast<uint32_t>(0);
}

// ============================================================================
// [asmjit::CodeEmitter - Relocations]
// ============================================================================

//! Encode MOD byte.
static ASMJIT_INLINE uint32_t x86EncodeMod(uint32_t m, uint32_t o, uint32_t rm) noexcept {
  return (m << 6) | (o << 3) | rm;
}

Error CodeHolder::newRelocEntry(RelocEntry** dst, uint32_t type, uint32_t size) noexcept {
  ASMJIT_PROPAGATE(_relocations.willGrow(&_baseHeap));

  size_t index = _relocations.getLength();
  if (ASMJIT_UNLIKELY(index > size_t(0xFFFFFFFFU)))
    return DebugUtils::errored(kErrorRelocIndexOverflow);

  RelocEntry* re = _baseHeap.allocZeroedT<RelocEntry>();
  if (ASMJIT_UNLIKELY(!re))
    return DebugUtils::errored(kErrorNoHeapMemory);

  re->_id = static_cast<uint32_t>(index);
  re->_type = static_cast<uint8_t>(type);
  re->_size = static_cast<uint8_t>(size);
  re->_sourceSectionId = SectionEntry::kInvalidId;
  re->_targetSectionId = SectionEntry::kInvalidId;
  _relocations.appendUnsafe(re);

  *dst = re;
  return kErrorOk;
}

// TODO: Support multiple sections, this only relocates the first.
// TODO: This should go to Runtime as it's responsible for relocating the
//       code, CodeHolder should just hold it.
size_t CodeHolder::relocate(void* _dst, uint64_t baseAddress) const noexcept {
  SectionEntry* section = _sections[0];
  ASMJIT_ASSERT(section != nullptr);

  uint8_t* dst = static_cast<uint8_t*>(_dst);
  if (baseAddress == Globals::kNoBaseAddress)
    baseAddress = static_cast<uint64_t>((uintptr_t)dst);

#if !defined(ASMJIT_DISABLE_LOGGING)
  Logger* logger = getLogger();
#endif // ASMJIT_DISABLE_LOGGING

  size_t minCodeSize = section->getBuffer().getLength(); // Minimum code size.
  size_t maxCodeSize = getCodeSize();                    // Includes all possible trampolines.

  // We will copy the exact size of the generated code. Extra code for trampolines
  // is generated on-the-fly by the relocator (this code doesn't exist at the moment).
  ::memcpy(dst, section->_buffer._data, minCodeSize);

  // Trampoline offset from the beginning of dst/baseAddress.
  size_t trampOffset = minCodeSize;

  // Relocate all recorded locations.
  size_t numRelocs = _relocations.getLength();
  const RelocEntry* const* reArray = _relocations.getData();

  for (size_t i = 0; i < numRelocs; i++) {
    const RelocEntry* re = reArray[i];

    // Possibly deleted or optimized out relocation entry.
    if (re->getType() == RelocEntry::kTypeNone)
      continue;

    uint64_t ptr = re->getData();
    size_t codeOffset = static_cast<size_t>(re->getSourceOffset());

    // Make sure that the `RelocEntry` is correct, we don't want to write
    // out of bounds in `dst`.
    if (ASMJIT_UNLIKELY(codeOffset + re->getSize() > maxCodeSize))
      return DebugUtils::errored(kErrorInvalidRelocEntry);

    // Whether to use trampoline, can be only used if relocation type is `kRelocTrampoline`.
    bool useTrampoline = false;

    switch (re->getType()) {
      case RelocEntry::kTypeAbsToAbs: {
        break;
      }

      case RelocEntry::kTypeRelToAbs: {
        ptr += baseAddress;
        break;
      }

      case RelocEntry::kTypeAbsToRel: {
        ptr -= baseAddress + re->getSourceOffset() + re->getSize();
        break;
      }

      case RelocEntry::kTypeTrampoline: {
        if (re->getSize() != 4)
          return DebugUtils::errored(kErrorInvalidRelocEntry);

        ptr -= baseAddress + re->getSourceOffset() + re->getSize();
        if (!Utils::isInt32(static_cast<int64_t>(ptr))) {
          ptr = (uint64_t)trampOffset - re->getSourceOffset() - re->getSize();
          useTrampoline = true;
        }
        break;
      }

      default:
        return DebugUtils::errored(kErrorInvalidRelocEntry);
    }

    switch (re->getSize()) {
      case 1:
        Utils::writeU8(dst + codeOffset, static_cast<uint32_t>(ptr & 0xFFU));
        break;

      case 4:
        Utils::writeU32u(dst + codeOffset, static_cast<uint32_t>(ptr & 0xFFFFFFFFU));
        break;

      case 8:
        Utils::writeU64u(dst + codeOffset, ptr);
        break;

      default:
        return DebugUtils::errored(kErrorInvalidRelocEntry);
    }

    // Handle the trampoline case.
    if (useTrampoline) {
      // Bytes that replace [REX, OPCODE] bytes.
      uint32_t byte0 = 0xFF;
      uint32_t byte1 = dst[codeOffset - 1];

      if (byte1 == 0xE8) {
        // Patch CALL/MOD byte to FF/2 (-> 0x15).
        byte1 = x86EncodeMod(0, 2, 5);
      }
      else if (byte1 == 0xE9) {
        // Patch JMP/MOD byte to FF/4 (-> 0x25).
        byte1 = x86EncodeMod(0, 4, 5);
      }
      else {
        return DebugUtils::errored(kErrorInvalidRelocEntry);
      }

      // Patch `jmp/call` instruction.
      ASMJIT_ASSERT(codeOffset >= 2);
      dst[codeOffset - 2] = static_cast<uint8_t>(byte0);
      dst[codeOffset - 1] = static_cast<uint8_t>(byte1);

      // Store absolute address and advance the trampoline pointer.
      Utils::writeU64u(dst + trampOffset, re->getData());
      trampOffset += 8;

#if !defined(ASMJIT_DISABLE_LOGGING)
      if (logger)
        logger->logf("[reloc] dq 0x%016llX ; Trampoline\n", re->getData());
#endif // !ASMJIT_DISABLE_LOGGING
    }
  }

  // If there are no trampolines this is the same as `minCodeSize`.
  return trampOffset;
}

} // asmjit namespace
} // namespace PLMD

// [Api-End]
#include "./asmjit_apiend.h"
#pragma GCC diagnostic pop
#endif // __PLUMED_HAS_ASMJIT
