aboutsummaryrefslogtreecommitdiffstats
path: root/library/cpp/containers/comptrie/prefix_iterator.h
blob: 23b9cbf7c6b63f75f5ef181147693c6ac43ee8b4 (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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
#pragma once 

#include "comptrie_trie.h" 
 
// Iterates over all prefixes of the given key in the trie. 
template <class TTrie> 
class TPrefixIterator { 
public: 
    using TSymbol = typename TTrie::TSymbol; 
    using TPacker = typename TTrie::TPacker; 
    using TData = typename TTrie::TData; 
 
private: 
    const TTrie& Trie; 
    const TSymbol* key; 
    size_t keylen; 
    const TSymbol* keyend; 
    size_t prefixLen; 
    const char* valuepos; 
    const char* datapos; 
    const char* dataend; 
    TPacker Packer; 
    const char* EmptyValue; 
    bool result; 
 
    bool Next(); 
 
public: 
    TPrefixIterator(const TTrie& trie, const TSymbol* aKey, size_t aKeylen)
        : Trie(trie) 
        , key(aKey)
        , keylen(aKeylen)
        , keyend(aKey + aKeylen)
        , prefixLen(0) 
        , valuepos(nullptr) 
        , datapos(trie.DataHolder.AsCharPtr()) 
        , dataend(datapos + trie.DataHolder.Length()) 
    { 
        result = Next(); 
    } 
 
    operator bool() const {
        return result; 
    } 
 
    TPrefixIterator& operator++() {
        result = Next(); 
        return *this; 
    } 
 
    size_t GetPrefixLen() const { 
        return prefixLen; 
    } 
 
    void GetValue(TData& to) const { 
        Trie.Packer.UnpackLeaf(valuepos, to); 
    } 
}; 
 
template <class TTrie> 
bool TPrefixIterator<TTrie>::Next() { 
    using namespace NCompactTrie; 
    if (!key || datapos == dataend) 
        return false; 
 
    if ((key == keyend - keylen) && !valuepos && Trie.EmptyValue) { 
        valuepos = Trie.EmptyValue; 
        return true; 
    } 
 
    while (datapos && key != keyend) { 
        TSymbol label = *(key++); 
        if (!Advance(datapos, dataend, valuepos, label, Packer)) { 
            return false; 
        } 
        if (valuepos) { // There is a value at the end of this symbol. 
            prefixLen = keylen - (keyend - key); 
            return true; 
        } 
    } 
 
    return false; 
} 
 
template <class TTrie> 
TPrefixIterator<TTrie> MakePrefixIterator(const TTrie& trie, const typename TTrie::TSymbol* key, size_t keylen) { 
    return TPrefixIterator<TTrie>(trie, key, keylen); 
}