Add possible empty unvoted list in the Reults to fix the Unvote

schulze.go

@@ -27,7 +27,8 @@ type Ballot[C comparable] map[C]int
 
 // Record represents a single vote with ranked choices. It is a list of Ballot
 // values. The first ballot is the list with the first choices, the second
-// ballot is the list with the second choices, and so on.
+// ballot is the list with the second choices, and so on. The last ballot is the
+// list of choices that are not ranked, which can be an empty list.
 type Record[C comparable] [][]C
 
 // Vote updates the preferences passed as the first argument with the Ballot
@@ -51,14 +52,16 @@ func Vote[C comparable](preferences []int, choices []C, b Ballot[C]) (Record[C],
 		}
 	}
 
+	ranksLen := len(ranks)
+
 	// set diagonal values as the values of the column of the least ranked
 	// choice to be able to have the correct preferences matrix when adding new
 	// choices
 	if hasUnrankedChoices {
 		// treat the diagonal values as one of the unranked choices,
 		// deprioritizing all choices except unranked as they are of the same
-		if l := len(ranks); l > 0 {
-			for _, choices1 := range ranks[:l-1] {
+		if ranksLen > 0 {
+			for _, choices1 := range ranks[:ranksLen-1] {
 				for _, i := range choices1 {
 					preferences[int(i)*choicesCount+int(i)] += 1
 				}
@@ -72,7 +75,14 @@ func Vote[C comparable](preferences []int, choices []C, b Ballot[C]) (Record[C],
 		}
 	}
 
-	r := make([][]C, len(ranks))
+	// prepare results capacity to avoid allocation on appending the potential
+	// unranked choices
+	resultsCap := ranksLen
+	if !hasUnrankedChoices {
+		resultsCap++
+	}
+
+	r := make([][]C, ranksLen, resultsCap)
 	for rank, indexes := range ranks {
 		if r[rank] == nil {
 			r[rank] = make([]C, 0, len(indexes))
@@ -82,6 +92,10 @@ func Vote[C comparable](preferences []int, choices []C, b Ballot[C]) (Record[C],
 		}
 	}
 
+	if !hasUnrankedChoices {
+		r = append(r, make([]C, 0))
+	}
+
 	return r, nil
 }
 
@@ -116,6 +130,9 @@ func Unvote[C comparable](preferences []int, choices []C, r Record[C]) error {
 	knownChoices := newBitset(uint64(choicesCount))
 	rankedChoices := newBitset(uint64(choicesCount))
 	// remove voting from the ranked choices of the Record
+	//
+	// it is essential to have the last record ballot as
+	// unranked choices even if it is empty
 	for _, choices1 := range r[:recordLength-1] {
 		for _, choice1 := range choices1 {
 			i := getChoiceIndex(choices, choice1)

schulze_test.go

@@ -380,7 +380,7 @@ func TestVoting(t *testing.T) {
 }
 
 func TestUnvote_afterSetChoices(t *testing.T) {
-	t.Run("add", func(t *testing.T) {
+	t.Run("add with unranked", func(t *testing.T) {
 		choices := []string{"A", "B", "C"}
 		preferences := schulze.NewPreferences(len(choices))
 
@@ -411,7 +411,38 @@ func TestUnvote_afterSetChoices(t *testing.T) {
 		}
 	})
 
-	t.Run("remove", func(t *testing.T) {
+	t.Run("add without unranked", func(t *testing.T) {
+		choices := []string{"A", "B", "C"}
+		preferences := schulze.NewPreferences(len(choices))
+
+		ballot := schulze.Ballot[string]{"A": 1, "B": 2, "C": 1}
+		record, err := schulze.Vote(preferences, choices, ballot)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		t.Logf("initial\n%v", sprintPreferences(choices, preferences))
+
+		updatedChoices := []string{"A", "B", "C", "D"}
+
+		updatedPreferences := schulze.SetChoices(preferences, choices, updatedChoices)
+
+		t.Logf("updated\n%v", sprintPreferences(updatedChoices, updatedPreferences))
+
+		if err := schulze.Unvote(updatedPreferences, updatedChoices, record); err != nil {
+			t.Fatal(err)
+		}
+
+		t.Logf("unvoted\n%v\n%v", sprintPreferences(updatedChoices, updatedPreferences), record)
+
+		wantPreferences := make([]int, len(updatedPreferences))
+
+		if !reflect.DeepEqual(updatedPreferences, wantPreferences) {
+			t.Errorf("got preferences %v, want %v", updatedPreferences, wantPreferences)
+		}
+	})
+
+	t.Run("remove with unranked", func(t *testing.T) {
 		choices := []string{"A", "B", "C"}
 		preferences := schulze.NewPreferences(len(choices))
 
@@ -436,6 +467,31 @@ func TestUnvote_afterSetChoices(t *testing.T) {
 		}
 	})
 
+	t.Run("remove without unranked", func(t *testing.T) {
+		choices := []string{"A", "B", "C"}
+		preferences := schulze.NewPreferences(len(choices))
+
+		ballot := schulze.Ballot[string]{"A": 1, "B": 2, "C": 1}
+		record, err := schulze.Vote(preferences, choices, ballot)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		updatedChoices := []string{"A", "C"}
+
+		updatedPreferences := schulze.SetChoices(preferences, choices, updatedChoices)
+
+		if err := schulze.Unvote(updatedPreferences, updatedChoices, record); err != nil {
+			t.Fatal(err)
+		}
+
+		wantPreferences := make([]int, len(updatedPreferences))
+
+		if !reflect.DeepEqual(updatedPreferences, wantPreferences) {
+			t.Errorf("got preferences %v, want %v", updatedPreferences, wantPreferences)
+		}
+	})
+
 	t.Run("rearrange", func(t *testing.T) {
 		choices := []string{"A", "B", "C"}
 		preferences := schulze.NewPreferences(len(choices))
@@ -461,7 +517,7 @@ func TestUnvote_afterSetChoices(t *testing.T) {
 		}
 	})
 
-	t.Run("complex", func(t *testing.T) {
+	t.Run("complex with unranked", func(t *testing.T) {
 		choices := []string{"A", "B", "C", "D", "E", "F", "G", "H", "I", "J"}
 		preferences := schulze.NewPreferences(len(choices))
 
@@ -491,6 +547,37 @@ func TestUnvote_afterSetChoices(t *testing.T) {
 			t.Errorf("got preferences %v, want %v", updatedPreferences, wantPreferences)
 		}
 	})
+
+	t.Run("complex without unranked", func(t *testing.T) {
+		choices := []string{"A", "B", "C", "D", "E", "F", "G", "H", "I", "J"}
+		preferences := schulze.NewPreferences(len(choices))
+
+		ballot := schulze.Ballot[string]{"A": 1, "B": 2, "C": 2, "D": 3, "E": 3, "F": 3, "G": 3, "H": 1, "I": 2, "J": 3}
+		record, err := schulze.Vote(preferences, choices, ballot)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		t.Logf("initial\n%v", sprintPreferences(choices, preferences))
+
+		updatedChoices := []string{"A", "K", "C", "E", "D", "G", "H", "J"}
+
+		updatedPreferences := schulze.SetChoices(preferences, choices, updatedChoices)
+
+		t.Logf("updated\n%v", sprintPreferences(updatedChoices, updatedPreferences))
+
+		if err := schulze.Unvote(updatedPreferences, updatedChoices, record); err != nil {
+			t.Fatal(err)
+		}
+
+		t.Logf("unvoted\n%v\n%v", sprintPreferences(updatedChoices, updatedPreferences), record)
+
+		wantPreferences := make([]int, len(updatedPreferences))
+
+		if !reflect.DeepEqual(updatedPreferences, wantPreferences) {
+			t.Errorf("got preferences %v, want %v", updatedPreferences, wantPreferences)
+		}
+	})
 }
 
 func TestDuel_Outcome(t *testing.T) {