Add validator grid mapper for neighbor discovery

internal/validator/grid.go

@@ -0,0 +1,197 @@
+package validator
+
+import (
+	"crypto/ed25519"
+	"fmt"
+	"github.com/eigerco/strawberry/internal/common"
+	"github.com/eigerco/strawberry/internal/crypto"
+	"github.com/eigerco/strawberry/internal/safrole"
+	"math"
+)
+
+// GridMapper manages the mapping between grid indices and validator information across epochs.
+// It maintains data for three sets of validators:
+// - Current validators: The active set in the current epoch
+// - Archived validators: The set from the previous epoch
+// - Queued validators: The set for the next epoch
+//
+// The grid structure arranges validators in a square-like grid where validators are considered
+// neighbors if they share the same row or column. This structure is used primarily for block
+// announcements and other network communications.
+type GridMapper struct {
+	currentValidators  safrole.ValidatorsData
+	archivedValidators safrole.ValidatorsData
+	queuedValidators   safrole.ValidatorsData
+}
+
+// NewGridMapper creates a new mapper instance using the provided validator state.
+// The state must contain information about current, archived, and queued validators.
+func NewGridMapper(state ValidatorState) GridMapper {
+	return GridMapper{
+		currentValidators:  state.CurrentValidators,
+		archivedValidators: state.ArchivedValidators,
+		queuedValidators:   state.QueuedValidators,
+	}
+}
+
+// GetAllEpochsNeighborValidators returns all neighbor validators across epochs for a given index.
+// This includes:
+// - All grid neighbors from the current epoch (same row or column)
+// - The validator with the same index from the previous epoch
+// - The validator with the same index from the next epoch
+func (m GridMapper) GetAllEpochsNeighborValidators(index uint16) ([]*crypto.ValidatorKey, error) {
+	neighborsSameEpoch, err := m.GetCurrentEpochNeighborValidators(index)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get current epoch neighbor validators: %w", err)
+	}
+
+	// Add the two validators (from the previous epoch and the next epoch) with the same index
+	neighbors := make([]*crypto.ValidatorKey, 0, len(neighborsSameEpoch)+2)
+	neighbors = append(neighbors, m.archivedValidators[index], m.queuedValidators[index])
+	neighbors = append(neighbors, neighborsSameEpoch...)
+
+	return neighbors, nil
+}
+
+// GetCurrentEpochNeighborValidators returns all grid neighbors for a validator
+// within the current epoch. Grid neighbors are validators that share either
+// the same row or column in the grid structure.
+func (m GridMapper) GetCurrentEpochNeighborValidators(index uint16) ([]*crypto.ValidatorKey, error) {
+	neighborIndices := getCurrentEpochNeighborIndices(index)
+	neighbors := make([]*crypto.ValidatorKey, 0, len(neighborIndices))
+
+	for _, idx := range neighborIndices {
+		neighbors = append(neighbors, m.currentValidators[idx])
+	}
+
+	return neighbors, nil
+}
+
+// IsNeighbor determines if two validators are neighbors based on their public keys.
+// Two validators are considered neighbors if either:
+// - They are in the same epoch and share the same row or column in the grid
+// - They are in different epochs but have the same grid index
+// Parameters:
+//   - key1, key2: The Ed25519 public keys of the two validators
+//   - sameEpoch: Whether both validators are from the same epoch
+func (m GridMapper) IsNeighbor(key1, key2 ed25519.PublicKey, sameEpoch bool) bool {
+	// Self-connections are not considered neighbors
+	if key1.Equal(key2) {
+		return false
+	}
+
+	if sameEpoch {
+		idx1, found1 := m.FindValidatorIndex(key1)
+		idx2, found2 := m.FindValidatorIndex(key2)
+
+		if !found1 || !found2 {
+			return false
+		}
+		return areGridNeighbors(idx1, idx2)
+	}
+
+	// For different epochs, validators are neighbors if they have the same index
+	indices1 := m.getValidatorIndices(key1)
+	if len(indices1) == 0 {
+		return false // key1 not found in any epoch
+	}
+
+	indices2 := m.getValidatorIndices(key2)
+	// Check if there are any matching indices between the two validators
+	for idx := range indices2 {
+		if indices1[idx] {
+			return true
+		}
+	}
+
+	return false
+}
+
+// FindValidatorIndex searches for a validator's grid index in the current validator set
+// by their Ed25519 public key. Returns the index and true if found, or 0 and false if not found.
+func (m GridMapper) FindValidatorIndex(key ed25519.PublicKey) (uint16, bool) {
+	return findValidatorIndexInSlice(m.currentValidators, key)
+}
+
+// FindValidatorIndexInArchived searches for a validator's grid index in the previous epoch's
+// validator set by their Ed25519 public key. Returns the index and true if found,
+// or 0 and false if not found.
+func (m GridMapper) FindValidatorIndexInArchived(key ed25519.PublicKey) (uint16, bool) {
+	return findValidatorIndexInSlice(m.archivedValidators, key)
+}
+
+// FindValidatorIndexInQueued searches for a validator's grid index in the next epoch's
+// validator set by their Ed25519 public key. Returns the index and true if found,
+// or 0 and false if not found.
+func (m GridMapper) FindValidatorIndexInQueued(key ed25519.PublicKey) (uint16, bool) {
+	return findValidatorIndexInSlice(m.queuedValidators, key)
+}
+
+// getValidatorIndices finds all possible indices for a validator across epochs.
+func (m GridMapper) getValidatorIndices(key ed25519.PublicKey) map[uint16]bool {
+	indices := make(map[uint16]bool)
+
+	if idx, found := m.FindValidatorIndexInArchived(key); found {
+		indices[idx] = true
+	}
+	if idx, found := m.FindValidatorIndex(key); found {
+		indices[idx] = true
+	}
+	if idx, found := m.FindValidatorIndexInQueued(key); found {
+		indices[idx] = true
+	}
+
+	return indices
+}
+
+// findValidatorIndexInSlice searches for a validator's grid index
+// in a given validator set by their Ed25519 public key. It returns the index and true if found,
+// or 0 and false if not found. The index can be used to determine the validator's position
+// in the grid structure.
+func findValidatorIndexInSlice(validators safrole.ValidatorsData, key ed25519.PublicKey) (uint16, bool) {
+	for i, validator := range validators {
+		if validator != nil && ed25519.PublicKey.Equal(validator.Ed25519, key) {
+			return uint16(i), true
+		}
+	}
+	return 0, false
+}
+
+// getCurrentEpochNeighborIndices returns all validator indices that are considered
+// neighbors within the same epoch based on the grid structure. This includes all
+// validators that share either:
+// - The same row (index / gridWidth)
+// - The same column (index % gridWidth)
+// The returned slice excludes the input validatorIndex itself.
+func getCurrentEpochNeighborIndices(validatorIndex uint16) []uint16 {
+	neighbors := make([]uint16, 0)
+
+	// Loop through all validators and check if they are neighbors
+	for i := uint16(0); i < common.NumberOfValidators; i++ {
+		if i != validatorIndex && areGridNeighbors(validatorIndex, i) {
+			neighbors = append(neighbors, i)
+		}
+	}
+
+	return neighbors
+}
+
+// areGridNeighbors determines if two validators within the same epoch are neighbors
+// in the grid structure by checking if they share the same row or column.
+// The grid width is calculated as floor(sqrt(total_validators)).
+func areGridNeighbors(validatorIndex1, validatorIndex2 uint16) bool {
+	gridWidth := getGridWidth()
+	row1, col1 := validatorIndex1/gridWidth, validatorIndex1%gridWidth
+	row2, col2 := validatorIndex2/gridWidth, validatorIndex2%gridWidth
+
+	return row1 == row2 || col1 == col2
+}
+
+// getGridWidth calculates the width of the validator grid.
+// The grid is arranged as a square-like structure with width = floor(sqrt(number_of_validators)).
+// This ensures the grid dimensions are as close to square as possible while accommodating
+// all validators.
+func getGridWidth() uint16 {
+	// floor(sqrt(numValidators))
+	return uint16(math.Floor(math.Sqrt(float64(common.NumberOfValidators))))
+}

