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

4 files changed, 60 insertions, 13 deletions

diff --git a/internal/flamegraph/livetrie.go b/internal/flamegraph/livetrie.go
index 5487b9f..a682a0a 100644
--- a/internal/flamegraph/livetrie.go
+++ b/internal/flamegraph/livetrie.go
@@ -63,7 +63,7 @@ func NewLiveTrie(fields []string, countField string) *LiveTrie {
 }
 
 func (lt *LiveTrie) addLocked(frames []string, value uint64) {
-	insertTriePath(lt.root, frames, value)
+	insertTriePath(lt.root, frames, value, value)
 	if len(frames) > lt.maxDepth {
 		lt.maxDepth = len(frames)
 	}
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)
diff --git a/internal/flamegraph/trie_insert.go b/internal/flamegraph/trie_insert.go
index 7748b4a..8c01792 100644
--- a/internal/flamegraph/trie_insert.go
+++ b/internal/flamegraph/trie_insert.go
@@ -1,7 +1,7 @@
 package flamegraph
 
-// insertTriePath follows or creates nodes for frames and adds value at the leaf.
-func insertTriePath(root *trieNode, frames []string, value uint64) {
+// insertTriePath follows or creates nodes for frames and adds values at the leaf.
+func insertTriePath(root *trieNode, frames []string, value, heightValue uint64) {
 	node := root
 	for _, frame := range frames {
 		if node.childMap == nil {
@@ -19,4 +19,5 @@ func insertTriePath(root *trieNode, frames []string, value uint64) {
 		node = child
 	}
 	node.value += value
+	node.heightValue += heightValue
 }
diff --git a/internal/flamegraph/trie_test.go b/internal/flamegraph/trie_test.go
index d7f27e6..47ef770 100644
--- a/internal/flamegraph/trie_test.go
+++ b/internal/flamegraph/trie_test.go
@@ -33,6 +33,10 @@ func TestTrieSingleStack(t *testing.T) {
 		t.Fatalf("expected total propagation to be 10, got root=%d a=%d b=%d c=%d",
 			tr.root.total, a.total, b.total, c.total)
 	}
+	if tr.root.heightTotal != 10 || a.heightTotal != 10 || b.heightTotal != 10 || c.heightTotal != 10 {
+		t.Fatalf("expected heightTotal propagation to be 10, got root=%d a=%d b=%d c=%d",
+			tr.root.heightTotal, a.heightTotal, b.heightTotal, c.heightTotal)
+	}
 }
 
 func TestTrieMultipleStacks(t *testing.T) {
@@ -68,9 +72,15 @@ func TestTrieOverlappingPaths(t *testing.T) {
 	if b.value != 8 {
 		t.Fatalf("expected merged leaf value=8, got %d", b.value)
 	}
+	if b.heightValue != 8 {
+		t.Fatalf("expected merged leaf heightValue=8, got %d", b.heightValue)
+	}
 	if tr.root.total != 8 {
 		t.Fatalf("expected root total=8, got %d", tr.root.total)
 	}
+	if tr.root.heightTotal != 8 {
+		t.Fatalf("expected root heightTotal=8, got %d", tr.root.heightTotal)
+	}
 }
 
 func TestTrieEmpty(t *testing.T) {
@@ -80,6 +90,9 @@ func TestTrieEmpty(t *testing.T) {
 	if tr.root.total != 0 {
 		t.Fatalf("expected root.total=0, got %d", tr.root.total)
 	}
+	if tr.root.heightTotal != 0 {
+		t.Fatalf("expected root.heightTotal=0, got %d", tr.root.heightTotal)
+	}
 }
 
 func TestTrieDeterministicChildOrder(t *testing.T) {
@@ -123,7 +136,34 @@ func TestTrieEmptyFramesAccumulateAtRoot(t *testing.T) {
 	if tr.root.value != 5 {
 		t.Fatalf("expected root.value=5, got %d", tr.root.value)
 	}
+	if tr.root.heightValue != 5 {
+		t.Fatalf("expected root.heightValue=5, got %d", tr.root.heightValue)
+	}
 	if tr.root.total != 5 {
 		t.Fatalf("expected root.total=5, got %d", tr.root.total)
 	}
+	if tr.root.heightTotal != 5 {
+		t.Fatalf("expected root.heightTotal=5, got %d", tr.root.heightTotal)
+	}
+}
+
+func TestInsertTriePathStoresValueAndHeightValueIndependently(t *testing.T) {
+	root := &trieNode{childMap: make(map[string]*trieNode)}
+	insertTriePath(root, []string{"a", "b"}, 7, 3)
+	insertTriePath(root, []string{"a", "b"}, 5, 2)
+
+	a := findChild(root, "a")
+	if a == nil {
+		t.Fatal("expected node a")
+	}
+	b := findChild(a, "b")
+	if b == nil {
+		t.Fatal("expected node b")
+	}
+	if b.value != 12 {
+		t.Fatalf("expected leaf value=12, got %d", b.value)
+	}
+	if b.heightValue != 5 {
+		t.Fatalf("expected leaf heightValue=5, got %d", b.heightValue)
+	}
 }