Preserve regex replace_all anchor semantics

Author: mgomes

vibes/builtins.go

@@ -11,7 +11,6 @@ import (
 	"strconv"
 	"strings"
 	"time"
-	"unicode/utf8"
 )
 
 const (
@@ -496,64 +495,83 @@ func builtinRegexReplaceInternal(args []Value, kwargs map[string]Value, block Va
 }
 
 func regexReplaceAllWithLimit(re *regexp.Regexp, text string, replacement string, method string) (string, error) {
-	out := make([]byte, 0, len(text))
-	lastAppended := 0
-	searchStart := 0
-	lastMatchEnd := -1
-	for searchStart <= len(text) {
-		loc := re.FindStringSubmatchIndex(text[searchStart:])
-		if loc == nil {
-			break
+	literalBytesPerMatch, refCountPerMatch := regexReplacementTemplateStats(replacement)
+	matchCount := int64(0)
+	totalMatchBytes := int64(0)
+	referenceMatchBytes := int64(0)
+
+	// Use regexp's native replacement scanner so anchoring/boundary semantics
+	// match the final replace_all behavior while estimating worst-case output.
+	re.ReplaceAllStringFunc(text, func(match string) string {
+		matchCount++
+		matchBytes := int64(len(match))
+		totalMatchBytes += matchBytes
+		referenceMatchBytes += matchBytes * int64(refCountPerMatch)
+		return match
+	})
+
+	upperBound := int64(len(text)) - totalMatchBytes + matchCount*int64(literalBytesPerMatch) + referenceMatchBytes
+	if upperBound > int64(maxRegexInputBytes) {
+		return "", fmt.Errorf("%s output exceeds limit %d bytes", method, maxRegexInputBytes)
+	}
+
+	out := re.ReplaceAllString(text, replacement)
+	if len(out) > maxRegexInputBytes {
+		return "", fmt.Errorf("%s output exceeds limit %d bytes", method, maxRegexInputBytes)
+	}
+	return out, nil
+}
+
+func regexReplacementTemplateStats(replacement string) (literalBytes int, referenceCount int) {
+	for i := 0; i < len(replacement); {
+		if replacement[i] != '$' {
+			literalBytes++
+			i++
+			continue
 		}
-		for i := range loc {
-			if loc[i] >= 0 {
-				loc[i] += searchStart
-			}
+		if i+1 >= len(replacement) {
+			literalBytes++
+			i++
+			continue
 		}
-		start := loc[0]
-		end := loc[1]
-		if start == end && start == lastMatchEnd {
-			if searchStart >= len(text) {
-				break
+		if replacement[i+1] == '$' {
+			literalBytes++
+			i += 2
+			continue
+		}
+		if replacement[i+1] == '{' {
+			j := i + 2
+			for j < len(replacement) && replacement[j] != '}' {
+				j++
 			}
-			_, size := utf8.DecodeRuneInString(text[searchStart:])
-			if size == 0 {
-				size = 1
+			if j < len(replacement) {
+				referenceCount++
+				i = j + 1
+				continue
 			}
-			searchStart += size
+			literalBytes++
+			i++
 			continue
 		}
 
-		segmentLen := start - lastAppended
-		if len(out) > maxRegexInputBytes-segmentLen {
-			return "", fmt.Errorf("%s output exceeds limit %d bytes", method, maxRegexInputBytes)
+		j := i + 1
+		for j < len(replacement) && isRegexReplacementNameByte(replacement[j]) {
+			j++
 		}
-		out = append(out, text[lastAppended:start]...)
-		out = re.ExpandString(out, replacement, text, loc)
-		if len(out) > maxRegexInputBytes {
-			return "", fmt.Errorf("%s output exceeds limit %d bytes", method, maxRegexInputBytes)
-		}
-		lastAppended = end
-		lastMatchEnd = end
-
-		if end > searchStart {
-			searchStart = end
+		if j == i+1 {
+			literalBytes++
+			i++
 			continue
 		}
-		if searchStart >= len(text) {
-			break
-		}
-		_, size := utf8.DecodeRuneInString(text[searchStart:])
-		if size == 0 {
-			size = 1
-		}
-		searchStart += size
+		referenceCount++
+		i = j
 	}
+	return literalBytes, referenceCount
+}
 
-	tailLen := len(text) - lastAppended
-	if len(out) > maxRegexInputBytes-tailLen {
-		return "", fmt.Errorf("%s output exceeds limit %d bytes", method, maxRegexInputBytes)
-	}
-	out = append(out, text[lastAppended:]...)
-	return string(out), nil
+func isRegexReplacementNameByte(char byte) bool {
+	return (char >= '0' && char <= '9') ||
+		(char >= 'a' && char <= 'z') ||
+		(char >= 'A' && char <= 'Z') ||
+		char == '_'
 }

vibes/runtime_test.go

@@ -1954,6 +1954,7 @@ func TestRegexBuiltins(t *testing.T) {
 	        match_empty: Regex.match("^", "ID-12"),
 	        replace_one: Regex.replace("ID-12 ID-34", "ID-[0-9]+", "X"),
 	        replace_all: Regex.replace_all("ID-12 ID-34", "ID-[0-9]+", "X"),
+	        replace_all_anchor: Regex.replace_all("abc", "^", "X"),
 	        replace_capture: Regex.replace("ID-12 ID-34", "ID-([0-9]+)", "X-$1"),
 	        replace_boundary: Regex.replace("ab", "\\Bb", "X")
 	      }
@@ -1984,6 +1985,9 @@ func TestRegexBuiltins(t *testing.T) {
 	if !out["replace_all"].Equal(NewString("X X")) {
 		t.Fatalf("replace_all mismatch: %v", out["replace_all"])
 	}
+	if !out["replace_all_anchor"].Equal(NewString("Xabc")) {
+		t.Fatalf("replace_all_anchor mismatch: %v", out["replace_all_anchor"])
+	}
 	if !out["replace_capture"].Equal(NewString("X-12 ID-34")) {
 		t.Fatalf("replace_capture mismatch: %v", out["replace_capture"])
 	}