mirror of
https://github.com/luxfi/util.git
synced 2026-10-30 15:13:56 +00:00
zeekayandHanzo Dev
b45161c43f
These three sit in luxfi/node under utils/, and P-Chain and X-Chain reach for them. Both chains have to leave that repository before it can be archived, and a chain should not have to depend on a node to measure a quantity in nLUX, keep a heap, or walk a sequence. units imports nothing at all; heap wanted only Zero, which is already here; iterator wanted heap and Zero. Three of iterator's tests exercise it against platformvm's state iterators. They stay with the node, and travel with platformvm, because a generic sequence should not be tested through one virtual machine's idea of state. What remains here covers the package on its own terms. Co-authored-by: Hanzo Dev <[email protected]>
90 lines
1.8 KiB
Go
90 lines
1.8 KiB
Go
// Copyright (C) 2019-2025, Lux Industries Inc. All rights reserved.
|
|
// See the file LICENSE for licensing terms.
|
|
|
|
package iterator
|
|
|
|
import (
|
|
"github.com/google/btree"
|
|
|
|
"github.com/luxfi/util/heap"
|
|
)
|
|
|
|
var _ Iterator[any] = (*merged[any])(nil)
|
|
|
|
type merged[T any] struct {
|
|
initialized bool
|
|
// heap only contains iterators that have been initialized and are not
|
|
// exhausted.
|
|
heap heap.Queue[Iterator[T]]
|
|
}
|
|
|
|
// Merge returns an iterator that returns all of the elements of [iterators] in
|
|
// order.
|
|
func Merge[T any](less btree.LessFunc[T], iterators ...Iterator[T]) Iterator[T] {
|
|
// Filter out iterators that are already exhausted.
|
|
i := 0
|
|
for i < len(iterators) {
|
|
it := iterators[i]
|
|
if it.Next() {
|
|
i++
|
|
continue
|
|
}
|
|
it.Release()
|
|
|
|
newLength := len(iterators) - 1
|
|
iterators[i] = iterators[newLength]
|
|
iterators[newLength] = nil
|
|
iterators = iterators[:newLength]
|
|
}
|
|
|
|
it := &merged[T]{
|
|
heap: heap.QueueOf(
|
|
func(a, b Iterator[T]) bool {
|
|
return less(a.Value(), b.Value())
|
|
},
|
|
iterators...,
|
|
),
|
|
}
|
|
|
|
return it
|
|
}
|
|
|
|
func (it *merged[_]) Next() bool {
|
|
if it.heap.Len() == 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 Merge.
|
|
it.initialized = true
|
|
return true
|
|
}
|
|
|
|
// Update the heap root.
|
|
current, _ := it.heap.Peek()
|
|
if current.Next() {
|
|
// Calling Next() above modifies [current] so we fix the heap.
|
|
it.heap.Fix(0)
|
|
return true
|
|
}
|
|
|
|
// The old root is exhausted. Remove it from the heap.
|
|
current.Release()
|
|
it.heap.Pop()
|
|
return it.heap.Len() > 0
|
|
}
|
|
|
|
func (it *merged[T]) Value() T {
|
|
peek, _ := it.heap.Peek()
|
|
return peek.Value()
|
|
}
|
|
|
|
func (it *merged[_]) Release() {
|
|
for it.heap.Len() > 0 {
|
|
removed, _ := it.heap.Pop()
|
|
removed.Release()
|
|
}
|
|
}
|