feat(tui): add flamegraph bytes metric toggle

2 files changed, 87 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 {
diff --git a/internal/flamegraph/livetrie_test.go b/internal/flamegraph/livetrie_test.go
index 71f645c..53bdf1f 100644
--- a/internal/flamegraph/livetrie_test.go
+++ b/internal/flamegraph/livetrie_test.go
@@ -223,6 +223,54 @@ func TestLiveTrieReconfigureRejectsInvalidFields(t *testing.T) {
 	}
 }
 
+func TestLiveTrieSetCountFieldSwitchesMetricAndResetsBaseline(t *testing.T) {
+	lt := NewLiveTrie([]string{"comm"}, "count")
+	lt.Ingest(newTestPair("svc", 42, 1001, "/tmp/a", 10, 1, 64))
+
+	initial := decodeLiveSnapshot(t, lt)
+	if got, want := initial.Total, uint64(1); got != want {
+		t.Fatalf("count snapshot total = %d, want %d", got, want)
+	}
+
+	if err := lt.SetCountField("bytes"); err != nil {
+		t.Fatalf("set count field: %v", err)
+	}
+	if got, want := lt.CountField(), "bytes"; got != want {
+		t.Fatalf("count field = %q, want %q", got, want)
+	}
+
+	empty := decodeLiveSnapshot(t, lt)
+	if got := empty.Total; got != 0 {
+		t.Fatalf("expected reset baseline after metric switch, total=%d", got)
+	}
+
+	lt.Ingest(newTestPair("svc", 42, 1002, "/tmp/b", 10, 1, 64))
+	bytesSnap := decodeLiveSnapshot(t, lt)
+	if got, want := bytesSnap.Total, uint64(64); got != want {
+		t.Fatalf("bytes snapshot total = %d, want %d", got, want)
+	}
+	leaf := findSnapshotPath(t, &bytesSnap, "svc")
+	if got, want := leaf.Total, uint64(64); got != want {
+		t.Fatalf("bytes leaf total = %d, want %d", got, want)
+	}
+}
+
+func TestLiveTrieSetCountFieldRejectsInvalidValue(t *testing.T) {
+	lt := NewLiveTrie([]string{"comm"}, "count")
+	lt.Ingest(newTestPair("svc", 42, 1001, "/tmp/a", 1, 1, 1))
+	beforeVersion := lt.Version()
+
+	if err := lt.SetCountField("bogus"); err == nil {
+		t.Fatalf("expected invalid count field error")
+	}
+	if got, want := lt.CountField(), "count"; got != want {
+		t.Fatalf("count field changed unexpectedly: got %q want %q", got, want)
+	}
+	if got := lt.Version(); got != beforeVersion {
+		t.Fatalf("version changed on invalid count field: got %d want %d", got, beforeVersion)
+	}
+}
+
 func TestLiveTrieSnapshotJSONCaching(t *testing.T) {
 	lt := NewLiveTrie([]string{"comm"}, "count")
 	lt.Ingest(newTestPair("svc", 42, 1001, "/tmp/a", 1, 1, 1))