crypto/merkle: hashing ~10% faster and ~60% fewer allocs

crypto/merkle/hash.go

@@ -1,6 +1,8 @@
 package merkle
 
 import (
+	"hash"
+
 	"github.com/lazyledger/lazyledger-core/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.go

@@ -50,13 +50,13 @@ func ProofsFromByteSlices(items [][]byte) (rootHash []byte, proofs []*Proof) {
 // Verify that the Proof proves the root hash.
 // Check sp.Index/sp.Total manually if needed
 func (sp *Proof) Verify(rootHash []byte, leaf []byte) error {
-	leafHash := leafHash(leaf)
 	if sp.Total < 0 {
 		return errors.New("proof total must be positive")
 	}
 	if sp.Index < 0 {
 		return errors.New("proof index cannot be negative")
 	}
+	leafHash := leafHash(leaf)
 	if !bytes.Equal(sp.LeafHash, leafHash) {
 		return fmt.Errorf("invalid leaf hash: wanted %X got %X", leafHash, sp.LeafHash)
 	}

crypto/merkle/proof_key_path_test.go

@@ -35,6 +35,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

@@ -171,7 +171,6 @@ func TestProofValidateBasic(t *testing.T) {
 	}
 }
 func TestVoteProtobuf(t *testing.T) {
-
 	_, proofs := ProofsFromByteSlices([][]byte{
 		[]byte("apple"),
 		[]byte("watermelon"),

crypto/merkle/tree.go

@@ -1,22 +1,27 @@
 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)
 	}
 }
 
@@ -60,10 +65,11 @@ func HashFromByteSlices(items [][]byte) []byte {
 // read, it might not be worthwhile to switch to a less intuitive
 // implementation for so little benefit.
 func HashFromByteSlicesIterative(input [][]byte) []byte {
+	sha := sha256.New()
 	items := make([][]byte, len(input))
 
 	for i, leaf := range input {
-		items[i] = leafHash(leaf)
+		items[i] = leafHashOpt(sha, leaf)
 	}
 
 	size := len(items)
@@ -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]

go.sum

@@ -613,8 +613,6 @@ github.com/lazyledger/go-verifcid v0.0.1-lazypatch h1:jAVwUw+DhuCzx5IcYpFh6d6HWx
 github.com/lazyledger/go-verifcid v0.0.1-lazypatch/go.mod h1:kXPYu0XqTNUKWA1h3M95UHjUqBzDwXVVt/RXZDjKJmQ=
 github.com/lazyledger/merkletree v0.0.0-20201214195110-6901c4c3c75f h1:jbyPAH6o6hGte4RtZBaqWs2n4Fl6hS7qJGXX3qnjiy4=
 github.com/lazyledger/merkletree v0.0.0-20201214195110-6901c4c3c75f/go.mod h1:10PA0NlnYtB8HrtwIDQAyTKWp8TEZ0zBZCGlYC/7+QE=
-github.com/lazyledger/nmt v0.4.0 h1:RcNmU7nfRtgsmIBN2T+5eOSfuOGLkX3F+cINdWdphrQ=
-github.com/lazyledger/nmt v0.4.0/go.mod h1:tY7ypPX26Sbkt6F8EbPl3AT33B5N0BJe4OVPbq849YI=
 github.com/lazyledger/nmt v0.5.0 h1:WEcFOx6rqJuI8wviyMgFH/ZOFL0jQowewxiHtErtV/I=
 github.com/lazyledger/nmt v0.5.0/go.mod h1:tY7ypPX26Sbkt6F8EbPl3AT33B5N0BJe4OVPbq849YI=
 github.com/lazyledger/rsmt2d v0.2.0 h1:RrKkAd9WTewCnOmAtvM5bz05PC+XzIYSrD+3V3odPtg=