refactor patch code

pkg/commands/patch/transform.go

@@ -0,0 +1,156 @@
+package patch
+
+import "github.com/samber/lo"
+
+type patchTransformer struct {
+	patch *Patch
+	opts  TransformOpts
+}
+
+type TransformOpts struct {
+	// Create a patch that will applied in reverse with `git apply --reverse`.
+	// This affects how unselected lines are treated when only parts of a hunk
+	// are selected: usually, for unselected lines we change '-' lines to
+	// context lines and remove '+' lines, but when Reverse is true we need to
+	// turn '+' lines into context lines and remove '-' lines.
+	Reverse bool
+
+	// If set, we will replace the original header with one referring to this file name.
+	// For staging/unstaging lines we don't want the original header because
+	// it makes git confused e.g. when dealing with deleted/added files
+	// but with building and applying patches the original header gives git
+	// information it needs to cleanly apply patches
+	FileNameOverride string
+
+	// The indices of lines that should be included in the patch.
+	IncludedLineIndices []int
+}
+
+func transform(patch *Patch, opts TransformOpts) *Patch {
+	transformer := &patchTransformer{
+		patch: patch,
+		opts:  opts,
+	}
+
+	return transformer.transform()
+}
+
+// helper function that takes a start and end index and returns a slice of all
+// indexes inbetween (inclusive)
+func ExpandRange(start int, end int) []int {
+	expanded := []int{}
+	for i := start; i <= end; i++ {
+		expanded = append(expanded, i)
+	}
+	return expanded
+}
+
+func (self *patchTransformer) transform() *Patch {
+	header := self.transformHeader()
+	hunks := self.transformHunks()
+
+	return &Patch{
+		header: header,
+		hunks:  hunks,
+	}
+}
+
+func (self *patchTransformer) transformHeader() []string {
+	if self.opts.FileNameOverride != "" {
+		return []string{
+			"--- a/" + self.opts.FileNameOverride,
+			"+++ b/" + self.opts.FileNameOverride,
+		}
+	} else {
+		return self.patch.header
+	}
+}
+
+func (self *patchTransformer) transformHunks() []*Hunk {
+	newHunks := make([]*Hunk, 0, len(self.patch.hunks))
+
+	startOffset := 0
+	var formattedHunk *Hunk
+	for i, hunk := range self.patch.hunks {
+		startOffset, formattedHunk = self.transformHunk(
+			hunk,
+			startOffset,
+			self.patch.HunkStartIdx(i),
+		)
+		if formattedHunk.containsChanges() {
+			newHunks = append(newHunks, formattedHunk)
+		}
+	}
+
+	return newHunks
+}
+
+func (self *patchTransformer) transformHunk(hunk *Hunk, startOffset int, firstLineIdx int) (int, *Hunk) {
+	newLines := self.transformHunkLines(hunk, firstLineIdx)
+	newNewStart, newStartOffset := self.transformHunkHeader(newLines, hunk.oldStart, startOffset)
+
+	newHunk := &Hunk{
+		bodyLines:     newLines,
+		oldStart:      hunk.oldStart,
+		newStart:      newNewStart,
+		headerContext: hunk.headerContext,
+	}
+
+	return newStartOffset, newHunk
+}
+
+func (self *patchTransformer) transformHunkLines(hunk *Hunk, firstLineIdx int) []*PatchLine {
+	skippedNewlineMessageIndex := -1
+	newLines := []*PatchLine{}
+
+	for i, line := range hunk.bodyLines {
+		lineIdx := i + firstLineIdx + 1 // plus one for header line
+		if line.Content == "" {
+			break
+		}
+		isLineSelected := lo.Contains(self.opts.IncludedLineIndices, lineIdx)
+
+		if isLineSelected || (line.Kind == NEWLINE_MESSAGE && skippedNewlineMessageIndex != lineIdx) || line.Kind == CONTEXT {
+			newLines = append(newLines, line)
+			continue
+		}
+
+		if (line.Kind == DELETION && !self.opts.Reverse) || (line.Kind == ADDITION && self.opts.Reverse) {
+			content := " " + line.Content[1:]
+			newLines = append(newLines, &PatchLine{
+				Kind:    CONTEXT,
+				Content: content,
+			})
+			continue
+		}
+
+		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
+		}
+	}
+
+	return newLines
+}
+
+func (self *patchTransformer) transformHunkHeader(newBodyLines []*PatchLine, oldStart int, startOffset int) (int, int) {
+	oldLength := nLinesWithKind(newBodyLines, []PatchLineKind{CONTEXT, DELETION})
+	newLength := nLinesWithKind(newBodyLines, []PatchLineKind{CONTEXT, ADDITION})
+
+	var newStartOffset int
+	// if the hunk went from zero to positive length, we need to increment the starting point by one
+	// if the hunk went from positive to zero length, we need to decrement the starting point by one
+	if oldLength == 0 {
+		newStartOffset = 1
+	} else if newLength == 0 {
+		newStartOffset = -1
+	} else {
+		newStartOffset = 0
+	}
+
+	newStart := oldStart + startOffset + newStartOffset
+
+	newStartOffset = startOffset + newLength - oldLength
+
+	return newStart, newStartOffset
+}