From e945d87d769aa9c13cd692d73f95fb6ff69369b0 Mon Sep 17 00:00:00 2001
From: Andrew Dunham <andrew@du.nham.ca>
Date: Tue, 30 Aug 2022 17:56:51 -0400
Subject: [PATCH] util/uniq: use generics instead of reflect (#5491)

This takes 75% less time per operation per some benchmarks on my mac.

Signed-off-by: Andrew Dunham <andrew@du.nham.ca>
---
 util/uniq/slice.go              | 62 +++++++++------------------------
 util/uniq/slice_test.go         | 12 +++++--
 wgengine/magicsock/magicsock.go |  2 +-
 3 files changed, 26 insertions(+), 50 deletions(-)

diff --git a/util/uniq/slice.go b/util/uniq/slice.go
index fc1ff0d7a..5438a2556 100644
--- a/util/uniq/slice.go
+++ b/util/uniq/slice.go
@@ -6,60 +6,30 @@
 // It is similar to the unix command uniq.
 package uniq
 
-import (
-	"fmt"
-	"reflect"
-)
-
-type badTypeError struct {
-	typ reflect.Type
-}
-
-func (e badTypeError) Error() string {
-	return fmt.Sprintf("uniq.ModifySlice's first argument must have type *[]T, got %v", e.typ)
-}
-
-// ModifySlice removes adjacent duplicate elements from the slice pointed to by sliceptr.
-// It adjusts the length of the slice appropriately and zeros the tail.
-// eq reports whether (*sliceptr)[i] and (*sliceptr)[j] are equal.
-// ModifySlice does O(len(*sliceptr)) operations.
-func ModifySlice(sliceptr any, eq func(i, j int) bool) {
-	rvp := reflect.ValueOf(sliceptr)
-	if rvp.Type().Kind() != reflect.Ptr {
-		panic(badTypeError{rvp.Type()})
-	}
-	rv := rvp.Elem()
-	if rv.Type().Kind() != reflect.Slice {
-		panic(badTypeError{rvp.Type()})
-	}
-
-	length := rv.Len()
+// ModifySlice removes adjacent duplicate elements from the given slice. It
+// adjusts the length of the slice appropriately and zeros the tail.
+//
+// ModifySlice does O(len(*slice)) operations.
+func ModifySlice[E comparable](slice *[]E) {
+	// Remove duplicates
 	dst := 0
-	for i := 1; i < length; i++ {
-		if eq(dst, i) {
+	for i := 1; i < len(*slice); i++ {
+		if (*slice)[i] == (*slice)[dst] {
 			continue
 		}
 		dst++
-		// slice[dst] = slice[i]
-		rv.Index(dst).Set(rv.Index(i))
+		(*slice)[dst] = (*slice)[i]
 	}
 
+	// Zero out the elements we removed at the end of the slice
 	end := dst + 1
-	var zero reflect.Value
-	if end < length {
-		zero = reflect.Zero(rv.Type().Elem())
-	}
-
-	// for i := range slice[end:] {
-	//   size[i] = 0/nil/{}
-	// }
-	for i := end; i < length; i++ {
-		// slice[i] = 0/nil/{}
-		rv.Index(i).Set(zero)
+	var zero E
+	for i := end; i < len(*slice); i++ {
+		(*slice)[i] = zero
 	}
 
-	// slice = slice[:end]
-	if end < length {
-		rv.SetLen(end)
+	// Truncate the slice
+	if end < len(*slice) {
+		*slice = (*slice)[:end]
 	}
 }
diff --git a/util/uniq/slice_test.go b/util/uniq/slice_test.go
index 2455cdc4c..313e2b435 100644
--- a/util/uniq/slice_test.go
+++ b/util/uniq/slice_test.go
@@ -12,7 +12,7 @@ import (
 	"tailscale.com/util/uniq"
 )
 
-func TestModifySlice(t *testing.T) {
+func runTests(t *testing.T, cb func(*[]int)) {
 	tests := []struct {
 		in   []int
 		want []int
@@ -29,7 +29,7 @@ func TestModifySlice(t *testing.T) {
 	for _, test := range tests {
 		in := make([]int, len(test.in))
 		copy(in, test.in)
-		uniq.ModifySlice(&test.in, func(i, j int) bool { return test.in[i] == test.in[j] })
+		cb(&test.in)
 		if !reflect.DeepEqual(test.in, test.want) {
 			t.Errorf("uniq.Slice(%v) = %v, want %v", in, test.in, test.want)
 		}
@@ -43,6 +43,12 @@ func TestModifySlice(t *testing.T) {
 	}
 }
 
+func TestModifySlice(t *testing.T) {
+	runTests(t, func(slice *[]int) {
+		uniq.ModifySlice(slice)
+	})
+}
+
 func Benchmark(b *testing.B) {
 	benches := []struct {
 		name  string
@@ -83,6 +89,6 @@ func benchmark(b *testing.B, size int64, reset func(s []byte)) {
 	for i := 0; i < b.N; i++ {
 		s = s[:size]
 		reset(s)
-		uniq.ModifySlice(&s, func(i, j int) bool { return s[i] == s[j] })
+		uniq.ModifySlice(&s)
 	}
 }
diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index ac9d6cdc4..11f6f7ad4 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -2843,7 +2843,7 @@ func (c *Conn) bindSocket(rucPtr **RebindingUDPConn, network string, curPortFate
 	}
 	ports = append(ports, 0)
 	// Remove duplicates. (All duplicates are consecutive.)
-	uniq.ModifySlice(&ports, func(i, j int) bool { return ports[i] == ports[j] })
+	uniq.ModifySlice(&ports)
 
 	var pconn nettype.PacketConn
 	for _, port := range ports {