From 513938f8e67dc6393ecdd7e0103ba2e8215c01c3 Mon Sep 17 00:00:00 2001
From: Andrey Vyazigin <duhaaaa@yandex.ru>
Date: Mon, 18 May 2020 13:55:33 +0700
Subject: [PATCH] Generate valid id for any custom alphabet (#15)

* Generate valid id for any alphabet
* Add more complex examples
---
 examples/simple_example.go | 22 ++++++++++++++++++++++
 gonanoid.go                | 12 +++++++-----
 gonanoid_test.go           |  2 +-
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/examples/simple_example.go b/examples/simple_example.go
index 0692d26..181e86a 100644
--- a/examples/simple_example.go
+++ b/examples/simple_example.go
@@ -7,9 +7,31 @@ import (
 )
 
 func main() {
+	// Simple usage
 	id, err := gonanoid.Nanoid()
 	if err != nil {
 		panic(err)
 	}
 	fmt.Printf("Generated id: %s\n", id)
+
+	// Custom length
+	id, err = gonanoid.ID(5)
+	if err != nil {
+		panic(err)
+	}
+	fmt.Printf("Generated id: %s\n", id)
+
+	// Custom alphabet
+	id, err = gonanoid.Generate("abcdefg", 10)
+	if err != nil {
+		panic(err)
+	}
+	fmt.Printf("Generated id: %s\n", id)
+
+	// Custom non ascii alphabet
+	id, err = gonanoid.Generate("こちんにабдежиклмнは你好喂שלום😯😪🥱😌😛äöüß", 10)
+	if err != nil {
+		panic(err)
+	}
+	fmt.Printf("Generated id: %s\n", id)
 }
diff --git a/gonanoid.go b/gonanoid.go
index 884a55b..2440cd1 100644
--- a/gonanoid.go
+++ b/gonanoid.go
@@ -7,8 +7,8 @@ import (
 	"math"
 )
 
+var defaultAlphabet = []rune("_-0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
 const (
-	defaultAlphabet = "_-0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" // len=64
 	defaultSize     = 21
 	defaultMaskSize = 5
 )
@@ -34,7 +34,7 @@ func initMasks(params ...int) []uint {
 	return masks
 }
 
-func getMask(alphabet string, masks []uint) int {
+func getMask(alphabet []rune, masks []uint) int {
 	for i := 0; i < len(masks); i++ {
 		curr := int(masks[i])
 		if curr >= len(alphabet)-1 {
@@ -45,7 +45,9 @@ func getMask(alphabet string, masks []uint) int {
 }
 
 // Generate is a low-level function to change alphabet and ID size.
-func Generate(alphabet string, size int) (string, error) {
+func Generate(rawAlphabet string, size int) (string, error) {
+	alphabet := []rune(rawAlphabet)
+
 	if len(alphabet) == 0 || len(alphabet) > 255 {
 		return "", fmt.Errorf("alphabet must not empty and contain no more than 255 chars. Current len is %d", len(alphabet))
 	}
@@ -58,7 +60,7 @@ func Generate(alphabet string, size int) (string, error) {
 	ceilArg := 1.6 * float64(mask*size) / float64(len(alphabet))
 	step := int(math.Ceil(ceilArg))
 
-	id := make([]byte, size)
+	id := make([]rune, size)
 	bytes := make([]byte, step)
 	for j := 0; ; {
 		_, err := BytesGenerator(bytes)
@@ -97,7 +99,7 @@ func Nanoid(param ...int) (string, error) {
 	if err != nil {
 		return "", err
 	}
-	id := make([]byte, size)
+	id := make([]rune, size)
 	for i := 0; i < size; i++ {
 		id[i] = defaultAlphabet[bytes[i]&63]
 	}
diff --git a/gonanoid_test.go b/gonanoid_test.go
index 97d7e5f..e400c98 100644
--- a/gonanoid_test.go
+++ b/gonanoid_test.go
@@ -28,7 +28,7 @@ func TestGeneratesURLFriendlyIDs(t *testing.T) {
 		runeID := []rune(id)
 
 		for j := 0; j < len(runeID); j++ {
-			res := strings.Contains(defaultAlphabet, string(runeID[j]))
+			res := strings.Contains(string(defaultAlphabet), string(runeID[j]))
 			if !res {
 				t.Errorf(
 					"GeneratesURLFriendlyIds error: char %v should be contained in %v",