EIP6110: add validator index cache (#13943)

* EIP6110: add validator index cache

* Add tests

* Radek's feedback

beacon-chain/state/interfaces.go

@@ -100,6 +100,7 @@ type WriteOnlyBeaconState interface {
 	SetLatestExecutionPayloadHeader(payload interfaces.ExecutionData) error
 	SetNextWithdrawalIndex(i uint64) error
 	SetNextWithdrawalValidatorIndex(i primitives.ValidatorIndex) error
+	SaveValidatorIndices()
 }
 
 // ReadOnlyValidator defines a struct which only has read access to validator methods.

beacon-chain/state/state-native/BUILD.bazel

@@ -37,6 +37,7 @@ go_library(
         "ssz.go",
         "state_trie.go",
         "types.go",
+        "validator_index_cache.go",
     ] + select({
         "//config:mainnet": ["beacon_state_mainnet.go"],
         "//config:minimal": ["beacon_state_minimal.go"],
@@ -102,6 +103,7 @@ go_test(
         "state_test.go",
         "state_trie_test.go",
         "types_test.go",
+        "validator_index_cache_test.go",
     ],
     data = glob(["testdata/**"]),
     embed = [":go_default_library"],

beacon-chain/state/state-native/beacon_state_mainnet.go

@@ -67,6 +67,7 @@ type BeaconState struct {
 	stateFieldLeaves      map[types.FieldIndex]*fieldtrie.FieldTrie
 	rebuildTrie           map[types.FieldIndex]bool
 	valMapHandler         *stateutil.ValidatorMapHandler
+	validatorIndexCache   *finalizedValidatorIndexCache
 	merkleLayers          [][][]byte
 	sharedFieldReferences map[types.FieldIndex]*stateutil.Reference
 }

beacon-chain/state/state-native/beacon_state_minimal.go

@@ -67,6 +67,7 @@ type BeaconState struct {
 	stateFieldLeaves      map[types.FieldIndex]*fieldtrie.FieldTrie
 	rebuildTrie           map[types.FieldIndex]bool
 	valMapHandler         *stateutil.ValidatorMapHandler
+	validatorIndexCache   *finalizedValidatorIndexCache
 	merkleLayers          [][][]byte
 	sharedFieldReferences map[types.FieldIndex]*stateutil.Reference
 }

beacon-chain/state/state-native/getters_validator.go

@@ -178,6 +178,10 @@ func (b *BeaconState) ValidatorIndexByPubkey(key [fieldparams.BLSPubkeyLength]by
 	b.lock.RLock()
 	defer b.lock.RUnlock()
 
+	if features.Get().EIP6110ValidatorIndexCache {
+		return b.getValidatorIndex(key)
+	}
+
 	var numOfVals int
 	if features.Get().EnableExperimentalState {
 		numOfVals = b.validatorsMultiValue.Len(b)

beacon-chain/state/state-native/readonly_validator.go

@@ -70,6 +70,11 @@ func (v readOnlyValidator) PublicKey() [fieldparams.BLSPubkeyLength]byte {
 	return pubkey
 }
 
+// publicKeySlice returns the public key in the slice form for the read only validator.
+func (v readOnlyValidator) publicKeySlice() []byte {
+	return v.validator.PublicKey
+}
+
 // WithdrawalCredentials returns the withdrawal credentials of the
 // read only validator.
 func (v readOnlyValidator) WithdrawalCredentials() []byte {

beacon-chain/state/state-native/setters_misc.go

@@ -4,6 +4,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/state/state-native/types"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/state/stateutil"
+	"github.com/prysmaticlabs/prysm/v5/config/features"
 	fieldparams "github.com/prysmaticlabs/prysm/v5/config/fieldparams"
 	"github.com/prysmaticlabs/prysm/v5/consensus-types/primitives"
 	"github.com/prysmaticlabs/prysm/v5/container/slice"
@@ -107,6 +108,18 @@ func (b *BeaconState) SetHistoricalRoots(val [][]byte) error {
 	return nil
 }
 
+// SaveValidatorIndices save validator indices of beacon chain to cache
+func (b *BeaconState) SaveValidatorIndices() {
+	if !features.Get().EIP6110ValidatorIndexCache {
+		return
+	}
+
+	b.lock.Lock()
+	defer b.lock.Unlock()
+
+	b.saveValidatorIndices()
+}
+
 // AppendHistoricalRoots for the beacon state. Appends the new value
 // to the end of list.
 func (b *BeaconState) AppendHistoricalRoots(root [32]byte) error {

beacon-chain/state/state-native/state_trie.go

@@ -759,7 +759,8 @@ func (b *BeaconState) Copy() state.BeaconState {
 		stateFieldLeaves: make(map[types.FieldIndex]*fieldtrie.FieldTrie, fieldCount),
 
 		// Share the reference to validator index map.
-		valMapHandler: b.valMapHandler,
+		valMapHandler:       b.valMapHandler,
+		validatorIndexCache: b.validatorIndexCache,
 	}
 
 	if features.Get().EnableExperimentalState {

beacon-chain/state/state-native/validator_index_cache.go

@@ -0,0 +1,96 @@
+package state_native
+
+import (
+	"bytes"
+	"sync"
+
+	"github.com/prysmaticlabs/prysm/v5/config/features"
+	fieldparams "github.com/prysmaticlabs/prysm/v5/config/fieldparams"
+	"github.com/prysmaticlabs/prysm/v5/consensus-types/primitives"
+	"github.com/prysmaticlabs/prysm/v5/encoding/bytesutil"
+	ethpb "github.com/prysmaticlabs/prysm/v5/proto/prysm/v1alpha1"
+)
+
+// finalizedValidatorIndexCache maintains a mapping from validator public keys to their indices within the beacon state.
+// It includes a lastFinalizedIndex to track updates up to the last finalized validator index,
+// and uses a mutex for concurrent read/write access to the cache.
+type finalizedValidatorIndexCache struct {
+	indexMap map[[fieldparams.BLSPubkeyLength]byte]primitives.ValidatorIndex // Maps finalized BLS public keys to validator indices.
+	sync.RWMutex
+}
+
+// newFinalizedValidatorIndexCache initializes a new validator index cache with an empty index map.
+func newFinalizedValidatorIndexCache() *finalizedValidatorIndexCache {
+	return &finalizedValidatorIndexCache{
+		indexMap: make(map[[fieldparams.BLSPubkeyLength]byte]primitives.ValidatorIndex),
+	}
+}
+
+// getValidatorIndex retrieves the validator index for a given public key from the cache.
+// If the public key is not found in the cache, it searches through the state starting from the last finalized index.
+func (b *BeaconState) getValidatorIndex(pubKey [fieldparams.BLSPubkeyLength]byte) (primitives.ValidatorIndex, bool) {
+	b.validatorIndexCache.RLock()
+	index, found := b.validatorIndexCache.indexMap[pubKey]
+	b.validatorIndexCache.RUnlock()
+	if found {
+		return index, true
+	}
+
+	validatorCount := len(b.validatorIndexCache.indexMap)
+	vals := b.validatorsReadOnlySinceIndex(validatorCount)
+	for i, val := range vals {
+		if bytes.Equal(bytesutil.PadTo(val.publicKeySlice(), 48), pubKey[:]) {
+			index := primitives.ValidatorIndex(validatorCount + i)
+			return index, true
+		}
+	}
+	return 0, false
+}
+
+// saveValidatorIndices updates the validator index cache with new indices.
+// It processes validator indices starting after the last finalized index and updates the tracker.
+func (b *BeaconState) saveValidatorIndices() {
+	b.validatorIndexCache.Lock()
+	defer b.validatorIndexCache.Unlock()
+
+	validatorCount := len(b.validatorIndexCache.indexMap)
+	vals := b.validatorsReadOnlySinceIndex(validatorCount)
+	for i, val := range vals {
+		b.validatorIndexCache.indexMap[val.PublicKey()] = primitives.ValidatorIndex(validatorCount + i)
+	}
+}
+
+// validatorsReadOnlySinceIndex constructs a list of read only validator references after a specified index.
+// The indices in the returned list correspond to their respective validator indices in the state.
+// It returns nil if the specified index is out of bounds. This function is read-only and does not use locks.
+func (b *BeaconState) validatorsReadOnlySinceIndex(index int) []readOnlyValidator {
+	totalValidators := b.validatorsLen()
+	if index >= totalValidators {
+		return nil
+	}
+
+	var v []*ethpb.Validator
+	if features.Get().EnableExperimentalState {
+		if b.validatorsMultiValue == nil {
+			return nil
+		}
+		v = b.validatorsMultiValue.Value(b)
+	} else {
+		if b.validators == nil {
+			return nil
+		}
+		v = b.validators
+	}
+
+	result := make([]readOnlyValidator, totalValidators-index)
+	for i := 0; i < len(result); i++ {
+		val := v[i+index]
+		if val == nil {
+			continue
+		}
+		result[i] = readOnlyValidator{
+			validator: val,
+		}
+	}
+	return result
+}

beacon-chain/state/state-native/validator_index_cache_test.go

@@ -0,0 +1,105 @@
+package state_native
+
+import (
+	"testing"
+
+	fieldparams "github.com/prysmaticlabs/prysm/v5/config/fieldparams"
+	"github.com/prysmaticlabs/prysm/v5/consensus-types/primitives"
+	ethpb "github.com/prysmaticlabs/prysm/v5/proto/prysm/v1alpha1"
+	"github.com/prysmaticlabs/prysm/v5/testing/require"
+)
+
+func Test_FinalizedValidatorIndexCache(t *testing.T) {
+	c := newFinalizedValidatorIndexCache()
+	b := &BeaconState{validatorIndexCache: c}
+
+	// What happens if you call getValidatorIndex with a public key that is not in the cache and state?
+	// The function will return 0 and false.
+	i, exists := b.getValidatorIndex([fieldparams.BLSPubkeyLength]byte{0})
+	require.Equal(t, primitives.ValidatorIndex(0), i)
+	require.Equal(t, false, exists)
+
+	// Validators are added to the state. They are [0, 1, 2]
+	b.validators = []*ethpb.Validator{
+		{PublicKey: []byte{1}},
+		{PublicKey: []byte{2}},
+		{PublicKey: []byte{3}},
+	}
+	// We should be able to retrieve these validators by public key even when they are not in the cache
+	i, exists = b.getValidatorIndex([fieldparams.BLSPubkeyLength]byte{1})
+	require.Equal(t, primitives.ValidatorIndex(0), i)
+	require.Equal(t, true, exists)
+	i, exists = b.getValidatorIndex([fieldparams.BLSPubkeyLength]byte{2})
+	require.Equal(t, primitives.ValidatorIndex(1), i)
+	require.Equal(t, true, exists)
+	i, exists = b.getValidatorIndex([fieldparams.BLSPubkeyLength]byte{3})
+	require.Equal(t, primitives.ValidatorIndex(2), i)
+	require.Equal(t, true, exists)
+
+	// State is finalized. We save [0, 1, 2 ] to the cache.
+	b.saveValidatorIndices()
+	require.Equal(t, 3, len(b.validatorIndexCache.indexMap))
+	i, exists = b.getValidatorIndex([fieldparams.BLSPubkeyLength]byte{1})
+	require.Equal(t, primitives.ValidatorIndex(0), i)
+	require.Equal(t, true, exists)
+	i, exists = b.getValidatorIndex([fieldparams.BLSPubkeyLength]byte{2})
+	require.Equal(t, primitives.ValidatorIndex(1), i)
+	require.Equal(t, true, exists)
+	i, exists = b.getValidatorIndex([fieldparams.BLSPubkeyLength]byte{3})
+	require.Equal(t, primitives.ValidatorIndex(2), i)
+	require.Equal(t, true, exists)
+
+	// New validators are added to the state. They are [4, 5]
+	b.validators = []*ethpb.Validator{
+		{PublicKey: []byte{1}},
+		{PublicKey: []byte{2}},
+		{PublicKey: []byte{3}},
+		{PublicKey: []byte{4}},
+		{PublicKey: []byte{5}},
+	}
+	// We should be able to retrieve these validators by public key even when they are not in the cache
+	i, exists = b.getValidatorIndex([fieldparams.BLSPubkeyLength]byte{4})
+	require.Equal(t, primitives.ValidatorIndex(3), i)
+	require.Equal(t, true, exists)
+	i, exists = b.getValidatorIndex([fieldparams.BLSPubkeyLength]byte{5})
+	require.Equal(t, primitives.ValidatorIndex(4), i)
+	require.Equal(t, true, exists)
+
+	// State is finalized. We save [4, 5] to the cache.
+	b.saveValidatorIndices()
+	require.Equal(t, 5, len(b.validatorIndexCache.indexMap))
+
+	// New validators are added to the state. They are [6]
+	b.validators = []*ethpb.Validator{
+		{PublicKey: []byte{1}},
+		{PublicKey: []byte{2}},
+		{PublicKey: []byte{3}},
+		{PublicKey: []byte{4}},
+		{PublicKey: []byte{5}},
+		{PublicKey: []byte{6}},
+	}
+	// We should be able to retrieve these validators by public key even when they are not in the cache
+	i, exists = b.getValidatorIndex([fieldparams.BLSPubkeyLength]byte{6})
+	require.Equal(t, primitives.ValidatorIndex(5), i)
+	require.Equal(t, true, exists)
+
+	// State is finalized. We save [6] to the cache.
+	b.saveValidatorIndices()
+	require.Equal(t, 6, len(b.validatorIndexCache.indexMap))
+
+	// Save a few more times.
+	b.saveValidatorIndices()
+	b.saveValidatorIndices()
+	require.Equal(t, 6, len(b.validatorIndexCache.indexMap))
+
+	// Can still retrieve the validators from the cache
+	i, exists = b.getValidatorIndex([fieldparams.BLSPubkeyLength]byte{1})
+	require.Equal(t, primitives.ValidatorIndex(0), i)
+	require.Equal(t, true, exists)
+	i, exists = b.getValidatorIndex([fieldparams.BLSPubkeyLength]byte{2})
+	require.Equal(t, primitives.ValidatorIndex(1), i)
+	require.Equal(t, true, exists)
+	i, exists = b.getValidatorIndex([fieldparams.BLSPubkeyLength]byte{3})
+	require.Equal(t, primitives.ValidatorIndex(2), i)
+	require.Equal(t, true, exists)
+}