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