flamegraph: add dual trie value/height totals (task po)

1 files changed, 16 insertions, 10 deletions

diff --git a/internal/flamegraph/trie.go b/internal/flamegraph/trie.go
index d7790c2..6b31d23 100644
--- a/internal/flamegraph/trie.go
+++ b/internal/flamegraph/trie.go
@@ -6,11 +6,13 @@ import (
 )
 
 type trieNode struct {
-	name     string
-	value    uint64
-	total    uint64
-	children []*trieNode
-	childMap map[string]*trieNode
+	name        string
+	value       uint64
+	heightValue uint64
+	total       uint64
+	heightTotal uint64
+	children    []*trieNode
+	childMap    map[string]*trieNode
 }
 
 type trie struct {
@@ -27,13 +29,13 @@ func newTrie() *trie {
 }
 
 func (t *trie) add(frames []string, value uint64) {
-	insertTriePath(t.root, frames, value)
+	insertTriePath(t.root, frames, value, value)
 }
 
 func (t *trie) computeTotals() {
 	t.maxDepth = 0
-	var walk func(node *trieNode, depth int) uint64
-	walk = func(node *trieNode, depth int) uint64 {
+	var walk func(node *trieNode, depth int) (uint64, uint64)
+	walk = func(node *trieNode, depth int) (uint64, uint64) {
 		if depth > t.maxDepth {
 			t.maxDepth = depth
 		}
@@ -43,12 +45,16 @@ func (t *trie) computeTotals() {
 		})
 
 		total := node.value
+		heightTotal := node.heightValue
 		for _, child := range node.children {
-			total += walk(child, depth+1)
+			childTotal, childHeightTotal := walk(child, depth+1)
+			total += childTotal
+			heightTotal += childHeightTotal
 		}
 		node.total = total
+		node.heightTotal = heightTotal
 		node.childMap = nil
-		return total
+		return total, heightTotal
 	}
 
 	walk(t.root, 0)