State: HasPendingBalanceToWithdraw (#14200)

* Implement HasPendingBalanceToWithdraw to improve the best / average case lookup

* Add tests for HasPendingBalanceToWithdraw

beacon-chain/core/blocks/exit.go

@@ -193,12 +193,12 @@ func verifyExitConditions(st state.ReadOnlyBeaconState, validator state.ReadOnly
 
 	if st.Version() >= version.Electra {
 		// Only exit validator if it has no pending withdrawals in the queue.
-		pbw, err := st.PendingBalanceToWithdraw(exit.ValidatorIndex)
+		ok, err := st.HasPendingBalanceToWithdraw(exit.ValidatorIndex)
 		if err != nil {
 			return fmt.Errorf("unable to retrieve pending balance to withdraw for validator %d: %w", exit.ValidatorIndex, err)
 		}
-		if pbw != 0 {
-			return fmt.Errorf("validator %d must have no pending balance to withdraw, got %d pending balance to withdraw", exit.ValidatorIndex, pbw)
+		if ok {
+			return fmt.Errorf("validator %d must have no pending balance to withdraw", exit.ValidatorIndex)
 		}
 	}
 

beacon-chain/state/interfaces.go

@@ -200,6 +200,7 @@ type ReadOnlyWithdrawals interface {
 	NextWithdrawalIndex() (uint64, error)
 	PendingBalanceToWithdraw(idx primitives.ValidatorIndex) (uint64, error)
 	NumPendingPartialWithdrawals() (uint64, error)
+	HasPendingBalanceToWithdraw(idx primitives.ValidatorIndex) (bool, error)
 }
 
 // ReadOnlyParticipation defines a struct which only has read access to participation methods.

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

@@ -501,3 +501,24 @@ func (b *BeaconState) PendingBalanceToWithdraw(idx primitives.ValidatorIndex) (u
 	}
 	return sum, nil
 }
+
+func (b *BeaconState) HasPendingBalanceToWithdraw(idx primitives.ValidatorIndex) (bool, error) {
+	if b.version < version.Electra {
+		return false, errNotSupported("HasPendingBalanceToWithdraw", b.version)
+	}
+
+	b.lock.RLock()
+	defer b.lock.RUnlock()
+
+	// TODO: Consider maintaining this value in the state, if it's a potential bottleneck.
+	// This is n*m complexity, but this method can only be called
+	// MAX_WITHDRAWAL_REQUESTS_PER_PAYLOAD per slot. A more optimized storage indexing such as a
+	// lookup map could be used to reduce the complexity marginally.
+	for _, w := range b.pendingPartialWithdrawals {
+		if w.Index == idx {
+			return true, nil
+		}
+	}
+
+	return false, nil
+}

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

@@ -162,3 +162,32 @@ func TestAggregateKeyFromIndices(t *testing.T) {
 
 	assert.Equal(t, true, aggKey.Equals(retKey), "unequal aggregated keys")
 }
+
+func TestHasPendingBalanceToWithdraw(t *testing.T) {
+	pb := &ethpb.BeaconStateElectra{
+		PendingPartialWithdrawals: []*ethpb.PendingPartialWithdrawal{
+			{
+				Amount: 100,
+				Index:  1,
+			},
+			{
+				Amount: 200,
+				Index:  2,
+			},
+			{
+				Amount: 300,
+				Index:  3,
+			},
+		},
+	}
+	state, err := statenative.InitializeFromProtoUnsafeElectra(pb)
+	require.NoError(t, err)
+
+	ok, err := state.HasPendingBalanceToWithdraw(1)
+	require.NoError(t, err)
+	require.Equal(t, true, ok)
+
+	ok, err = state.HasPendingBalanceToWithdraw(5)
+	require.NoError(t, err)
+	require.Equal(t, false, ok)
+}