1//  Copyright (c) 2017 Couchbase, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// 		http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package vellum
16
17import (
18	"bytes"
19)
20
21// Iterator represents a means of visity key/value pairs in order.
22type Iterator interface {
23
24	// Current() returns the key/value pair currently pointed to.
25	// The []byte of the key is ONLY guaranteed to be valid until
26	// another call to Next/Seek/Close.  If you need it beyond that
27	// point you MUST make a copy.
28	Current() ([]byte, uint64)
29
30	// Next() advances the iterator to the next key/value pair.
31	// If no more key/value pairs exist, ErrIteratorDone is returned.
32	Next() error
33
34	// Seek() advances the iterator the specified key, or the next key
35	// if it does not exist.
36	// If no keys exist after that point, ErrIteratorDone is returned.
37	Seek(key []byte) error
38
39	// Reset resets the Iterator' internal state to allow for iterator
40	// reuse (e.g. pooling).
41	Reset(f *FST, startKeyInclusive, endKeyExclusive []byte, aut Automaton) error
42
43	// Close() frees any resources held by this iterator.
44	Close() error
45}
46
47// FSTIterator is a structure for iterating key/value pairs in this FST in
48// lexicographic order.  Iterators should be constructed with the FSTIterator
49// method on the parent FST structure.
50type FSTIterator struct {
51	f   *FST
52	aut Automaton
53
54	startKeyInclusive []byte
55	endKeyExclusive   []byte
56
57	statesStack    []fstState
58	keysStack      []byte
59	keysPosStack   []int
60	valsStack      []uint64
61	autStatesStack []int
62
63	nextStart []byte
64}
65
66func newIterator(f *FST, startKeyInclusive, endKeyExclusive []byte,
67	aut Automaton) (*FSTIterator, error) {
68
69	rv := &FSTIterator{}
70	err := rv.Reset(f, startKeyInclusive, endKeyExclusive, aut)
71	if err != nil {
72		return nil, err
73	}
74	return rv, nil
75}
76
77// Reset resets the Iterator' internal state to allow for iterator
78// reuse (e.g. pooling).
79func (i *FSTIterator) Reset(f *FST,
80	startKeyInclusive, endKeyExclusive []byte, aut Automaton) error {
81	if aut == nil {
82		aut = alwaysMatchAutomaton
83	}
84
85	i.f = f
86	i.startKeyInclusive = startKeyInclusive
87	i.endKeyExclusive = endKeyExclusive
88	i.aut = aut
89
90	return i.pointTo(startKeyInclusive)
91}
92
93// pointTo attempts to point us to the specified location
94func (i *FSTIterator) pointTo(key []byte) error {
95	// tried to seek before start
96	if bytes.Compare(key, i.startKeyInclusive) < 0 {
97		key = i.startKeyInclusive
98	}
99
100	// tried to see past end
101	if i.endKeyExclusive != nil &&
102		bytes.Compare(key, i.endKeyExclusive) > 0 {
103		key = i.endKeyExclusive
104	}
105
106	// reset any state, pointTo always starts over
107	i.statesStack = i.statesStack[:0]
108	i.keysStack = i.keysStack[:0]
109	i.keysPosStack = i.keysPosStack[:0]
110	i.valsStack = i.valsStack[:0]
111	i.autStatesStack = i.autStatesStack[:0]
112
113	root, err := i.f.decoder.stateAt(i.f.decoder.getRoot(), nil)
114	if err != nil {
115		return err
116	}
117
118	autStart := i.aut.Start()
119
120	maxQ := -1
121	// root is always part of the path
122	i.statesStack = append(i.statesStack, root)
123	i.autStatesStack = append(i.autStatesStack, autStart)
124	for j := 0; j < len(key); j++ {
125		keyJ := key[j]
126		curr := i.statesStack[len(i.statesStack)-1]
127		autCurr := i.autStatesStack[len(i.autStatesStack)-1]
128
129		pos, nextAddr, nextVal := curr.TransitionFor(keyJ)
130		if nextAddr == noneAddr {
131			// needed transition doesn't exist
132			// find last trans before the one we needed
133			for q := curr.NumTransitions() - 1; q >= 0; q-- {
134				if curr.TransitionAt(q) < keyJ {
135					maxQ = q
136					break
137				}
138			}
139			break
140		}
141		autNext := i.aut.Accept(autCurr, keyJ)
142
143		next, err := i.f.decoder.stateAt(nextAddr, nil)
144		if err != nil {
145			return err
146		}
147
148		i.statesStack = append(i.statesStack, next)
149		i.keysStack = append(i.keysStack, keyJ)
150		i.keysPosStack = append(i.keysPosStack, pos)
151		i.valsStack = append(i.valsStack, nextVal)
152		i.autStatesStack = append(i.autStatesStack, autNext)
153		continue
154	}
155
156	if !i.statesStack[len(i.statesStack)-1].Final() ||
157		!i.aut.IsMatch(i.autStatesStack[len(i.autStatesStack)-1]) ||
158		bytes.Compare(i.keysStack, key) < 0 {
159		return i.next(maxQ)
160	}
161
162	return nil
163}
164
165// Current returns the key and value currently pointed to by the iterator.
166// If the iterator is not pointing at a valid value (because Iterator/Next/Seek)
167// returned an error previously, it may return nil,0.
168func (i *FSTIterator) Current() ([]byte, uint64) {
169	curr := i.statesStack[len(i.statesStack)-1]
170	if curr.Final() {
171		var total uint64
172		for _, v := range i.valsStack {
173			total += v
174		}
175		total += curr.FinalOutput()
176		return i.keysStack, total
177	}
178	return nil, 0
179}
180
181// Next advances this iterator to the next key/value pair.  If there is none
182// or the advancement goes beyond the configured endKeyExclusive, then
183// ErrIteratorDone is returned.
184func (i *FSTIterator) Next() error {
185	return i.next(-1)
186}
187
188func (i *FSTIterator) next(lastOffset int) error {
189	// remember where we started
190	i.nextStart = append(i.nextStart[:0], i.keysStack...)
191
192	nextOffset := lastOffset + 1
193
194OUTER:
195	for true {
196		curr := i.statesStack[len(i.statesStack)-1]
197		autCurr := i.autStatesStack[len(i.autStatesStack)-1]
198
199		if curr.Final() && i.aut.IsMatch(autCurr) &&
200			bytes.Compare(i.keysStack, i.nextStart) > 0 {
201			// in final state greater than start key
202			return nil
203		}
204
205		numTrans := curr.NumTransitions()
206
207	INNER:
208		for nextOffset < numTrans {
209			t := curr.TransitionAt(nextOffset)
210			autNext := i.aut.Accept(autCurr, t)
211			if !i.aut.CanMatch(autNext) {
212				nextOffset += 1
213				continue INNER
214			}
215
216			pos, nextAddr, v := curr.TransitionFor(t)
217
218			// the next slot in the statesStack might have an
219			// fstState instance that we can reuse
220			var nextPrealloc fstState
221			if len(i.statesStack) < cap(i.statesStack) {
222				nextPrealloc = i.statesStack[0:cap(i.statesStack)][len(i.statesStack)]
223			}
224
225			// push onto stack
226			next, err := i.f.decoder.stateAt(nextAddr, nextPrealloc)
227			if err != nil {
228				return err
229			}
230
231			i.statesStack = append(i.statesStack, next)
232			i.keysStack = append(i.keysStack, t)
233			i.keysPosStack = append(i.keysPosStack, pos)
234			i.valsStack = append(i.valsStack, v)
235			i.autStatesStack = append(i.autStatesStack, autNext)
236
237			// check to see if new keystack might have gone too far
238			if i.endKeyExclusive != nil &&
239				bytes.Compare(i.keysStack, i.endKeyExclusive) >= 0 {
240				return ErrIteratorDone
241			}
242
243			nextOffset = 0
244			continue OUTER
245		}
246
247		if len(i.statesStack) <= 1 {
248			// stack len is 1 (root), can't go back further, we're done
249			break
250		}
251
252		// no transitions, and still room to pop
253		i.statesStack = i.statesStack[:len(i.statesStack)-1]
254		i.keysStack = i.keysStack[:len(i.keysStack)-1]
255
256		nextOffset = i.keysPosStack[len(i.keysPosStack)-1] + 1
257
258		i.keysPosStack = i.keysPosStack[:len(i.keysPosStack)-1]
259		i.valsStack = i.valsStack[:len(i.valsStack)-1]
260		i.autStatesStack = i.autStatesStack[:len(i.autStatesStack)-1]
261	}
262
263	return ErrIteratorDone
264}
265
266// Seek advances this iterator to the specified key/value pair.  If this key
267// is not in the FST, Current() will return the next largest key.  If this
268// seek operation would go past the last key, or outside the configured
269// startKeyInclusive/endKeyExclusive then ErrIteratorDone is returned.
270func (i *FSTIterator) Seek(key []byte) error {
271	return i.pointTo(key)
272}
273
274// Close will free any resources held by this iterator.
275func (i *FSTIterator) Close() error {
276	// at the moment we don't do anything,
277	// but wanted this for API completeness
278	return nil
279}
280