Optimize validator assignment lookups with O(1) reverse index map

beacon-chain/core/helpers/beacon_committee.go

@@ -406,6 +406,29 @@ func AssignmentForValidator(
 	return &LiteAssignment{} // validator is not scheduled this epoch
 }
 
+// BuildValidatorAssignmentMap creates a reverse index map from validator index to assignment
+// for O(1) lookups instead of O(slots * committees * committee_size) linear search.
+func BuildValidatorAssignmentMap(
+	bySlot [][][]primitives.ValidatorIndex,
+	startSlot primitives.Slot,
+) map[primitives.ValidatorIndex]*LiteAssignment {
+	validatorToAssignment := make(map[primitives.ValidatorIndex]*LiteAssignment)
+	
+	for relativeSlot, committees := range bySlot {
+		for cIdx, committee := range committees {
+			for pos, vIdx := range committee {
+				validatorToAssignment[vIdx] = &LiteAssignment{
+					AttesterSlot:            startSlot + primitives.Slot(relativeSlot),
+					CommitteeIndex:          primitives.CommitteeIndex(cIdx),
+					CommitteeLength:         uint64(len(committee)),
+					ValidatorCommitteeIndex: uint64(pos),
+				}
+			}
+		}
+	}
+	return validatorToAssignment
+}
+
 // CommitteeAssignments calculates committee assignments for each validator during the specified epoch.
 // It retrieves active validator indices, determines the number of committees per slot, and computes
 // assignments for each validator based on their presence in the provided validators slice.

beacon-chain/core/helpers/beacon_committee_test.go

@@ -917,6 +917,52 @@ func TestAssignmentForValidator(t *testing.T) {
 	})
 }
 
