perf(core): Hand-rolled UnmarshalBinary for AggregatedBloomFilter

core/aggregated_bloom_filter.go

@@ -72,8 +72,15 @@
 
 const (
 	NumBlocksPerFilter uint64 = 8192
+	// MaxBlockOffsetPerFilter is the last block a filter covers, relative to
+	// its fromBlock: toBlock == fromBlock + MaxBlockOffsetPerFilter.
+	MaxBlockOffsetPerFilter = NumBlocksPerFilter - 1
 )
 
+// wordsPerFilterRow is the uint64 word count bitset writes for one
+// NumBlocksPerFilter-bit row: (bits + 63) / 64.
+const wordsPerFilterRow = int((NumBlocksPerFilter + 63) / 64)
+
 var (
 	ErrAggregatedBloomFilterBlockOutOfRange error = errors.New("block number is not within range")
 	ErrBloomFilterSizeMismatch              error = errors.New("bloom filter len mismatch")
@@ -92,7 +99,7 @@
 	return AggregatedBloomFilter{
 		bitmap:    bitmap,
 		fromBlock: fromBlock,
-		toBlock:   fromBlock + NumBlocksPerFilter - 1,
+		toBlock:   fromBlock + MaxBlockOffsetPerFilter,
 	}
 }
 
@@ -251,43 +258,86 @@
 	return buf.Bytes(), nil
 }
 
