sequenceTrie.hpp

Go to the documentation of this file.

#pragma once
 
#include "exceptions.hpp"
#include <algorithm>
#include <array>
#include <atomic>
#include <cstdint>
#include <mutex>
#include <vector>
 
namespace jaffarCommon
{
 
namespace sequenceTrie
{
 
template <typename Element>
class SequenceTrie
{
public:
  using nodeId_t = uint32_t;
 
  static constexpr nodeId_t ROOT = 0;
 
  static constexpr nodeId_t NONE = 0xFFFFFFFFu;
 
  explicit SequenceTrie(uint32_t numShards = 1, uint32_t chunkSizeLog2 = 20)
      : _chunkBits(chunkSizeLog2), _chunkSize(1u << chunkSizeLog2), _chunkMask((1u << chunkSizeLog2) - 1), _shards(numShards == 0 ? 1 : numShards)
  {
    for (auto& c : _chunks) c.store(nullptr, std::memory_order_relaxed);
    for (auto& s : _shards) s.head = NONE;
    _bump.store(0, std::memory_order_relaxed);
 
    // Allocate ROOT (id 0). It carries a permanent self-reference so it is never recycled.
    const nodeId_t root = allocNode(0);
    Node&          r    = node(root);
    r.parent            = NONE;
    r.element           = Element{};
    r.refCount.store(1, std::memory_order_relaxed);
  }
 
  ~SequenceTrie()
  {
    for (auto& c : _chunks)
    {
      Node* p = c.load(std::memory_order_relaxed);
      delete[] p;
    }
  }
 
  SequenceTrie(const SequenceTrie&)            = delete;
  SequenceTrie& operator=(const SequenceTrie&) = delete;
 
  nodeId_t extend(nodeId_t parent, Element element, uint32_t shard = 0)
  {
    const nodeId_t id = allocNode(shard);
    Node&          n  = node(id);
    n.parent          = parent;
    n.element         = element;
    n.refCount.store(1, std::memory_order_relaxed); // the caller's reference
 
    // Account for the new child edge on the parent (kept alive by its children).
    incRef(node(parent));
    return id;
  }
 
  void acquire(nodeId_t id) { incRef(node(id)); }
 
  void release(nodeId_t id, uint32_t shard = 0)
  {
    while (id != NONE)
    {
      Node& n = node(id);
      // release-acquire so a thread that observes the count hit zero sees all prior writes to the node.
      if (n.refCount.fetch_sub(1, std::memory_order_acq_rel) != 1) break; // still referenced -> stop
 
      const nodeId_t parent = n.parent; // read before the slot is handed to the free list
      freeNode(id, shard);
      id = parent; // the freed node was a child of `parent`; drop that child edge too
    }
  }
 
  template <typename OutElem = Element>
  void reconstruct(nodeId_t id, std::vector<OutElem>& out) const
  {
    out.clear();
    for (nodeId_t cur = id; cur != ROOT && cur != NONE;)
    {
      const Node& n = node(cur);
      out.push_back(static_cast<OutElem>(n.element)); // widen the 16-bit stored element to the caller's type
      cur = n.parent;
    }
    std::reverse(out.begin(), out.end());
  }
 
  size_t getDepth(nodeId_t id) const
  {
    size_t d = 0;
    for (nodeId_t cur = id; cur != ROOT && cur != NONE; cur = node(cur).parent) d++;
    return d;
  }
 
  size_t getAllocatedNodeCount() const { return _bump.load(std::memory_order_relaxed); }
 
  size_t getApproxMemoryBytes() const { return getAllocatedNodeCount() * sizeof(Node); }
 
  size_t getMaxMemoryBytes() const { return MAX_CHUNKS * (size_t)_chunkSize * sizeof(Node); }
 
private:
  struct Node
  {
    uint32_t              parent;   
    Element               element;  
    std::atomic<uint16_t> refCount; 
  };
 
  // Fixed cap on chunks so the chunk-pointer array never reallocates (handles stay valid lock-free).
  static constexpr size_t MAX_CHUNKS = 4096;
 
  __attribute__((always_inline)) Node& node(nodeId_t id) const { return _chunks[id >> _chunkBits].load(std::memory_order_acquire)[id & _chunkMask]; }
 
  __attribute__((always_inline)) void incRef(Node& n) const
  {
    if (n.refCount.fetch_add(1, std::memory_order_relaxed) == UINT16_MAX)
      JAFFAR_THROW_RUNTIME("SequenceTrie node refCount overflowed 16 bits (out-degree + holders > %u); the "
                           "input set is too large for the 16-bit-refCount trie node.",
                           (unsigned)UINT16_MAX);
  }
 
  // Obtain a node slot: reuse one recycled into this shard's (thread-private) free list if available,
  // otherwise bump-allocate a fresh id. The per-shard free list needs no atomics because a given shard
  // is only ever touched by its owning thread -- removing the single-free-list contention that otherwise
  // serializes all workers. Only the rare fresh-id bump (free list empty) touches a shared atomic.
  nodeId_t allocNode(uint32_t shard)
  {
    FreeShard& fs = _shards[shard];
    if (fs.head != NONE)
    {
      const nodeId_t top = fs.head;
      fs.head            = node(top).parent; // free-list "next" stored in parent
      return top;
    }
 
    // Free list empty: bump-allocate a fresh id, faulting in its chunk on first use.
    const nodeId_t id    = (nodeId_t)_bump.fetch_add(1, std::memory_order_relaxed);
    const size_t   chunk = id >> _chunkBits;
    if (chunk >= MAX_CHUNKS) JAFFAR_THROW_RUNTIME("SequenceTrie exceeded its maximum capacity (%zu chunks)", MAX_CHUNKS);
    if (_chunks[chunk].load(std::memory_order_acquire) == nullptr) ensureChunk(chunk);
    return id;
  }
 
  void freeNode(nodeId_t id, uint32_t shard)
  {
    FreeShard& fs   = _shards[shard];
    node(id).parent = fs.head; // push onto this shard's private stack
    fs.head         = id;
  }
 
  void ensureChunk(size_t chunk)
  {
    std::lock_guard<std::mutex> lock(_growMutex);
    if (_chunks[chunk].load(std::memory_order_relaxed) != nullptr) return; // another thread won the race
    Node* p = new Node[_chunkSize];
    _chunks[chunk].store(p, std::memory_order_release);
  }
 
  struct alignas(64) FreeShard
  {
    nodeId_t head = NONE;
  };
 
  const uint32_t _chunkBits;
  const uint32_t _chunkSize;
  const uint32_t _chunkMask;
 
  mutable std::array<std::atomic<Node*>, MAX_CHUNKS> _chunks;
  std::atomic<uint64_t>                              _bump; 
  std::vector<FreeShard>                             _shards;
  std::mutex                                         _growMutex;
};
 
} // namespace sequenceTrie
 
} // namespace jaffarCommon