{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples #-}

-- | Contains a simulated-time implementaion. It can be used directly,
-- to implement the simulated versions of getCurrentTime and
-- threadDelay, or can be used via monad transformers defined in
-- `Control.Monad.Trans.SimulatedTime`
module Test.SimulatedTime
  ( TimeEnv,
    create,
    advance,
    triggerEvents,
    getSimulatedTime,
    threadDelay',
  )
where

import Control.Concurrent hiding (threadDelay)
import qualified Control.Concurrent
import Control.Concurrent.STM
import Control.Monad (unless, when)
import Data.Foldable (forM_)
import Data.Function (on)
import Data.Functor (void)
import Data.List (insertBy, partition)
import Data.Maybe (listToMaybe)
import Data.Time hiding (getCurrentTime)
import qualified Data.Time
import Data.Bifunctor (Bifunctor(first))

-- | It remembers the offset from system time and keeps a list of
-- sleeping threads (threads in 'threadDelay'' call)
data TimeEnv
  = TimeEnv
      { TimeEnv -> TVar NominalDiffTime
offset :: TVar NominalDiffTime,
        TimeEnv -> TVar [(UTCTime, MVar ())]
events :: TVar [(UTCTime, MVar ())]
      }

-- | Create the simulated time env from the given time start point. For example
--
-- @
-- timeEnv <- create (fromGregorian 2000 1 1) 0
-- getSimulatedTime timeEnv
-- @
create :: UTCTime -> IO TimeEnv
create :: UTCTime -> IO TimeEnv
create UTCTime
epoch = do
  UTCTime
now <- IO UTCTime
Data.Time.getCurrentTime
  TVar NominalDiffTime -> TVar [(UTCTime, MVar ())] -> TimeEnv
TimeEnv (TVar NominalDiffTime -> TVar [(UTCTime, MVar ())] -> TimeEnv)
-> IO (TVar NominalDiffTime)
-> IO (TVar [(UTCTime, MVar ())] -> TimeEnv)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NominalDiffTime -> IO (TVar NominalDiffTime)
forall a. a -> IO (TVar a)
newTVarIO (UTCTime -> UTCTime -> NominalDiffTime
diffUTCTime UTCTime
epoch UTCTime
now) IO (TVar [(UTCTime, MVar ())] -> TimeEnv)
-> IO (TVar [(UTCTime, MVar ())]) -> IO TimeEnv
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [(UTCTime, MVar ())] -> IO (TVar [(UTCTime, MVar ())])
forall a. a -> IO (TVar a)
newTVarIO []

-- | The current simulated time
getSimulatedTime :: TimeEnv -> IO UTCTime
getSimulatedTime :: TimeEnv -> IO UTCTime
getSimulatedTime TimeEnv
t = do
  UTCTime
now <- IO UTCTime
Data.Time.getCurrentTime
  NominalDiffTime
offset <- TVar NominalDiffTime -> IO NominalDiffTime
forall a. TVar a -> IO a
readTVarIO (TimeEnv -> TVar NominalDiffTime
offset TimeEnv
t)
  UTCTime -> IO UTCTime
forall (m :: * -> *) a. Monad m => a -> m a
return (UTCTime -> IO UTCTime) -> UTCTime -> IO UTCTime
forall a b. (a -> b) -> a -> b
$ NominalDiffTime -> UTCTime -> UTCTime
addUTCTime NominalDiffTime
offset UTCTime
now

-- | Wake up due sleeping threads, based on TimeEnv. This happens on
-- its own, but this call can quicken things up. Under the hood it is
-- just @advance tenv 0@
triggerEvents :: TimeEnv -> IO ()
triggerEvents :: TimeEnv -> IO ()
triggerEvents TimeEnv
t = TimeEnv -> NominalDiffTime -> IO ()
advance TimeEnv
t NominalDiffTime
0

-- | Move the simulated time by a delta
advance :: TimeEnv -> NominalDiffTime -> IO ()
advance :: TimeEnv -> NominalDiffTime -> IO ()
advance TimeEnv
t NominalDiffTime
diff = do
  UTCTime
now <- IO UTCTime
Data.Time.getCurrentTime
  [MVar ()]
due <- STM [MVar ()] -> IO [MVar ()]
forall a. STM a -> IO a
atomically (STM [MVar ()] -> IO [MVar ()]) -> STM [MVar ()] -> IO [MVar ()]
forall a b. (a -> b) -> a -> b
$ do
    Bool -> STM () -> STM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (NominalDiffTime
diff NominalDiffTime -> NominalDiffTime -> Bool
forall a. Eq a => a -> a -> Bool
== NominalDiffTime
0) (STM () -> STM ()) -> STM () -> STM ()
forall a b. (a -> b) -> a -> b
$ TVar NominalDiffTime
-> (NominalDiffTime -> NominalDiffTime) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar (TimeEnv -> TVar NominalDiffTime
offset TimeEnv
t) (NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
+ NominalDiffTime
diff)
    TVar [(UTCTime, MVar ())]
-> ([(UTCTime, MVar ())] -> [(UTCTime, MVar ())]) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar (TimeEnv -> TVar [(UTCTime, MVar ())]
events TimeEnv
t) (((UTCTime, MVar ()) -> (UTCTime, MVar ()))
-> [(UTCTime, MVar ())] -> [(UTCTime, MVar ())]
forall a b. (a -> b) -> [a] -> [b]
map ((UTCTime -> UTCTime) -> (UTCTime, MVar ()) -> (UTCTime, MVar ())
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (NominalDiffTime -> UTCTime -> UTCTime
addUTCTime (- NominalDiffTime
diff))))
    TimeEnv -> UTCTime -> STM [MVar ()]
extractDueEvents TimeEnv
t UTCTime
now
  [MVar ()] -> (MVar () -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [MVar ()]
due ((MVar () -> () -> IO ()) -> () -> MVar () -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar ())
  TimeEnv -> UTCTime -> IO ()
scheduleNextWakeUp TimeEnv
t UTCTime
now

-- the 'now' is the system time now, as opposed to simulated now
extractDueEvents :: TimeEnv -> UTCTime -> STM [MVar ()]
extractDueEvents :: TimeEnv -> UTCTime -> STM [MVar ()]
extractDueEvents TimeEnv
t UTCTime
now = do
  ([(UTCTime, MVar ())]
due, [(UTCTime, MVar ())]
inTheFuture) <- ((UTCTime, MVar ()) -> Bool)
-> [(UTCTime, MVar ())]
-> ([(UTCTime, MVar ())], [(UTCTime, MVar ())])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition ((UTCTime -> UTCTime -> Bool
forall a. Ord a => a -> a -> Bool
<= UTCTime
now) (UTCTime -> Bool)
-> ((UTCTime, MVar ()) -> UTCTime) -> (UTCTime, MVar ()) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (UTCTime, MVar ()) -> UTCTime
forall a b. (a, b) -> a
fst) ([(UTCTime, MVar ())]
 -> ([(UTCTime, MVar ())], [(UTCTime, MVar ())]))
-> STM [(UTCTime, MVar ())]
-> STM ([(UTCTime, MVar ())], [(UTCTime, MVar ())])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar [(UTCTime, MVar ())] -> STM [(UTCTime, MVar ())]
forall a. TVar a -> STM a
readTVar (TimeEnv -> TVar [(UTCTime, MVar ())]
events TimeEnv
t)
  TVar [(UTCTime, MVar ())] -> [(UTCTime, MVar ())] -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (TimeEnv -> TVar [(UTCTime, MVar ())]
events TimeEnv
t) [(UTCTime, MVar ())]
inTheFuture
  [MVar ()] -> STM [MVar ()]
forall (m :: * -> *) a. Monad m => a -> m a
return (((UTCTime, MVar ()) -> MVar ())
-> [(UTCTime, MVar ())] -> [MVar ()]
forall a b. (a -> b) -> [a] -> [b]
map (UTCTime, MVar ()) -> MVar ()
forall a b. (a, b) -> b
snd [(UTCTime, MVar ())]
due)

scheduleNextWakeUp :: TimeEnv -> UTCTime -> IO ()
scheduleNextWakeUp :: TimeEnv -> UTCTime -> IO ()
scheduleNextWakeUp TimeEnv
t UTCTime
now = do
  Maybe UTCTime
nextEvent <- [UTCTime] -> Maybe UTCTime
forall a. [a] -> Maybe a
listToMaybe ([UTCTime] -> Maybe UTCTime)
-> ([(UTCTime, MVar ())] -> [UTCTime])
-> [(UTCTime, MVar ())]
-> Maybe UTCTime
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((UTCTime, MVar ()) -> UTCTime)
-> [(UTCTime, MVar ())] -> [UTCTime]
forall a b. (a -> b) -> [a] -> [b]
map (UTCTime, MVar ()) -> UTCTime
forall a b. (a, b) -> a
fst ([(UTCTime, MVar ())] -> Maybe UTCTime)
-> IO [(UTCTime, MVar ())] -> IO (Maybe UTCTime)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar [(UTCTime, MVar ())] -> IO [(UTCTime, MVar ())]
forall a. TVar a -> IO a
readTVarIO (TimeEnv -> TVar [(UTCTime, MVar ())]
events TimeEnv
t)
  case Maybe UTCTime
nextEvent of
    Maybe UTCTime
Nothing -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    Just UTCTime
tm -> do
      let us :: Int
us = NominalDiffTime -> Int
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (UTCTime -> UTCTime -> NominalDiffTime
diffUTCTime UTCTime
tm UTCTime
now NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
* NominalDiffTime
1000000)
      IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO () -> IO ThreadId
forkIO (Int -> IO ()
Control.Concurrent.threadDelay Int
us IO () -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> TimeEnv -> IO ()
triggerEvents TimeEnv
t))

-- | Simulated alternative to `Control.Concurrent.threadDelay`, sleep
-- for the given number of microseconds. Threads wakes up according to
-- the simulated time.
threadDelay' :: TimeEnv -> Int -> IO ()
threadDelay' :: TimeEnv -> Int -> IO ()
threadDelay' TimeEnv
t Int
us = do
  MVar ()
mvar <- IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar
  UTCTime
now <- IO UTCTime
Data.Time.getCurrentTime
  Bool
shouldScheduleWakeUp <- STM Bool -> IO Bool
forall a. STM a -> IO a
atomically (STM Bool -> IO Bool) -> STM Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ do
    Maybe UTCTime
nextWakeUp <- [UTCTime] -> Maybe UTCTime
forall a. [a] -> Maybe a
listToMaybe ([UTCTime] -> Maybe UTCTime)
-> ([(UTCTime, MVar ())] -> [UTCTime])
-> [(UTCTime, MVar ())]
-> Maybe UTCTime
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((UTCTime, MVar ()) -> UTCTime)
-> [(UTCTime, MVar ())] -> [UTCTime]
forall a b. (a -> b) -> [a] -> [b]
map (UTCTime, MVar ()) -> UTCTime
forall a b. (a, b) -> a
fst ([(UTCTime, MVar ())] -> Maybe UTCTime)
-> STM [(UTCTime, MVar ())] -> STM (Maybe UTCTime)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar [(UTCTime, MVar ())] -> STM [(UTCTime, MVar ())]
forall a. TVar a -> STM a
readTVar (TimeEnv -> TVar [(UTCTime, MVar ())]
events TimeEnv
t)
    let eventTime :: UTCTime
eventTime = NominalDiffTime -> UTCTime -> UTCTime
addUTCTime (Int -> NominalDiffTime
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
us NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Fractional a => a -> a -> a
/ (NominalDiffTime
1000 NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
* NominalDiffTime
1000)) UTCTime
now
    TVar [(UTCTime, MVar ())]
-> ([(UTCTime, MVar ())] -> [(UTCTime, MVar ())]) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar (TimeEnv -> TVar [(UTCTime, MVar ())]
events TimeEnv
t) (((UTCTime, MVar ()) -> (UTCTime, MVar ()) -> Ordering)
-> (UTCTime, MVar ())
-> [(UTCTime, MVar ())]
-> [(UTCTime, MVar ())]
forall a. (a -> a -> Ordering) -> a -> [a] -> [a]
insertBy (UTCTime -> UTCTime -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (UTCTime -> UTCTime -> Ordering)
-> ((UTCTime, MVar ()) -> UTCTime)
-> (UTCTime, MVar ())
-> (UTCTime, MVar ())
-> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (UTCTime, MVar ()) -> UTCTime
forall a b. (a, b) -> a
fst) (UTCTime
eventTime, MVar ()
mvar))
    Bool -> STM Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> (UTCTime -> Bool) -> Maybe UTCTime -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True (UTCTime
eventTime UTCTime -> UTCTime -> Bool
forall a. Ord a => a -> a -> Bool
<) Maybe UTCTime
nextWakeUp)
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
shouldScheduleWakeUp (TimeEnv -> UTCTime -> IO ()
scheduleNextWakeUp TimeEnv
t UTCTime
now)
  MVar () -> IO ()
forall a. MVar a -> IO a
readMVar MVar ()
mvar