Handle wrapped lines in patch explorer state

So far, lines in the view corresponded 1:1 to lines in the patch. Once we turn
on wrapping for the staging view (which we don't do yet), this is no longer
true, so we need to convert from view lines to patch lines or vice versa all
over the place.

pkg/gui/context/patch_explorer_context.go

@@ -106,13 +106,13 @@ func (self *PatchExplorerContext) FocusSelection() {
 	state := self.GetState()
 	bufferHeight := view.InnerHeight()
 	_, origin := view.Origin()
-	numLines := view.LinesHeight()
+	numLines := view.ViewLinesHeight()
 
 	newOriginY := state.CalculateOrigin(origin, bufferHeight, numLines)
 
 	view.SetOriginY(newOriginY)
 
-	startIdx, endIdx := state.SelectedRange()
+	startIdx, endIdx := state.SelectedViewRange()
 	// As far as the view is concerned, we are always selecting a range
 	view.SetRangeSelectStart(startIdx)
 	view.SetCursorY(endIdx - newOriginY)

pkg/gui/controllers/helpers/patch_building_helper.go

@@ -91,7 +91,7 @@ func (self *PatchBuildingHelper) RefreshPatchBuildingPanel(opts types.OnFocusOpt
 
 	oldState := context.GetState()
 
-	state := patch_exploring.NewState(diff, selectedLineIdx, oldState)
+	state := patch_exploring.NewState(diff, selectedLineIdx, context.GetView(), oldState)
 	context.SetState(state)
 	if state == nil {
 		self.Escape()

pkg/gui/controllers/helpers/staging_helper.go

@@ -63,11 +63,11 @@ func (self *StagingHelper) RefreshStagingPanel(focusOpts types.OnFocusOpts) {
 	secondaryContext.GetMutex().Lock()
 
 	mainContext.SetState(
-		patch_exploring.NewState(mainDiff, mainSelectedLineIdx, mainContext.GetState()),
+		patch_exploring.NewState(mainDiff, mainSelectedLineIdx, mainContext.GetView(), mainContext.GetState()),
 	)
 
 	secondaryContext.SetState(
-		patch_exploring.NewState(secondaryDiff, secondarySelectedLineIdx, secondaryContext.GetState()),
+		patch_exploring.NewState(secondaryDiff, secondarySelectedLineIdx, secondaryContext.GetView(), secondaryContext.GetState()),
 	)
 
 	mainState := mainContext.GetState()

pkg/gui/controllers/patch_building_controller.go

@@ -136,13 +136,13 @@ func (self *PatchBuildingController) toggleSelection() error {
 	if err != nil {
 		return err
 	}
-	currentLineIsStaged := lo.Contains(includedLineIndices, state.GetSelectedLineIdx())
+	currentLineIsStaged := lo.Contains(includedLineIndices, state.GetSelectedPatchLineIdx())
 	if currentLineIsStaged {
 		toggleFunc = self.c.Git().Patch.PatchBuilder.RemoveFileLineRange
 	}
 
 	// add range of lines to those set for the file
-	firstLineIdx, lastLineIdx := state.SelectedRange()
+	firstLineIdx, lastLineIdx := state.SelectedPatchRange()
 
 	if err := toggleFunc(filename, firstLineIdx, lastLineIdx); err != nil {
 		// might actually want to return an error here

pkg/gui/controllers/patch_explorer_controller.go

@@ -170,9 +170,9 @@ func (self *PatchExplorerController) GetMouseKeybindings(opts types.KeybindingsO
 }
 
 func (self *PatchExplorerController) HandlePrevLine() error {
-	before := self.context.GetState().GetSelectedLineIdx()
+	before := self.context.GetState().GetSelectedViewLineIdx()
 	self.context.GetState().CycleSelection(false)
-	after := self.context.GetState().GetSelectedLineIdx()
+	after := self.context.GetState().GetSelectedViewLineIdx()
 
 	if self.context.GetState().SelectingLine() {
 		checkScrollUp(self.context.GetViewTrait(), self.c.UserConfig(), before, after)
@@ -182,9 +182,9 @@ func (self *PatchExplorerController) HandlePrevLine() error {
 }
 
 func (self *PatchExplorerController) HandleNextLine() error {
-	before := self.context.GetState().GetSelectedLineIdx()
+	before := self.context.GetState().GetSelectedViewLineIdx()
 	self.context.GetState().CycleSelection(true)
-	after := self.context.GetState().GetSelectedLineIdx()
+	after := self.context.GetState().GetSelectedViewLineIdx()
 
 	if self.context.GetState().SelectingLine() {
 		checkScrollDown(self.context.GetViewTrait(), self.c.UserConfig(), before, after)

pkg/gui/controllers/staging_controller.go

@@ -220,7 +220,7 @@ func (self *StagingController) applySelection(reverse bool) error {
 		return nil
 	}
 
-	firstLineIdx, lastLineIdx := state.SelectedRange()
+	firstLineIdx, lastLineIdx := state.SelectedPatchRange()
 	patchToApply := patch.
 		Parse(state.GetDiff()).
 		Transform(patch.TransformOpts{
@@ -249,7 +249,7 @@ func (self *StagingController) applySelection(reverse bool) error {
 	}
 
 	if state.SelectingRange() {
-		firstLine, _ := state.SelectedRange()
+		firstLine, _ := state.SelectedViewRange()
 		state.SelectLine(firstLine)
 	}
 
@@ -290,7 +290,7 @@ func (self *StagingController) editHunk() error {
 	}
 
 	lineOffset := 3
-	lineIdxInHunk := state.GetSelectedLineIdx() - hunkStartIdx
+	lineIdxInHunk := state.GetSelectedPatchLineIdx() - hunkStartIdx
 	if err := self.c.Helpers().Files.EditFileAtLineAndWait(patchFilepath, lineIdxInHunk+lineOffset); err != nil {
 		return err
 	}

pkg/gui/patch_exploring/state.go

@@ -1,14 +1,19 @@
 package patch_exploring
 
 import (
+	"strings"
+
 	"github.com/jesseduffield/generics/set"
+	"github.com/jesseduffield/gocui"
 	"github.com/jesseduffield/lazygit/pkg/commands/patch"
+	"github.com/jesseduffield/lazygit/pkg/utils"
 )
 
 // State represents the current state of the patch explorer context i.e. when
 // you're staging a file or you're building a patch from an existing commit
 // this struct holds the info about the diff you're interacting with and what's currently selected.
 type State struct {
+	// These are in terms of view lines (wrapped), not patch lines
 	selectedLineIdx   int
 	rangeStartLineIdx int
 	// If a range is sticky, it means we expand the range when we move up or down.
@@ -17,6 +22,11 @@ type State struct {
 	diff          string
 	patch         *patch.Patch
 	selectMode    selectMode
+
+	// Array of indices of the wrapped lines indexed by a patch line index
+	viewLineIndices []int
+	// Array of indices of the original patch lines indexed by a wrapped view line index
+	patchLineIndices []int
 }
 
 // these represent what select mode we're in
@@ -28,7 +38,7 @@ const (
 	HUNK
 )
 
-func NewState(diff string, selectedLineIdx int, oldState *State) *State {
+func NewState(diff string, selectedLineIdx int, view *gocui.View, oldState *State) *State {
 	if oldState != nil && diff == oldState.diff && selectedLineIdx == -1 {
 		// if we're here then we can return the old state. If selectedLineIdx was not -1
 		// then that would mean we were trying to click and potentiall drag a range, which
@@ -42,6 +52,8 @@ func NewState(diff string, selectedLineIdx int, oldState *State) *State {
 		return nil
 	}
 
+	viewLineIndices, patchLineIndices := wrapPatchLines(diff, view)
+
 	rangeStartLineIdx := 0
 	if oldState != nil {
 		rangeStartLineIdx = oldState.rangeStartLineIdx
@@ -50,6 +62,10 @@ func NewState(diff string, selectedLineIdx int, oldState *State) *State {
 	selectMode := LINE
 	// if we have clicked from the outside to focus the main view we'll pass in a non-negative line index so that we can instantly select that line
 	if selectedLineIdx >= 0 {
+		// Clamp to the number of wrapped view lines; index might be out of
+		// bounds if a custom pager is being used which produces more lines
+		selectedLineIdx = min(selectedLineIdx, len(viewLineIndices)-1)
+
 		selectMode = RANGE
 		rangeStartLineIdx = selectedLineIdx
 	} else if oldState != nil {
@@ -57,9 +73,9 @@ func NewState(diff string, selectedLineIdx int, oldState *State) *State {
 		if oldState.selectMode == HUNK {
 			selectMode = HUNK
 		}
-		selectedLineIdx = patch.GetNextChangeIdx(oldState.selectedLineIdx)
+		selectedLineIdx = viewLineIndices[patch.GetNextChangeIdx(oldState.patchLineIndices[oldState.selectedLineIdx])]
 	} else {
-		selectedLineIdx = patch.GetNextChangeIdx(0)
+		selectedLineIdx = viewLineIndices[patch.GetNextChangeIdx(0)]
 	}
 
 	return &State{
@@ -69,10 +85,16 @@ func NewState(diff string, selectedLineIdx int, oldState *State) *State {
 		rangeStartLineIdx: rangeStartLineIdx,
 		rangeIsSticky:     false,
 		diff:              diff,
+		viewLineIndices:   viewLineIndices,
+		patchLineIndices:  patchLineIndices,
 	}
 }
 
-func (s *State) GetSelectedLineIdx() int {
+func (s *State) GetSelectedPatchLineIdx() int {
+	return s.patchLineIndices[s.selectedLineIdx]
+}
+
+func (s *State) GetSelectedViewLineIdx() int {
 	return s.selectedLineIdx
 }
 
@@ -142,8 +164,8 @@ func (s *State) SelectLine(newSelectedLineIdx int) {
 func (s *State) selectLineWithoutRangeCheck(newSelectedLineIdx int) {
 	if newSelectedLineIdx < 0 {
 		newSelectedLineIdx = 0
-	} else if newSelectedLineIdx > s.patch.LineCount()-1 {
-		newSelectedLineIdx = s.patch.LineCount() - 1
+	} else if newSelectedLineIdx > len(s.patchLineIndices)-1 {
+		newSelectedLineIdx = len(s.patchLineIndices) - 1
 	}
 
 	s.selectedLineIdx = newSelectedLineIdx
@@ -177,12 +199,12 @@ func (s *State) CycleHunk(forward bool) {
 		change = -1
 	}
 
-	hunkIdx := s.patch.HunkContainingLine(s.selectedLineIdx)
+	hunkIdx := s.patch.HunkContainingLine(s.patchLineIndices[s.selectedLineIdx])
 	if hunkIdx != -1 {
 		newHunkIdx := hunkIdx + change
 		if newHunkIdx >= 0 && newHunkIdx < s.patch.HunkCount() {
 			start := s.patch.HunkStartIdx(newHunkIdx)
-			s.selectedLineIdx = s.patch.GetNextChangeIdx(start)
+			s.selectedLineIdx = s.viewLineIndices[s.patch.GetNextChangeIdx(start)]
 		}
 	}
 }
@@ -215,16 +237,17 @@ func (s *State) CycleRange(forward bool) {
 
 // returns first and last patch line index of current hunk
 func (s *State) CurrentHunkBounds() (int, int) {
-	hunkIdx := s.patch.HunkContainingLine(s.selectedLineIdx)
+	hunkIdx := s.patch.HunkContainingLine(s.patchLineIndices[s.selectedLineIdx])
 	start := s.patch.HunkStartIdx(hunkIdx)
 	end := s.patch.HunkEndIdx(hunkIdx)
 	return start, end
 }
 
-func (s *State) SelectedRange() (int, int) {
+func (s *State) SelectedViewRange() (int, int) {
 	switch s.selectMode {
 	case HUNK:
-		return s.CurrentHunkBounds()
+		start, end := s.CurrentHunkBounds()
+		return s.viewLineIndices[start], s.viewLineIndices[end]
 	case RANGE:
 		if s.rangeStartLineIdx > s.selectedLineIdx {
 			return s.selectedLineIdx, s.rangeStartLineIdx
@@ -239,8 +262,13 @@ func (s *State) SelectedRange() (int, int) {
 	}
 }
 
+func (s *State) SelectedPatchRange() (int, int) {
+	start, end := s.SelectedViewRange()
+	return s.patchLineIndices[start], s.patchLineIndices[end]
+}
+
 func (s *State) CurrentLineNumber() int {
-	return s.patch.LineNumberOfLine(s.selectedLineIdx)
+	return s.patch.LineNumberOfLine(s.patchLineIndices[s.selectedLineIdx])
 }
 
 func (s *State) AdjustSelectedLineIdx(change int) {
@@ -256,13 +284,13 @@ func (s *State) RenderForLineIndices(includedLineIndices []int) string {
 }
 
 func (s *State) PlainRenderSelected() string {
-	firstLineIdx, lastLineIdx := s.SelectedRange()
+	firstLineIdx, lastLineIdx := s.SelectedPatchRange()
 	return s.patch.FormatRangePlain(firstLineIdx, lastLineIdx)
 }
 
 func (s *State) SelectBottom() {
 	s.DismissHunkSelectMode()
-	s.SelectLine(s.patch.LineCount() - 1)
+	s.SelectLine(len(s.patchLineIndices) - 1)
 }
 
 func (s *State) SelectTop() {
@@ -271,7 +299,13 @@ func (s *State) SelectTop() {
 }
 
 func (s *State) CalculateOrigin(currentOrigin int, bufferHeight int, numLines int) int {
-	firstLineIdx, lastLineIdx := s.SelectedRange()
+	firstLineIdx, lastLineIdx := s.SelectedViewRange()
 
-	return calculateOrigin(currentOrigin, bufferHeight, numLines, firstLineIdx, lastLineIdx, s.GetSelectedLineIdx(), s.selectMode)
+	return calculateOrigin(currentOrigin, bufferHeight, numLines, firstLineIdx, lastLineIdx, s.GetSelectedViewLineIdx(), s.selectMode)
+}
+
+func wrapPatchLines(diff string, view *gocui.View) ([]int, []int) {
+	_, viewLineIndices, patchLineIndices := utils.WrapViewLinesToWidth(
+		view.Wrap, strings.TrimSuffix(diff, "\n"), view.InnerWidth())
+	return viewLineIndices, patchLineIndices
 }