refactor

3 files changed, 136 insertions, 52 deletions

diff --git a/pkg/gui/mergeconflicts/find_conflicts.go b/pkg/gui/mergeconflicts/find_conflicts.go
new file mode 100644
index 000000000..661cc5cc0
--- /dev/null
+++ b/pkg/gui/mergeconflicts/find_conflicts.go
@@ -0,0 +1,63 @@
+package mergeconflicts
+
+import (
+	"regexp"
+	"strings"
+
+	"github.com/jesseduffield/lazygit/pkg/utils"
+)
+
+// LineType tells us whether a given line is a start/middle/end marker of a conflict,
+// or if it's not a marker at all
+type LineType int
+
+const (
+	START LineType = iota
+	MIDDLE
+	END
+	NOT_A_MARKER
+)
+
+func findConflicts(content string) []*mergeConflict {
+	conflicts := make([]*mergeConflict, 0)
+
+	if content == "" {
+		return conflicts
+	}
+
+	var newConflict *mergeConflict
+	for i, line := range utils.SplitLines(content) {
+		switch determineLineType(line) {
+		case START:
+			newConflict = &mergeConflict{start: i}
+		case MIDDLE:
+			newConflict.middle = i
+		case END:
+			newConflict.end = i
+			conflicts = append(conflicts, newConflict)
+		default:
+			// line isn't a merge conflict marker so we just continue
+		}
+	}
+
+	return conflicts
+}
+
+func determineLineType(line string) LineType {
+	trimmedLine := strings.TrimPrefix(line, "++")
+
+	mapping := map[string]LineType{
+		"^<<<<<<< (HEAD|MERGE_HEAD|Updated upstream|ours)(:.+)?$": START,
+		"^=======$":    MIDDLE,
+		"^>>>>>>> .*$": END,
+	}
+
+	for regexp_str, lineType := range mapping {
+		match, _ := regexp.MatchString(regexp_str, trimmedLine)
+		if match {
+			return lineType
+		}
+	}
+
+	return NOT_A_MARKER
+}
diff --git a/pkg/gui/mergeconflicts/find_conflicts_test.go b/pkg/gui/mergeconflicts/find_conflicts_test.go
new file mode 100644
index 000000000..c0a8e3e5b
--- /dev/null
+++ b/pkg/gui/mergeconflicts/find_conflicts_test.go
@@ -0,0 +1,57 @@
+package mergeconflicts
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDetermineLineType(t *testing.T) {
+	type scenario struct {
+		line     string
+		expected LineType
+	}
+
+	scenarios := []scenario{
+		{
+			line:     "",
+			expected: NOT_A_MARKER,
+		},
+		{
+			line:     "blah",
+			expected: NOT_A_MARKER,
+		},
+		{
+			line:     "<<<<<<< HEAD",
+			expected: START,
+		},
+		{
+			line:     "<<<<<<< HEAD:my_branch",
+			expected: START,
+		},
+		{
+			line:     "<<<<<<< MERGE_HEAD:my_branch",
+			expected: START,
+		},
+		{
+			line:     "<<<<<<< Updated upstream:my_branch",
+			expected: START,
+		},
+		{
+			line:     "<<<<<<< ours:my_branch",
+			expected: START,
+		},
+		{
+			line:     "=======",
+			expected: MIDDLE,
+		},
+		{
+			line:     ">>>>>>> blah",
+			expected: END,
+		},
+	}
+
+	for _, s := range scenarios {
+		assert.EqualValues(t, s.expected, determineLineType(s.line))
+	}
+}
diff --git a/pkg/gui/mergeconflicts/state.go b/pkg/gui/mergeconflicts/state.go
index 383d71a36..9d8320ffb 100644
--- a/pkg/gui/mergeconflicts/state.go
+++ b/pkg/gui/mergeconflicts/state.go
@@ -1,9 +1,6 @@
 package mergeconflicts
 
 import (
-	"bufio"
-	"os"
-	"strings"
 	"sync"
 
 	"github.com/golang-collections/collections/stack"
@@ -19,7 +16,7 @@ const (
 )
 
 // mergeConflict : A git conflict with a start middle and end corresponding to line
-// numbers in the file where the conflict bars appear
+// numbers in the file where the conflict markers appear
 type mergeConflict struct {
 	start  int
 	middle int
@@ -88,36 +85,6 @@ func (s *State) SetConflictsFromCat(cat string) {
 	s.setConflicts(findConflicts(cat))
 }
 
-func findConflicts(content string) []*mergeConflict {
-	conflicts := make([]*mergeConflict, 0)
-
-	if content == "" {
-		return conflicts
-	}
-
-	var newConflict *mergeConflict
-	for i, line := range utils.SplitLines(content) {
-		trimmedLine := strings.TrimPrefix(line, "++")
-		switch trimmedLine {
-		case "<<<<<<< HEAD", "<<<<<<< MERGE_HEAD", "<<<<<<< Updated upstream", "<<<<<<< ours":
-			newConflict = &mergeConflict{start: i}
-		case "=======":
-			newConflict.middle = i
-		default:
-			// Sometimes these lines look like "<<<<<<< HEAD:foo/bar/baz.go" so handle that case as well.
-			if strings.HasPrefix(trimmedLine, "<<<<<<< HEAD:") {
-				newConflict = &mergeConflict{start: i}
-			}
-			if strings.HasPrefix(trimmedLine, ">>>>>>> ") {
-				newConflict.end = i
-				conflicts = append(conflicts, newConflict)
-			}
-		}
-
-	}
-	return conflicts
-}
-
 func (s *State) setConflicts(conflicts []*mergeConflict) {
 	s.conflicts = conflicts
 
@@ -158,32 +125,29 @@ func (s *State) ContentAfterConflictResolve(path string, selection Selection) (b
 		return false, "", nil
 	}
 
-	file, err := os.Open(path)
-	if err != nil {
-		return false, "", err
-	}
-	defer file.Close()
-
-	reader := bufio.NewReader(file)
 	content := ""
-	for i := 0; true; i++ {
-		line, err := reader.ReadString('\n')
-		if err != nil {
-			break
-		}
+	err := utils.ForEachLineInFile(path, func(line string, i int) {
 		if !isIndexToDelete(i, conflict, selection) {
 			content += line
 		}
+	})
+
+	if err != nil {
+		return false, "", err
 	}
 
 	return true, content, nil
 }
 
 func isIndexToDelete(i int, conflict *mergeConflict, selection Selection) bool {
-	return i == conflict.middle ||
-		i == conflict.start ||
-		i == conflict.end ||
-		selection != BOTH &&
-			(selection == BOTTOM && i > conflict.start && i < conflict.middle) ||
-		(selection == TOP && i > conflict.middle && i < conflict.end)
+	isMarkerLine :=
+		i == conflict.middle ||
+			i == conflict.start ||
+			i == conflict.end
+
+	isUnwantedContent :=
+		(selection == BOTTOM && conflict.start < i && i < conflict.middle) ||
+			(selection == TOP && conflict.middle < i && i < conflict.end)
+
+	return isMarkerLine || isUnwantedContent
 }