diff --git a/triedb/pathdb/history_trienode_utils.go b/triedb/pathdb/history_trienode_utils.go index 0513343404..241b8a7d3c 100644 --- a/triedb/pathdb/history_trienode_utils.go +++ b/triedb/pathdb/history_trienode_utils.go @@ -19,6 +19,7 @@ package pathdb import ( "encoding/binary" "fmt" + "math/bits" "slices" ) @@ -81,3 +82,25 @@ func isBitSet(b []byte, index int) bool { func setBit(b []byte, index int) { b[index/8] |= 1 << (7 - index%8) } + +// bitPosTwoBytes returns the positions of set bits in a 2-byte bitmap. +// +// The bitmap is interpreted as a big-endian uint16. Bit positions are +// numbered from 0 to 15, where position 0 corresponds to the most +// significant bit of b[0], and position 15 corresponds to the least +// significant bit of b[1]. +func bitPosTwoBytes(b []byte) []int { + if len(b) != 2 { + panic("expect 2 bytes") + } + var ( + pos []int + mask = binary.BigEndian.Uint16(b) + ) + for mask != 0 { + p := bits.LeadingZeros16(mask) + pos = append(pos, p) + mask &^= 1 << (15 - p) + } + return pos +} diff --git a/triedb/pathdb/history_trienode_utils_test.go b/triedb/pathdb/history_trienode_utils_test.go index 17eabb2a98..c3bd0d5b1f 100644 --- a/triedb/pathdb/history_trienode_utils_test.go +++ b/triedb/pathdb/history_trienode_utils_test.go @@ -18,6 +18,7 @@ package pathdb import ( "bytes" + "reflect" "testing" ) @@ -79,3 +80,47 @@ func TestBitmapSet(t *testing.T) { } } } + +func TestBitPositions(t *testing.T) { + suites := []struct { + input []byte + expect []int + }{ + { + []byte{0b10000000, 0x0}, []int{0}, + }, + { + []byte{0b01000000, 0x0}, []int{1}, + }, + { + []byte{0b00000001, 0x0}, []int{7}, + }, + { + []byte{0b00000000, 0b10000000}, []int{8}, + }, + { + []byte{0b00000000, 0b00000001}, []int{15}, + }, + { + []byte{0b10000000, 0b00000001}, []int{0, 15}, + }, + { + []byte{0b10000001, 0b00000001}, []int{0, 7, 15}, + }, + { + []byte{0b10000001, 0b10000001}, []int{0, 7, 8, 15}, + }, + { + []byte{0b11111111, 0b11111111}, []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}, + }, + { + []byte{0x0, 0x0}, nil, + }, + } + for _, tc := range suites { + got := bitPosTwoBytes(tc.input) + if !reflect.DeepEqual(got, tc.expect) { + t.Fatalf("Unexpected position set, want: %v, got: %v", tc.expect, got) + } + } +} diff --git a/triedb/pathdb/nodes.go b/triedb/pathdb/nodes.go index 4eede439e4..b7290ed235 100644 --- a/triedb/pathdb/nodes.go +++ b/triedb/pathdb/nodes.go @@ -500,8 +500,7 @@ func encodeNodeCompressed(addExtension bool, elements [][]byte, indices []int) [ flag |= 1 << 4 // Use the reserved flagE continue } - bitIndex := uint(pos % 8) - bitmap[pos/8] |= 1 << bitIndex + setBit(bitmap, pos) } enc = append(enc, flag) enc = append(enc, bitmap...) @@ -553,14 +552,7 @@ func decodeNodeCompressed(data []byte) ([][]byte, []int, error) { return nil, nil, errors.New("invalid data: too short") } bitmap := data[1:3] - for index, b := range bitmap { - for bitIdx := 0; bitIdx < 8; bitIdx++ { - if b&(1<