/
groupby.go
105 lines (92 loc) · 1.97 KB
/
groupby.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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
package itertools
type groupBy struct {
curKey interface{}
curKeyOk bool
curValue interface{}
keyFunc func(interface{}) interface{}
input Iterator
groupKey interface{}
groupKeyOk bool
}
type Group interface {
Iterator
Key() interface{}
}
type group struct {
gb *groupBy
key interface{}
first bool
stopped bool
}
func (me *group) Stop() {
me.stopped = true
}
func (me *group) Next() (ok bool) {
if me.stopped {
return false
}
if me.first {
me.first = false
return true
}
me.gb.advance()
if !me.gb.curKeyOk || me.gb.curKey != me.key {
me.Stop()
return
}
ok = true
return
}
func (me group) Value() (ret interface{}) {
if me.stopped {
panic("iterator stopped")
}
ret = me.gb.curValue
return
}
func (me group) Key() interface{} {
return me.key
}
func (me *groupBy) advance() {
me.curKeyOk = me.input.Next()
if me.curKeyOk {
me.curValue = me.input.Value()
me.curKey = me.keyFunc(me.curValue)
}
}
func (me *groupBy) Next() (ok bool) {
for me.curKey == me.groupKey {
ok = me.input.Next()
if !ok {
return
}
me.curValue = me.input.Value()
me.curKey = me.keyFunc(me.curValue)
me.curKeyOk = true
}
me.groupKey = me.curKey
me.groupKeyOk = true
return true
}
func (me *groupBy) Value() (ret interface{}) {
return &group{me, me.groupKey, true, false}
}
func (me *groupBy) Stop() {
}
// Allows use of nil as a return from the key func.
var uniqueKey = new(int)
// Group by returns an iterator of iterators over the values of the input
// iterator that consecutively return the same value when input to the key
// function. Note that repeated calls to each value of the GroupBy Iterator
// does not return a new iterator over the values for that key.
func GroupBy(input Iterator, keyFunc func(interface{}) interface{}) Iterator {
if keyFunc == nil {
keyFunc = func(a interface{}) interface{} { return a }
}
return &groupBy{
input: input,
keyFunc: keyFunc,
groupKey: uniqueKey,
curKey: uniqueKey,
}
}