+func TestBuildValidatorAssignmentMap(t *testing.T) {
+	start := primitives.Slot(200)
+	bySlot := [][][]primitives.ValidatorIndex{
+		{{1, 2, 3}},
+		{{7, 8, 9}},
+	}
+
+	assignmentMap := helpers.BuildValidatorAssignmentMap(bySlot, start)
+
+	// Test validator 8 assignment (should match TestAssignmentForValidator)
+	vIdx := primitives.ValidatorIndex(8)
+	got, exists := assignmentMap[vIdx]
+	assert.Equal(t, true, exists)
+	require.NotNil(t, got)
+	require.Equal(t, start+1, got.AttesterSlot)
+	require.Equal(t, primitives.CommitteeIndex(0), got.CommitteeIndex)
+	require.Equal(t, uint64(3), got.CommitteeLength)
+	require.Equal(t, uint64(1), got.ValidatorCommitteeIndex)
+
+	// Test validator 1 assignment
+	vIdx1 := primitives.ValidatorIndex(1)
+	got1, exists1 := assignmentMap[vIdx1]
+	assert.Equal(t, true, exists1)
+	require.NotNil(t, got1)
+	require.Equal(t, start, got1.AttesterSlot)
+	require.Equal(t, primitives.CommitteeIndex(0), got1.CommitteeIndex)
+	require.Equal(t, uint64(3), got1.CommitteeLength)
+	require.Equal(t, uint64(0), got1.ValidatorCommitteeIndex)
+
+	// Test non-existent validator
+	_, exists99 := assignmentMap[primitives.ValidatorIndex(99)]
+	assert.Equal(t, false, exists99)
+
+	// Verify that we get the same results as the linear search
+	for _, committees := range bySlot {
+		for _, committee := range committees {
+			for _, validatorIdx := range committee {
+				linearResult := helpers.AssignmentForValidator(bySlot, start, validatorIdx)
+				mapResult, mapExists := assignmentMap[validatorIdx]
+				assert.Equal(t, true, mapExists)
+				require.DeepEqual(t, linearResult, mapResult)
+			}
+		}
+	}
+}
+
 // Regression for #15450
 func TestInitializeProposerLookahead_RegressionTest(t *testing.T) {
 	ctx := t.Context()

beacon-chain/rpc/prysm/v1alpha1/validator/duties_v2.go

@@ -12,6 +12,7 @@ import (
 	"github.com/OffchainLabs/prysm/v6/encoding/bytesutil"
 	"github.com/OffchainLabs/prysm/v6/monitoring/tracing/trace"
 	ethpb "github.com/OffchainLabs/prysm/v6/proto/prysm/v1alpha1"
+	"github.com/OffchainLabs/prysm/v6/time"
 	"github.com/OffchainLabs/prysm/v6/time/slots"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -82,13 +83,30 @@ func (vs *Server) dutiesv2(ctx context.Context, req *ethpb.DutiesRequest) (*ethp
 			continue
 		}
 
-		meta.current.liteAssignment = helpers.AssignmentForValidator(meta.current.committeesBySlot, meta.current.startSlot, idx)
-		meta.next.liteAssignment = helpers.AssignmentForValidator(meta.next.committeesBySlot, meta.next.startSlot, idx)
+		timeNow := time.Now()
+		// O(1) lookup instead of O(slots * committees * committee_size) linear search
+		if assignment, exists := meta.current.validatorAssignmentMap[idx]; exists {
+			meta.current.liteAssignment = assignment
+		} else {
+			meta.current.liteAssignment = &helpers.LiteAssignment{}
+		}
+		log.Info("Time took to compute current epoch assignment", time.Since(timeNow))
 
+		timeNow = time.Now()
+		// O(1) lookup instead of O(slots * committees * committee_size) linear search
+		if assignment, exists := meta.next.validatorAssignmentMap[idx]; exists {
+			meta.next.liteAssignment = assignment
+		} else {
+			meta.next.liteAssignment = &helpers.LiteAssignment{}
+		}
+		log.Info("Time took to compute next epoch assignment", time.Since(timeNow))
+
+		timeNow = time.Now()
 		assignment, nextAssignment, err := vs.buildValidatorDuty(pubKey, idx, s, req.Epoch, meta)
 		if err != nil {
 			return nil, err
 		}
+		log.Info("Time took to build validator duty", time.Since(timeNow))
 		validatorAssignments = append(validatorAssignments, assignment)
 		nextValidatorAssignments = append(nextValidatorAssignments, nextAssignment)
 	}
@@ -143,11 +161,12 @@ type dutiesMetadata struct {
 }
 
 type metadata struct {
-	committeesAtSlot uint64
-	proposalSlots    map[primitives.ValidatorIndex][]primitives.Slot
-	startSlot        primitives.Slot
-	committeesBySlot [][][]primitives.ValidatorIndex
-	liteAssignment   *helpers.LiteAssignment
+	committeesAtSlot       uint64
+	proposalSlots          map[primitives.ValidatorIndex][]primitives.Slot
+	startSlot              primitives.Slot
+	committeesBySlot       [][][]primitives.ValidatorIndex
+	validatorAssignmentMap map[primitives.ValidatorIndex]*helpers.LiteAssignment
+	liteAssignment         *helpers.LiteAssignment // Deprecated: use validatorAssignmentMap for O(1) lookups
 }
 
 func loadDutiesMetadata(ctx context.Context, s state.BeaconState, reqEpoch primitives.Epoch) (*dutiesMetadata, error) {
@@ -193,6 +212,9 @@ func loadMetadata(ctx context.Context, s state.BeaconState, reqEpoch primitives.
 		return nil, err
 	}
 
+	// Build reverse index map for O(1) validator assignment lookups
+	meta.validatorAssignmentMap = helpers.BuildValidatorAssignmentMap(meta.committeesBySlot, meta.startSlot)
+
 	return meta, nil
 }
 
@@ -208,19 +230,24 @@ func (vs *Server) buildValidatorDuty(
 	assignment := &ethpb.DutiesV2Response_Duty{PublicKey: pubKey}
 	nextAssignment := &ethpb.DutiesV2Response_Duty{PublicKey: pubKey}
 
+	timeNow := time.Now()
 	statusEnum := assignmentStatus(s, idx)
 	assignment.ValidatorIndex = idx
 	assignment.Status = statusEnum
 	assignment.CommitteesAtSlot = meta.current.committeesAtSlot
 	assignment.ProposerSlots = meta.current.proposalSlots[idx]
 	populateCommitteeFields(assignment, meta.current.liteAssignment)
+	log.Info("Time took to compute current epoch assignment status", time.Since(timeNow))
 
+	timeNow = time.Now()
 	nextAssignment.ValidatorIndex = idx
 	nextAssignment.Status = statusEnum
 	nextAssignment.CommitteesAtSlot = meta.next.committeesAtSlot
 	populateCommitteeFields(nextAssignment, meta.next.liteAssignment)
+	log.Info("Time took to compute next epoch assignment status", time.Since(timeNow))
 
 	// Sync committee flags
+	timeNow = time.Now()
 	if coreTime.HigherEqualThanAltairVersionAndEpoch(s, reqEpoch) {
 		inSync, err := helpers.IsCurrentPeriodSyncCommittee(s, idx)
 		if err != nil {
@@ -256,6 +283,7 @@ func (vs *Server) buildValidatorDuty(
 			}
 		}
 	}
+	log.Info("Time took to compute sync committee flags", time.Since(timeNow))
 
 	return assignment, nextAssignment, nil
 }