content: module Test.Driver.Dialogue.TestClock (
      TestClock -- opaque
    , Tick      -- opaque
    , new
      -- * Current time
    , waitForTick
    , advance
      -- * Green light
    , waitForGreenLight
    , giveGreenLight
      -- * Interleavings
    , interleave
    , assignTimings
    ) where
  
  import Prelude hiding (id)
  
  import Control.Concurrent.STM
  import Control.Exception
  import Control.Monad
  import Control.Monad.IO.Class
  import Data.List.NonEmpty (NonEmpty(..))
  import Data.List.NonEmpty qualified as NE
  import Data.Map (Map)
  import Data.Map qualified as Map
  import Data.Maybe (mapMaybe)
  import Data.Set (Set)
  import Data.Set qualified as Set
  import GHC.Stack
  import Test.QuickCheck (Gen, choose)
  
  {-------------------------------------------------------------------------------
    Definition
  -------------------------------------------------------------------------------}
  
  -- | Test clock
  --
  -- When we are testing concurrent code, we may get different results depending
  -- on the exact way that the various threads are scheduled, potentially making
  -- test outcomes difficult to specify. Moreover, even if the final outcome does
  -- not depend on the exact scheduling, debugging tests is much simpler if they
  -- are deterministic and do the same thing each time. Consider tracing a
  -- particular test execution, or inspecting it through Wireshark: if the test
  -- execution is different each time, debugging becomes much more difficult.
  --
  -- We therefore introduce a global \"test clock\", and randomly assign each
  -- action to be executed a \"clock tick\" (formally, we are exploring a /random
  -- interleaving/). This then begs the question of when the clock gets advanced.
  -- We could just have a thread tick the test clock every @n@ milliseconds on the
  -- wallclock, but choosing @n@ is hard: too low, and we /still/ execute actions
  -- concurrently; too high, and the tests run much slower than they need to.
  -- Instead, we say that the actions /themselves/ are responsible for advancing
  -- the test clock when appropriate.
  --
  -- In a concurrent setting, the \"when appropriate\" part is however not always
  -- easy to determine. Suppose we have thread A sending messages to thread B,
  -- then there is no need to wait for each message to have been /received/ by B
  -- (this would reduce the test to effectively using the network layer in a
  -- synchronous manner); this would suggest we can tick the clock immediately
  -- after sending (rather than waiting for the message to be received first).
  -- However, if at some point later A throws a (deliberate) exception, then A
  -- /should/ wait for B to have received all messages first, or risk that the
  -- exception could \"overtake\" some message, leading to flaky tests.
  --
  -- Getting this right is surprisingly subtle. A knows to wait for B, but only
  -- once it gets to the action that throws the exception; at that point it wants
  -- to know that B received the messages it sent /in some previous action/. The
  -- approach we take attempts to avoid this \"spooky action at a distance\" as
  -- follows: when A is about to throw a deliberate exception at tick T, it waits
  -- to get the green light for T; on B's side, when B is ready to /react/ to this
  -- exception (implying any previous actions have been verified already), it
  -- /gives/ the green light for T. Since A knows when it needs to wait for the
  -- green light (e.g., throwing a deliberate exception) and when it doesn't
  -- (e.g., sending a message), we avoid over-sychronization. Meanwhile, B giving
  -- the green light for an action rather than acknowledging some previous action
  -- keeps the tests easier to understand: A and B are both talking about the same
  -- tick T.
  newtype TestClock = TestClock (TVar State)
  
  -- | State of the clock
  data State = State {
        stateNow   :: Tick
      , stateGreen :: Set Tick
      }
  
  -- | Clock tick
  newtype Tick = Tick Word
    deriving stock (Show)
    deriving newtype (Eq, Ord, Enum)
  
  -- | Start the clock
  new :: IO TestClock
  new = TestClock <$> newTVarIO initState
    where
      initState :: State
      initState = State {
          stateNow   = Tick 0
        , stateGreen = Set.empty
        }
  
  {-------------------------------------------------------------------------------
    Current time
  -------------------------------------------------------------------------------}
  
  -- | Thrown by 'waitForTick'
  data TimePassed = TimePassed {
        timePassedNow    :: Tick
      , timePassedWanted :: Tick
      , timePassedAt     :: CallStack
      }
    deriving stock (Show)
    deriving anyclass (Exception)
  
  -- | Wait for specified clock tick
  --
  -- If the clock has already gone past the specified time, throws 'TimePassed'.
  waitForTick :: MonadIO m => TestClock -> Tick -> m ()
  waitForTick (TestClock clock) t = liftIO . atomically $ do
      State{stateNow} <- readTVar clock
      if | stateNow < t  -> retry
         | stateNow == t -> return ()
         | otherwise     -> throwSTM $ TimePassed {
                                timePassedNow    = stateNow
                              , timePassedWanted = t
                              , timePassedAt     = callStack
                              }
  
  -- | Advance the clock by one tick
  --
  -- There should only ever be /one/ thread responsible for advancing the clock
  -- at any given time.
  advance :: MonadIO m => TestClock -> m ()
  advance (TestClock clock) = liftIO . atomically $ do
      modifyTVar clock $ \st@State{stateNow} ->
        st{stateNow = succ stateNow}
  
  {-------------------------------------------------------------------------------
    Green light
  -------------------------------------------------------------------------------}
  
  -- | Wait for green light
  --
  -- See 'TestClock' for discussion.
  waitForGreenLight :: MonadIO m => TestClock -> Tick -> m ()
  waitForGreenLight (TestClock clock) t = liftIO . atomically $ do
      State{stateGreen} <- readTVar clock
      unless (t `Set.member` stateGreen) retry
  
  -- | Give green light
  --
  -- See 'TestClock' for discussion.
  giveGreenLight :: MonadIO m => TestClock -> Tick -> m ()
  giveGreenLight (TestClock clock) t = liftIO . atomically $ do
      modifyTVar clock $ \st@State{stateGreen} ->
        st{stateGreen = Set.insert t stateGreen}
  
  {-------------------------------------------------------------------------------
    Computing interleavings
  
    We assign timings /from/ interleavings. This is more suitable to shrinking;
    during generation we pick an arbitrary interleaving, then as we shrink, we
    leave the interleaving (mostly) alone, but re-assign timings.
  -------------------------------------------------------------------------------}
  
  -- | Pick an arbitrary interleaving
  --
  -- This flattens the list; each item in the result is an element from /one/
  -- of the input lists, annotated with the index of that particular list.
  --
  -- > ghci> sample (interleave ["abc", "de"])
  -- > [(0,'a'),(1,'d'),(0,'b'),(0,'c'),(1,'e')]
  -- > [(1,'d'),(0,'a'),(0,'b'),(1,'e'),(0,'c')]
  -- > [(1,'d'),(1,'e'),(0,'a'),(0,'b'),(0,'c')]
  interleave :: forall a. [[a]] -> Gen [(Int, a)]
  interleave =
      go . mapMaybe (\(i, xs) -> (i,) <$> NE.nonEmpty xs) . zip [0..]
    where
      go :: [(Int, NonEmpty a)] -> Gen [(Int, a)]
      go [] = return []
      go xs = do
          (as, (i, b :| bs), cs) <- isolate xs
          fmap ((i, b) :) $
            case bs of
              []       -> go (as ++ cs)
              b' : bs' -> go (as ++ [(i, b' :| bs')] ++ cs)
  
  -- | Assign timings, given an interleaving
  --
  -- In some sense this is an inverse to 'interleave':
  --
  -- >    assignTimings [(0,'a'),(0,'b'),(1,'d'),(1,'e'),(0,'c')]
  -- > == [ [ (Tick 0,'a')
  -- >      , (Tick 2,'b')
  -- >      , (Tick 8,'c')
  -- >      ]
  -- >    , [ (Tick 4,'d')
  -- >      ,( Tick 6,'e')
  -- >      ]
  -- >    ]
  --
  -- Put another way
  --
  -- > map (map snd) . assignTimings <$> interleave xs
  --
  -- will just generate @xs@.
  assignTimings :: forall id a. Ord id => [(id, a)] -> [[(Tick, a)]]
  assignTimings = separate . assignClockTicks
    where
      separate :: [(id, x)] -> [[x]]
      separate = go Map.empty
        where
          go :: Map id [x] -> [(id, x)] -> [[x]]
          go acc []             = map reverse $ Map.elems acc
          go acc ((id, x) : xs) = go (Map.alter (Just . maybe [x] (x:)) id acc) xs
  
      assignClockTicks :: [(id, a)] -> [(id, (Tick, a))]
      assignClockTicks = go (Tick 0)
        where
          go :: Tick -> [(id, a)] -> [(id, (Tick, a))]
          go _ []           = []
          go t ((id, a):as) = (id, (t, a)) : go (succ t) as
  
  {-------------------------------------------------------------------------------
    Internal auxiliary
  -------------------------------------------------------------------------------}
  
  isolate :: [a] -> Gen ([a], a, [a])
  isolate = \case
      [] -> error "isolate: empty list"
      xs -> do n <- choose (0, length xs - 1)
               return $ go [] xs n
    where
      go :: [a] -> [a] -> Int -> ([a], a, [a])
      go _    []     _ = error "isolate: impossible"
      go prev (x:xs) 0 = (reverse prev, x, xs)
      go prev (x:xs) n = go (x:prev) xs (pred n)
  