Add ensure helper, tests

Author: earthboundkid

crockford.go

@@ -104,16 +104,26 @@ func AppendRandom(e *base32.Encoding, dst []byte) []byte {
 // AppendMD5 appends 26 encoded bytes generated by MD5 hashing src onto dst.
 func AppendMD5(e *base32.Encoding, dst, src []byte) []byte {
 	//16 bytes -> 26 base32 characters
-	var (
-		buf [16]byte
-		enc [26]byte
-	)
+	var buf [md5.Size]byte
+
 	h := md5.New()
 	h.Write(src)
 	h.Sum(buf[:0])
 
 	// Ensure dst has 26 bytes capacity
-	dst = append(dst, enc[:]...)
-	e.Encode(dst[len(dst)-26:], buf[:])
-	return dst
+	ret, tar := ensure(26, dst)
+	e.Encode(tar, buf[:])
+	return ret
+}
+
+func ensure(size int, b []byte) (ret, tar []byte) {
+	ret = b
+	newLen := len(b) + size
+	if cap(b) >= newLen {
+		ret = b[:newLen]
+	} else {
+		ret = append(b, make([]byte, size)...)
+	}
+	tar = ret[len(b):newLen]
+	return
 }

crockford_test.go

@@ -1,6 +1,7 @@
 package crockford_test
 
 import (
+	"bytes"
 	"testing"
 
 	"github.com/carlmjohnson/crockford"
@@ -15,6 +16,38 @@ func EqBytes(t *testing.T) func(want string, got []byte) {
 	}
 }
 
+func TestEnsure(t *testing.T) {
+	cases := map[string]struct {
+		size int
+		b    []byte
+	}{
+		"0-nil":    {0, nil},
+		"4-nil":    {4, nil},
+		"0-sliced": {4, []byte("1234")[:0]},
+		"2-sliced": {2, []byte("1234")[:2]},
+		"overflow": {2, []byte("1234")},
+	}
+	for name, tc := range cases {
+		t.Run(name, func(t *testing.T) {
+			ret, tar := crockford.Ensure(tc.size, tc.b)
+			if len(tar) != tc.size {
+				t.Fatalf("bad target: %q", tar)
+			}
+			if len(ret) != len(tc.b)+tc.size {
+				t.Fatalf("bad return: %q", ret)
+			}
+			if cap(tc.b)-len(tc.b) >= tc.size {
+				if bytes.ContainsAny(tar, "\x00") {
+					t.Fatalf("overwrote existing cap: %q", ret)
+				}
+			}
+			if !bytes.HasPrefix(ret, tc.b) {
+				t.Fatalf("lost prefix: %q", ret)
+			}
+		})
+	}
+}
+
 func TestAppendMD5(t *testing.T) {
 	cases := map[string]struct {
 		in   string
@@ -41,10 +74,11 @@ func TestAppendMD5(t *testing.T) {
 
 			r := testing.Benchmark(func(b *testing.B) {
 				for i := 0; i < b.N; i++ {
+					dst = dst[:0]
 					dst = crockford.AppendMD5(crockford.Lower, dst, in)
 				}
 			})
-			if r.AllocsPerOp() > 0 {
+			if r.AllocsPerOp() != 0 {
 				t.Errorf("benchmark regression %q: %v", dst, r.MemString())
 			}
 		})

export_test.go

@@ -0,0 +1,3 @@
+package crockford
+
+var Ensure = ensure