charset.go

package id

import (
	crand "crypto/rand"
	"errors"
	"io"
	"math/bits"
	"strings"
)

// New is a convenient package level function that generate an id with the default charset.
func New(n int) (string, error) {
	return defaultCharset.ID(n)
}

// New64 is a convenient package level function that generate an id with the charset 64.
func New64(n int) (string, error) {
	return charset64.ID(n)
}

// New32 is a convenient package level function that generate an id with the charset 32.
func New32(n int) (string, error) {
	return charset32.ID(n)
}

// New16 is a convenient package level function that generate an id with the charset 16.
func New16(n int) (string, error) {
	return charset16.ID(n)
}

// Must is a convenient package level function that generate an id with the default charset.
// It panics if any error occur.
func Must(n int) string {
	return defaultCharset.Must(n)
}

// Must64 is a convenient package level function that generate an id with the charset 64.
// It panics if any error occur.
func Must64(n int) string {
	return charset64.Must(n)
}

// Must32 is a convenient package level function that generate an id with the charset 32.
// It panics if any error occur.
func Must32(n int) string {
	return charset32.Must(n)
}

// Must16 is a convenient package level function that generate an id with the charset 16.
// It panics if any error occur.
func Must16(n int) string {
	return charset16.Must(n)
}

// SetDefault sets the default charset for package level functions.
// It is not safe for concurrent use.
func SetDefault(c *Charset) {
	defaultCharset = c
}

// charsets for package level functions.
var (
	charset64      = MustNewCharset(Charset64())
	charset32      = MustNewCharset(Charset32())
	charset16      = MustNewCharset(Charset16())
	defaultCharset = charset64
)

// Charset represents the allowed characters to generate ids.
// Fields are all private, you must use NewCharset to create a new charset.
type Charset struct {
	charset []rune
	mask    int
}

// ErrInvalidCharset is returned when the charset is not valid.
// A valid charset must not contains duplicate and the bitmask should not exceed the charset size.
var ErrInvalidCharset = errors.New("invalid charset")

// NewCharset returns a new and ready to use charset to generate ids.
// It returns an error if the charset is not valid.
// A valid charset must not contains duplicate and the bitmask should not exceed the charset size.
// Also the charset size must be a power of two.
func NewCharset(charset []rune) (*Charset, error) {
	// First a bitmask is necessary to generate ids afterwards. The bitmask makes bytes
	// values closer to the charset size. The bitmask calculates the closest `2^31 - 1` number.
	l := len(charset) - 1
	mask := (2 << (31 - bits.LeadingZeros32(uint32(l)|1))) - 1
	// Ensure that the bitmask do not exceed the charset size
	// to avoid extra magic or extra random generator call in ID method.
	if mask > l {
		return nil, ErrInvalidCharset
	}
	// Ensure that the charset contains no duplicate.
	chars := make(map[rune]struct{}, len(charset))
	for _, c := range charset {
		if _, ok := chars[c]; ok {
			return nil, ErrInvalidCharset
		}
		chars[c] = struct{}{}
	}
	return &Charset{
		charset: charset,
		mask:    mask,
	}, nil
}

// MustNewCharset is equivalent to NewCharset but panics if any error occurs.
func MustNewCharset(charset []rune) *Charset {
	c, err := NewCharset(charset)
	if err != nil {
		panic(err)
	}
	return c
}

// ID generates an ID with flat characters distribution.
func (c *Charset) ID(n int) (string, error) {
	buf := make([]byte, n)
	_, err := io.ReadFull(crand.Reader, buf)
	if err != nil {
		return "", err
	}
	// using string.Builder instead of []byte avoid one alloc.
	b := strings.Builder{}
	b.Grow(n)
	for i := 0; i < n; i++ {
		// The bitmask cannot exceed the charset size since this is ensured by
		// the charset creation (if not this could result in a out of range panic).
		// So adding hacks, such as empty string fallback or magic numbers
		// is unneccessary because the bitmask trims bytes down to the alphabet size.
		b.WriteRune(c.charset[int(buf[i])&c.mask])
	}
	return b.String(), nil
}

// Must generates an ID and panics if any error occur.
func (c *Charset) Must(n int) string {
	id, err := c.ID(n)
	if err != nil {
		panic(err)
	}
	return id
}

// Len returns the length of the charset.
func (c *Charset) Len() int {
	return len(c.charset)
}

// Charset16 returns a valid charset with 16 characters that can be used with NewCharset.
// All the characters in this charset are URL friendly.
func Charset16() []rune {
	return []rune{
		'0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
		'a', 'b', 'c', 'd', 'e', 'f',
	}
}

// Charset32 returns a valid charset with 32 characters that can be used with NewCharset.
// All the charaters in this charset are URL friendly.
func Charset32() []rune {
	return []rune{
		'0', '1', '2', '3', '4', '5',
		'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
		'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',
	}
}

// Charset64 returns a valid charset with 64 characters that can be used with NewCharset.
// All the charaters in this charset are URL friendly.
func Charset64() []rune {
	return []rune{
		'_', '-', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
		'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
		'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',
		'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M',
		'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
	}
}