#include "make_fast_layout.h"
#include "node.h"
#include "writeable_node.h"
#include "write_trie_backwards.h"
#include "comptrie_impl.h"
#include <util/generic/hash.h>
#include <util/generic/utility.h>
// Lay the trie in memory in such a way that there are less cache misses when jumping from root to leaf.
// The trie becomes about 2% larger, but the access became about 25% faster in our experiments.
// Can be called on minimized and non-minimized tries, in the first case in requires half a trie more memory.
// Calling it the second time on the same trie does nothing.
//
// The algorithm is based on van Emde Boas layout as described in the yandex data school lectures on external memory algoritms
// by Maxim Babenko and Ivan Puzyrevsky. The difference is that when we cut the tree into levels
// two nodes connected by a forward link are put into the same level (because they usually lie near each other in the original tree).
// The original paper (describing the layout in Section 2.1) is:
// Michael A. Bender, Erik D. Demaine, Martin Farach-Colton. Cache-Oblivious B-Trees // SIAM Journal on Computing, volume 35, number 2, 2005, pages 341–358.
// Available on the web: http://erikdemaine.org/papers/CacheObliviousBTrees_SICOMP/
// Or: Michael A. Bender, Erik D. Demaine, and Martin Farach-Colton. Cache-Oblivious B-Trees // Proceedings of the 41st Annual Symposium
// on Foundations of Computer Science (FOCS 2000), Redondo Beach, California, November 12–14, 2000, pages 399–409.
// Available on the web: http://erikdemaine.org/papers/FOCS2000b/
// (there is not much difference between these papers, actually).
//
namespace NCompactTrie {
static size_t FindSupportingPowerOf2(size_t n) {
size_t result = 1ull << (8 * sizeof(size_t) - 1);
while (result > n) {
result >>= 1;
}
return result;
}
namespace {
class TTrieNodeSet {
public:
TTrieNodeSet() = default;
explicit TTrieNodeSet(const TOpaqueTrie& trie)
: Body(trie.Length / (8 * MinNodeSize) + 1, 0)
{
}
bool Has(size_t offset) const {
const size_t reducedOffset = ReducedOffset(offset);
return OffsetCell(reducedOffset) & OffsetMask(reducedOffset);
}
void Add(size_t offset) {
const size_t reducedOffset = ReducedOffset(offset);
OffsetCell(reducedOffset) |= OffsetMask(reducedOffset);
}
void Remove(size_t offset) {
const size_t reducedOffset = ReducedOffset(offset);
OffsetCell(reducedOffset) &= ~OffsetMask(reducedOffset);
}
void Swap(TTrieNodeSet& other) {
Body.swap(other.Body);
}
private:
static const size_t MinNodeSize = 2;
TVector<ui8> Body;
static size_t ReducedOffset(size_t offset) {
return offset / MinNodeSize;
}
static ui8 OffsetMask(size_t reducedOffset) {
return 1 << (reducedOffset % 8);
}
ui8& OffsetCell(size_t reducedOffset) {
return Body.at(reducedOffset / 8);
}
const ui8& OffsetCell(size_t reducedOffset) const {
return Body.at(reducedOffset / 8);
}
};
//---------------------------------------------------------------------
class TTrieNodeCounts {
public:
TTrieNodeCounts() = default;
explicit TTrieNodeCounts(const TOpaqueTrie& trie)
: Body(trie.Length / MinNodeSize, 0)
, IsTree(false)
{
}
size_t Get(size_t offset) const {
return IsTree ? 1 : Body.at(offset / MinNodeSize);
}
void Inc(size_t offset) {
if (IsTree) {
return;
}
ui8& count = Body.at(offset / MinNodeSize);
if (count != MaxCount) {
++count;
}
}
size_t Dec(size_t offset) {
if (IsTree) {
return 0;
}
ui8& count = Body.at(offset / MinNodeSize);
Y_ASSERT(count > 0);
if (count != MaxCount) {
--count;
}
return count;
}
void Swap(TTrieNodeCounts& other) {
Body.swap(other.Body);
::DoSwap(IsTree, other.IsTree);
}
void SetTreeMode() {
IsTree = true;
Body = TVector<ui8>();
}
private:
static const size_t MinNodeSize = 2;
static const ui8 MaxCount = 255;
TVector<ui8> Body;
bool IsTree = false;
};
//----------------------------------------------------------
class TOffsetIndex {
public:
// In all methods:
// Key --- offset from the beginning of the old trie.
// Value --- offset from the end of the new trie.
explicit TOffsetIndex(TTrieNodeCounts& counts) {
ParentCounts.Swap(counts);
}
void Add(size_t key, size_t value) {
Data[key] = value;
}
size_t Size() const {
return Data.size();
}
size_t Get(size_t key) {
auto pos = Data.find(key);
if (pos == Data.end()) {
ythrow yexception() << "Bad node walking order: trying to get node offset too early or too many times!";
}
size_t result = pos->second;
if (ParentCounts.Dec(key) == 0) {
// We don't need this offset any more.
Data.erase(pos);
}
return result;
}
private:
TTrieNodeCounts ParentCounts;
THashMap<size_t, size_t> Data;
};
//---------------------------------------------------------------------------------------
class TTrieMeasurer {
public:
TTrieMeasurer(const TOpaqueTrie& trie, bool verbose);
void Measure();
size_t GetDepth() const {
return Depth;
}
size_t GetNodeCount() const {
return NodeCount;
}
size_t GetUnminimizedNodeCount() const {
return UnminimizedNodeCount;
}
bool IsTree() const {
return NodeCount == UnminimizedNodeCount;
}
TTrieNodeCounts& GetParentCounts() {
return ParentCounts;
}
const TOpaqueTrie& GetTrie() const {
return Trie;
}
private:
const TOpaqueTrie& Trie;
size_t Depth;
TTrieNodeCounts ParentCounts;
size_t NodeCount;
size_t UnminimizedNodeCount;
const bool Verbose;
// returns depth, increments NodeCount.
size_t MeasureSubtrie(size_t rootOffset, bool isNewPath);
};
TTrieMeasurer::TTrieMeasurer(const TOpaqueTrie& trie, bool verbose)
: Trie(trie)
, Depth(0)
, ParentCounts(trie)
, NodeCount(0)
, UnminimizedNodeCount(0)
, Verbose(verbose)
{
Y_ASSERT(Trie.Data);
}
void TTrieMeasurer::Measure() {
if (Verbose) {
Cerr << "Measuring the trie..." << Endl;
}
NodeCount = 0;
UnminimizedNodeCount = 0;
Depth = MeasureSubtrie(0, true);
if (IsTree()) {
ParentCounts.SetTreeMode();
}
if (Verbose) {
Cerr << "Unminimized node count: " << UnminimizedNodeCount << Endl;
Cerr << "Trie depth: " << Depth << Endl;
Cerr << "Node count: " << NodeCount << Endl;
}
}
// A chain of nodes linked by forward links
// is considered one node with many left and right children
// for depth measuring here and in
// TVanEmdeBoasReverseNodeEnumerator::FindDescendants.
size_t TTrieMeasurer::MeasureSubtrie(size_t rootOffset, bool isNewPath) {
Y_ASSERT(rootOffset < Trie.Length);
TNode node(Trie.Data, rootOffset, Trie.SkipFunction);
size_t depth = 0;
for (;;) {
++UnminimizedNodeCount;
if (Verbose) {
ShowProgress(UnminimizedNodeCount);
}
if (isNewPath) {
if (ParentCounts.Get(node.GetOffset()) > 0) {
isNewPath = false;
} else {
++NodeCount;
}
ParentCounts.Inc(node.GetOffset());
}
if (node.GetLeftOffset()) {
depth = Max(depth, 1 + MeasureSubtrie(node.GetLeftOffset(), isNewPath));
}
if (node.GetRightOffset()) {
depth = Max(depth, 1 + MeasureSubtrie(node.GetRightOffset(), isNewPath));
}
if (node.GetForwardOffset()) {
node = TNode(Trie.Data, node.GetForwardOffset(), Trie.SkipFunction);
} else {
break;
}
}
return depth;
}
//--------------------------------------------------------------------------------------
using TLevelNodes = TVector<size_t>;
struct TLevel {
size_t Depth;
TLevelNodes Nodes;
explicit TLevel(size_t depth)
: Depth(depth)
{
}
};
//----------------------------------------------------------------------------------------
class TVanEmdeBoasReverseNodeEnumerator: public TReverseNodeEnumerator {
public:
TVanEmdeBoasReverseNodeEnumerator(TTrieMeasurer& measurer, bool verbose)
: Fresh(true)
, Trie(measurer.GetTrie())
, Depth(measurer.GetDepth())
, MaxIndexSize(0)
, BackIndex(measurer.GetParentCounts())
, ProcessedNodes(measurer.GetTrie())
, Verbose(verbose)
{
}
bool Move() override {
if (!Fresh) {
CentralWord.pop_back();
}
if (CentralWord.empty()) {
return MoveCentralWordStart();
}
return true;
}
const TNode& Get() const {
return CentralWord.back();
}
size_t GetLeafLength() const override {
return Get().GetLeafLength();
}
// Returns recalculated offset from the end of the current node.
size_t PrepareOffset(size_t absoffset, size_t resultLength) {
if (!absoffset)
return NPOS;
return resultLength - BackIndex.Get(absoffset);
}
size_t RecreateNode(char* buffer, size_t resultLength) override {
TWriteableNode newNode(Get(), Trie.Data);
newNode.ForwardOffset = PrepareOffset(Get().GetForwardOffset(), resultLength);
newNode.LeftOffset = PrepareOffset(Get().GetLeftOffset(), resultLength);
newNode.RightOffset = PrepareOffset(Get().GetRightOffset(), resultLength);
const size_t len = newNode.Pack(buffer);
ProcessedNodes.Add(Get().GetOffset());
BackIndex.Add(Get().GetOffset(), resultLength + len);
MaxIndexSize = Max(MaxIndexSize, BackIndex.Size());
return len;
}
private:
bool Fresh;
TOpaqueTrie Trie;
size_t Depth;
size_t MaxIndexSize;
TVector<TLevel> Trace;
TOffsetIndex BackIndex;
TVector<TNode> CentralWord;
TTrieNodeSet ProcessedNodes;
const bool Verbose;
private:
bool IsVisited(size_t offset) const {
return ProcessedNodes.Has(offset);
}
void AddCascade(size_t step) {
Y_ASSERT(!(step & (step - 1))); // Should be a power of 2.
while (step > 0) {
size_t root = Trace.back().Nodes.back();
TLevel level(Trace.back().Depth + step);
Trace.push_back(level);
size_t actualStep = FindSupportingPowerOf2(FindDescendants(root, step, Trace.back().Nodes));
if (actualStep != step) {
// Retry with a smaller step.
Y_ASSERT(actualStep < step);
step = actualStep;
Trace.pop_back();
} else {
step /= 2;
}
}
}
void FillCentralWord() {
CentralWord.clear();
CentralWord.push_back(TNode(Trie.Data, Trace.back().Nodes.back(), Trie.SkipFunction));
// Do not check for epsilon links, as the traversal order now is different.
while (CentralWord.back().GetForwardOffset() && !IsVisited(CentralWord.back().GetForwardOffset())) {
CentralWord.push_back(TNode(Trie.Data, CentralWord.back().GetForwardOffset(), Trie.SkipFunction));
}
}
bool MoveCentralWordStart() {
do {
if (Fresh) {
TLevel root(0);
Trace.push_back(root);
Trace.back().Nodes.push_back(0);
const size_t sectionDepth = FindSupportingPowerOf2(Depth);
AddCascade(sectionDepth);
Fresh = false;
} else {
Trace.back().Nodes.pop_back();
if (Trace.back().Nodes.empty() && Trace.size() == 1) {
if (Verbose) {
Cerr << "Max index size: " << MaxIndexSize << Endl;
Cerr << "Current index size: " << BackIndex.Size() << Endl;
}
// We just popped the root.
return false;
}
size_t lastStep = Trace.back().Depth - Trace[Trace.size() - 2].Depth;
if (Trace.back().Nodes.empty()) {
Trace.pop_back();
}
AddCascade(lastStep / 2);
}
} while (IsVisited(Trace.back().Nodes.back()));
FillCentralWord();
return true;
}
// Returns the maximal depth it has reached while searching.
// This is a method because it needs OffsetIndex to skip processed nodes.
size_t FindDescendants(size_t rootOffset, size_t depth, TLevelNodes& result) const {
if (depth == 0) {
result.push_back(rootOffset);
return 0;
}
size_t actualDepth = 0;
TNode node(Trie.Data, rootOffset, Trie.SkipFunction);
for (;;) {
if (node.GetLeftOffset() && !IsVisited(node.GetLeftOffset())) {
actualDepth = Max(actualDepth,
1 + FindDescendants(node.GetLeftOffset(), depth - 1, result));
}
if (node.GetRightOffset() && !IsVisited(node.GetRightOffset())) {
actualDepth = Max(actualDepth,
1 + FindDescendants(node.GetRightOffset(), depth - 1, result));
}
if (node.GetForwardOffset() && !IsVisited(node.GetForwardOffset())) {
node = TNode(Trie.Data, node.GetForwardOffset(), Trie.SkipFunction);
} else {
break;
}
}
return actualDepth;
}
};
} // Anonymous namespace.
//-----------------------------------------------------------------------------------
size_t RawCompactTrieFastLayoutImpl(IOutputStream& os, const TOpaqueTrie& trie, bool verbose) {
if (!trie.Data || !trie.Length) {
return 0;
}
TTrieMeasurer mes(trie, verbose);
mes.Measure();
TVanEmdeBoasReverseNodeEnumerator enumerator(mes, verbose);
return WriteTrieBackwards(os, enumerator, verbose);
}
}