{-# LANGUAGE BlockArguments   #-}
{-# LANGUAGE RecordWildCards  #-}
{-# LANGUAGE TypeApplications #-}
{-|
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               (threadDelay)
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 :: IO Word256
currentTimestamp = POSIXTime -> Word256
toMillisWord256 (POSIXTime -> Word256) -> IO POSIXTime -> IO Word256
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO POSIXTime
getPOSIXTime
	where
		toMillisWord256 :: POSIXTime -> Word256
toMillisWord256 = POSIXTime -> Word256
forall a b. (RealFrac a, Integral b) => a -> b
round (POSIXTime -> Word256)
-> (POSIXTime -> POSIXTime) -> POSIXTime -> Word256
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (POSIXTime -> POSIXTime -> POSIXTime
forall a. Num a => a -> a -> a
*POSIXTime
1000)
{-# INLINE currentTimestamp #-}

currentTimestampBits :: Word8 -> IO Word256
currentTimestampBits :: Word8 -> IO Word256
currentTimestampBits Word8
n = (Word256 -> Integer -> Word256
forall a bitCount.
(Num a, Bits a, Integral bitCount) =>
a -> bitCount -> a
`cutBits` Word8 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
n) (Word256 -> Word256) -> IO Word256 -> IO Word256
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Word256
currentTimestamp
{-# INLINE currentTimestampBits #-}

-- | Create a new generator. Takes a configuration and node id.
newSnowcheckedGen :: (MonadIO io) => SnowcheckedConfig -> Word256 -> io SnowcheckedGen
newSnowcheckedGen :: SnowcheckedConfig -> Word256 -> io SnowcheckedGen
newSnowcheckedGen conf :: SnowcheckedConfig
conf@SnowcheckedConfig{Word8
confCheckBits :: SnowcheckedConfig -> Word8
confNodeBits :: SnowcheckedConfig -> Word8
confCountBits :: SnowcheckedConfig -> Word8
confTimeBits :: SnowcheckedConfig -> Word8
confCheckBits :: Word8
confNodeBits :: Word8
confCountBits :: Word8
confTimeBits :: Word8
..} Word256
nodeId = IO SnowcheckedGen -> io SnowcheckedGen
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO SnowcheckedGen -> io SnowcheckedGen)
-> IO SnowcheckedGen -> io SnowcheckedGen
forall a b. (a -> b) -> a -> b
$ do
	Word256
startTimeBits <- Word8 -> IO Word256
currentTimestampBits Word8
confTimeBits
	MVar Flake -> SnowcheckedGen
SnowcheckedGen (MVar Flake -> SnowcheckedGen)
-> IO (MVar Flake) -> IO SnowcheckedGen
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Flake -> IO (MVar Flake)
forall a. a -> IO (MVar a)
newMVar Flake :: Word256 -> Word256 -> Word256 -> SnowcheckedConfig -> Flake
Flake
		{ flakeTime :: Word256
flakeTime = Word256
startTimeBits
		, flakeCount :: Word256
flakeCount = Word256
0
		, flakeNodeId :: Word256
flakeNodeId = Word256 -> Word8 -> Word256
forall a bitCount.
(Num a, Bits a, Integral bitCount) =>
a -> bitCount -> a
cutBits Word256
nodeId Word8
confNodeBits
		, flakeConfig :: SnowcheckedConfig
flakeConfig = SnowcheckedConfig
conf
		}
{-# INLINEABLE newSnowcheckedGen #-}
{-# SPECIALIZE newSnowcheckedGen :: SnowcheckedConfig -> Word256 -> IO SnowcheckedGen #-}

snowcheckedConfigBitCount :: SnowcheckedConfig -> Word32
snowcheckedConfigBitCount :: SnowcheckedConfig -> Word32
snowcheckedConfigBitCount SnowcheckedConfig{Word8
confCheckBits :: Word8
confNodeBits :: Word8
confCountBits :: Word8
confTimeBits :: Word8
confCheckBits :: SnowcheckedConfig -> Word8
confNodeBits :: SnowcheckedConfig -> Word8
confCountBits :: SnowcheckedConfig -> Word8
confTimeBits :: SnowcheckedConfig -> Word8
..} = (Word8 -> Word32 -> Word32) -> Word32 -> [Word8] -> Word32
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Word8 -> Word32 -> Word32
foldFunc Word32
0
		[ Word8
confTimeBits
		, Word8
confCountBits
		, Word8
confNodeBits
		, Word8
confCheckBits
		]
	where
		foldFunc :: Word8 -> Word32 -> Word32
		foldFunc :: Word8 -> Word32 -> Word32
foldFunc Word8
nxt Word32
memo = Word32
memo Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word8 -> Word32
forall a. Integral a => a -> Word32
toWord32 Word8
nxt
{-# INLINEABLE snowcheckedConfigBitCount #-}

-- | Generates the next id.
nextFlake :: (MonadIO io) => SnowcheckedGen -> io Flake
nextFlake :: SnowcheckedGen -> io Flake
nextFlake SnowcheckedGen{MVar Flake
genLastFlake :: SnowcheckedGen -> MVar Flake
genLastFlake :: MVar Flake
..} = IO Flake -> io Flake
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Flake -> io Flake) -> IO Flake -> io Flake
forall a b. (a -> b) -> a -> b
$ MVar Flake -> (Flake -> IO (Flake, Flake)) -> IO Flake
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar MVar Flake
genLastFlake Flake -> IO (Flake, Flake)
mkNextFlake
	where
		-- TODO: Special case when confTimeBits is 0.
		-- TODO: Track the number of flakes generated and error out if we've exhausted them.
		mkNextFlake :: Flake -> IO (Flake, Flake)
mkNextFlake flake :: Flake
flake@Flake{Word256
SnowcheckedConfig
flakeConfig :: SnowcheckedConfig
flakeNodeId :: Word256
flakeCount :: Word256
flakeTime :: Word256
flakeConfig :: Flake -> SnowcheckedConfig
flakeNodeId :: Flake -> Word256
flakeCount :: Flake -> Word256
flakeTime :: Flake -> Word256
..} =
			let SnowcheckedConfig{Word8
confCheckBits :: Word8
confNodeBits :: Word8
confCountBits :: Word8
confTimeBits :: Word8
confCheckBits :: SnowcheckedConfig -> Word8
confNodeBits :: SnowcheckedConfig -> Word8
confCountBits :: SnowcheckedConfig -> Word8
confTimeBits :: SnowcheckedConfig -> Word8
..} = SnowcheckedConfig
flakeConfig in
			Word8 -> IO Word256
currentTimestampBits Word8
confTimeBits IO Word256 -> (Word256 -> IO (Flake, Flake)) -> IO (Flake, Flake)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Word256
currentTimeBits ->
				if Word256
flakeTime Word256 -> Word256 -> Bool
forall a. Ord a => a -> a -> Bool
< Word256
currentTimeBits then
					let newFlake :: Flake
newFlake = Flake
flake
						{ flakeTime :: Word256
flakeTime = Word256
currentTimeBits
						, flakeCount :: Word256
flakeCount = Word256
0
						}
					in (Flake, Flake) -> IO (Flake, Flake)
forall (m :: * -> *) a. Monad m => a -> m a
return (Flake
newFlake, Flake
newFlake)
				else if Word8
confCountBits Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0 then
					Int -> IO ()
threadDelay Int
1000 IO () -> IO (Flake, Flake) -> IO (Flake, Flake)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Flake -> IO (Flake, Flake)
mkNextFlake Flake
flake
				else
					let nextCount :: Word256
nextCount = Word256 -> Word8 -> Word256
forall a bitCount.
(Num a, Bits a, Integral bitCount) =>
a -> bitCount -> a
cutBits (Word256
flakeCount Word256 -> Word256 -> Word256
forall a. Num a => a -> a -> a
+ Word256
1) Word8
confCountBits in
					if Word256
flakeCount Word256 -> Word256 -> Bool
forall a. Ord a => a -> a -> Bool
< Word256
nextCount then
						let newFlake :: Flake
newFlake = Flake
flake { flakeCount :: Word256
flakeCount = Word256
nextCount }
						in (Flake, Flake) -> IO (Flake, Flake)
forall (m :: * -> *) a. Monad m => a -> m a
return (Flake
newFlake, Flake
newFlake)
					else
						-- The count wrapped and we need to wait for the time to change.
						-- This assumes that the next millisecond will give us a new time.
						Int -> IO ()
threadDelay Int
1000 IO () -> IO (Flake, Flake) -> IO (Flake, Flake)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Flake -> IO (Flake, Flake)
mkNextFlake Flake
flake
{-# INLINEABLE nextFlake #-}
{-# SPECIALIZE nextFlake :: SnowcheckedGen -> IO Flake #-}


-- | Provides the count of total number of unique flakes possibly generated by this configuration.
uniqueFlakeCount :: SnowcheckedConfig -> Integer
uniqueFlakeCount :: SnowcheckedConfig -> Integer
uniqueFlakeCount SnowcheckedConfig{Word8
confCheckBits :: Word8
confNodeBits :: Word8
confCountBits :: Word8
confTimeBits :: Word8
confCheckBits :: SnowcheckedConfig -> Word8
confNodeBits :: SnowcheckedConfig -> Word8
confCountBits :: SnowcheckedConfig -> Word8
confTimeBits :: SnowcheckedConfig -> Word8
..} = Word8 -> Integer
forall a. Integral a => a -> Integer
toInteger Word8
confCountBits Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Word8 -> Integer
forall a. Integral a => a -> Integer
toInteger Word8
confTimeBits
{-# INLINE uniqueFlakeCount #-}