Implement ReplaceFunc (#105)

README.md

@@ -37,7 +37,6 @@ behavior differences.
 All APIs found in `regexp` are available except
 
 - `*Reader`: re2 does not support streaming input
-- `*Func`: re2 does not support replacement with callback functions
 
 Note that unlike many packages that wrap C++ libraries, there is no added `Close` type of method.
 See the [rationale](./RATIONALE.md) for more details.

all_test.go

@@ -244,6 +244,18 @@ var replaceLiteralTests = []ReplaceTest{
 	{"a+", "$", "aaa", "$"},
 }
 
+type ReplaceFuncTest struct {
+	pattern       string
+	replacement   func(string) string
+	input, output string
+}
+
+var replaceFuncTests = []ReplaceFuncTest{
+	{"[a-c]", func(s string) string { return "x" + s + "y" }, "defabcdef", "defxayxbyxcydef"},
+	{"[a-c]+", func(s string) string { return "x" + s + "y" }, "defabcdef", "defxabcydef"},
+	{"[a-c]*", func(s string) string { return "x" + s + "y" }, "defabcdef", "xydxyexyfxabcydxyexyfxy"},
+}
+
 func TestReplaceAll(t *testing.T) {
 	for _, tc := range replaceTests {
 		re, err := Compile(tc.pattern)
@@ -310,6 +322,27 @@ func TestReplaceAllLiteral(t *testing.T) {
 	}
 }
 
+func TestReplaceAllFunc(t *testing.T) {
+	for _, tc := range replaceFuncTests {
+		re, err := Compile(tc.pattern)
+		if err != nil {
+			t.Errorf("Unexpected error compiling %q: %v", tc.pattern, err)
+			continue
+		}
+		actual := re.ReplaceAllStringFunc(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.ReplaceAllFunc([]byte(tc.input), func(s []byte) []byte { return []byte(tc.replacement(string(s))) }))
+		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

example_test.go

@@ -2,6 +2,7 @@ package re2_test
 
 import (
 	"fmt"
+	"strings"
 
 	regexp "github.com/wasilibs/go-re2"
 )
@@ -255,13 +256,12 @@ func ExampleRegexp_ReplaceAllString() {
 	// -W-xxW-
 }
 
-// GAP: We cannot pass callback functions to re2
-//func ExampleRegexp_ReplaceAllStringFunc() {
-//	re := regexp.MustCompile(`[^aeiou]`)
-//	fmt.Println(re.ReplaceAllStringFunc("seafood fool", strings.ToUpper))
-//	// Output:
-//	// SeaFooD FooL
-//}
+func ExampleRegexp_ReplaceAllStringFunc() {
+	re := regexp.MustCompile(`[^aeiou]`)
+	fmt.Println(re.ReplaceAllStringFunc("seafood fool", strings.ToUpper))
+	// Output:
+	// SeaFooD FooL
+}
 
 func ExampleRegexp_SubexpNames() {
 	re := regexp.MustCompile(`(?P<first>[a-zA-Z]+) (?P<last>[a-zA-Z]+)`)

internal/re2.go

@@ -829,6 +829,21 @@ func (re *Regexp) ReplaceAll(src, repl []byte) []byte {
 	return b
 }
 
+// ReplaceAllFunc 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 matched byte slice. The replacement returned by repl is substituted
+// directly, without using [Regexp.Expand].
+func (re *Regexp) ReplaceAllFunc(src []byte, repl func([]byte) []byte) []byte {
+	alloc := re.abi.startOperation(len(src) + 8*re.numMatches + 8)
+	defer re.abi.endOperation(alloc)
+
+	cs := alloc.newCStringFromBytes(src)
+
+	return re.replaceAll(&alloc, src, "", cs, func(dst []byte, m []int) []byte {
+		return append(dst, repl(src[m[0]:m[1]])...)
+	})
+}
+
 // ReplaceAllLiteral returns a copy of src, replacing matches of the Regexp
 // with the replacement bytes repl. The replacement repl is substituted directly,
 // without using Expand.
@@ -875,6 +890,23 @@ func (re *Regexp) ReplaceAllString(src, repl string) string {
 	return string(b)
 }
 
+// ReplaceAllStringFunc 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 matched substring. The replacement returned by repl is substituted
+// directly, without using [Regexp.Expand].
+func (re *Regexp) ReplaceAllStringFunc(src string, repl func(string) string) string {
+	alloc := re.abi.startOperation(len(src) + 8*re.numMatches + 8)
+	defer re.abi.endOperation(alloc)
+
+	cs := alloc.newCString(src)
+
+	b := re.replaceAll(&alloc, nil, src, cs, func(dst []byte, m []int) []byte {
+		return append(dst, repl(src[m[0]:m[1]])...)
+	})
+
+	return string(b)
+}
+
 func (re *Regexp) replaceAll(alloc *allocation, bsrc []byte, src string, cs cString, repl func(dst []byte, m []int) []byte) []byte {
 	lastMatchEnd := 0
 	var buf []byte