/* +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
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 "./constpool.h"
#include "./utils.h"

#include <algorithm>

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

namespace PLMD {
namespace asmjit {

// Binary tree code is based on Julienne Walker's "Andersson Binary Trees"
// article and implementation. However, only three operations are implemented -
// get, insert and traverse.

// ============================================================================
// [asmjit::ConstPool::Tree - Ops]
// ============================================================================

//! \internal
//!
//! Remove left horizontal links.
static ASMJIT_INLINE ConstPool::Node* ConstPoolTree_skewNode(ConstPool::Node* node) noexcept {
  ConstPool::Node* link = node->_link[0];
  uint32_t level = node->_level;

  if (level != 0 && link && link->_level == level) {
    node->_link[0] = link->_link[1];
    link->_link[1] = node;

    node = link;
  }

  return node;
}

//! \internal
//!
//! Remove consecutive horizontal links.
static ASMJIT_INLINE ConstPool::Node* ConstPoolTree_splitNode(ConstPool::Node* node) noexcept {
  ConstPool::Node* link = node->_link[1];
  uint32_t level = node->_level;

  if (level != 0 && link && link->_link[1] && link->_link[1]->_level == level) {
    node->_link[1] = link->_link[0];
    link->_link[0] = node;

    node = link;
    node->_level++;
  }

  return node;
}

ConstPool::Node* ConstPool::Tree::get(const void* data) noexcept {
  ConstPool::Node* node = _root;
  size_t dataSize = _dataSize;

  while (node) {
    int c = ::memcmp(node->getData(), data, dataSize);
    if (c == 0)
      return node;
    node = node->_link[c < 0];
  }

  return nullptr;
}

void ConstPool::Tree::put(ConstPool::Node* newNode) noexcept {
  size_t dataSize = _dataSize;
  _length++;

  if (!_root) {
    _root = newNode;
    return;
  }

  ConstPool::Node* node = _root;
  ConstPool::Node* stack[kHeightLimit];

  unsigned int top = 0;
  unsigned int dir;

  // Find a spot and save the stack.
  for (;;) {
    stack[top++] = node;
    dir = ::memcmp(node->getData(), newNode->getData(), dataSize) < 0;

    ConstPool::Node* link = node->_link[dir];
    if (!link) break;

    node = link;
  }

  // Link and rebalance.
  node->_link[dir] = newNode;

  while (top > 0) {
    // Which child?
    node = stack[--top];

    if (top != 0) {
      dir = stack[top - 1]->_link[1] == node;
    }

    node = ConstPoolTree_skewNode(node);
    node = ConstPoolTree_splitNode(node);

    // Fix the parent.
    if (top != 0)
      stack[top - 1]->_link[dir] = node;
    else
      _root = node;
  }
}

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

ConstPool::ConstPool(Zone* zone) noexcept { reset(zone); }
ConstPool::~ConstPool() noexcept {}

// ============================================================================
// [asmjit::ConstPool - Reset]
// ============================================================================

void ConstPool::reset(Zone* zone) noexcept {
  _zone = zone;

  size_t dataSize = 1;
  for (size_t i = 0; i < ASMJIT_ARRAY_SIZE(_tree); i++) {
    _tree[i].reset();
    _tree[i].setDataSize(dataSize);
    _gaps[i] = nullptr;
    dataSize <<= 1;
  }

  _gapPool = nullptr;
  _size = 0;
  _alignment = 0;
}

// ============================================================================
// [asmjit::ConstPool - Ops]
// ============================================================================

static ASMJIT_INLINE ConstPool::Gap* ConstPool_allocGap(ConstPool* self) noexcept {
  ConstPool::Gap* gap = self->_gapPool;
  if (!gap) return self->_zone->allocT<ConstPool::Gap>();

  self->_gapPool = gap->_next;
  return gap;
}

static ASMJIT_INLINE void ConstPool_freeGap(ConstPool* self,  ConstPool::Gap* gap) noexcept {
  gap->_next = self->_gapPool;
  self->_gapPool = gap;
}

static void ConstPool_addGap(ConstPool* self, size_t offset, size_t length) noexcept {
  ASMJIT_ASSERT(length > 0);

  while (length > 0) {
    size_t gapIndex;
    size_t gapLength;

      gapIndex = ConstPool::kIndex16;
    if (length >= 16 && Utils::isAligned<size_t>(offset, 16)) {
      gapLength = 16;
    }
    else if (length >= 8 && Utils::isAligned<size_t>(offset, 8)) {
      gapIndex = ConstPool::kIndex8;
      gapLength = 8;
    }
    else if (length >= 4 && Utils::isAligned<size_t>(offset, 4)) {
      gapIndex = ConstPool::kIndex4;
      gapLength = 4;
    }
    else if (length >= 2 && Utils::isAligned<size_t>(offset, 2)) {
      gapIndex = ConstPool::kIndex2;
      gapLength = 2;
    }
    else {
      gapIndex = ConstPool::kIndex1;
      gapLength = 1;
    }

    // We don't have to check for errors here, if this failed nothing really
    // happened (just the gap won't be visible) and it will fail again at
    // place where checking will cause kErrorNoHeapMemory.
    ConstPool::Gap* gap = ConstPool_allocGap(self);
    if (!gap) return;

    gap->_next = self->_gaps[gapIndex];
    self->_gaps[gapIndex] = gap;

    gap->_offset = offset;
    gap->_length = gapLength;

    offset += gapLength;
    length -= gapLength;
  }
}

Error ConstPool::add(const void* data, size_t size, size_t& dstOffset) noexcept {
  size_t treeIndex;

  if (size == 32)
    treeIndex = kIndex32;
  else if (size == 16)
    treeIndex = kIndex16;
  else if (size == 8)
    treeIndex = kIndex8;
  else if (size == 4)
    treeIndex = kIndex4;
  else if (size == 2)
    treeIndex = kIndex2;
  else if (size == 1)
    treeIndex = kIndex1;
  else
    return DebugUtils::errored(kErrorInvalidArgument);

  ConstPool::Node* node = _tree[treeIndex].get(data);
  if (node) {
    dstOffset = node->_offset;
    return kErrorOk;
  }

  // Before incrementing the current offset try if there is a gap that can
  // be used for the requested data.
  size_t offset = ~static_cast<size_t>(0);
  size_t gapIndex = treeIndex;

  while (gapIndex != kIndexCount - 1) {
    ConstPool::Gap* gap = _gaps[treeIndex];

    // Check if there is a gap.
    if (gap) {
      size_t gapOffset = gap->_offset;
      size_t gapLength = gap->_length;

      // Destroy the gap for now.
      _gaps[treeIndex] = gap->_next;
      ConstPool_freeGap(this, gap);

      offset = gapOffset;
      ASMJIT_ASSERT(Utils::isAligned<size_t>(offset, size));

      gapLength -= size;
      if (gapLength > 0)
        ConstPool_addGap(this, gapOffset, gapLength);
    }

    gapIndex++;
  }

  if (offset == ~static_cast<size_t>(0)) {
    // Get how many bytes have to be skipped so the address is aligned accordingly
    // to the 'size'.
    size_t diff = Utils::alignDiff<size_t>(_size, size);

    if (diff != 0) {
      ConstPool_addGap(this, _size, diff);
      _size += diff;
    }

    offset = _size;
    _size += size;
  }

  // Add the initial node to the right index.
  node = ConstPool::Tree::_newNode(_zone, data, size, offset, false);
  if (!node) return DebugUtils::errored(kErrorNoHeapMemory);

  _tree[treeIndex].put(node);
  _alignment = std::max<size_t>(_alignment, size);

  dstOffset = offset;

  // Now create a bunch of shared constants that are based on the data pattern.
  // We stop at size 4, it probably doesn't make sense to split constants down
  // to 1 byte.
  size_t pCount = 1;
  while (size > 4) {
    size >>= 1;
    pCount <<= 1;

    ASMJIT_ASSERT(treeIndex != 0);
    treeIndex--;

    const uint8_t* pData = static_cast<const uint8_t*>(data);
    for (size_t i = 0; i < pCount; i++, pData += size) {
      node = _tree[treeIndex].get(pData);
      if (node) continue;

      node = ConstPool::Tree::_newNode(_zone, pData, size, offset + (i * size), true);
      _tree[treeIndex].put(node);
    }
  }

  return kErrorOk;
}

// ============================================================================
// [asmjit::ConstPool - Reset]
// ============================================================================

struct ConstPoolFill {
  ASMJIT_INLINE ConstPoolFill(uint8_t* dst, size_t dataSize) noexcept :
    _dst(dst),
    _dataSize(dataSize) {}

  ASMJIT_INLINE void visit(const ConstPool::Node* node) noexcept {
    if (!node->_shared)
      ::memcpy(_dst + node->_offset, node->getData(), _dataSize);
  }

  uint8_t* _dst;
  size_t _dataSize;
};

void ConstPool::fill(void* dst) const noexcept {
  // Clears possible gaps, asmjit should never emit garbage to the output.
  ::memset(dst, 0, _size);

  ConstPoolFill filler(static_cast<uint8_t*>(dst), 1);
  for (size_t i = 0; i < ASMJIT_ARRAY_SIZE(_tree); i++) {
    _tree[i].iterate(filler);
    filler._dataSize <<= 1;
  }
}

// ============================================================================
// [asmjit::ConstPool - Test]
// ============================================================================

#if defined(ASMJIT_TEST)
UNIT(base_constpool) {
  Zone zone(32384 - Zone::kZoneOverhead);
  ConstPool pool(&zone);

  uint32_t i;
  uint32_t kCount = 1000000;

  INFO("Adding %u constants to the pool.", kCount);
  {
    size_t prevOffset;
    size_t curOffset;
    uint64_t c = ASMJIT_UINT64_C(0x0101010101010101);

    EXPECT(pool.add(&c, 8, prevOffset) == kErrorOk,
      "pool.add() - Returned error");
    EXPECT(prevOffset == 0,
      "pool.add() - First constant should have zero offset");

    for (i = 1; i < kCount; i++) {
      c++;
      EXPECT(pool.add(&c, 8, curOffset) == kErrorOk,
        "pool.add() - Returned error");
      EXPECT(prevOffset + 8 == curOffset,
        "pool.add() - Returned incorrect curOffset");
      EXPECT(pool.getSize() == (i + 1) * 8,
        "pool.getSize() - Reported incorrect size");
      prevOffset = curOffset;
    }

    EXPECT(pool.getAlignment() == 8,
      "pool.getAlignment() - Expected 8-byte alignment");
  }

  INFO("Retrieving %u constants from the pool.", kCount);
  {
    uint64_t c = ASMJIT_UINT64_C(0x0101010101010101);

    for (i = 0; i < kCount; i++) {
      size_t offset;
      EXPECT(pool.add(&c, 8, offset) == kErrorOk,
        "pool.add() - Returned error");
      EXPECT(offset == i * 8,
        "pool.add() - Should have reused constant");
      c++;
    }
  }

  INFO("Checking if the constants were split into 4-byte patterns");
  {
    uint32_t c = 0x01010101;
    for (i = 0; i < kCount; i++) {
      size_t offset;
      EXPECT(pool.add(&c, 4, offset) == kErrorOk,
        "pool.add() - Returned error");
      EXPECT(offset == i * 8,
        "pool.add() - Should reuse existing constant");
      c++;
    }
  }

  INFO("Adding 2 byte constant to misalign the current offset");
  {
    uint16_t c = 0xFFFF;
    size_t offset;

    EXPECT(pool.add(&c, 2, offset) == kErrorOk,
      "pool.add() - Returned error");
    EXPECT(offset == kCount * 8,
      "pool.add() - Didn't return expected position");
    EXPECT(pool.getAlignment() == 8,
      "pool.getAlignment() - Expected 8-byte alignment");
  }

  INFO("Adding 8 byte constant to check if pool gets aligned again");
  {
    uint64_t c = ASMJIT_UINT64_C(0xFFFFFFFFFFFFFFFF);
    size_t offset;

    EXPECT(pool.add(&c, 8, offset) == kErrorOk,
      "pool.add() - Returned error");
    EXPECT(offset == kCount * 8 + 8,
      "pool.add() - Didn't return aligned offset");
  }

  INFO("Adding 2 byte constant to verify the gap is filled");
  {
    uint16_t c = 0xFFFE;
    size_t offset;

    EXPECT(pool.add(&c, 2, offset) == kErrorOk,
      "pool.add() - Returned error");
    EXPECT(offset == kCount * 8 + 2,
      "pool.add() - Didn't fill the gap");
    EXPECT(pool.getAlignment() == 8,
      "pool.getAlignment() - Expected 8-byte alignment");
  }

  INFO("Checking reset functionality");
  {
    pool.reset(&zone);
    zone.reset();

    EXPECT(pool.getSize() == 0,
      "pool.getSize() - Expected pool size to be zero");
    EXPECT(pool.getAlignment() == 0,
      "pool.getSize() - Expected pool alignment to be zero");
  }

  INFO("Checking pool alignment when combined constants are added");
  {
    uint8_t bytes[32] = { 0 };
    size_t offset;

    pool.add(bytes, 1, offset);

    EXPECT(pool.getSize() == 1,
      "pool.getSize() - Expected pool size to be 1 byte");
    EXPECT(pool.getAlignment() == 1,
      "pool.getSize() - Expected pool alignment to be 1 byte");
    EXPECT(offset == 0,
      "pool.getSize() - Expected offset returned to be zero");

    pool.add(bytes, 2, offset);

    EXPECT(pool.getSize() == 4,
      "pool.getSize() - Expected pool size to be 4 bytes");
    EXPECT(pool.getAlignment() == 2,
      "pool.getSize() - Expected pool alignment to be 2 bytes");
    EXPECT(offset == 2,
      "pool.getSize() - Expected offset returned to be 2");

    pool.add(bytes, 4, offset);

    EXPECT(pool.getSize() == 8,
      "pool.getSize() - Expected pool size to be 8 bytes");
    EXPECT(pool.getAlignment() == 4,
      "pool.getSize() - Expected pool alignment to be 4 bytes");
    EXPECT(offset == 4,
      "pool.getSize() - Expected offset returned to be 4");

    pool.add(bytes, 4, offset);

    EXPECT(pool.getSize() == 8,
      "pool.getSize() - Expected pool size to be 8 bytes");
    EXPECT(pool.getAlignment() == 4,
      "pool.getSize() - Expected pool alignment to be 4 bytes");
    EXPECT(offset == 4,
      "pool.getSize() - Expected offset returned to be 8");

    pool.add(bytes, 32, offset);
    EXPECT(pool.getSize() == 64,
      "pool.getSize() - Expected pool size to be 64 bytes");
    EXPECT(pool.getAlignment() == 32,
      "pool.getSize() - Expected pool alignment to be 32 bytes");
    EXPECT(offset == 32,
      "pool.getSize() - Expected offset returned to be 32");
  }
}
#endif // ASMJIT_TEST

} // asmjit namespace
} // namespace PLMD

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