{-# LANGUAGE ScopedTypeVariables #-}
module Test.Util.Tracer (
    recordingTracerIORef
  , recordingTracerM
  , recordingTracerTVar
  ) where

import           Control.Tracer
import           Data.IORef
import           Ouroboros.Consensus.Util.IOLike
import           System.IO.Unsafe (unsafePerformIO)

-- | Create a 'Tracer' that stores all events in an 'IORef' that is atomically
-- updated. The second return value lets you obtain the events recorded so far
-- (from oldest to newest). Obtaining the events does not erase them.
recordingTracerIORef :: IO (Tracer IO ev, IO [ev])
recordingTracerIORef :: forall ev. IO (Tracer IO ev, IO [ev])
recordingTracerIORef = [ev] -> IO (IORef [ev])
forall a. a -> IO (IORef a)
newIORef [] IO (IORef [ev])
-> (IORef [ev] -> IO (Tracer IO ev, IO [ev]))
-> IO (Tracer IO ev, IO [ev])
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \IORef [ev]
ref -> (Tracer IO ev, IO [ev]) -> IO (Tracer IO ev, IO [ev])
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return
    ( (ev -> IO ()) -> Tracer IO ev
forall (m :: * -> *) a. (a -> m ()) -> Tracer m a
Tracer ((ev -> IO ()) -> Tracer IO ev) -> (ev -> IO ()) -> Tracer IO ev
forall a b. (a -> b) -> a -> b
$ \ev
ev -> IORef [ev] -> ([ev] -> ([ev], ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef [ev]
ref (([ev] -> ([ev], ())) -> IO ()) -> ([ev] -> ([ev], ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \[ev]
evs -> (ev
evev -> [ev] -> [ev]
forall a. a -> [a] -> [a]
:[ev]
evs, ())
    , [ev] -> [ev]
forall a. [a] -> [a]
reverse ([ev] -> [ev]) -> IO [ev] -> IO [ev]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef [ev] -> IO [ev]
forall a. IORef a -> IO a
readIORef IORef [ev]
ref
    )

-- | Create a 'Tracer' that stores all events in a 'TVar' that is atomically
-- updated. The second return value lets you obtain the events recorded so far
-- (from oldest to newest). Obtaining the events does not erase them.
recordingTracerTVar :: MonadSTM m => m (Tracer m ev, m [ev])
recordingTracerTVar :: forall (m :: * -> *) ev. MonadSTM m => m (Tracer m ev, m [ev])
recordingTracerTVar = [ev] -> m (StrictTVar m [ev])
forall (m :: * -> *) a. MonadSTM m => a -> m (StrictTVar m a)
uncheckedNewTVarM [] m (StrictTVar m [ev])
-> (StrictTVar m [ev] -> m (Tracer m ev, m [ev]))
-> m (Tracer m ev, m [ev])
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \StrictTVar m [ev]
ref -> (Tracer m ev, m [ev]) -> m (Tracer m ev, m [ev])
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return
    ( (ev -> m ()) -> Tracer m ev
forall (m :: * -> *) a. (a -> m ()) -> Tracer m a
Tracer ((ev -> m ()) -> Tracer m ev) -> (ev -> m ()) -> Tracer m ev
forall a b. (a -> b) -> a -> b
$ \ev
ev -> STM m () -> m ()
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m () -> m ()) -> STM m () -> m ()
forall a b. (a -> b) -> a -> b
$ StrictTVar m [ev] -> ([ev] -> [ev]) -> STM m ()
forall (m :: * -> *) a.
MonadSTM m =>
StrictTVar m a -> (a -> a) -> STM m ()
modifyTVar StrictTVar m [ev]
ref (ev
evev -> [ev] -> [ev]
forall a. a -> [a] -> [a]
:)
    , STM m [ev] -> m [ev]
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m [ev] -> m [ev]) -> STM m [ev] -> m [ev]
forall a b. (a -> b) -> a -> b
$ [ev] -> [ev]
forall a. [a] -> [a]
reverse ([ev] -> [ev]) -> STM m [ev] -> STM m [ev]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StrictTVar m [ev] -> STM m [ev]
forall (m :: * -> *) a. MonadSTM m => StrictTVar m a -> STM m a
readTVar StrictTVar m [ev]
ref
    )

-- | Like 'recordingTracerIORef', but lifts IO to an arbitrary applicative.
-- This is useful to record events without changing the scheduling during a
-- test.
recordingTracerM :: forall m ev. Monad m => m (Tracer m ev, m [ev])
recordingTracerM :: forall (m :: * -> *) ev. Monad m => m (Tracer m ev, m [ev])
recordingTracerM = do
  (Tracer IO ev
tr, IO [ev]
get) <- IO (Tracer IO ev, IO [ev]) -> m (Tracer IO ev, IO [ev])
forall a. IO a -> m a
liftIOtoM IO (Tracer IO ev, IO [ev])
forall ev. IO (Tracer IO ev, IO [ev])
recordingTracerIORef
  (Tracer m ev, m [ev]) -> m (Tracer m ev, m [ev])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((forall a. IO a -> m a) -> Tracer IO ev -> Tracer m ev
forall (m :: * -> *) (n :: * -> *) s.
(forall x. m x -> n x) -> Tracer m s -> Tracer n s
natTracer IO x -> m x
forall a. IO a -> m a
liftIOtoM Tracer IO ev
tr, IO [ev] -> m [ev]
forall a. IO a -> m a
liftIOtoM IO [ev]
get)
  where
    liftIOtoM :: IO a -> m a
    liftIOtoM :: forall a. IO a -> m a
liftIOtoM IO a
m = do
      -- The ficticious state is only used to force unsafePerformIO to run @m@
      -- every time @liftIOtoM m@ is evaluated.
      Int
s <- m Int
getStateM
      a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> m a) -> a -> m a
forall a b. (a -> b) -> a -> b
$! (Int, a) -> a
forall a b. (a, b) -> b
snd ((Int, a) -> a) -> (Int, a) -> a
forall a b. (a -> b) -> a -> b
$ IO (Int, a) -> (Int, a)
forall a. IO a -> a
unsafePerformIO (IO (Int, a) -> (Int, a)) -> IO (Int, a) -> (Int, a)
forall a b. (a -> b) -> a -> b
$ do
        a
r <- IO a
m
        (Int, a) -> IO (Int, a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
s, a
r)

    -- We mark this function as NOINLINE to ensure the compiler cannot reason
    -- that two calls of @getStateM@ might yield the same value.
    {-# NOINLINE getStateM #-}
    getStateM :: m Int
    getStateM :: m Int
getStateM = Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
0