From 47fea26b6282e6a340ff45c2f07befb74c55efde Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Thu, 7 Mar 2024 09:36:40 +0530
Subject: [PATCH] Add an IndexByte implementation useful for benchmarking
 against stdlib SIMD implementation

---
 tools/simdstring/generate.go        | 22 ++++++++++++++++++++++
 tools/simdstring/intrinsics.go      | 11 +++++++++++
 tools/simdstring/intrinsics_test.go | 20 ++++++++++++++++++++
 tools/simdstring/scalar.go          | 18 ++++++++++++++++++
 4 files changed, 71 insertions(+)

diff --git a/tools/simdstring/generate.go b/tools/simdstring/generate.go
index e24461c1f..df3e776be 100644
--- a/tools/simdstring/generate.go
+++ b/tools/simdstring/generate.go
@@ -1451,6 +1451,27 @@ func (s *State) indexbyte2() {
 
 }
 
+func (s *State) indexbyte_body(f *Function) {
+	b := f.Vec()
+	f.Set1Epi8("b", b)
+	test_bytes := func(bytes_to_test, test_ans Register) {
+		f.CmpEqEpi8(bytes_to_test, b, test_ans)
+	}
+	s.index_func(f, test_bytes)
+}
+
+func (s *State) indexbyte() {
+	f := s.NewFunction("index_byte_asm", "Find the index of a byte", []FunctionParam{{"data", ByteSlice}, {"b", types.Byte}}, []FunctionParam{{"ans", types.Int}})
+	if s.ISA.HasSIMD {
+		s.indexbyte_body(f)
+	}
+	f = s.NewFunction("index_byte_string_asm", "Find the index of a byte", []FunctionParam{{"data", types.String}, {"b", types.Byte}}, []FunctionParam{{"ans", types.Int}})
+	if s.ISA.HasSIMD {
+		s.indexbyte_body(f)
+	}
+
+}
+
 func (s *State) indexc0_body(f *Function) {
 	lower := f.Vec()
 	upper := f.Vec()
@@ -1494,6 +1515,7 @@ func (s *State) Generate() {
 
 	s.indexbyte2()
 	s.indexc0()
+	s.indexbyte()
 
 	s.OutputFunction()
 }
diff --git a/tools/simdstring/intrinsics.go b/tools/simdstring/intrinsics.go
index acf33e034..413a9e3a3 100644
--- a/tools/simdstring/intrinsics.go
+++ b/tools/simdstring/intrinsics.go
@@ -12,6 +12,13 @@ var Have128bit = false
 var Have256bit = false
 var VectorSize = 1
 
+// Return the index at which b first occurs in data. If not found -1 is returned.
+var IndexByte func(data []byte, b byte) int = index_byte_scalar
+
+// Return the index at which either a or b first occurs in text. If neither is
+// found -1 is returned.
+var IndexByteString func(text string, b byte) int = index_byte_string_scalar
+
 // Return the index at which either a or b first occurs in data. If neither is
 // found -1 is returned.
 var IndexByte2 func(data []byte, a, b byte) int = index_byte2_scalar
@@ -38,12 +45,16 @@ func init() {
 		Have256bit = HasSIMD256Code
 	}
 	if Have256bit {
+		IndexByte = index_byte_asm_256
+		IndexByteString = index_byte_string_asm_256
 		IndexByte2 = index_byte2_asm_256
 		IndexByte2String = index_byte2_string_asm_256
 		IndexC0 = index_c0_asm_256
 		IndexC0String = index_c0_string_asm_256
 		VectorSize = 32
 	} else if Have128bit {
+		IndexByte = index_byte_asm_128
+		IndexByteString = index_byte_string_asm_128
 		IndexByte2 = index_byte2_asm_128
 		IndexByte2String = index_byte2_string_asm_128
 		IndexC0 = index_c0_asm_128
diff --git a/tools/simdstring/intrinsics_test.go b/tools/simdstring/intrinsics_test.go
index 79e8203db..e848ddc41 100644
--- a/tools/simdstring/intrinsics_test.go
+++ b/tools/simdstring/intrinsics_test.go
@@ -219,6 +219,26 @@ func TestSIMDStringOps(t *testing.T) {
 	c0tests("afsgdfg\x7f")
 	c0tests("afgd\x1bfgd\t")
 	c0tests("a\x00")
+
+	index_test := func(haystack []byte, needle byte) {
+		var actual int
+		expected := index_byte_scalar(haystack, needle)
+
+		for _, sz := range sizes {
+			switch sz {
+			case 16:
+				actual = index_byte_asm_128(haystack, needle)
+			case 32:
+				actual = index_byte_asm_256(haystack, needle)
+			}
+			if actual != expected {
+				t.Fatalf("index failed in: %#v (%d != %d) at size: %d with needle: %#v", string(haystack), expected, actual, sz, needle)
+			}
+		}
+	}
+	index_test([]byte("abc"), 'x')
+	index_test([]byte("abc"), 'b')
+
 }
 
 func TestIntrinsics(t *testing.T) {
diff --git a/tools/simdstring/scalar.go b/tools/simdstring/scalar.go
index d00130328..5320aebe8 100644
--- a/tools/simdstring/scalar.go
+++ b/tools/simdstring/scalar.go
@@ -2,6 +2,24 @@
 
 package simdstring
 
+func index_byte_scalar(data []byte, b byte) int {
+	for i, ch := range data {
+		if ch == b {
+			return i
+		}
+	}
+	return -1
+}
+
+func index_byte_string_scalar(data string, b byte) int {
+	for i := 0; i < len(data); i++ {
+		if data[i] == b {
+			return i
+		}
+	}
+	return -1
+}
+
 func index_byte2_scalar(data []byte, a, b byte) int {
 	for i, ch := range data {
 		switch ch {