github.com/grafana/pyroscope@v1.18.0/pkg/util/loser/tree.go

// Loser tree, from https://en.wikipedia.org/wiki/K-way_merge_algorithm#Tournament_Tree

package loser

type Sequence interface {
	Next() bool // Advances and returns true if there is a value at this new position.
	Err() error // Returns any error encountered while advancing.
}

// New returns a new loser tree that merges the given sequences.
// The sequences must be sorted according to the less function already.
// The maxVal is used to initialize the tree.
// The at function returns the current value of the sequence.
// The less function compares two values.
// The close function is called on each sequence when the tree is closed or when a sequence returns false from Next().
// If any sequence returns an error from Err() after Next() = false, the tree will stop and return that error in Err().
// Examples:
//
//	tree := loser.New(...)
//	defer tree.Close()
//	for tree.Next() {
//		value := tree.Winner().At()
//		...
//	}
//	if err := tree.Err(); err != nil {
//		...
//	}
func New[E any, S Sequence](sequences []S, maxVal E, at func(S) E, less func(E, E) bool, close func(S)) *Tree[E, S] {
	nSequences := len(sequences)
	t := Tree[E, S]{
		maxVal: maxVal,
		at:     at,
		less:   less,
		close:  close,
		nodes:  make([]node[E, S], nSequences*2),
	}
	for i, s := range sequences {
		t.nodes[i+nSequences].items = s
		if !t.moveNext(i + nSequences) { // Must call Next on each item so that At() has a value.
			if t.err != nil {
				// error during initialize, requires us to close sequences not touched yet and mark nodes as uninitialized
				for j := i + 1; j < nSequences; j++ {
					t.close(sequences[j])
					t.nodes[j+nSequences].index = -1
				}
				break
			}
		}
	}
	if nSequences > 0 {
		t.nodes[0].index = -1 // flag to be initialized on first call to Next().
	}
	return &t
}

// Call the close function on all sequences that are still open.
func (t *Tree[E, S]) Close() {
	for _, e := range t.nodes[len(t.nodes)/2 : len(t.nodes)] {
		if e.index == -1 {
			continue
		}
		t.close(e.items)
	}
}

// A loser tree is a binary tree laid out such that nodes N and N+1 have parent N/2.
// We store M leaf nodes in positions M...2M-1, and M-1 internal nodes in positions 1..M-1.
// Node 0 is a special node, containing the winner of the contest.
type Tree[E any, S Sequence] struct {
	maxVal E
	at     func(S) E
	less   func(E, E) bool
	close  func(S) // Called when Next() returns false.
	nodes  []node[E, S]

	err error
}

type node[E any, S Sequence] struct {
	index int // This is the loser for all nodes except the 0th, where it is the winner.
	value E   // Value copied from the loser node, or winner for node 0.
	items S   // Only populated for leaf nodes.
}

func (t *Tree[E, S]) moveNext(index int) bool {
	n := &t.nodes[index]
	if n.items.Next() {
		n.value = t.at(n.items)
		return true
	}
	t.close(n.items) // Next() returned false; close it and mark as finished.
	t.err = n.items.Err()
	n.value = t.maxVal
	n.index = -1
	return false
}

func (t *Tree[E, S]) Winner() S {
	return t.nodes[t.nodes[0].index].items
}

func (t *Tree[E, S]) Next() bool {
	if len(t.nodes) == 0 || t.err != nil {
		return false
	}
	if t.nodes[0].index == -1 { // If tree has not been initialized yet, do that.
		t.initialize()
		return t.nodes[t.nodes[0].index].index != -1
	}
	if t.nodes[t.nodes[0].index].index == -1 { // already exhausted
		return false
	}
	if !t.moveNext(t.nodes[0].index) && t.err != nil {
		return false
	}
	t.replayGames(t.nodes[0].index)
	return t.nodes[t.nodes[0].index].index != -1
}

func (t *Tree[E, S]) Err() error {
	return t.err
}

func (t *Tree[E, S]) initialize() {
	winners := make([]int, len(t.nodes))
	// Initialize leaf nodes as winners to start.
	for i := len(t.nodes) / 2; i < len(t.nodes); i++ {
		winners[i] = i
	}
	for i := len(t.nodes) - 2; i > 0; i -= 2 {
		// At each stage the winners play each other, and we record the loser in the node.
		loser, winner := t.playGame(winners[i], winners[i+1])
		p := parent(i)
		t.nodes[p].index = loser
		t.nodes[p].value = t.nodes[loser].value
		winners[p] = winner
	}
	t.nodes[0].index = winners[1]
	t.nodes[0].value = t.nodes[winners[1]].value
}

// Starting at pos, re-consider all values up to the root.
func (t *Tree[E, S]) replayGames(pos int) {
	// At the start, pos is a leaf node, and is the winner at that level.
	n := parent(pos)
	for n != 0 {
		if t.less(t.nodes[n].value, t.nodes[pos].value) {
			loser := pos
			// Record pos as the loser here, and the old loser is the new winner.
			pos = t.nodes[n].index
			t.nodes[n].index = loser
			t.nodes[n].value = t.nodes[loser].value
		}
		n = parent(n)
	}
	// pos is now the winner; store it in node 0.
	t.nodes[0].index = pos
	t.nodes[0].value = t.nodes[pos].value
}

func (t *Tree[E, S]) playGame(a, b int) (loser, winner int) {
	if t.less(t.nodes[a].value, t.nodes[b].value) {
		return b, a
	}
	return a, b
}

func parent(i int) int { return i / 2 }

// Add a new sequence to the merge set
func (t *Tree[E, S]) Push(sequence S) error {
	// First, see if we can replace one that was previously finished.
	for newPos := len(t.nodes) / 2; newPos < len(t.nodes); newPos++ {
		if t.nodes[newPos].index == -1 {
			t.nodes[newPos].index = newPos
			t.nodes[newPos].items = sequence
			if !t.moveNext(newPos) {
				if t.err != nil {
					return t.err
				}
			}
			t.nodes[0].index = -1 // flag for re-initialize on next call to Next()
			return nil
		}
	}
	// We need to expand the tree. Pick the next biggest power of 2 to amortise resizing cost.
	size := 1
	for ; size <= len(t.nodes)/2; size *= 2 {
	}
	newPos := size + len(t.nodes)/2
	newNodes := make([]node[E, S], size*2)
	// Copy data over and fix up the indexes.
	for i, n := range t.nodes[len(t.nodes)/2:] {
		newNodes[i+size] = n
		newNodes[i+size].index = i + size
	}
	t.nodes = newNodes
	t.nodes[newPos].index = newPos
	t.nodes[newPos].items = sequence
	// Mark all the empty nodes we have added as finished.
	for i := newPos + 1; i < len(t.nodes); i++ {
		t.nodes[i].index = -1
		t.nodes[i].value = t.maxVal
	}
	if !t.moveNext(newPos) {
		if t.err != nil {
			return t.err
		}
	}
	t.nodes[0].index = -1 // flag for re-initialize on next call to Next()
	return nil
}