regexp: add ReplaceAllSubmatchFunc

Fixes #5690

Author: enumag

src/regexp/all_test.go

@@ -263,6 +263,22 @@ var replaceFuncTests = []ReplaceFuncTest{
 	{"[a-c]*", func(s string) string { return "x" + s + "y" }, "defabcdef", "xydxyexyfxabcydxyexyfxy"},
 }
 
+type ReplaceFuncSubmatchTest struct {
+	pattern       string
+	replacement   func([]string) string
+	input, output string
+}
+
+var replaceFuncSubmatchTests = []ReplaceFuncSubmatchTest{
+	{"[a-c]+", func(g []string) string { return "x" + g[0] + "y" }, "defdef", "defdef"},
+	{"[a-c]", func(g []string) string { return "x" + g[0] + "y" }, "defabcdef", "defxayxbyxcydef"},
+	{"[a-c]+", func(g []string) string { return "x" + g[0] + "y" }, "defabcdef", "defxabcydef"},
+	{"[a-c]*", func(g []string) string { return "x" + g[0] + "y" }, "defabcdef", "xydxyexyfxabcydxyexyfxy"},
+	{"<<placeholder\\.(\\d+)>>", func(g []string) string { return "-" + g[1] + "-" }, "a<<placeholder.1>>b<<placeholder.2>>c<<placeholder.3>>d", "a-1-b-2-c-3-d"},
+	{"(\\w+)\\s+(\\w+)", func(g []string) string { return g[2] + " " + g[1] }, "hello world", "world hello"},
+	{"[aeiou]", func(g []string) string { return "" }, "hello", "hll"},
+}
+
 func TestReplaceAll(t *testing.T) {
 	for _, tc := range replaceTests {
 		re, err := Compile(tc.pattern)
@@ -350,6 +366,36 @@ func TestReplaceAllFunc(t *testing.T) {
 	}
 }
 
+func TestReplaceAllSubmatchFunc(t *testing.T) {
+	for _, tc := range replaceFuncSubmatchTests {
+		re, err := Compile(tc.pattern)
+		if err != nil {
+			t.Errorf("Unexpected error compiling %q: %v", tc.pattern, err)
+			continue
+		}
+		actual := re.ReplaceAllStringSubmatchFunc(tc.input, tc.replacement)
+		if actual != tc.output {
+			t.Errorf("%q.ReplaceFunc(%q,fn) = %q; want %q",
+				tc.pattern, tc.input, actual, tc.output)
+		}
+		// now try bytes
+		actual = string(re.ReplaceAllSubmatchFunc(
+			[]byte(tc.input),
+			func(g [][]byte) []byte {
+				stringGroups := make([]string, len(g))
+				for i, group := range g {
+					stringGroups[i] = string(group)
+				}
+				return []byte(tc.replacement(stringGroups))
+			},
+		))
+		if actual != tc.output {
+			t.Errorf("%q.ReplaceFunc(%q,fn) = %q; want %q",
+				tc.pattern, tc.input, actual, tc.output)
+		}
+	}
+}
+
 type MetaTest struct {
 	pattern, output, literal string
 	isLiteral                bool

src/regexp/regexp.go

@@ -580,6 +580,25 @@ func (re *Regexp) ReplaceAllStringFunc(src string, repl func(string) string) str
 	return string(b)
 }
 
+// ReplaceAllStringSubmatchFunc returns a copy of src in which all matches of the
+// [Regexp] have been replaced by the return value of function repl applied to the
+// submatch. The replacement returned by repl is substituted directly, without using
+// [Regexp.Expand].
+func (re *Regexp) ReplaceAllStringSubmatchFunc(src string, repl func(groups []string) string) string {
+	// Submatch positions are needed in this function, similar to when ReplaceAllString finds a $ in the replacement string.
+	n := 2 * (re.numSubexp + 1)
+	b := re.replaceAll(nil, src, n, func(dst []byte, match []int) []byte {
+		groups := make([]string, len(match)/2)
+		for i := 0; i < len(groups); i++ {
+			if match[2*i] >= 0 {
+				groups[i] = src[match[2*i]:match[2*i+1]]
+			}
+		}
+		return append(dst, repl(groups)...)
+	})
+	return string(b)
+}
+
 func (re *Regexp) replaceAll(bsrc []byte, src string, nmatch int, repl func(dst []byte, m []int) []byte) []byte {
 	lastMatchEnd := 0 // end position of the most recent match
 	searchPos := 0    // position where we next look for a match
@@ -682,6 +701,25 @@ func (re *Regexp) ReplaceAllFunc(src []byte, repl func([]byte) []byte) []byte {
 	})
 }
 
+// ReplaceAllSubmatchFunc returns a copy of src in which all matches of the
+// [Regexp] have been replaced by the return value of function repl applied to the
+// submatch. The replacement returned by repl is substituted directly, without using
+// [Regexp.Expand].
+func (re *Regexp) ReplaceAllSubmatchFunc(src []byte, repl func(groups [][]byte) []byte) []byte {
+	// Submatch positions are needed in this function, similar to when ReplaceAll finds a $ in the replacement string.
+	n := 2 * (re.numSubexp + 1)
+	b := re.replaceAll(src, "", n, func(dst []byte, match []int) []byte {
+		groups := make([][]byte, len(match)/2)
+		for i := 0; i < len(groups); i++ {
+			if match[2*i] >= 0 {
+				groups[i] = src[match[2*i]:match[2*i+1]:match[2*i+1]]
+			}
+		}
+		return append(dst, repl(groups)...)
+	})
+	return b
+}
+
 // Bitmap used by func special to check whether a character needs to be escaped.
 var specialBytes [16]byte