Preserve per-byte error replacements

4 files changed, 38 insertions, 0 deletions

diff --git a/tokenizer/bpe/split/fsa.go b/tokenizer/bpe/split/fsa.go
index a0c4cd1..b85254b 100644
--- a/tokenizer/bpe/split/fsa.go
+++ b/tokenizer/bpe/split/fsa.go
@@ -155,6 +155,13 @@ func (f *FSA) Read(next rune) bool {
 }
 
 func (f *FSA) FindAll(s string) []string {
+	if !utf8.ValidString(s) {
+		// s = strings.ToValidUTF8(s, string(utf8.RuneError)) replaces with a single utf8.RuneError
+		// s = strings.Map(func(r rune) rune { return r }, s) should be identical to s = string([]rune(s))
+
+		s = string([]rune(s)) // replaces each invalid byte with utf8.RuneError during decoding
+	}
+
 	var findAll func(start int, matches []string) []string
 
 	findAll = func(start int, matches []string) []string {
diff --git a/tokenizer/bpe/split/fsa_test.go b/tokenizer/bpe/split/fsa_test.go
index be11e79..9188d99 100644
--- a/tokenizer/bpe/split/fsa_test.go
+++ b/tokenizer/bpe/split/fsa_test.go
@@ -38,3 +38,30 @@ func BenchmarkFSA_FindAllReference(b *testing.B) {
 		_ = f.FindAll(s)
 	}
 }
+
+func FuzzFSA_FindAll(f *testing.F) {
+	fsa := NewFSA()
+
+	ref := mbpe.NewFSA()
+
+	f.Add("foo")
+	f.Add(" ")
+	f.Add("\n")
+	f.Add("   bar")
+
+	f.Fuzz(func(t *testing.T, s string) {
+		out := fsa.FindAll(s)
+
+		expected := ref.FindAll(s)
+
+		if len(out) != len(expected) {
+			t.Fatalf("expected %s but got %s", expected, out)
+		}
+
+		for i, m := range out {
+			if m != expected[i] {
+				t.Errorf("expected [%s] but got [%s]", expected[i], m)
+			}
+		}
+	})
+}
diff --git a/tokenizer/bpe/split/testdata/fuzz/FuzzFSA_FindAll/0c88d15026412e02 b/tokenizer/bpe/split/testdata/fuzz/FuzzFSA_FindAll/0c88d15026412e02
new file mode 100644
index 0000000..c5f9645
--- /dev/null
+++ b/tokenizer/bpe/split/testdata/fuzz/FuzzFSA_FindAll/0c88d15026412e02
@@ -0,0 +1,2 @@
+go test fuzz v1
+string("\xa2\x9d")
diff --git a/tokenizer/bpe/split/testdata/fuzz/FuzzFSA_FindAll/674ae72cbe3ff0c3 b/tokenizer/bpe/split/testdata/fuzz/FuzzFSA_FindAll/674ae72cbe3ff0c3
new file mode 100644
index 0000000..c854a03
--- /dev/null
+++ b/tokenizer/bpe/split/testdata/fuzz/FuzzFSA_FindAll/674ae72cbe3ff0c3
@@ -0,0 +1,2 @@
+go test fuzz v1
+string("\xef")