diff options
Diffstat (limited to 'go/snowflake.go')
| -rw-r--r-- | go/snowflake.go | 196 |
1 files changed, 196 insertions, 0 deletions
diff --git a/go/snowflake.go b/go/snowflake.go new file mode 100644 index 0000000..92a6124 --- /dev/null +++ b/go/snowflake.go @@ -0,0 +1,196 @@ +// Package snowflake generates random, URL-safe identifiers. +// +// Each ID is 16 bytes of cryptographically secure randomness, encoded as a +// base62 string. IDs implement encoding.TextMarshaler, json.Marshaler, +// and database/sql.Scanner/Valuer for convenient serialization. +package snowflake + +import ( + "crypto/rand" + "database/sql/driver" + "encoding/hex" + "errors" + "fmt" +) + +const ( + // IDByteSize is the size of a Snowflake ID in bytes. + IDByteSize = 16 + // Alphabet is the base62 character set used for encoding. + Alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" +) + +var ( + ErrInvalidIDLength = errors.New("snowflake: invalid ID byte length") + ErrInvalidEncoding = errors.New("snowflake: invalid encoding string") + decodeLUT [256]int8 +) + +func init() { + for i := range decodeLUT { + decodeLUT[i] = -1 + } + + for i := 0; i < len(Alphabet); i++ { + decodeLUT[Alphabet[i]] = int8(i) + } +} + +// Snowflake is a 16-byte random identifier. +type Snowflake struct { + bytes [IDByteSize]byte +} + +// New returns a new random Snowflake ID. +func New() Snowflake { + var s Snowflake + if _, err := rand.Read(s.bytes[:]); err != nil { + panic(err) + } + return s +} + +// Zero returns the zero value of Snowflake. +func Zero() Snowflake { + return Snowflake{} +} + +// IsZero reports whether s is the zero value. +func (s Snowflake) IsZero() bool { + return s == Snowflake{} +} + +// String returns the base62 encoding of the ID. +func (s Snowflake) String() string { + base := uint64(len(Alphabet)) + var res [64]byte + i := len(res) + + tmp := s.bytes + for { + var remainder uint64 + allZero := true + for j := 0; j < IDByteSize; j++ { + accumulator := uint64(tmp[j]) + (remainder << 8) + tmp[j] = byte(accumulator / base) + remainder = accumulator % base + if tmp[j] > 0 { + allZero = false + } + } + i-- + res[i] = Alphabet[remainder] + if allZero { + break + } + } + return string(res[i:]) +} + +// FromString parses a base62-encoded Snowflake ID. +func FromString(s string) (Snowflake, error) { + var res Snowflake + base := uint64(len(Alphabet)) + + for i := 0; i < len(s); i++ { + charVal := decodeLUT[s[i]] + if charVal < 0 { + return res, ErrInvalidEncoding + } + + carry := uint64(charVal) + for j := IDByteSize - 1; j >= 0; j-- { + accumulator := uint64(res.bytes[j])*base + carry + res.bytes[j] = byte(accumulator) + carry = accumulator >> 8 + } + + if carry > 0 { + return Snowflake{}, ErrInvalidEncoding + } + } + return res, nil +} + +// Hex returns the hex-encoded ID. +func (s Snowflake) Hex() string { + return hex.EncodeToString(s.bytes[:]) +} + +// FromHex parses a hex-encoded Snowflake ID. +func FromHex(h string) (Snowflake, error) { + decoded, err := hex.DecodeString(h) + if err != nil { + return Snowflake{}, err + } + if len(decoded) != IDByteSize { + return Snowflake{}, ErrInvalidIDLength + } + var s Snowflake + copy(s.bytes[:], decoded) + return s, nil +} + +// MarshalText implements encoding.TextMarshaler. +func (s Snowflake) MarshalText() ([]byte, error) { + return []byte(s.String()), nil +} + +// UnmarshalText implements encoding.TextUnmarshaler. +func (s *Snowflake) UnmarshalText(text []byte) error { + id, err := FromString(string(text)) + if err != nil { + return err + } + *s = id + return nil +} + +// MarshalJSON implements json.Marshaler. +func (s Snowflake) MarshalJSON() ([]byte, error) { + str := s.String() + res := make([]byte, 0, len(str)+2) + res = append(res, '"') + res = append(res, str...) + res = append(res, '"') + return res, nil +} + +// UnmarshalJSON implements json.Unmarshaler. +func (s *Snowflake) UnmarshalJSON(data []byte) error { + if len(data) < 2 || data[0] != '"' || data[len(data)-1] != '"' { + return ErrInvalidEncoding + } + return s.UnmarshalText(data[1 : len(data)-1]) +} + +// Value implements driver.Valuer for database/sql support. +func (s Snowflake) Value() (driver.Value, error) { + b := make([]byte, IDByteSize) + copy(b, s.bytes[:]) + return b, nil +} + +// Scan implements sql.Scanner for database/sql support. +func (s *Snowflake) Scan(value interface{}) error { + if value == nil { + s.bytes = [IDByteSize]byte{} + return nil + } + + var source []byte + switch t := value.(type) { + case []byte: + source = t + case string: + source = []byte(t) + default: + return fmt.Errorf("snowflake: cannot scan type %T", value) + } + + if len(source) != IDByteSize { + return ErrInvalidIDLength + } + copy(s.bytes[:], source) + return nil +} |