Merge pull request nsqio#1019 from ploxiln/uniq_rands_v2

internal/util: new UniqRands() implementation

Author: ploxiln

internal/util/rand.go

@@ -4,19 +4,22 @@ import (
 	"math/rand"
 )
 
-func UniqRands(l int, n int) []int {
-	set := make(map[int]struct{})
-	nums := make([]int, 0, l)
-	for {
-		num := rand.Intn(n)
-		if _, ok := set[num]; !ok {
-			set[num] = struct{}{}
-			nums = append(nums, num)
-		}
-		if len(nums) == l {
-			goto exit
-		}
+func UniqRands(quantity int, maxval int) []int {
+	if maxval < quantity {
+		quantity = maxval
 	}
-exit:
-	return nums
+
+	intSlice := make([]int, maxval)
+	for i := 0; i < maxval; i++ {
+		intSlice[i] = i
+	}
+
+	for i := 0; i < quantity; i++ {
+		j := rand.Int()%maxval + i
+		// swap
+		intSlice[i], intSlice[j] = intSlice[j], intSlice[i]
+		maxval--
+
+	}
+	return intSlice[0:quantity]
 }

internal/util/util_test.go

@@ -0,0 +1,36 @@
+package util
+
+import (
+	"testing"
+
+	"github.com/nsqio/nsq/internal/test"
+)
+
+func BenchmarkUniqRands5of5(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		UniqRands(5, 5)
+	}
+}
+func BenchmarkUniqRands20of20(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		UniqRands(20, 20)
+	}
+}
+
+func BenchmarkUniqRands20of50(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		UniqRands(20, 50)
+	}
+}
+
+func TestUniqRands(t *testing.T) {
+	var x []int
+	x = UniqRands(3, 10)
+	test.Equal(t, 3, len(x))
+
+	x = UniqRands(10, 5)
+	test.Equal(t, 5, len(x))
+
+	x = UniqRands(10, 20)
+	test.Equal(t, 10, len(x))
+}