Fix staging only some lines of a block of consecutive changes

This fixes the problem; a consequence of this change is that given the following
scenario:

@@ -1,3 +1,3 @@
 1
-2
+2b
 3

staging only the line `+2b` will put it *before* the unchanged `2` line, rather
than after it as you might expect (the changed unit tests demonstrate this).
Since this should be a pretty uncommon scenario, I guess it is an ok compromise.

As you can see in the changed tests, while the behavior of what gets staged is
fixed now, it doesn't always correctly select the next line to stage. We'll
address this in the next commit.

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/integration/tests/staging/select_next_line_after_staging_in_two_hunk_diff.go

@@ -54,8 +54,15 @@ var SelectNextLineAfterStagingInTwoHunkDiff = NewIntegrationTest(NewIntegrationT
 			PressPrimaryAction().
 			SelectedLine(Contains("+1b")).
 			PressPrimaryAction().
+			/* EXPECTED:
 			SelectedLine(Contains("+2b")).
+			ACTUAL: */
+			SelectedLine(Contains("-1")).
+			NavigateToLine(Contains("+2b")).
 			PressPrimaryAction().
+			/* EXPECTED:
 			SelectedLine(Contains("-3"))
+			ACTUAL: */
+			SelectedLine(Contains("-1"))
 	},
 })

pkg/integration/tests/staging/stage_partial_block_of_changes_first_lines.go

@@ -50,12 +50,12 @@ var StagePartialBlockOfChangesFirstLines = NewIntegrationTest(NewIntegrationTest
 			PressPrimaryAction().
 			NavigateToLine(Contains("+2b")).
 			PressPrimaryAction().
-			SelectedLines(Contains("+3b")).
+			// TODO: we expect +3b to be selected, but it isn't, so navigate to it
+			NavigateToLine(Contains("+3b")).
 			PressPrimaryAction()
 
 		t.Views().StagingSecondary().
 			ContainsLines(
-				/* EXPECTED:
 				Contains(" 1"),
 				Contains("-2"),
 				Contains("-3"),
@@ -64,17 +64,6 @@ var StagePartialBlockOfChangesFirstLines = NewIntegrationTest(NewIntegrationTest
 				Contains(" 4"),
 				Contains(" 5"),
 				Contains(" 6"),
-				ACTUAL: */
-				Contains(" 1"),
-				Contains("-2"),
-				Contains("-3"),
-				Contains(" 4"),
-				Contains(" 5"),
-				Contains(" 6"),
-				Contains(" 7"),
-				Contains("+2b"),
-				Contains("+3b"),
-				Contains(" 8"),
 			)
 	},
 })