Update base64 string handling and adjust LUT (#51)

* add internal string to bytes conversion

* use a slice for lookup tables

* use pointer for lut

* add test to verify issue #50

Author: kalamay

base64/base64_amd64.go

@@ -2,19 +2,19 @@ package base64
 
 import (
 	"encoding/base64"
-	"unsafe"
 
 	"github.com/segmentio/asm/cpu"
+	"github.com/segmentio/asm/internal/unsafebytes"
 )
 
 // An Encoding is a radix 64 encoding/decoding scheme, defined by a
 // 64-character alphabet.
 type Encoding struct {
-	enc    func(dst []byte, src []byte, lut [16]int8) (int, int)
-	enclut [16]int8
+	enc    func(dst []byte, src []byte, lut *int8) (int, int)
+	enclut [32]int8
 
-	dec    func(dst []byte, src []byte, lut [32]int8) (int, int)
-	declut [32]int8
+	dec    func(dst []byte, src []byte, lut *int8) (int, int)
+	declut [48]int8
 
 	base *base64.Encoding
 }
@@ -42,7 +42,7 @@ func (e *Encoding) enableEncodeAVX2(encoder string) {
 	// [52..61]  [48..57]    -4  [2..11]  0123456789
 	// [62]      [43]       -19       12  +
 	// [63]      [47]       -16       13  /
-	tab := [16]int8{int8(encoder[0]), int8(encoder[letterRange]) - letterRange}
+	tab := [32]int8{int8(encoder[0]), int8(encoder[letterRange]) - letterRange}
 	for i, ch := range encoder[2*letterRange:] {
 		tab[2+i] = int8(ch) - 2*letterRange - int8(i)
 	}
@@ -67,7 +67,7 @@ func (e *Encoding) enableDecodeAVX2(encoder string) {
 	// [48..57]   [52..61]   +4        3  0123456789
 	// [65..90]   [0..25]   -65      4,5  ABCDEFGHIJKLMNOPQRSTUVWXYZ
 	// [97..122]  [26..51]  -71      6,7  abcdefghijklmnopqrstuvwxyz
-	tab := [32]int8{
+	tab := [48]int8{
 		0, 63 - c63, 62 - c62, 4, -65, -65, -71, -71,
 		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
 		0x15, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11,
@@ -104,7 +104,7 @@ func (enc Encoding) Strict() *Encoding {
 // This will write EncodedLen(len(src)) bytes to dst.
 func (enc *Encoding) Encode(dst, src []byte) {
 	if len(src) >= minEncodeLen && enc.enc != nil {
-		d, s := enc.enc(dst, src, enc.enclut)
+		d, s := enc.enc(dst, src, &enc.enclut[0])
 		dst = dst[d:]
 		src = src[s:]
 	}
@@ -131,7 +131,7 @@ func (enc *Encoding) EncodedLen(n int) int {
 func (enc *Encoding) Decode(dst, src []byte) (n int, err error) {
 	var d, s int
 	if len(src) >= minDecodeLen && enc.dec != nil {
-		d, s = enc.dec(dst, src, enc.declut)
+		d, s = enc.dec(dst, src, &enc.declut[0])
 		dst = dst[d:]
 		src = src[s:]
 	}
@@ -143,7 +143,7 @@ func (enc *Encoding) Decode(dst, src []byte) (n int, err error) {
 // DecodeString decodes the base64 encoded string s, returns the decoded
 // value as bytes.
 func (enc *Encoding) DecodeString(s string) ([]byte, error) {
-	src := *(*[]byte)(unsafe.Pointer(&s))
+	src := unsafebytes.BytesOf(s)
 	dst := make([]byte, enc.base.DecodedLen(len(s)))
 	n, err := enc.Decode(dst, src)
 	return dst[:n], err

base64/base64_amd64_test.go

@@ -46,7 +46,7 @@ func TestEncodeAVX2(t *testing.T) {
 				}
 				defer dst.Release()
 
-				_, ns := enc.candidate.enc(dst.ProtectTail(), buf, enc.candidate.enclut)
+				_, ns := enc.candidate.enc(dst.ProtectTail(), buf, &enc.candidate.enclut[0])
 
 				if len(buf)-ns >= 32 {
 					t.Errorf("encode remain should be less than 32, but is %d", len(buf)-ns)
@@ -86,7 +86,7 @@ func TestDecodeAVX2(t *testing.T) {
 
 				enc.candidate.Encode(src, buf)
 
-				_, ns := enc.candidate.dec(dst.ProtectTail(), src, enc.candidate.declut)
+				_, ns := enc.candidate.dec(dst.ProtectTail(), src, &enc.candidate.declut[0])
 
 				if len(buf)-ns >= 45 {
 					t.Errorf("decode remain should be less than 45, but is %d", len(buf)-ns)

base64/base64_test.go

@@ -79,6 +79,22 @@ func TestEncoding(t *testing.T) {
 				if !bytes.Equal(decExpect, decActual) {
 					t.Fatalf("failed decode:\n\texpect = %v\n\tactual = %v", decExpect, decActual)
 				}
+
+				encString := enc.control.EncodeToString(src)
+				decExpect, errControl = enc.control.DecodeString(encString)
+				decActual, errCandidate = enc.candidate.DecodeString(encString)
+
+				if errControl != nil {
+					t.Fatalf("control decode error: %v", errControl)
+				}
+
+				if errCandidate != nil {
+					t.Fatalf("candidate decode error: %v", errCandidate)
+				}
+
+				if !bytes.Equal(decExpect, decActual) {
+					t.Fatalf("failed decode:\n\texpect = %v\n\tactual = %v", decExpect, decActual)
+				}
 			}
 		})
 	}

base64/decode_amd64.go

@@ -32,20 +32,21 @@ DATA b64_dec_shuf<>+16(SB)/8, $0x0c0d0e08090a0405
 DATA b64_dec_shuf<>+24(SB)/8, $0x0000000000000000
 GLOBL b64_dec_shuf<>(SB), RODATA|NOPTR, $32
 
-// func decodeAVX2(dst []byte, src []byte, lut [32]int8) (int, int)
+// func decodeAVX2(dst []byte, src []byte, lut *int8) (int, int)
 // Requires: AVX, AVX2, SSE4.1
-TEXT ·decodeAVX2(SB), NOSPLIT, $0-96
+TEXT ·decodeAVX2(SB), NOSPLIT, $0-72
 	MOVQ         dst_base+0(FP), AX
 	MOVQ         src_base+24(FP), DX
-	MOVQ         src_len+32(FP), SI
+	MOVQ         lut+48(FP), SI
+	MOVQ         src_len+32(FP), DI
 	MOVB         $0x2f, CL
 	PINSRB       $0x00, CX, X8
 	VPBROADCASTB X8, Y8
 	XORQ         CX, CX
 	XORQ         BX, BX
 	VPXOR        Y7, Y7, Y7
-	VPERMQ       $0x44, lut_0+48(FP), Y6
-	VPERMQ       $0x44, lut_0+64(FP), Y4
+	VPERMQ       $0x44, (SI), Y6
+	VPERMQ       $0x44, 16(SI), Y4
 	VMOVDQA      b64_dec_lut_hi<>+0(SB), Y5
 
 loop:
@@ -71,20 +72,20 @@ loop:
 	VMOVDQU      Y1, (AX)(CX*1)
 	ADDQ         $0x18, CX
 	ADDQ         $0x20, BX
-	SUBQ         $0x20, SI
-	CMPQ         SI, $0x2d
+	SUBQ         $0x20, DI
+	CMPQ         DI, $0x2d
 	JB           done
 	JMP          loop
 
 done:
-	MOVQ CX, ret+80(FP)
-	MOVQ BX, ret1+88(FP)
+	MOVQ CX, ret+56(FP)
+	MOVQ BX, ret1+64(FP)
 	VZEROUPPER
 	RET
 
-// func decodeAVX2URI(dst []byte, src []byte, lut [32]int8) (int, int)
+// func decodeAVX2URI(dst []byte, src []byte, lut *int8) (int, int)
 // Requires: AVX, AVX2, SSE4.1
-TEXT ·decodeAVX2URI(SB), NOSPLIT, $0-96
+TEXT ·decodeAVX2URI(SB), NOSPLIT, $0-72
 	MOVB         $0x2f, AL
 	PINSRB       $0x00, AX, X0
 	VPBROADCASTB X0, Y0
@@ -93,15 +94,16 @@ TEXT ·decodeAVX2URI(SB), NOSPLIT, $0-96
 	VPBROADCASTB X1, Y1
 	MOVQ         dst_base+0(FP), AX
 	MOVQ         src_base+24(FP), DX
-	MOVQ         src_len+32(FP), SI
+	MOVQ         lut+48(FP), SI
+	MOVQ         src_len+32(FP), DI
 	MOVB         $0x2f, CL
 	PINSRB       $0x00, CX, X10
 	VPBROADCASTB X10, Y10
 	XORQ         CX, CX
 	XORQ         BX, BX
 	VPXOR        Y9, Y9, Y9
-	VPERMQ       $0x44, lut_0+48(FP), Y8
-	VPERMQ       $0x44, lut_0+64(FP), Y6
+	VPERMQ       $0x44, (SI), Y8
+	VPERMQ       $0x44, 16(SI), Y6
 	VMOVDQA      b64_dec_lut_hi<>+0(SB), Y7
 
 loop:
@@ -129,13 +131,13 @@ loop:
 	VMOVDQU      Y3, (AX)(CX*1)
 	ADDQ         $0x18, CX
 	ADDQ         $0x20, BX
-	SUBQ         $0x20, SI
-	CMPQ         SI, $0x2d
+	SUBQ         $0x20, DI
+	CMPQ         DI, $0x2d
 	JB           done
 	JMP          loop
 
 done:
-	MOVQ CX, ret+80(FP)
-	MOVQ BX, ret1+88(FP)
+	MOVQ CX, ret+56(FP)
+	MOVQ BX, ret1+64(FP)
 	VZEROUPPER
 	RET

base64/decode_amd64.s

@@ -4,12 +4,13 @@
 
 #include "textflag.h"
 
-// func encodeAVX2(dst []byte, src []byte, lut [16]int8) (int, int)
+// func encodeAVX2(dst []byte, src []byte, lut *int8) (int, int)
 // Requires: AVX, AVX2, SSE4.1
-TEXT ·encodeAVX2(SB), NOSPLIT, $0-80
+TEXT ·encodeAVX2(SB), NOSPLIT, $0-72
 	MOVQ         dst_base+0(FP), AX
 	MOVQ         src_base+24(FP), DX
-	MOVQ         src_len+32(FP), SI
+	MOVQ         lut+48(FP), SI
+	MOVQ         src_len+32(FP), DI
 	MOVB         $0x33, CL
 	PINSRB       $0x00, CX, X4
 	VPBROADCASTB X4, Y4
@@ -20,7 +21,7 @@ TEXT ·encodeAVX2(SB), NOSPLIT, $0-80
 	XORQ         BX, BX
 
 	// Load the 16-byte LUT into both lanes of the register
-	VPERMQ $0x44, lut_0+48(FP), Y3
+	VPERMQ $0x44, (SI), Y3
 
 	// Load the first block using a mask to avoid potential fault
 	VMOVDQU    b64_enc_load<>+0(SB), Y0
@@ -43,15 +44,15 @@ loop:
 	VMOVDQU  Y0, (AX)(CX*1)
 	ADDQ     $0x20, CX
 	ADDQ     $0x18, BX
-	SUBQ     $0x18, SI
-	CMPQ     SI, $0x20
+	SUBQ     $0x18, DI
+	CMPQ     DI, $0x20
 	JB       done
 	VMOVDQU  -4(DX)(BX*1), Y0
 	JMP      loop
 
 done:
-	MOVQ CX, ret+64(FP)
-	MOVQ BX, ret1+72(FP)
+	MOVQ CX, ret+56(FP)
+	MOVQ BX, ret1+64(FP)
 	VZEROUPPER
 	RET
 

base64/encode_amd64.go

@@ -47,12 +47,12 @@ func init() {
 }
 
 func main() {
-	TEXT("decodeAVX2", NOSPLIT, "func(dst, src []byte, lut [32]int8) (int, int)")
+	TEXT("decodeAVX2", NOSPLIT, "func(dst, src []byte, lut *int8) (int, int)")
 	createDecode(Param("dst"), Param("src"), Param("lut"), func(m Mem, r VecVirtual) {
 		VMOVDQU(m, r)
 	})
 
-	TEXT("decodeAVX2URI", NOSPLIT, "func(dst, src []byte, lut [32]int8) (int, int)")
+	TEXT("decodeAVX2URI", NOSPLIT, "func(dst, src []byte, lut *int8) (int, int)")
 	slash := VecBroadcast(U8('/'), YMM())
 	underscore := VecBroadcast(U8('_'), YMM())
 	createDecode(Param("dst"), Param("src"), Param("lut"), func(m Mem, r VecVirtual) {
@@ -68,11 +68,8 @@ func main() {
 func createDecode(pdst, psrc, plut Component, load func(m Mem, r VecVirtual)) {
 	dst := Mem{Base: Load(pdst.Base(), GP64()), Index: GP64(), Scale: 1}
 	src := Mem{Base: Load(psrc.Base(), GP64()), Index: GP64(), Scale: 1}
+	lut := Mem{Base: Load(plut, GP64())}
 	rem := Load(psrc.Len(), GP64())
-	lut, err := plut.Index(0).Resolve()
-	if err != nil {
-		panic(err)
-	}
 
 	rsrc := YMM()
 	rdst := YMM()
@@ -93,8 +90,8 @@ func createDecode(pdst, psrc, plut Component, load func(m Mem, r VecVirtual)) {
 	XORQ(src.Index, src.Index)
 	VPXOR(zero, zero, zero)
 
-	VPERMQ(Imm(1<<6|1<<2), lut.Addr, lutr)
-	VPERMQ(Imm(1<<6|1<<2), lut.Addr.Offset(16), lutl)
+	VPERMQ(Imm(1<<6|1<<2), lut, lutr)
+	VPERMQ(Imm(1<<6|1<<2), lut.Offset(16), lutl)
 	VMOVDQA(lutHi, luth)
 
 	Label("loop")

base64/encode_amd64.s

@@ -21,12 +21,12 @@ func init() {
 }
 
 func main() {
-	TEXT("encodeAVX2", NOSPLIT, "func(dst, src []byte, lut [16]int8) (int, int)")
+	TEXT("encodeAVX2", NOSPLIT, "func(dst, src []byte, lut *int8) (int, int)")
 
 	dst := Mem{Base: Load(Param("dst").Base(), GP64()), Index: GP64(), Scale: 1}
 	src := Mem{Base: Load(Param("src").Base(), GP64()), Index: GP64(), Scale: 1}
+	lut := Mem{Base: Load(Param("lut"), GP64())}
 	rem := Load(Param("src").Len(), GP64())
-	lut, _ := Param("lut").Index(0).Resolve()
 
 	rsrc := YMM()
 	rdst := YMM()
@@ -47,7 +47,7 @@ func main() {
 	XORQ(src.Index, src.Index)
 
 	Comment("Load the 16-byte LUT into both lanes of the register")
-	VPERMQ(Imm(1<<6|1<<2), lut.Addr, xtab)
+	VPERMQ(Imm(1<<6|1<<2), lut, xtab)
 
 	Comment("Load the first block using a mask to avoid potential fault")
 	VMOVDQU(ConstLoadMask32("b64_enc_load",

build/base64/decode_asm.go

@@ -9,3 +9,12 @@ func Pointer(b []byte) *byte {
 func String(b []byte) string {
 	return *(*string)(unsafe.Pointer(&b))
 }
+
+func BytesOf(s string) []byte {
+	return *(*[]byte)(unsafe.Pointer(&sliceHeader{str: s, cap: len(s)}))
+}
+
+type sliceHeader struct {
+	str string
+	cap int
+}