Custom Block HTR (#5219)

* add custom htr

* fix root

* fix everything

* Apply suggestions from code review

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

* Update beacon-chain/blockchain/receive_block.go

Co-Authored-By: terence tsao <terence@prysmaticlabs.com>

* Update beacon-chain/blockchain/process_block.go

Co-Authored-By: terence tsao <terence@prysmaticlabs.com>

* terence's review

Co-authored-by: terence tsao <terence@prysmaticlabs.com>
Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>

beacon-chain/blockchain/BUILD.bazel

@@ -40,6 +40,7 @@ go_library(
         "//beacon-chain/powchain:go_default_library",
         "//beacon-chain/state:go_default_library",
         "//beacon-chain/state/stategen:go_default_library",
+        "//beacon-chain/state/stateutil:go_default_library",
         "//proto/beacon/p2p/v1:go_default_library",
         "//shared/attestationutil:go_default_library",
         "//shared/bytesutil:go_default_library",

beacon-chain/blockchain/process_block.go

@@ -5,9 +5,9 @@ import (
 	"encoding/hex"
 	"fmt"
 
+	"github.com/prysmaticlabs/prysm/beacon-chain/state/stateutil"
 	"github.com/pkg/errors"
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
-	"github.com/prysmaticlabs/go-ssz"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/helpers"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/state"
 	"github.com/prysmaticlabs/prysm/beacon-chain/flags"
@@ -67,7 +67,7 @@ func (s *Service) onBlock(ctx context.Context, signed *ethpb.SignedBeaconBlock)
 	}
 	preStateValidatorCount := preState.NumValidators()
 
-	root, err := ssz.HashTreeRoot(b)
+	root, err := stateutil.BlockRoot(b)
 	if err != nil {
 		return nil, errors.Wrapf(err, "could not get signing root of block %d", b.Slot)
 	}
@@ -213,7 +213,7 @@ func (s *Service) onBlockInitialSyncStateTransition(ctx context.Context, signed
 	if err := s.beaconDB.SaveBlock(ctx, signed); err != nil {
 		return errors.Wrapf(err, "could not save block from slot %d", b.Slot)
 	}
-	root, err := ssz.HashTreeRoot(b)
+	root, err := stateutil.BlockRoot(b)
 	if err != nil {
 		return errors.Wrapf(err, "could not get signing root of block %d", b.Slot)
 	}

beacon-chain/blockchain/receive_block.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/hex"
 
+	"github.com/prysmaticlabs/prysm/beacon-chain/state/stateutil"
 	"github.com/pkg/errors"
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
 	"github.com/prysmaticlabs/go-ssz"
@@ -88,7 +89,7 @@ func (s *Service) ReceiveBlockNoPubsub(ctx context.Context, block *ethpb.SignedB
 	defer s.epochParticipationLock.Unlock()
 	s.epochParticipation[helpers.SlotToEpoch(blockCopy.Block.Slot)] = precompute.Balances
 
-	root, err := ssz.HashTreeRoot(blockCopy.Block)
+	root, err := stateutil.BlockRoot(blockCopy.Block)
 	if err != nil {
 		return errors.Wrap(err, "could not get signing root on received block")
 	}
@@ -139,7 +140,7 @@ func (s *Service) ReceiveBlockNoPubsubForkchoice(ctx context.Context, block *eth
 		return err
 	}
 
-	root, err := ssz.HashTreeRoot(blockCopy.Block)
+	root, err := stateutil.BlockRoot(blockCopy.Block)
 	if err != nil {
 		return errors.Wrap(err, "could not get signing root on received block")
 	}
@@ -191,7 +192,7 @@ func (s *Service) ReceiveBlockNoVerify(ctx context.Context, block *ethpb.SignedB
 		return err
 	}
 
-	root, err := ssz.HashTreeRoot(blockCopy.Block)
+	root, err := stateutil.BlockRoot(blockCopy.Block)
 	if err != nil {
 		return errors.Wrap(err, "could not get signing root on received blockCopy")
 	}

beacon-chain/core/blocks/block_operations.go

@@ -58,6 +58,25 @@ func verifySigningRoot(obj interface{}, pub []byte, signature []byte, domain uin
 	return nil
 }
 
+func verifyBlockRoot(blk *ethpb.BeaconBlock, pub []byte, signature []byte, domain uint64) error {
+	publicKey, err := bls.PublicKeyFromBytes(pub)
+	if err != nil {
+		return errors.Wrap(err, "could not convert bytes to public key")
+	}
+	sig, err := bls.SignatureFromBytes(signature)
+	if err != nil {
+		return errors.Wrap(err, "could not convert bytes to signature")
+	}
+	root, err := stateutil.BlockRoot(blk)
+	if err != nil {
+		return errors.Wrap(err, "could not get signing root")
+	}
+	if !sig.Verify(root[:], publicKey, domain) {
+		return ErrSigFailedToVerify
+	}
+	return nil
+}
+
 // Deprecated: This method uses deprecated ssz.SigningRoot.
 func verifyDepositDataSigningRoot(obj *ethpb.Deposit_Data, pub []byte, signature []byte, domain uint64) error {
 	publicKey, err := bls.PublicKeyFromBytes(pub)
@@ -223,7 +242,7 @@ func ProcessBlockHeader(
 	if err != nil {
 		return nil, err
 	}
-	if err := verifySigningRoot(block.Block, proposer.PublicKey, block.Signature, domain); err != nil {
+	if err := verifyBlockRoot(block.Block, proposer.PublicKey, block.Signature, domain); err != nil {
 		return nil, ErrSigFailedToVerify
 	}
 
@@ -286,7 +305,7 @@ func ProcessBlockHeaderNoVerify(
 		return nil, fmt.Errorf("proposer at index %d was previously slashed", idx)
 	}
 
-	bodyRoot, err := ssz.HashTreeRoot(block.Body)
+	bodyRoot, err := stateutil.BlockBodyRoot(block.Body)
 	if err != nil {
 		return nil, err
 	}

beacon-chain/db/kv/BUILD.bazel

@@ -29,6 +29,7 @@ go_library(
         "//beacon-chain/db/filters:go_default_library",
         "//beacon-chain/db/iface:go_default_library",
         "//beacon-chain/state:go_default_library",
+        "//beacon-chain/state/stateutil:go_default_library",
         "//proto/beacon/db:go_default_library",
         "//proto/beacon/p2p/v1:go_default_library",
         "//shared/bytesutil:go_default_library",

beacon-chain/db/kv/blocks.go

@@ -7,6 +7,8 @@ import (
 	"math"
 	"strconv"
 
+	"github.com/prysmaticlabs/prysm/beacon-chain/state/stateutil"
+
 	"github.com/pkg/errors"
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
 	"github.com/prysmaticlabs/go-ssz"
@@ -190,7 +192,7 @@ func (k *Store) DeleteBlocks(ctx context.Context, blockRoots [][32]byte) error {
 func (k *Store) SaveBlock(ctx context.Context, signed *ethpb.SignedBeaconBlock) error {
 	ctx, span := trace.StartSpan(ctx, "BeaconDB.SaveBlock")
 	defer span.End()
-	blockRoot, err := ssz.HashTreeRoot(signed.Block)
+	blockRoot, err := stateutil.BlockRoot(signed.Block)
 	if err != nil {
 		return err
 	}

beacon-chain/rpc/beacon/BUILD.bazel

@@ -35,6 +35,7 @@ go_library(
         "//beacon-chain/powchain:go_default_library",
         "//beacon-chain/state:go_default_library",
         "//beacon-chain/state/stategen:go_default_library",
+        "//beacon-chain/state/stateutil:go_default_library",
         "//proto/beacon/p2p/v1:go_default_library",
         "//shared/attestationutil:go_default_library",
         "//shared/bytesutil:go_default_library",

beacon-chain/rpc/beacon/blocks.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"strconv"
 
+	"github.com/prysmaticlabs/prysm/beacon-chain/state/stateutil"
 	ptypes "github.com/gogo/protobuf/types"
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
 	"github.com/prysmaticlabs/go-ssz"
@@ -241,7 +242,7 @@ func (bs *Server) chainHeadRetrieval(ctx context.Context) (*ethpb.ChainHead, err
 	if headBlock == nil {
 		return nil, status.Error(codes.Internal, "Head block of chain was nil")
 	}
-	headBlockRoot, err := ssz.HashTreeRoot(headBlock.Block)
+	headBlockRoot, err := stateutil.BlockRoot(headBlock.Block)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "Could not get head block root: %v", err)
 	}

beacon-chain/state/stateutil/BUILD.bazel

@@ -30,12 +30,14 @@ go_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",
+        "@com_github_prysmaticlabs_go_ssz//:go_default_library",
     ],
 )
 
 go_test(
     name = "go_default_test",
     srcs = [
+        "blocks_test.go",
         "state_root_cache_fuzz_test.go",
         "state_root_test.go",
         "trie_helpers_test.go",
@@ -53,3 +55,26 @@ go_test(
         "@com_github_prysmaticlabs_go_ssz//:go_default_library",
     ],
 )
+
+go_test(
+    name = "go_benchmark_test",
+    size = "small",
+    srcs = ["benchmark_test.go"],
+    args = [
+        "-test.bench=.",
+        "-test.benchmem",
+        "-test.v",
+    ],
+    local = True,
+    tags = [
+        "benchmark",
+        "manual",
+        "no-cache",
+    ],
+    deps = [
+        "//beacon-chain/state/stateutil:go_default_library",
+        "//shared/testutil:go_default_library",
+        "@com_github_prysmaticlabs_ethereumapis//eth/v1alpha1:go_default_library",
+        "@com_github_prysmaticlabs_go_ssz//:go_default_library",
+    ],
+)

beacon-chain/state/stateutil/attestations.go

@@ -28,7 +28,7 @@ func PendingAttestationRoot(att *pb.PendingAttestation) ([32]byte, error) {
 	fieldRoots := [][32]byte{}
 	if att != nil {
 		// Bitfield.
-		aggregationRoot, err := bitlistRoot(att.AggregationBits, 2048)
+		aggregationRoot, err := bitlistRoot(att.AggregationBits, params.BeaconConfig().MaxValidatorsPerCommittee)
 		if err != nil {
 			return [32]byte{}, err
 		}
@@ -88,6 +88,64 @@ func marshalAttestationData(data *ethpb.AttestationData) []byte {
 	return enc
 }
 
+func attestationRoot(att *ethpb.Attestation) ([32]byte, error) {
+	fieldRoots := make([][32]byte, 3)
+
+	// Bitfield.
+	aggregationRoot, err := bitlistRoot(att.AggregationBits, params.BeaconConfig().MaxValidatorsPerCommittee)
+	if err != nil {
+		return [32]byte{}, err
+	}
+	fieldRoots[0] = aggregationRoot
+
+	dataRoot, err := attestationDataRoot(att.Data)
+	if err != nil {
+		return [32]byte{}, err
+	}
+	fieldRoots[1] = dataRoot
+
+	signatureBuf := bytesutil.ToBytes96(att.Signature)
+	packedSig, err := pack([][]byte{signatureBuf[:]})
+	if err != nil {
+		return [32]byte{}, err
+	}
+	sigRoot, err := bitwiseMerkleize(packedSig, uint64(len(packedSig)), uint64(len(packedSig)))
+	if err != nil {
+		return [32]byte{}, err
+	}
+	fieldRoots[2] = sigRoot
+	return bitwiseMerkleizeArrays(fieldRoots, uint64(len(fieldRoots)), uint64(len(fieldRoots)))
+}
+
+func blockAttestationRoot(atts []*ethpb.Attestation) ([32]byte, error) {
+	roots := make([][]byte, len(atts))
+	for i := 0; i < len(atts); i++ {
+		pendingRoot, err := attestationRoot(atts[i])
+		if err != nil {
+			return [32]byte{}, errors.Wrap(err, "could not attestation merkleization")
+		}
+		roots[i] = pendingRoot[:]
+	}
+
+	attsRootsRoot, err := bitwiseMerkleize(
+		roots,
+		uint64(len(roots)),
+		params.BeaconConfig().MaxAttestations,
+	)
+	if err != nil {
+		return [32]byte{}, errors.Wrap(err, "could not compute block attestations merkleization")
+	}
+	attsLenBuf := new(bytes.Buffer)
+	if err := binary.Write(attsLenBuf, binary.LittleEndian, uint64(len(atts))); err != nil {
+		return [32]byte{}, errors.Wrap(err, "could not marshal epoch attestations length")
+	}
+	// We need to mix in the length of the slice.
+	attsLenRoot := make([]byte, 32)
+	copy(attsLenRoot, attsLenBuf.Bytes())
+	res := mixInLength(attsRootsRoot, attsLenRoot)
+	return res, nil
+}
+
 func attestationDataRoot(data *ethpb.AttestationData) ([32]byte, error) {
 	fieldRoots := make([][]byte, 5)
 

beacon-chain/state/stateutil/benchmark_test.go

@@ -0,0 +1,49 @@
+package stateutil_benchmark
+
+import (
+	"testing"
+
+	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
+	"github.com/prysmaticlabs/go-ssz"
+	"github.com/prysmaticlabs/prysm/beacon-chain/state/stateutil"
+	"github.com/prysmaticlabs/prysm/shared/testutil"
+)
+
+func BenchmarkBlockHTR(b *testing.B) {
+	genState, keys := testutil.DeterministicGenesisState(b, 200)
+	conf := testutil.DefaultBlockGenConfig()
+	blk, err := testutil.GenerateFullBlock(genState, keys, conf, 10)
+	if err != nil {
+		b.Fatal(err)
+	}
+	atts := make([]*ethpb.Attestation, 0, 128)
+	for i := 0; i < 128; i++ {
+		atts = append(atts, blk.Block.Body.Attestations[0])
+	}
+	deposits, _, err := testutil.DeterministicDepositsAndKeys(16)
+	if err != nil {
+		b.Fatal(err)
+	}
+	blk.Block.Body.Attestations = atts
+	blk.Block.Body.Deposits = deposits
+
+	b.Run("SSZ_HTR", func(b *testing.B) {
+		b.ResetTimer()
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			if _, err := ssz.HashTreeRoot(blk.Block); err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+
+	b.Run("Custom_SSZ_HTR", func(b *testing.B) {
+		b.ResetTimer()
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			if _, err := stateutil.BlockRoot(blk.Block); err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+}