groupby.go 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105
  1. package iter
  2. type groupBy struct {
  3. curKey interface{}
  4. curKeyOk bool
  5. curValue interface{}
  6. keyFunc func(interface{}) interface{}
  7. input Iterator
  8. groupKey interface{}
  9. groupKeyOk bool
  10. }
  11. type Group interface {
  12. Iterator
  13. Key() interface{}
  14. }
  15. type group struct {
  16. gb *groupBy
  17. key interface{}
  18. first bool
  19. stopped bool
  20. }
  21. func (me *group) Stop() {
  22. me.stopped = true
  23. }
  24. func (me *group) Next() (ok bool) {
  25. if me.stopped {
  26. return false
  27. }
  28. if me.first {
  29. me.first = false
  30. return true
  31. }
  32. me.gb.advance()
  33. if !me.gb.curKeyOk || me.gb.curKey != me.key {
  34. me.Stop()
  35. return
  36. }
  37. ok = true
  38. return
  39. }
  40. func (me group) Value() (ret interface{}) {
  41. if me.stopped {
  42. panic("iterator stopped")
  43. }
  44. ret = me.gb.curValue
  45. return
  46. }
  47. func (me group) Key() interface{} {
  48. return me.key
  49. }
  50. func (me *groupBy) advance() {
  51. me.curKeyOk = me.input.Next()
  52. if me.curKeyOk {
  53. me.curValue = me.input.Value()
  54. me.curKey = me.keyFunc(me.curValue)
  55. }
  56. }
  57. func (me *groupBy) Next() (ok bool) {
  58. for me.curKey == me.groupKey {
  59. ok = me.input.Next()
  60. if !ok {
  61. return
  62. }
  63. me.curValue = me.input.Value()
  64. me.curKey = me.keyFunc(me.curValue)
  65. me.curKeyOk = true
  66. }
  67. me.groupKey = me.curKey
  68. me.groupKeyOk = true
  69. return true
  70. }
  71. func (me *groupBy) Value() (ret interface{}) {
  72. return &group{me, me.groupKey, true, false}
  73. }
  74. func (me *groupBy) Stop() {
  75. }
  76. // Allows use of nil as a return from the key func.
  77. var uniqueKey = new(int)
  78. // Group by returns an iterator of iterators over the values of the input
  79. // iterator that consecutively return the same value when input to the key
  80. // function. Note that repeated calls to each value of the GroupBy Iterator
  81. // does not return a new iterator over the values for that key.
  82. func GroupBy(input Iterator, keyFunc func(interface{}) interface{}) Iterator {
  83. if keyFunc == nil {
  84. keyFunc = func(a interface{}) interface{} { return a }
  85. }
  86. return &groupBy{
  87. input: input,
  88. keyFunc: keyFunc,
  89. groupKey: uniqueKey,
  90. curKey: uniqueKey,
  91. }
  92. }