-- |
-- Module      : Data.UUID.V7
-- License     : MIT
-- Maintainer  : mmzk1526@outlook.com
-- Portability : GHC
--
-- UUIDv7 implementation.
--
-- UUIDv7 is not currently present in the uuid package, therefore I have to
-- make a quick patch of my own.
--
-- Note that since the specification for v7 is not yet finalised, this module's
-- implementation may change in the future according to the potential
-- adjustments in the specification.
--
-- WARNING: The 'nil' re-export will be removed in the next major version.
--
module Data.UUID.V7
  (
  -- * Data type
    UUID
  -- * 'UUID'v7 generation
  , nil
  , genUUID
  , genUUID'
  , genUUIDs
  -- Validation
  , validate
  , validateWithTime
  -- * Miscellaneous helpers
  , getTime
  , getEpochMilli
  ) where

import           Control.Monad
import           Control.Monad.IO.Class
import           Data.Binary
import           Data.Binary.Get
import           Data.Binary.Put
import           Data.Bits
import qualified Data.ByteString.Lazy as BSL
import           Data.IORef
import           Data.Time.Clock.POSIX
import           Data.UUID.Types.Internal
import           System.Entropy
import           System.IO.Unsafe (unsafePerformIO)

-- | Generate a 'UUID'v7.
genUUID :: MonadIO m => m UUID
genUUID :: forall (m :: * -> *). MonadIO m => m UUID
genUUID = forall a. [a] -> a
head forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadIO m => Word16 -> m [UUID]
genUUIDs Word16
1
{-# INLINE genUUID #-}

-- | Generate a stateless 'UUID'v7.
--
-- It is faster than 'genUUID' but it is not guaranteed to be monotonically
-- increasing if multiple 'UUID's are generated at the same timestamp.
--
-- In use cases where the ordering is not important, this function is could be
-- preferred.
genUUID' :: MonadIO m => m UUID
genUUID' :: forall (m :: * -> *). MonadIO m => m UUID
genUUID' = do
  Word64
timestamp <- forall (m :: * -> *). MonadIO m => m Word64
getEpochMilli
  Word16
entropy16 <- forall (m :: * -> *). MonadIO m => m Word16
getEntropyWord16
  Word64
entropy64 <- forall (m :: * -> *). MonadIO m => m Word64
getEntropyWord64
  let bs :: ByteString
bs = Put -> ByteString
runPut do
        Word64 -> Put
fillTime Word64
timestamp
        Word16 -> Put
fillVerAndRandA Word16
entropy16
        Word16 -> Word64 -> Put
fillVarAndRandB Word16
entropy16 Word64
entropy64
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Word64 -> Word64 -> UUID
UUID forall a b. (a -> b) -> a -> b
$ forall a. Get a -> ByteString -> a
runGet (forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (,)) Get Word64
getWord64be) ByteString
bs

-- | Generate a list of 'UUID'v7s.
--
-- It tries its best to generate 'UUID's at the same timestamp, but it may not
-- be possible if we are asking too many 'UUID's at the same time.
--
-- It is guaranteed that the first 32768 'UUID's are generated at the same
-- timestamp.
genUUIDs :: MonadIO m => Word16 -> m [UUID]
genUUIDs :: forall (m :: * -> *). MonadIO m => Word16 -> m [UUID]
genUUIDs Word16
0 = forall (f :: * -> *) a. Applicative f => a -> f a
pure []
genUUIDs Word16
n = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
  Word64
timestamp <- forall (m :: * -> *). MonadIO m => m Word64
getEpochMilli
  -- We set the first bit of the entropy to 0 to ensure that there's enough
  -- room for incrementing the sequence number.
  Word16
entropy16 <- (forall a. Bits a => a -> a -> a
.&. Word16
0x7FFF) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadIO m => m Word16
getEntropyWord16
  -- Calculate the maximum number of slots we can use for the current timestamp
  -- before the sequence number overflows.
  let getMaxSlots :: b -> b -> (b, b)
getMaxSlots b
num b
seqNo = if b
0xFFFF forall a. Num a => a -> a -> a
- b
seqNo forall a. Ord a => a -> a -> Bool
< b
num
        then (b
0xFFFF forall a. Num a => a -> a -> a
- b
seqNo, b
0xFFFF)
        else (b
num, b
seqNo forall a. Num a => a -> a -> a
+ b
num)
  -- Get the sequence number corresponding to the current timestamp and the
  -- number of UUIDs we can generate.
  (Word16
n', Word16
seqNo)  <- forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef (Word64, Word16)
__state__ forall a b. (a -> b) -> a -> b
$ \(Word64
ts, Word16
seqNo) -> if
    | Word64
