Correctly handle choice removal in Unvote

janos · Feb 3, 2023 · ba2a0d7 · ba2a0d7
1 parent df9c037
commit ba2a0d7
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 16 deletions.
diff --git a/schulze.go b/schulze.go
@@ -71,9 +71,15 @@ func Unvote[C comparable](preferences []int, choices []C, r Record[C]) error {
 		rest := r[rank+1:]
 		for _, choice1 := range choices1 {
 			i := getChoiceIndex(choices, choice1)
+			if i < 0 {
+				continue
+			}
 			for _, choices1 := range rest {
 				for _, choice2 := range choices1 {
 					j := getChoiceIndex(choices, choice2)
+					if j < 0 {
+						continue
+					}
 					preferences[int(i)*choicesCount+int(j)] -= 1
 				}
 			}

diff --git a/schulze_test.go b/schulze_test.go
@@ -316,28 +316,80 @@ func TestVoting(t *testing.T) {
 }
 
 func TestUnvote_afterSetChoices(t *testing.T) {
-	choices := []string{"A", "B", "C"}
-	preferences := schulze.NewPreferences(len(choices))
+	t.Run("add", func(t *testing.T) {
+		choices := []string{"A", "B", "C"}
+		preferences := schulze.NewPreferences(len(choices))
 
-	ballot := schulze.Ballot[string]{"A": 1, "B": 2}
-	report, err := schulze.Vote(preferences, choices, ballot)
-	if err != nil {
-		t.Fatal(err)
-	}
+		ballot := schulze.Ballot[string]{"A": 1, "B": 2}
+		record, err := schulze.Vote(preferences, choices, ballot)
+		if err != nil {
+			t.Fatal(err)
+		}
 
-	updatedChoices := []string{"A", "D", "B", "C"}
+		updatedChoices := []string{"A", "D", "B", "C"}
 
-	updatedPreferences := schulze.SetChoices(preferences, choices, updatedChoices)
+		updatedPreferences := schulze.SetChoices(preferences, choices, updatedChoices)
 
-	if err := schulze.Unvote(updatedPreferences, updatedChoices, report); err != nil {
-		t.Fatal(err)
-	}
+		if err := schulze.Unvote(updatedPreferences, updatedChoices, record); err != nil {
+			t.Fatal(err)
+		}
 
-	wantPreferences := make([]int, len(updatedPreferences))
+		wantPreferences := make([]int, len(updatedPreferences))
 
-	if !reflect.DeepEqual(updatedPreferences, wantPreferences) {
-		t.Errorf("got preferences %v, want %v", updatedPreferences, wantPreferences)
-	}
+		if !reflect.DeepEqual(updatedPreferences, wantPreferences) {
+			t.Errorf("got preferences %v, want %v", updatedPreferences, wantPreferences)
+		}
+	})
+
+	t.Run("remove", func(t *testing.T) {
+		choices := []string{"A", "B", "C"}
+		preferences := schulze.NewPreferences(len(choices))
+
+		ballot := schulze.Ballot[string]{"A": 1, "B": 2}
+		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))
+
+		ballot := schulze.Ballot[string]{"A": 1, "B": 2}
+		record, err := schulze.Vote(preferences, choices, ballot)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		updatedChoices := []string{"B", "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)
+		}
+	})
 }
 
 func TestSetChoices(t *testing.T) {