lazygit/pkg/gui/controllers/helpers/fixup_helper.go

package helpers

import (
	"errors"
	"fmt"
	"regexp"
	"strings"

	"github.com/jesseduffield/generics/set"
	"github.com/jesseduffield/lazygit/pkg/commands/models"
	"github.com/jesseduffield/lazygit/pkg/gui/types"
	"github.com/jesseduffield/lazygit/pkg/utils"
	"github.com/samber/lo"
	"golang.org/x/sync/errgroup"
)

type FixupHelper struct {
	c *HelperCommon
}

func NewFixupHelper(
	c *HelperCommon,
) *FixupHelper {
	return &FixupHelper{
		c: c,
	}
}

// hunk describes the lines in a diff hunk. Used for two distinct cases:
//
// - when the hunk contains some deleted lines. Because we're diffing with a
// context of 0, all deleted lines always come first, and then the added lines
// (if any). In this case, numLines is only the number of deleted lines, we
// ignore whether there are also some added lines in the hunk, as this is not
// relevant for our algorithm.
//
// - when the hunk contains only added lines, in which case (obviously) numLines
// is the number of added lines.
type hunk struct {
	filename     string
	startLineIdx int
	numLines     int
}

func (self *FixupHelper) HandleFindBaseCommitForFixupPress() error {
	diff, hasStagedChanges, err := self.getDiff()
	if err != nil {
		return err
	}

	deletedLineHunks, addedLineHunks := parseDiff(diff)

	commits := self.c.Model().Commits

	var hashes []string
	warnAboutAddedLines := false

	if len(deletedLineHunks) > 0 {
		hashes, err = self.blameDeletedLines(deletedLineHunks)
		warnAboutAddedLines = len(addedLineHunks) > 0
	} else if len(addedLineHunks) > 0 {
		hashes, err = self.blameAddedLines(commits, addedLineHunks)
	} else {
		return errors.New(self.c.Tr.NoChangedFiles)
	}

	if err != nil {
		return err
	}

	if len(hashes) == 0 {
		// This should never happen
		return errors.New(self.c.Tr.NoBaseCommitsFound)
	}

	// If a commit can't be found, and the last known commit is already merged,
	// we know that the commit we're looking for is also merged. Otherwise we
	// can't tell.
	notFoundMeansMerged := len(commits) > 0 && commits[len(commits)-1].Status == models.StatusMerged

	const (
		MERGED int = iota
		NOT_MERGED
		CANNOT_TELL
	)

	// Group the hashes into buckets by merged status
	hashGroups := lo.GroupBy(hashes, func(hash string) int {
		commit, _, ok := self.findCommit(commits, hash)
		if ok {
			return lo.Ternary(commit.Status == models.StatusMerged, MERGED, NOT_MERGED)
		}
		return lo.Ternary(notFoundMeansMerged, MERGED, CANNOT_TELL)
	})

	if len(hashGroups[CANNOT_TELL]) > 0 {
		// If we have any commits that we can't tell if they're merged, just
		// show the generic "not in current view" error. This can only happen if
		// a feature branch has more than 300 commits, or there is no main
		// branch. Both are so unlikely that we don't bother returning a more
		// detailed error message (e.g. we could say something about the commits
		// that *are* in the current branch, but it's not worth it).
		return errors.New(self.c.Tr.BaseCommitIsNotInCurrentView)
	}

	if len(hashGroups[NOT_MERGED]) == 0 {
		// If all the commits are merged, show the "already on main branch"
		// error. It isn't worth doing a detailed report of which commits we
		// found.
		return errors.New(self.c.Tr.BaseCommitIsAlreadyOnMainBranch)
	}

	foundCommits := getCommitsForHashes(commits, hashGroups[NOT_MERGED])
	// If there are multiple commits that could be the base commit, remove all
	// those that are fixups for the last one.
	foundCommits = removeFixupCommits(foundCommits)

	if len(foundCommits) > 1 {
		// If there are still multiple commits that could be the base commit, list
		// them in the error message. But only the candidates from the current
		// branch, not including any that are already merged.
		subjects := getHashesAndSubjects(foundCommits)
		message := lo.Ternary(hasStagedChanges,
			self.c.Tr.MultipleBaseCommitsFoundStaged,
			self.c.Tr.MultipleBaseCommitsFoundUnstaged)
		return fmt.Errorf("%s\n\n%s", message, subjects)
	}

	// Now we know that foundCommits has exactly one commit, so find its index
	_, index, _ := self.findCommit(commits, foundCommits[0].Hash())

	return self.c.ConfirmIf(warnAboutAddedLines, types.ConfirmOpts{
		Title:  self.c.Tr.FindBaseCommitForFixup,
		Prompt: self.c.Tr.HunksWithOnlyAddedLinesWarning,
		HandleConfirm: func() error {
			if !hasStagedChanges {
				if err := self.c.Git().WorkingTree.StageAll(true); err != nil {
					return err
				}
				self.c.Refresh(types.RefreshOptions{Scope: []types.RefreshableView{types.FILES}})
			}

			self.c.Contexts().LocalCommits.SetSelection(index)
			self.c.Context().Push(self.c.Contexts().LocalCommits, types.OnFocusOpts{})
			return nil
		},
	})
}