ts forall a. Ord a => a -> a -> Bool
< Word64
timestamp -> let (Word16
n', Word16
entropy16') = forall {b}. (Ord b, Num b) => b -> b -> (b, b)
getMaxSlots Word16
n Word16
entropy16
                        in  ((Word64
timestamp, Word16
entropy16'), (Word16
n', Word16
entropy16 forall a. Num a => a -> a -> a
+ Word16
1))
    | Word64
ts forall a. Ord a => a -> a -> Bool
> Word64
timestamp -> ((Word64
ts, Word16
seqNo), (Word16
0, Word16
0))
    | Bool
otherwise      -> let (Word16
n', Word16
entropy16') = forall {b}. (Ord b, Num b) => b -> b -> (b, b)
getMaxSlots Word16
n Word16
seqNo
                        in  ((Word64
timestamp, Word16
entropy16'), (Word16
n', Word16
seqNo forall a. Num a => a -> a -> a
+ Word16
1))
  -- If we can't generate any UUIDs, we try again, hoping that the timestamp
  -- has changed.
  if Word16
n' forall a. Eq a => a -> a -> Bool
== Word16
0
    then forall (m :: * -> *). MonadIO m => Word16 -> m [UUID]
genUUIDs Word16
n
    else do
      [UUID]
uuids <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Word16
0..(Word16
n' forall a. Num a => a -> a -> a
- Word16
1)] forall a b. (a -> b) -> a -> b
$ \Word16
curN -> do
        Word64
entropy64 <- forall (m :: * -> *). MonadIO m => m Word64
getEntropyWord64
        let bs :: ByteString
bs = Put -> ByteString
runPut do
              Word64 -> Put
fillTime Word64
timestamp
              Word16 -> Put
fillVerAndRandA (Word16
seqNo forall a. Num a => a -> a -> a
+ Word16
curN)
              Word16 -> Word64 -> Put
fillVarAndRandB (Word16
seqNo forall a. Num a => a -> a -> a
+ Word16
curN) Word64
entropy64
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Word64 -> Word64 -> UUID
UUID forall a b. (a -> b) -> a -> b
$ forall a. Get a -> ByteString -> a
runGet (forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (,)) Get Word64
getWord64be) ByteString
bs
      if Word16
