From a685984f2e394da50aeed9aa08133ff107c19a0b Mon Sep 17 00:00:00 2001
From: Bruce Marriner <bruce@sqls.net>
Date: Mon, 19 Mar 2018 21:07:06 +0000
Subject: [PATCH] Allow custom NodeBits and StepBits to be set.

You can now set the snowflake.NodeBits or snowflake.StepBits to
custom values.  Please keep in mind you have a total of 22 bits
available to use between these two values.  Setting these to
inappropriate values may break things.  So take care to understand
what you're doing here.

This is not the best way of doing this but it does not break
API compatibility with existing users of this package. I will
release a version 2 of this package which will implement this
feature better.
---
 snowflake.go | 40 ++++++++++++++++++++++++++--------------
 1 file changed, 26 insertions(+), 14 deletions(-)

diff --git a/snowflake.go b/snowflake.go
index e560164..4cc5c16 100644
--- a/snowflake.go
+++ b/snowflake.go
@@ -11,15 +11,24 @@ import (
 	"time"
 )
 
-const (
-	spareBits       = 22
-	nodeBits        = 10
-	stepBits        = spareBits - nodeBits
-	nodeMax         = -1 ^ (-1 << nodeBits)
-	nodeMask        = nodeMax << stepBits
-	stepMask  int64 = -1 ^ (-1 << stepBits)
-	timeShift uint8 = nodeBits + stepBits
-	nodeShift uint8 = stepBits
+var (
+	// Epoch is set to the twitter snowflake epoch of Nov 04 2010 01:42:54 UTC
+	// You may customize this to set a different epoch for your application.
+	Epoch int64 = 1288834974657
+
+	// Number of bits to use for Node
+	// Remember, you have a total 22 bits to share between Node/Step
+	NodeBits uint8 = 10
+
+	// Number of bits to use for Step
+	// Remember, you have a total 22 bits to share between Node/Step
+	StepBits uint8 = 12
+
+	nodeMax   int64 = -1 ^ (-1 << NodeBits)
+	nodeMask  int64 = nodeMax << StepBits
+	stepMask  int64 = -1 ^ (-1 << StepBits)
+	timeShift uint8 = NodeBits + StepBits
+	nodeShift uint8 = StepBits
 )
 
 const encodeBase32Map = "ybndrfg8ejkmcpqxot1uwisza345h769"
@@ -63,10 +72,6 @@ var ErrInvalidBase58 = errors.New("invalid base58")
 // ErrInvalidBase32 is returned by ParseBase32 when given an invalid []byte
 var ErrInvalidBase32 = errors.New("invalid base32")
 
-// Epoch is set to the twitter snowflake epoch of 2006-03-21:20:50:14 GMT
-// You may customize this to set a different epoch for your application.
-var Epoch int64 = 1288834974657
-
 // A Node struct holds the basic information needed for a snowflake generator
 // node
 type Node struct {
@@ -88,6 +93,13 @@ func NewNode(node int64) (*Node, error) {
 		return nil, errors.New("Node number must be between 0 and 1023")
 	}
 
+	// re-calc in case custom NodeBits or StepBits were set
+	nodeMax = -1 ^ (-1 << NodeBits)
+	nodeMask = nodeMax << StepBits
+	stepMask = -1 ^ (-1 << StepBits)
+	timeShift = NodeBits + StepBits
+	nodeShift = StepBits
+
 	return &Node{
 		time: 0,
 		node: node,
@@ -242,7 +254,7 @@ func (f ID) IntBytes() [8]byte {
 
 // Time returns an int64 unix timestamp of the snowflake ID time
 func (f ID) Time() int64 {
-	return (int64(f) >> 22) + Epoch
+	return (int64(f) >> timeShift) + Epoch
 }
 
 // Node returns an int64 of the snowflake ID node number