Assign instead of append to keys slice

x/merkledb/codec.go

@@ -5,7 +5,6 @@ package merkledb
 
 import (
 	"bytes"
-	"crypto/sha256"
 	"encoding/binary"
 	"errors"
 	"io"
@@ -74,62 +73,6 @@ func encodedDBNodeSize(n *dbNode) int {
 	return size
 }
 
-// Returns the canonical hash of [n].
-//
-// Assumes [n] is non-nil.
-// This method is performance critical. It is not expected to perform any memory
-// allocations.
-func hashNode(n *node) ids.ID {
-	var (
-		// sha.Write always returns nil, so we ignore its return values.
-		sha  = sha256.New()
-		hash ids.ID
-		// The hash length is larger than the maximum Uvarint length. This
-		// ensures binary.AppendUvarint doesn't perform any memory allocations.
-		emptyHashBuffer = hash[:0]
-	)
-
-	// By directly calling sha.Write rather than passing sha around as an
-	// io.Writer, the compiler can perform sufficient escape analysis to avoid
-	// allocating buffers on the heap.
-	numChildren := len(n.children)
-	_, _ = sha.Write(binary.AppendUvarint(emptyHashBuffer, uint64(numChildren)))
-
-	// Avoid allocating keys entirely if the node doesn't have any children.
-	if numChildren != 0 {
-		// By allocating BranchFactorLargest rather than len(n.children), this
-		// slice is allocated on the stack rather than the heap.
-		// BranchFactorLargest is at least len(n.children) which avoids memory
-		// allocations.
-		keys := make([]byte, 0, BranchFactorLargest)
-		for k := range n.children {
-			keys = append(keys, k)
-		}
-
-		// Ensure that the order of entries is correct.
-		slices.Sort(keys)
-		for _, index := range keys {
-			entry := n.children[index]
-			_, _ = sha.Write(binary.AppendUvarint(emptyHashBuffer, uint64(index)))
-			_, _ = sha.Write(entry.id[:])
-		}
-	}
-
-	if n.valueDigest.HasValue() {
-		_, _ = sha.Write(trueBytes)
-		value := n.valueDigest.Value()
-		_, _ = sha.Write(binary.AppendUvarint(emptyHashBuffer, uint64(len(value))))
-		_, _ = sha.Write(value)
-	} else {
-		_, _ = sha.Write(falseBytes)
-	}
-
-	_, _ = sha.Write(binary.AppendUvarint(emptyHashBuffer, uint64(n.key.length)))
-	_, _ = sha.Write(n.key.Bytes())
-	sha.Sum(emptyHashBuffer)
-	return hash
-}
-
 // Assumes [n] is non-nil.
 func encodeDBNode(n *dbNode) []byte {
 	length := encodedDBNodeSize(n)

x/merkledb/codec_test.go

@@ -18,74 +18,6 @@ import (
 )
 
 var (
-	hashNodeTests = []struct {
-		name         string
-		n            *node
-		expectedHash string
-	}{
-		{
-			name:         "empty node",
-			n:            newNode(Key{}),
-			expectedHash: "rbhtxoQ1DqWHvb6w66BZdVyjmPAneZUSwQq9uKj594qvFSdav",
-		},
-		{
-			name: "has value",
-			n: func() *node {
-				n := newNode(Key{})
-				n.setValue(maybe.Some([]byte("value1")))
-				return n
-			}(),
-			expectedHash: "2vx2xueNdWoH2uB4e8hbMU5jirtZkZ1c3ePCWDhXYaFRHpCbnQ",
-		},
-		{
-			name:         "has key",
-			n:            newNode(ToKey([]byte{0, 1, 2, 3, 4, 5, 6, 7})),
-			expectedHash: "2vA8ggXajhFEcgiF8zHTXgo8T2ALBFgffp1xfn48JEni1Uj5uK",
-		},
-		{
-			name: "1 child",
-			n: func() *node {
-				n := newNode(Key{})
-				childNode := newNode(ToKey([]byte{255}))
-				childNode.setValue(maybe.Some([]byte("value1")))
-				n.addChildWithID(childNode, 4, hashNode(childNode))
-				return n
-			}(),
-			expectedHash: "YfJRufqUKBv9ez6xZx6ogpnfDnw9fDsyebhYDaoaH57D3vRu3",
-		},
-		{
-			name: "2 children",
-			n: func() *node {
-				n := newNode(Key{})
-
-				childNode1 := newNode(ToKey([]byte{255}))
-				childNode1.setValue(maybe.Some([]byte("value1")))
-
-				childNode2 := newNode(ToKey([]byte{237}))
-				childNode2.setValue(maybe.Some([]byte("value2")))
-
-				n.addChildWithID(childNode1, 4, hashNode(childNode1))
-				n.addChildWithID(childNode2, 4, hashNode(childNode2))
-				return n
-			}(),
-			expectedHash: "YVmbx5MZtSKuYhzvHnCqGrswQcxmozAkv7xE1vTA2EiGpWUkv",
-		},
-		{
-			name: "16 children",
-			n: func() *node {
-				n := newNode(Key{})
-
-				for i := byte(0); i < 16; i++ {
-					childNode := newNode(ToKey([]byte{i << 4}))
-					childNode.setValue(maybe.Some([]byte("some value")))
-
-					n.addChildWithID(childNode, 4, hashNode(childNode))
-				}
-				return n
-			}(),
-			expectedHash: "5YiFLL7QV3f441See9uWePi3wVKsx9fgvX5VPhU8PRxtLqhwY",
-		},
-	}
 	encodeDBNodeTests = []struct {
 		name          string
 		n             *dbNode
@@ -613,70 +545,6 @@ func TestCodecDecodeDBNode_TooShort(t *testing.T) {
 	require.ErrorIs(err, io.ErrUnexpectedEOF)
 }
 
-// Ensure that hashNode is deterministic
-func FuzzHashNode(f *testing.F) {
-	f.Fuzz(
-		func(
-			t *testing.T,
-			randSeed int,
-		) {
-			require := require.New(t)
-			for _, bf := range validBranchFactors { // Create a random node
-				r := rand.New(rand.NewSource(int64(randSeed))) // #nosec G404
-
-				children := map[byte]*child{}
-				numChildren := r.Intn(int(bf)) // #nosec G404
-				for i := 0; i < numChildren; i++ {
-					compressedKeyLen := r.Intn(32) // #nosec G404
-					compressedKeyBytes := make([]byte, compressedKeyLen)
-					_, _ = r.Read(compressedKeyBytes) // #nosec G404
-
-					children[byte(i)] = &child{
-						compressedKey: ToKey(compressedKeyBytes),
-						id:            ids.GenerateTestID(),
-						hasValue:      r.Intn(2) == 1, // #nosec G404
-					}
-				}
-
-				hasValue := r.Intn(2) == 1 // #nosec G404
-				value := maybe.Nothing[[]byte]()
-				if hasValue {
-					valueBytes := make([]byte, r.Intn(64)) // #nosec G404
-					_, _ = r.Read(valueBytes)              // #nosec G404
-					value = maybe.Some(valueBytes)
-				}
-
-				key := make([]byte, r.Intn(32)) // #nosec G404
-				_, _ = r.Read(key)              // #nosec G404
-
-				hv := &node{
-					key: ToKey(key),
-					dbNode: dbNode{
-						children: children,
-						value:    value,
-					},
-				}
-
-				// Hash hv multiple times
-				hash1 := hashNode(hv)
-				hash2 := hashNode(hv)
-
-				// Make sure they're the same
-				require.Equal(hash1, hash2)
-			}
-		},
-	)
-}
-
-func TestHashNode(t *testing.T) {
-	for _, test := range hashNodeTests {
-		t.Run(test.name, func(t *testing.T) {
-			hash := hashNode(test.n)
-			require.Equal(t, test.expectedHash, hash.String())
-		})
-	}
-}
-
 func TestEncodeDBNode(t *testing.T) {
 	for _, test := range encodeDBNodeTests {
 		t.Run(test.name, func(t *testing.T) {
@@ -744,16 +612,6 @@ func TestUintSize(t *testing.T) {
 	}
 }
 
-func Benchmark_HashNode(b *testing.B) {
-	for _, benchmark := range hashNodeTests {
-		b.Run(benchmark.name, func(b *testing.B) {
-			for i := 0; i < b.N; i++ {
-				hashNode(benchmark.n)
-			}
-		})
-	}
-}
-
 func Benchmark_EncodeDBNode(b *testing.B) {
 	for _, benchmark := range encodeDBNodeTests {
 		b.Run(benchmark.name, func(b *testing.B) {

x/merkledb/db.go

@@ -163,6 +163,11 @@ type Config struct {
 	// BranchFactor determines the number of children each node can have.
 	BranchFactor BranchFactor
 
+	// Hasher defines the hash function to use when hashing the trie.
+	//
+	// If not specified, [SHA256Hasher] will be used.
+	Hasher Hasher
+
 	// RootGenConcurrency is the number of goroutines to use when
 	// generating a new state root.
 	//
@@ -234,6 +239,8 @@ type merkleDB struct {
 	hashNodesKeyPool *bytesPool
 
 	tokenSize int
+
+	hasher Hasher
 }
 
 // New returns a new merkle database.
@@ -255,6 +262,11 @@ func newDatabase(
 		return nil, err
 	}
 
+	hasher := config.Hasher
+	if hasher == nil {
+		hasher = SHA256Hasher
+	}
+
 	rootGenConcurrency := runtime.NumCPU()
 	if config.RootGenConcurrency != 0 {
 		rootGenConcurrency = int(config.RootGenConcurrency)
@@ -274,17 +286,23 @@ func newDatabase(
 			int(config.IntermediateNodeCacheSize),
 			int(config.IntermediateWriteBufferSize),
 			int(config.IntermediateWriteBatchSize),
-			BranchFactorToTokenSize[config.BranchFactor]),
-		valueNodeDB: newValueNodeDB(db,
+			BranchFactorToTokenSize[config.BranchFactor],
+			hasher,
+		),
+		valueNodeDB: newValueNodeDB(
+			db,
 			bufferPool,
 			metrics,
-			int(config.ValueNodeCacheSize)),
+			int(config.ValueNodeCacheSize),
+			hasher,
+		),
 		history:          newTrieHistory(int(config.HistoryLength)),
 		debugTracer:      getTracerIfEnabled(config.TraceLevel, DebugTrace, config.Tracer),
 		infoTracer:       getTracerIfEnabled(config.TraceLevel, InfoTrace, config.Tracer),
 		childViews:       make([]*view, 0, defaultPreallocationSize),
 		hashNodesKeyPool: newBytesPool(rootGenConcurrency),
 		tokenSize:        BranchFactorToTokenSize[config.BranchFactor],
+		hasher:           hasher,
 	}
 
 	shutdownType, err := trieDB.baseDB.Get(cleanShutdownKey)
@@ -1237,7 +1255,8 @@ func (db *merkleDB) initializeRoot() error {
 		}
 	}
 
-	db.rootID = root.calculateID(db.metrics)
+	db.metrics.HashCalculated()
+	db.rootID = db.hasher.HashNode(root)
 	db.root = maybe.Some(root)
 	return nil
 }

x/merkledb/db_test.go

@@ -48,6 +48,7 @@ func newDefaultConfig() Config {
 		Reg:                         prometheus.NewRegistry(),
 		Tracer:                      trace.Noop,
 		BranchFactor:                BranchFactor16,
+		Hasher:                      SHA256Hasher,
 	}
 }
 
@@ -962,6 +963,7 @@ func runRandDBTest(require *require.Assertions, r *rand.Rand, rt randTest, token
 				end,
 				root,
 				tokenSize,
+				db.hasher,
 			))
 		case opGenerateChangeProof:
 			root, err := db.GetMerkleRoot(context.Background())