diff --git a/internal/encoding/bench_test.go b/internal/encoding/bench_test.go new file mode 100644 index 000000000..03a2f320b --- /dev/null +++ b/internal/encoding/bench_test.go @@ -0,0 +1,59 @@ +package encoding + +import ( + "bytes" + "encoding/binary" + "fmt" + "math" + "testing" +) + +var encValues = []int64{ + -1, -100, -1 << 32, + 0, 1, 100, 1 << 32, + -1 << 52, 1 << 52, 17, + 19, 28, 37, 388888888, + -99999999999, 99999999999, + math.MaxInt64, math.MinInt64, +} + +// This tests that the results from directly invoking binary.PutVarint match +// exactly those that we get from invoking EncodeVarint and its internals. +func TestEncodeVarintParity(t *testing.T) { + buf := new(bytes.Buffer) + var board [binary.MaxVarintLen64]byte + + for _, val := range encValues { + val := val + name := fmt.Sprintf("%d", val) + + buf.Reset() + t.Run(name, func(t *testing.T) { + if err := EncodeVarint(buf, val); err != nil { + t.Fatal(err) + } + + n := binary.PutVarint(board[:], val) + got := buf.Bytes() + want := board[:n] + if !bytes.Equal(got, want) { + t.Fatalf("Result mismatch\n\tGot: %d\n\tWant: %d", got, want) + } + }) + } +} + +func BenchmarkEncodeVarint(b *testing.B) { + buf := new(bytes.Buffer) + b.ReportAllocs() + b.ResetTimer() + + for i := 0; i < b.N; i++ { + for _, val := range encValues { + if err := EncodeVarint(buf, val); err != nil { + b.Fatal(err) + } + buf.Reset() + } + } +} diff --git a/internal/encoding/encoding.go b/internal/encoding/encoding.go index 78a4f9899..6e390c36c 100644 --- a/internal/encoding/encoding.go +++ b/internal/encoding/encoding.go @@ -155,6 +155,10 @@ func EncodeUvarintSize(u uint64) int { // EncodeVarint writes a varint-encoded integer to an io.Writer. func EncodeVarint(w io.Writer, i int64) error { + if bw, ok := w.(io.ByteWriter); ok { + return fVarintEncode(bw, i) + } + // Use a pool here to reduce allocations. // // Though this allocates just 10 bytes on the stack, doing allocation for every calls @@ -173,6 +177,26 @@ func EncodeVarint(w io.Writer, i int64) error { return err } +func fVarintEncode(bw io.ByteWriter, x int64) error { + // Firstly convert it into a uvarint + ux := uint64(x) << 1 + if x < 0 { + ux = ^ux + } + for ux >= 0x80 { // While there are 7 or more bits in the value, keep going + // Convert it into a byte then toggle the + // 7th bit to indicate that more bytes coming. + // byte(x & 0x7f) is redundant but useful for illustrative + // purposes when translating to other languages + if err := bw.WriteByte(byte(ux&0x7f) | 0x80); err != nil { + return err + } + ux >>= 7 + } + + return bw.WriteByte(byte(ux & 0x7f)) +} + // EncodeVarintSize returns the byte size of the given integer as a varint. func EncodeVarintSize(i int64) int { ux := uint64(i) << 1