From a6a53e1fb9cf062bebc4f72785fd8dfdde9a14b2 Mon Sep 17 00:00:00 2001
From: pigwantacat <79680042+pigwantacat@users.noreply.github.com>
Date: Fri, 20 Sep 2024 06:16:24 +0800
Subject: [PATCH] refactor:refactor RandomString function (#524)

* refactor:refactor RandomString function

* feat:Upgrade math/rand to math/rand/v2

* Update string.go

Co-authored-by: ccoVeille <3875889+ccoVeille@users.noreply.github.com>

* feat: improve comments for RandomString function

* Update string.go

Co-authored-by: Samuel Berthe <dev@samuel-berthe.fr>

* feat:Upgrade math/rand/v2 to github.com/samber/lo/internal/rand

* Update string.go

* Update ordered_go118.go

---------

Co-authored-by: zhuhebin <zhuhebin@fengtaisec.com>
Co-authored-by: ccoVeille <3875889+ccoVeille@users.noreply.github.com>
Co-authored-by: Samuel Berthe <dev@samuel-berthe.fr>
---
 internal/rand/ordered_go118.go | 12 ++++++++
 internal/rand/ordered_go122.go |  4 +++
 string.go                      | 56 +++++++++++++++++++++++++++++-----
 3 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/internal/rand/ordered_go118.go b/internal/rand/ordered_go118.go
index a31bb9f2..9fbc5385 100644
--- a/internal/rand/ordered_go118.go
+++ b/internal/rand/ordered_go118.go
@@ -12,3 +12,15 @@ func IntN(n int) int {
 	// bearer:disable go_gosec_crypto_weak_random
 	return rand.Intn(n)
 }
+
+func Int64() int64 {
+	// bearer:disable go_gosec_crypto_weak_random
+	n := rand.Int63()
+    
+	// bearer:disable go_gosec_crypto_weak_random
+	if rand.Intn(2) == 0 {
+		return -n
+	}
+
+	return n
+}
diff --git a/internal/rand/ordered_go122.go b/internal/rand/ordered_go122.go
index 532ed339..65abf51a 100644
--- a/internal/rand/ordered_go122.go
+++ b/internal/rand/ordered_go122.go
@@ -11,3 +11,7 @@ func Shuffle(n int, swap func(i, j int)) {
 func IntN(n int) int {
 	return rand.IntN(n)
 }
+
+func Int64() int64 {
+	return rand.Int64()
+}
diff --git a/string.go b/string.go
index 94dddde4..51b09899 100644
--- a/string.go
+++ b/string.go
@@ -1,13 +1,13 @@
 package lo
 
 import (
+	"github.com/samber/lo/internal/rand"
+	"math"
 	"regexp"
 	"strings"
 	"unicode"
 	"unicode/utf8"
 
-	"github.com/samber/lo/internal/rand"
-
 	"golang.org/x/text/cases"
 	"golang.org/x/text/language"
 )
@@ -25,6 +25,7 @@ var (
 	splitWordReg = regexp.MustCompile(`([a-z])([A-Z0-9])|([a-zA-Z])([0-9])|([0-9])([a-zA-Z])|([A-Z])([A-Z])([a-z])`)
 	// bearer:disable go_lang_permissive_regex_validation
 	splitNumberLetterReg = regexp.MustCompile(`([0-9])([a-zA-Z])`)
+	maximumCapacity      = math.MaxInt>>1 + 1
 )
 
 // RandomString return a random string.
@@ -37,12 +38,53 @@ func RandomString(size int, charset []rune) string {
 		panic("lo.RandomString: Charset parameter must not be empty")
 	}
 
-	b := make([]rune, size)
-	possibleCharactersCount := len(charset)
-	for i := range b {
-		b[i] = charset[rand.IntN(possibleCharactersCount)]
+	// see https://stackoverflow.com/questions/22892120/how-to-generate-a-random-string-of-a-fixed-length-in-go
+	sb := strings.Builder{}
+	sb.Grow(size)
+	// Calculate the number of bits required to represent the charset,
+	// e.g., for 62 characters, it would need 6 bits (since 62 -> 64 = 2^6)
+	letterIdBits := int(math.Log2(float64(nearestPowerOfTwo(len(charset)))))
+	// Determine the corresponding bitmask,
+	// e.g., for 62 characters, the bitmask would be 111111.
+	var letterIdMask int64 = 1<<letterIdBits - 1
+	// Available count, since rand.Int64() returns a non-negative number, the first bit is fixed, so there are 63 random bits
+	// e.g., for 62 characters, this value is 10 (63 / 6).
+	letterIdMax := 63 / letterIdBits
+	// Generate the random string in a loop.
+	for i, cache, remain := size-1, rand.Int64(), letterIdMax; i >= 0; {
+		// Regenerate the random number if all available bits have been used
+		if remain == 0 {
+			cache, remain = rand.Int64(), letterIdMax
+		}
+		// Select a character from the charset
+		if idx := int(cache & letterIdMask); idx < len(charset) {
+			sb.WriteRune(charset[idx])
+			i--
+		}
+		// Shift the bits to the right to prepare for the next character selection,
+		// e.g., for 62 characters, shift by 6 bits.
+		cache >>= letterIdBits
+		// Decrease the remaining number of uses for the current random number.
+		remain--
 	}
-	return string(b)
+	return sb.String()
+}
+
+// nearestPowerOfTwo returns the nearest power of two.
+func nearestPowerOfTwo(cap int) int {
+	n := cap - 1
+	n |= n >> 1
+	n |= n >> 2
+	n |= n >> 4
+	n |= n >> 8
+	n |= n >> 16
+	if n < 0 {
+		return 1
+	}
+	if n >= maximumCapacity {
+		return maximumCapacity
+	}
+	return n + 1
 }
 
 // Substring return part of a string.