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.
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]
}
}

@ -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)
}
}

@ -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 {

Loading…
Cancel
Save