fix: update epp to not take deneb state

benchmark_test.go

@@ -4,19 +4,20 @@ import (
 	"testing"
 
 	"github.com/Layr-Labs/eigenpod-proofs-generation/beacon"
+	"github.com/attestantio/go-eth2-client/spec"
 	"github.com/attestantio/go-eth2-client/spec/phase0"
 	"github.com/stretchr/testify/assert"
 )
 
 func BenchmarkComputeBeaconStateRoot(b *testing.B) {
-	computed, err := epp.ComputeBeaconStateRoot(beaconState.Deneb)
+	computed, err := epp.ComputeBeaconStateRoot(beaconState)
 	if err != nil {
 		b.Fatal(err)
 	}
 
 	var cached phase0.Root
 	for i := 0; i < b.N; i++ {
-		cached, err = epp.ComputeBeaconStateRoot(beaconState.Deneb)
+		cached, err = epp.ComputeBeaconStateRoot(beaconState)
 		if err != nil {
 			b.Fatal(err)
 		}
@@ -42,6 +43,11 @@ func BenchmarkComputeBeaconStateTopLevelRoots(b *testing.B) {
 }
 
 func BenchmarkComputeValidatorTree(b *testing.B) {
+	// If beacon state is not Deneb, then skip this test
+	if beaconState.Version != spec.DataVersionDeneb {
+		b.Skip("skipping test for non-Deneb beacon state")
+	}
+
 	computed, err := epp.ComputeValidatorTree(beaconState.Deneb.Slot, beaconState.Deneb.Validators)
 	if err != nil {
 		b.Fatal(err)
@@ -59,6 +65,11 @@ func BenchmarkComputeValidatorTree(b *testing.B) {
 }
 
 func BenchmarkComputeValidatorBalancesTree(b *testing.B) {
+	// If beacon state is not Deneb, then skip this test
+	if beaconState.Version != spec.DataVersionDeneb {
+		b.Skip("skipping test for non-Deneb beacon state")
+	}
+
 	computed, err := epp.ComputeValidatorBalancesTree(beaconState.Deneb.Slot, beaconState.Deneb.Balances)
 	if err != nil {
 		b.Fatal(err)

data/electra_mekong_beacon_headers_654719.json

@@ -0,0 +1,7 @@
+{
+    "slot": "654719",
+    "proposer_index": "88253",
+    "parent_root": "0x0a22809eecc5109a76e40aef633602d80c3b9706962c32830d46053a3d9d8f24",
+    "state_root": "0x875eabbd4e15c3e5203e79a0ace1190d1339c154c0f6432a5bc9c41b7a78962e",
+    "body_root": "0xb7d119550f4f56e3d3563c05da56ea05c8fd73abfab58700fb3d083ab048f09c"
+}

eigen_pod_proofs.go

@@ -2,10 +2,10 @@ package eigenpodproofs
 
 import (
 	"errors"
+	"fmt"
 	"time"
 
 	"github.com/attestantio/go-eth2-client/spec"
-	"github.com/attestantio/go-eth2-client/spec/deneb"
 	"github.com/attestantio/go-eth2-client/spec/phase0"
 	expirable "github.com/hashicorp/golang-lru/v2/expirable"
 
@@ -64,27 +64,49 @@ func (epp *EigenPodProofs) PrecomputeCache(state *spec.VersionedBeaconState) err
 		return err
 	}
 
-	epp.ComputeBeaconStateRoot(state.Deneb)
+	epp.ComputeBeaconStateRoot(state)
 	epp.ComputeBeaconStateTopLevelRoots(state)
 	epp.ComputeVersionedBeaconStateTopLevelRoots(state)
 	epp.ComputeValidatorTree(slot, validators)
 	epp.ComputeValidatorBalancesTree(slot, balances)
 	return nil
 }
 
-func (epp *EigenPodProofs) ComputeBeaconStateRoot(beaconState *deneb.BeaconState) (phase0.Root, error) {
-	beaconStateRoot, err := epp.loadOrComputeBeaconStateRoot(
-		beaconState.Slot,
-		func() (phase0.Root, error) {
-			stateRoot, err := beaconState.HashTreeRoot()
-			if err != nil {
-				return phase0.Root{}, err
-			}
-			return stateRoot, nil
-		},
-	)
+func (epp *EigenPodProofs) ComputeBeaconStateRoot(state *spec.VersionedBeaconState) (phase0.Root, error) {
+
+	var beaconStateRoot phase0.Root
+	var err error
+	switch state.Version {
+	case spec.DataVersionElectra:
+		beaconState := state.Electra
+		beaconStateRoot, err = epp.loadOrComputeBeaconStateRoot(
+			beaconState.Slot,
+			func() (phase0.Root, error) {
+				stateRoot, err := beaconState.HashTreeRoot()
+				if err != nil {
+					return phase0.Root{}, err
+				}
+				return stateRoot, nil
+			},
+		)
+	case spec.DataVersionDeneb:
+		beaconState := state.Deneb
+		beaconStateRoot, err = epp.loadOrComputeBeaconStateRoot(
+			beaconState.Slot,
+			func() (phase0.Root, error) {
+				stateRoot, err := beaconState.HashTreeRoot()
+				if err != nil {
+					return phase0.Root{}, err
+				}
+				return stateRoot, nil
+			},
+		)
+	default:
+		return phase0.Root{}, errors.New("unsupported beacon state version")
+	}
+
 	if err != nil {
-		return phase0.Root{}, err
+		return phase0.Root{}, fmt.Errorf("failed to compute beacon state root: %w", err)
 	}
 
 	return beaconStateRoot, nil

onchain_test.go

@@ -16,7 +16,7 @@ func TestValidatorContainersProofOnChain(t *testing.T) {
 	}
 
 	validatorIndices := []uint64{}
-	for i := int(0); i < len(validators); i += 100000 {
+	for i := int(0); i < len(validators); i += 100000000 {
 		validatorIndices = append(validatorIndices, uint64(i))
 	}
 
@@ -65,7 +65,7 @@ func TestValidatorBalancesProofOnChain(t *testing.T) {
 	}
 
 	validatorIndices := []uint64{}
-	for i := int(0); i < len(validators); i += 100000 {
+	for i := int(0); i < len(validators); i += 100000000 {
 		validatorIndices = append(validatorIndices, uint64(i))
 	}
 

prove_validator_test.go

@@ -6,6 +6,7 @@ import (
 	eigenpodproofs "github.com/Layr-Labs/eigenpod-proofs-generation"
 	"github.com/Layr-Labs/eigenpod-proofs-generation/beacon"
 	"github.com/Layr-Labs/eigenpod-proofs-generation/common"
+	"github.com/attestantio/go-eth2-client/spec"
 	"github.com/attestantio/go-eth2-client/spec/deneb"
 	"github.com/attestantio/go-eth2-client/spec/phase0"
 	"github.com/stretchr/testify/assert"
@@ -27,10 +28,10 @@ func TestProveValidatorContainers(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	assert.True(t, verifyStateRootAgainstBlockHeader(t, epp, beaconHeader, beaconState.Deneb, verifyValidatorFieldsCallParams.StateRootProof.Proof))
+	assert.True(t, verifyStateRootAgainstBlockHeader(t, epp, beaconHeader, beaconState, verifyValidatorFieldsCallParams.StateRootProof.Proof))
 
 	for i := 0; i < len(verifyValidatorFieldsCallParams.ValidatorFields); i++ {
-		assert.True(t, verifyValidatorAgainstBeaconState(t, epp, beaconState.Deneb, verifyValidatorFieldsCallParams.ValidatorFieldsProofs[i], validatorIndices[i]))
+		assert.True(t, verifyValidatorAgainstBeaconState(t, epp, beaconState, verifyValidatorFieldsCallParams.ValidatorFieldsProofs[i], validatorIndices[i]))
 	}
 }
 
@@ -57,7 +58,7 @@ func TestProveValidatorBalances(t *testing.T) {
 	}
 }
 
-func verifyStateRootAgainstBlockHeader(t *testing.T, epp *eigenpodproofs.EigenPodProofs, oracleBlockHeader *phase0.BeaconBlockHeader, oracleState *deneb.BeaconState, proof common.Proof) bool {
+func verifyStateRootAgainstBlockHeader(t *testing.T, epp *eigenpodproofs.EigenPodProofs, oracleBlockHeader *phase0.BeaconBlockHeader, oracleState *spec.VersionedBeaconState, proof common.Proof) bool {
 	root, err := oracleBlockHeader.HashTreeRoot()
 	if err != nil {
 		t.Fatal(err)
@@ -71,8 +72,18 @@ func verifyStateRootAgainstBlockHeader(t *testing.T, epp *eigenpodproofs.EigenPo
 	return common.ValidateProof(root, proof, leaf, beacon.STATE_ROOT_INDEX)
 }
 
-func verifyValidatorAgainstBeaconState(t *testing.T, epp *eigenpodproofs.EigenPodProofs, oracleState *deneb.BeaconState, proof common.Proof, validatorIndex uint64) bool {
-	leaf, err := oracleState.Validators[validatorIndex].HashTreeRoot()
+func verifyValidatorAgainstBeaconState(t *testing.T, epp *eigenpodproofs.EigenPodProofs, oracleState *spec.VersionedBeaconState, proof common.Proof, validatorIndex uint64) bool {
+	var leaf phase0.Root
+	var err error
+	switch oracleState.Version {
+	case spec.DataVersionElectra:
+		leaf, err = oracleState.Electra.Validators[validatorIndex].HashTreeRoot()
+	case spec.DataVersionDeneb:
+		leaf, err = oracleState.Deneb.Validators[validatorIndex].HashTreeRoot()
+	default:
+		t.Fatal("unsupported beacon state version")
+	}
+
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -91,8 +102,18 @@ func verifyValidatorBalancesRootAgainstBlockHeader(t *testing.T, epp *eigenpodpr
 	if err != nil {
 		t.Fatal(err)
 	}
-	/// TODO: update for pectra
-	return common.ValidateProof(root, proof.Proof, proof.ValidatorBalancesRoot, beacon.STATE_ROOT_INDEX<<beacon.BEACON_STATE_TREE_HEIGHT_DENEB|beacon.BALANCES_INDEX)
+
+	var beaconStateTreeHeight uint64
+	switch beaconState.Version {
+	case spec.DataVersionElectra:
+		beaconStateTreeHeight = beacon.BEACON_STATE_TREE_HEIGHT_ELECTRA
+	case spec.DataVersionDeneb:
+		beaconStateTreeHeight = beacon.BEACON_STATE_TREE_HEIGHT_DENEB
+	default:
+		t.Fatal("unsupported beacon state version")
+	}
+
+	return common.ValidateProof(root, proof.Proof, proof.ValidatorBalancesRoot, beacon.STATE_ROOT_INDEX<<beaconStateTreeHeight|beacon.BALANCES_INDEX)
 }
 
 func verifyValidatorBalanceAgainstValidatorBalancesRoot(t *testing.T, epp *eigenpodproofs.EigenPodProofs, oracleState *deneb.BeaconState, validatorBalancesRoot phase0.Root, proof *eigenpodproofs.BalanceProof, validatorIndex uint64) bool {