Rewrite CommitteeAssignments to not return every validator

beacon-chain/core/helpers/beacon_committee.go

@@ -143,105 +143,138 @@ func BeaconCommittee(
 	return ComputeCommittee(validatorIndices, seed, indexOffset, count)
 }
 
-// CommitteeAssignmentContainer represents a committee list, committee index, and to be attested slot for a given epoch.
-type CommitteeAssignmentContainer struct {
+// CommitteeAssignment represents committee list, committee index, and to be attested slot for a given epoch.
+type CommitteeAssignment struct {
 	Committee      []primitives.ValidatorIndex
 	AttesterSlot   primitives.Slot
 	CommitteeIndex primitives.CommitteeIndex
 }
 
-// CommitteeAssignments is a map of validator indices pointing to the appropriate committee
-// assignment for the given epoch.
-//
-// 1. Determine the proposer validator index for each slot.
-// 2. Compute all committees.
-// 3. Determine the attesting slot for each committee.
-// 4. Construct a map of validator indices pointing to the respective committees.
-func CommitteeAssignments(
-	ctx context.Context,
-	state state.BeaconState,
-	epoch primitives.Epoch,
-) (map[primitives.ValidatorIndex]*CommitteeAssignmentContainer, map[primitives.ValidatorIndex][]primitives.Slot, error) {
+// verifyAssignmentEpoch verifies if the given epoch is valid for assignment based on the provided state.
+// It checks if the epoch is not greater than the next epoch, and if the start slot of the epoch is greater
+// than or equal to the minimum valid start slot calculated based on the state's current slot and historical roots.
+func verifyAssignmentEpoch(epoch primitives.Epoch, state state.BeaconState) error {
 	nextEpoch := time.NextEpoch(state)
 	if epoch > nextEpoch {
-		return nil, nil, fmt.Errorf(
-			"epoch %d can't be greater than next epoch %d",
-			epoch,
-			nextEpoch,
-		)
+		return fmt.Errorf("epoch %d can't be greater than next epoch %d", epoch, nextEpoch)
 	}
 
-	// We determine the slots in which proposers are supposed to act.
-	// Some validators may need to propose multiple times per epoch, so
-	// we use a map of proposer idx -> []slot to keep track of this possibility.
 	startSlot, err := slots.EpochStart(epoch)
 	if err != nil {
-		return nil, nil, err
+		return err
 	}
 	minValidStartSlot := primitives.Slot(0)
-	if state.Slot() >= params.BeaconConfig().SlotsPerHistoricalRoot {
-		minValidStartSlot = state.Slot() - params.BeaconConfig().SlotsPerHistoricalRoot
+	if stateSlot := state.Slot(); stateSlot >= params.BeaconConfig().SlotsPerHistoricalRoot {
+		minValidStartSlot = stateSlot - params.BeaconConfig().SlotsPerHistoricalRoot
 	}
 	if startSlot < minValidStartSlot {
-		return nil, nil, fmt.Errorf("start slot %d is smaller than the minimum valid start slot %d", startSlot, minValidStartSlot)
+		return fmt.Errorf("start slot %d is smaller than the minimum valid start slot %d", startSlot, minValidStartSlot)
+	}
+	return nil
+}
+
+// ProposerAssignments calculates proposer assignments for each validator during the specified epoch.
+// It verifies the validity of the epoch, then iterates through each slot in the epoch to determine the
+// proposer for that slot and assigns them accordingly.
+func ProposerAssignments(ctx context.Context, state state.BeaconState, epoch primitives.Epoch) (map[primitives.ValidatorIndex][]primitives.Slot, error) {
+	// Verify if the epoch is valid for assignment based on the provided state.
+	if err := verifyAssignmentEpoch(epoch, state); err != nil {
+		return nil, err
+	}
+	startSlot, err := slots.EpochStart(epoch)
+	if err != nil {
+		return nil, err
 	}
 
-	proposerIndexToSlots := make(map[primitives.ValidatorIndex][]primitives.Slot, params.BeaconConfig().SlotsPerEpoch)
+	proposerAssignments := make(map[primitives.ValidatorIndex][]primitives.Slot)
+
+	originalStateSlot := state.Slot()
+
 	for slot := startSlot; slot < startSlot+params.BeaconConfig().SlotsPerEpoch; slot++ {
 		// Skip proposer assignment for genesis slot.
 		if slot == 0 {
 			continue
 		}
+		// Set the state's current slot.
 		if err := state.SetSlot(slot); err != nil {
-			return nil, nil, err
+			return nil, err
 		}
+
+		// Determine the proposer index for the current slot.
 		i, err := BeaconProposerIndex(ctx, state)
 		if err != nil {
-			return nil, nil, errors.Wrapf(err, "could not check proposer at slot %d", state.Slot())
+			return nil, errors.Wrapf(err, "could not check proposer at slot %d", state.Slot())
 		}
-		proposerIndexToSlots[i] = append(proposerIndexToSlots[i], slot)
-	}
 
-	// If previous proposer indices computation is outside if current proposal epoch range,
-	// we need to reset state slot back to start slot so that we can compute the correct committees.
-	currentProposalEpoch := epoch < nextEpoch
-	if !currentProposalEpoch {
-		if err := state.SetSlot(state.Slot() - params.BeaconConfig().SlotsPerEpoch); err != nil {
-			return nil, nil, err
+		// Append the slot to the proposer's assignments.
+		if _, ok := proposerAssignments[i]; !ok {
+			proposerAssignments[i] = make([]primitives.Slot, 0)
 		}
+		proposerAssignments[i] = append(proposerAssignments[i], slot)
 	}
 
+	// Reset state back to its original slot.
+	if err := state.SetSlot(originalStateSlot); err != nil {
+		return nil, err
+	}
+
+	return proposerAssignments, nil
+}
+
+// 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.
+func CommitteeAssignments(ctx context.Context, state state.BeaconState, epoch primitives.Epoch, validators []primitives.ValidatorIndex) (map[primitives.ValidatorIndex]*CommitteeAssignment, error) {
+	// Verify if the epoch is valid for assignment based on the provided state.
+	if err := verifyAssignmentEpoch(epoch, state); err != nil {
+		return nil, err
+	}
+
+	// Retrieve active validator indices for the specified epoch.
 	activeValidatorIndices, err := ActiveValidatorIndices(ctx, state, epoch)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
-	// Each slot in an epoch has a different set of committees. This value is derived from the
-	// active validator set, which does not change.
+
+	// Determine the number of committees per slot based on the number of active validator indices.
 	numCommitteesPerSlot := SlotCommitteeCount(uint64(len(activeValidatorIndices)))
-	validatorIndexToCommittee := make(map[primitives.ValidatorIndex]*CommitteeAssignmentContainer, len(activeValidatorIndices))
 
-	// Compute all committees for all slots.
+	startSlot, err := slots.EpochStart(epoch)
+	if err != nil {
+		return nil, err
+	}
+
+	assignments := make(map[primitives.ValidatorIndex]*CommitteeAssignment)
+	vals := make(map[primitives.ValidatorIndex]struct{})
+	for _, v := range validators {
+		vals[v] = struct{}{}
+	}
+
+	// Compute committee assignments for each slot in the epoch.
 	for i := primitives.Slot(0); i < params.BeaconConfig().SlotsPerEpoch; i++ {
-		// Compute committees.
+		// Compute committees for the current slot.
 		for j := uint64(0); j < numCommitteesPerSlot; j++ {
 			slot := startSlot + i
-			committee, err := BeaconCommitteeFromState(ctx, state, slot, primitives.CommitteeIndex(j) /*committee index*/)
+			committee, err := BeaconCommitteeFromState(ctx, state, slot, primitives.CommitteeIndex(j))
 			if err != nil {
-				return nil, nil, err
+				return nil, err
 			}
 
-			cac := &CommitteeAssignmentContainer{
-				Committee:      committee,
-				CommitteeIndex: primitives.CommitteeIndex(j),
-				AttesterSlot:   slot,
-			}
 			for _, vIndex := range committee {
-				validatorIndexToCommittee[vIndex] = cac
+				if _, ok := vals[vIndex]; !ok { // Skip if the validator is not in the provided validators slice.
+					continue
+				}
+				if _, ok := assignments[vIndex]; !ok {
+					assignments[vIndex] = &CommitteeAssignment{}
+				}
+				assignments[vIndex].Committee = committee
+				assignments[vIndex].AttesterSlot = slot
+				assignments[vIndex].CommitteeIndex = primitives.CommitteeIndex(j)
 			}
 		}
 	}
 
-	return validatorIndexToCommittee, proposerIndexToSlots, nil
+	return assignments, nil
 }
 
 // VerifyBitfieldLength verifies that a bitfield length matches the given committee size.

beacon-chain/core/helpers/beacon_committee_test.go

@@ -104,7 +104,7 @@ func TestCommitteeAssignments_CannotRetrieveFutureEpoch(t *testing.T) {
 		Slot: 0, // Epoch 0.
 	})
 	require.NoError(t, err)
-	_, _, err = helpers.CommitteeAssignments(context.Background(), state, epoch+1)
+	_, err = helpers.Assignments(context.Background(), state, epoch+1)
 	assert.ErrorContains(t, "can't be greater than next epoch", err)
 }
 
@@ -128,10 +128,10 @@ func TestCommitteeAssignments_NoProposerForSlot0(t *testing.T) {
 		RandaoMixes: make([][]byte, params.BeaconConfig().EpochsPerHistoricalVector),
 	})
 	require.NoError(t, err)
-	_, proposerIndexToSlots, err := helpers.CommitteeAssignments(context.Background(), state, 0)
-	require.NoError(t, err, "Failed to determine CommitteeAssignments")
-	for _, ss := range proposerIndexToSlots {
-		for _, s := range ss {
+	assignments, err := helpers.Assignments(context.Background(), state, 0)
+	require.NoError(t, err, "Failed to determine Assignments")
+	for _, a := range assignments {
+		for _, s := range a.ProposerSlots {
 			assert.NotEqual(t, uint64(0), s, "No proposer should be assigned to slot 0")
 		}
 	}
@@ -201,14 +201,14 @@ func TestCommitteeAssignments_CanRetrieve(t *testing.T) {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
 			helpers.ClearCache()
 
-			validatorIndexToCommittee, proposerIndexToSlots, err := helpers.CommitteeAssignments(context.Background(), state, slots.ToEpoch(tt.slot))
-			require.NoError(t, err, "Failed to determine CommitteeAssignments")
-			cac := validatorIndexToCommittee[tt.index]
+			assignments, err := helpers.Assignments(context.Background(), state, slots.ToEpoch(tt.slot))
+			require.NoError(t, err, "Failed to determine Assignments")
+			cac := assignments[tt.index]
 			assert.Equal(t, tt.committeeIndex, cac.CommitteeIndex, "Unexpected committeeIndex for validator index %d", tt.index)
 			assert.Equal(t, tt.slot, cac.AttesterSlot, "Unexpected slot for validator index %d", tt.index)
-			if len(proposerIndexToSlots[tt.index]) > 0 && proposerIndexToSlots[tt.index][0] != tt.proposerSlot {
+			if len(assignments[tt.index].ProposerSlots) > 0 && assignments[tt.index].ProposerSlots[0] != tt.proposerSlot {
 				t.Errorf("wanted proposer slot %d, got proposer slot %d for validator index %d",
-					tt.proposerSlot, proposerIndexToSlots[tt.index][0], tt.index)
+					tt.proposerSlot, assignments[tt.index].ProposerSlots[0], tt.index)
 			}
 			assert.DeepEqual(t, tt.committee, cac.Committee, "Unexpected committee for validator index %d", tt.index)
 		})
@@ -238,13 +238,13 @@ func TestCommitteeAssignments_CannotRetrieveFuture(t *testing.T) {
 		RandaoMixes: make([][]byte, params.BeaconConfig().EpochsPerHistoricalVector),
 	})
 	require.NoError(t, err)
-	_, proposerIndxs, err := helpers.CommitteeAssignments(context.Background(), state, time.CurrentEpoch(state))
+	assignments, err := helpers.Assignments(context.Background(), state, time.CurrentEpoch(state))
 	require.NoError(t, err)
-	require.NotEqual(t, 0, len(proposerIndxs), "wanted non-zero proposer index set")
+	require.NotEqual(t, 0, len(assignments), "wanted non-zero proposer index set")
 
-	_, proposerIndxs, err = helpers.CommitteeAssignments(context.Background(), state, time.CurrentEpoch(state)+1)
+	assignments, err = helpers.Assignments(context.Background(), state, time.CurrentEpoch(state)+1)
 	require.NoError(t, err)
-	require.NotEqual(t, 0, len(proposerIndxs), "wanted non-zero proposer index set")
+	require.NotEqual(t, 0, len(assignments), "wanted non-zero proposer index set")
 }
 
 func TestCommitteeAssignments_CannotRetrieveOlderThanSlotsPerHistoricalRoot(t *testing.T) {
@@ -264,7 +264,7 @@ func TestCommitteeAssignments_CannotRetrieveOlderThanSlotsPerHistoricalRoot(t *t
 		RandaoMixes: make([][]byte, params.BeaconConfig().EpochsPerHistoricalVector),
 	})
 	require.NoError(t, err)
