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