Optimize Archival Assignment Retrieval

beacon-chain/rpc/beacon/assignments.go

@@ -2,7 +2,6 @@ package beacon
 
 import (
 	"context"
-	"fmt"
 	"strconv"
 
 	"github.com/pkg/errors"
@@ -106,12 +105,22 @@ func (bs *Server) ListValidatorAssignments(
 	proposerIndexToSlot := map[uint64]uint64{}
 	archivedInfo := &pb.ArchivedCommitteeInfo{}
 	archivedBalances := []uint64{}
+	archivedAssingments := make(map[uint64]*ethpb.ValidatorAssignments_CommitteeAssignment)
 
 	if shouldFetchFromArchive {
 		archivedInfo, archivedBalances, err = bs.archivedCommitteeData(ctx, requestedEpoch)
 		if err != nil {
 			return nil, err
 		}
+		archivedAssingments, err = archivedValidatorCommittee(
+			requestedEpoch,
+			archivedInfo,
+			activeIndices,
+			archivedBalances,
+		)
+		if err != nil {
+			return nil, status.Errorf(codes.Internal, "Could not retrieve archived assignment for epoch %d: %v", requestedEpoch, err)
+		}
 	} else {
 		committeeAssignments, proposerIndexToSlot, err = helpers.CommitteeAssignments(headState, requestedEpoch)
 		if err != nil {
@@ -125,24 +134,12 @@ func (bs *Server) ListValidatorAssignments(
 				index, len(headState.Validators))
 		}
 		if shouldFetchFromArchive {
-			committee, committeeIndex, attesterSlot, proposerSlot, err := archivedValidatorCommittee(
-				requestedEpoch,
-				index,
-				archivedInfo,
-				activeIndices,
-				archivedBalances,
-			)
-			if err != nil {
-				return nil, status.Errorf(codes.Internal, "Could not retrieve archived assignment for validator %d: %v", index, err)
-			}
-			assign := &ethpb.ValidatorAssignments_CommitteeAssignment{
-				BeaconCommittees: committee,
-				CommitteeIndex:   committeeIndex,
-				AttesterSlot:     attesterSlot,
-				ProposerSlot:     proposerSlot,
-				PublicKey:        headState.Validators[index].PublicKey,
+			assignment, ok := archivedAssingments[index]
+			if !ok {
+				return nil, status.Errorf(codes.Internal, "Could not get archived committee assignment for index %d", index)
 			}
-			res = append(res, assign)
+			assignment.PublicKey = headState.Validators[index].PublicKey
+			res = append(res, assignment)
 			continue
 		}
 		comAssignment := committeeAssignments[index]
@@ -168,25 +165,31 @@ func (bs *Server) ListValidatorAssignments(
 // information, archived balances, and a set of active validators.
 func archivedValidatorCommittee(
 	epoch uint64,
-	validatorIndex uint64,
 	archivedInfo *pb.ArchivedCommitteeInfo,
 	activeIndices []uint64,
 	archivedBalances []uint64,
-) ([]uint64, uint64, uint64, uint64, error) {
+) (map[uint64]*ethpb.ValidatorAssignments_CommitteeAssignment, error) {
 	proposerSeed := bytesutil.ToBytes32(archivedInfo.ProposerSeed)
 	attesterSeed := bytesutil.ToBytes32(archivedInfo.AttesterSeed)
 
 	startSlot := helpers.StartSlot(epoch)
 	proposerIndexToSlot := make(map[uint64]uint64)
+	activeVals := make([]*ethpb.Validator, len(archivedBalances))
+	for i, bal := range archivedBalances {
+		activeVals[i] = &ethpb.Validator{EffectiveBalance: bal}
+	}
+
 	for slot := startSlot; slot < startSlot+params.BeaconConfig().SlotsPerEpoch; slot++ {
 		seedWithSlot := append(proposerSeed[:], bytesutil.Bytes8(slot)...)
 		seedWithSlotHash := hashutil.Hash(seedWithSlot)
-		i, err := archivedProposerIndex(activeIndices, archivedBalances, seedWithSlotHash)
+		i, err := helpers.ComputeProposerIndex(activeVals, activeIndices, seedWithSlotHash)
 		if err != nil {
-			return nil, 0, 0, 0, errors.Wrapf(err, "could not check proposer at slot %d", slot)
+			return nil, errors.Wrapf(err, "could not check proposer at slot %d", slot)
 		}
 		proposerIndexToSlot[i] = slot
 	}
+
+	assingmentMap := make(map[uint64]*ethpb.ValidatorAssignments_CommitteeAssignment)
 	for slot := startSlot; slot < startSlot+params.BeaconConfig().SlotsPerEpoch; slot++ {
 		var countAtSlot = uint64(len(activeIndices)) / params.BeaconConfig().SlotsPerEpoch / params.BeaconConfig().TargetCommitteeSize
 		if countAtSlot > params.BeaconConfig().MaxCommitteesPerSlot {
@@ -196,21 +199,21 @@ func archivedValidatorCommittee(
 			countAtSlot = 1
 		}
 		for i := uint64(0); i < countAtSlot; i++ {
-			epochOffset := i + (slot%params.BeaconConfig().SlotsPerEpoch)*countAtSlot
-			totalCount := countAtSlot * params.BeaconConfig().SlotsPerEpoch
-			committee, err := helpers.ComputeCommittee(activeIndices, attesterSeed, epochOffset, totalCount)
+			committee, err := helpers.BeaconCommittee(activeIndices, attesterSeed, slot, i)
 			if err != nil {
-				return nil, 0, 0, 0, errors.Wrap(err, "could not compute committee")
+				return nil, errors.Wrap(err, "could not compute committee")
 			}
 			for _, index := range committee {
-				if validatorIndex == index {
-					proposerSlot, _ := proposerIndexToSlot[validatorIndex]
-					return committee, i, slot, proposerSlot, nil
+				assingmentMap[index] = &ethpb.ValidatorAssignments_CommitteeAssignment{
+					BeaconCommittees: committee,
+					CommitteeIndex:   i,
+					AttesterSlot:     slot,
+					ProposerSlot:     proposerIndexToSlot[index],
 				}
 			}
 		}
 	}
-	return nil, 0, 0, 0, fmt.Errorf("could not find committee for validator index %d", validatorIndex)
+	return assingmentMap, nil
 }
 
 func (bs *Server) archivedCommitteeData(ctx context.Context, requestedEpoch uint64) (*pb.ArchivedCommitteeInfo,
@@ -247,12 +250,3 @@ func (bs *Server) archivedCommitteeData(ctx context.Context, requestedEpoch uint
 	}
 	return archivedInfo, archivedBalances, nil
 }
-
-// helpers.ComputeProposerIndex wrapper.
-func archivedProposerIndex(activeIndices []uint64, activeBalances []uint64, seed [32]byte) (uint64, error) {
-	validators := make([]*ethpb.Validator, len(activeBalances))
-	for i, bal := range activeBalances {
-		validators[i] = &ethpb.Validator{EffectiveBalance: bal}
-	}
-	return helpers.ComputeProposerIndex(validators, activeIndices, seed)
-}