blob: 96fee5152b1cbb5b041e80c094c56c004b8907e6 [file] [log] [blame]
Matthias Andreas Benkard832a54e2019-01-29 09:27:38 +01001package diskv
2
3import (
4 "sync"
5
6 "github.com/google/btree"
7)
8
9// Index is a generic interface for things that can
10// provide an ordered list of keys.
11type Index interface {
12 Initialize(less LessFunction, keys <-chan string)
13 Insert(key string)
14 Delete(key string)
15 Keys(from string, n int) []string
16}
17
18// LessFunction is used to initialize an Index of keys in a specific order.
19type LessFunction func(string, string) bool
20
21// btreeString is a custom data type that satisfies the BTree Less interface,
22// making the strings it wraps sortable by the BTree package.
23type btreeString struct {
24 s string
25 l LessFunction
26}
27
28// Less satisfies the BTree.Less interface using the btreeString's LessFunction.
29func (s btreeString) Less(i btree.Item) bool {
30 return s.l(s.s, i.(btreeString).s)
31}
32
33// BTreeIndex is an implementation of the Index interface using google/btree.
34type BTreeIndex struct {
35 sync.RWMutex
36 LessFunction
37 *btree.BTree
38}
39
40// Initialize populates the BTree tree with data from the keys channel,
41// according to the passed less function. It's destructive to the BTreeIndex.
42func (i *BTreeIndex) Initialize(less LessFunction, keys <-chan string) {
43 i.Lock()
44 defer i.Unlock()
45 i.LessFunction = less
46 i.BTree = rebuild(less, keys)
47}
48
49// Insert inserts the given key (only) into the BTree tree.
50func (i *BTreeIndex) Insert(key string) {
51 i.Lock()
52 defer i.Unlock()
53 if i.BTree == nil || i.LessFunction == nil {
54 panic("uninitialized index")
55 }
56 i.BTree.ReplaceOrInsert(btreeString{s: key, l: i.LessFunction})
57}
58
59// Delete removes the given key (only) from the BTree tree.
60func (i *BTreeIndex) Delete(key string) {
61 i.Lock()
62 defer i.Unlock()
63 if i.BTree == nil || i.LessFunction == nil {
64 panic("uninitialized index")
65 }
66 i.BTree.Delete(btreeString{s: key, l: i.LessFunction})
67}
68
69// Keys yields a maximum of n keys in order. If the passed 'from' key is empty,
70// Keys will return the first n keys. If the passed 'from' key is non-empty, the
71// first key in the returned slice will be the key that immediately follows the
72// passed key, in key order.
73func (i *BTreeIndex) Keys(from string, n int) []string {
74 i.RLock()
75 defer i.RUnlock()
76
77 if i.BTree == nil || i.LessFunction == nil {
78 panic("uninitialized index")
79 }
80
81 if i.BTree.Len() <= 0 {
82 return []string{}
83 }
84
85 btreeFrom := btreeString{s: from, l: i.LessFunction}
86 skipFirst := true
87 if len(from) <= 0 || !i.BTree.Has(btreeFrom) {
88 // no such key, so fabricate an always-smallest item
89 btreeFrom = btreeString{s: "", l: func(string, string) bool { return true }}
90 skipFirst = false
91 }
92
93 keys := []string{}
94 iterator := func(i btree.Item) bool {
95 keys = append(keys, i.(btreeString).s)
96 return len(keys) < n
97 }
98 i.BTree.AscendGreaterOrEqual(btreeFrom, iterator)
99
100 if skipFirst && len(keys) > 0 {
101 keys = keys[1:]
102 }
103
104 return keys
105}
106
107// rebuildIndex does the work of regenerating the index
108// with the given keys.
109func rebuild(less LessFunction, keys <-chan string) *btree.BTree {
110 tree := btree.New(2)
111 for key := range keys {
112 tree.ReplaceOrInsert(btreeString{s: key, l: less})
113 }
114 return tree
115}