internal/validator/grid_test.go

@@ -0,0 +1,166 @@
+package validator
+
+import (
+	"crypto/ed25519"
+	"math"
+	"testing"
+
+	"github.com/eigerco/strawberry/internal/common"
+	"github.com/eigerco/strawberry/internal/crypto"
+	"github.com/eigerco/strawberry/internal/safrole"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewGridMapper(t *testing.T) {
+	state := ValidatorState{
+		CurrentValidators:  safrole.ValidatorsData{},
+		ArchivedValidators: safrole.ValidatorsData{},
+		QueuedValidators:   safrole.ValidatorsData{},
+	}
+	mapper := NewGridMapper(state)
+
+	assert.Equal(t, state.CurrentValidators, mapper.currentValidators)
+	assert.Equal(t, state.ArchivedValidators, mapper.archivedValidators)
+	assert.Equal(t, state.QueuedValidators, mapper.queuedValidators)
+}
+
+func TestGetAllEpochsNeighborValidators(t *testing.T) {
+	validators := safrole.ValidatorsData{}
+	mapper := GridMapper{
+		currentValidators:  validators,
+		archivedValidators: validators,
+		queuedValidators:   validators,
+	}
+
+	neighbors, err := mapper.GetAllEpochsNeighborValidators(0)
+	assert.NoError(t, err)
+	assert.Len(t, neighbors, 64) //62 neighbors + 1 archived + 1 queued
+}
+
+func TestFindValidatorIndex(t *testing.T) {
+	key := ed25519.PublicKey("key")
+	validators := safrole.ValidatorsData{}
+	validators[42] = &crypto.ValidatorKey{Ed25519: key}
+	mapper := GridMapper{currentValidators: validators}
+	index, found := mapper.FindValidatorIndex(key)
+	assert.True(t, found)
+	assert.Equal(t, uint16(42), index)
+
+	_, found = mapper.FindValidatorIndex(ed25519.PublicKey("missing"))
+	assert.False(t, found)
+}
+
+func TestIsNeighbor(t *testing.T) {
+	// Setup basic validator data structures
+	currentValidators := safrole.ValidatorsData{}
+	archivedValidators := safrole.ValidatorsData{}
+	queuedValidators := safrole.ValidatorsData{}
+
+	// Create test keys
+	key1 := ed25519.PublicKey("key1")
+	key2 := ed25519.PublicKey("key2")
+	key3 := ed25519.PublicKey("key3")
+	key4 := ed25519.PublicKey("key4")
+	key5 := ed25519.PublicKey("key5")
+	keyNotValidator := ed25519.PublicKey("notvalidator")
+
+	// Calculate grid width based on total validator count
+	gridWidth := uint16(math.Floor(math.Sqrt(float64(common.NumberOfValidators))))
+
+	// Setup indices for different test scenarios
+	sameRowIdx1 := uint16(0)
+	sameRowIdx2 := uint16(1)
+	sameColIdx2 := gridWidth
+	differentIdx := gridWidth + 1
+	crossEpochIdx := uint16(42)
+
+	// Setup validators in current epoch
+	currentValidators[sameRowIdx1] = &crypto.ValidatorKey{Ed25519: key1}
+	currentValidators[sameRowIdx2] = &crypto.ValidatorKey{Ed25519: key2}
+	currentValidators[sameColIdx2] = &crypto.ValidatorKey{Ed25519: key3}
+	currentValidators[differentIdx] = &crypto.ValidatorKey{Ed25519: key4}
+	currentValidators[crossEpochIdx] = &crypto.ValidatorKey{Ed25519: key5}
+
+	// Setup validators in archived and queued epochs
+	archivedValidators[crossEpochIdx] = &crypto.ValidatorKey{Ed25519: key1}
+	queuedValidators[crossEpochIdx] = &crypto.ValidatorKey{Ed25519: key2}
+
+	mapper := GridMapper{
+		currentValidators:  currentValidators,
+		archivedValidators: archivedValidators,
+		queuedValidators:   queuedValidators,
+	}
+
+	tests := []struct {
+		name      string
+		key1      ed25519.PublicKey
+		key2      ed25519.PublicKey
+		sameEpoch bool
+		want      bool
+		reason    string
+	}{
+		{
+			name:      "same epoch - same row",
+			key1:      key1,
+			key2:      key2,
+			sameEpoch: true,
+			want:      true,
+			reason:    "validators in same row should be neighbors",
+		},
+		{
+			name:      "same epoch - same column",
+			key1:      key1,
+			key2:      key3,
+			sameEpoch: true,
+			want:      true,
+			reason:    "validators in same column should be neighbors",
+		},
+		{
+			name:      "same epoch - different row and column",
+			key1:      key1,
+			key2:      key4,
+			sameEpoch: true,
+			want:      false,
+			reason:    "validators in different rows and columns should not be neighbors",
+		},
+		{
+			name:      "same epoch - self connection",
+			key1:      key1,
+			key2:      key1,
+			sameEpoch: true,
+			want:      false,
+			reason:    "validator should not be neighbor with itself",
+		},
+		{
+			name:      "different epochs - same index",
+			key1:      key1,
+			key2:      key2,
+			sameEpoch: false,
+			want:      true,
+			reason:    "validators with same index in different epochs should be neighbors",
+		},
+		{
+			name:      "different epochs - self connection",
+			key1:      key1,
+			key2:      key1,
+			sameEpoch: false,
+			want:      false,
+			reason:    "validator should not be neighbor with itself across epochs",
+		},
+		{
+			name:      "non-validator key",
+			key1:      key1,
+			key2:      keyNotValidator,
+			sameEpoch: true,
+			want:      false,
+			reason:    "non-validator key should not be neighbor with any validator",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := mapper.IsNeighbor(tt.key1, tt.key2, tt.sameEpoch)
+			assert.Equal(t, tt.want, got, tt.reason)
+		})
+	}
+}