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: I0b00920e5b3eea846b6cedc30bd34d978a049fd3main
Andrew Dunham
3 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