Implement get_next_epoch_committee_assignment helper function

beacon-chain/core/helpers/committee.go

@@ -360,3 +360,75 @@ func AttestationParticipants(
 	}
 	return participants, nil
 }
+
+// NextEpochCommitteeAssignment query slots in the next epoch
+// for it to discover which shard and slot a validator gets assigned.
+//
+// Spec pseudocode definition:
+//   def get_next_epoch_committee_assignment(
+//        state: BeaconState,
+//        validator_index: ValidatorIndex,
+//        registry_change: bool) -> Tuple[List[ValidatorIndex], ShardNumber, SlotNumber, bool]:
+//    """
+//    Return the committee assignment in the next epoch for ``validator_index`` and ``registry_change``.
+//    ``assignment`` returned is a tuple of the following form:
+//        * ``assignment[0]`` is the list of validators in the committee
+//        * ``assignment[1]`` is the shard to which the committee is assigned
+//        * ``assignment[2]`` is the slot at which the committee is assigned
+//        * ``assignment[3]`` is a bool signalling if the validator is expected to propose
+//            a beacon block at the assigned slot.
+//    """
+//    current_epoch = get_current_epoch(state)
+//    next_epoch = current_epoch + 1
+//    next_epoch_start_slot = get_epoch_start_slot(next_epoch)
+//    for slot in range(next_epoch_start_slot, next_epoch_start_slot + EPOCH_LENGTH):
+//        crosslink_committees = get_crosslink_committees_at_slot(
+//            state,
+//            slot,
+//            registry_change=registry_change,
+//        )
+//        selected_committees = [
+//            committee  # Tuple[List[ValidatorIndex], ShardNumber]
+//            for committee in crosslink_committees
+//            if validator_index in committee[0]
+//        ]
+//        if len(selected_committees) > 0:
+//            validators = selected_committees[0][0]
+//            shard = selected_committees[0][1]
+//            first_committee_at_slot = crosslink_committees[0][0]  # List[ValidatorIndex]
+//            is_proposer = first_committee_at_slot[slot % len(first_committee_at_slot)] == validator_index
+//
+//            assignment = (validators, shard, slot, is_proposer)
+//            return assignment
+func NextEpochCommitteeAssignment(
+	state *pb.BeaconState,
+	index uint64,
+	registryChange bool) ([]uint64, uint64, uint64, bool, error) {
+	var selectedCommittees []*CrosslinkCommittee
+	nextEpoch := NextEpoch(state)
+	nextEpochStartSlot := StartSlot(nextEpoch)
+	for slot := nextEpochStartSlot; slot < nextEpochStartSlot+params.BeaconConfig().EpochLength; slot++ {
+		crosslinkCommittees, err := CrosslinkCommitteesAtSlot(
+			state, slot, registryChange)
+		if err != nil {
+			return []uint64{}, 0, 0, false, fmt.Errorf("could not get crosslink committee: %v", err)
+		}
+		for _, committee := range crosslinkCommittees {
+			for _, idx := range committee.Committee {
+				if idx == index {
+					selectedCommittees = append(selectedCommittees, committee)
+				}
+
+				if len(selectedCommittees) > 0 {
+					validators := selectedCommittees[0].Committee
+					shard := selectedCommittees[0].Shard
+					firstCommitteeAtSlot := crosslinkCommittees[0].Committee
+					isProposer := firstCommitteeAtSlot[slot%
+						uint64(len(firstCommitteeAtSlot))] == index
+					return validators, shard, slot, isProposer, nil
+				}
+			}
+		}
+	}
+	return []uint64{}, 0, 0, false, fmt.Errorf("could not get assignment validator %d", index)
+}

beacon-chain/core/helpers/committee_test.go

@@ -314,3 +314,93 @@ func TestAttestationParticipants_IncorrectBitfield(t *testing.T) {
 		t.Error("attestation participants should have failed with incorrect bitfield")
 	}
 }
+
+func TestNextEpochCommitteeAssignment_ok(t *testing.T) {
+	// Initialize test with 128 validators, each slot and each shard gets 2 validators.
+	validators := make([]*pb.Validator, 2*params.BeaconConfig().EpochLength)
+	for i := 0; i < len(validators); i++ {
+		validators[i] = &pb.Validator{
+			ExitEpoch: params.BeaconConfig().FarFutureEpoch,
+		}
+	}
+	state := &pb.BeaconState{
+		ValidatorRegistry: validators,
+		Slot:              params.BeaconConfig().EpochLength,
+	}
+
+	tests := []struct {
+		index      uint64
+		slot       uint64
+		committee  []uint64
+		shard      uint64
+		isProposer bool
+	}{
+		{
+			index:      0,
+			slot:       146,
+			committee:  []uint64{105, 0},
+			shard:      18,
+			isProposer: false,
+		},
+		{
+			index:      105,
+			slot:       146,
+			committee:  []uint64{105, 0},
+			shard:      18,
+			isProposer: true,
+		},
+		{
+			index:      64,
+			slot:       139,
+			committee:  []uint64{64, 52},
+			shard:      11,
+			isProposer: false,
+		},
+		{
+			index:      11,
+			slot:       130,
+			committee:  []uint64{11, 121},
+			shard:      2,
+			isProposer: true,
+		},
+	}
+
+	for _, tt := range tests {
+		committee, shard, slot, isProposer, err := NextEpochCommitteeAssignment(
+			state, tt.index, false)
+		if err != nil {
+			t.Fatalf("failed to execute NextEpochCommitteeAssignment: %v", err)
+		}
+		if shard != tt.shard {
+			t.Errorf("wanted shard %d, got shard %d",
+				tt.shard, shard)
+		}
+		if slot != tt.slot {
+			t.Errorf("wanted slot %d, got slot %d",
+				tt.slot, slot)
+		}
+		if isProposer != tt.isProposer {
+			t.Errorf("wanted isProposer %v, got isProposer %v",
+				tt.isProposer, isProposer)
+		}
+		if !reflect.DeepEqual(committee, tt.committee) {
+			t.Errorf("wanted committee %v, got committee %v",
+				tt.committee, committee)
+		}
+	}
+}
+
+func TestNextEpochCommitteeAssignment_CantFindValidator(t *testing.T) {
+	state := &pb.BeaconState{
+		Slot: params.BeaconConfig().EpochLength,
+	}
+	index := uint64(10000)
+	want := fmt.Sprintf(
+		"could not get assignment validator %d",
+		index,
+	)
+	if _, _, _, _, err := NextEpochCommitteeAssignment(
+		state, index, false); !strings.Contains(err.Error(), want) {
+		t.Errorf("Expected %s, received %v", want, err)
+	}
+}