func getCommitsForHashes(commits []*models.Commit, hashes []string) []*models.Commit {
	// This is called only for the NOT_MERGED commits, and we know that all of them are contained in
	// the commits slice.
	commitsSet := set.NewFromSlice(hashes)
	result := make([]*models.Commit, 0, len(hashes))
	for _, c := range commits {
		if commitsSet.Includes(c.Hash()) {
			result = append(result, c)
			commitsSet.Remove(c.Hash())
			if commitsSet.Len() == 0 {
				break
			}
		}
	}
	return result
}

func getHashesAndSubjects(commits []*models.Commit) string {
	subjects := lo.Map(commits, func(c *models.Commit, _ int) string {
		return fmt.Sprintf("%s %s", c.ShortRefName(), c.Name)
	})
	return strings.Join(subjects, "\n")
}

func removeFixupCommits(commits []*models.Commit) []*models.Commit {
	if len(commits) <= 1 {
		return commits
	}

	// If the last found commit is itself a fixup, don't eliminate anything.
	baseSubject, lastIsFixup := IsFixupCommit(commits[len(commits)-1].Name)
	if lastIsFixup {
		return commits
	}

	// need to go backwards because we are mutating the slice as we go
	for i := len(commits) - 2; i >= 0; i-- {
		subject, isFixup := IsFixupCommit(commits[i].Name)
		if isFixup && subject == baseSubject {
			commits = utils.Remove(commits, i)
		}
	}

	return commits
}

func (self *FixupHelper) getDiff() (string, bool, error) {
	args := []string{"-U0", "--ignore-submodules=all", "HEAD", "--"}

	// Try staged changes first
	hasStagedChanges := true
	diff, err := self.c.Git().Diff.DiffIndexCmdObj(append([]string{"--cached"}, args...)...).DontLog().RunWithOutput()

	if err == nil && diff == "" {
		hasStagedChanges = false
		// If there are no staged changes, try unstaged changes
		diff, err = self.c.Git().Diff.DiffIndexCmdObj(args...).DontLog().RunWithOutput()
	}

	return diff, hasStagedChanges, err
}

// Parse the diff output into hunks, and return two lists of hunks: the first
// are ones that contain deleted lines, the second are ones that contain only
// added lines.
func parseDiff(diff string) ([]*hunk, []*hunk) {
	lines := strings.Split(strings.TrimSuffix(diff, "\n"), "\n")

	deletedLineHunks := []*hunk{}
	addedLineHunks := []*hunk{}

	hunkHeaderRegexp := regexp.MustCompile(`@@ -(\d+)(?:,\d+)? \+\d+(?:,\d+)? @@`)

	var filename string
	var currentHunk *hunk
	numDeletedLines := 0
	numAddedLines := 0
	finishHunk := func() {
		if currentHunk != nil {
			if numDeletedLines > 0 {
				currentHunk.numLines = numDeletedLines
				deletedLineHunks = append(deletedLineHunks, currentHunk)
			} else if numAddedLines > 0 {
				currentHunk.numLines = numAddedLines
				addedLineHunks = append(addedLineHunks, currentHunk)
			}
		}
		numDeletedLines = 0
		numAddedLines = 0
	}
	for _, line := range lines {
		if strings.HasPrefix(line, "diff --git") {
			finishHunk()
			currentHunk = nil
		} else if currentHunk == nil && strings.HasPrefix(line, "--- ") {
			// For some reason, the line ends with a tab character if the file
			// name contains spaces
			filename = strings.TrimRight(line[6:], "\t")
		} else if strings.HasPrefix(line, "@@ ") {
			finishHunk()
			match := hunkHeaderRegexp.FindStringSubmatch(line)
			startIdx := utils.MustConvertToInt(match[1])
			currentHunk = &hunk{filename, startIdx, 0}
		} else if currentHunk != nil && line[0] == '-' {
			numDeletedLines++
		} else if currentHunk != nil && line[0] == '+' {
			numAddedLines++
		}
	}
	finishHunk()

	return deletedLineHunks, addedLineHunks
}

