1
2
3
4
5
6
7 package testutil
8
9 import (
10 "fmt"
11 "math/rand"
12 "sort"
13 "strings"
14
15 "github.com/syndtr/goleveldb/leveldb/util"
16 )
17
18 type KeyValueEntry struct {
19 key, value []byte
20 }
21
22 type KeyValue struct {
23 entries []KeyValueEntry
24 nbytes int
25 }
26
27 func (kv *KeyValue) Put(key, value []byte) {
28 if n := len(kv.entries); n > 0 && cmp.Compare(kv.entries[n-1].key, key) >= 0 {
29 panic(fmt.Sprintf("Put: keys are not in increasing order: %q, %q", kv.entries[n-1].key, key))
30 }
31 kv.entries = append(kv.entries, KeyValueEntry{key, value})
32 kv.nbytes += len(key) + len(value)
33 }
34
35 func (kv *KeyValue) PutString(key, value string) {
36 kv.Put([]byte(key), []byte(value))
37 }
38
39 func (kv *KeyValue) PutU(key, value []byte) bool {
40 if i, exist := kv.Get(key); !exist {
41 if i < kv.Len() {
42 kv.entries = append(kv.entries[:i+1], kv.entries[i:]...)
43 kv.entries[i] = KeyValueEntry{key, value}
44 } else {
45 kv.entries = append(kv.entries, KeyValueEntry{key, value})
46 }
47 kv.nbytes += len(key) + len(value)
48 return true
49 } else {
50 kv.nbytes += len(value) - len(kv.ValueAt(i))
51 kv.entries[i].value = value
52 }
53 return false
54 }
55
56 func (kv *KeyValue) PutUString(key, value string) bool {
57 return kv.PutU([]byte(key), []byte(value))
58 }
59
60 func (kv *KeyValue) Delete(key []byte) (exist bool, value []byte) {
61 i, exist := kv.Get(key)
62 if exist {
63 value = kv.entries[i].value
64 kv.DeleteIndex(i)
65 }
66 return
67 }
68
69 func (kv *KeyValue) DeleteIndex(i int) bool {
70 if i < kv.Len() {
71 kv.nbytes -= len(kv.KeyAt(i)) + len(kv.ValueAt(i))
72 kv.entries = append(kv.entries[:i], kv.entries[i+1:]...)
73 return true
74 }
75 return false
76 }
77
78 func (kv KeyValue) Len() int {
79 return len(kv.entries)
80 }
81
82 func (kv *KeyValue) Size() int {
83 return kv.nbytes
84 }
85
86 func (kv KeyValue) KeyAt(i int) []byte {
87 return kv.entries[i].key
88 }
89
90 func (kv KeyValue) ValueAt(i int) []byte {
91 return kv.entries[i].value
92 }
93
94 func (kv KeyValue) Index(i int) (key, value []byte) {
95 if i < 0 || i >= len(kv.entries) {
96 panic(fmt.Sprintf("Index #%d: out of range", i))
97 }
98 return kv.entries[i].key, kv.entries[i].value
99 }
100
101 func (kv KeyValue) IndexInexact(i int) (key_, key, value []byte) {
102 key, value = kv.Index(i)
103 var key0 []byte
104 var key1 = kv.KeyAt(i)
105 if i > 0 {
106 key0 = kv.KeyAt(i - 1)
107 }
108 key_ = BytesSeparator(key0, key1)
109 return
110 }
111
112 func (kv KeyValue) IndexOrNil(i int) (key, value []byte) {
113 if i >= 0 && i < len(kv.entries) {
114 return kv.entries[i].key, kv.entries[i].value
115 }
116 return nil, nil
117 }
118
119 func (kv KeyValue) IndexString(i int) (key, value string) {
120 key_, _value := kv.Index(i)
121 return string(key_), string(_value)
122 }
123
124 func (kv KeyValue) Search(key []byte) int {
125 return sort.Search(kv.Len(), func(i int) bool {
126 return cmp.Compare(kv.KeyAt(i), key) >= 0
127 })
128 }
129
130 func (kv KeyValue) SearchString(key string) int {
131 return kv.Search([]byte(key))
132 }
133
134 func (kv KeyValue) Get(key []byte) (i int, exist bool) {
135 i = kv.Search(key)
136 if i < kv.Len() && cmp.Compare(kv.KeyAt(i), key) == 0 {
137 exist = true
138 }
139 return
140 }
141
142 func (kv KeyValue) GetString(key string) (i int, exist bool) {
143 return kv.Get([]byte(key))
144 }
145
146 func (kv KeyValue) Iterate(fn func(i int, key, value []byte)) {
147 for i, x := range kv.entries {
148 fn(i, x.key, x.value)
149 }
150 }
151
152 func (kv KeyValue) IterateString(fn func(i int, key, value string)) {
153 kv.Iterate(func(i int, key, value []byte) {
154 fn(i, string(key), string(value))
155 })
156 }
157
158 func (kv KeyValue) IterateShuffled(rnd *rand.Rand, fn func(i int, key, value []byte)) {
159 ShuffledIndex(rnd, kv.Len(), 1, func(i int) {
160 fn(i, kv.entries[i].key, kv.entries[i].value)
161 })
162 }
163
164 func (kv KeyValue) IterateShuffledString(rnd *rand.Rand, fn func(i int, key, value string)) {
165 kv.IterateShuffled(rnd, func(i int, key, value []byte) {
166 fn(i, string(key), string(value))
167 })
168 }
169
170 func (kv KeyValue) IterateInexact(fn func(i int, key_, key, value []byte)) {
171 for i := range kv.entries {
172 key_, key, value := kv.IndexInexact(i)
173 fn(i, key_, key, value)
174 }
175 }
176
177 func (kv KeyValue) IterateInexactString(fn func(i int, key_, key, value string)) {
178 kv.IterateInexact(func(i int, key_, key, value []byte) {
179 fn(i, string(key_), string(key), string(value))
180 })
181 }
182
183 func (kv KeyValue) Clone() KeyValue {
184 return KeyValue{append([]KeyValueEntry{}, kv.entries...), kv.nbytes}
185 }
186
187 func (kv KeyValue) Slice(start, limit int) KeyValue {
188 if start < 0 || limit > kv.Len() {
189 panic(fmt.Sprintf("Slice %d .. %d: out of range", start, limit))
190 } else if limit < start {
191 panic(fmt.Sprintf("Slice %d .. %d: invalid range", start, limit))
192 }
193 return KeyValue{append([]KeyValueEntry{}, kv.entries[start:limit]...), kv.nbytes}
194 }
195
196 func (kv KeyValue) SliceKey(start, limit []byte) KeyValue {
197 start_ := 0
198 limit_ := kv.Len()
199 if start != nil {
200 start_ = kv.Search(start)
201 }
202 if limit != nil {
203 limit_ = kv.Search(limit)
204 }
205 return kv.Slice(start_, limit_)
206 }
207
208 func (kv KeyValue) SliceKeyString(start, limit string) KeyValue {
209 return kv.SliceKey([]byte(start), []byte(limit))
210 }
211
212 func (kv KeyValue) SliceRange(r *util.Range) KeyValue {
213 if r != nil {
214 return kv.SliceKey(r.Start, r.Limit)
215 }
216 return kv.Clone()
217 }
218
219 func (kv KeyValue) Range(start, limit int) (r util.Range) {
220 if kv.Len() > 0 {
221 if start == kv.Len() {
222 r.Start = BytesAfter(kv.KeyAt(start - 1))
223 } else {
224 r.Start = kv.KeyAt(start)
225 }
226 }
227 if limit < kv.Len() {
228 r.Limit = kv.KeyAt(limit)
229 }
230 return
231 }
232
233 func KeyValue_EmptyKey() *KeyValue {
234 kv := &KeyValue{}
235 kv.PutString("", "v")
236 return kv
237 }
238
239 func KeyValue_EmptyValue() *KeyValue {
240 kv := &KeyValue{}
241 kv.PutString("abc", "")
242 kv.PutString("abcd", "")
243 return kv
244 }
245
246 func KeyValue_OneKeyValue() *KeyValue {
247 kv := &KeyValue{}
248 kv.PutString("abc", "v")
249 return kv
250 }
251
252 func KeyValue_BigValue() *KeyValue {
253 kv := &KeyValue{}
254 kv.PutString("big1", strings.Repeat("1", 200000))
255 return kv
256 }
257
258 func KeyValue_SpecialKey() *KeyValue {
259 kv := &KeyValue{}
260 kv.PutString("\xff\xff", "v3")
261 return kv
262 }
263
264 func KeyValue_MultipleKeyValue() *KeyValue {
265 kv := &KeyValue{}
266 kv.PutString("a", "v")
267 kv.PutString("aa", "v1")
268 kv.PutString("aaa", "v2")
269 kv.PutString("aaacccccccccc", "v2")
270 kv.PutString("aaaccccccccccd", "v3")
271 kv.PutString("aaaccccccccccf", "v4")
272 kv.PutString("aaaccccccccccfg", "v5")
273 kv.PutString("ab", "v6")
274 kv.PutString("abc", "v7")
275 kv.PutString("abcd", "v8")
276 kv.PutString("accccccccccccccc", "v9")
277 kv.PutString("b", "v10")
278 kv.PutString("bb", "v11")
279 kv.PutString("bc", "v12")
280 kv.PutString("c", "v13")
281 kv.PutString("c1", "v13")
282 kv.PutString("czzzzzzzzzzzzzz", "v14")
283 kv.PutString("fffffffffffffff", "v15")
284 kv.PutString("g11", "v15")
285 kv.PutString("g111", "v15")
286 kv.PutString("g111\xff", "v15")
287 kv.PutString("zz", "v16")
288 kv.PutString("zzzzzzz", "v16")
289 kv.PutString("zzzzzzzzzzzzzzzz", "v16")
290 return kv
291 }
292
293 var keymap = []byte("012345678ABCDEFGHIJKLMNOPQRSTUVWXYabcdefghijklmnopqrstuvwxy")
294
295 func KeyValue_Generate(rnd *rand.Rand, n, incr, minlen, maxlen, vminlen, vmaxlen int) *KeyValue {
296 if rnd == nil {
297 rnd = NewRand()
298 }
299 if maxlen < minlen {
300 panic("max len should >= min len")
301 }
302
303 rrand := func(min, max int) int {
304 if min == max {
305 return max
306 }
307 return rnd.Intn(max-min) + min
308 }
309
310 kv := &KeyValue{}
311 endC := byte(len(keymap) - incr)
312 gen := make([]byte, 0, maxlen)
313 for i := 0; i < n; i++ {
314 m := rrand(minlen, maxlen)
315 last := gen
316 retry:
317 gen = last[:m]
318 if k := len(last); m > k {
319 for j := k; j < m; j++ {
320 gen[j] = 0
321 }
322 } else {
323 for j := m - 1; j >= 0; j-- {
324 c := last[j]
325 if c == endC {
326 continue
327 }
328 gen[j] = c + byte(incr)
329 for j++; j < m; j++ {
330 gen[j] = 0
331 }
332 goto ok
333 }
334 if m < maxlen {
335 m++
336 goto retry
337 }
338 panic(fmt.Sprintf("only able to generate %d keys out of %d keys, try increasing max len", kv.Len(), n))
339 ok:
340 }
341 key := make([]byte, m)
342 for j := 0; j < m; j++ {
343 key[j] = keymap[gen[j]]
344 }
345 value := make([]byte, rrand(vminlen, vmaxlen))
346 for n := copy(value, fmt.Sprintf("v%d", i)); n < len(value); n++ {
347 value[n] = 'x'
348 }
349 kv.Put(key, value)
350 }
351 return kv
352 }
353
View as plain text