intermediate changes

ref:cde9a383711a11544ce7e107a78147fb96cc4029

1 files changed, 467 insertions, 0 deletions

diff --git a/library/cpp/containers/comptrie/make_fast_layout.cpp b/library/cpp/containers/comptrie/make_fast_layout.cpp
new file mode 100644
index 0000000000..ade78d7899
--- /dev/null
+++ b/library/cpp/containers/comptrie/make_fast_layout.cpp
@@ -0,0 +1,467 @@
+#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);
+    }
+
+}