Snowchecked.hs

{-# LANGUAGE BlockArguments  #-}
{-# LANGUAGE RecordWildCards #-}
{-|
Description : Unique id generator derived from Twitter's Snowflake.
License     : Apache 2.0
Maintainer  : smokejumperit@gmail.com
Stability   : experimental

This generates unique (guaranteed) identifiers build from a timestamp,
counter, and node id.  Identifiers are convertible to values which are
monotonically increasing with respect to time.
-}
module Data.Snowchecked
( newSnowcheckedGen
, nextFlake
, SnowcheckedConfig(..)
, SnowcheckedGen
, Flake
, snowcheckedConfigBitCount
, uniqueFlakeCount
) where

import           Control.Concurrent.MVar
import           Control.Monad.IO.Class           (MonadIO, liftIO)
import           Data.Snowchecked.Internal.Import
import           Data.Time.Clock.POSIX            (getPOSIXTime)


currentTimestamp :: IO Word256
currentTimestamp = toMillisWord256 <$> getPOSIXTime
	where
		toMillisWord256 = round . (*1000)
{-# INLINE currentTimestamp #-}

currentTimestampBits :: Word8 -> IO Word256
currentTimestampBits n = 
	if n == 0 then
		return 0
	else
		(`cutBits` toInt n) <$> currentTimestamp
{-# INLINE currentTimestampBits #-}

-- | Create a new generator. Takes a configuration and node id.  The node id may be any
-- value that fits in a 'Word256', but it will be truncated to the number of bits specified
-- in the provided configuration.
newSnowcheckedGen :: (MonadIO io) => SnowcheckedConfig -> Word256 -> io SnowcheckedGen
newSnowcheckedGen conf@SnowcheckedConfig{..} nodeId = liftIO $
	SnowcheckedGen <$> newMVar Flake
		{ flakeTime = 0
		, flakeCount = 0
		, flakeNodeId = cutBits nodeId (toInt confNodeBits)
		, flakeConfig = conf
		}
{-# INLINEABLE newSnowcheckedGen #-}
{-# SPECIALIZE newSnowcheckedGen :: SnowcheckedConfig -> Word256 -> IO SnowcheckedGen #-}

-- | Calculates the number of bits in each 'Flake' generated using a given configuration.
-- It returns a 'Word32' because there are 4 fields and the bitlength of each field fits
-- in a 'Word8', so the total bit count must fit within a 'Word32'.
snowcheckedConfigBitCount :: SnowcheckedConfig -> Word32
snowcheckedConfigBitCount SnowcheckedConfig{..} = foldr foldFunc 0
		[ confTimeBits
		, confCountBits
		, confNodeBits
		, confCheckBits
		]
	where
		foldFunc :: Word8 -> Word32 -> Word32
		foldFunc nxt = (+ toWord32 nxt)
{-# INLINEABLE snowcheckedConfigBitCount #-}

-- | Generates the next id.
nextFlake :: (MonadIO io) => SnowcheckedGen -> io Flake
nextFlake SnowcheckedGen{..} = liftIO $ modifyMVar genLastFlake mkNextFlake
	where
		-- TODO: Track the number of flakes generated and error out if we've exhausted them.
		mkNextFlake flake@Flake{..} =
			let SnowcheckedConfig{..} = flakeConfig in
			currentTimestampBits confTimeBits >>= \currentTimeBits ->
				if flakeTime < currentTimeBits then
					let newFlake = flake { 
						flakeTime = currentTimeBits, 
						flakeCount = 0 
					} in return (newFlake, newFlake)
				else if confTimeBits == 0 then
					let newFlake = flake
						{ flakeTime = 0
						, flakeCount = flakeCount + 1
						}
					in return (newFlake, newFlake)
				else if confCountBits == 0 then
					let newFlake = flake
						{ flakeTime = flakeTime + 1
						, flakeCount = 0
						}
					in return (newFlake, newFlake)
				else
					let nextCount = cutBits (flakeCount + 1) (toInt confCountBits) in
					if nextCount == 0 then
						let newFlake = flake
							{ flakeTime = flakeTime + 1
							, flakeCount = 0
							}
						in return (newFlake, newFlake)
					else
						let newFlake = flake { flakeCount = nextCount }
						in return (newFlake, newFlake)
{-# INLINEABLE nextFlake #-}
{-# SPECIALIZE nextFlake :: SnowcheckedGen -> IO Flake #-}


-- | Provides the count of total number of unique flakes possibly generated by a node using
-- this configuration.
uniqueFlakeCount :: SnowcheckedConfig -> Integer
uniqueFlakeCount SnowcheckedConfig{..}
	| confCountBits == 0 = 2^confTimeBits
	| confTimeBits  == 0 = 2^confCountBits
	| otherwise          = 2^confCountBits * 2^confTimeBits
{-# INLINE uniqueFlakeCount #-}