From a0c07062242a2580fe0425692d5007e0b9dcc700 Mon Sep 17 00:00:00 2001
From: terencechain <terence@prysmaticlabs.com>
Date: Thu, 24 Nov 2022 09:54:55 -0800
Subject: [PATCH] Add Capella state changes (#11688)

* Add Capella state changes

* Use params.configs
---
 beacon-chain/state/state-native/hasher.go          |  6 +++---
 .../state/state-native/setters_withdrawal.go       |  3 +++
 .../state/state-native/setters_withdrawal_test.go  |  6 ++++++
 beacon-chain/state/types/types.go                  | 14 ++++++++++++--
 4 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/beacon-chain/state/state-native/hasher.go b/beacon-chain/state/state-native/hasher.go
index 0377e1b65..ce5ece6eb 100644
--- a/beacon-chain/state/state-native/hasher.go
+++ b/beacon-chain/state/state-native/hasher.go
@@ -252,9 +252,9 @@ func ComputeFieldRootsWithHasher(ctx context.Context, state *BeaconState) ([][]b
 		fieldRoots[nativetypes.NextWithdrawalIndex.RealPosition()] = nextWithdrawalIndexRoot
 
 		// Next partial withdrawal validator index root.
-		nextPartialWithdrawalValidatorIndexRoot := make([]byte, 32)
-		binary.LittleEndian.PutUint64(nextPartialWithdrawalValidatorIndexRoot, uint64(state.nextWithdrawalValidatorIndex))
-		fieldRoots[nativetypes.NextWithdrawalValidatorIndex.RealPosition()] = nextPartialWithdrawalValidatorIndexRoot
+		nextWithdrawalValidatorIndexRoot := make([]byte, 32)
+		binary.LittleEndian.PutUint64(nextWithdrawalValidatorIndexRoot, uint64(state.nextWithdrawalValidatorIndex))
+		fieldRoots[nativetypes.NextWithdrawalValidatorIndex.RealPosition()] = nextWithdrawalValidatorIndexRoot
 	}
 
 	return fieldRoots, nil
diff --git a/beacon-chain/state/state-native/setters_withdrawal.go b/beacon-chain/state/state-native/setters_withdrawal.go
index bb6ee91a6..53be0e7ae 100644
--- a/beacon-chain/state/state-native/setters_withdrawal.go
+++ b/beacon-chain/state/state-native/setters_withdrawal.go
@@ -1,6 +1,7 @@
 package state_native
 
 import (
+	nativetypes "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/state-native/types"
 	types "github.com/prysmaticlabs/prysm/v3/consensus-types/primitives"
 	"github.com/prysmaticlabs/prysm/v3/runtime/version"
 )
@@ -15,6 +16,7 @@ func (b *BeaconState) SetNextWithdrawalIndex(i uint64) error {
 	defer b.lock.Unlock()
 
 	b.nextWithdrawalIndex = i
+	b.markFieldAsDirty(nativetypes.NextWithdrawalIndex)
 	return nil
 }
 
@@ -29,5 +31,6 @@ func (b *BeaconState) SetNextWithdrawalValidatorIndex(i types.ValidatorIndex) er
 	defer b.lock.Unlock()
 
 	b.nextWithdrawalValidatorIndex = i
+	b.markFieldAsDirty(nativetypes.NextWithdrawalValidatorIndex)
 	return nil
 }
diff --git a/beacon-chain/state/state-native/setters_withdrawal_test.go b/beacon-chain/state/state-native/setters_withdrawal_test.go
index 6cb28a7b5..e40c628e5 100644
--- a/beacon-chain/state/state-native/setters_withdrawal_test.go
+++ b/beacon-chain/state/state-native/setters_withdrawal_test.go
@@ -3,6 +3,7 @@ package state_native
 import (
 	"testing"
 
+	nativetypes "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/state-native/types"
 	types "github.com/prysmaticlabs/prysm/v3/consensus-types/primitives"
 	"github.com/prysmaticlabs/prysm/v3/runtime/version"
 	"github.com/prysmaticlabs/prysm/v3/testing/require"
@@ -12,15 +13,20 @@ func TestSetNextWithdrawalIndex(t *testing.T) {
 	s := BeaconState{
 		version:             version.Capella,
 		nextWithdrawalIndex: 3,
+		dirtyFields:         make(map[nativetypes.FieldIndex]bool),
 	}
 	require.NoError(t, s.SetNextWithdrawalIndex(5))
 	require.Equal(t, uint64(5), s.nextWithdrawalIndex)
+	require.Equal(t, true, s.dirtyFields[nativetypes.NextWithdrawalIndex])
 }
+
 func TestSetLastWithdrawalValidatorIndex(t *testing.T) {
 	s := BeaconState{
 		version:                      version.Capella,
 		nextWithdrawalValidatorIndex: 3,
+		dirtyFields:                  make(map[nativetypes.FieldIndex]bool),
 	}
 	require.NoError(t, s.SetNextWithdrawalValidatorIndex(5))
 	require.Equal(t, types.ValidatorIndex(5), s.nextWithdrawalValidatorIndex)
+	require.Equal(t, true, s.dirtyFields[nativetypes.NextWithdrawalValidatorIndex])
 }
diff --git a/beacon-chain/state/types/types.go b/beacon-chain/state/types/types.go
index 3fc4a48b0..51cbb0ef8 100644
--- a/beacon-chain/state/types/types.go
+++ b/beacon-chain/state/types/types.go
@@ -67,12 +67,12 @@ func (f FieldIndex) String(stateVersion int) string {
 	case Slashings:
 		return "slashings"
 	case PreviousEpochAttestations:
-		if version.Altair == stateVersion || version.Bellatrix == stateVersion {
+		if stateVersion > version.Phase0 {
 			return "previousEpochParticipationBits"
 		}
 		return "previousEpochAttestations"
 	case CurrentEpochAttestations:
-		if version.Altair == stateVersion || version.Bellatrix == stateVersion {
+		if stateVersion > version.Phase0 {
 			return "currentEpochParticipationBits"
 		}
 		return "currentEpochAttestations"
@@ -92,6 +92,12 @@ func (f FieldIndex) String(stateVersion int) string {
 		return "nextSyncCommittee"
 	case LatestExecutionPayloadHeader:
 		return "latestExecutionPayloadHeader"
+	case LatestExecutionPayloadHeaderCapella:
+		return "latestExecutionPayloadHeaderCapella"
+	case NextWithdrawalIndex:
+		return "nextWithdrawalIndex"
+	case NextWithdrawalValidatorIndex:
+		return "nextWithdrawalValidatorIndex"
 	default:
 		return ""
 	}
@@ -152,6 +158,10 @@ const (
 	NextSyncCommittee
 	// State fields added in Bellatrix.
 	LatestExecutionPayloadHeader
+	// State fields added in Capella
+	LatestExecutionPayloadHeaderCapella
+	NextWithdrawalIndex
+	NextWithdrawalValidatorIndex
 )
 
 // Altair fields which replaced previous phase 0 fields.