Fix staging only some lines of a block of consecutive changes

pkg/commands/patch/patch.go

@@ -51,6 +51,16 @@ func (self *Patch) Lines() []*PatchLine {
 	return lines
 }
 
+// Returns the old-file starting line number of the hunk containing the given
+// patch line index. Returns 0 if the line is not inside any hunk.
+func (self *Patch) HunkOldStartForLine(idx int) int {
+	hunkIdx := self.HunkContainingLine(idx)
+	if hunkIdx == -1 {
+		return 0
+	}
+	return self.hunks[hunkIdx].oldStart
+}
+
 // Returns the patch line index of the first line in the given hunk
 func (self *Patch) HunkStartIdx(hunkIndex int) int {
 	hunkIndex = lo.Clamp(hunkIndex, 0, len(self.hunks)-1)

pkg/commands/patch/patch_line.go

@@ -22,6 +22,14 @@ func (self *PatchLine) IsChange() bool {
 	return self.Kind == ADDITION || self.Kind == DELETION
 }
 
+func (self *PatchLine) IsAddition() bool {
+	return self.Kind == ADDITION
+}
+
+func (self *PatchLine) IsDeletion() bool {
+	return self.Kind == DELETION
+}
+
 // Returns the number of lines in the given slice that have one of the given kinds
 func nLinesWithKind(lines []*PatchLine, kinds []PatchLineKind) int {
 	return lo.CountBy(lines, func(line *PatchLine) bool {

pkg/commands/patch/patch_test.go

@@ -215,8 +215,8 @@ func TestTransform(t *testing.T) {
 +++ b/filename
 @@ -1,5 +1,6 @@
  apple
- orange
 +grape
+ orange
  ...
  ...
  ...
@@ -354,8 +354,8 @@ func TestTransform(t *testing.T) {
  ...
  ...
  ...
- last line
 +last line
+ last line
 \ No newline at end of file
 `,
 		},
@@ -412,8 +412,8 @@ func TestTransform(t *testing.T) {
 +++ b/filename
 @@ -1,5 +1,6 @@
  apple
- grape
 +orange
+ grape
  ...
  ...
  ...

pkg/commands/patch/transform.go

@@ -125,6 +125,22 @@ func (self *patchTransformer) transformHunk(hunk *Hunk, startOffset int, firstLi
 func (self *patchTransformer) transformHunkLines(hunk *Hunk, firstLineIdx int) []*PatchLine {
 	skippedNewlineMessageIndex := -1
 	newLines := []*PatchLine{}
+	// Unselected "old-file" lines (deletions when staging, additions when
+	// reverse-staging) are converted to context but buffered here rather than
+	// appended immediately. This ensures they end up after any selected additions
+	// in the same change block, giving the correct output ordering:
+	//   [selected deletions] [selected additions] [context from unselected deletions]
+	// Exception: if unselected new-file lines have been skipped earlier in the
+	// current change block, the selected addition comes "later" in the block. In
+	// that case the pending context (from unselected deletions before it) must be
+	// flushed first so those context lines appear before the addition in the output.
+	pendingContext := []*PatchLine{}
+	didSeeUnselectedNewFileLine := false
+
+	flushPendingContext := func() {
+		newLines = append(newLines, pendingContext...)
+		pendingContext = pendingContext[:0]
+	}
 
 	for i, line := range hunk.bodyLines {
 		lineIdx := i + firstLineIdx + 1 // plus one for header line
@@ -133,26 +149,58 @@ func (self *patchTransformer) transformHunkLines(hunk *Hunk, firstLineIdx int) [
 		}
 		isLineSelected := lo.Contains(self.opts.IncludedLineIndices, lineIdx)
 
-		if isLineSelected || (line.Kind == NEWLINE_MESSAGE && skippedNewlineMessageIndex != lineIdx) || line.Kind == CONTEXT {
+		if line.Kind == CONTEXT {
+			flushPendingContext()
+			didSeeUnselectedNewFileLine = false
+			newLines = append(newLines, line)
+			continue
+		}
+
+		if line.Kind == NEWLINE_MESSAGE {
+			if skippedNewlineMessageIndex != lineIdx {
+				flushPendingContext()
+				newLines = append(newLines, line)
+			}
+			continue
+		}
+
+		isOldFileLine := (line.Kind == DELETION && !self.opts.Reverse) || (line.Kind == ADDITION && self.opts.Reverse)
+
+		if isLineSelected {
+			// Selected "old-file" lines must flush pending context first to preserve
+			// the correct ordering of old-file lines (deletions and context) relative
+			// to each other.
+			if isOldFileLine ||
+				// Some new-file lines were skipped earlier in this change block, meaning
+				// this selected addition comes after them positionally. Flush pending
+				// context first so the unselected deletion context lines appear before
+				// this addition rather than after it.
+				didSeeUnselectedNewFileLine {
+				flushPendingContext()
+			}
 			newLines = append(newLines, line)
 			continue
 		}
 
-		if (line.Kind == DELETION && !self.opts.Reverse) || (line.Kind == ADDITION && self.opts.Reverse) {
+		if isOldFileLine {
 			content := " " + line.Content[1:]
-			newLines = append(newLines, &PatchLine{
+			pendingContext = append(pendingContext, &PatchLine{
 				Kind:    CONTEXT,
 				Content: content,
 			})
 			continue
 		}
 
+		didSeeUnselectedNewFileLine = true
+
 		if line.Kind == ADDITION {
 			// we don't want to include the 'newline at end of file' line if it involves an addition we're not including
 			skippedNewlineMessageIndex = lineIdx + 1
 		}
 	}
 
+	flushPendingContext()
+
 	return newLines
 }
 

pkg/gui/patch_exploring/state.go

@@ -89,7 +89,24 @@ func NewState(diff string, selectedLineIdx int, view *gocui.View, oldState *Stat
 		if oldState.selectMode != RANGE {
 			selectMode = oldState.selectMode
 		}
-		selectedLineIdx = viewLineIndices[patch.GetNextChangeIdx(oldState.patchLineIndices[oldState.selectedLineIdx])]
+		oldPatchLineIdx := oldState.patchLineIndices[oldState.selectedLineIdx]
+		newPatchLineIdx := patch.GetNextChangeIdx(oldPatchLineIdx)
+		// When staging an addition from a consecutive changes block, the unselected deletions get
+		// reordered to appear before the remaining additions in the new diff. This can cause the
+		// cursor to land on a deletion at the same patch line index where the staged addition used
+		// to be. In that case, skip forward past any deletions, then call GetNextChangeIdx from the
+		// first non-deletion position, which correctly lands on the next meaningful change.
+		newLines := patch.Lines()
+		if newPatchLineIdx == oldPatchLineIdx &&
+			oldState.patch.Lines()[oldPatchLineIdx].IsAddition() &&
+			newLines[newPatchLineIdx].IsDeletion() &&
+			patch.HunkOldStartForLine(newPatchLineIdx) == oldState.patch.HunkOldStartForLine(oldPatchLineIdx) {
+			for newPatchLineIdx < len(newLines) && newLines[newPatchLineIdx].IsDeletion() {
+				newPatchLineIdx++
+			}
+			newPatchLineIdx = patch.GetNextChangeIdx(newPatchLineIdx)
+		}
+		selectedLineIdx = viewLineIndices[newPatchLineIdx]
 	} else {
 		selectedLineIdx = viewLineIndices[patch.GetNextChangeIdx(0)]
 	}

pkg/integration/tests/staging/select_next_line_after_staging_in_two_hunk_diff.go

@@ -0,0 +1,61 @@
+package staging
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+// Tests that after staging individual lines from a consecutive changes block,
+// the cursor advances to the correct next change. The file has two separate
+// hunks so that we can verify the cursor crosses hunk boundaries correctly.
+var SelectNextLineAfterStagingInTwoHunkDiff = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "After staging lines from a two-hunk diff, the cursor advances correctly",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig: func(config *config.AppConfig) {
+		config.GetUserConfig().Gui.UseHunkModeInStagingView = false
+	},
+	SetupRepo: func(shell *Shell) {
+		// Use 7 context lines between the two change blocks so that git creates
+		// two separate hunks.
+		shell.CreateFileAndAdd("file1", "1\n2\na\nb\nc\nd\ne\nf\ng\n3\n4\n")
+		shell.Commit("one")
+
+		shell.UpdateFile("file1", "1b\n2b\na\nb\nc\nd\ne\nf\ng\n3b\n4b\n")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Files().
+			IsFocused().
+			Lines(
+				Contains("file1").IsSelected(),
+			).
+			PressEnter()
+
+		t.Views().Staging().
+			IsFocused().
+			ContainsLines(
+				Contains("-1"),
+				Contains("-2"),
+				Contains("+1b"),
+				Contains("+2b"),
+				Contains(" a"),
+				Contains(" b"),
+				Contains(" c"),
+				Contains("@@"),
+				Contains(" e"),
+				Contains(" f"),
+				Contains(" g"),
+				Contains("-3"),
+				Contains("-4"),
+				Contains("+3b"),
+				Contains("+4b"),
+			).
+			NavigateToLine(Contains("-2")).
+			PressPrimaryAction().
+			SelectedLine(Contains("+1b")).
+			PressPrimaryAction().
+			SelectedLine(Contains("+2b")).
+			PressPrimaryAction().
+			SelectedLine(Contains("-3"))
+	},
+})

pkg/integration/tests/staging/select_next_line_after_staging_isolated_added_line.go

@@ -0,0 +1,51 @@
+package staging
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+// Tests that after staging an isolated addition (one that is alone in its block of changes), the
+// cursor stays at the first change of the next block of changes which moves up to the same line,
+// even if that block starts with a deletion.
+var SelectNextLineAfterStagingIsolatedAddedLine = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "After staging an isolated added line, the cursor advances to the next hunk's first change",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig: func(config *config.AppConfig) {
+		config.GetUserConfig().Gui.UseHunkModeInStagingView = false
+	},
+	SetupRepo: func(shell *Shell) {
+		shell.CreateFileAndAdd("file1", "1\n2\n3\n4\n5\n6\n7\n8\n9\n")
+		shell.Commit("one")
+
+		shell.UpdateFile("file1", "1\n2\n3\nnew\n4\n5\n6\n7b\n8\n9\n")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Files().
+			IsFocused().
+			Lines(
+				Contains("file1").IsSelected(),
+			).
+			PressEnter()
+
+		t.Views().Staging().
+			IsFocused().
+			ContainsLines(
+				Contains(" 1"),
+				Contains(" 2"),
+				Contains(" 3"),
+				Contains("+new"),
+				Contains(" 4"),
+				Contains(" 5"),
+				Contains(" 6"),
+				Contains("-7"),
+				Contains("+7b"),
+				Contains(" 8"),
+				Contains(" 9"),
+			).
+			SelectedLine(Contains("+new")).
+			PressPrimaryAction().
+			SelectedLine(Contains("-7"))
+	},
+})

pkg/integration/tests/staging/stage_partial_block_of_changes_first_lines.go

@@ -0,0 +1,68 @@
+package staging
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var StagePartialBlockOfChangesFirstLines = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Stage only the first few lines of a block of consecutive changes",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig: func(config *config.AppConfig) {
+		config.GetUserConfig().Gui.UseHunkModeInStagingView = false
+	},
+	SetupRepo: func(shell *Shell) {
+		shell.CreateFileAndAdd("file1", "1\n2\n3\n4\n5\n6\n7\n8\n")
+		shell.Commit("one")
+
+		shell.UpdateFile("file1", "1\n2b\n3b\n4b\n5b\n6b\n7b\n8\n")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Files().
+			IsFocused().
+			Lines(
+				Contains("file1").IsSelected(),
+			).
+			PressEnter()
+
+		t.Views().Staging().
+			IsFocused().
+			ContainsLines(
+				Contains(" 1"),
+				Contains("-2"),
+				Contains("-3"),
+				Contains("-4"),
+				Contains("-5"),
+				Contains("-6"),
+				Contains("-7"),
+				Contains("+2b"),
+				Contains("+3b"),
+				Contains("+4b"),
+				Contains("+5b"),
+				Contains("+6b"),
+				Contains("+7b"),
+				Contains(" 8"),
+			).
+			SelectedLines(Contains("-2")).
+			PressPrimaryAction().
+			SelectedLines(Contains("-3")).
+			PressPrimaryAction().
+			NavigateToLine(Contains("+2b")).
+			PressPrimaryAction().
+			SelectedLines(Contains("+3b")).
+			PressPrimaryAction()
+
+		t.Views().StagingSecondary().
+			ContainsLines(
+				Contains(" 1"),
+				Contains("-2"),
+				Contains("-3"),
+				Contains("+2b"),
+				Contains("+3b"),
+				Contains(" 4"),
+				Contains(" 5"),
+				Contains(" 6"),
+			)
+	},
+})