feat: statistics tests (#189)

* feat: statistics tests

* feat: use a set for ed25519 public key

internal/crypto/sets.go

@@ -0,0 +1,14 @@
+package crypto
+
+import "crypto/ed25519"
+
+type ED25519PublicKeySet map[[Ed25519PublicSize]byte]struct{}
+
+func (set ED25519PublicKeySet) Add(key ed25519.PublicKey) {
+	set[[Ed25519PublicSize]byte(key)] = struct{}{}
+}
+
+func (set ED25519PublicKeySet) Has(key ed25519.PublicKey) bool {
+	_, ok := set[[Ed25519PublicSize]byte(key)]
+	return ok
+}

internal/polkavm/host_call/accumulate_functions.go

@@ -76,7 +76,7 @@ func Designate(gas Gas, regs Registers, mem Memory, ctxPair AccumulateContextPai
 
 	const (
 		bandersnatch = crypto.BandersnatchSize
-		ed25519      = bandersnatch + 32
+		ed25519      = bandersnatch + crypto.Ed25519PublicSize
 		bls          = ed25519 + crypto.BLSSize
 		metadata     = bls + crypto.MetadataSize
 	)

internal/statetransition/state_transition.go

@@ -67,15 +67,15 @@ func UpdateState(s *state.State, newBlock block.Block) error {
 		return err
 	}
 
-	newValidatorStatistics := calculateNewValidatorStatistics(newBlock, newTimeState, s.ValidatorStatistics)
-
 	intermediateCoreAssignments := calculateIntermediateCoreAssignmentsFromExtrinsics(newBlock.Extrinsic.ED, s.CoreAssignments)
 	intermediateCoreAssignments, err = CalculateIntermediateCoreFromAssurances(newValidatorState.CurrentValidators, intermediateCoreAssignments, newBlock.Header, newBlock.Extrinsic.EA)
 	if err != nil {
 		return err
 	}
 
-	newCoreAssignments := calculateNewCoreAssignments(newBlock.Extrinsic.EG, intermediateCoreAssignments, s.ValidatorState, newTimeState, newEntropyPool)
+	newCoreAssignments, reporters := calculateNewCoreAssignments(newBlock.Extrinsic.EG, intermediateCoreAssignments, s.ValidatorState, newTimeState, newEntropyPool)
+
+	newValidatorStatistics := CalculateNewValidatorStatistics(newBlock, s.TimeslotIndex, s.ValidatorStatistics, reporters, s.ValidatorState.CurrentValidators)
 
 	workReports := getAvailableWorkReports(intermediateCoreAssignments)
 
@@ -1168,8 +1168,8 @@ func calculateNewCoreAssignments(
 	validatorState validator.ValidatorState,
 	newTimeslot jamtime.Timeslot,
 	entropyPool state.EntropyPool,
-) state.CoreAssignments {
-	newAssignments := intermediateAssignments
+) (newAssignments state.CoreAssignments, reporters crypto.ED25519PublicKeySet) {
+	newAssignments = intermediateAssignments
 	sortedGuarantees := sortGuaranteesByCoreIndex(guarantees.Guarantees)
 
 	for _, guarantee := range sortedGuarantees {
@@ -1191,7 +1191,9 @@ func calculateNewCoreAssignments(
 				validatorState.ArchivedValidators,
 			)
 
-			if verifyGuaranteeCredentials(guarantee, validators, entropyPool, newTimeslot) {
+			var ok bool
+			ok, reporters = verifyGuaranteeCredentials(guarantee, validators, entropyPool, newTimeslot)
+			if ok {
 				newAssignments[coreIndex] = &state.Assignment{
 					WorkReport: &guarantee.WorkReport,
 					Time:       newTimeslot,
@@ -1200,7 +1202,7 @@ func calculateNewCoreAssignments(
 		}
 	}
 
-	return newAssignments
+	return newAssignments, reporters
 }
 
 // generateRefinementContextID serializes the RefinementContext and returns its SHA-256 hash as a hex string.
@@ -1474,11 +1476,12 @@ func verifyGuaranteeCredentials(
 	validators safrole.ValidatorsData,
 	entropyPool state.EntropyPool,
 	currentTimeslot jamtime.Timeslot,
-) bool {
+) (bool, crypto.ED25519PublicKeySet) {
+	reporters := make(crypto.ED25519PublicKeySet)
 	// Verify that credentials are ordered by validator index (equation 138)
 	for i := 1; i < len(guarantee.Credentials); i++ {
 		if guarantee.Credentials[i-1].ValidatorIndex >= guarantee.Credentials[i].ValidatorIndex {
-			return false
+			return false, reporters
 		}
 	}
 
@@ -1499,38 +1502,39 @@ func verifyGuaranteeCredentials(
 	}
 
 	if err != nil {
-		return false
+		return false, reporters
 	}
 
 	// Verify the signatures using the correct validator keys (Equation 140 v0.4.5)
 	for _, credential := range guarantee.Credentials {
 		if credential.ValidatorIndex >= uint16(len(validators)) {
-			return false
+			return false, reporters
 		}
 
 		validatorKey := validators[credential.ValidatorIndex]
 		// Check if the validator key is valid
 		if len(validatorKey.Ed25519) != ed25519.PublicKeySize {
-			return false
+			return false, reporters
 		}
 
 		// Check if the validator is assigned to the core specified in the work report
 		if !isValidatorAssignedToCore(credential.ValidatorIndex, guarantee.WorkReport.CoreIndex, coreAssignments) {
-			return false
+			return false, reporters
 		}
 
 		reportBytes, err := jam.Marshal(guarantee.WorkReport)
 		if err != nil {
-			return false
+			return false, reporters
 		}
 		hashed := crypto.HashData(reportBytes)
 		message := append([]byte(signatureContextGuarantee), hashed[:]...)
 		if !ed25519.Verify(validatorKey.Ed25519, message, credential.Signature[:]) {
-			return false
+			return false, reporters
 		}
+		reporters.Add(validatorKey.Ed25519)
 	}
 
-	return true
+	return true, reporters
 }
 
 // isValidatorAssignedToCore checks if a validator is assigned to a specific core.
@@ -2104,9 +2108,9 @@ func buildServiceAccumulationCommitments(accumResults map[block.ServiceId]state.
 	return commitments
 }
 
-// calculateNewValidatorStatistics implements equation 30:
+// CalculateNewValidatorStatistics implements equation 30:
 // π′ ≺ (EG, EP, EA, ET, τ, κ′, π, H)
-func calculateNewValidatorStatistics(block block.Block, timeslot jamtime.Timeslot, validatorStatistics validator.ValidatorStatisticsState) validator.ValidatorStatisticsState {
+func CalculateNewValidatorStatistics(block block.Block, timeslot jamtime.Timeslot, validatorStatistics validator.ValidatorStatisticsState, reporters crypto.ED25519PublicKeySet, currValidators safrole.ValidatorsData) validator.ValidatorStatisticsState {
 	newStats := validatorStatistics
 
 	// Implements equations 170-171:
@@ -2120,7 +2124,7 @@ func calculateNewValidatorStatistics(block block.Block, timeslot jamtime.Timeslo
 	}
 
 	// Implements equation 172: ∀v ∈ NV
-	for v := uint16(0); v < uint16(len(newStats)); v++ {
+	for v := uint16(0); v < uint16(len(newStats[0])); v++ {
 		// π′₀[v]b ≡ a[v]b + (v = Hi)
 		if v == block.Header.BlockAuthorIndex {
 			newStats[1][v].NumOfBlocks++
@@ -2142,11 +2146,9 @@ func calculateNewValidatorStatistics(block block.Block, timeslot jamtime.Timeslo
 
 		// π′₀[v]g ≡ a[v]g + (κ′v ∈ R)
 		// Where R is the set of reporter keys defined in eq 139
-		for _, guarantee := range block.Extrinsic.EG.Guarantees {
-			for _, credential := range guarantee.Credentials {
-				if credential.ValidatorIndex == v {
-					newStats[1][v].NumOfGuaranteedReports++
-				}
+		for reporter := range reporters {
+			if currValidators[v] != nil && slices.Equal(currValidators[v].Ed25519, reporter[:]) {
+				newStats[1][v].NumOfGuaranteedReports++
 			}
 		}
 

internal/statetransition/state_transition_test.go

@@ -484,7 +484,7 @@ func TestCalculateNewCoreAssignments(t *testing.T) {
 
 		intermediateAssignments := state.CoreAssignments{}
 
-		newAssignments := calculateNewCoreAssignments(
+		newAssignments, _ := calculateNewCoreAssignments(
 			guarantees,
 			intermediateAssignments,
 			validatorState,
@@ -551,7 +551,7 @@ func TestCalculateNewCoreAssignments(t *testing.T) {
 
 		intermediateAssignments := state.CoreAssignments{}
 
-		newAssignments := calculateNewCoreAssignments(
+		newAssignments, _ := calculateNewCoreAssignments(
 			guarantees,
 			intermediateAssignments,
 			validatorState,
@@ -616,7 +616,7 @@ func TestCalculateNewCoreAssignments(t *testing.T) {
 
 		intermediateAssignments := state.CoreAssignments{}
 
-		newAssignments := calculateNewCoreAssignments(
+		newAssignments, _ := calculateNewCoreAssignments(
 			guarantees,
 			intermediateAssignments,
 			validatorState,
@@ -683,7 +683,7 @@ func TestCalculateNewCoreAssignments(t *testing.T) {
 
 		intermediateAssignments := state.CoreAssignments{}
 
-		newAssignments := calculateNewCoreAssignments(
+		newAssignments, _ := calculateNewCoreAssignments(
 			guarantees,
 			intermediateAssignments,
 			validatorState,
@@ -840,7 +840,7 @@ func TestCalculateNewValidatorStatistics(t *testing.T) {
 			},
 		}
 
-		newStats := calculateNewValidatorStatistics(block, jamtime.Timeslot(599), initialStats)
+		newStats := CalculateNewValidatorStatistics(block, jamtime.Timeslot(599), initialStats, make(crypto.ED25519PublicKeySet), safrole.ValidatorsData{})
 
 		// Check that stats were rotated correctly
 		assert.Equal(t, uint32(10), newStats[0][0].NumOfBlocks, "Previous current stats should become history")
@@ -870,7 +870,7 @@ func TestCalculateNewValidatorStatistics(t *testing.T) {
 			},
 		}
 
-		newStats := calculateNewValidatorStatistics(block, jamtime.Timeslot(5), initialStats)
+		newStats := CalculateNewValidatorStatistics(block, jamtime.Timeslot(5), initialStats, make(crypto.ED25519PublicKeySet), safrole.ValidatorsData{})
 
 		// Check block author stats
 		assert.Equal(t, uint32(1), newStats[1][1].NumOfBlocks, "Block count should increment")
@@ -913,11 +913,15 @@ func TestCalculateNewValidatorStatistics(t *testing.T) {
 				},
 			},
 		}
-
-		newStats := calculateNewValidatorStatistics(block, jamtime.Timeslot(5), initialStats)
+		ed25519key1 := testutils.RandomED25519PublicKey(t)
+		ed25519key2 := testutils.RandomED25519PublicKey(t)
+		reporters := make(crypto.ED25519PublicKeySet)
+		reporters.Add(ed25519key1)
+		reporters.Add(ed25519key2)
+		newStats := CalculateNewValidatorStatistics(block, jamtime.Timeslot(5), initialStats, reporters, safrole.ValidatorsData{{Ed25519: ed25519key1}, {Ed25519: ed25519key2}})
 
 		// Check guarantees and assurances
-		assert.Equal(t, uint64(2), newStats[1][0].NumOfGuaranteedReports, "Should count all guarantees for validator 0")
+		assert.Equal(t, uint64(1), newStats[1][0].NumOfGuaranteedReports, "Should count all guarantees for validator 0")
 		assert.Equal(t, uint64(1), newStats[1][1].NumOfGuaranteedReports, "Should count all guarantees for validator 1")
 		assert.Equal(t, uint64(1), newStats[1][0].NumOfAvailabilityAssurances, "Should count assurance for validator 0")
 		assert.Equal(t, uint64(1), newStats[1][1].NumOfAvailabilityAssurances, "Should count assurance for validator 1")
@@ -964,7 +968,12 @@ func TestCalculateNewValidatorStatistics(t *testing.T) {
 			},
 		}
 
-		newStats := calculateNewValidatorStatistics(block, jamtime.Timeslot(5), initialStats)
+		ed25519key1 := testutils.RandomED25519PublicKey(t)
+		ed25519key2 := testutils.RandomED25519PublicKey(t)
+		reporters := make(crypto.ED25519PublicKeySet)
+		reporters.Add(ed25519key1)
+		reporters.Add(ed25519key2)
+		newStats := CalculateNewValidatorStatistics(block, jamtime.Timeslot(5), initialStats, reporters, safrole.ValidatorsData{{Ed25519: ed25519key1}, {Ed25519: ed25519key2}})
 
 		expected := validator.ValidatorStatistics{
 			NumOfBlocks:                 6,

internal/validator/validators.go

@@ -32,12 +32,12 @@ type ValidatorStatistics struct {
 // validator data matches the offenders list (ψ′), all the keys for that
 // validator are zero'd out.
 func NullifyOffenders(queuedValidators safrole.ValidatorsData, offenders []ed25519.PublicKey) safrole.ValidatorsData {
-	offenderMap := make(map[string]struct{})
+	offenderMap := make(crypto.ED25519PublicKeySet)
 	for _, key := range offenders {
-		offenderMap[string(key)] = struct{}{}
+		offenderMap.Add(key)
 	}
 	for i, validator := range queuedValidators {
-		if _, found := offenderMap[string(validator.Ed25519)]; found {
+		if offenderMap.Has(validator.Ed25519) {
 			queuedValidators[i] = &crypto.ValidatorKey{
 				// Ensure these 32 bytes are zero'd out, and not just nil.  TODO
 				// - maybe use a custom wrapper type for [32]byte ?