forked from hashicorp/nomad
-
Notifications
You must be signed in to change notification settings - Fork 0
/
select.go
85 lines (72 loc) · 1.76 KB
/
select.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
package scheduler
// LimitIterator is a RankIterator used to limit the number of options
// that are returned before we artifically end the stream.
type LimitIterator struct {
ctx Context
source RankIterator
limit int
seen int
}
// NewLimitIterator is returns a LimitIterator with a fixed limit of returned options
func NewLimitIterator(ctx Context, source RankIterator, limit int) *LimitIterator {
iter := &LimitIterator{
ctx: ctx,
source: source,
limit: limit,
}
return iter
}
func (iter *LimitIterator) SetLimit(limit int) {
iter.limit = limit
}
func (iter *LimitIterator) Next() *RankedNode {
if iter.seen == iter.limit {
return nil
}
option := iter.source.Next()
if option == nil {
return nil
}
iter.seen += 1
return option
}
func (iter *LimitIterator) Reset() {
iter.source.Reset()
iter.seen = 0
}
// MaxScoreIterator is a RankIterator used to return only a single result
// of the item with the highest score. This iterator will consume all of the
// possible inputs and only returns the highest ranking result.
type MaxScoreIterator struct {
ctx Context
source RankIterator
max *RankedNode
}
// MaxScoreIterator returns a MaxScoreIterator over the given source
func NewMaxScoreIterator(ctx Context, source RankIterator) *MaxScoreIterator {
iter := &MaxScoreIterator{
ctx: ctx,
source: source,
}
return iter
}
func (iter *MaxScoreIterator) Next() *RankedNode {
// Check if we've found the max, return nil
if iter.max != nil {
return nil
}
// Consume and determine the max
for {
option := iter.source.Next()
if option == nil {
return iter.max
}
if iter.max == nil || option.Score > iter.max.Score {
iter.max = option
}
}
}
func (iter *MaxScoreIterator) Reset() {
iter.source.Reset()
iter.max = nil
}