slices: Introduce slices.UniqueFunc()

pkg/slices/slices.go

@@ -49,6 +49,29 @@ func Unique[S ~[]T, T comparable](s S) S {
 	return s[:last]
 }
 
+// UniqueFunc deduplicates the elements in the input slice like Unique, but takes a
+// function to extract the comparable "key" to compare T. This is slower than Unique,
+// but can be used with non-comparable elements.
+func UniqueFunc[S ~[]T, T any, K comparable](s S, key func(i int) K) S {
+	if len(s) < 2 {
+		return s
+	}
+
+	last := 0
+
+	set := make(map[K]struct{}, len(s))
+	for i := 0; i < len(s); i++ {
+		if _, ok := set[key(i)]; ok {
+			continue
+		}
+		set[key(i)] = struct{}{}
+		s[last] = s[i]
+		last++
+	}
+
+	return s[:last]
+}
+
 // SortedUnique sorts and dedup the input slice in place.
 // It uses the < operator to compare the elements in the slice and thus requires
 // the elements to satisfies contraints.Ordered.

pkg/slices/slices_test.go

@@ -59,6 +59,20 @@ func TestUnique(t *testing.T) {
 	}
 }
 
+func TestUniqueFunc(t *testing.T) {
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			got := UniqueFunc(
+				tc.input,
+				func(i int) int {
+					return tc.input[i]
+				},
+			)
+			assert.ElementsMatch(t, tc.expected, got)
+		})
+	}
+}
+
 func TestSortedUnique(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
@@ -143,6 +157,14 @@ func TestUniqueKeepOrdering(t *testing.T) {
 // The relevant differences between the two approaches in terms of memory are shown in the previous
 // two columns.
 func BenchmarkUnique(b *testing.B) {
+	benchmarkUnique(b, false)
+}
+
+func BenchmarkUniqueFunc(b *testing.B) {
+	benchmarkUnique(b, true)
+}
+
+func benchmarkUnique(b *testing.B, benchUniqueFunc bool) {
 	var benchCases = [...]int{96, 128, 160, 192, 256, 512, 1024}
 
 	r := rand.New(rand.NewSource(time.Now().Unix()))
@@ -162,6 +184,11 @@ func BenchmarkUnique(b *testing.B) {
 				orig = append(orig, next)
 			}
 			values := make([]int, len(orig))
+
+			key := func(i int) int {
+				return values[i]
+			}
+
 			b.ResetTimer()
 
 			for i := 0; i < b.N; i++ {
@@ -171,9 +198,13 @@ func BenchmarkUnique(b *testing.B) {
 				rand.Shuffle(len(orig), func(i, j int) {
 					orig[i], orig[j] = orig[j], orig[i]
 				})
-				b.StartTimer()
-
-				Unique(values)
+				if benchUniqueFunc {
+					b.StartTimer()
+					UniqueFunc(values, key)
+				} else {
+					b.StartTimer()
+					Unique(values)
+				}
 			}
 		})
 	}