diff --git a/benchmark/slice_benchmark_test.go b/benchmark/slice_benchmark_test.go index 8df0acd80..7a0d7f374 100644 --- a/benchmark/slice_benchmark_test.go +++ b/benchmark/slice_benchmark_test.go @@ -7,6 +7,7 @@ import ( "testing" "github.com/samber/lo" + "github.com/samber/lo/mutable" ) var lengths = []int{10, 100, 1000} @@ -248,3 +249,25 @@ func BenchmarkFilterTakeVsFilterAndTake(b *testing.B) { } }) } + +func BenchmarkShuffle(b *testing.B) { + for _, n := range lengths { + ints := genSliceInt(n) + b.Run(fmt.Sprintf("ints_%d", n), func(b *testing.B) { + b.ReportAllocs() // This reports memory allocations + for i := 0; i < b.N; i++ { + mutable.Shuffle(ints) + } + }) + } + + for _, n := range lengths { + strs := genSliceString(n) + b.Run(fmt.Sprintf("strings_%d", n), func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + mutable.Shuffle(strs) + } + }) + } +} diff --git a/internal/xrand/ordered_go118.go b/internal/xrand/ordered_go118.go index 410a5f342..b6aac42ec 100644 --- a/internal/xrand/ordered_go118.go +++ b/internal/xrand/ordered_go118.go @@ -2,10 +2,44 @@ package xrand -import "math/rand" +import ( + "math/rand" + "sync" +) + +var ( + mu sync.Mutex + seededRand *rand.Rand +) + +// SetSeed sets a custom seed for the random number generator. +// This allows for reproducible random sequences, which is useful for testing. +// Pass a negative value to reset to the default (non-reproducible) behavior. +func SetSeed(seed int64) { + mu.Lock() + defer mu.Unlock() + + if seed < 0 { + seededRand = nil + return + } + + //nolint:staticcheck // rand.NewSource is fine here for backward compatibility + seededRand = rand.New(rand.NewSource(seed)) +} + +// ResetSeed resets the random number generator to its default (non-reproducible) behavior. +func ResetSeed() { + SetSeed(-1) +} // Shuffle returns a slice of shuffled values. Uses the Fisher-Yates shuffle algorithm. func Shuffle(n int, swap func(i, j int)) { + if seededRand != nil { + seededRand.Shuffle(n, swap) + return + } + rand.Shuffle(n, swap) } @@ -13,13 +47,27 @@ func Shuffle(n int, swap func(i, j int)) { // from the default Source. // It panics if n <= 0. func IntN(n int) int { + if seededRand != nil { + return seededRand.Intn(n) + } + // bearer:disable go_gosec_crypto_weak_random return rand.Intn(n) } -// Int64 returns a non-negative pseudo-random 63-bit integer as an int64 +// Int64 returns a pseudo-random 63-bit integer as an int64 // from the default Source. +// For Go < 1.22, this simulates the full int64 range by randomly +// negating the result of Int63(). func Int64() int64 { + if seededRand != nil { + n := seededRand.Int63() + if seededRand.Intn(2) == 0 { + return -n + } + return n + } + // bearer:disable go_gosec_crypto_weak_random n := rand.Int63() diff --git a/internal/xrand/ordered_go122.go b/internal/xrand/ordered_go122.go index 6bca250e2..b30074710 100644 --- a/internal/xrand/ordered_go122.go +++ b/internal/xrand/ordered_go122.go @@ -2,10 +2,44 @@ package xrand -import "math/rand/v2" +import ( + "math/rand/v2" + "sync" +) + +var ( + mu sync.Mutex + seededRand *rand.Rand +) + +// SetSeed sets a custom seed for the random number generator. +// This allows for reproducible random sequences, which is useful for testing. +// Pass a negative value to reset to the default (non-reproducible) behavior. +func SetSeed(seed int64) { + mu.Lock() + defer mu.Unlock() + + if seed < 0 { + seededRand = nil + return + } + + uSeed := uint64(seed) + seededRand = rand.New(rand.NewPCG(uSeed, 0)) +} + +// ResetSeed resets the random number generator to its default (non-reproducible) behavior. +func ResetSeed() { + SetSeed(-1) +} // Shuffle returns a slice of shuffled values. Uses the Fisher-Yates shuffle algorithm. func Shuffle(n int, swap func(i, j int)) { + if seededRand != nil { + seededRand.Shuffle(n, swap) + return + } + rand.Shuffle(n, swap) } @@ -13,11 +47,19 @@ func Shuffle(n int, swap func(i, j int)) { // from the default Source. // It panics if n <= 0. func IntN(n int) int { + if seededRand != nil { + return seededRand.IntN(n) + } + return rand.IntN(n) } // Int64 returns a non-negative pseudo-random 63-bit integer as an int64 // from the default Source. func Int64() int64 { + if seededRand != nil { + return seededRand.Int64() + } + return rand.Int64() } diff --git a/rand.go b/rand.go new file mode 100644 index 000000000..5be8e2568 --- /dev/null +++ b/rand.go @@ -0,0 +1,38 @@ +package lo + +import "github.com/samber/lo/internal/xrand" + +// SetRandomSeed sets a custom seed for the random number generator used by +// RandomString, Shuffle, and other functions that rely on random numbers. +// This allows for reproducible random sequences, which is useful for testing. +// +// Pass a negative value to reset to the default (non-reproducible) behavior. +// +// Example: +// +// lo.SetRandomSeed(42) +// s1 := lo.RandomString(10, lo.AlphanumericCharset) // Always same result with seed 42 +// +// lo.SetRandomSeed(42) +// s2 := lo.RandomString(10, lo.AlphanumericCharset) // s1 == s2 +// +// lo.ResetRandomSeed() // Back to default random behavior +// +// Note: This function is NOT safe to call concurrently with functions that use +// the random number generator. It is intended to be called once at the start +// of a test or program. +func SetRandomSeed(seed int64) { + xrand.SetSeed(seed) +} + +// ResetRandomSeed resets the random number generator to its default +// (non-reproducible) behavior. This is equivalent to calling SetRandomSeed(-1). +// +// Example: +// +// lo.SetRandomSeed(42) +// // ... do some reproducible random operations ... +// lo.ResetRandomSeed() // Back to default random behavior +func ResetRandomSeed() { + xrand.ResetSeed() +} diff --git a/rand_test.go b/rand_test.go new file mode 100644 index 000000000..e0247a704 --- /dev/null +++ b/rand_test.go @@ -0,0 +1,109 @@ +package lo + +import ( + "testing" + + "github.com/samber/lo/mutable" + "github.com/stretchr/testify/assert" +) + +func TestSetRandomSeed(t *testing.T) { //nolint:paralleltest + // t.Parallel() + t.Cleanup(func() { + ResetRandomSeed() + }) + + is := assert.New(t) + + t.Run("reproducible RandomString", func(t *testing.T) { + SetRandomSeed(42) + s1 := RandomString(20, AlphanumericCharset) + + SetRandomSeed(42) + s2 := RandomString(20, AlphanumericCharset) + + is.Equal(s1, s2, "RandomString should produce the same result with the same seed") + }) + + t.Run("different seeds produce different results", func(t *testing.T) { + SetRandomSeed(42) + s1 := RandomString(20, AlphanumericCharset) + + SetRandomSeed(123) + s2 := RandomString(20, AlphanumericCharset) + + is.NotEqual(s1, s2, "Different seeds should produce different results") + }) + + t.Run("reproducible Shuffle", func(t *testing.T) { + original := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10} + + SetRandomSeed(42) + slice1 := make([]int, len(original)) + copy(slice1, original) + mutable.Shuffle(slice1) + + SetRandomSeed(42) + slice2 := make([]int, len(original)) + copy(slice2, original) + mutable.Shuffle(slice2) + + is.Equal(slice1, slice2, "mutable.Shuffle should produce the same result with the same seed") + }) + + t.Run("reset returns to non-reproducible behavior", func(t *testing.T) { + SetRandomSeed(42) + s1 := RandomString(20, AlphanumericCharset) + + ResetRandomSeed() + s2 := RandomString(20, AlphanumericCharset) + s3 := RandomString(20, AlphanumericCharset) + + // After reset, consecutive calls should produce different results (with very high probability) + // Note: There's an astronomically small chance this could fail if truly random + is.NotEqual(s2, s3, "After reset, RandomString should produce different results") + + // And it should not match the seeded result + SetRandomSeed(42) + s4 := RandomString(20, AlphanumericCharset) + is.Equal(s1, s4, "Re-seeding with 42 should reproduce the original result") + }) + + t.Run("negative seed resets", func(t *testing.T) { + SetRandomSeed(42) + s1 := RandomString(20, AlphanumericCharset) + + SetRandomSeed(-1) // Reset via negative seed + + SetRandomSeed(42) + s2 := RandomString(20, AlphanumericCharset) + + is.Equal(s1, s2, "Negative seed should reset, then re-seeding should work") + }) +} + +func TestSetRandomSeed_MultipleOperations(t *testing.T) { //nolint:paralleltest + // t.Parallel() + t.Cleanup(func() { + ResetRandomSeed() + }) + + is := assert.New(t) + + // Test that a sequence of operations is reproducible + SetRandomSeed(999) + results1 := []string{ + RandomString(10, AlphanumericCharset), + RandomString(5, LowerCaseLettersCharset), + RandomString(15, NumbersCharset), + } + + SetRandomSeed(999) + results2 := []string{ + RandomString(10, AlphanumericCharset), + RandomString(5, LowerCaseLettersCharset), + RandomString(15, NumbersCharset), + } + + is.Equal(results1, results2, "A sequence of operations should be reproducible") +}