From 277f241004def543f547bdb65aa4688216994bbc Mon Sep 17 00:00:00 2001
From: Lee Bousfield <ljbousfield@gmail.com>
Date: Tue, 23 Jan 2024 23:40:07 -0700
Subject: [PATCH] Pack more bits into blobs

---
 util/blobs/blobs.go      | 74 ++++++++++++++++++++++++++++++++--------
 util/blobs/blobs_test.go | 52 ++++++++++++++++++++++++++++
 2 files changed, 112 insertions(+), 14 deletions(-)
 create mode 100644 util/blobs/blobs_test.go

diff --git a/util/blobs/blobs.go b/util/blobs/blobs.go
index 55df57f9d1..2852f2b29f 100644
--- a/util/blobs/blobs.go
+++ b/util/blobs/blobs.go
@@ -14,6 +14,44 @@ import (
 	"github.com/ethereum/go-ethereum/rlp"
 )
 
+func fillBlobBytes(blob []byte, data []byte) []byte {
+	for fieldElement := 0; fieldElement < params.BlobTxFieldElementsPerBlob; fieldElement++ {
+		startIdx := fieldElement*32 + 1
+		copy(blob[startIdx:startIdx+31], data)
+		if len(data) <= 31 {
+			return nil
+		}
+		data = data[31:]
+	}
+	return data
+}
+
+// The number of bits in a BLS scalar that aren't part of a whole byte.
+const spareBlobBits = 6 // = math.floor(math.log2(BLS_MODULUS)) % 8
+
+func fillBlobBits(blob []byte, data []byte) ([]byte, error) {
+	var acc uint16
+	accBits := 0
+	for fieldElement := 0; fieldElement < params.BlobTxFieldElementsPerBlob; fieldElement++ {
+		if accBits < spareBlobBits && len(data) > 0 {
+			acc |= uint16(data[0]) << accBits
+			accBits += 8
+			data = data[1:]
+		}
+		blob[fieldElement*32] = uint8(acc & ((1 << spareBlobBits) - 1))
+		accBits -= spareBlobBits
+		if accBits < 0 {
+			// We're out of data
+			break
+		}
+		acc >>= spareBlobBits
+	}
+	if accBits > 0 {
+		return nil, fmt.Errorf("somehow ended up with %v spare accBits", accBits)
+	}
+	return data, nil
+}
+
 // EncodeBlobs takes in raw bytes data to convert into blobs used for KZG commitment EIP-4844
 // transactions on Ethereum.
 func EncodeBlobs(data []byte) ([]kzg4844.Blob, error) {
@@ -21,21 +59,15 @@ func EncodeBlobs(data []byte) ([]kzg4844.Blob, error) {
 	if err != nil {
 		return nil, err
 	}
-	blobs := []kzg4844.Blob{{}}
-	blobIndex := 0
-	fieldIndex := -1
-	for i := 0; i < len(data); i += 31 {
-		fieldIndex++
-		if fieldIndex == params.BlobTxFieldElementsPerBlob {
-			blobs = append(blobs, kzg4844.Blob{})
-			blobIndex++
-			fieldIndex = 0
-		}
-		max := i + 31
-		if max > len(data) {
-			max = len(data)
+	var blobs []kzg4844.Blob
+	for len(data) > 0 {
+		var b kzg4844.Blob
+		data = fillBlobBytes(b[:], data)
+		data, err = fillBlobBits(b[:], data)
+		if err != nil {
+			return nil, err
 		}
-		copy(blobs[blobIndex][fieldIndex*32+1:], data[i:max])
+		blobs = append(blobs, b)
 	}
 	return blobs, nil
 }
@@ -47,6 +79,20 @@ func DecodeBlobs(blobs []kzg4844.Blob) ([]byte, error) {
 		for fieldIndex := 0; fieldIndex < params.BlobTxFieldElementsPerBlob; fieldIndex++ {
 			rlpData = append(rlpData, blob[fieldIndex*32+1:(fieldIndex+1)*32]...)
 		}
+		var acc uint16
+		accBits := 0
+		for fieldIndex := 0; fieldIndex < params.BlobTxFieldElementsPerBlob; fieldIndex++ {
+			acc |= uint16(blob[fieldIndex*32]) << accBits
+			accBits += spareBlobBits
+			if accBits >= 8 {
+				rlpData = append(rlpData, uint8(acc))
+				acc >>= 8
+				accBits -= 8
+			}
+		}
+		if accBits != 0 {
+			return nil, fmt.Errorf("somehow ended up with %v spare accBits", accBits)
+		}
 	}
 	var outputData []byte
 	err := rlp.Decode(bytes.NewReader(rlpData), &outputData)
diff --git a/util/blobs/blobs_test.go b/util/blobs/blobs_test.go
new file mode 100644
index 0000000000..753b50a489
--- /dev/null
+++ b/util/blobs/blobs_test.go
@@ -0,0 +1,52 @@
+// Copyright 2024, Offchain Labs, Inc.
+// For license information, see https://github.com/nitro/blob/master/LICENSE
+
+package blobs
+
+import (
+	"bytes"
+	"math/big"
+	"math/rand"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/params"
+)
+
+const bytesEncodedPerBlob = 254 * 4096 / 8
+
+var blsModulus, _ = new(big.Int).SetString("52435875175126190479447740508185965837690552500527637822603658699938581184513", 10)
+
+func TestBlobEncoding(t *testing.T) {
+	r := rand.New(rand.NewSource(1))
+outer:
+	for i := 0; i < 40; i++ {
+		data := make([]byte, r.Int()%bytesEncodedPerBlob*3)
+		_, err := r.Read(data)
+		if err != nil {
+			t.Fatalf("failed to generate random bytes: %v", err)
+		}
+		enc, err := EncodeBlobs(data)
+		if err != nil {
+			t.Errorf("failed to encode blobs for length %v: %v", len(data), err)
+			continue
+		}
+		for _, b := range enc {
+			for fieldElement := 0; fieldElement < params.BlobTxFieldElementsPerBlob; fieldElement++ {
+				bigInt := new(big.Int).SetBytes(b[fieldElement*32 : (fieldElement+1)*32])
+				if bigInt.Cmp(blsModulus) >= 0 {
+					t.Errorf("for length %v blob %v has field element %v value %v >= modulus %v", len(data), b, fieldElement, bigInt, blsModulus)
+					continue outer
+				}
+			}
+		}
+		dec, err := DecodeBlobs(enc)
+		if err != nil {
+			t.Errorf("failed to decode blobs for length %v: %v", len(data), err)
+			continue
+		}
+		if !bytes.Equal(data, dec) {
+			t.Errorf("got different decoding for length %v", len(data))
+			continue
+		}
+	}
+}