jesseduffield.lazygit/pkg/gui/controllers/helpers/suggestions_helper.go

package helpers

import (
	"fmt"
	"strings"
	"sync/atomic"

	"github.com/jesseduffield/generics/set"
	"github.com/jesseduffield/lazygit/pkg/commands/git_commands"
	"github.com/jesseduffield/lazygit/pkg/commands/models"
	"github.com/jesseduffield/lazygit/pkg/gocui"
	"github.com/jesseduffield/lazygit/pkg/gui/presentation"
	"github.com/jesseduffield/lazygit/pkg/gui/types"
	"github.com/jesseduffield/lazygit/pkg/utils"
	"github.com/samber/lo"
	"golang.org/x/exp/slices"
	"gopkg.in/ozeidan/fuzzy-patricia.v3/patricia"
)

// Thinking out loud: I'm typically a staunch advocate of organising code by feature rather than type,
// because colocating code that relates to the same feature means far less effort
// to get all the context you need to work on any particular feature. But the one
// major benefit of grouping by type is that it makes it makes it less likely that
// somebody will re-implement the same logic twice, because they can quickly see
// if a certain method has been used for some use case, given that as a starting point
// they know about the type. In that vein, I'm including all our functions for
// finding suggestions in this file, so that it's easy to see if a function already
// exists for fetching a particular model.

type SuggestionsHelper struct {
	c *HelperCommon

	// filesTrie holds the repo's file paths for file-path suggestions. It's
	// rebuilt asynchronously and read from the suggestions worker goroutine, so
	// it lives here as an atomic pointer rather than in the (UI-thread-only)
	// model.
	filesTrie atomic.Pointer[patricia.Trie]
}

func NewSuggestionsHelper(
	c *HelperCommon,
) *SuggestionsHelper {
	self := &SuggestionsHelper{c: c}
	self.filesTrie.Store(patricia.NewTrie())
	return self
}

func (self *SuggestionsHelper) getRemoteNames() []string {
	return lo.Map(self.c.Model().Remotes, func(remote *models.Remote, _ int) string {
		return remote.Name
	})
}

func matchesToSuggestions(matches []string) []*types.Suggestion {
	return lo.Map(matches, func(match string, _ int) *types.Suggestion {
		return &types.Suggestion{
			Value: match,
			Label: match,
		}
	})
}

func (self *SuggestionsHelper) GetRemoteSuggestionsFunc() func(string) []*types.Suggestion {
	remoteNames := self.getRemoteNames()

	return FilterFunc(remoteNames, self.c.UserConfig().Gui.UseFuzzySearch())
}

func (self *SuggestionsHelper) getBranchNames() []string {
	return lo.Map(self.c.Model().Branches, func(branch *models.Branch, _ int) string {
		return branch.Name
	})
}

func (self *SuggestionsHelper) GetBranchNameSuggestionsFunc() func(string) []*types.Suggestion {
	branchNames := self.getBranchNames()

	return func(input string) []*types.Suggestion {
		var matchingBranchNames []string
		if input == "" {
			matchingBranchNames = branchNames
		} else {
			matchingBranchNames = utils.FilterStrings(input, branchNames, self.c.UserConfig().Gui.UseFuzzySearch())
		}

		return lo.Map(matchingBranchNames, func(branchName string, _ int) *types.Suggestion {
			return &types.Suggestion{
				Value: branchName,
				Label: presentation.GetBranchTextStyle(branchName).Sprint(branchName),
			}
		})
	}
}

// GetWorktreeBranchNameSuggestionsFunc suggests branches you can base a new
// worktree on: local branches that aren't checked out in any worktree (you can't
// make a second worktree for them), plus remote branches that don't yet have a
// local branch of the same name. Picking a remote branch creates a new local
// tracking branch, which would fail if that local branch already existed (whether
// or not it's checked out), so we leave those out and you reach the branch via its
// local entry instead.
func (self *SuggestionsHelper) GetWorktreeBranchNameSuggestionsFunc() func(string) []*types.Suggestion {
	localBranchNames := lo.FilterMap(self.c.Model().Branches, func(branch *models.Branch, _ int) (string, bool) {
		_, checkedOut := git_commands.WorktreeForBranch(branch, self.c.Model().Worktrees)
		return branch.Name, !checkedOut
	})

	existingLocalBranches := set.NewFromSlice(self.getBranchNames())
	remoteBranchNames := lo.Filter(self.getRemoteBranchNames("/"), func(remoteBranchName string, _ int) bool {
		_, branchName, _ := strings.Cut(remoteBranchName, "/")
		return !existingLocalBranches.Includes(branchName)
	})

	return FilterFunc(append(localBranchNames, remoteBranchNames...), self.c.UserConfig().Gui.UseFuzzySearch())
}

