1package store
2
3import (
4	"bytes"
5
6	"github.com/siddontang/ledisdb/store/driver"
7)
8
9const (
10	IteratorForward  uint8 = 0
11	IteratorBackward uint8 = 1
12)
13
14const (
15	RangeClose uint8 = 0x00
16	RangeLOpen uint8 = 0x01
17	RangeROpen uint8 = 0x10
18	RangeOpen  uint8 = 0x11
19)
20
21// min must less or equal than max
22//
23// range type:
24//
25//  close: [min, max]
26//  open: (min, max)
27//  lopen: (min, max]
28//  ropen: [min, max)
29//
30type Range struct {
31	Min []byte
32	Max []byte
33
34	Type uint8
35}
36
37type Limit struct {
38	Offset int
39	Count  int
40}
41
42type Iterator struct {
43	it driver.IIterator
44	st *Stat
45}
46
47// Returns a copy of key.
48func (it *Iterator) Key() []byte {
49	k := it.it.Key()
50	if k == nil {
51		return nil
52	}
53
54	return append([]byte{}, k...)
55}
56
57// Returns a copy of value.
58func (it *Iterator) Value() []byte {
59	v := it.it.Value()
60	if v == nil {
61		return nil
62	}
63
64	return append([]byte{}, v...)
65}
66
67// Returns a reference of key.
68// you must be careful that it will be changed after next iterate.
69func (it *Iterator) RawKey() []byte {
70	return it.it.Key()
71}
72
73// Returns a reference of value.
74// you must be careful that it will be changed after next iterate.
75func (it *Iterator) RawValue() []byte {
76	return it.it.Value()
77}
78
79// Copy key to b, if b len is small or nil, returns a new one.
80func (it *Iterator) BufKey(b []byte) []byte {
81	k := it.RawKey()
82	if k == nil {
83		return nil
84	}
85	if b == nil {
86		b = []byte{}
87	}
88
89	b = b[0:0]
90	return append(b, k...)
91}
92
93// Copy value to b, if b len is small or nil, returns a new one.
94func (it *Iterator) BufValue(b []byte) []byte {
95	v := it.RawValue()
96	if v == nil {
97		return nil
98	}
99
100	if b == nil {
101		b = []byte{}
102	}
103
104	b = b[0:0]
105	return append(b, v...)
106}
107
108func (it *Iterator) Close() {
109	if it.it != nil {
110		it.st.IterCloseNum.Add(1)
111		it.it.Close()
112		it.it = nil
113	}
114}
115
116func (it *Iterator) Valid() bool {
117	return it.it.Valid()
118}
119
120func (it *Iterator) Next() {
121	it.st.IterSeekNum.Add(1)
122	it.it.Next()
123}
124
125func (it *Iterator) Prev() {
126	it.st.IterSeekNum.Add(1)
127	it.it.Prev()
128}
129
130func (it *Iterator) SeekToFirst() {
131	it.st.IterSeekNum.Add(1)
132	it.it.First()
133}
134
135func (it *Iterator) SeekToLast() {
136	it.st.IterSeekNum.Add(1)
137	it.it.Last()
138}
139
140func (it *Iterator) Seek(key []byte) {
141	it.st.IterSeekNum.Add(1)
142	it.it.Seek(key)
143}
144
145// Finds by key, if not found, nil returns.
146func (it *Iterator) Find(key []byte) []byte {
147	it.Seek(key)
148	if it.Valid() {
149		k := it.RawKey()
150		if k == nil {
151			return nil
152		} else if bytes.Equal(k, key) {
153			return it.Value()
154		}
155	}
156
157	return nil
158}
159
160// Finds by key, if not found, nil returns, else a reference of value returns.
161// you must be careful that it will be changed after next iterate.
162func (it *Iterator) RawFind(key []byte) []byte {
163	it.Seek(key)
164	if it.Valid() {
165		k := it.RawKey()
166		if k == nil {
167			return nil
168		} else if bytes.Equal(k, key) {
169			return it.RawValue()
170		}
171	}
172
173	return nil
174}
175
176type RangeLimitIterator struct {
177	it *Iterator
178
179	r *Range
180	l *Limit
181
182	step int
183
184	//0 for IteratorForward, 1 for IteratorBackward
185	direction uint8
186}
187
188func (it *RangeLimitIterator) Key() []byte {
189	return it.it.Key()
190}
191
192func (it *RangeLimitIterator) Value() []byte {
193	return it.it.Value()
194}
195
196func (it *RangeLimitIterator) RawKey() []byte {
197	return it.it.RawKey()
198}
199
200func (it *RangeLimitIterator) RawValue() []byte {
201	return it.it.RawValue()
202}
203
204func (it *RangeLimitIterator) BufKey(b []byte) []byte {
205	return it.it.BufKey(b)
206}
207
208func (it *RangeLimitIterator) BufValue(b []byte) []byte {
209	return it.it.BufValue(b)
210}
211
212func (it *RangeLimitIterator) Valid() bool {
213	if it.l.Offset < 0 {
214		return false
215	} else if !it.it.Valid() {
216		return false
217	} else if it.l.Count >= 0 && it.step >= it.l.Count {
218		return false
219	}
220
221	if it.direction == IteratorForward {
222		if it.r.Max != nil {
223			r := bytes.Compare(it.it.RawKey(), it.r.Max)
224			if it.r.Type&RangeROpen > 0 {
225				return !(r >= 0)
226			} else {
227				return !(r > 0)
228			}
229		}
230	} else {
231		if it.r.Min != nil {
232			r := bytes.Compare(it.it.RawKey(), it.r.Min)
233			if it.r.Type&RangeLOpen > 0 {
234				return !(r <= 0)
235			} else {
236				return !(r < 0)
237			}
238		}
239	}
240
241	return true
242}
243
244func (it *RangeLimitIterator) Next() {
245	it.step++
246
247	if it.direction == IteratorForward {
248		it.it.Next()
249	} else {
250		it.it.Prev()
251	}
252}
253
254func (it *RangeLimitIterator) Close() {
255	it.it.Close()
256}
257
258func NewRangeLimitIterator(i *Iterator, r *Range, l *Limit) *RangeLimitIterator {
259	return rangeLimitIterator(i, r, l, IteratorForward)
260}
261
262func NewRevRangeLimitIterator(i *Iterator, r *Range, l *Limit) *RangeLimitIterator {
263	return rangeLimitIterator(i, r, l, IteratorBackward)
264}
265
266func NewRangeIterator(i *Iterator, r *Range) *RangeLimitIterator {
267	return rangeLimitIterator(i, r, &Limit{0, -1}, IteratorForward)
268}
269
270func NewRevRangeIterator(i *Iterator, r *Range) *RangeLimitIterator {
271	return rangeLimitIterator(i, r, &Limit{0, -1}, IteratorBackward)
272}
273
274func rangeLimitIterator(i *Iterator, r *Range, l *Limit, direction uint8) *RangeLimitIterator {
275	it := new(RangeLimitIterator)
276
277	it.it = i
278
279	it.r = r
280	it.l = l
281	it.direction = direction
282
283	it.step = 0
284
285	if l.Offset < 0 {
286		return it
287	}
288
289	if direction == IteratorForward {
290		if r.Min == nil {
291			it.it.SeekToFirst()
292		} else {
293			it.it.Seek(r.Min)
294
295			if r.Type&RangeLOpen > 0 {
296				if it.it.Valid() && bytes.Equal(it.it.RawKey(), r.Min) {
297					it.it.Next()
298				}
299			}
300		}
301	} else {
302		if r.Max == nil {
303			it.it.SeekToLast()
304		} else {
305			it.it.Seek(r.Max)
306
307			if !it.it.Valid() {
308				it.it.SeekToLast()
309			} else {
310				if !bytes.Equal(it.it.RawKey(), r.Max) {
311					it.it.Prev()
312				}
313			}
314
315			if r.Type&RangeROpen > 0 {
316				if it.it.Valid() && bytes.Equal(it.it.RawKey(), r.Max) {
317					it.it.Prev()
318				}
319			}
320		}
321	}
322
323	for i := 0; i < l.Offset; i++ {
324		if it.it.Valid() {
325			if it.direction == IteratorForward {
326				it.it.Next()
327			} else {
328				it.it.Prev()
329			}
330		}
331	}
332
333	return it
334}
335