refactor cheatsheet generator

2 files changed, 351 insertions, 123 deletions

diff --git a/pkg/cheatsheet/generate.go b/pkg/cheatsheet/generate.go
index 04d8d3fd5..d20a0c71a 100644
--- a/pkg/cheatsheet/generate.go
+++ b/pkg/cheatsheet/generate.go
@@ -21,6 +21,8 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
 	"github.com/jesseduffield/lazygit/pkg/i18n"
 	"github.com/jesseduffield/lazygit/pkg/integration"
+	"github.com/jesseduffield/lazygit/pkg/utils"
+	"github.com/samber/lo"
 )
 
 type bindingSection struct {
@@ -28,6 +30,17 @@ type bindingSection struct {
 	bindings []*types.Binding
 }
 
+type header struct {
+	// priority decides the order of the headers in the cheatsheet (lower means higher)
+	priority int
+	title    string
+}
+
+type headerWithBindings struct {
+	header   header
+	bindings []*types.Binding
+}
+
 func CommandToRun() string {
 	return "go run scripts/cheatsheet/main.go generate"
 }
@@ -49,7 +62,8 @@ func generateAtDir(cheatsheetDir string) {
 			panic(err)
 		}
 
-		bindingSections := getBindingSections(mApp)
+		bindings := mApp.Gui.GetCheatsheetKeybindings()
+		bindingSections := getBindingSections(bindings, mApp.Tr)
 		content := formatSections(mApp.Tr, bindingSections)
 		content = fmt.Sprintf("_This file is auto-generated. To update, make the changes in the "+
 			"pkg/i18n directory and then run `%s` from the project root._\n\n%s", CommandToRun(), content)
@@ -68,9 +82,7 @@ func writeString(file *os.File, str string) {
 	}
 }
 
