Add missing bin header functions

There is no bytes equivalent of `(m *Reader) ReadBytesHeader()` and `(mw *Writer) WriteBytesHeader`.

Add `ReadBytesHeader(b []byte)` and `AppendBytesHeader(b []byte, sz uint32)`.

Author: klauspost

msgp/read_bytes.go

@@ -253,6 +253,46 @@ func ReadArrayHeaderBytes(b []byte) (sz uint32, o []byte, err error) {
 	}
 }
 
+// ReadBytesHeader reads the 'bin' header size
+// off of 'b' and returns the size and remaining bytes.
+// Possible errors:
+// - ErrShortBytes (too few bytes)
+// - TypeError{} (not a bin object)
+func ReadBytesHeader(b []byte) (sz uint32, o []byte, err error) {
+	if len(b) < 1 {
+		return 0, nil, ErrShortBytes
+	}
+	switch b[0] {
+	case mbin8:
+		if len(b) < 2 {
+			err = ErrShortBytes
+			return
+		}
+		sz = uint32(b[1])
+		o = b[2:]
+		return
+	case mbin16:
+		if len(b) < 3 {
+			err = ErrShortBytes
+			return
+		}
+		sz = uint32(big.Uint16(b[1:]))
+		o = b[3:]
+		return
+	case mbin32:
+		if len(b) < 5 {
+			err = ErrShortBytes
+			return
+		}
+		sz = big.Uint32(b[1:])
+		o = b[5:]
+		return
+	default:
+		err = badPrefix(BinType, b[0])
+		return
+	}
+}
+
 // ReadNilBytes tries to read a "nil" byte
 // off of 'b' and return the remaining bytes.
 // Possible errors:

msgp/read_bytes_test.go

@@ -98,6 +98,50 @@ func BenchmarkReadArrayHeaderBytes(b *testing.B) {
 	}
 }
 
+func TestReadBytesHeader(t *testing.T) {
+	var buf bytes.Buffer
+	en := NewWriter(&buf)
+
+	tests := []uint32{0, 1, 5, 49082, 1 << 16, math.MaxUint32}
+
+	for i, v := range tests {
+		buf.Reset()
+		en.WriteBytesHeader(v)
+		en.Flush()
+
+		out, left, err := ReadBytesHeader(buf.Bytes())
+		if err != nil {
+			t.Errorf("test case %d: %s", i, err)
+		}
+
+		if len(left) != 0 {
+			t.Errorf("expected 0 bytes left; found %d", len(left))
+		}
+
+		if out != v {
+			t.Errorf("%d in; %d out", v, out)
+		}
+	}
+}
+
+func BenchmarkTestReadBytesHeader(b *testing.B) {
+	sizes := []uint32{1, 100, tuint16, tuint32}
+	buf := make([]byte, 0, 5*len(sizes))
+	for _, sz := range sizes {
+		buf = AppendBytesHeader(buf, sz)
+	}
+	b.SetBytes(int64(len(buf) / len(sizes)))
+	b.ReportAllocs()
+	b.ResetTimer()
+	o := buf
+	for i := 0; i < b.N; i++ {
+		_, buf, _ = ReadBytesHeader(buf)
+		if len(buf) == 0 {
+			buf = o
+		}
+	}
+}
+
 func TestReadNilBytes(t *testing.T) {
 	var buf bytes.Buffer
 	en := NewWriter(&buf)

msgp/write_bytes.go

@@ -193,6 +193,26 @@ func AppendBytes(b []byte, bts []byte) []byte {
 	return o[:n+copy(o[n:], bts)]
 }
 
+// AppendBytesHeader appends an 'bin' header with
+// the given size to the slice.
+func AppendBytesHeader(b []byte, sz uint32) []byte {
+	var o []byte
+	var n int
+	switch {
+	case sz <= math.MaxUint8:
+		o, n = ensure(b, 2)
+		prefixu8(o[n:], mbin8, uint8(sz))
+		return o
+	case sz <= math.MaxUint16:
+		o, n = ensure(b, 3)
+		prefixu16(o[n:], mbin16, uint16(sz))
+		return o
+	}
+	o, n = ensure(b, 5)
+	prefixu32(o[n:], mbin32, sz)
+	return o
+}
+
 // AppendBool appends a bool to the slice
 func AppendBool(b []byte, t bool) []byte {
 	if t {

msgp/write_bytes_test.go

@@ -93,6 +93,37 @@ func BenchmarkAppendArrayHeader(b *testing.B) {
 	}
 }
 
+func TestAppendBytesHeader(t *testing.T) {
+	szs := []uint32{0, 1, uint32(tint8), uint32(tint16), tuint32, math.MaxUint32}
+	var buf bytes.Buffer
+	en := NewWriter(&buf)
+
+	var bts []byte
+	for _, sz := range szs {
+		buf.Reset()
+		en.WriteBytesHeader(sz)
+		en.Flush()
+		bts = AppendBytesHeader(bts[0:0], sz)
+
+		if !bytes.Equal(buf.Bytes(), bts) {
+			t.Errorf("for size %d, encoder wrote %q and append wrote %q", sz, buf.Bytes(), bts)
+		}
+	}
+}
+
+func BenchmarkAppendBytesHeader(b *testing.B) {
+	buf := make([]byte, 0, 9)
+	N := b.N / 4
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < N; i++ {
+		AppendBytesHeader(buf[:0], 0)
+		AppendBytesHeader(buf[:0], uint32(tint8))
+		AppendBytesHeader(buf[:0], tuint16)
+		AppendBytesHeader(buf[:0], tuint32)
+	}
+}
+
 func TestAppendNil(t *testing.T) {
 	var bts []byte
 	bts = AppendNil(bts[0:0])