match beginning and end of line correctly in ReplaceRegex

Author: matthias314

internal/buffer/search.go

@@ -122,6 +122,26 @@ func (b *Buffer) findUp(r *regexp.Regexp, start, end Loc) ([2]Loc, bool) {
 	return [2]Loc{}, false
 }
 
+func (b *Buffer) findAll(r *regexp.Regexp, start, end Loc) [][2]Loc {
+	matches := [][2]Loc{}
+	loc := start
+	for {
+		match, found := b.findDown(r, loc, end)
+		if !found {
+			break
+		}
+		matches = append(matches, match)
+		if match[0] != match[1] {
+			loc = match[1]
+		} else if match[1] != end {
+			loc = match[1].Move(1, b)
+		} else {
+			break
+		}
+	}
+	return matches
+}
+
 // FindNext finds the next occurrence of a given string in the buffer
 // It returns the start and end location of the match (if found) and
 // a boolean indicating if it was found
@@ -165,53 +185,73 @@ func (b *Buffer) FindNext(s string, start, end, from Loc, down bool, useRegex bo
 }
 
 // ReplaceRegex replaces all occurrences of 'search' with 'replace' in the given area
-// and returns the number of replacements made and the number of runes
+// and returns the number of replacements made and the number of characters
 // added or removed on the last line of the range
 func (b *Buffer) ReplaceRegex(start, end Loc, search *regexp.Regexp, replace []byte, captureGroups bool) (int, int) {
 	if start.GreaterThan(end) {
 		start, end = end, start
 	}
 
-	netrunes := 0
-
+	charsEnd := util.CharacterCount(b.LineBytes(end.Y))
 	found := 0
 	var deltas []Delta
-	for i := start.Y; i <= end.Y; i++ {
-		l := b.lines[i].data
-		charpos := 0
-
-		if start.Y == end.Y && i == start.Y {
-			l = util.SliceStart(l, end.X)
-			l = util.SliceEnd(l, start.X)
-			charpos = start.X
-		} else if i == start.Y {
-			l = util.SliceEnd(l, start.X)
-			charpos = start.X
-		} else if i == end.Y {
-			l = util.SliceStart(l, end.X)
-		}
-		newText := search.ReplaceAllFunc(l, func(in []byte) []byte {
-			var result []byte
+
+	// This replacement function works in general, but it creates a separate
+	// modification for each match. We only use it for the first and last lines,
+	// which may use padded regexps
+	replaceFirstLast := func(start, end Loc) []Delta {
+		matches := b.findAll(search, start, end)
+		for j := len(matches) - 1; j >= 0; j-- {
+			// if we counted upwards, the different deltas would interfere
+			match := matches[j]
+			var newText []byte
 			if captureGroups {
-				for _, submatches := range search.FindAllSubmatchIndex(in, -1) {
-					result = search.Expand(result, replace, in, submatches)
-				}
+				newText = search.ReplaceAll(b.Substr(match[0], match[1]), replace)
 			} else {
-				result = replace
+				newText = replace
 			}
-			found++
-			if i == end.Y {
-				netrunes += util.CharacterCount(result) - util.CharacterCount(in)
-			}
-			return result
-		})
+			deltas = append(deltas, Delta{newText, match[0], match[1]})
+		}
+		found += len(matches)
+		return deltas
+	}
+
+	replaceMiddle := func(in []byte) []byte {
+		found++
+		var result []byte
+		if captureGroups {
+			match := search.FindSubmatchIndex(in)
+			result = search.Expand(result, replace, in, match)
+		} else {
+			result = replace
+		}
+		return result
+	}
 
-		from := Loc{charpos, i}
-		to := Loc{charpos + util.CharacterCount(l), i}
+	// first line (if different from last line)
+	if start.Y < end.Y {
+		n := util.CharacterCount(b.LineBytes(start.Y))
+		startEnd := Loc{n, start.Y}
+		deltas = replaceFirstLast(start, startEnd)
+	}
 
-		deltas = append(deltas, Delta{newText, from, to})
+	// middle lines
+	for i := start.Y + 1; i < end.Y; i++ {
+		l := b.LineBytes(i)
+		n := util.CharacterCount(l)
+		newLine := search.ReplaceAllFunc(l, replaceMiddle)
+		deltas = append(deltas, Delta{newLine, Loc{0, i}, Loc{n, i}})
 	}
+
+	// last line
+	endStart := Loc{0, end.Y}
+	if start.Y == end.Y {
+		endStart = start
+	}
+	deltas = replaceFirstLast(endStart, end)
+
 	b.MultipleReplace(deltas)
 
-	return found, netrunes
+	deltaEndX := util.CharacterCount(b.LineBytes(end.Y)) - charsEnd
+	return found, deltaEndX
 }