105 lines
2.4 KiB
Go
105 lines
2.4 KiB
Go
|
/*
|
||
|
|
* Copyright 2019 Dgraph Labs, Inc. and Contributors
|
||
|
|
*
|
||
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
* you may not use this file except in compliance with the License.
|
||
|
|
* You may obtain a copy of the License at
|
||
|
|
*
|
||
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
*
|
||
|
|
* Unless required by applicable law or agreed to in writing, software
|
||
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
* See the License for the specific language governing permissions and
|
||
|
|
* limitations under the License.
|
||
|
|
*/
|
||
|
|
|
||
|
|
package trie
|
||
|
|
|
||
|
|
type node struct {
|
||
|
|
children map[byte]*node
|
||
|
|
ids []uint64
|
||
|
|
}
|
||
|
|
|
||
|
|
func newNode() *node {
|
||
|
|
return &node{
|
||
|
|
children: make(map[byte]*node),
|
||
|
|
ids: []uint64{},
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Trie datastructure.
|
||
|
|
type Trie struct {
|
||
|
|
root *node
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewTrie returns Trie.
|
||
|
|
func NewTrie() *Trie {
|
||
|
|
return &Trie{
|
||
|
|
root: newNode(),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Add adds the id in the trie for the given prefix path.
|
||
|
|
func (t *Trie) Add(prefix []byte, id uint64) {
|
||
|
|
node := t.root
|
||
|
|
for _, val := range prefix {
|
||
|
|
child, ok := node.children[val]
|
||
|
|
if !ok {
|
||
|
|
child = newNode()
|
||
|
|
node.children[val] = child
|
||
|
|
}
|
||
|
|
node = child
|
||
|
|
}
|
||
|
|
// We only need to add the id to the last node of the given prefix.
|
||
|
|
node.ids = append(node.ids, id)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Get returns prefix matched ids for the given key.
|
||
|
|
func (t *Trie) Get(key []byte) map[uint64]struct{} {
|
||
|
|
out := make(map[uint64]struct{})
|
||
|
|
node := t.root
|
||
|
|
// If root has ids that means we have subscribers for "nil/[]byte{}"
|
||
|
|
// prefix. Add them to the list.
|
||
|
|
if len(node.ids) > 0 {
|
||
|
|
for _, i := range node.ids {
|
||
|
|
out[i] = struct{}{}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
for _, val := range key {
|
||
|
|
child, ok := node.children[val]
|
||
|
|
if !ok {
|
||
|
|
break
|
||
|
|
}
|
||
|
|
// We need ids of the all the node in the matching key path.
|
||
|
|
for _, id := range child.ids {
|
||
|
|
out[id] = struct{}{}
|
||
|
|
}
|
||
|
|
node = child
|
||
|
|
}
|
||
|
|
return out
|
||
|
|
}
|
||
|
|
|
||
|
|
// Delete will delete the id if the id exist in the given index path.
|
||
|
|
func (t *Trie) Delete(index []byte, id uint64) {
|
||
|
|
node := t.root
|
||
|
|
for _, val := range index {
|
||
|
|
child, ok := node.children[val]
|
||
|
|
if !ok {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
node = child
|
||
|
|
}
|
||
|
|
// We're just removing the id not the hanging path.
|
||
|
|
out := node.ids[:0]
|
||
|
|
for _, val := range node.ids {
|
||
|
|
if val != id {
|
||
|
|
out = append(out, val)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
for i := len(out); i < len(node.ids); i++ {
|
||
|
|
node.ids[i] = 0 // garbage collecting
|
||
|
|
}
|
||
|
|
node.ids = out
|
||
|
|
}
|