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>
pull/5504/head
Andrew Dunham 2 years ago committed by GitHub
parent 1ac4a26fee
commit e945d87d76
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -6,60 +6,30 @@
// It is similar to the unix command uniq. // It is similar to the unix command uniq.
package uniq package uniq
import ( // ModifySlice removes adjacent duplicate elements from the given slice. It
"fmt" // adjusts the length of the slice appropriately and zeros the tail.
"reflect" //
) // ModifySlice does O(len(*slice)) operations.
func ModifySlice[E comparable](slice *[]E) {
type badTypeError struct { // Remove duplicates
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()
dst := 0 dst := 0
for i := 1; i < length; i++ { for i := 1; i < len(*slice); i++ {
if eq(dst, i) { if (*slice)[i] == (*slice)[dst] {
continue continue
} }
dst++ dst++
// slice[dst] = slice[i] (*slice)[dst] = (*slice)[i]
rv.Index(dst).Set(rv.Index(i))
} }
// Zero out the elements we removed at the end of the slice
end := dst + 1 end := dst + 1
var zero reflect.Value var zero E
if end < length { for i := end; i < len(*slice); i++ {
zero = reflect.Zero(rv.Type().Elem()) (*slice)[i] = zero
}
// for i := range slice[end:] {
// size[i] = 0/nil/{}
// }
for i := end; i < length; i++ {
// slice[i] = 0/nil/{}
rv.Index(i).Set(zero)
} }
// slice = slice[:end] // Truncate the slice
if end < length { if end < len(*slice) {
rv.SetLen(end) *slice = (*slice)[:end]
} }
} }

@ -12,7 +12,7 @@ import (
"tailscale.com/util/uniq" "tailscale.com/util/uniq"
) )
func TestModifySlice(t *testing.T) { func runTests(t *testing.T, cb func(*[]int)) {
tests := []struct { tests := []struct {
in []int in []int
want []int want []int
@ -29,7 +29,7 @@ func TestModifySlice(t *testing.T) {
for _, test := range tests { for _, test := range tests {
in := make([]int, len(test.in)) in := make([]int, len(test.in))
copy(in, 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) { if !reflect.DeepEqual(test.in, test.want) {
t.Errorf("uniq.Slice(%v) = %v, want %v", in, 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) { func Benchmark(b *testing.B) {
benches := []struct { benches := []struct {
name string name string
@ -83,6 +89,6 @@ func benchmark(b *testing.B, size int64, reset func(s []byte)) {
for i := 0; i < b.N; i++ { for i := 0; i < b.N; i++ {
s = s[:size] s = s[:size]
reset(s) reset(s)
uniq.ModifySlice(&s, func(i, j int) bool { return s[i] == s[j] }) uniq.ModifySlice(&s)
} }
} }

@ -2843,7 +2843,7 @@ func (c *Conn) bindSocket(rucPtr **RebindingUDPConn, network string, curPortFate
} }
ports = append(ports, 0) ports = append(ports, 0)
// Remove duplicates. (All duplicates are consecutive.) // 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 var pconn nettype.PacketConn
for _, port := range ports { for _, port := range ports {

Loading…
Cancel
Save