-- |
-- 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.
--
module Data.UUID.V7
  (
  -- * Data type
    UUID
  -- * 'UUID'v7 generation
  , 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           Data.UUID.Versions
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 = [UUID] -> UUID
forall a. HasCallStack => [a] -> a
head ([UUID] -> UUID) -> m [UUID] -> m UUID
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word16 -> m [UUID]
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 <- m Word64
forall (m :: * -> *). MonadIO m => m Word64
getEpochMilli
  Word16
entropy16 <- m Word16
forall (m :: * -> *). MonadIO m => m Word16
getEntropyWord16
  Word64
entropy64 <- m Word64
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
  UUID -> m UUID
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (UUID -> m UUID)
-> ((Word64, Word64) -> UUID) -> (Word64, Word64) -> m UUID
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word64 -> Word64 -> UUID) -> (Word64, Word64) -> UUID
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Word64 -> Word64 -> UUID
UUID ((Word64, Word64) -> m UUID) -> (Word64, Word64) -> m UUID
forall a b. (a -> b) -> a -> b
$ Get (Word64, Word64) -> ByteString -> (Word64, Word64)
forall a. Get a -> ByteString -> a
runGet ((Get Word64 -> Get Word64 -> Get (Word64, Word64))
-> Get Word64 -> Get (Word64, Word64)
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join ((Word64 -> Word64 -> (Word64, Word64))
-> Get Word64 -> Get Word64 -> Get (Word64, Word64)
forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (,)) Get Word64
getWord64be) ByteString
bs
{-# INLINE genUUID' #-}

-- | 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 = [UUID] -> m [UUID]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
genUUIDs Word16
n = IO [UUID] -> m [UUID]
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
  Word64
timestamp <- IO Word64
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 <- (Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.&. Word16
0x7FFF) (Word16 -> Word16) -> IO Word16 -> IO Word16
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Word16
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 b -> b -> b
forall a. Num a => a -> a -> a
- b
seqNo b -> b -> Bool
forall a. Ord a => a -> a -> Bool
< b
num
        then (b
0xFFFF b -> b -> b
forall a. Num a => a -> a -> a
- b
seqNo, b
0xFFFF)
        else (b
num, b
seqNo b -> b -> b
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)  <- IORef (Word64, Word16)
-> ((Word64, Word16) -> ((Word64, Word16), (Word16, Word16)))
-> IO (Word16, Word16)
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef (Word64, Word16)
__state__ (((Word64, Word16) -> ((Word64, Word16), (Word16, Word16)))
 -> IO (Word16, Word16))
-> ((Word64, Word16) -> ((Word64, Word16), (Word16, Word16)))
-> IO (Word16, Word16)
forall a b. (a -> b) -> a -> b
$ \(Word64
ts, Word16
seqNo) -> if
    | Word64
ts Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
timestamp -> let (Word16
n', Word16
entropy16') = Word16 -> Word16 -> (Word16, Word16)
forall {b}. (Ord b, Num b) => b -> b -> (b, b)
getMaxSlots Word16
n Word16
entropy16
                        in  ((Word64
timestamp, Word16
entropy16'), (Word16
n', Word16
entropy16 Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
+ Word16
1))
    | Word64
ts Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> Word64
timestamp -> ((Word64
ts, Word16
seqNo), (Word16
0, Word16
0))
    | Bool
otherwise      -> let (Word16
n', Word16
entropy16') = Word16 -> Word16 -> (Word16, Word16)
forall {b}. (Ord b, Num b) => b -> b -> (b, b)
getMaxSlots Word16
n Word16
seqNo
                        in  ((Word64
timestamp, Word16
entropy16'), (Word16
n', Word16
seqNo Word16 -> Word16 -> Word16
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' Word16 -> Word16 -> Bool
forall a. Eq a => a -> a -> Bool
== Word16
0
    then Word16 -> IO [UUID]
forall (m :: * -> *). MonadIO m => Word16 -> m [UUID]
genUUIDs Word16
n
    else do
      [UUID]
uuids <- [Word16] -> (Word16 -> IO UUID) -> IO [UUID]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Word16
0..(Word16
n' Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
- Word16
1)] ((Word16 -> IO UUID) -> IO [UUID])
-> (Word16 -> IO UUID) -> IO [UUID]
forall a b. (a -> b) -> a -> b
$ \Word16
curN -> do
        Word64
entropy64 <- IO Word64
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 Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
+ Word16
curN)
              Word16 -> Word64 -> Put
fillVarAndRandB (Word16
seqNo Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
+ Word16
curN) Word64
entropy64
        UUID -> IO UUID
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (UUID -> IO UUID)
-> ((Word64, Word64) -> UUID) -> (Word64, Word64) -> IO UUID
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word64 -> Word64 -> UUID) -> (Word64, Word64) -> UUID
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Word64 -> Word64 -> UUID
UUID ((Word64, Word64) -> IO UUID) -> (Word64, Word64) -> IO UUID
forall a b. (a -> b) -> a -> b
$ Get (Word64, Word64) -> ByteString -> (Word64, Word64)
forall a. Get a -> ByteString -> a
runGet ((Get Word64 -> Get Word64 -> Get (Word64, Word64))
-> Get Word64 -> Get (Word64, Word64)
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join ((Word64 -> Word64 -> (Word64, Word64))
-> Get Word64 -> Get Word64 -> Get (Word64, Word64)
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' Word16 -> Word16 -> Bool
forall a. Eq a => a -> a -> Bool
== Word16
n then [UUID] -> IO [UUID]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [UUID]
uuids else ([UUID]
uuids [UUID] -> [UUID] -> [UUID]
forall a. [a] -> [a] -> [a]
++) ([UUID] -> [UUID]) -> IO [UUID] -> IO [UUID]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word16 -> IO [UUID]
forall (m :: * -> *). MonadIO m => Word16 -> m [UUID]
genUUIDs (Word16
n Word16 -> Word16 -> Word16
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 -> UUIDVersion -> Bool) -> UUIDVersion -> UUID -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip UUID -> UUIDVersion -> Bool
validateWithVersion UUIDVersion
V7
{-# 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 <- m Word64
forall (m :: * -> *). MonadIO m => m Word64
getEpochMilli
  Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ UUID -> Bool
validate UUID
uuid Bool -> Bool -> Bool
&& (UUID -> Word64
getTime UUID
uuid Word64 -> Word64 -> Bool
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 = IO Word64 -> m Word64
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
  POSIXTime
t <- IO POSIXTime
getPOSIXTime
  Word64 -> IO Word64
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word64 -> IO Word64)
-> (POSIXTime -> Word64) -> POSIXTime -> IO Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. POSIXTime -> Word64
forall b. Integral b => POSIXTime -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (POSIXTime -> IO Word64) -> POSIXTime -> IO Word64
forall a b. (a -> b) -> a -> b
$ POSIXTime
t POSIXTime -> POSIXTime -> POSIXTime
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 Word64 -> Int -> Word64
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__ = IO (IORef (Word64, Word16)) -> IORef (Word64, Word16)
forall a. IO a -> a
unsafePerformIO ((Word64, Word16) -> IO (IORef (Word64, Word16))
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
  (Word16 -> Put) -> [Word16] -> Put
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 Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
`shiftR` Int
4
  let randAWithVer :: Word16
randAWithVer = Word16
seqNoRandA Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.|. (Word16
0x7 Word16 -> Int -> Word16
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 Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.&. Word16
0xF
  let randBWithVar :: Word64
randBWithVar = Word16 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word16
seqNoRandB Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.|. (Word16
0x2 Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
`shiftL` Int
4))
  Word64 -> Put
putWord64be (Word64 -> Put) -> Word64 -> Put
forall a b. (a -> b) -> a -> b
$ (Word64
entropy Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
0x3FFFFFFFFFFFFFF) Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Word64
randBWithVar Word64 -> Int -> Word64
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 = Word64 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
n Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
0xFFFF)
      b1 :: Word16
b1 = Word64 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word64
n Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
16) Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
0xFFFF)
      b2 :: Word16
b2 = Word64 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word64
n Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
32) Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
0xFFFF)
      b3 :: Word16
b3 = Word64 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word64
n Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
48) Word64 -> Word64 -> Word64
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 = IO Word16 -> m Word16
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
  ByteString
bs <- ByteString -> ByteString
BSL.fromStrict (ByteString -> ByteString) -> IO ByteString -> IO ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO ByteString
getEntropy Int
2
  Word16 -> IO Word16
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word16 -> IO Word16) -> Word16 -> IO Word16
forall a b. (a -> b) -> a -> b
$ Get Word16 -> ByteString -> Word16
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 = IO Word64 -> m Word64
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
  ByteString
bs <- ByteString -> ByteString
BSL.fromStrict (ByteString -> ByteString) -> IO ByteString -> IO ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO ByteString
getEntropy Int
8
  Word64 -> IO Word64
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word64 -> IO Word64) -> Word64 -> IO Word64
forall a b. (a -> b) -> a -> b
$ Get Word64 -> ByteString -> Word64
forall a. Get a -> ByteString -> a
runGet Get Word64
getWord64host ByteString
bs
{-# INLINE getEntropyWord64 #-}