-
Notifications
You must be signed in to change notification settings - Fork 672
/
merged_iterator.go
112 lines (93 loc) · 2.43 KB
/
merged_iterator.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
106
107
108
109
110
111
112
// Copyright (C) 2019-2023, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package state
import (
"container/heap"
)
var (
_ StakerIterator = (*mergedIterator)(nil)
_ heap.Interface = (*mergedIterator)(nil)
)
type mergedIterator struct {
initialized bool
// heap only contains iterators that have been initialized and are not
// exhausted.
heap []StakerIterator
}
// Returns an iterator that returns all of the elements of [stakers] in order.
func NewMergedIterator(stakers ...StakerIterator) StakerIterator {
// Filter out iterators that are already exhausted.
i := 0
for i < len(stakers) {
staker := stakers[i]
if staker.Next() {
i++
continue
}
staker.Release()
newLength := len(stakers) - 1
stakers[i] = stakers[newLength]
stakers[newLength] = nil
stakers = stakers[:newLength]
}
it := &mergedIterator{
heap: stakers,
}
heap.Init(it)
return it
}
func (it *mergedIterator) Next() bool {
if len(it.heap) == 0 {
return false
}
if !it.initialized {
// Note that on the first call to Next() (i.e. here) we don't call
// Next() on the current iterator. This is because we already called
// Next() on each iterator in NewMergedIterator.
it.initialized = true
return true
}
// Update the heap root.
current := it.heap[0]
if current.Next() {
// Calling Next() above modifies [current] so we fix the heap.
heap.Fix(it, 0)
return true
}
// The old root is exhausted. Remove it from the heap.
current.Release()
heap.Pop(it)
return len(it.heap) > 0
}
func (it *mergedIterator) Value() *Staker {
return it.heap[0].Value()
}
// When Release() returns, Release() has been called on each element of
// [stakers].
func (it *mergedIterator) Release() {
for _, it := range it.heap {
it.Release()
}
it.heap = nil
}
// Returns the number of sub-iterators in [it].
func (it *mergedIterator) Len() int {
return len(it.heap)
}
func (it *mergedIterator) Less(i, j int) bool {
return it.heap[i].Value().Less(it.heap[j].Value())
}
func (it *mergedIterator) Swap(i, j int) {
it.heap[j], it.heap[i] = it.heap[i], it.heap[j]
}
// Push is never actually used - but we need it to implement heap.Interface.
func (it *mergedIterator) Push(value interface{}) {
it.heap = append(it.heap, value.(StakerIterator))
}
func (it *mergedIterator) Pop() interface{} {
newLength := len(it.heap) - 1
value := it.heap[newLength]
it.heap[newLength] = nil
it.heap = it.heap[:newLength]
return value
}