Fix staging only some lines of a block of consecutive changes (#5396)

Given a block of consecutive changed lines, staging only some of those
changes didn't always work. Consider the following scenario:
```diff
@@ -1,5 +1,5 @@
 a
-b
-c
-d
+b'
+c'
+d'
 e
```
Trying to stage only the change to "d" worked already as expected, but
staging either "b" or "c" didn't; the order of the deleted and added
lines in staged changes were wrong. We don't solve the "c" case in this
PR, as it's probably uncommon enough so that it doesn't matter too much,
and it seems to be quite a bit harder than the other two (and git gui
doesn't support it either, for the record). But we fix the "b" case
here, which is probably the most common one.

Fixes #5103.

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 == DELETION && !self.opts.Reverse) || (line.Kind == ADDITION && self.opts.Reverse) {
+		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 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"),
+			)
+	},
+})

pkg/integration/tests/staging/stage_partial_block_of_changes_last_lines.go

@@ -0,0 +1,68 @@
+package staging
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var StagePartialBlockOfChangesLastLines = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Stage only the last few lines of a consecutive block of 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"),
+			).
+			NavigateToLine(Contains("-6")).
+			PressPrimaryAction().
+			SelectedLines(Contains("-7")).
+			PressPrimaryAction().
+			NavigateToLine(Contains("+6b")).
+			PressPrimaryAction().
+			SelectedLines(Contains("+7b")).
+			PressPrimaryAction()
+
+		t.Views().StagingSecondary().
+			ContainsLines(
+				Contains(" 3"),
+				Contains(" 4"),
+				Contains(" 5"),
+				Contains("-6"),
+				Contains("-7"),
+				Contains("+6b"),
+				Contains("+7b"),
+				Contains(" 8"),
+			)
+	},
+})

pkg/integration/tests/staging/stage_partial_block_of_changes_middle_lines.go

@@ -0,0 +1,74 @@
+package staging
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var StagePartialBlockOfChangesMiddleLines = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Stage only the middle lines of a consecutive block of 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"),
+			).
+			NavigateToLine(Contains("-4")).
+			PressPrimaryAction().
+			SelectedLines(Contains("-5")).
+			PressPrimaryAction().
+			NavigateToLine(Contains("+4b")).
+			PressPrimaryAction().
+			SelectedLines(Contains("+5b")).
+			PressPrimaryAction()
+
+		t.Views().StagingSecondary().
+			// This is not the desired result, ideally the added lines would come right after the
+			// deleted lines. However, this is hard to do, and it's a lot less common than staging
+			// either the first lines or last lines of a block of changes, so we live with the
+			// imperfection for now (but document it with a test here).
+			ContainsLines(
+				Contains(" 1"),
+				Contains(" 2"),
+				Contains(" 3"),
+				Contains("-4"),
+				Contains("-5"),
+				Contains(" 6"),
+				Contains(" 7"),
+				Contains("+4b"),
+				Contains("+5b"),
+				Contains(" 8"),
+			)
+	},
+})

pkg/integration/tests/test_list.go

@@ -377,8 +377,13 @@ var tests = []*components.IntegrationTest{
 	staging.DiffContextChange,
 	staging.DiscardAllChanges,
 	staging.Search,
+	staging.SelectNextLineAfterStagingInTwoHunkDiff,
+	staging.SelectNextLineAfterStagingIsolatedAddedLine,
 	staging.StageHunks,
 	staging.StageLines,
+	staging.StagePartialBlockOfChangesFirstLines,
+	staging.StagePartialBlockOfChangesLastLines,
+	staging.StagePartialBlockOfChangesMiddleLines,
 	staging.StageRanges,
 	stash.Apply,
 	stash.ApplyPatch,