Cleanup Attestation Helper Function and Test (#1220)

beacon-chain/core/epoch/epoch_processing.go

@@ -278,3 +278,29 @@ func ProcessPartialValidatorRegistry(
 	}
 	return state, nil
 }
+
+// CleanupAttestations removes any attestation in state's latest attestations
+// such that the attestation slot is lower than state slot minus epoch length.
+// Spec pseudocode definition:
+// 		Remove any attestation in state.latest_attestations such
+// 		that attestation.data.slot < state.slot - EPOCH_LENGTH
+func CleanupAttestations(state *pb.BeaconState) *pb.BeaconState {
+	epochLength := params.BeaconConfig().EpochLength
+	var earliestSlot uint64
+
+	// If the state slot is less than epochLength, then the earliestSlot would
+	// result in a negative number. Therefore we should default to
+	// earliestSlot = 0 in this case.
+	if state.Slot > epochLength {
+		earliestSlot = state.Slot - epochLength
+	}
+
+	var latestAttestations []*pb.PendingAttestationRecord
+	for _, attestation := range state.LatestAttestations {
+		if attestation.Data.Slot >= earliestSlot {
+			latestAttestations = append(latestAttestations, attestation)
+		}
+	}
+	state.LatestAttestations = latestAttestations
+	return state
+}

beacon-chain/core/epoch/epoch_processing_test.go

@@ -2,6 +2,7 @@ package epoch
 
 import (
 	"bytes"
+	"reflect"
 	"testing"
 
 	"github.com/gogo/protobuf/proto"
@@ -484,3 +485,38 @@ func TestProcessPartialValidatorRegistry_ReachedUpperBound(t *testing.T) {
 		t.Fatalf("ProcessValidatorRegistry should have failed with upperbound")
 	}
 }
+
+func TestCleanupAttestations(t *testing.T) {
+	if params.BeaconConfig().EpochLength != 64 {
+		t.Errorf("EpochLength should be 64 for these tests to pass")
+	}
+	epochLength := params.BeaconConfig().EpochLength
+	state := &pb.BeaconState{
+		Slot: 2 * epochLength,
+		LatestAttestations: []*pb.PendingAttestationRecord{
+			{Data: &pb.AttestationData{Slot: 1}},
+			{Data: &pb.AttestationData{Slot: epochLength - 10}},
+			{Data: &pb.AttestationData{Slot: epochLength}},
+			{Data: &pb.AttestationData{Slot: epochLength + 1}},
+			{Data: &pb.AttestationData{Slot: epochLength + 20}},
+			{Data: &pb.AttestationData{Slot: 32}},
+			{Data: &pb.AttestationData{Slot: 33}},
+			{Data: &pb.AttestationData{Slot: 2 * epochLength}},
+		},
+	}
+	wanted := &pb.BeaconState{
+		Slot: 2 * epochLength,
+		LatestAttestations: []*pb.PendingAttestationRecord{
+			{Data: &pb.AttestationData{Slot: epochLength}},
+			{Data: &pb.AttestationData{Slot: epochLength + 1}},
+			{Data: &pb.AttestationData{Slot: epochLength + 20}},
+			{Data: &pb.AttestationData{Slot: 2 * epochLength}},
+		},
+	}
+	newState := CleanupAttestations(state)
+
+	if !reflect.DeepEqual(newState, wanted) {
+		t.Errorf("Wanted state: %v, got state: %v ",
+			wanted, newState)
+	}
+}