1package iterator
2
3import (
4	"context"
5	"fmt"
6
7	"github.com/cayleygraph/cayley/graph"
8)
9
10var _ graph.Iterator = &Limit{}
11
12// Limit iterator will stop iterating if certain a number of values were encountered.
13// Zero and negative limit values means no limit.
14type Limit struct {
15	uid       uint64
16	limit     int64
17	count     int64
18	primaryIt graph.Iterator
19}
20
21func NewLimit(primaryIt graph.Iterator, limit int64) *Limit {
22	return &Limit{
23		uid:       NextUID(),
24		limit:     limit,
25		primaryIt: primaryIt,
26	}
27}
28
29func (it *Limit) UID() uint64 {
30	return it.uid
31}
32
33// Reset resets the internal iterators and the iterator itself.
34func (it *Limit) Reset() {
35	it.count = 0
36	it.primaryIt.Reset()
37}
38
39func (it *Limit) Tagger() *graph.Tagger {
40	return it.primaryIt.Tagger()
41}
42
43func (it *Limit) TagResults(dst map[string]graph.Value) {
44	it.primaryIt.TagResults(dst)
45}
46
47func (it *Limit) Clone() graph.Iterator {
48	return NewLimit(it.primaryIt.Clone(), it.limit)
49}
50
51// SubIterators returns a slice of the sub iterators.
52func (it *Limit) SubIterators() []graph.Iterator {
53	return []graph.Iterator{it.primaryIt}
54}
55
56// Next advances the Limit iterator. It will stop iteration if limit was reached.
57func (it *Limit) Next(ctx context.Context) bool {
58	graph.NextLogIn(it)
59	if it.limit > 0 && it.count >= it.limit {
60		return graph.NextLogOut(it, false)
61	}
62	if it.primaryIt.Next(ctx) {
63		it.count++
64		return graph.NextLogOut(it, true)
65	}
66	return graph.NextLogOut(it, false)
67}
68
69func (it *Limit) Err() error {
70	return it.primaryIt.Err()
71}
72
73func (it *Limit) Result() graph.Value {
74	return it.primaryIt.Result()
75}
76
77func (it *Limit) Contains(ctx context.Context, val graph.Value) bool {
78	return it.primaryIt.Contains(ctx, val) // FIXME(dennwc): limit is ignored in this case
79}
80
81// NextPath checks whether there is another path. Will call primary iterator
82// if limit is not reached yet.
83func (it *Limit) NextPath(ctx context.Context) bool {
84	if it.limit > 0 && it.count >= it.limit {
85		return false
86	}
87	if it.primaryIt.NextPath(ctx) {
88		it.count++
89		return true
90	}
91	return false
92}
93
94// Close closes the primary and all iterators.  It closes all subiterators
95// it can, but returns the first error it encounters.
96func (it *Limit) Close() error {
97	return it.primaryIt.Close()
98}
99
100func (it *Limit) Type() graph.Type { return graph.Limit }
101
102func (it *Limit) Optimize() (graph.Iterator, bool) {
103	optimizedPrimaryIt, optimized := it.primaryIt.Optimize()
104	if it.limit <= 0 { // no limit
105		return optimizedPrimaryIt, true
106	}
107	it.primaryIt = optimizedPrimaryIt
108	return it, optimized
109}
110
111func (it *Limit) Stats() graph.IteratorStats {
112	primaryStats := it.primaryIt.Stats()
113	if it.limit > 0 && primaryStats.Size > it.limit {
114		primaryStats.Size = it.limit
115	}
116	return primaryStats
117}
118
119func (it *Limit) Size() (int64, bool) {
120	primarySize, exact := it.primaryIt.Size()
121	if it.limit > 0 && primarySize > it.limit {
122		primarySize = it.limit
123	}
124	return primarySize, exact
125}
126
127func (it *Limit) String() string {
128	return fmt.Sprintf("Limit(%d)", it.limit)
129}
130