tui: fix flamegraph click re-rooting

4 files changed, 183 insertions, 39 deletions

diff --git a/internal/tui/flamegraph/model.go b/internal/tui/flamegraph/model.go
index c001e98..d73bd65 100644
--- a/internal/tui/flamegraph/model.go
+++ b/internal/tui/flamegraph/model.go
@@ -528,24 +528,12 @@ func (m *Model) zoomIn() {
 		m.statusMessage = "Zoom unchanged: selected frame is current view root"
 		return
 	}
-	target := findNodeByPath(m.snapshot, selectedPath)
-	if target == nil {
+	prevRootPath := m.zoomPath
+	if !m.setZoomPath(selectedPath) {
 		m.statusMessage = "Zoom failed: selected node is unavailable"
 		return
 	}
-	selectedWidth := m.frames[m.selectedIdx].Width
-	if selectedWidth < 1 {
-		selectedWidth = 1
-	}
-	m.zoomStack = append(m.zoomStack, zoomState{
-		path:                m.zoomPath,
-		previousSelectedIdx: m.selectedIdx,
-		lineWidth:           m.zoomLineWidth,
-	})
-	m.zoomRoot = target
-	m.zoomPath = selectedPath
-	m.zoomLineWidth = selectedWidth
-	m.rebuildFrames(false)
+	m.zoomStack = append(m.zoomStack, zoomState{path: prevRootPath})
 	m.statusMessage = "Zoom: " + compactFramePath(selectedPath)
 }
 
@@ -554,18 +542,13 @@ func (m *Model) zoomUndo() {
 		m.statusMessage = "Zoom undo unavailable"
 		return
 	}
