diff --git a/z/buffer.go b/z/buffer.go index edd2f852..3fc8e84a 100644 --- a/z/buffer.go +++ b/z/buffer.go @@ -28,6 +28,8 @@ import ( "github.com/pkg/errors" ) +const padding = 8 + // Buffer is equivalent of bytes.Buffer without the ability to read. It is NOT thread-safe. // // In UseCalloc mode, z.Calloc is used to allocate memory, which depending upon how the code is @@ -113,7 +115,7 @@ func NewBufferWith(sz, maxSz int, bufType BufferType) (*Buffer, error) { } b := &Buffer{ - offset: 1, + offset: padding, // Use 8 bytes of padding so that the elements are aligned. curSz: sz, maxSz: maxSz, bufType: UseCalloc, // by default. @@ -135,17 +137,24 @@ func NewBufferWith(sz, maxSz int, bufType BufferType) (*Buffer, error) { } func (b *Buffer) IsEmpty() bool { - return b.offset == 1 + return b.offset == b.StartOffset() } -// Len would return the number of bytes written to the buffer so far. -func (b *Buffer) Len() int { +// LenWithPadding would return the number of bytes written to the buffer so far +// plus the padding at the start of the buffer. +func (b *Buffer) LenWithPadding() int { return b.offset } +// LenNoPadding would return the number of bytes written to the buffer so far +// (without the padding). +func (b *Buffer) LenNoPadding() int { + return b.offset - padding +} + // Bytes would return all the written bytes as a slice. func (b *Buffer) Bytes() []byte { - return b.buf[1:b.offset] + return b.buf[padding:b.offset] } func (b *Buffer) AutoMmapAfter(size int) { @@ -232,13 +241,15 @@ func (b *Buffer) SliceAllocate(sz int) []byte { return b.Allocate(sz) } +func (b *Buffer) StartOffset() int { return padding } + func (b *Buffer) WriteSlice(slice []byte) { dst := b.SliceAllocate(len(slice)) copy(dst, slice) } func (b *Buffer) SliceIterate(f func(slice []byte) error) error { - slice, next := []byte{}, 1 + slice, next := []byte{}, b.StartOffset() for next != 0 { slice, next = b.Slice(next) if err := f(slice); err != nil { @@ -357,7 +368,7 @@ func (s *sortHelper) sort(lo, hi int) []byte { // SortSlice is like SortSliceBetween but sorting over the entire buffer. func (b *Buffer) SortSlice(less func(left, right []byte) bool) { - b.SortSliceBetween(1, b.offset, less) + b.SortSliceBetween(b.StartOffset(), b.offset, less) } func (b *Buffer) SortSliceBetween(start, end int, less LessFunc) { if start >= end { @@ -422,7 +433,7 @@ func (b *Buffer) Slice(offset int) ([]byte, int) { // SliceOffsets is an expensive function. Use sparingly. func (b *Buffer) SliceOffsets() []int { - next := 1 + next := b.StartOffset() var offsets []int for next != 0 { offsets = append(offsets, next) @@ -448,7 +459,7 @@ func (b *Buffer) Write(p []byte) (n int, err error) { // Reset would reset the buffer to be reused. func (b *Buffer) Reset() { - b.offset = 1 + b.offset = b.StartOffset() } // Release would free up the memory allocated by the buffer. Once the usage of buffer is done, it is diff --git a/z/buffer_test.go b/z/buffer_test.go index 178ea826..936ad3bf 100644 --- a/z/buffer_test.go +++ b/z/buffer_test.go @@ -104,7 +104,7 @@ func TestBufferAutoMmap(t *testing.T) { b := buf.SliceAllocate(len(wb)) copy(b, wb[:]) } - t.Logf("Buffer size: %d\n", buf.Len()) + t.Logf("Buffer size: %d\n", buf.LenWithPadding()) buf.SortSlice(func(l, r []byte) bool { return bytes.Compare(l, r) < 0 @@ -113,18 +113,18 @@ func TestBufferAutoMmap(t *testing.T) { var count int var last []byte - slice, next := []byte{}, 1 - for next != 0 { - slice, next = buf.Slice(next) + buf.SliceIterate(func(slice []byte) error { require.True(t, bytes.Compare(slice, last) >= 0) last = append(last[:0], slice...) count++ - } + return nil + }) require.Equal(t, N, count) } func TestBufferSimpleSort(t *testing.T) { buf := NewBuffer(1 << 20) + defer buf.Release() for i := 0; i < 25600; i++ { b := buf.SliceAllocate(4) binary.BigEndian.PutUint32(b, uint32(rand.Int31n(256000))) @@ -136,9 +136,7 @@ func TestBufferSimpleSort(t *testing.T) { }) var last uint32 var i int - slice, next := []byte{}, 1 - for next != 0 { - slice, next = buf.Slice(next) + buf.SliceIterate(func(slice []byte) error { num := binary.BigEndian.Uint32(slice) if num < last { fmt.Printf("num: %d idx: %d last: %d\n", num, i, last) @@ -147,7 +145,8 @@ func TestBufferSimpleSort(t *testing.T) { require.GreaterOrEqual(t, num, last) last = num // fmt.Printf("Got number: %d\n", num) - } + return nil + }) } func TestBufferSlice(t *testing.T) { @@ -175,19 +174,18 @@ func TestBufferSlice(t *testing.T) { } compare := func() { - next, i := 1, 0 - for next != 0 { + i := 0 + buf.SliceIterate(func(slice []byte) error { // All the slices returned by the buffer should be equal to what we // inserted earlier. - var slice []byte - slice, next = buf.Slice(next) if !bytes.Equal(exp[i], slice) { fmt.Printf("exp: %s got: %s\n", hex.Dump(exp[i]), hex.Dump(slice)) t.Fail() } require.Equal(t, exp[i], slice) i++ - } + return nil + }) require.Equal(t, len(exp), i) } compare() // same order as inserted. @@ -223,8 +221,8 @@ func TestBufferSort(t *testing.T) { } test := func(start, end int) { - start = 1 + 12*start - end = 1 + 12*end + start = padding + 12*start + end = padding + 12*end buf.SortSliceBetween(start, end, func(ls, rs []byte) bool { lhs := binary.BigEndian.Uint64(ls) rhs := binary.BigEndian.Uint64(rs)