{-# 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
ev ev -> [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
ev ev -> [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
  (tr, 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
  pure (natTracer liftIOtoM tr, liftIOtoM 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.
    s <- m Int
getStateM
    pure $! snd $ unsafePerformIO $ do
      r <- m
      pure (s, 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