aboutsummaryrefslogtreecommitdiffstats
path: root/library/cpp/containers/comptrie/first_symbol_iterator.h
blob: d3c521e7e64696eebb3b4fb202949e2e4b0763a2 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
#pragma once

#include "opaque_trie_iterator.h"
#include <util/generic/ptr.h>

namespace NCompactTrie {
    // Iterates over possible first symbols in a trie.
    // Allows one to get the symbol and the subtrie starting from it.
    template <class TTrie> 
    class TFirstSymbolIterator {
    public:
        using TSymbol = typename TTrie::TSymbol;
        using TData = typename TTrie::TData;

        void SetTrie(const TTrie& trie, const ILeafSkipper& skipper) {
            Trie = trie;
            Impl.Reset(new TOpaqueTrieIterator(
                TOpaqueTrie(Trie.Data().AsCharPtr(), Trie.Data().Size(), skipper),
                nullptr,
                false,
                sizeof(TSymbol)));
            if (Impl->MeasureKey<TSymbol>() == 0) {
                MakeStep();
            }
        }

        const TTrie& GetTrie() const {
            return Trie;
        }

        bool AtEnd() const {
            return *Impl == TOpaqueTrieIterator(Impl->GetTrie(), nullptr, true, sizeof(TSymbol));
        }

        TSymbol GetKey() const {
            return Impl->GetKey<TSymbol>()[0];
        }

        TTrie GetTails() const {
            const TNode& node = Impl->GetNode();
            const size_t forwardOffset = node.GetForwardOffset();
            const char* emptyValue = node.IsFinal() ? Trie.Data().AsCharPtr() + node.GetLeafOffset() : nullptr;
            if (forwardOffset) {
                const char* start = Trie.Data().AsCharPtr() + forwardOffset;
                TBlob body = TBlob::NoCopy(start, Trie.Data().Size() - forwardOffset);
                return TTrie(body, emptyValue, Trie.GetPacker());
            } else {
                return TTrie(emptyValue);
            }
        }

        void MakeStep() {
            Impl->Forward();
        }

    private:
        TTrie Trie;
        TCopyPtr<TOpaqueTrieIterator> Impl;
    };

}