+// UnmarshalBinary decodes MarshalBinary's output in place. Returns an error
+// in case of DB corruption. On error f may be partially written and must be
+// discarded.
 func (f *AggregatedBloomFilter) UnmarshalBinary(data []byte) error {
-	r := bytes.NewReader(data)
+	const (
+		bytesUint32 = 4
+		bytesUint64 = 8
+		// Header layout (see MarshalBinary): fromBlock + toBlock + count.
+		headerSize    = bytesUint64 + bytesUint64 + bytesUint32
+		rowLenSize    = bytesUint32
+		bitsetLenSize = bytesUint64
+		// A canonical row is a bitset length prefix plus wordsPerFilterRow words.
+		blobLenWant = bitsetLenSize + wordsPerFilterRow*bytesUint64
+		rowSize     = rowLenSize + blobLenWant
+	)
+
+	if len(data) < headerSize {
+		return io.ErrUnexpectedEOF
+	}
 
-	if err := binary.Read(r, binary.BigEndian, &f.fromBlock); err != nil {
-		return err
+	count := int(binary.BigEndian.Uint32(data[2*bytesUint64 : headerSize]))
+	// Consumers index the matrix with these fixed constants, so a header that
+	// disagrees would panic or silently corrupt results on later use.
+	if count != EventsBloomLength {
+		return ErrBloomFilterSizeMismatch
 	}
-	if err := binary.Read(r, binary.BigEndian, &f.toBlock); err != nil {
-		return err
+	// Reject data too short for count rows before allocating, so truncated
+	// input fails fast rather than driving an oversized allocation.
+	if count > (len(data)-headerSize)/rowSize {
+		return io.ErrUnexpectedEOF
 	}
 
-	var count uint32
-	if err := binary.Read(r, binary.BigEndian, &count); err != nil {
-		return err
+	f.fromBlock = binary.BigEndian.Uint64(data[0:bytesUint64])
+	f.toBlock = binary.BigEndian.Uint64(data[bytesUint64 : 2*bytesUint64])
+	if f.toBlock != f.fromBlock+MaxBlockOffsetPerFilter {
+		return ErrBloomFilterSizeMismatch
 	}
 
+	backing := make([]uint64, count*wordsPerFilterRow)
 	f.bitmap = make([]bitset.BitSet, count)
+
+	offset := headerSize
 	for i := range count {
-		var length uint32
-		if err := binary.Read(r, binary.BigEndian, &length); err != nil {
-			return err
+		if offset+rowLenSize > len(data) {
+			return io.ErrUnexpectedEOF
 		}
-
-		b := make([]byte, length)
-		if _, err := io.ReadFull(r, b); err != nil {
-			return err
+		blobLen := int(binary.BigEndian.Uint32(data[offset:]))
+		offset += rowLenSize
+		if offset+blobLen > len(data) {
+			return io.ErrUnexpectedEOF
+		}
+		// bitsetLen and blobLen are independent fields, so both are checked.
+		if blobLen != blobLenWant {
+			return ErrBloomFilterSizeMismatch
+		}
+		if bitsetLen := binary.BigEndian.Uint64(data[offset:]); bitsetLen != NumBlocksPerFilter {
+			return ErrBloomFilterSizeMismatch
 		}
 
-		if err := f.bitmap[i].UnmarshalBinary(b); err != nil {
-			return err
+		rowStart := i * wordsPerFilterRow
+		row := backing[rowStart : rowStart+wordsPerFilterRow : rowStart+wordsPerFilterRow]
+		wordsAt := offset + bitsetLenSize
+		for w := range wordsPerFilterRow {
+			row[w] = binary.BigEndian.Uint64(data[wordsAt+w*bytesUint64:])
 		}
+		f.bitmap[i] = *bitset.FromWithLength(uint(NumBlocksPerFilter), row)
+
+		offset += blobLen
+	}
+
+	// Trailing bytes mean framing corruption; a canonical blob is consumed exactly.
+	if offset != len(data) {
+		return io.ErrUnexpectedEOF
 	}
 	return nil
 }
 
 func makeBitset() bitset.BitSet {
 	b := bitset.BitSet{}
-	b.Set(uint(NumBlocksPerFilter - 1))
-	b.Clear(uint(NumBlocksPerFilter - 1))
+	b.Set(uint(MaxBlockOffsetPerFilter))
+	b.Clear(uint(MaxBlockOffsetPerFilter))
 	return b
 }

core/aggregated_bloom_filter_bench_test.go

@@ -0,0 +1,27 @@
+package core_test
+
+import (
+	"testing"
+
+	"github.com/NethermindEth/juno/core"
+	"github.com/bits-and-blooms/bloom/v3"
+	"github.com/stretchr/testify/require"
+)
+
+func BenchmarkAggregatedBloomFilterUnmarshal(b *testing.B) {
+	filter := core.NewAggregatedFilter(0)
+	bf := bloom.New(core.EventsBloomLength, core.EventsBloomHashFuncs)
+	bf.Add([]byte{0x01})
+	require.NoError(b, filter.Insert(bf, 0))
+
+	data, err := filter.MarshalBinary()
+	require.NoError(b, err)
+
+	b.ReportAllocs()
+	for b.Loop() {
+		var decoded core.AggregatedBloomFilter
+		if err := decoded.UnmarshalBinary(data); err != nil {
+			b.Fatal(err)
+		}
+	}
+}

core/aggregated_bloom_filter_test.go

@@ -1,6 +1,8 @@
 package core_test
 
 import (
+	"encoding/binary"
+	"io"
 	"testing"
 
 	"github.com/NethermindEth/juno/core"
@@ -30,7 +32,7 @@ func TestAggregatedBloomFilter_Insert(t *testing.T) {
 			},
 			{
 				description: "last block in range",
-				blockNumber: 200 + core.NumBlocksPerFilter - 1,
+				blockNumber: 200 + core.MaxBlockOffsetPerFilter,
 			},
 		}
 
@@ -148,3 +150,172 @@ func TestAggregatedBloomFilter_Serialise(t *testing.T) {
 	require.NoError(t, filter2.UnmarshalBinary(data))
 	require.Equal(t, filter, *filter2)
 }
+
+func TestAggregatedBloomFilter_UnmarshalBinary_Compat(t *testing.T) {
+	// Build a filter with a couple of keys at known blocks.
+	filter := core.NewAggregatedFilter(0)
+	keyA := []byte{0x33}
+	keyB := []byte{0xAB, 0xCD}
+
+	bloomA := bloom.New(core.EventsBloomLength, core.EventsBloomHashFuncs)
+	bloomA.Add(keyA)
+	require.NoError(t, filter.Insert(bloomA, 0))
+
+	bloomB := bloom.New(core.EventsBloomLength, core.EventsBloomHashFuncs)
+	bloomB.Add(keyB)
+	require.NoError(t, filter.Insert(bloomB, 5))
+
+	data, err := filter.MarshalBinary()
+	require.NoError(t, err)
+
+	// Round-trips into an equal value.
+	var decoded core.AggregatedBloomFilter
+	require.NoError(t, decoded.UnmarshalBinary(data))
+	require.Equal(t, filter, decoded)
+
+	// Matching behavior is preserved after decode.
+	require.True(t, decoded.BlocksForKeys([][]byte{keyA}).Test(0))
+	require.True(t, decoded.BlocksForKeys([][]byte{keyB}).Test(5))
+	require.False(t, decoded.BlocksForKeys([][]byte{keyB}).Test(0))
+
+	// Short/corrupt input returns a non-nil, non-panicking result.
+	require.Error(t, decoded.UnmarshalBinary(data[:10]))
+	require.Error(t, decoded.UnmarshalBinary(nil))
+}
+
+// Round-trips a filter whose range does not start at block 0, so a byte-offset
+// bug in the header parse cannot hide behind fromBlock == 0.
+func TestAggregatedBloomFilter_UnmarshalBinary_NonZeroRange(t *testing.T) {
+	const from = 3 * core.NumBlocksPerFilter // distinctive, non-zero start
+	filter := core.NewAggregatedFilter(from)
+	key := []byte{0x5e}
+	b := bloom.New(core.EventsBloomLength, core.EventsBloomHashFuncs)
+	b.Add(key)
+	require.NoError(t, filter.Insert(b, from+7))
+
+	data, err := filter.MarshalBinary()
+	require.NoError(t, err)
+
+	var decoded core.AggregatedBloomFilter
+	require.NoError(t, decoded.UnmarshalBinary(data))
+	require.Equal(t, filter, decoded)
+	require.Equal(t, from, decoded.FromBlock())
+	require.Equal(t, from+core.MaxBlockOffsetPerFilter, decoded.ToBlock())
+	require.True(t, decoded.BlocksForKeys([][]byte{key}).Test(7))
+}
+
+// Pins the on-disk header layout (fromBlock:uint64, toBlock:uint64,
+// count:uint32, all big-endian) that UnmarshalBinary relies on, guarding
+// against a future change to MarshalBinary silently breaking the wire format.
+func TestAggregatedBloomFilter_MarshalBinary_HeaderLayout(t *testing.T) {
+	const from = 0x0102030405060708
+	filter := core.NewAggregatedFilter(from)
+
+	data, err := filter.MarshalBinary()
+	require.NoError(t, err)
+	require.GreaterOrEqual(t, len(data), 20)
+
+	require.Equal(t, uint64(from), binary.BigEndian.Uint64(data[0:8]))
+	require.Equal(t, from+core.MaxBlockOffsetPerFilter, binary.BigEndian.Uint64(data[8:16]))
+	require.Equal(t, uint32(core.EventsBloomLength), binary.BigEndian.Uint32(data[16:20]))
+}
+
+// Exercises the decode validation and bounds-checking branches: a row with a
+// wrong bit-length header is rejected, and truncation mid-row does not panic.
+func TestAggregatedBloomFilter_UnmarshalBinary_Corrupt(t *testing.T) {
+	filter := core.NewAggregatedFilter(0)
+	b := bloom.New(core.EventsBloomLength, core.EventsBloomHashFuncs)
+	b.Add([]byte{0x01})
+	require.NoError(t, filter.Insert(b, 0))
+
+	data, err := filter.MarshalBinary()
+	require.NoError(t, err)
+
+	t.Run("wrong row bit-length", func(t *testing.T) {
+		corrupt := make([]byte, len(data))
+		copy(corrupt, data)
+		// Row 0 blob starts after the 20-byte header + 4-byte blob length;
+		// its first 8 bytes are the bitset bit-length header.
+		binary.BigEndian.PutUint64(corrupt[24:], core.NumBlocksPerFilter+1)
+
+		var decoded core.AggregatedBloomFilter
+		require.ErrorIs(t, decoded.UnmarshalBinary(corrupt), core.ErrBloomFilterSizeMismatch)
+	})
+
+	t.Run("truncated mid-row", func(t *testing.T) {
+		var decoded core.AggregatedBloomFilter
+		require.ErrorIs(t, decoded.UnmarshalBinary(data[:30]), io.ErrUnexpectedEOF)
+	})
+
+	t.Run("wrong row count", func(t *testing.T) {
+		corrupt := make([]byte, len(data))
+		copy(corrupt, data)
+		binary.BigEndian.PutUint32(corrupt[16:20], core.EventsBloomLength-1)
+
+		var decoded core.AggregatedBloomFilter
+		require.ErrorIs(t, decoded.UnmarshalBinary(corrupt), core.ErrBloomFilterSizeMismatch)
+	})
+
+	t.Run("toBlock not matching range", func(t *testing.T) {
+		corrupt := make([]byte, len(data))
+		copy(corrupt, data)
+		binary.BigEndian.PutUint64(corrupt[8:16], 1<<40)
+
+		var decoded core.AggregatedBloomFilter
+		require.ErrorIs(t, decoded.UnmarshalBinary(corrupt), core.ErrBloomFilterSizeMismatch)
+	})
+
+	t.Run("trailing bytes", func(t *testing.T) {
+		withJunk := append([]byte{}, data...)
+		withJunk = append(withJunk, 0x00, 0x01, 0x02, 0x03)
+		var decoded core.AggregatedBloomFilter
+		require.ErrorIs(t, decoded.UnmarshalBinary(withJunk), io.ErrUnexpectedEOF)
+	})
+}
+
+// UnmarshalBinary decodes untrusted DB bytes and must never panic on arbitrary
+// input, only return an error.
+func FuzzAggregatedBloomFilterUnmarshal(f *testing.F) {
+	filter := core.NewAggregatedFilter(0)
+	b := bloom.New(core.EventsBloomLength, core.EventsBloomHashFuncs)
+	b.Add([]byte{0x01})
+	require.NoError(f, filter.Insert(b, 0))
+	valid, err := filter.MarshalBinary()
+	require.NoError(f, err)
+
+	f.Add([]byte(nil))
+	f.Add([]byte{0x00})
+	f.Add(valid)
+	f.Add(valid[:25])
+	// Header claiming a huge row count must not drive an unbounded allocation.
+	hugeCount := make([]byte, 20)
+	binary.BigEndian.PutUint32(hugeCount[16:], 0xFFFFFFFF)
+	f.Add(hugeCount)
+
+	f.Fuzz(func(t *testing.T, data []byte) {
+		var decoded core.AggregatedBloomFilter
+		_ = decoded.UnmarshalBinary(data) // must not panic
+	})
+}
+
+func TestAggregatedBloomFilter_UnmarshalBinary_Allocs(t *testing.T) {
+	// A fully populated filter is the worst case: all EventsBloomLength rows present.
+	filter := core.NewAggregatedFilter(0)
+	b := bloom.New(core.EventsBloomLength, core.EventsBloomHashFuncs)
+	b.Add([]byte{0x01})
+	require.NoError(t, filter.Insert(b, 0))
+
+	data, err := filter.MarshalBinary()
+	require.NoError(t, err)
+
+	var decoded core.AggregatedBloomFilter
+	allocs := testing.AllocsPerRun(20, func() {
+		if err := decoded.UnmarshalBinary(data); err != nil {
+			t.Fatal(err)
+		}
+	})
+	// One backing []uint64 + one []bitset.BitSet header. Allow slack for
+	// interface boxing in AllocsPerRun itself; current code is ~2.
+	require.LessOrEqual(t, allocs, float64(8),
+		"decode should allocate O(1) buffers, got %.0f", allocs)
+}

core/running_event_filter.go

@@ -254,7 +254,7 @@ func (f *RunningEventFilter) onReorg(writer db.KeyValueWriter) error {
 		}
 
 		rangeStartAligned := curBlock - (curBlock % NumBlocksPerFilter)
-		rangeEndAligned := rangeStartAligned + NumBlocksPerFilter - 1
+		rangeEndAligned := rangeStartAligned + MaxBlockOffsetPerFilter
 
 		lastStoredFilter, err := GetAggregatedBloomFilter(
 			f.database,
@@ -364,7 +364,7 @@ func rebuildRunningEventFilter(
 	latest uint64,
 ) (*RunningEventFilter, error) {
 	rangeStartAligned := latest - latest%NumBlocksPerFilter
-	lastStoredFilterRangeEnd := rangeStartAligned + NumBlocksPerFilter - 1
+	lastStoredFilterRangeEnd := rangeStartAligned + MaxBlockOffsetPerFilter
 
 	var continueFrom uint64
 	for {