match beginning and end of line correctly in ReplaceRegex

internal/buffer/loc.go

@@ -47,6 +47,16 @@ func (l Loc) LessEqual(b Loc) bool {
 	return l == b
 }
 
+// Clamp clamps a loc between start and end
+func (l Loc) Clamp(start, end Loc) Loc {
+	if l.GreaterEqual(end) {
+		return end
+	} else if l.LessThan(start) {
+		return start
+	}
+	return l
+}
+
 // The following functions require a buffer to know where newlines are
 
 // Diff returns the distance between two locations
@@ -139,10 +149,5 @@ func ByteOffset(pos Loc, buf *Buffer) int {
 
 // clamps a loc within a buffer
 func clamp(pos Loc, la *LineArray) Loc {
-	if pos.GreaterEqual(la.End()) {
-		return la.End()
-	} else if pos.LessThan(la.Start()) {
-		return la.Start()
-	}
-	return pos
+	return pos.Clamp(la.Start(), la.End())
 }

internal/buffer/search.go

@@ -127,6 +127,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
@@ -170,53 +190,58 @@ 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
-			if captureGroups {
-				for _, submatches := range search.FindAllSubmatchIndex(in, -1) {
-					result = search.Expand(result, replace, in, submatches)
+		l := b.LineBytes(i)
+		charCount := util.CharacterCount(l)
+		if (i == start.Y && start.X > 0) || (i == end.Y && end.X < charCount) {
+			// This replacement code 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
+
+			from := Loc{0, i}.Clamp(start, end)
+			to := Loc{charCount, i}.Clamp(start, end)
+			matches := b.findAll(search, from, to)
+			found += len(matches)
+
+			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 {
+					newText = search.ReplaceAll(b.Substr(match[0], match[1]), replace)
+				} else {
+					newText = replace
 				}
-			} else {
-				result = replace
-			}
-			found++
-			if i == end.Y {
-				netrunes += util.CharacterCount(result) - util.CharacterCount(in)
+				deltas = append(deltas, Delta{newText, match[0], match[1]})
 			}
-			return result
-		})
-
-		from := Loc{charpos, i}
-		to := Loc{charpos + util.CharacterCount(l), i}
-
-		deltas = append(deltas, Delta{newText, from, to})
+		} else {
+			newLine := search.ReplaceAllFunc(l, 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
+			})
+			deltas = append(deltas, Delta{newLine, Loc{0, i}, Loc{charCount, i}})
+		}
 	}
+
 	b.MultipleReplace(deltas)
 
-	return found, netrunes
+	return found, util.CharacterCount(b.LineBytes(end.Y)) - charsEnd
 }