util/uniq: use generics instead of reflect (tailscale#5491)

andrew-d · web-flow · commit e945d87d769a · 2022-08-30T17:56:51.000-04:00
This takes 75% less time per operation per some benchmarks on my mac.

Signed-off-by: Andrew Dunham &lt;andrew@du.nham.ca&gt;
diff --git 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
@@ -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
@@ -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 {

Original file line number	Diff line number	Diff line change
`@@ -12,7 +12,7 @@ import (`
`12`	`12`	`"tailscale.com/util/uniq"`
`13`	`13`	`)`
`14`	`14`
`15`		`-func TestModifySlice(t *testing.T) {`
	`15`	`+func runTests(t testing.T, cb func([]int)) {`
`16`	`16`	`tests := []struct {`
`17`	`17`	`in []int`
`18`	`18`	`want []int`
`@@ -29,7 +29,7 @@ func TestModifySlice(t *testing.T) {`
`29`	`29`	`for _, test := range tests {`
`30`	`30`	`in := make([]int, len(test.in))`
`31`	`31`	`copy(in, test.in)`
`32`		`- uniq.ModifySlice(&test.in, func(i, j int) bool { return test.in[i] == test.in[j] })`
	`32`	`+ cb(&test.in)`
`33`	`33`	`if !reflect.DeepEqual(test.in, test.want) {`
`34`	`34`	`t.Errorf("uniq.Slice(%v) = %v, want %v", in, test.in, test.want)`
`35`	`35`	`}`
`@@ -43,6 +43,12 @@ func TestModifySlice(t *testing.T) {`
`43`	`43`	`}`
`44`	`44`	`}`
`45`	`45`
	`46`	`+func TestModifySlice(t *testing.T) {`
	`47`	`+ runTests(t, func(slice *[]int) {`
	`48`	`+ uniq.ModifySlice(slice)`
	`49`	`+ })`
	`50`	`+}`
	`51`	`+`
`46`	`52`	`func Benchmark(b *testing.B) {`
`47`	`53`	`benches := []struct {`
`48`	`54`	`name string`
`@@ -83,6 +89,6 @@ func benchmark(b *testing.B, size int64, reset func(s []byte)) {`
`83`	`89`	`for i := 0; i < b.N; i++ {`
`84`	`90`	`s = s[:size]`
`85`	`91`	`reset(s)`
`86`		`- uniq.ModifySlice(&s, func(i, j int) bool { return s[i] == s[j] })`
	`92`	`+ uniq.ModifySlice(&s)`
`87`	`93`	`}`
`88`	`94`	`}`
Original file line number	Diff line number	Diff line change
`@@ -2843,7 +2843,7 @@ func (c Conn) bindSocket(rucPtr *RebindingUDPConn, network string, curPortFate`
`2843`	`2843`	`}`
`2844`	`2844`	`ports = append(ports, 0)`
`2845`	`2845`	`// Remove duplicates. (All duplicates are consecutive.)`
`2846`		`- uniq.ModifySlice(&ports, func(i, j int) bool { return ports[i] == ports[j] })`
	`2846`	`+ uniq.ModifySlice(&ports)`
`2847`	`2847`
`2848`	`2848`	`var pconn nettype.PacketConn`
`2849`	`2849`	`for _, port := range ports {`