flamegraph: add trie aggregation and unit tests

1 files changed, 65 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)
+}