// here we asynchronously fetch the latest set of paths in the repo and store in
// self.c.Model().FilesTrie. On the main thread we'll be doing a fuzzy search via
// self.c.Model().FilesTrie. So if we've looked for a file previously, we'll start with
// the old trie and eventually it'll be swapped out for the new one.
func (self *SuggestionsHelper) GetFilePathSuggestionsFunc() func(string) []*types.Suggestion {
	_ = self.c.WithWaitingStatus(self.c.Tr.LoadingFileSuggestions, func(gocui.Task) error {
		trie := patricia.NewTrie()

		// load every file in the repo
		files, err := self.c.Git().WorkingTree.AllRepoFiles()
		if err != nil {
			return err
		}

		seen := set.New[string]()
		for _, file := range files {
			// For every file we also want to add its parent directories, but only once.
			for i := range len(file) {
				if file[i] == '/' {
					dir := file[:i]
					if !seen.Includes(dir) {
						trie.Insert(patricia.Prefix(dir), dir)
						seen.Add(dir)
					}
				}
			}

			trie.Insert(patricia.Prefix(file), file)
		}

		// cache the trie for future use
		self.filesTrie.Store(trie)
		self.c.OnUIThread(func() error {
			self.c.Contexts().Suggestions.RefreshSuggestions()
			return nil
		})

		return err
	})

	return func(input string) []*types.Suggestion {
		filesTrie := self.filesTrie.Load()
		matchingNames := []string{}
		if self.c.UserConfig().Gui.UseFuzzySearch() {
			_ = filesTrie.VisitFuzzy(patricia.Prefix(input), true, func(prefix patricia.Prefix, item patricia.Item, skipped int) error {
				matchingNames = append(matchingNames, item.(string))
				return nil
			})

			// doing another fuzzy search for good measure
			matchingNames = utils.FilterStrings(input, matchingNames, true)
		} else {
			substrings := strings.Fields(input)
			_ = filesTrie.Visit(func(prefix patricia.Prefix, item patricia.Item) error {
				for _, sub := range substrings {
					if !utils.CaseAwareContains(item.(string), sub) {
						return nil
					}
				}
				matchingNames = append(matchingNames, item.(string))
				return nil
			})
		}

		return matchesToSuggestions(matchingNames)
	}
}

func (self *SuggestionsHelper) getRemoteBranchNames(separator string) []string {
	return lo.FlatMap(self.c.Model().Remotes, func(remote *models.Remote, _ int) []string {
		return lo.Map(remote.Branches, func(branch *models.RemoteBranch, _ int) string {
			return fmt.Sprintf("%s%s%s", remote.Name, separator, branch.Name)
		})
	})
}

func (self *SuggestionsHelper) getRemoteBranchNamesForRemote(remoteName string) []string {
	remote, ok := lo.Find(self.c.Model().Remotes, func(remote *models.Remote) bool {
		return remote.Name == remoteName
	})
	if ok {
		return lo.Map(remote.Branches, func(branch *models.RemoteBranch, _ int) string {
			return branch.Name
		})
	}
	return nil
}

func (self *SuggestionsHelper) GetRemoteBranchesSuggestionsFunc(separator string) func(string) []*types.Suggestion {
	return FilterFunc(self.getRemoteBranchNames(separator), self.c.UserConfig().Gui.UseFuzzySearch())
}

func (self *SuggestionsHelper) GetRemoteBranchesForRemoteSuggestionsFunc(remoteName string) func(string) []*types.Suggestion {
	return FilterFunc(self.getRemoteBranchNamesForRemote(remoteName), self.c.UserConfig().Gui.UseFuzzySearch())
}

func (self *SuggestionsHelper) getTagNames() []string {
	return lo.Map(self.c.Model().Tags, func(tag *models.Tag, _ int) string {
		return tag.Name
	})
}

func (self *SuggestionsHelper) GetTagsSuggestionsFunc() func(string) []*types.Suggestion {
	tagNames := self.getTagNames()

	return FilterFunc(tagNames, self.c.UserConfig().Gui.UseFuzzySearch())
}

func (self *SuggestionsHelper) GetRefsSuggestionsFunc() func(string) []*types.Suggestion {
	remoteBranchNames := self.getRemoteBranchNames("/")
	localBranchNames := self.getBranchNames()
	tagNames := self.getTagNames()
	additionalRefNames := []string{"HEAD", "FETCH_HEAD", "MERGE_HEAD", "ORIG_HEAD"}

	refNames := append(append(append(remoteBranchNames, localBranchNames...), tagNames...), additionalRefNames...)

	return FilterFunc(refNames, self.c.UserConfig().Gui.UseFuzzySearch())
}

func (self *SuggestionsHelper) GetAuthorsSuggestionsFunc() func(string) []*types.Suggestion {
	authors := lo.Map(lo.Values(self.c.Model().Authors), func(author *models.Author, _ int) string {
		return author.Combined()
	})

	slices.Sort(authors)

	return FilterFunc(authors, self.c.UserConfig().Gui.UseFuzzySearch())
}

func FilterFunc(options []string, useFuzzySearch bool) func(string) []*types.Suggestion {
	return func(input string) []*types.Suggestion {
		var matches []string
		if input == "" {
			matches = options
		} else {
			matches = utils.FilterStrings(input, options, useFuzzySearch)
		}

		return matchesToSuggestions(matches)
	}
}