ydb/library/cpp/containers/comptrie/make_fast_layout.cpp

#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);
    }

}