1 files changed, 39 insertions, 0 deletions
diff --git a/internal/flamegraph/livetrie.go b/internal/flamegraph/livetrie.go
index 13d7de9..9f1fd91 100644
--- a/internal/flamegraph/livetrie.go
+++ b/internal/flamegraph/livetrie.go
@@ -42,6 +42,9 @@ type LiveTrie struct {
 
 // NewLiveTrie constructs an empty live trie with the configured frame/count fields.
 func NewLiveTrie(fields []string, countField string) *LiveTrie {
+	if !isLiveTrieCountField(countField) {
+		countField = "count"
+	}
 	return &LiveTrie{
 		root: &trieNode{
 			childMap: make(map[string]*trieNode),
@@ -123,6 +126,33 @@ func (lt *LiveTrie) Fields() []string {
 	return out
 }
 
+// CountField returns the active metric used to aggregate node values.
+func (lt *LiveTrie) CountField() string {
+	lt.mu.RLock()
+	field := lt.countField
+	lt.mu.RUnlock()
+	return field
+}
+
+// SetCountField changes the active aggregation metric and starts a new baseline.
+func (lt *LiveTrie) SetCountField(countField string) error {
+	field := strings.TrimSpace(countField)
+	if !isLiveTrieCountField(field) {
+		return fmt.Errorf("invalid count field %q", countField)
+	}
+
+	lt.mu.Lock()
+	if lt.countField == field {
+		lt.mu.Unlock()
+		return nil
+	}
+	lt.countField = field
+	lt.resetLocked()
+	lt.mu.Unlock()
+	lt.invalidateCache()
+	return nil
+}
+
 // Reconfigure changes frame fields and clears accumulated data for a new baseline.
 func (lt *LiveTrie) Reconfigure(fields []string) error {
 	normalized, err := normalizeLiveTrieFields(fields)
@@ -239,6 +269,15 @@ func isLiveTrieField(field string) bool {
 	}
 }
 
+func isLiveTrieCountField(field string) bool {
+	switch field {
+	case "count", "duration", "durationToPrev", "bytes":
+		return true
+	default:
+		return false
+	}
+}
+
 func subtreeTotal(node *trieNode) uint64 {
 	total := node.value
 	for _, child := range node.children {