refactor patch code

4 files changed, 60 insertions, 35 deletions

diff --git a/pkg/gui/commits_panel.go b/pkg/gui/commits_panel.go
index 2b297c90d..b56b3af2b 100644
--- a/pkg/gui/commits_panel.go
+++ b/pkg/gui/commits_panel.go
@@ -52,7 +52,7 @@ func (gui *Gui) branchCommitsRenderToMain() error {
 
 func (gui *Gui) secondaryPatchPanelUpdateOpts() *types.ViewUpdateOpts {
 	if gui.git.Patch.PatchManager.Active() {
-		patch := gui.git.Patch.PatchManager.RenderAggregatedPatchColored(false)
+		patch := gui.git.Patch.PatchManager.RenderAggregatedPatch(false)
 
 		return &types.ViewUpdateOpts{
 			Task:  types.NewRenderStringWithoutScrollTask(patch),
diff --git a/pkg/gui/controllers/staging_controller.go b/pkg/gui/controllers/staging_controller.go
index 78c271640..0738ae1dc 100644
--- a/pkg/gui/controllers/staging_controller.go
+++ b/pkg/gui/controllers/staging_controller.go
@@ -181,10 +181,16 @@ func (self *StagingController) applySelection(reverse bool) error {
 	}
 
 	firstLineIdx, lastLineIdx := state.SelectedRange()
-	patch := patch.ModifiedPatchForRange(self.c.Log, path, state.GetDiff(), firstLineIdx, lastLineIdx,
-		patch.PatchOptions{Reverse: reverse, KeepOriginalHeader: false})
-
-	if patch == "" {
+	patchToApply := patch.
+		Parse(state.GetDiff()).
+		Transform(patch.TransformOpts{
+			Reverse:             reverse,
+			IncludedLineIndices: patch.ExpandRange(firstLineIdx, lastLineIdx),
+			FileNameOverride:    path,
+		}).
+		FormatPlain()
+
+	if patchToApply == "" {
 		return nil
 	}
 
@@ -198,7 +204,7 @@ func (self *StagingController) applySelection(reverse bool) error {
 		applyFlags = append(applyFlags, "cached")
 	}
 	self.c.LogAction(self.c.Tr.Actions.ApplyPatch)
-	err := self.git.WorkingTree.ApplyPatch(patch, applyFlags...)
+	err := self.git.WorkingTree.ApplyPatch(patchToApply, applyFlags...)
 	if err != nil {
 		return self.c.Error(err)
 	}
@@ -229,18 +235,23 @@ func (self *StagingController) editHunk() error {
 		return nil
 	}
 
-	hunk := state.CurrentHunk()
-	patchText := patch.ModifiedPatchForRange(
-		self.c.Log, path, state.GetDiff(), hunk.FirstLineIdx, hunk.LastLineIdx(),
-		patch.PatchOptions{Reverse: self.staged, KeepOriginalHeader: false},
-	)
+	hunkStartIdx, hunkEndIdx := state.CurrentHunkBounds()
+	patchText := patch.
+		Parse(state.GetDiff()).
+		Transform(patch.TransformOpts{
+			Reverse:             self.staged,
+			IncludedLineIndices: patch.ExpandRange(hunkStartIdx, hunkEndIdx),
+			FileNameOverride:    path,
+		}).
+		FormatPlain()
+
 	patchFilepath, err := self.git.WorkingTree.SaveTemporaryPatch(patchText)
 	if err != nil {
 		return err
 	}
 
 	lineOffset := 3
-	lineIdxInHunk := state.GetSelectedLineIdx() - hunk.FirstLineIdx
+	lineIdxInHunk := state.GetSelectedLineIdx() - hunkStartIdx
 	if err := self.helpers.Files.EditFileAtLine(patchFilepath, lineIdxInHunk+lineOffset); err != nil {
 		return err
 	}
@@ -253,10 +264,13 @@ func (self *StagingController) editHunk() error {
 	self.c.LogAction(self.c.Tr.Actions.ApplyPatch)
 
 	lineCount := strings.Count(editedPatchText, "\n") + 1
-	newPatchText := patch.ModifiedPatchForRange(
-		self.c.Log, path, editedPatchText, 0, lineCount,
-		patch.PatchOptions{KeepOriginalHeader: false},
-	)
+	newPatchText := patch.
+		Parse(editedPatchText).
+		Transform(patch.TransformOpts{
+			IncludedLineIndices: patch.ExpandRange(0, lineCount),
+			FileNameOverride:    path,
+		}).
+		FormatPlain()
 
 	applyFlags := []string{"cached"}
 	if self.staged {
diff --git a/pkg/gui/custom_patch_options_panel.go b/pkg/gui/custom_patch_options_panel.go
index 6d640157b..2dfcf4e1e 100644
--- a/pkg/gui/custom_patch_options_panel.go
+++ b/pkg/gui/custom_patch_options_panel.go
@@ -202,7 +202,7 @@ func (gui *Gui) handleApplyPatch(reverse bool) error {
 }
 
 func (gui *Gui) copyPatchToClipboard() error {
-	patch := gui.git.Patch.PatchManager.RenderAggregatedPatchColored(true)
+	patch := gui.git.Patch.PatchManager.RenderAggregatedPatch(true)
 
 	gui.c.LogAction(gui.c.Tr.Actions.CopyPatchToClipboard)
 	if err := gui.os.CopyToClipboard(patch); err != nil {
diff --git a/pkg/gui/patch_exploring/state.go b/pkg/gui/patch_exploring/state.go
index 008338326..ad5426805 100644
--- a/pkg/gui/patch_exploring/state.go
+++ b/pkg/gui/patch_exploring/state.go
@@ -1,6 +1,7 @@
 package patch_exploring
 
 import (
+	"github.com/jesseduffield/generics/set"
 	"github.com/jesseduffield/lazygit/pkg/commands/patch"
 	"github.com/sirupsen/logrus"
 )
@@ -12,7 +13,7 @@ type State struct {
 	selectedLineIdx   int
 	rangeStartLineIdx int
 	diff              string
-	patchParser       *patch.PatchParser
+	patch             *patch.Patch
 	selectMode        selectMode
 }
 
@@ -33,9 +34,9 @@ func NewState(diff string, selectedLineIdx int, oldState *State, log *logrus.Ent
 		return oldState
 	}
 
-	patchParser := patch.NewPatchParser(log, diff)
+	patch := patch.Parse(diff)
 
-	if len(patchParser.StageableLines) == 0 {
+	if !patch.ContainsChanges() {
 		return nil
 	}
 
@@ -54,13 +55,13 @@ func NewState(diff string, selectedLineIdx int, oldState *State, log *logrus.Ent
 		if oldState.selectMode == HUNK {
 			selectMode = HUNK
 		}
-		selectedLineIdx = patchParser.GetNextStageableLineIndex(oldState.selectedLineIdx)
+		selectedLineIdx = patch.GetNextChangeIdx(oldState.selectedLineIdx)
 	} else {
-		selectedLineIdx = patchParser.StageableLines[0]
+		selectedLineIdx = patch.GetNextChangeIdx(0)
 	}
 
 	return &State{
-		patchParser:       patchParser,
+		patch:             patch,
 		selectedLineIdx:   selectedLineIdx,
 		selectMode:        selectMode,
 		rangeStartLineIdx: rangeStartLineIdx,
@@ -112,8 +113,8 @@ func (s *State) SetLineSelectMode() {
 func (s *State) SelectLine(newSelectedLineIdx int) {
 	if newSelectedLineIdx < 0 {
 		newSelectedLineIdx = 0
-	} else if newSelectedLineIdx > len(s.patchParser.PatchLines)-1 {
-		newSelectedLineIdx = len(s.patchParser.PatchLines) - 1
+	} else if newSelectedLineIdx > s.patch.LineCount()-1 {
+		newSelectedLineIdx = s.patch.LineCount() - 1
 	}
 
 	s.selectedLineIdx = newSelectedLineIdx
@@ -141,8 +142,9 @@ func (s *State) CycleHunk(forward bool) {
 		change = -1
 	}
 
-	newHunk := s.patchParser.GetHunkContainingLine(s.selectedLineIdx, change)
-	s.selectedLineIdx = s.patchParser.GetNextStageableLineIndex(newHunk.FirstLineIdx)
+	hunkIdx := s.patch.HunkContainingLine(s.selectedLineIdx)
+	start := s.patch.HunkStartIdx(hunkIdx + change)
+	s.selectedLineIdx = s.patch.GetNextChangeIdx(start)
 }
 
 func (s *State) CycleLine(forward bool) {
@@ -154,15 +156,18 @@ func (s *State) CycleLine(forward bool) {
 	s.SelectLine(s.selectedLineIdx + change)
 }
 
-func (s *State) CurrentHunk() *patch.PatchHunk {
-	return s.patchParser.GetHunkContainingLine(s.selectedLineIdx, 0)
+// returns first and last patch line index of current hunk
+func (s *State) CurrentHunkBounds() (int, int) {
+	hunkIdx := s.patch.HunkContainingLine(s.selectedLineIdx)
+	start := s.patch.HunkStartIdx(hunkIdx)
+	end := s.patch.HunkEndIdx(hunkIdx)
+	return start, end
 }
 
 func (s *State) SelectedRange() (int, int) {
 	switch s.selectMode {
 	case HUNK:
-		hunk := s.CurrentHunk()
-		return hunk.FirstLineIdx, hunk.LastLineIdx()
+		return s.CurrentHunkBounds()
 	case RANGE:
 		if s.rangeStartLineIdx > s.selectedLineIdx {
 			return s.selectedLineIdx, s.rangeStartLineIdx
@@ -178,7 +183,7 @@ func (s *State) SelectedRange() (int, int) {
 }
 
 func (s *State) CurrentLineNumber() int {
-	return s.CurrentHunk().LineNumberOfLine(s.selectedLineIdx)
+	return s.patch.LineNumberOfLine(s.selectedLineIdx)
 }
 
 func (s *State) AdjustSelectedLineIdx(change int) {
@@ -187,17 +192,23 @@ func (s *State) AdjustSelectedLineIdx(change int) {
 
 func (s *State) RenderForLineIndices(isFocused bool, includedLineIndices []int) string {
 	firstLineIdx, lastLineIdx := s.SelectedRange()
-	return s.patchParser.Render(isFocused, firstLineIdx, lastLineIdx, includedLineIndices)
+	includedLineIndicesSet := set.NewFromSlice(includedLineIndices)
+	return s.patch.FormatView(patch.FormatViewOpts{
+		IsFocused:      isFocused,
+		FirstLineIndex: firstLineIdx,
+		LastLineIndex:  lastLineIdx,
+		IncLineIndices: includedLineIndicesSet,
+	})
 }
 
 func (s *State) PlainRenderSelected() string {
 	firstLineIdx, lastLineIdx := s.SelectedRange()
-	return s.patchParser.RenderLinesPlain(firstLineIdx, lastLineIdx)
+	return s.patch.FormatRangePlain(firstLineIdx, lastLineIdx)
 }
 
 func (s *State) SelectBottom() {
 	s.SetLineSelectMode()
-	s.SelectLine(len(s.patchParser.PatchLines) - 1)
+	s.SelectLine(s.patch.LineCount() - 1)
 }
 
 func (s *State) SelectTop() {