-func localisedTitle(mApp *app.App, str string) string {
-	tr := mApp.Tr
-
+func localisedTitle(tr *i18n.TranslationSet, str string) string {
 	contextTitleMap := map[string]string{
 		"global":         tr.GlobalTitle,
 		"navigation":     tr.NavigationTitle,
@@ -110,142 +122,66 @@ func localisedTitle(mApp *app.App, str string) string {
 	return title
 }
 
-func formatTitle(title string) string {
-	return fmt.Sprintf("\n## %s\n\n", title)
-}
-
-func formatBinding(binding *types.Binding) string {
-	if binding.Alternative != "" {
-		return fmt.Sprintf("  <kbd>%s</kbd>: %s (%s)\n", gui.GetKeyDisplay(binding.Key), binding.Description, binding.Alternative)
-	}
-	return fmt.Sprintf("  <kbd>%s</kbd>: %s\n", gui.GetKeyDisplay(binding.Key), binding.Description)
-}
-
-func getBindingSections(mApp *app.App) []*bindingSection {
-	bindingSections := []*bindingSection{}
-
-	bindings := mApp.Gui.GetCheatsheetKeybindings()
-
-	type contextAndViewType struct {
-		subtitle string
-		title    string
-	}
-
-	contextAndViewBindingMap := map[contextAndViewType][]*types.Binding{}
-
-outer:
-	for _, binding := range bindings {
-		if binding.Tag == "navigation" {
-			key := contextAndViewType{subtitle: "", title: "navigation"}
-			existing := contextAndViewBindingMap[key]
-			if existing == nil {
-				contextAndViewBindingMap[key] = []*types.Binding{binding}
-			} else {
-				if !slices.Some(contextAndViewBindingMap[key], func(navBinding *types.Binding) bool {
-					return navBinding.Description == binding.Description
-				}) {
-					contextAndViewBindingMap[key] = append(contextAndViewBindingMap[key], binding)
-				}
-			}
-
-			continue outer
-		}
-
-		contexts := []string{}
-		if len(binding.Contexts) == 0 {
-			contexts = append(contexts, "")
-		} else {
-			contexts = append(contexts, binding.Contexts...)
-		}
+func getBindingSections(bindings []*types.Binding, tr *i18n.TranslationSet) []*bindingSection {
+	bindingsToDisplay := slices.Filter(bindings, func(binding *types.Binding) bool {
+		return binding.Description != "" || binding.Alternative != ""
+	})
 
-		for _, context := range contexts {
-			key := contextAndViewType{subtitle: context, title: binding.ViewName}
-			existing := contextAndViewBindingMap[key]
-			if existing == nil {
-				contextAndViewBindingMap[key] = []*types.Binding{binding}
-			} else {
-				contextAndViewBindingMap[key] = append(contextAndViewBindingMap[key], binding)
-			}
-		}
-	}
+	bindingsByHeader := utils.MuiltiGroupBy(bindingsToDisplay, func(binding *types.Binding) []header {
+		return getHeaders(binding, tr)
+	})
 
-	type groupedBindingsType struct {
-		contextAndView contextAndViewType
-		bindings       []*types.Binding
-	}
+	bindingGroups := maps.MapToSlice(bindingsByHeader, func(header header, hBindings []*types.Binding) headerWithBindings {
+		uniqBindings := lo.UniqBy(hBindings, func(binding *types.Binding) string {
+			return binding.Description + gui.GetKeyDisplay(binding.Key)
+		})
 
-	groupedBindings := maps.MapToSlice(
-		contextAndViewBindingMap,
-		func(contextAndView contextAndViewType, contextBindings []*types.Binding) groupedBindingsType {
-			return groupedBindingsType{contextAndView: contextAndView, bindings: contextBindings}
-		},
-	)
-
-	slices.SortFunc(groupedBindings, func(a, b groupedBindingsType) bool {
-		first := a.contextAndView
-		second := b.contextAndView
-		if first.title == "" {
-			return true
-		}
-		if second.title == "" {
-			return false
+		return headerWithBindings{
+			header:   header,
+			bindings: uniqBindings,
 		}
-		if first.title == "navigation" {
-			return true
-		}
-		if second.title == "navigation" {
-			return false
-		}
-		return first.title < second.title || (first.title == second.title && first.subtitle < second.subtitle)
 	})
 
-	for _, group := range groupedBindings {
-		contextAndView := group.contextAndView
-		contextBindings := group.bindings
-		mApp.Log.Info("viewname: " + contextAndView.title + ", context: " + contextAndView.subtitle)
-		viewName := contextAndView.title
-		if viewName == "" {
-			viewName = "global"
-		}
-		translatedView := localisedTitle(mApp, viewName)
-		var title string
-		if contextAndView.subtitle == "" {
-			addendum := " " + mApp.Tr.Panel
-			if viewName == "global" || viewName == "navigation" {
-				addendum = ""
-			}
-			title = fmt.Sprintf("%s%s", translatedView, addendum)
-		} else {
-			translatedContextName := localisedTitle(mApp, contextAndView.subtitle)
-			title = fmt.Sprintf("%s %s (%s)", translatedView, mApp.Tr.Panel, translatedContextName)
+	slices.SortFunc(bindingGroups, func(a, b headerWithBindings) bool {
+		if a.header.priority != b.header.priority {
+			return a.header.priority > b.header.priority
 		}
+		return a.header.title < b.header.title
+	})
 
-		for _, binding := range contextBindings {
-			bindingSections = addBinding(title, bindingSections, binding)
+	return slices.Map(bindingGroups, func(hb headerWithBindings) *bindingSection {
+		return &bindingSection{
+			title:    hb.header.title,
+			bindings: hb.bindings,
 		}
-	}
-
-	return bindingSections
+	})
 }
 
-func addBinding(title string, bindingSections []*bindingSection, binding *types.Binding) []*bindingSection {
-	if binding.Description == "" && binding.Alternative == "" {
-		return bindingSections
+// a binding may belong to multiple headers if it is applicable to multiple contexts,
+// for example the copy-to-clipboard binding.
+func getHeaders(binding *types.Binding, tr *i18n.TranslationSet) []header {
+	if binding.Tag == "navigation" {
+		return []header{{priority: 2, title: localisedTitle(tr, "navigation")}}
 	}
 
-	for _, section := range bindingSections {
-		if title == section.title {
-			section.bindings = append(section.bindings, binding)
-			return bindingSections
-		}
+	if binding.ViewName == "" {
+		return []header{{priority: 3, title: localisedTitle(tr, "global")}}
 	}
 
-	section := &bindingSection{
-		title:    title,
-		bindings: []*types.Binding{binding},
+	if len(binding.Contexts) == 0 {
+		translatedView := localisedTitle(tr, binding.ViewName)
+		title := fmt.Sprintf("%s %s", translatedView, tr.Panel)
+
+		return []header{{priority: 1, title: title}}
 	}
 
-	return append(bindingSections, section)
+	return slices.Map(binding.Contexts, func(context string) header {
+		translatedView := localisedTitle(tr, binding.ViewName)
+		translatedContextName := localisedTitle(tr, context)
+		title := fmt.Sprintf("%s %s (%s)", translatedView, tr.Panel, translatedContextName)
+
+		return header{priority: 1, title: title}
+	})
 }
 
 func formatSections(tr *i18n.TranslationSet, bindingSections []*bindingSection) string {
@@ -262,3 +198,14 @@ func formatSections(tr *i18n.TranslationSet, bindingSections []*bindingSection)
 
 	return content
 }
+
+func formatTitle(title string) string {
+	return fmt.Sprintf("\n## %s\n\n", title)
+}
+
+func formatBinding(binding *types.Binding) string {
+	if binding.Alternative != "" {
+		return fmt.Sprintf("  <kbd>%s</kbd>: %s (%s)\n", gui.GetKeyDisplay(binding.Key), binding.Description, binding.Alternative)
+	}
+	return fmt.Sprintf("  <kbd>%s</kbd>: %s\n", gui.GetKeyDisplay(binding.Key), binding.Description)
+}
diff --git a/pkg/cheatsheet/generate_test.go b/pkg/cheatsheet/generate_test.go
new file mode 100644
index 000000000..94b571454
--- /dev/null
+++ b/pkg/cheatsheet/generate_test.go
@@ -0,0 +1,281 @@
+package cheatsheet
+
+import (
+	"testing"
+
+	"github.com/jesseduffield/lazygit/pkg/gui/types"
+	"github.com/jesseduffield/lazygit/pkg/i18n"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetBindingSections(t *testing.T) {
+	tr := i18n.EnglishTranslationSet()
+
+	tests := []struct {
+		testName string
+		bindings []*types.Binding
+		expected []*bindingSection
+	}{
+		{
+			testName: "no bindings",
+			bindings: []*types.Binding{},
+			expected: []*bindingSection{},
+		},
+		{
+			testName: "one binding",
+			bindings: []*types.Binding{
+				{
+					ViewName:    "files",
+					Description: "stage file",
+				},
+			},
+			expected: []*bindingSection{
+				{
+					title: "Files Panel",
+					bindings: []*types.Binding{
+						{
+							ViewName:    "files",
+							Description: "stage file",
+						},
+					},
+				},
+			},
+		},
+		{
+			testName: "one binding with context",
+			bindings: []*types.Binding{
+				{
+					ViewName:    "files",
+					Description: "stage file",
+					Contexts:    []string{"submodules"},
+				},
+			},
+			expected: []*bindingSection{
+				{
+					title: "Files Panel (Submodules)",
+					bindings: []*types.Binding{
+						{
+							ViewName:    "files",
+							Description: "stage file",
+							Contexts:    []string{"submodules"},
+						},
+					},
+				},
+			},
+		},
+		{
+			testName: "global binding",
+			bindings: []*types.Binding{
+				{
+					ViewName:    "",
+					Description: "quit",
+				},
+			},
+			expected: []*bindingSection{
+				{
+					title: "Global Keybindings",
+					bindings: []*types.Binding{
+						{
+							ViewName:    "",
+							Description: "quit",
+						},
+					},
+				},
+			},
+		},
+		{
+			testName: "grouped bindings",
+			bindings: []*types.Binding{
+				{
+					ViewName:    "files",
+					Description: "stage file",
+					Contexts:    []string{"files"},
+				},
+				{
+					ViewName:    "files",
+					Description: "unstage file",
+					Contexts:    []string{"files"},
+				},
+				{
+					ViewName:    "files",
+					Description: "drop submodule",
+					Contexts:    []string{"submodules"},
+				},
+				{
+					ViewName:    "commits",
+					Description: "revert commit",
+				},
+			},
+			expected: []*bindingSection{
+				{
+					title: "Commits Panel",
+					bindings: []*types.Binding{
+						{
+							ViewName:    "commits",
+							Description: "revert commit",
+						},
+					},
+				},
+				{
+					title: "Files Panel (Files)",
+					bindings: []*types.Binding{
+						{
+							ViewName:    "files",
+							Description: "stage file",
+							Contexts:    []string{"files"},
+						},
+						{
+							ViewName:    "files",
+							Description: "unstage file",
+							Contexts:    []string{"files"},
+						},
+					},
+				},
+				{
+					title: "Files Panel (Submodules)",
+					bindings: []*types.Binding{
+						{
+							ViewName:    "files",
+							Description: "drop submodule",
+							Contexts:    []string{"submodules"},
+						},
+					},
+				},
+			},
+		},
+		{
+			testName: "with navigation bindings",
+			bindings: []*types.Binding{
+				{
+					ViewName:    "files",
+					Description: "stage file",
+				},
+				{
+					ViewName:    "files",
+					Description: "unstage file",
+				},
+				{
+					ViewName:    "files",
+					Description: "scroll",
+					Tag:         "navigation",
+				},
+				{
+					ViewName:    "commits",
+					Description: "revert commit",
+				},
+			},
+			expected: []*bindingSection{
+				{
+					title: "List Panel Navigation",
+					bindings: []*types.Binding{
+						{
+							ViewName:    "files",
+							Description: "scroll",
+							Tag:         "navigation",
+						},
+					},
+				},
+				{
+					title: "Commits Panel",
+					bindings: []*types.Binding{
+						{
+							ViewName:    "commits",
+							Description: "revert commit",
+						},
+					},
+				},
+				{
+					title: "Files Panel",
+					bindings: []*types.Binding{
+						{
+							ViewName:    "files",
+							Description: "stage file",
+						},
+						{
+							ViewName:    "files",
+							Description: "unstage file",
+						},
+					},
+				},
+			},
+		},
+		{
+			testName: "with duplicate navigation bindings",
+			bindings: []*types.Binding{
+				{
+					ViewName:    "files",
+					Description: "stage file",
+				},
+				{
+					ViewName:    "files",
+					Description: "unstage file",
+				},
+				{
+					ViewName:    "files",
+					Description: "scroll",
+					Tag:         "navigation",
+				},
+				{
+					ViewName:    "commits",
+					Description: "revert commit",
+				},
+				{
+					ViewName:    "commits",
+					Description: "scroll",
+					Tag:         "navigation",
+				},
+				{
+					ViewName:    "commits",
+					Description: "page up",
+					Tag:         "navigation",
+				},
+			},
+			expected: []*bindingSection{
+				{
+					title: "List Panel Navigation",
+					bindings: []*types.Binding{
+						{
+							ViewName:    "files",
+							Description: "scroll",
+							Tag:         "navigation",
+						},
+						{
+							ViewName:    "commits",
+							Description: "page up",
+							Tag:         "navigation",
+						},
+					},
+				},
+				{
+					title: "Commits Panel",
+					bindings: []*types.Binding{
+						{
+							ViewName:    "commits",
+							Description: "revert commit",
+						},
+					},
+				},
+				{
+					title: "Files Panel",
+					bindings: []*types.Binding{
+						{
+							ViewName:    "files",
+							Description: "stage file",
+						},
+						{
+							ViewName:    "files",
+							Description: "unstage file",
+						},
+					},
+				},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.testName, func(t *testing.T) {
+			actual := getBindingSections(test.bindings, &tr)
+			assert.EqualValues(t, test.expected, actual)
+		})
+	}
+}