task 358: add flamegraph keyboard navigation

4 files changed, 233 insertions, 7 deletions

diff --git a/internal/tui/flamegraph/model.go b/internal/tui/flamegraph/model.go
index f13298e..8c8d434 100644
--- a/internal/tui/flamegraph/model.go
+++ b/internal/tui/flamegraph/model.go
@@ -6,7 +6,9 @@ import (
 	"image/color"
 	coreflamegraph "ior/internal/flamegraph"
 	common "ior/internal/tui/common"
+	"sort"
 
+	"charm.land/bubbles/v2/key"
 	tea "charm.land/bubbletea/v2"
 )
 
@@ -24,6 +26,22 @@ type zoomState struct {
 
 type frameSpring struct{}
 
+type flameKeyMap struct {
+	MoveShallower key.Binding
+	MoveDeeper    key.Binding
+	PrevSibling   key.Binding
+	NextSibling   key.Binding
+}
+
+func defaultFlameKeyMap() flameKeyMap {
+	return flameKeyMap{
+		MoveShallower: key.NewBinding(key.WithKeys("j", "down")),
+		MoveDeeper:    key.NewBinding(key.WithKeys("k", "up")),
+		PrevSibling:   key.NewBinding(key.WithKeys("h", "left")),
+		NextSibling:   key.NewBinding(key.WithKeys("l", "right")),
+	}
+}
+
 // Model is the Bubble Tea model for the TUI flamegraph tab.
 type Model struct {
 	liveTrie    *coreflamegraph.LiveTrie
@@ -42,6 +60,7 @@ type Model struct {
 	searchActive bool
 	searchQuery  string
 	matchIndices map[int]bool
+	subtreeSet   map[int]bool
 
 	fieldPresets [][]string
 	fieldIndex   int
@@ -49,6 +68,7 @@ type Model struct {
 	springs []frameSpring
 	paused  bool
 	isDark  bool
+	keys    flameKeyMap
 }
 
 // tuiFrame stores one terminal flamegraph frame cell.
@@ -69,12 +89,14 @@ func NewModel(liveTrie *coreflamegraph.LiveTrie) Model {
 	return Model{
 		liveTrie:     liveTrie,
 		matchIndices: make(map[int]bool),
+		subtreeSet:   make(map[int]bool),
 		fieldPresets: [][]string{
 			{"comm", "path"},
 			{"tracepoint", "comm", "path"},
 			{"pid", "tid", "comm", "path"},
 		},
 		isDark: true,
+		keys:   defaultFlameKeyMap(),
 	}
 }
 
@@ -84,13 +106,30 @@ func (m Model) Init() tea.Cmd {
 }
 
 // Update handles incoming messages.
-func (m Model) Update(tea.Msg) (tea.Model, tea.Cmd) {
+func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	switch msg := msg.(type) {
+	case tea.KeyPressMsg:
+		prev := m.selectedIdx
+		switch {
+		case key.Matches(msg, m.keys.MoveShallower):
+			m.moveVertical(-1)
+		case key.Matches(msg, m.keys.MoveDeeper):
+			m.moveVertical(1)
+		case key.Matches(msg, m.keys.PrevSibling):
+			m.moveSibling(-1)
+		case key.Matches(msg, m.keys.NextSibling):
+			m.moveSibling(1)
+		}
+		if m.selectedIdx != prev {
+			m.subtreeSet = computeSubtreeSet(m.frames, m.selectedIdx)
+		}
+	}
 	return m, nil
 }
 
 // View renders the flamegraph viewport.
 func (m Model) View() tea.View {
-	content := RenderTerminalView(m.frames, m.width, m.height, m.selectedIdx, m.matchIndices, m.isDark)
+	content := RenderTerminalView(m.frames, m.width, m.height, m.selectedIdx, m.subtreeSet, m.matchIndices, m.isDark)
 	if m.snapshot != nil && len(m.frames) == 0 {
 		content = common.PanelStyle.Render(fmt.Sprintf("Flame: snapshot v%d has no visible frames", m.lastVersion))
 	}
@@ -102,6 +141,10 @@ func (m *Model) SetLiveTrie(liveTrie *coreflamegraph.LiveTrie) {
 	m.liveTrie = liveTrie
 	m.lastVersion = 0
 	m.snapshot = nil
+	m.selectedIdx = 0
+	m.frames = nil
+	m.targetFrames = nil
+	m.subtreeSet = make(map[int]bool)
 }
 
 // RefreshFromLiveTrie loads a new snapshot when the source version changes.
@@ -121,6 +164,8 @@ func (m *Model) RefreshFromLiveTrie() bool {
 	}
 	m.targetFrames = BuildTerminalLayout(&snapshot, m.width, m.height)
 	m.frames = append(m.frames[:0], m.targetFrames...)
+	m.clampSelection()
+	m.subtreeSet = computeSubtreeSet(m.frames, m.selectedIdx)
 	m.snapshot = &snapshot
 	m.lastVersion = version
 	return true
@@ -135,9 +180,105 @@ func (m Model) LastVersion() uint64 {
 func (m *Model) SetViewport(width, height int) {
 	m.width = width
 	m.height = height
+	m.targetFrames = BuildTerminalLayout(m.snapshot, width, height)
+	m.frames = append(m.frames[:0], m.targetFrames...)
+	m.clampSelection()
+	m.subtreeSet = computeSubtreeSet(m.frames, m.selectedIdx)
 }
 
 // SetDarkMode sets the active color theme mode.
 func (m *Model) SetDarkMode(isDark bool) {
 	m.isDark = isDark
 }
+
+func (m *Model) moveVertical(delta int) {
+	if len(m.frames) == 0 {
+		return
+	}
+	m.clampSelection()
+	current := m.frames[m.selectedIdx]
+	targetDepth := current.Depth + delta
+	targets := framesAtDepth(m.frames, targetDepth)
+	if len(targets) == 0 {
+		return
+	}
+	best := targets[0]
+	bestDist := abs(m.frames[best].Col - current.Col)
+	for _, idx := range targets[1:] {
+		dist := abs(m.frames[idx].Col - current.Col)
+		if dist < bestDist {
+			best = idx
+			bestDist = dist
+		}
+	}
+	m.selectedIdx = best
+}
+
+func (m *Model) moveSibling(delta int) {
+	if len(m.frames) == 0 {
+		return
+	}
+	m.clampSelection()
+	current := m.frames[m.selectedIdx]
+	siblings := framesAtDepth(m.frames, current.Depth)
+	if len(siblings) <= 1 {
+		return
+	}
+	pos := indexOf(siblings, m.selectedIdx)
+	if pos < 0 {
+		return
+	}
+	next := pos + delta
+	if next < 0 {
+		next = 0
+	}
+	if next >= len(siblings) {
+		next = len(siblings) - 1
+	}
+	m.selectedIdx = siblings[next]
+}
+
+func framesAtDepth(frames []tuiFrame, depth int) []int {
+	if depth < 0 {
+		return nil
+	}
+	indices := make([]int, 0)
+	for idx, frame := range frames {
+		if frame.Depth == depth {
+			indices = append(indices, idx)
+		}
+	}
+	sort.Slice(indices, func(i, j int) bool {
+		return frames[indices[i]].Col < frames[indices[j]].Col
+	})
+	return indices
+}
+
+func indexOf(values []int, target int) int {
+	for idx, value := range values {
+		if value == target {
+			return idx
+		}
+	}
+	return -1
+}
+
+func (m *Model) clampSelection() {
+	if len(m.frames) == 0 {
+		m.selectedIdx = 0
+		return
+	}
+	if m.selectedIdx < 0 {
+		m.selectedIdx = 0
+	}
+	if m.selectedIdx >= len(m.frames) {
+		m.selectedIdx = len(m.frames) - 1
+	}
+}
+
+func abs(v int) int {
+	if v < 0 {
+		return -v
+	}
+	return v
+}
diff --git a/internal/tui/flamegraph/model_test.go b/internal/tui/flamegraph/model_test.go
index 1e472ae..4569296 100644
--- a/internal/tui/flamegraph/model_test.go
+++ b/internal/tui/flamegraph/model_test.go
@@ -3,6 +3,8 @@ package flamegraph
 import (
 	coreflamegraph "ior/internal/flamegraph"
 	"testing"
+
+	tea "charm.land/bubbletea/v2"
 )
 
 func TestNewModelDefaults(t *testing.T) {
@@ -48,3 +50,84 @@ func TestRefreshFromLiveTrieTracksVersionAndSnapshot(t *testing.T) {
 		t.Fatalf("expected no refresh when version is unchanged")
 	}
 }
+
+func TestKeyboardNavigationDeepNarrowTree(t *testing.T) {
+	m := NewModel(nil)
+	m.frames = []tuiFrame{
+		{Name: "root", Depth: 0, Col: 0, Path: "root"},
+		{Name: "child", Depth: 1, Col: 0, Path: "root" + pathSeparator + "child"},
+		{Name: "leaf", Depth: 2, Col: 0, Path: "root" + pathSeparator + "child" + pathSeparator + "leaf"},
+	}
+
+	m = pressFlameKey(t, m, tea.KeyPressMsg{Code: []rune{'k'}[0], Text: "k"})
+	if m.selectedIdx != 1 {
+		t.Fatalf("expected selection to move deeper to idx 1, got %d", m.selectedIdx)
+	}
+	m = pressFlameKey(t, m, tea.KeyPressMsg{Code: []rune{'k'}[0], Text: "k"})
+	if m.selectedIdx != 2 {
+		t.Fatalf("expected selection to move deeper to idx 2, got %d", m.selectedIdx)
+	}
+	m = pressFlameKey(t, m, tea.KeyPressMsg{Code: []rune{'j'}[0], Text: "j"})
+	if m.selectedIdx != 1 {
+		t.Fatalf("expected selection to move shallower to idx 1, got %d", m.selectedIdx)
+	}
+}
+
+func TestKeyboardNavigationShallowWideSiblings(t *testing.T) {
+	m := NewModel(nil)
+	m.frames = []tuiFrame{
+		{Name: "root", Depth: 0, Col: 0, Path: "root"},
+		{Name: "A", Depth: 1, Col: 0, Path: "root" + pathSeparator + "A"},
+		{Name: "B", Depth: 1, Col: 30, Path: "root" + pathSeparator + "B"},
+		{Name: "C", Depth: 1, Col: 60, Path: "root" + pathSeparator + "C"},
+	}
+
+	m = pressFlameKey(t, m, tea.KeyPressMsg{Code: []rune{'k'}[0], Text: "k"})
+	if m.selectedIdx != 1 {
+		t.Fatalf("expected first deeper frame to be A, got idx %d", m.selectedIdx)
+	}
+	m = pressFlameKey(t, m, tea.KeyPressMsg{Code: []rune{'l'}[0], Text: "l"})
+	if m.selectedIdx != 2 {
+		t.Fatalf("expected next sibling B, got idx %d", m.selectedIdx)
+	}
+	m = pressFlameKey(t, m, tea.KeyPressMsg{Code: []rune{'l'}[0], Text: "l"})
+	if m.selectedIdx != 3 {
+		t.Fatalf("expected next sibling C, got idx %d", m.selectedIdx)
+	}
+	m = pressFlameKey(t, m, tea.KeyPressMsg{Code: []rune{'l'}[0], Text: "l"})
+	if m.selectedIdx != 3 {
+		t.Fatalf("expected selection to clamp at last sibling, got idx %d", m.selectedIdx)
+	}
+	m = pressFlameKey(t, m, tea.KeyPressMsg{Code: []rune{'h'}[0], Text: "h"})
+	if m.selectedIdx != 2 {
+		t.Fatalf("expected previous sibling B, got idx %d", m.selectedIdx)
+	}
+}
+
+func TestKeyboardNavigationSingleNodeClamped(t *testing.T) {
+	m := NewModel(nil)
+	m.frames = []tuiFrame{{Name: "root", Depth: 0, Col: 0, Path: "root"}}
+
+	keys := []tea.KeyPressMsg{
+		{Code: []rune{'j'}[0], Text: "j"},
+		{Code: []rune{'k'}[0], Text: "k"},
+		{Code: []rune{'h'}[0], Text: "h"},
+		{Code: []rune{'l'}[0], Text: "l"},
+		{Code: tea.KeyDown},
+		{Code: tea.KeyUp},
+		{Code: tea.KeyLeft},
+		{Code: tea.KeyRight},
+	}
+	for _, keyMsg := range keys {
+		m = pressFlameKey(t, m, keyMsg)
+		if m.selectedIdx != 0 {
+			t.Fatalf("expected single-node selection to stay at idx 0, got %d", m.selectedIdx)
+		}
+	}
+}
+
+func pressFlameKey(t *testing.T, m Model, keyMsg tea.KeyPressMsg) Model {
+	t.Helper()
+	next, _ := m.Update(keyMsg)
+	return next.(Model)
+}
diff --git a/internal/tui/flamegraph/renderer.go b/internal/tui/flamegraph/renderer.go
index f837bfc..8b2aff0 100644
--- a/internal/tui/flamegraph/renderer.go
+++ b/internal/tui/flamegraph/renderer.go
@@ -106,7 +106,7 @@ func terminalFrameColor(name string) color.Color {
 }
 
 // RenderTerminalView renders a terminal flamegraph viewport from laid out frames.
-func RenderTerminalView(frames []tuiFrame, width, height, selectedIdx int, matchSet map[int]bool, isDark bool) string {
+func RenderTerminalView(frames []tuiFrame, width, height, selectedIdx int, subtreeSet, matchSet map[int]bool, isDark bool) string {
 	if width < minFlameWidth {
 		return common.PanelStyle.Render("Flame: terminal too narrow (need >= 60 columns)")
 	}
@@ -130,7 +130,9 @@ func RenderTerminalView(frames []tuiFrame, width, height, selectedIdx int, match
 		selectedIdx = 0
 	}
 	selected := frames[selectedIdx]
-	subtreeSet := computeSubtreeSet(frames, selectedIdx)
+	if subtreeSet == nil {
+		subtreeSet = computeSubtreeSet(frames, selectedIdx)
+	}
 
 	toolbar := fmt.Sprintf("Flame | frames:%d | rows:%d", len(frames), availableRows)
 	if truncated {
diff --git a/internal/tui/flamegraph/renderer_test.go b/internal/tui/flamegraph/renderer_test.go
index 0827a1a..f1f8c69 100644
--- a/internal/tui/flamegraph/renderer_test.go
+++ b/internal/tui/flamegraph/renderer_test.go
@@ -109,7 +109,7 @@ func TestBuildTerminalLayoutUsesPathSeparatorAndColor(t *testing.T) {
 }
 
 func TestRenderTerminalViewShowsNarrowMessage(t *testing.T) {
-	out := RenderTerminalView(nil, 50, 10, 0, nil, true)
+	out := RenderTerminalView(nil, 50, 10, 0, nil, nil, true)
 	if !strings.Contains(out, "terminal too narrow") {
 		t.Fatalf("expected narrow terminal warning, got %q", out)
 	}
@@ -125,7 +125,7 @@ func TestRenderTerminalViewIncludesToolbarAndStatus(t *testing.T) {
 	}
 	frames := BuildTerminalLayout(snapshot, 80, 6)
 
-	out := RenderTerminalView(frames, 80, 6, 1, nil, true)
+	out := RenderTerminalView(frames, 80, 6, 1, nil, nil, true)
 	if !strings.Contains(out, "Flame | frames:2") {
 		t.Fatalf("expected toolbar to include frame count, got %q", out)
 	}
@@ -161,7 +161,7 @@ func TestRenderTerminalViewShowsDeepLevelTruncationHint(t *testing.T) {
 		},
 	}
 	frames := BuildTerminalLayout(snapshot, 80, 10)
-	out := RenderTerminalView(frames, 80, 4, 0, nil, true)
+	out := RenderTerminalView(frames, 80, 4, 0, nil, nil, true)
 	if !strings.Contains(out, "showing deepest levels") {
 		t.Fatalf("expected truncation hint in toolbar, got %q", out)
 	}