-	_, _, err = helpers.CommitteeAssignments(context.Background(), state, 0)
+	_, err = helpers.Assignments(context.Background(), state, 0)
 	require.ErrorContains(t, "start slot 0 is smaller than the minimum valid start slot 1", err)
 }
 
@@ -286,12 +286,12 @@ func TestCommitteeAssignments_EverySlotHasMin1Proposer(t *testing.T) {
 	})
 	require.NoError(t, err)
 	epoch := primitives.Epoch(1)
-	_, proposerIndexToSlots, err := helpers.CommitteeAssignments(context.Background(), state, epoch)
-	require.NoError(t, err, "Failed to determine CommitteeAssignments")
+	assignments, err := helpers.Assignments(context.Background(), state, epoch)
+	require.NoError(t, err, "Failed to determine Assignments")
 
 	slotsWithProposers := make(map[primitives.Slot]bool)
-	for _, proposerSlots := range proposerIndexToSlots {
-		for _, slot := range proposerSlots {
+	for _, a := range assignments {
+		for _, slot := range a.ProposerSlots {
 			slotsWithProposers[slot] = true
 		}
 	}

beacon-chain/rpc/eth/validator/handlers.go

@@ -665,7 +665,7 @@ func (s *Server) GetAttesterDuties(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	committeeAssignments, _, err := helpers.CommitteeAssignments(ctx, st, requestedEpoch)
+	assignments, err := helpers.CommitteeAssignments(ctx, st, requestedEpoch, requestedValIndices)
 	if err != nil {
 		httputil.HandleError(w, "Could not compute committee assignments: "+err.Error(), http.StatusInternalServerError)
 		return
@@ -685,7 +685,7 @@ func (s *Server) GetAttesterDuties(w http.ResponseWriter, r *http.Request) {
 			httputil.HandleError(w, fmt.Sprintf("Invalid validator index %d", index), http.StatusBadRequest)
 			return
 		}
-		committee := committeeAssignments[index]
+		committee := assignments[index]
 		if committee == nil {
 			continue
 		}
@@ -794,19 +794,19 @@ func (s *Server) GetProposerDuties(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
-	var proposals map[primitives.ValidatorIndex][]primitives.Slot
+	var assignments map[primitives.ValidatorIndex][]primitives.Slot
 	if nextEpochLookahead {
-		_, proposals, err = helpers.CommitteeAssignments(ctx, st, nextEpoch)
+		assignments, err = helpers.ProposerAssignments(ctx, st, nextEpoch)
 	} else {
-		_, proposals, err = helpers.CommitteeAssignments(ctx, st, requestedEpoch)
+		assignments, err = helpers.ProposerAssignments(ctx, st, requestedEpoch)
 	}
 	if err != nil {
 		httputil.HandleError(w, "Could not compute committee assignments: "+err.Error(), http.StatusInternalServerError)
 		return
 	}
 
 	duties := make([]*structs.ProposerDuty, 0)
-	for index, proposalSlots := range proposals {
+	for index, proposalSlots := range assignments {
 		val, err := st.ValidatorAtIndexReadOnly(index)
 		if err != nil {
 			httputil.HandleError(w, fmt.Sprintf("Could not get validator at index %d: %v", index, err), http.StatusInternalServerError)

beacon-chain/rpc/prysm/v1alpha1/beacon/assignments.go

@@ -105,23 +105,28 @@ func (bs *Server) ListValidatorAssignments(
 	}
 
 	// Initialize all committee related data.
-	committeeAssignments, proposerIndexToSlots, err := helpers.CommitteeAssignments(ctx, requestedState, requestedEpoch)
+	assignments, err := helpers.CommitteeAssignments(ctx, requestedState, requestedEpoch, filteredIndices[start:end])
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "Could not compute committee assignments: %v", err)
 	}
 
+	proposalSlots, err := helpers.ProposerAssignments(ctx, requestedState, requestedEpoch)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "Could not compute proposer slots: %v", err)
+	}
+
 	for _, index := range filteredIndices[start:end] {
 		if uint64(index) >= uint64(requestedState.NumValidators()) {
 			return nil, status.Errorf(codes.OutOfRange, "Validator index %d >= validator count %d",
 				index, requestedState.NumValidators())
 		}
-		comAssignment := committeeAssignments[index]
+		a := assignments[index]
 		pubkey := requestedState.PubkeyAtIndex(index)
 		assign := &ethpb.ValidatorAssignments_CommitteeAssignment{
-			BeaconCommittees: comAssignment.Committee,
-			CommitteeIndex:   comAssignment.CommitteeIndex,
-			AttesterSlot:     comAssignment.AttesterSlot,
-			ProposerSlots:    proposerIndexToSlots[index],
+			BeaconCommittees: a.Committee,
+			CommitteeIndex:   a.CommitteeIndex,
+			AttesterSlot:     a.AttesterSlot,
+			ProposerSlots:    proposalSlots[index],
 			PublicKey:        pubkey[:],
 			ValidatorIndex:   index,
 		}