n' forall a. Eq a => a -> a -> Bool
== Word16
n then forall (f :: * -> *) a. Applicative f => a -> f a
pure [UUID]
uuids else ([UUID]
uuids forall a. [a] -> [a] -> [a]
++) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadIO m => Word16 -> m [UUID]
genUUIDs (Word16
n forall a. Num a => a -> a -> a
- Word16
n')

-- | Validate the version and variant of the 'UUID'v7.
validate :: UUID -> Bool
validate :: UUID -> Bool
validate (UUID Word64
w1 Word64
w2)
  = (Word64
w1 forall a. Bits a => a -> Int -> a
`shiftR` Int
12) forall a. Bits a => a -> a -> a
.&. Word64
0xF forall a. Eq a => a -> a -> Bool
== Word64
0x7 Bool -> Bool -> Bool
&& (Word64
w2 forall a. Bits a => a -> Int -> a
`shiftR` Int
62) forall a. Bits a => a -> a -> a
.&. Word64
0x3 forall a. Eq a => a -> a -> Bool
== Word64
0x2
{-# INLINE validate #-}

-- | Validate the version and variant of the 'UUID'v7 as well as its timestamp
-- is no greater than the current time.
validateWithTime :: MonadIO m => UUID -> m Bool
validateWithTime :: forall (m :: * -> *). MonadIO m => UUID -> m Bool
validateWithTime UUID
uuid = do
  Word64
curTime <- forall (m :: * -> *). MonadIO m => m Word64
getEpochMilli
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ UUID -> Bool
validate UUID
uuid Bool -> Bool -> Bool
&& (UUID -> Word64
getTime UUID
uuid forall a. Ord a => a -> a -> Bool
<= Word64
curTime)
{-# INLINE validateWithTime #-}
  
-- | Get the current time in milliseconds since the Unix epoch.
getEpochMilli :: MonadIO m => m Word64
getEpochMilli :: forall (m :: * -> *). MonadIO m => m Word64
getEpochMilli = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
  POSIXTime
t <- IO POSIXTime
getPOSIXTime
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (RealFrac a, Integral b) => a -> b
round forall a b. (a -> b) -> a -> b
$ POSIXTime
t forall a. Num a => a -> a -> a
* POSIXTime
1000
{-# INLINE getEpochMilli #-}

-- | Get the time field (unix_ts_ms) of a 'UUID'v7.
getTime :: UUID -> Word64
getTime :: UUID -> Word64
getTime (UUID Word64
w1 Word64
_) = Word64
w1 forall a. Bits a => a -> Int -> a
`shiftR` Int
16
{-# INLINE getTime #-}

-- | The global mutable state of (timestamp, sequence number).
--
-- The "NOINLINE" pragma is IMPORTANT! The logic would be flawed if it is
-- is inlined by its definition.
__state__ :: IORef (Word64, Word16)
__state__ :: IORef (Word64, Word16)
__state__ = forall a. IO a -> a
unsafePerformIO (forall a. a -> IO (IORef a)
newIORef (Word64
0, Word16
0))
{-# NOINLINE __state__ #-}

-- | Fill in the 48-bit time field (unix_ts_ms) of a 'UUID'v7 with the given
-- time.
fillTime :: Word64 -> Put
fillTime :: Word64 -> Put
fillTime Word64
timestamp = do
  let (Word16
_, Word16
p2, Word16
p1, Word16
p0) = Word64 -> (Word16, Word16, Word16, Word16)
splitWord64ToWord16s Word64
timestamp
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Word16 -> Put
putWord16be [Word16
p2, Word16
p1, Word16
p0]
{-# INLINE fillTime #-}

-- | Fill in the version and rand_a part of a 'UUID'v7 with the given sequence
-- number.
--
-- The sequence number is a 16-bit integer, of which the first 12 bits are used
-- here in rand_a, and the last 4 bits are used in rand_b. The version is 7.
fillVerAndRandA :: Word16 -> Put
fillVerAndRandA :: Word16 -> Put
fillVerAndRandA Word16
seqNo = do
  let seqNoRandA :: Word16
seqNoRandA   = Word16
seqNo forall a. Bits a => a -> Int -> a
`shiftR` Int
4
  let randAWithVer :: Word16
randAWithVer = Word16
seqNoRandA forall a. Bits a => a -> a -> a
.|. (Word16
0x7 forall a. Bits a => a -> Int -> a
`shiftL` Int
12)
  Word16 -> Put
putWord16be Word16
randAWithVer
{-# INLINE fillVerAndRandA #-}

-- | Fill in the variant and rand_b part of a 'UUID'v7 with the given sequence
-- number and random number. The variant is 2.
--
-- The sequence number is a 16-bit integer, of which the last 4 bits are used
-- here in rand_b while the first 12 bits are used in rand_a.
--
-- The random number is a 64-bit integer of which the last 58 bits are used
-- while the rest are replaced by the variant bits and the last 4 bits of the
-- sequence number.
fillVarAndRandB :: Word16 -> Word64 -> Put
fillVarAndRandB :: Word16 -> Word64 -> Put
fillVarAndRandB Word16
seqNo Word64
entropy = do
  let seqNoRandB :: Word16
seqNoRandB   = Word16
seqNo forall a. Bits a => a -> a -> a
.&. Word16
0xF
  let randBWithVar :: Word64
randBWithVar = forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word16
seqNoRandB forall a. Bits a => a -> a -> a
.|. (Word16
0x2 forall a. Bits a => a -> Int -> a
`shiftL` Int
4))
  Word64 -> Put
putWord64be forall a b. (a -> b) -> a -> b
$ (Word64
entropy forall a. Bits a => a -> a -> a
.&. Word64
0x3FFFFFFFFFFFFFF) forall a. Bits a => a -> a -> a
.|. (Word64
randBWithVar forall a. Bits a => a -> Int -> a
`shiftL` Int
58)
{-# INLINE fillVarAndRandB #-}

splitWord64ToWord16s :: Word64 -> (Word16, Word16, Word16, Word16)
splitWord64ToWord16s :: Word64 -> (Word16, Word16, Word16, Word16)
splitWord64ToWord16s Word64
n =
  let b0 :: Word16
b0 = forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
n forall a. Bits a => a -> a -> a
.&. Word64
0xFFFF)
      b1 :: Word16
b1 = forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word64
n forall a. Bits a => a -> Int -> a
`shiftR` Int
16) forall a. Bits a => a -> a -> a
.&. Word64
0xFFFF)
      b2 :: Word16
b2 = forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word64
n forall a. Bits a => a -> Int -> a
`shiftR` Int
32) forall a. Bits a => a -> a -> a
.&. Word64
0xFFFF)
      b3 :: Word16
b3 = forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word64
n forall a. Bits a => a -> Int -> a
`shiftR` Int
48) forall a. Bits a => a -> a -> a
.&. Word64
0xFFFF)
  in (Word16
b3, Word16
b2, Word16
b1, Word16
b0)
{-# INLINE splitWord64ToWord16s #-}

getEntropyWord16 :: MonadIO m => m Word16
getEntropyWord16 :: forall (m :: * -> *). MonadIO m => m Word16
getEntropyWord16 = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
  ByteString
bs <- ByteString -> ByteString
BSL.fromStrict forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO ByteString
getEntropy Int
2
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. Get a -> ByteString -> a
runGet Get Word16
getWord16host ByteString
bs
{-# INLINE getEntropyWord16 #-}

getEntropyWord64 :: MonadIO m => m Word64
getEntropyWord64 :: forall (m :: * -> *). MonadIO m => m Word64
getEntropyWord64 = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
  ByteString
bs <- ByteString -> ByteString
BSL.fromStrict forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO ByteString
getEntropy Int
8
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. Get a -> ByteString -> a
runGet Get Word64
getWord64host ByteString
bs
{-# INLINE getEntropyWord64 #-}