{-# LANGUAGE LambdaCase #-}

module Skeletest.Internal.TestInfo (
  TestInfo (..),
  withTestInfo,
  getTestInfo,
  lookupTestInfo,
) where

import Control.Monad.IO.Class (MonadIO)
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Text (Text)
import System.IO.Unsafe (unsafePerformIO)
import UnliftIO (MonadUnliftIO)
import UnliftIO.Concurrent (ThreadId, myThreadId)
import UnliftIO.Exception (bracket_)
import UnliftIO.IORef (IORef, modifyIORef, newIORef, readIORef)

import Skeletest.Internal.Error (invariantViolation)
import Skeletest.Internal.Markers (SomeMarker)

data TestInfo = TestInfo
  { TestInfo -> Text
testModule :: Text
  , TestInfo -> [Text]
testContexts :: [Text]
  , TestInfo -> Text
testName :: Text
  , TestInfo -> [SomeMarker]
testMarkers :: [SomeMarker]
  , TestInfo -> FilePath
testFile :: FilePath
  -- ^ Relative to CWD
  }
  deriving (Int -> TestInfo -> ShowS
[TestInfo] -> ShowS
TestInfo -> FilePath
(Int -> TestInfo -> ShowS)
-> (TestInfo -> FilePath) -> ([TestInfo] -> ShowS) -> Show TestInfo
forall a.
(Int -> a -> ShowS) -> (a -> FilePath) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> TestInfo -> ShowS
showsPrec :: Int -> TestInfo -> ShowS
$cshow :: TestInfo -> FilePath
show :: TestInfo -> FilePath
$cshowList :: [TestInfo] -> ShowS
showList :: [TestInfo] -> ShowS
Show)

type TestInfoMap = Map ThreadId TestInfo

testInfoMapRef :: IORef TestInfoMap
testInfoMapRef :: IORef TestInfoMap
testInfoMapRef = IO (IORef TestInfoMap) -> IORef TestInfoMap
forall a. IO a -> a
unsafePerformIO (IO (IORef TestInfoMap) -> IORef TestInfoMap)
-> IO (IORef TestInfoMap) -> IORef TestInfoMap
forall a b. (a -> b) -> a -> b
$ TestInfoMap -> IO (IORef TestInfoMap)
forall (m :: * -> *) a. MonadIO m => a -> m (IORef a)
newIORef TestInfoMap
forall k a. Map k a
Map.empty
{-# NOINLINE testInfoMapRef #-}

withTestInfo :: (MonadUnliftIO m) => TestInfo -> m a -> m a
withTestInfo :: forall (m :: * -> *) a. MonadUnliftIO m => TestInfo -> m a -> m a
withTestInfo TestInfo
info m a
m = do
  tid <- m ThreadId
forall (m :: * -> *). MonadIO m => m ThreadId
myThreadId
  bracket_ (set tid) (unset tid) m
  where
    set :: ThreadId -> m ()
set ThreadId
tid = IORef TestInfoMap -> (TestInfoMap -> TestInfoMap) -> m ()
forall (m :: * -> *) a. MonadIO m => IORef a -> (a -> a) -> m ()
modifyIORef IORef TestInfoMap
testInfoMapRef ((TestInfoMap -> TestInfoMap) -> m ())
-> (TestInfoMap -> TestInfoMap) -> m ()
forall a b. (a -> b) -> a -> b
$ ThreadId -> TestInfo -> TestInfoMap -> TestInfoMap
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert ThreadId
tid TestInfo
info
    unset :: ThreadId -> m ()
unset ThreadId
tid = IORef TestInfoMap -> (TestInfoMap -> TestInfoMap) -> m ()
forall (m :: * -> *) a. MonadIO m => IORef a -> (a -> a) -> m ()
modifyIORef IORef TestInfoMap
testInfoMapRef ((TestInfoMap -> TestInfoMap) -> m ())
-> (TestInfoMap -> TestInfoMap) -> m ()
forall a b. (a -> b) -> a -> b
$ ThreadId -> TestInfoMap -> TestInfoMap
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete ThreadId
tid

lookupTestInfo :: (MonadIO m) => m (Maybe TestInfo)
lookupTestInfo :: forall (m :: * -> *). MonadIO m => m (Maybe TestInfo)
lookupTestInfo = do
  tid <- m ThreadId
forall (m :: * -> *). MonadIO m => m ThreadId
myThreadId
  Map.lookup tid <$> readIORef testInfoMapRef

getTestInfo :: (MonadIO m) => m TestInfo
getTestInfo :: forall (m :: * -> *). MonadIO m => m TestInfo
getTestInfo =
  m (Maybe TestInfo)
forall (m :: * -> *). MonadIO m => m (Maybe TestInfo)
lookupTestInfo m (Maybe TestInfo) -> (Maybe TestInfo -> m TestInfo) -> m TestInfo
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just TestInfo
info -> TestInfo -> m TestInfo
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TestInfo
info
    -- it's not possible for a user to write code that's executed within a test,
    -- because we define the entire main function.
    Maybe TestInfo
Nothing -> FilePath -> m TestInfo
forall a. FilePath -> a
invariantViolation FilePath
"test info not initialized"