// returns the list of commit hashes that introduced the lines which have now been deleted
func (self *FixupHelper) blameDeletedLines(deletedLineHunks []*hunk) ([]string, error) {
	errg := errgroup.Group{}
	hashChan := make(chan string)

	for _, h := range deletedLineHunks {
		errg.Go(func() error {
			blameOutput, err := self.c.Git().Blame.BlameLineRange(h.filename, "HEAD", h.startLineIdx, h.numLines)
			if err != nil {
				return err
			}
			blameLines := strings.SplitSeq(strings.TrimSuffix(blameOutput, "\n"), "\n")
			for line := range blameLines {
				hashChan <- strings.Split(line, " ")[0]
			}
			return nil
		})
	}

	go func() {
		// We don't care about the error here, we'll check it later (in the
		// return statement below). Here we only wait for all the goroutines to
		// finish so that we can close the channel.
		_ = errg.Wait()
		close(hashChan)
	}()

	result := set.New[string]()
	for hash := range hashChan {
		result.Add(hash)
	}

	return result.ToSlice(), errg.Wait()
}

func (self *FixupHelper) blameAddedLines(commits []*models.Commit, addedLineHunks []*hunk) ([]string, error) {
	errg := errgroup.Group{}
	hashesChan := make(chan []string)

	for _, h := range addedLineHunks {
		errg.Go(func() error {
			result := make([]string, 0, 2)

			appendBlamedLine := func(blameOutput string) {
				blameLines := strings.Split(strings.TrimSuffix(blameOutput, "\n"), "\n")
				if len(blameLines) == 1 {
					result = append(result, strings.Split(blameLines[0], " ")[0])
				}
			}

			// Blame the line before this hunk, if there is one
			if h.startLineIdx > 0 {
				blameOutput, err := self.c.Git().Blame.BlameLineRange(h.filename, "HEAD", h.startLineIdx, 1)
				if err != nil {
					return err
				}
				appendBlamedLine(blameOutput)
			}

			// Blame the line after this hunk. We don't know how many lines the
			// file has, so we can't check if there is a line after the hunk;
			// let the error tell us.
			blameOutput, err := self.c.Git().Blame.BlameLineRange(h.filename, "HEAD", h.startLineIdx+1, 1)
			if err != nil {
				// If this fails, we're probably at the end of the file (we
				// could have checked this beforehand, but it's expensive). If
				// there was a line before this hunk, this is fine, we'll just
				// return that one; if not, the hunk encompasses the entire
				// file, and we can't blame the lines before and after the hunk.
				// This is an error.
				if h.startLineIdx == 0 {
					return errors.New("Entire file") // TODO i18n
				}
			} else {
				appendBlamedLine(blameOutput)
			}

			hashesChan <- result
			return nil
		})
	}

	go func() {
		// We don't care about the error here, we'll check it later (in the
		// return statement below). Here we only wait for all the goroutines to
		// finish so that we can close the channel.
		_ = errg.Wait()
		close(hashesChan)
	}()

	result := set.New[string]()
	for hashes := range hashesChan {
		if len(hashes) == 1 {
			result.Add(hashes[0])
		} else if len(hashes) > 1 {
			if hashes[0] == hashes[1] {
				result.Add(hashes[0])
			} else {
				_, index1, ok1 := self.findCommit(commits, hashes[0])
				_, index2, ok2 := self.findCommit(commits, hashes[1])
				if ok1 && ok2 {
					result.Add(lo.Ternary(index1 < index2, hashes[0], hashes[1]))
				} else if ok1 {
					result.Add(hashes[0])
				} else if ok2 {
					result.Add(hashes[1])
				} else {
					return nil, errors.New(self.c.Tr.NoBaseCommitsFound)
				}
			}
		}
	}

	return result.ToSlice(), errg.Wait()
}

func (self *FixupHelper) findCommit(commits []*models.Commit, hash string) (*models.Commit, int, bool) {
	return lo.FindIndexOf(commits, func(commit *models.Commit) bool {
		return commit.Hash() == hash
	})
}

// Check whether the given subject line is the subject of a fixup commit, and
// returns (trimmedSubject, true) if so (where trimmedSubject is the subject
// with all fixup prefixes removed), or (subject, false) if not.
func IsFixupCommit(subject string) (string, bool) {
	prefixes := []string{"fixup! ", "squash! ", "amend! "}
	trimPrefix := func(s string) (string, bool) {
		for _, prefix := range prefixes {
			if trimmedSubject, ok := strings.CutPrefix(s, prefix); ok {
				return trimmedSubject, true
			}
		}
		return s, false
	}

	if subject, wasTrimmed := trimPrefix(subject); wasTrimmed {
		for {
			// handle repeated prefixes like "fixup! amend! fixup! Subject"
			if subject, wasTrimmed = trimPrefix(subject); !wasTrimmed {
				break
			}
		}
		return subject, true
	}

	return subject, false
}