From aafe4c38ace8ac66ed8c83253299aa60703334cf Mon Sep 17 00:00:00 2001
From: Wael Nasreddine <wael.nasreddine@gmail.com>
Date: Fri, 16 Jan 2026 19:27:34 -0800
Subject: [PATCH] fix: optimize golomb encoding/decoding and improve error
 handling

This commit addresses performance concerns raised in PR 564.

Changes:
- Optimizes `EncodeBig` and `DecodeBig` in `pkg/golomb` to use `uint64` arithmetic for quotient operations when possible, avoiding expensive `big.Int` allocations and operations for common cases.
- Updates `pkg/golomb/golomb_test.go` to properly handle errors from `Flush()`.

This significantly reduces the overhead of Golomb coding for standard integer sizes while maintaining correctness for arbitrarily large numbers.
---
 pkg/golomb/golomb.go      | 51 ++++++++++++++++++++++++++++++---------
 pkg/golomb/golomb_test.go |  2 +-
 2 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/pkg/golomb/golomb.go b/pkg/golomb/golomb.go
index a7e7d171..d0821270 100644
--- a/pkg/golomb/golomb.go
+++ b/pkg/golomb/golomb.go
@@ -284,18 +284,29 @@ func (ge *Encoder) EncodeBig(d *big.Int) error {
 	// Go's big.Int doesn't support "iterate up to value" easily without loop/cmp.
 	// But we can check BitLen.
 
-	zero := big.NewInt(0)
-	one := big.NewInt(1)
-
 	// Write q ones
-	// We decrement q until 0
-	currQ := new(big.Int).Set(q)
-	for currQ.Cmp(zero) > 0 {
-		if err := ge.BitWriter.WriteBit(true); err != nil {
-			return err
+	if q.IsUint64() {
+		qVal := q.Uint64()
+		for i := uint64(0); i < qVal; i++ {
+			if err := ge.BitWriter.WriteBit(true); err != nil {
+				return err
+			}
+		}
+	} else {
+		// This path is for q > 2^64, which will produce an enormous output
+		// and be extremely slow. It's kept for correctness with very large numbers,
+		// but in practice, `k` should be chosen to keep `q` small.
+		zero := big.NewInt(0)
+		one := big.NewInt(1)
+
+		currQ := new(big.Int).Set(q)
+		for currQ.Cmp(zero) > 0 {
+			if err := ge.BitWriter.WriteBit(true); err != nil {
+				return err
+			}
+
+			currQ.Sub(currQ, one)
 		}
-
-		currQ.Sub(currQ, one)
 	}
 
 	// Write zero delimiter
@@ -312,7 +323,10 @@ func (ge *Encoder) EncodeBig(d *big.Int) error {
 func (gd *Decoder) DecodeBig() (*big.Int, error) {
 	// Decode unary q
 	q := new(big.Int)
-	one := big.NewInt(1)
+
+	var qUint64 uint64
+
+	useBigInt := false
 
 	for {
 		bit, err := gd.br.ReadBit()
@@ -324,7 +338,20 @@ func (gd *Decoder) DecodeBig() (*big.Int, error) {
 			break
 		}
 
-		q.Add(q, one)
+		if useBigInt {
+			q.Add(q, big.NewInt(1)) // This is slow, but q is already huge.
+		} else if qUint64 < ^uint64(0) {
+			qUint64++
+		} else {
+			useBigInt = true
+
+			q.SetUint64(qUint64)
+			q.Add(q, big.NewInt(1))
+		}
+	}
+
+	if !useBigInt {
+		q.SetUint64(qUint64)
 	}
 
 	// Decode binary r: k bits
diff --git a/pkg/golomb/golomb_test.go b/pkg/golomb/golomb_test.go
index 23332dc8..b44379cd 100644
--- a/pkg/golomb/golomb_test.go
+++ b/pkg/golomb/golomb_test.go
@@ -92,7 +92,7 @@ func TestGolombExample(t *testing.T) {
 	err = enc.Encode(1000)
 	require.NoError(t, err)
 
-	_ = enc.Flush()
+	require.NoError(t, enc.Flush())
 
 	// 1110 1110 1000 0000 (padded) -> EE 80
 	decodedBytes := buf.Bytes()