Improve HTR of State (#5058)

* add cache
* Update beacon-chain/state/stateutil/blocks.go

Co-Authored-By: terence tsao <terence@prysmaticlabs.com>
* Update beacon-chain/state/stateutil/blocks.go

Co-Authored-By: terence tsao <terence@prysmaticlabs.com>
* Update beacon-chain/state/stateutil/hash_function.go

Co-Authored-By: terence tsao <terence@prysmaticlabs.com>
* Merge branch 'master' into improveHTR
* add back string casting
* fix imports

beacon-chain/state/stateutil/BUILD.bazel

@@ -6,6 +6,7 @@ go_library(
         "arrays.go",
         "attestations.go",
         "blocks.go",
+        "hash_function.go",
         "helpers.go",
         "state_root.go",
         "validators.go",
@@ -24,7 +25,6 @@ go_library(
         "@com_github_dgraph_io_ristretto//:go_default_library",
         "@com_github_minio_sha256_simd//:go_default_library",
         "@com_github_pkg_errors//:go_default_library",
-        "@com_github_protolambda_zssz//htr:go_default_library",
         "@com_github_protolambda_zssz//merkle:go_default_library",
         "@com_github_prysmaticlabs_ethereumapis//eth/v1alpha1:go_default_library",
         "@com_github_prysmaticlabs_go_bitfield//:go_default_library",

beacon-chain/state/stateutil/blocks.go

@@ -7,6 +7,8 @@ import (
 	"github.com/pkg/errors"
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
 	"github.com/prysmaticlabs/prysm/shared/bytesutil"
+	"github.com/prysmaticlabs/prysm/shared/featureconfig"
+	"github.com/prysmaticlabs/prysm/shared/hashutil"
 	"github.com/prysmaticlabs/prysm/shared/params"
 )
 
@@ -34,6 +36,7 @@ func BlockHeaderRoot(header *ethpb.BeaconBlockHeader) ([32]byte, error) {
 // a BeaconBlockHeader struct according to the eth2
 // Simple Serialize specification.
 func Eth1Root(eth1Data *ethpb.Eth1Data) ([32]byte, error) {
+	enc := make([]byte, 0, 96)
 	fieldRoots := make([][]byte, 3)
 	for i := 0; i < len(fieldRoots); i++ {
 		fieldRoots[i] = make([]byte, 32)
@@ -42,31 +45,54 @@ func Eth1Root(eth1Data *ethpb.Eth1Data) ([32]byte, error) {
 		if len(eth1Data.DepositRoot) > 0 {
 			depRoot := bytesutil.ToBytes32(eth1Data.DepositRoot)
 			fieldRoots[0] = depRoot[:]
+			enc = append(enc, depRoot[:]...)
 		}
 		eth1DataCountBuf := make([]byte, 8)
 		binary.LittleEndian.PutUint64(eth1DataCountBuf, eth1Data.DepositCount)
 		eth1CountRoot := bytesutil.ToBytes32(eth1DataCountBuf)
 		fieldRoots[1] = eth1CountRoot[:]
+		enc = append(enc, eth1CountRoot[:]...)
 		if len(eth1Data.BlockHash) > 0 {
 			blockHash := bytesutil.ToBytes32(eth1Data.BlockHash)
 			fieldRoots[2] = blockHash[:]
+			enc = append(enc, blockHash[:]...)
+		}
+		if featureconfig.Get().EnableSSZCache {
+			if found, ok := cachedHasher.rootsCache.Get(string(enc)); ok && found != nil {
+				return found.([32]byte), nil
+			}
 		}
 	}
-	return bitwiseMerkleize(fieldRoots, uint64(len(fieldRoots)), uint64(len(fieldRoots)))
+	root, err := bitwiseMerkleize(fieldRoots, uint64(len(fieldRoots)), uint64(len(fieldRoots)))
+	if err != nil {
+		return [32]byte{}, err
+	}
+	if featureconfig.Get().EnableSSZCache {
+		cachedHasher.rootsCache.Set(string(enc), root, 32)
+	}
+	return root, nil
 }
 
 // Eth1DataVotesRoot computes the HashTreeRoot Merkleization of
 // a list of Eth1Data structs according to the eth2
 // Simple Serialize specification.
 func Eth1DataVotesRoot(eth1DataVotes []*ethpb.Eth1Data) ([32]byte, error) {
 	eth1VotesRoots := make([][]byte, 0)
+	enc := make([]byte, len(eth1DataVotes)*32)
 	for i := 0; i < len(eth1DataVotes); i++ {
 		eth1, err := Eth1Root(eth1DataVotes[i])
 		if err != nil {
 			return [32]byte{}, errors.Wrap(err, "could not compute eth1data merkleization")
 		}
+		copy(enc[(i*32):(i+1)*32], eth1[:])
 		eth1VotesRoots = append(eth1VotesRoots, eth1[:])
 	}
+	hashKey := hashutil.FastSum256(enc)
+	if featureconfig.Get().EnableSSZCache {
+		if found, ok := cachedHasher.rootsCache.Get(string(hashKey[:])); ok && found != nil {
+			return found.([32]byte), nil
+		}
+	}
 	eth1Chunks, err := pack(eth1VotesRoots)
 	if err != nil {
 		return [32]byte{}, errors.Wrap(err, "could not chunk eth1 votes roots")
@@ -82,5 +108,9 @@ func Eth1DataVotesRoot(eth1DataVotes []*ethpb.Eth1Data) ([32]byte, error) {
 	// We need to mix in the length of the slice.
 	eth1VotesRootBufRoot := make([]byte, 32)
 	copy(eth1VotesRootBufRoot, eth1VotesRootBuf.Bytes())
-	return mixInLength(eth1VotesRootsRoot, eth1VotesRootBufRoot), nil
+	root := mixInLength(eth1VotesRootsRoot, eth1VotesRootBufRoot)
+	if featureconfig.Get().EnableSSZCache {
+		cachedHasher.rootsCache.Set(string(hashKey[:]), root, 32)
+	}
+	return root, nil
 }

beacon-chain/state/stateutil/hash_function.go

@@ -0,0 +1,26 @@
+package stateutil
+
+import "encoding/binary"
+
+// HashFn describes a hash function and its associated bytes buffer
+type HashFn struct {
+	f           func(input []byte) [32]byte
+	bytesBuffer [64]byte
+}
+
+// Combi describes a method which merges two 32-byte arrays and hashes
+// them.
+func (h HashFn) Combi(a [32]byte, b [32]byte) [32]byte {
+	copy(h.bytesBuffer[:32], a[:])
+	copy(h.bytesBuffer[32:], b[:])
+	return h.f(h.bytesBuffer[:])
+}
+
+// MixIn describes a method where we add in the provided
+// integer to the end of the byte array and hash it.
+func (h HashFn) MixIn(a [32]byte, i uint64) [32]byte {
+	copy(h.bytesBuffer[:32], a[:])
+	copy(h.bytesBuffer[32:], make([]byte, 32, 32))
+	binary.LittleEndian.PutUint64(h.bytesBuffer[32:], i)
+	return h.f(h.bytesBuffer[:])
+}

beacon-chain/state/stateutil/helpers.go

@@ -6,7 +6,6 @@ import (
 
 	"github.com/minio/sha256-simd"
 	"github.com/pkg/errors"
-	"github.com/protolambda/zssz/htr"
 	"github.com/protolambda/zssz/merkle"
 	"github.com/prysmaticlabs/go-bitfield"
 	"github.com/prysmaticlabs/prysm/shared/hashutil"
@@ -47,23 +46,27 @@ func bitwiseMerkleize(chunks [][]byte, count uint64, limit uint64) ([32]byte, er
 	if count > limit {
 		return [32]byte{}, errors.New("merkleizing list that is too large, over limit")
 	}
-	hasher := htr.HashFn(hashutil.CustomSHA256Hasher())
+	hashFn := &HashFn{
+		f: hashutil.CustomSHA256Hasher(),
+	}
 	leafIndexer := func(i uint64) []byte {
 		return chunks[i]
 	}
-	return merkle.Merkleize(hasher, count, limit, leafIndexer), nil
+	return merkle.Merkleize(hashFn.f, count, limit, leafIndexer), nil
 }
 
 // bitwiseMerkleizeArrays is used when a set of 32-byte root chunks are provided.
 func bitwiseMerkleizeArrays(chunks [][32]byte, count uint64, limit uint64) ([32]byte, error) {
 	if count > limit {
 		return [32]byte{}, errors.New("merkleizing list that is too large, over limit")
 	}
-	hasher := htr.HashFn(hashutil.CustomSHA256Hasher())
+	hashFn := &HashFn{
+		f: hashutil.CustomSHA256Hasher(),
+	}
 	leafIndexer := func(i uint64) []byte {
 		return chunks[i][:]
 	}
-	return merkle.Merkleize(hasher, count, limit, leafIndexer), nil
+	return merkle.Merkleize(hashFn.f, count, limit, leafIndexer), nil
 }
 
 func pack(serializedItems [][]byte) ([][]byte, error) {