mirror of https://github.com/tailscale/tailscale/
util/slicesx: add package for generic slice functions, use
Now that we're using rand.Shuffle in a few locations, create a generic shuffle function and use it instead. While we're at it, move the interleaveSlices function to the same package for use. Signed-off-by: Andrew Dunham <andrew@du.nham.ca> Change-Id: I0b00920e5b3eea846b6cedc30bd34d978a049fd3pull/7447/head
Andrew Dunham
2 years ago
parent
88c7d19d54
commit
73fa7dd7af
@ -0,0 +1,44 @@
|
||||
// Copyright (c) Tailscale Inc & AUTHORS
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
// Package slicesx contains some helpful generic slice functions.
|
||||
package slicesx
|
||||
|
||||
import "math/rand"
|
||||
|
||||
// Interleave combines two slices of the form [a, b, c] and [x, y, z] into a
|
||||
// slice with elements interleaved; i.e. [a, x, b, y, c, z].
|
||||
func Interleave[S ~[]T, T any](a, b S) S {
|
||||
// Avoid allocating an empty slice.
|
||||
if a == nil && b == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
var (
|
||||
i int
|
||||
ret = make([]T, 0, len(a)+len(b))
|
||||
)
|
||||
for i = 0; i < len(a) && i < len(b); i++ {
|
||||
ret = append(ret, a[i], b[i])
|
||||
}
|
||||
ret = append(ret, a[i:]...)
|
||||
ret = append(ret, b[i:]...)
|
||||
return ret
|
||||
}
|
||||
|
||||
// Shuffle randomly shuffles a slice in-place, similar to rand.Shuffle.
|
||||
func Shuffle[S ~[]T, T any](s S) {
|
||||
// TODO(andrew): use a pooled Rand?
|
||||
|
||||
// This is the same Fisher-Yates shuffle implementation as rand.Shuffle
|
||||
n := len(s)
|
||||
i := n - 1
|
||||
for ; i > 1<<31-1-1; i-- {
|
||||
j := int(rand.Int63n(int64(i + 1)))
|
||||
s[i], s[j] = s[j], s[i]
|
||||
}
|
||||
for ; i > 0; i-- {
|
||||
j := int(rand.Int31n(int32(i + 1)))
|
||||
s[i], s[j] = s[j], s[i]
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,66 @@
|
||||
// Copyright (c) Tailscale Inc & AUTHORS
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
package slicesx
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
"golang.org/x/exp/slices"
|
||||
)
|
||||
|
||||
func TestInterleave(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
a, b []int
|
||||
want []int
|
||||
}{
|
||||
{name: "equal", a: []int{1, 3, 5}, b: []int{2, 4, 6}, want: []int{1, 2, 3, 4, 5, 6}},
|
||||
{name: "short_b", a: []int{1, 3, 5}, b: []int{2, 4}, want: []int{1, 2, 3, 4, 5}},
|
||||
{name: "short_a", a: []int{1, 3}, b: []int{2, 4, 6}, want: []int{1, 2, 3, 4, 6}},
|
||||
{name: "len_1", a: []int{1}, b: []int{2, 4, 6}, want: []int{1, 2, 4, 6}},
|
||||
{name: "nil_a", a: nil, b: []int{2, 4, 6}, want: []int{2, 4, 6}},
|
||||
{name: "nil_all", a: nil, b: nil, want: nil},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
merged := Interleave(tc.a, tc.b)
|
||||
if !reflect.DeepEqual(merged, tc.want) {
|
||||
t.Errorf("got %v; want %v", merged, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkInterleave(b *testing.B) {
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
Interleave(
|
||||
[]int{1, 2, 3},
|
||||
[]int{9, 8, 7},
|
||||
)
|
||||
}
|
||||
}
|
||||
func TestShuffle(t *testing.T) {
|
||||
var sl []int
|
||||
for i := 0; i < 100; i++ {
|
||||
sl = append(sl, i)
|
||||
}
|
||||
|
||||
var wasShuffled bool
|
||||
for try := 0; try < 10; try++ {
|
||||
shuffled := slices.Clone(sl)
|
||||
Shuffle(shuffled)
|
||||
if !reflect.DeepEqual(shuffled, sl) {
|
||||
wasShuffled = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if !wasShuffled {
|
||||
t.Errorf("expected shuffle after 10 tries")
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue