flamegraph: add trie aggregation and unit tests

2 files changed, 194 insertions, 0 deletions

diff --git a/internal/flamegraph/trie.go b/internal/flamegraph/trie.go
new file mode 100644
index 0000000..dbd3de6
--- /dev/null
+++ b/internal/flamegraph/trie.go
@@ -0,0 +1,65 @@
+package flamegraph
+
+import "sort"
+
+type trieNode struct {
+	name     string
+	value    uint64
+	total    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) {
+	node := t.root
+	for _, frame := range frames {
+		child, ok := node.childMap[frame]
+		if !ok {
+			child = &trieNode{
+				name:     frame,
+				childMap: make(map[string]*trieNode),
+			}
+			node.children = append(node.children, child)
+			node.childMap[frame] = child
+		}
+		node = child
+	}
+	node.value += value
+}
+
+func (t *trie) computeTotals() {
+	t.maxDepth = 0
+	var walk func(node *trieNode, depth int) uint64
+	walk = func(node *trieNode, depth int) uint64 {
+		if depth > t.maxDepth {
+			t.maxDepth = depth
+		}
+
+		sort.Slice(node.children, func(i, j int) bool {
+			return node.children[i].name < node.children[j].name
+		})
+
+		total := node.value
+		for _, child := range node.children {
+			total += walk(child, depth+1)
+		}
+		node.total = total
+		node.childMap = nil
+		return total
+	}
+
+	walk(t.root, 0)
+}
diff --git a/internal/flamegraph/trie_test.go b/internal/flamegraph/trie_test.go
new file mode 100644
index 0000000..d7f27e6
--- /dev/null
+++ b/internal/flamegraph/trie_test.go
@@ -0,0 +1,129 @@
+package flamegraph
+
+import "testing"
+
+func findChild(node *trieNode, name string) *trieNode {
+	for _, child := range node.children {
+		if child.name == name {
+			return child
+		}
+	}
+	return nil
+}
+
+func TestTrieSingleStack(t *testing.T) {
+	tr := newTrie()
+	tr.add([]string{"a", "b", "c"}, 10)
+	tr.computeTotals()
+
+	a := findChild(tr.root, "a")
+	if a == nil {
+		t.Fatal("expected node a")
+	}
+	b := findChild(a, "b")
+	if b == nil {
+		t.Fatal("expected node b")
+	}
+	c := findChild(b, "c")
+	if c == nil {
+		t.Fatal("expected node c")
+	}
+
+	if tr.root.total != 10 || a.total != 10 || b.total != 10 || c.total != 10 {
+		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)
+	}
+}
+
+func TestTrieMultipleStacks(t *testing.T) {
+	tr := newTrie()
+	tr.add([]string{"a", "b"}, 5)
+	tr.add([]string{"a", "c"}, 3)
+	tr.computeTotals()
+
+	a := findChild(tr.root, "a")
+	if a == nil {
+		t.Fatal("expected node a")
+	}
+	if a.total != 8 {
+		t.Fatalf("expected a.total=8, got %d", a.total)
+	}
+}
+
+func TestTrieOverlappingPaths(t *testing.T) {
+	tr := newTrie()
+	tr.add([]string{"a", "b"}, 5)
+	tr.add([]string{"a", "b"}, 3)
+	tr.computeTotals()
+
+	a := findChild(tr.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 != 8 {
+		t.Fatalf("expected merged leaf value=8, got %d", b.value)
+	}
+	if tr.root.total != 8 {
+		t.Fatalf("expected root total=8, got %d", tr.root.total)
+	}
+}
+
+func TestTrieEmpty(t *testing.T) {
+	tr := newTrie()
+	tr.computeTotals()
+
+	if tr.root.total != 0 {
+		t.Fatalf("expected root.total=0, got %d", tr.root.total)
+	}
+}
+
+func TestTrieDeterministicChildOrder(t *testing.T) {
+	tr := newTrie()
+	tr.add([]string{"root", "z"}, 1)
+	tr.add([]string{"root", "a"}, 1)
+	tr.add([]string{"root", "m"}, 1)
+	tr.computeTotals()
+
+	root := findChild(tr.root, "root")
+	if root == nil {
+		t.Fatal("expected node root")
+	}
+	if len(root.children) != 3 {
+		t.Fatalf("expected 3 children, got %d", len(root.children))
+	}
+
+	if root.children[0].name != "a" || root.children[1].name != "m" || root.children[2].name != "z" {
+		t.Fatalf("expected alphabetical child order [a m z], got [%s %s %s]",
+			root.children[0].name, root.children[1].name, root.children[2].name)
+	}
+}
+
+func TestTrieMaxDepth(t *testing.T) {
+	tr := newTrie()
+	tr.add([]string{"a"}, 1)
+	tr.add([]string{"a", "b", "c", "d"}, 1)
+	tr.computeTotals()
+
+	if tr.maxDepth != 4 {
+		t.Fatalf("expected maxDepth=4, got %d", tr.maxDepth)
+	}
+}
+
+func TestTrieEmptyFramesAccumulateAtRoot(t *testing.T) {
+	tr := newTrie()
+	tr.add(nil, 2)
+	tr.add([]string{}, 3)
+	tr.computeTotals()
+
+	if tr.root.value != 5 {
+		t.Fatalf("expected root.value=5, got %d", tr.root.value)
+	}
+	if tr.root.total != 5 {
+		t.Fatalf("expected root.total=5, got %d", tr.root.total)
+	}
+}