From 3b2b42f68104ad749eeece3ba4bd4e7133e2ce4d Mon Sep 17 00:00:00 2001
From: Julien Pivotto <291750+roidelapluie@users.noreply.github.com>
Date: Fri, 20 Mar 2026 13:42:03 +0100
Subject: [PATCH] tsdb/chunkenc: add writeBits benchmarks, clarify comments,
 and simplify encodeJoint

Signed-off-by: Julien Pivotto <291750+roidelapluie@users.noreply.github.com>
---
 tsdb/chunkenc/bstream.go      |  9 ++++----
 tsdb/chunkenc/bstream_test.go | 39 +++++++++++++++++++++++++++++++++++
 tsdb/chunkenc/xor2.go         |  9 ++++----
 3 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/tsdb/chunkenc/bstream.go b/tsdb/chunkenc/bstream.go
index 2ac92b69c8..4fd37a140f 100644
--- a/tsdb/chunkenc/bstream.go
+++ b/tsdb/chunkenc/bstream.go
@@ -101,6 +101,7 @@ func (b *bstream) writeByte(byt byte) {
 
 // writeBits writes the nbits right-most bits of u to the stream
 // in left-to-right order.
+// TODO: Once XOR2 stabilizes, replace writeBits with the writeBitsFast implementation and remove writeBitsFast.
 func (b *bstream) writeBits(u uint64, nbits int) {
 	u <<= 64 - uint(nbits)
 	for nbits >= 8 {
@@ -368,9 +369,8 @@ func (b *bstreamReader) readXOR2Control() (uint8, error) {
 }
 
 // readUvarint decodes a varint-encoded uint64 using direct method calls,
-// avoiding the io.ByteReader interface dispatch used by binary.ReadUvarint.
-// This prevents interior pointer references on goroutine stacks that the GC
-// must trace via findObject, reducing GC overhead.
+// avoiding the io.ByteReader interface dispatch used by binary.ReadUvarint,
+// which causes the receiver to escape to the heap.
 func (b *bstreamReader) readUvarint() (uint64, error) {
 	var x uint64
 	var s uint
@@ -389,7 +389,8 @@ func (b *bstreamReader) readUvarint() (uint64, error) {
 }
 
 // readVarint decodes a varint-encoded int64 using direct method calls,
-// avoiding the io.ByteReader interface dispatch used by binary.ReadVarint.
+// avoiding the io.ByteReader interface dispatch used by binary.ReadVarint,
+// which causes the receiver to escape to the heap.
 func (b *bstreamReader) readVarint() (int64, error) {
 	ux, err := b.readUvarint()
 	x := int64(ux >> 1)
diff --git a/tsdb/chunkenc/bstream_test.go b/tsdb/chunkenc/bstream_test.go
index 3098be5945..0b6a0e9b35 100644
--- a/tsdb/chunkenc/bstream_test.go
+++ b/tsdb/chunkenc/bstream_test.go
@@ -14,6 +14,7 @@
 package chunkenc
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -32,6 +33,44 @@ func TestBstream_Reset(t *testing.T) {
 	}, bs)
 }
 
+// BenchmarkWriteBits benchmarks writeBits for various bit widths.
+func BenchmarkWriteBits(b *testing.B) {
+	sizes := []int{1, 8, 17, 32, 52, 64}
+	for _, nbits := range sizes {
+		b.Run(fmt.Sprintf("nbits=%d", nbits), func(b *testing.B) {
+			b.ReportAllocs()
+			var bs bstream
+			bs.stream = make([]byte, 0, 1024)
+			for range b.N {
+				bs.stream = bs.stream[:0]
+				bs.count = 0
+				for j := range 100 {
+					bs.writeBits(uint64(j), nbits)
+				}
+			}
+		})
+	}
+}
+
+// BenchmarkWriteBitsFast benchmarks writeBitsFast for various bit widths.
+func BenchmarkWriteBitsFast(b *testing.B) {
+	sizes := []int{1, 8, 17, 32, 52, 64}
+	for _, nbits := range sizes {
+		b.Run(fmt.Sprintf("nbits=%d", nbits), func(b *testing.B) {
+			b.ReportAllocs()
+			var bs bstream
+			bs.stream = make([]byte, 0, 1024)
+			for range b.N {
+				bs.stream = bs.stream[:0]
+				bs.count = 0
+				for j := range 100 {
+					bs.writeBitsFast(uint64(j), nbits)
+				}
+			}
+		})
+	}
+}
+
 func TestBstreamReader(t *testing.T) {
 	// Write to the bit stream.
 	w := bstream{}
diff --git a/tsdb/chunkenc/xor2.go b/tsdb/chunkenc/xor2.go
index 85db376ffb..defe1e8102 100644
--- a/tsdb/chunkenc/xor2.go
+++ b/tsdb/chunkenc/xor2.go
@@ -304,12 +304,13 @@ func (a *xor2Appender) encodeJoint(dod int64, v float64) {
 			a.b.writeBitsFast(0b11111, 5)
 			return
 		}
-		if vbits := math.Float64bits(v) ^ math.Float64bits(a.v); vbits == 0 {
+		vbits := math.Float64bits(v) ^ math.Float64bits(a.v)
+		if vbits == 0 {
 			a.b.writeBit(zero)
-		} else {
-			a.b.writeBitsFast(0b10, 2)
-			a.writeVDeltaKnownNonZero(vbits)
+			return
 		}
+		a.b.writeBitsFast(0b10, 2)
+		a.writeVDeltaKnownNonZero(vbits)
 		return
 	}