agreement.go

package election

import (
	"errors"
	"fmt"
	"log"
	"sync"

	"github.com/gonum/stat/combin"
	"github.com/google/uuid"
)

type BallotList map[string]uuid.UUID

type Agreement struct {
	sync.Mutex

	// The minimum servers required for being able to tally
	required int
	// Map of ballot lists receveid from servers
	ourList     BallotList
	ballotLists map[UniqueID]BallotList
}

func NewAgreement(required int) *Agreement {
	return &Agreement{
		Mutex:       sync.Mutex{},
		required:    required,
		ballotLists: make(map[UniqueID]BallotList),
	}
}

func (a *Agreement) AddList(serverID UniqueID, ballotList map[string]uuid.UUID) error {
	a.Lock()
	defer a.Unlock()

	if _, ok := a.ballotLists[serverID]; ok {
		return fmt.Errorf("server %s already sent ballot list", serverID)
	}

	a.ballotLists[serverID] = ballotList
	return nil
}

func (a *Agreement) TallyList() ([]uuid.UUID, error) {
	a.Lock()
	defer a.Unlock()

	// First create a list of received ballot lists
	bls := make([]BallotList, 0, len(a.ballotLists))
	for _, b := range a.ballotLists {
		bls = append(bls, b)
	}

	// Find the maximum subset
	hashes := MaxSubset(a.required, bls)
	if len(hashes) == 0 {
		return nil, errors.New("cannot agree on ballots")
	}

	// We know create a list of ballots we have in the subset.
	list := make([]uuid.UUID, 0, len(a.ourList))
	for h := range hashes {
		if id, ok := a.ourList[h]; ok {
			list = append(list, id)
			delete(hashes, h) // Delete hashes that we have.
		}
	}

	// If hashes is non-empty then we are not part of the max-subset.
	if len(hashes) != 0 {
		return nil, fmt.Errorf("missing %d ballots to tally", len(hashes))
	}

	return list, nil
}

func MaxSubset(t int, ballotLists []BallotList) map[string]bool {
	n := len(ballotLists)
	combis := combin.Combinations(n, t)
	log.Printf("N: %d, combis: %d\n", n, len(combis))
	subsets := make([]map[string]bool, len(combis), n)
	for i, ls := range combis {
		sets := make([]BallotList, 0, len(ls))
		for _, i := range ls {
			sets = append(sets, ballotLists[i])
		}
		subsets[i] = Intersection(sets)
	}
	max := 0
	for i := 0; i < len(subsets); i++ {
		if len(subsets[i]) > max {
			max = i
		}
	}
	return subsets[max]
}

func Intersection(superset []BallotList) map[string]bool {
	section := make(map[string]int, len(superset[0]))
	for _, set := range superset {
		for hash := range set {
			section[hash]++
		}
	}
	set := make(map[string]bool)
	for hash, num := range section {
		if num >= len(superset) {
			set[hash] = true
		}
	}
	return set
}