-
Notifications
You must be signed in to change notification settings - Fork 0
/
prefixstorebytetrie.go
151 lines (133 loc) · 4.46 KB
/
prefixstorebytetrie.go
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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
// Tripod provides search-by-prefix utilities.
package tripod
import (
"container/list"
"fmt"
)
// Represents the PrefixStore which uses an in-memory trie data-structure to
// store keys and efficiently return them when searched by prefix.
// PrefixStoreByteTrie is optimized for type []byte.
type PrefixStoreByteTrie struct {
isLast bool
children map[int]*PrefixStoreByteTrie
maxKeySizeInBytes int
}
// Creates and returns reference to a new instance of PrefixStoreByteTrie.
// maxKeySizeInBytes is the maximum size of the key ([]byte) that should be
// allowed to be added to the PrefixStore. When tried to put key of length more
// than maxKeySizeInBytes, the method will return the error.
func CreatePrefixStoreByteTrie(maxKeySizeInBytes int) *PrefixStoreByteTrie {
return &PrefixStoreByteTrie{
children: make(map[int]*PrefixStoreByteTrie),
maxKeySizeInBytes: maxKeySizeInBytes,
}
}
// Adds the []byte key to the PrefixStore and returns if key was succesfully
// added and any error encountered.
// A non nil error is returned if len(key) > maxKeySizeInBytes
func (t *PrefixStoreByteTrie) Put(key []byte) (bool, error) {
if len(key) > t.maxKeySizeInBytes {
return false, fmt.Errorf("max size of key should be %d (%d > %d)",
t.maxKeySizeInBytes, len(key), t.maxKeySizeInBytes)
}
current_node := t
for _, b := range key {
child := current_node.children[int(b)]
if child == nil {
child = CreatePrefixStoreByteTrie(t.maxKeySizeInBytes)
current_node.children[int(b)] = child
}
current_node = child
}
// If key is empty then current_node == t, and since nothing was added
// hence returning false.
if current_node == t {
return false, nil
}
newlyAdded := current_node.isLast == false
current_node.isLast = true
return newlyAdded, nil
}
// Checks and returns if given key is present in the PrefixStore
func (t *PrefixStoreByteTrie) Exists(key []byte) bool {
if len(key) > t.maxKeySizeInBytes {
// Shorting the lookup, since Put method does not allow to put key of
// size > maxKeySizeInBytes, hence shorting this evaluation.
return false
}
current_node := t
for _, b := range key {
child := current_node.children[int(b)]
if child == nil {
return false
}
current_node = child
}
return current_node.isLast
}
// For a given instance of PrefixStore t, this method returns a reference to
// subPrefixStore that ends at the key.
func (t *PrefixStoreByteTrie) get(key []byte) *PrefixStoreByteTrie {
if len(key) > t.maxKeySizeInBytes {
// Shorting the lookup, since Put method does not allow to put key of
// size > maxKeySizeInBytes, hence shorting this evaluation.
return nil
}
current_node := t
for _, b := range key {
child := current_node.children[int(b)]
if child == nil {
return nil
}
current_node = child
}
return current_node
}
// Does a Depth First Search traversal on the PrefixStoreByteTrie and returns a
// List (Double Linked List) containing the keys present in the
// PrefixStore for the given prefix.
func (t *PrefixStoreByteTrie) list(prefix []byte) *list.List {
entries := list.New()
buffer := make([]byte, 0, t.maxKeySizeInBytes)
for ch, tt := range t.children {
buffer = append(buffer, byte(ch))
_dfs(tt, prefix, buffer, entries)
buffer = buffer[:len(buffer)-1]
}
return entries
}
// The DFS Function which recursively calls the children and as it encounters
// a valid existing key, creates a copy of the temporary buffer
// and appends prefix + copyOfBuffer to the linked list; thus forming the
// complete key
func _dfs(t *PrefixStoreByteTrie, prefix []byte, buffer []byte, entries *list.List) {
if t.isLast {
copyOfBuffer := make([]byte, 0, len(buffer)+len(prefix))
// TODO: Check if any optimizations can be done here
copyOfBuffer = append(copyOfBuffer, prefix...)
copyOfBuffer = append(copyOfBuffer, buffer...)
entries.PushBack(copyOfBuffer)
}
for ch, tt := range t.children {
buffer = append(buffer, byte(ch))
_dfs(tt, prefix, buffer, entries)
buffer = buffer[:len(buffer)-1]
}
}
// Does the prefix search on the PrefixStore and returns a reference to
// list (*list.List) containings all entries from the store for the given
// prefix. Each element of the list is []byte.
func (t *PrefixStoreByteTrie) PrefixSearch(prefix []byte) *list.List {
if len(prefix) > t.maxKeySizeInBytes {
return list.New()
}
subTrie := t.get(prefix)
if subTrie == nil {
return list.New()
}
entries := subTrie.list(prefix)
if subTrie.isLast {
entries.PushFront(prefix)
}
return entries
}