From 56479224d3912172d2582d832151bb027e3868bd Mon Sep 17 00:00:00 2001
From: bddjr <bddjrcn@outlook.com>
Date: Mon, 16 Dec 2024 22:47:25 +0800
Subject: [PATCH] Optimize

---
 rand.go | 97 ++++++++++++++++++++++++++++-----------------------------
 1 file changed, 47 insertions(+), 50 deletions(-)

diff --git a/rand.go b/rand.go
index 523eca5..37f72bf 100644
--- a/rand.go
+++ b/rand.go
@@ -3,6 +3,7 @@ package cryptorandstr
 import (
 	"crypto/rand"
 	"fmt"
+	"unsafe"
 )
 
 const Chars64 = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ-_"
@@ -31,7 +32,7 @@ func MustRand64(strLen int) string {
 	return MustRand(strLen, Chars64)
 }
 
-// panic when error.
+// Panic when error.
 func MustRand(strLen int, chars string) string {
 	s, err := Rand(strLen, chars)
 	if err != nil {
@@ -67,27 +68,57 @@ func Rand(strLen int, chars string) (string, error) {
 		return "", fmt.Errorf("cryptorandstr: chars length %d > 256", len(chars))
 	}
 
-	chunk := make([]byte, 1)
-	maxB := byte(len(chars) - 1)
-	r := randWithGC{
-		outBitLen: bitLen(maxB),
-	}
-	difference := 8 - r.outBitLen
-
-	out := make([]byte, strLen)
-	i := 0
-	for i < strLen {
-		b, err := r.byte(chunk, difference)
-		if err != nil {
-			return "", err
+	var (
+		out = make([]byte, strLen)
+		i   int
+
+		chunk = make([]byte, 1)
+		n     int
+		err   error
+
+		cache       uint16
+		cacheBitLen byte
+
+		maxB       = byte(len(chars) - 1)
+		outBitLen  = bitLen(maxB)
+		difference = 8 - outBitLen
+		b          byte
+	)
+
+	for {
+		if cacheBitLen >= outBitLen {
+			// read cache
+			b = byte(cache)
+			cache >>= outBitLen
+			cacheBitLen -= outBitLen
+		} else {
+			// random
+			for {
+				n, err = rand.Reader.Read(chunk)
+				if err != nil {
+					return "", err
+				}
+				if n == 1 {
+					break
+				}
+			}
+			b = chunk[0]
+			// write cache
+			cache <<= difference
+			cache |= uint16(b >> outBitLen)
+			cacheBitLen += difference
 		}
+
+		b <<= difference
+		b >>= difference
 		if b <= maxB {
 			out[i] = chars[b]
 			i++
+			if i == strLen {
+				return *(*string)(unsafe.Pointer(&out)), nil
+			}
 		}
 	}
-
-	return string(out), nil
 }
 
 func bitLen(b byte) (l byte) {
@@ -97,37 +128,3 @@ func bitLen(b byte) (l byte) {
 	}
 	return
 }
-
-type randWithGC struct {
-	outBitLen   byte
-	cacheBitLen byte
-	cache       uint16
-}
-
-func (r *randWithGC) byte(chunk []byte, difference byte) (byte, error) {
-	// read cache
-	if r.cacheBitLen >= r.outBitLen {
-		b := (byte(r.cache) << difference) >> difference
-		r.cache >>= r.outBitLen
-		r.cacheBitLen -= r.outBitLen
-		return b, nil
-	}
-
-	// random
-	for {
-		n, err := rand.Reader.Read(chunk)
-		if err != nil {
-			return 0, err
-		}
-		if n == 1 {
-			break
-		}
-	}
-
-	// write cache
-	r.cache <<= difference
-	r.cache |= uint16(chunk[0] >> r.outBitLen)
-	r.cacheBitLen += difference
-
-	return (chunk[0] << difference) >> difference, nil
-}