Fix Go unaligned index implementation

tools/simdstring/generate.go

@@ -424,10 +424,18 @@ func (f *Function) MaskForCountDestructive(vec, ans Register) {
 	}
 }
 
-func (f *Function) ShiftSelfRight(self, amt any) {
-	op := "SHRQ"
-	if f.ISA.Goarch == ARM64 {
-		op = "LSR"
+func (f *Function) shift_self(right bool, self, amt any) {
+	op := ""
+	if right {
+		op = "SHRQ"
+		if f.ISA.Goarch == ARM64 {
+			op = "LSR"
+		}
+	} else {
+		op = "SHLQ"
+		if f.ISA.Goarch == ARM64 {
+			op = "LSL"
+		}
 	}
 	switch v := amt.(type) {
 	case Register:
@@ -440,13 +448,31 @@ func (f *Function) ShiftSelfRight(self, amt any) {
 	}
 }
 
-func (f *Function) ShiftMaskRightDestructive(mask, amt Register) {
+func (f *Function) ShiftSelfRight(self, amt any) {
+	f.shift_self(true, self, amt)
+}
+
+func (f *Function) ShiftSelfLeft(self, amt any) {
+	f.shift_self(false, self, amt)
+}
+
+func (f *Function) ShiftMaskRightDestructive(mask, amt any) {
 	// The amt register is clobbered by this function
-	if f.ISA.Goarch == ARM64 {
-		f.Comment("The mask has 4 bits per byte, so multiply", amt, "by 4")
-		f.ShiftSelfRight(amt, 2)
+	switch n := amt.(type) {
+	case Register:
+		if f.ISA.Goarch == ARM64 {
+			f.Comment("The mask has 4 bits per byte, so multiply", n, "by 4")
+			f.ShiftSelfLeft(n, 2)
+		}
+		f.ShiftSelfRight(mask, n)
+	case int:
+		if f.ISA.Goarch == ARM64 {
+			n <<= 2
+		}
+		f.ShiftSelfRight(mask, n)
+	default:
+		panic(fmt.Sprintf("Cannot shift by: %s", amt))
 	}
-	f.ShiftSelfRight(mask, amt)
 }
 
 func (f *Function) CountLeadingZeroBytesInMask(src, ans Register) {
@@ -1101,6 +1127,15 @@ func (f *Function) AndSelf(self Register, val any) {
 	f.AddTrailingComment(self, "&=", val)
 }
 
+func (f *Function) NegateSelf(self Register) {
+	if f.ISA.Goarch == "ARM64" {
+		f.instr("NEG", self, self)
+	} else {
+		f.instr("NEGQ", self)
+	}
+	f.AddTrailingComment(self, "*= -1")
+}
+
 func (f *Function) AddToSelf(self Register, val any) {
 	f.instr(f.ISA.NativeAdd(), val_repr_for_arithmetic(val), self) // pos += sizeof(vec)
 	f.AddTrailingComment(self, "+=", val)
@@ -1379,7 +1414,9 @@ func (s *State) index_func(f *Function, test_bytes func(bytes_to_test, test_ans
 		f.MaskForCountDestructive(test_ans, mask)
 		f.Comment("We need to shift out the possible extra bytes at the start of the string caused by the unaligned read")
 		f.ShiftMaskRightDestructive(mask, unaligned_bytes)
-		f.JumpIfNonZero(mask, "byte_found_in_mask")
+		f.JumpIfZero(mask, "loop_start")
+		f.CopyRegister(data_start, pos)
+		f.JumpTo("byte_found_in_mask")
 	}()
 
 	f.Comment("Now loop over aligned blocks")

tools/simdstring/intrinsics_test.go

@@ -5,6 +5,7 @@ package simdstring
 import (
 	"bytes"
 	"fmt"
+	"kitty/tools/utils"
 	"runtime"
 	"strings"
 	"testing"
@@ -119,21 +120,25 @@ func addressof_data(b []byte) uintptr {
 	return uintptr(unsafe.Pointer(&b[0]))
 }
 
-func aligned_slice(sz, alignment int) []byte {
-	ans := make([]byte, sz+alignment)
+func aligned_slice(sz, alignment int) ([]byte, []byte) {
+	ans := make([]byte, sz+alignment+512)
 	a := addressof_data(ans)
 	a &= uintptr(alignment - 1)
 	extra := uintptr(alignment) - a
-	return ans[extra : extra+uintptr(sz)]
+	utils.Memset(ans, '<')
+	utils.Memset(ans[extra+uintptr(sz):], '>')
+	return ans[extra : extra+uintptr(sz)], ans
 }
 
 func TestSIMDStringOps(t *testing.T) {
 	sizes := get_sizes(t)
-	test := func(haystack []byte, a, b byte) {
+	test := func(haystack []byte, a, b byte, dont_pad ...bool) {
 		var actual int
-		safe_haystack := append(bytes.Repeat([]byte{'<'}, 64), haystack...)
-		safe_haystack = append(safe_haystack, bytes.Repeat([]byte{'>'}, 64)...)
-		haystack = safe_haystack[64 : 64+len(haystack)]
+		if len(dont_pad) == 0 || !dont_pad[0] {
+			safe_haystack := append(bytes.Repeat([]byte{'<'}, 64), haystack...)
+			safe_haystack = append(safe_haystack, bytes.Repeat([]byte{'>'}, 64)...)
+			haystack = safe_haystack[64 : 64+len(haystack)]
+		}
 		expected := index_byte2_scalar(haystack, a, b)
 
 		for _, sz := range sizes {
@@ -144,21 +149,24 @@ func TestSIMDStringOps(t *testing.T) {
 				actual = index_byte2_asm_256(haystack, a, b)
 			}
 			if actual != expected {
-				t.Fatalf("Failed to find '%c' or '%c' in: %#v (%d != %d) at size: %d", a, b, string(haystack), expected, actual, sz)
+				t.Fatalf("Failed to find '%c' or '%c' in: %#v at align: %d (expected: %d != actual: %d) at size: %d",
+					a, b, string(haystack), addressof_data(haystack)&uintptr(sz-1), expected, actual, sz)
 			}
 		}
 
 	}
 	// test alignment issues
 	for sz := 0; sz < 32; sz++ {
-		q := aligned_slice(sz+3, 64)[sz:]
+		as, _ := aligned_slice(sz+3, 32)
+		q := as[sz:]
 		q[0] = 'a'
 		q[1] = 'b'
 		q[2] = 'c'
-		test(q, '<', '>')
-		test(q, '<', 'a')
-		test(q, '<', 'b')
-		test(q, 'c', '>')
+		test(q, '<', '>', true)
+		test(q, ' ', 'b', true)
+		test(q, '<', 'a', true)
+		test(q, '<', 'b', true)
+		test(q, 'c', '>', true)
 	}
 
 	tests := func(h string, a, b byte) {

tools/simdstring/test.sh

@@ -7,6 +7,6 @@
 #
 
 echo -e "\x1b[32mtesting amd64\x1b[m" && go test -v &&
-    echo -e "\x1b[32mtesting arm64\x1b[m" && GOARCH=arm64 go test -c && qemu-aarch64-static simdstring.test -test.v
+    echo -e "\x1b[32mtesting arm64\x1b[m" && GOARCH=arm64 go test -v -exec qemu-aarch64-static