tui: fix flamegraph click re-rooting

1 files changed, 18 insertions, 4 deletions

diff --git a/internal/tui/flamegraph/renderer.go b/internal/tui/flamegraph/renderer.go
index 80390fe..f9f6a89 100644
--- a/internal/tui/flamegraph/renderer.go
+++ b/internal/tui/flamegraph/renderer.go
@@ -38,11 +38,11 @@ func buildTerminalLayoutWithPath(snapshot *snapshotNode, width, height int, root
 		rootName = rootPath
 	}
 	frames := make([]tuiFrame, 0, len(snapshot.Children)+1)
-	collectTerminalLayout(&frames, snapshot, rootTotal, height, 0, 0, rootName, width)
+	collectTerminalLayout(&frames, snapshot, rootTotal, height, 0, 0, rootName, width, rootPath != "")
 	return frames
 }
 
-func collectTerminalLayout(out *[]tuiFrame, node *snapshotNode, rootTotal uint64, height, depth, col int, path string, span int) {
+func collectTerminalLayout(out *[]tuiFrame, node *snapshotNode, rootTotal uint64, height, depth, col int, path string, span int, normalizeRootChildren bool) {
 	if node == nil || depth >= height {
 		return
 	}
@@ -68,7 +68,13 @@ func collectTerminalLayout(out *[]tuiFrame, node *snapshotNode, rootTotal uint64
 		return
 	}
 
-	childWidths := allocateChildWidths(node.Children, total, span)
+	layoutTotal := total
+	if normalizeRootChildren && depth == 0 {
+		if childrenTotal := childSnapshotTotal(node.Children); childrenTotal > 0 {
+			layoutTotal = childrenTotal
+		}
+	}
+	childWidths := allocateChildWidths(node.Children, layoutTotal, span)
 	cursor := col
 	for idx, child := range node.Children {
 		childWidth := childWidths[idx]
@@ -77,11 +83,19 @@ func collectTerminalLayout(out *[]tuiFrame, node *snapshotNode, rootTotal uint64
 		}
 		childName := frameName(child.Name, depth+1)
 		childPath := strings.Join([]string{path, childName}, pathSeparator)
-		collectTerminalLayout(out, child, rootTotal, height, depth+1, cursor, childPath, childWidth)
+		collectTerminalLayout(out, child, rootTotal, height, depth+1, cursor, childPath, childWidth, false)
 		cursor += childWidth
 	}
 }
 
+func childSnapshotTotal(children []*snapshotNode) uint64 {
+	total := uint64(0)
+	for _, child := range children {
+		total += snapshotTotal(child)
+	}
+	return total
+}
+
 func allocateChildWidths(children []*snapshotNode, parentTotal uint64, span int) []int {
 	widths := make([]int, len(children))
 	if span <= 0 || parentTotal == 0 || len(children) == 0 {