-	last := m.zoomStack[len(m.zoomStack)-1]
-	m.zoomStack = m.zoomStack[:len(m.zoomStack)-1]
-	m.zoomPath = last.path
-	if m.zoomPath == "" {
-		m.zoomRoot = nil
-		m.zoomLineWidth = 0
-	} else {
-		m.zoomRoot = findNodeByPath(m.snapshot, m.zoomPath)
-		m.zoomLineWidth = last.lineWidth
+	lastIdx := len(m.zoomStack) - 1
+	last := m.zoomStack[lastIdx]
+	m.zoomStack = m.zoomStack[:lastIdx]
+	if !m.setZoomPath(last.path) {
+		m.statusMessage = "Zoom undo unavailable"
+		return
 	}
-	m.selectedIdx = last.previousSelectedIdx
-	m.rebuildFrames(false)
 	if m.zoomPath == "" {
 		m.statusMessage = "Zoom: root"
 		return
@@ -1058,22 +1041,76 @@ func (m *Model) handleMouseClick(msg tea.MouseClickMsg) bool {
 	}
 	clickedPath := m.frames[idx].Path
 	currentRoot := m.currentRootPath()
-	if m.zoomPath != "" && (clickedPath == currentRoot || hasPathBoundaryPrefix(currentRoot, clickedPath)) {
-		for steps := 0; steps < len(m.zoomStack)+1 && m.currentRootPath() != clickedPath; steps++ {
-			m.zoomUndo()
-		}
-		if sel := m.frameIndexByPath(clickedPath); sel >= 0 {
-			m.selectedIdx = sel
-		}
+	if clickedPath == currentRoot {
+		m.selectedIdx = idx
 		m.subtreeSet = computeSubtreeSetInto(m.frames, m.selectedIdx, m.subtreeSet)
 		return true
 	}
-	m.selectedIdx = idx
+	if m.zoomPath != "" && hasPathBoundaryPrefix(currentRoot, clickedPath) {
+		if !m.setZoomPath(clickedPath) {
+			return false
+		}
+		m.zoomStack = buildZoomStack(clickedPath)
+	} else {
+		prevRootPath := m.zoomPath
+		if !m.setZoomPath(clickedPath) {
+			return false
+		}
+		m.zoomStack = append(m.zoomStack, zoomState{path: prevRootPath})
+	}
+	if sel := m.frameIndexByPath(clickedPath); sel >= 0 {
+		m.selectedIdx = sel
+	}
 	m.subtreeSet = computeSubtreeSetInto(m.frames, m.selectedIdx, m.subtreeSet)
-	m.zoomIn()
+	m.statusMessage = "Zoom: " + compactFramePath(clickedPath)
 	return true
 }
 
+func (m *Model) setZoomPath(path string) bool {
+	if m.snapshot == nil {
+		return false
+	}
+	rootPath := m.rootSnapshotPath()
+	if path == "" || path == rootPath {
+		m.zoomRoot = nil
+		m.zoomPath = ""
+		m.zoomLineWidth = 0
+		m.rebuildFrames(false)
+		return true
+	}
+	target := findNodeByPath(m.snapshot, path)
+	if target == nil {
+		return false
+	}
+	m.zoomRoot = target
+	m.zoomPath = path
+	m.zoomLineWidth = 0
+	m.rebuildFrames(false)
+	return true
+}
+
+func (m Model) rootSnapshotPath() string {
+	if m.snapshot != nil {
+		return frameName(m.snapshot.Name, 0)
+	}
+	if len(m.frames) > 0 {
+		return m.frames[0].Path
+	}
+	return ""
+}
+
+func buildZoomStack(path string) []zoomState {
+	parts := strings.Split(path, pathSeparator)
+	if len(parts) <= 1 {
+		return nil
+	}
+	stack := []zoomState{{path: ""}}
+	for idx := 1; idx < len(parts)-1; idx++ {
+		stack = append(stack, zoomState{path: strings.Join(parts[:idx+1], pathSeparator)})
+	}
+	return stack
+}
+
 func (m Model) frameIndexAt(x, y int) int {
 	if len(m.frames) == 0 || m.width <= 0 || m.height <= 0 {
 		return -1
diff --git a/internal/tui/flamegraph/model_test.go b/internal/tui/flamegraph/model_test.go
index 69c1387..2f98b73 100644
--- a/internal/tui/flamegraph/model_test.go
+++ b/internal/tui/flamegraph/model_test.go
@@ -378,6 +378,68 @@ func TestMouseClickOnLineageAncestorUndoesToThatZoomLevel(t *testing.T) {
 	}
 }
 
+func TestMouseClickDirectDeepZoomThenAncestorClickReRootsToAncestor(t *testing.T) {
+	m := newZoomModel()
+	deepPath := "root" + pathSeparator + "A" + pathSeparator + "A1"
+	deepIdx := mustFrameIndex(t, m.frames, deepPath)
+	x, y, ok := firstClickablePointForFrame(m, deepIdx)
+	if !ok {
+		t.Fatalf("expected clickable point for %q", deepPath)
+	}
+
+	next, _ := m.Update(tea.MouseClickMsg{X: x, Y: y, Button: tea.MouseLeft})
+	m = next.(Model)
+	if got := m.zoomPath; got != deepPath {
+		t.Fatalf("expected direct deep click to zoom into %q, got %q", deepPath, got)
+	}
+
+	ancestorPath := "root" + pathSeparator + "A"
+	ancestorIdx := mustFrameIndex(t, m.frames, ancestorPath)
+	x, y, ok = firstClickablePointForFrame(m, ancestorIdx)
+	if !ok {
+		t.Fatalf("expected clickable point for ancestor %q", ancestorPath)
+	}
+
+	next, _ = m.Update(tea.MouseClickMsg{X: x, Y: y, Button: tea.MouseLeft})
+	m = next.(Model)
+	if got := m.zoomPath; got != ancestorPath {
+		t.Fatalf("expected ancestor click to re-root to %q, got %q", ancestorPath, got)
+	}
+	if got := m.currentRootPath(); got != ancestorPath {
+		t.Fatalf("expected current root %q after ancestor click, got %q", ancestorPath, got)
+	}
+	if idx := mustFrameIndex(t, m.frames, ancestorPath); m.frames[idx].Width != m.width {
+		t.Fatalf("expected ancestor %q to span full width %d, got %d", ancestorPath, m.width, m.frames[idx].Width)
+	}
+}
+
+func TestMouseClickDirectDeepZoomUndoReturnsToRoot(t *testing.T) {
+	m := newZoomModel()
+	deepPath := "root" + pathSeparator + "A" + pathSeparator + "A1"
+	deepIdx := mustFrameIndex(t, m.frames, deepPath)
+	x, y, ok := firstClickablePointForFrame(m, deepIdx)
+	if !ok {
+		t.Fatalf("expected clickable point for %q", deepPath)
+	}
+
+	next, _ := m.Update(tea.MouseClickMsg{X: x, Y: y, Button: tea.MouseLeft})
+	m = next.(Model)
+	if got, want := m.zoomPath, deepPath; got != want {
+		t.Fatalf("expected direct deep click to zoom into %q, got %q", want, got)
+	}
+	if got, want := len(m.zoomStack), 1; got != want {
+		t.Fatalf("expected single undo step after direct deep click, got %d", got)
+	}
+
+	m.zoomUndo()
+	if m.zoomPath != "" {
+		t.Fatalf("expected undo after direct deep click to return to root, got %q", m.zoomPath)
+	}
+	if len(m.zoomStack) != 0 {
+		t.Fatalf("expected zoom stack cleared after undo to root, got %d", len(m.zoomStack))
+	}
+}
+
 func TestStaticFixtureArrowTraversalVisitsAllFrames(t *testing.T) {
 	trie := coreflamegraph.NewLiveTrie([]string{"comm", "path", "tracepoint"}, "count")
 	coreflamegraph.SeedTestFlameData(trie)
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 {
diff --git a/internal/tui/flamegraph/renderer_test.go b/internal/tui/flamegraph/renderer_test.go
index 7adc82d..354b40a 100644
--- a/internal/tui/flamegraph/renderer_test.go
+++ b/internal/tui/flamegraph/renderer_test.go
@@ -294,6 +294,37 @@ func TestRenderTerminalViewFilterKeepsNonMatchingBranchesVisible(t *testing.T) {
 	}
 }
 
+func TestBuildTerminalLayoutWithPathNormalizesZoomRootChildrenToFullWidth(t *testing.T) {
+	snapshot := &snapshotNode{
+		Name:  "root",
+		Total: 100,
+		Children: []*snapshotNode{
+			{
+				Name:  "zoom",
+				Total: 80,
+				Children: []*snapshotNode{
+					{Name: "left", Total: 10},
+					{Name: "right", Total: 10},
+				},
+			},
+			{Name: "other", Total: 20},
+		},
+	}
+
+	frames := buildTerminalLayoutWithPath(snapshot.Children[0], 120, 12, "root"+pathSeparator+"zoom")
+	left := mustFindFrame(t, frames, "root"+pathSeparator+"zoom"+pathSeparator+"left")
+	right := mustFindFrame(t, frames, "root"+pathSeparator+"zoom"+pathSeparator+"right")
+	if got := left.Width + right.Width; got != 120 {
+		t.Fatalf("expected zoom-root children to fill full width 120, got %d", got)
+	}
+	if left.Col != 0 {
+		t.Fatalf("expected left child to start at column 0, got %d", left.Col)
+	}
+	if right.Col != left.Width {
+		t.Fatalf("expected right child to start after left child, got %d want %d", right.Col, left.Width)
+	}
+}
+
 func TestFilterSampleCoverageAvoidsDoubleCountingNestedMatches(t *testing.T) {
 	frames := []tuiFrame{
 		{Path: "root", Total: 100},