Introduce ConcurrentMap util

3 files changed, 53 insertions, 30 deletions

diff --git a/pkg/commands/git_commands/commit_loader.go b/pkg/commands/git_commands/commit_loader.go
index 14a6c5449..27f7d0c99 100644
--- a/pkg/commands/git_commands/commit_loader.go
+++ b/pkg/commands/git_commands/commit_loader.go
@@ -568,38 +568,26 @@ func (self *CommitLoader) getMergeBase(refName string) string {
 }
 
 func (self *CommitLoader) getExistingMainBranches() []string {
-	var existingBranches []string
-	var wg sync.WaitGroup
-
 	mainBranches := self.UserConfig.Git.MainBranches
-	existingBranches = make([]string, len(mainBranches))
-
-	for i, branchName := range mainBranches {
-		wg.Add(1)
-		i := i
-		branchName := branchName
-		go utils.Safe(func() {
-			defer wg.Done()
-
-			// Try to determine upstream of local main branch
-			if ref, err := self.cmd.New(
-				NewGitCmd("rev-parse").Arg("--symbolic-full-name", branchName+"@{u}").ToArgv(),
-			).DontLog().RunWithOutput(); err == nil {
-				existingBranches[i] = strings.TrimSpace(ref)
-				return
-			}
 
-			// If this failed, fallback to the local branch
-			ref := "refs/heads/" + branchName
-			if err := self.cmd.New(
-				NewGitCmd("rev-parse").Arg("--verify", "--quiet", ref).ToArgv(),
-			).DontLog().Run(); err == nil {
-				existingBranches[i] = ref
-			}
-		})
-	}
+	existingBranches := utils.ConcurrentMap(mainBranches, func(mainBranch string) string {
+		// Try to determine upstream of local main branch
+		if ref, err := self.cmd.New(
+			NewGitCmd("rev-parse").Arg("--symbolic-full-name", mainBranch+"@{u}").ToArgv(),
+		).DontLog().RunWithOutput(); err == nil {
+			return strings.TrimSpace(ref)
+		}
 
-	wg.Wait()
+		// If this failed, fallback to the local branch
+		ref := "refs/heads/" + mainBranch
+		if err := self.cmd.New(
+			NewGitCmd("rev-parse").Arg("--verify", "--quiet", ref).ToArgv(),
+		).DontLog().Run(); err == nil {
+			return ref
+		}
+
+		return ""
+	})
 
 	existingBranches = lo.Filter(existingBranches, func(branch string, _ int) bool {
 		return branch != ""
diff --git a/pkg/utils/slice.go b/pkg/utils/slice.go
index 0dc9e64e4..b2f629a0b 100644
--- a/pkg/utils/slice.go
+++ b/pkg/utils/slice.go
@@ -1,6 +1,10 @@
 package utils
 
-import "golang.org/x/exp/slices"
+import (
+	"sync"
+
+	"golang.org/x/exp/slices"
+)
 
 // NextIndex returns the index of the element that comes after the given number
 func NextIndex(numbers []int, currentNumber int) int {
@@ -179,3 +183,24 @@ func Shift[T any](slice []T) (T, []T) {
 	slice = slice[1:]
 	return value, slice
 }
+
+// Map function which handles each element in its own goroutine
+func ConcurrentMap[T any, V any](items []T, fn func(T) V) []V {
+	results := make([]V, len(items))
+
+	var wg sync.WaitGroup
+	for i, item := range items {
+		i := i
+		item := item
+
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			results[i] = fn(item)
+		}()
+	}
+
+	wg.Wait()
+
+	return results
+}
diff --git a/pkg/utils/slice_test.go b/pkg/utils/slice_test.go
index b3cad8885..6c77b4047 100644
--- a/pkg/utils/slice_test.go
+++ b/pkg/utils/slice_test.go
@@ -315,3 +315,13 @@ func TestMoveElement(t *testing.T) {
 		})
 	})
 }
+
+func TestConcurrentMap(t *testing.T) {
+	in := []int{1, 2, 3}
+
+	out := ConcurrentMap(in, func(i int) int {
+		return i * 2
+	})
+
+	assert.EqualValues(t, []int{2, 4, 6}, out)
+}