Make KeyBytes allocates less memory (#418)

If the segments is less than the layouts, the we don't need to allocate
full KeyFormat.length, just allocate enough memory to store the
segments.

name                                                old time/op    new time/op    delta
KeyFormat_KeyBytesOneSegment-8                        19.2ns ± 0%    16.4ns ± 1%  -14.62%  (p=0.000 n=9+10)
KeyFormat_KeyBytesThreeSegment-8                      25.6ns ± 0%    25.9ns ± 1%   +1.37%  (p=0.000 n=10+10)
KeyFormat_KeyBytesOneSegmentWithVariousLayouts-8      22.2ns ± 1%    16.4ns ± 0%  -26.14%  (p=0.000 n=10+8)
KeyFormat_KeyBytesThreeSegmentWithVariousLayouts-8    28.8ns ± 3%    28.7ns ± 0%     ~     (p=0.588 n=10+9)

name                                                old alloc/op   new alloc/op   delta
KeyFormat_KeyBytesOneSegment-8                         32.0B ± 0%     16.0B ± 0%  -50.00%  (p=0.000 n=10+10)
KeyFormat_KeyBytesThreeSegment-8                       32.0B ± 0%     32.0B ± 0%     ~     (all equal)
KeyFormat_KeyBytesOneSegmentWithVariousLayouts-8       64.0B ± 0%     16.0B ± 0%  -75.00%  (p=0.000 n=10+10)
KeyFormat_KeyBytesThreeSegmentWithVariousLayouts-8     64.0B ± 0%     64.0B ± 0%     ~     (all equal)

name                                                old allocs/op  new allocs/op  delta
KeyFormat_KeyBytesOneSegment-8                          1.00 ± 0%      1.00 ± 0%     ~     (all equal)
KeyFormat_KeyBytesThreeSegment-8                        1.00 ± 0%      1.00 ± 0%     ~     (all equal)
KeyFormat_KeyBytesOneSegmentWithVariousLayouts-8        1.00 ± 0%      1.00 ± 0%     ~     (all equal)
KeyFormat_KeyBytesThreeSegmentWithVariousLayouts-8      1.00 ± 0%      1.00 ± 0%     ~     (all equal)

Fixes #417

key_format.go

@@ -42,7 +42,18 @@ func NewKeyFormat(prefix byte, layout ...int) *KeyFormat {
 
 // Format the byte segments into the key format - will panic if the segment lengths do not match the layout.
 func (kf *KeyFormat) KeyBytes(segments ...[]byte) []byte {
-	key := make([]byte, kf.length)
+	keyLen := kf.length
+	// In case segments length is less than layouts length,
+	// we don't have to allocate the whole kf.length, just
+	// enough space to store the segments.
+	if len(segments) < len(kf.layout) {
+		keyLen = 1
+		for i := range segments {
+			keyLen += kf.layout[i]
+		}
+	}
+
+	key := make([]byte, keyLen)
 	key[0] = kf.prefix
 	n := 1
 	for i, s := range segments {

key_format_test.go

@@ -68,3 +68,35 @@ func TestOverflow(t *testing.T) {
 	assert.Equal(t, a, *ao)
 	assert.Equal(t, int64(b), *bo)
 }
+
+func benchmarkKeyFormatBytes(b *testing.B, kf *KeyFormat, segments ...[]byte) {
+	for i := 0; i < b.N; i++ {
+		kf.KeyBytes(segments...)
+	}
+}
+
+func BenchmarkKeyFormat_KeyBytesOneSegment(b *testing.B) {
+	benchmarkKeyFormatBytes(b, NewKeyFormat('e', 8, 8, 8), nil)
+}
+
+func BenchmarkKeyFormat_KeyBytesThreeSegment(b *testing.B) {
+	segments := [][]byte{
+		{1, 2, 3, 4, 5, 6, 7, 8},
+		{1, 2, 3, 4, 5, 6, 7, 8},
+		{1, 1, 2, 2, 3, 3},
+	}
+	benchmarkKeyFormatBytes(b, NewKeyFormat('e', 8, 8, 8), segments...)
+}
+
+func BenchmarkKeyFormat_KeyBytesOneSegmentWithVariousLayouts(b *testing.B) {
+	benchmarkKeyFormatBytes(b, NewKeyFormat('e', 8, 16, 32), nil)
+}
+
+func BenchmarkKeyFormat_KeyBytesThreeSegmentWithVariousLayouts(b *testing.B) {
+	segments := [][]byte{
+		{1, 2, 3, 4, 5, 6, 7, 8},
+		{1, 2, 3, 4, 5, 6, 7, 8},
+		{1, 1, 2, 2, 3, 3},
+	}
+	benchmarkKeyFormatBytes(b, NewKeyFormat('e', 8, 16, 32), segments...)
+}