coder · Emyrk · Aug 29, 2024 · Aug 28, 2024 · johnstcn · Aug 28, 2024
diff --git a/coderd/util/slice/example_test.go b/coderd/util/slice/example_test.go
@@ -10,8 +10,8 @@ import (
 func ExampleSymmetricDifference() {
 	// The goal of this function is to find the elements to add & remove from
 	// set 'a' to make it equal to set 'b'.
-	a := []int{1, 2, 5, 6}
-	b := []int{2, 3, 4, 5}
+	a := []int{1, 2, 5, 6, 6, 6}
+	b := []int{2, 3, 3, 3, 4, 5}
 	add, remove := slice.SymmetricDifference(a, b)
 	fmt.Println("Elements to add:", add)
 	fmt.Println("Elements to remove:", remove)

diff --git a/coderd/util/slice/slice.go b/coderd/util/slice/slice.go
@@ -62,6 +62,20 @@ func Overlap[T comparable](a []T, b []T) bool {
 	})
 }
 
+func UniqueFunc[T any](a []T, equal func(a, b T) bool) []T {
+	cpy := make([]T, 0, len(a))
+
+	for _, v := range a {
+		if ContainsCompare(cpy, v, equal) {
+			continue
+		}
+
+		cpy = append(cpy, v)
+	}
+
+	return cpy
+}
+
 // Unique returns a new slice with all duplicate elements removed.
 func Unique[T comparable](a []T) []T {
 	cpy := make([]T, 0, len(a))
@@ -109,15 +123,15 @@ func Descending[T constraints.Ordered](a, b T) int {
 }
 
 // SymmetricDifference returns the elements that need to be added and removed
-// to get from set 'a' to set 'b'.
+// to get from set 'a' to set 'b'. Note that duplicates are ignored in sets.
 // In classical set theory notation, SymmetricDifference returns
 // all elements of {add} and {remove} together. It is more useful to
 // return them as their own slices.
 // Notation: A Δ B = (A\B) ∪ (B\A)
 // Example:
 //
 //	a := []int{1, 3, 4}
-//	b := []int{1, 2}
+//	b := []int{1, 2, 2, 2}
 //	add, remove := SymmetricDifference(a, b)
 //	fmt.Println(add)    // [2]
 //	fmt.Println(remove) // [3, 4]
@@ -127,6 +141,8 @@ func SymmetricDifference[T comparable](a, b []T) (add []T, remove []T) {
 }
 
 func SymmetricDifferenceFunc[T any](a, b []T, equal func(a, b T) bool) (add []T, remove []T) {
+	// Ignore all duplicates
+	a, b = UniqueFunc(a, equal), UniqueFunc(b, equal)
 	return DifferenceFunc(b, a, equal), DifferenceFunc(a, b, equal)
 }
 

diff --git a/coderd/util/slice/slice_test.go b/coderd/util/slice/slice_test.go
@@ -52,6 +52,22 @@ func TestUnique(t *testing.T) {
 		slice.Unique([]string{
 			"a", "a", "a",
 		}))
+
+	require.ElementsMatch(t,
+		[]int{1, 2, 3, 4, 5},
+		slice.UniqueFunc([]int{
+			1, 2, 3, 4, 5, 1, 2, 3, 4, 5,
+		}, func(a, b int) bool {
+			return a == b
+		}))
+
+	require.ElementsMatch(t,
+		[]string{"a"},
+		slice.UniqueFunc([]string{
+			"a", "a", "a",
+		}, func(a, b string) bool {
+			return a == b
+		}))
 }
 
 func TestContains(t *testing.T) {
@@ -230,4 +246,15 @@ func TestSymmetricDifference(t *testing.T) {
 		require.ElementsMatch(t, []int{1, 2, 3}, add)
 		require.ElementsMatch(t, []int{}, remove)
 	})
+
+	t.Run("Duplicates", func(t *testing.T) {
+		t.Parallel()
+
+		add, remove := slice.SymmetricDifference(
+			[]int{5, 5, 5, 1, 1, 1, 3, 3, 3, 5, 5, 5},
+			[]int{2, 2, 2, 1, 1, 1, 2, 4, 4, 4, 5, 5, 5, 1, 1},
+		)
+		require.ElementsMatch(t, []int{2, 4}, add)
+		require.ElementsMatch(t, []int{3}, remove)
+	})
 }