summaryrefslogtreecommitdiff
path: root/internal/flamegraph/trie.go
blob: 6b31d23267ff7391a6439fd1f846ecc5fe25d88e (plain)
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
package flamegraph

import (
	"cmp"
	"slices"
)

type trieNode struct {
	name        string
	value       uint64
	heightValue uint64
	total       uint64
	heightTotal uint64
	children    []*trieNode
	childMap    map[string]*trieNode
}

type trie struct {
	root     *trieNode
	maxDepth int
}

func newTrie() *trie {
	return &trie{
		root: &trieNode{
			childMap: make(map[string]*trieNode),
		},
	}
}

func (t *trie) add(frames []string, value uint64) {
	insertTriePath(t.root, frames, value, value)
}

func (t *trie) computeTotals() {
	t.maxDepth = 0
	var walk func(node *trieNode, depth int) (uint64, uint64)
	walk = func(node *trieNode, depth int) (uint64, uint64) {
		if depth > t.maxDepth {
			t.maxDepth = depth
		}

		slices.SortFunc(node.children, func(a, b *trieNode) int {
			return cmp.Compare(a.name, b.name)
		})

		total := node.value
		heightTotal := node.heightValue
		for _, child := range node.children {
			childTotal, childHeightTotal := walk(child, depth+1)
			total += childTotal
			heightTotal += childHeightTotal
		}
		node.total = total
		node.heightTotal = heightTotal
		node.childMap = nil
		return total, heightTotal
	}

	walk(t.root, 0)
}