crypto/merkle: optimize merkle tree hashing (#6513) (#9446) (backport #25)

CHANGELOG.md

@@ -2,8 +2,11 @@
 
 ## Unreleased
 
+## v0.37.4-v24-osmo-3
+
 * [#21](https://github.com/osmosis-labs/cometbft/pull/21) Move websocket logs to Debug
 * [#22](https://github.com/osmosis-labs/cometbft/pull/22) Fix txSearch pagination performance issues
+* [#25](https://github.com/osmosis-labs/cometbft/pull/25) Optimize merkle tree hashing
 
 ## v0.37.4-v24-osmo-2
 

crypto/merkle/hash.go

@@ -1,6 +1,8 @@
 package merkle
 
 import (
+	"hash"
+
 	"github.com/cometbft/cometbft/crypto/tmhash"
 )
 
@@ -20,7 +22,23 @@ func leafHash(leaf []byte) []byte {
 	return tmhash.Sum(append(leafPrefix, leaf...))
 }
 
+// returns tmhash(0x00 || leaf)
+func leafHashOpt(s hash.Hash, leaf []byte) []byte {
+	s.Reset()
+	s.Write(leafPrefix)
+	s.Write(leaf)
+	return s.Sum(nil)
+}
+
 // returns tmhash(0x01 || left || right)
 func innerHash(left []byte, right []byte) []byte {
 	return tmhash.Sum(append(innerPrefix, append(left, right...)...))
 }
+
+func innerHashOpt(s hash.Hash, left []byte, right []byte) []byte {
+	s.Reset()
+	s.Write(innerPrefix)
+	s.Write(left)
+	s.Write(right)
+	return s.Sum(nil)
+}

crypto/merkle/proof_key_path_test.go

@@ -36,6 +36,7 @@ func TestKeyPath(t *testing.T) {
 
 		res, err := KeyPathToKeys(path.String())
 		require.Nil(t, err)
+		require.Equal(t, len(keys), len(res))
 
 		for i, key := range keys {
 			require.Equal(t, key, res[i])

crypto/merkle/proof_test.go

@@ -173,12 +173,12 @@ func TestProofValidateBasic(t *testing.T) {
 	}
 }
 func TestVoteProtobuf(t *testing.T) {
-
 	_, proofs := ProofsFromByteSlices([][]byte{
 		[]byte("apple"),
 		[]byte("watermelon"),
 		[]byte("kiwi"),
 	})
+
 	testCases := []struct {
 		testName string
 		v1       *Proof

crypto/merkle/tree.go

@@ -1,22 +1,28 @@
 package merkle
 
 import (
+	"crypto/sha256"
+	"hash"
 	"math/bits"
 )
 
 // HashFromByteSlices computes a Merkle tree where the leaves are the byte slice,
 // in the provided order. It follows RFC-6962.
 func HashFromByteSlices(items [][]byte) []byte {
+	return hashFromByteSlices(sha256.New(), items)
+}
+
+func hashFromByteSlices(sha hash.Hash, items [][]byte) []byte {
 	switch len(items) {
 	case 0:
 		return emptyHash()
 	case 1:
-		return leafHash(items[0])
+		return leafHashOpt(sha, items[0])
 	default:
 		k := getSplitPoint(int64(len(items)))
-		left := HashFromByteSlices(items[:k])
-		right := HashFromByteSlices(items[k:])
-		return innerHash(left, right)
+		left := hashFromByteSlices(sha, items[:k])
+		right := hashFromByteSlices(sha, items[k:])
+		return innerHashOpt(sha, left, right)
 	}
 }
 
@@ -61,7 +67,7 @@ func HashFromByteSlices(items [][]byte) []byte {
 // implementation for so little benefit.
 func HashFromByteSlicesIterative(input [][]byte) []byte {
 	items := make([][]byte, len(input))
-
+	sha := sha256.New()
 	for i, leaf := range input {
 		items[i] = leafHash(leaf)
 	}
@@ -78,7 +84,7 @@ func HashFromByteSlicesIterative(input [][]byte) []byte {
 			wp := 0 // write position
 			for rp < size {
 				if rp+1 < size {
-					items[wp] = innerHash(items[rp], items[rp+1])
+					items[wp] = innerHashOpt(sha, items[rp], items[rp+1])
 					rp += 2
 				} else {
 					items[wp] = items[rp]