feat: versioned pectra compatability

Author: ypatil12

beacon/beacon_state_top_level_roots.go

@@ -1,13 +1,32 @@
 package beacon
 
 import (
+	"errors"
 	"reflect"
 
 	"github.com/Layr-Labs/eigenpod-proofs-generation/common"
+	"github.com/attestantio/go-eth2-client/spec"
 	"github.com/attestantio/go-eth2-client/spec/phase0"
 )
 
-type BeaconStateTopLevelRoots struct {
+type VersionedBeaconStateTopLevelRoots struct {
+	Version spec.DataVersion
+	Deneb   *BeaconStateTopLevelRootsDeneb
+	Electra *BeaconStateTopLevelRootsElectra
+}
+
+func (v *VersionedBeaconStateTopLevelRoots) GetBalancesRoot() (*phase0.Root, error) {
+	switch v.Version {
+	case spec.DataVersionDeneb:
+		return v.Deneb.BalancesRoot, nil
+	case spec.DataVersionElectra:
+		return v.Electra.BalancesRoot, nil
+	default:
+		return nil, errors.New("unsupported beacon state version")
+	}
+}
+
+type BeaconStateTopLevelRootsDeneb struct {
 	GenesisTimeRoot                  *phase0.Root
 	GenesisValidatorsRoot            *phase0.Root
 	SlotRoot                         *phase0.Root
@@ -38,18 +57,68 @@ type BeaconStateTopLevelRoots struct {
 	HistoricalSummariesRoot          *phase0.Root
 }
 
-func ProveBeaconTopLevelRootAgainstBeaconState(beaconTopLevelRoots *BeaconStateTopLevelRoots, index uint64) (common.Proof, error) {
-	v := reflect.ValueOf(*beaconTopLevelRoots)
+type BeaconStateTopLevelRootsElectra struct {
+	GenesisTimeRoot                   *phase0.Root
+	GenesisValidatorsRoot             *phase0.Root
+	SlotRoot                          *phase0.Root
+	ForkRoot                          *phase0.Root
+	LatestBlockHeaderRoot             *phase0.Root
+	BlockRootsRoot                    *phase0.Root
+	StateRootsRoot                    *phase0.Root
+	HistoricalRootsRoot               *phase0.Root
+	ETH1DataRoot                      *phase0.Root
+	ETH1DataVotesRoot                 *phase0.Root
+	ETH1DepositIndexRoot              *phase0.Root
+	ValidatorsRoot                    *phase0.Root
+	BalancesRoot                      *phase0.Root
+	RANDAOMixesRoot                   *phase0.Root
+	SlashingsRoot                     *phase0.Root
+	PreviousEpochParticipationRoot    *phase0.Root
+	CurrentEpochParticipationRoot     *phase0.Root
+	JustificationBitsRoot             *phase0.Root
+	PreviousJustifiedCheckpointRoot   *phase0.Root
+	CurrentJustifiedCheckpointRoot    *phase0.Root
+	FinalizedCheckpointRoot           *phase0.Root
+	InactivityScoresRoot              *phase0.Root
+	CurrentSyncCommitteeRoot          *phase0.Root
+	NextSyncCommitteeRoot             *phase0.Root
+	LatestExecutionPayloadHeaderRoot  *phase0.Root
+	NextWithdrawalIndexRoot           *phase0.Root
+	NextWithdrawalValidatorIndexRoot  *phase0.Root
+	HistoricalSummariesRoot           *phase0.Root
+	DepositRequestsStartIndexRoot     *phase0.Root
+	DepositBalanceToConsumeRoot       *phase0.Root
+	ExitBalanceToConsumeRoot          *phase0.Root
+	EarliestExitEpochRoot             *phase0.Root
+	ConsolidationBalanceToConsumeRoot *phase0.Root
+	EarliestConsolidationEpochRoot    *phase0.Root
+	PendingDepositsRoot               *phase0.Root
+	PendingPartialWithdrawalsRoot     *phase0.Root
+	PendingConsolidationsRoot         *phase0.Root
+}
+
+func ProveBeaconTopLevelRootAgainstBeaconState(beaconTopLevelRoots *VersionedBeaconStateTopLevelRoots, index uint64) (common.Proof, error) {
+	var v reflect.Value
+	switch beaconTopLevelRoots.Version {
+	case spec.DataVersionDeneb:
+		v = reflect.ValueOf(*beaconTopLevelRoots.Deneb)
+	case spec.DataVersionElectra:
+		v = reflect.ValueOf(*beaconTopLevelRoots.Electra)
+	default:
+		return nil, errors.New("unsupported beacon state version")
+	}
+
 	beaconTopLevelRootsList := make([]interface{}, v.NumField())
 	for i := 0; i < v.NumField(); i++ {
 		r := v.Field(i).Interface()
 		typedR := r.(*phase0.Root)
 		beaconTopLevelRootsList[i] = *typedR
 	}
+
 	roots := make([]phase0.Root, len(beaconTopLevelRootsList))
 	for i, v := range beaconTopLevelRootsList {
 		roots[i] = v.(phase0.Root)
 	}
 
-	return common.GetProof(roots, index, BEACON_STATE_TREE_HEIGHT)
+	return common.GetProof(roots, index, BEACON_STATE_TREE_HEIGHT_DENEB)
 }

beacon/constants.go

@@ -4,7 +4,8 @@ const (
 	BEACON_BLOCK_HEADER_NUM_FIELDS = uint64(5)
 
 	BEACON_BLOCK_HEADER_TREE_HEIGHT = uint64(3)
-	BEACON_STATE_TREE_HEIGHT        = uint64(5)
+	BEACON_STATE_TREE_HEIGHT_DENEB  = uint64(5)
+	BEACON_STATE_TREE_HEIGHT_PECTRA = uint64(6)
 	BALANCE_TREE_HEIGHT             = uint64(38)
 	VALIDATOR_TREE_HEIGHT           = uint64(40)
 

beacon/deneb.go

@@ -2,16 +2,17 @@ package beacon
 
 import (
 	"github.com/Layr-Labs/eigenpod-proofs-generation/common"
+	"github.com/attestantio/go-eth2-client/spec"
 	"github.com/attestantio/go-eth2-client/spec/altair"
 	"github.com/attestantio/go-eth2-client/spec/deneb"
 	"github.com/attestantio/go-eth2-client/spec/phase0"
 	ssz "github.com/ferranbt/fastssz"
 )
 
 // taken from https://github.com/attestantio/go-eth2-client/blob/21f7dd480fed933d8e0b1c88cee67da721c80eb2/spec/deneb/beaconstate_ssz.go#L640
-func ComputeBeaconStateTopLevelRootsDeneb(b *deneb.BeaconState) (*BeaconStateTopLevelRoots, error) {
+func ComputeBeaconStateTopLevelRootsDeneb(b *deneb.BeaconState) (*VersionedBeaconStateTopLevelRoots, error) {
 	var err error
-	beaconStateTopLevelRoots := &BeaconStateTopLevelRoots{}
+	beaconStateTopLevelRoots := &BeaconStateTopLevelRootsDeneb{}
 
 	hh := ssz.NewHasher()
 
@@ -409,5 +410,8 @@ func ComputeBeaconStateTopLevelRootsDeneb(b *deneb.BeaconState) (*BeaconStateTop
 		hh.Reset()
 	}
 
-	return beaconStateTopLevelRoots, nil
+	return &VersionedBeaconStateTopLevelRoots{
+		Version: spec.DataVersionDeneb,
+		Deneb:   beaconStateTopLevelRoots,
+	}, nil
 }