module Convex.Tasty.Streaming.TreeMap (
  buildTestMap,
) where

import Convex.Tasty.Streaming.Types (TestInfo (..))
import Data.IORef
import Data.IntMap.Strict (IntMap)
import Data.IntMap.Strict qualified as IntMap
import Data.Text qualified as Text
import Test.Tasty (TestTree)
import Test.Tasty.Options (OptionSet)
import Test.Tasty.Runners (Ap (..), TreeFold (..), foldTestTree, trivialFold)

{- | Build a mapping from StatusMap integer indices to test metadata.

The indices correspond to the order tests appear during a fold of the TestTree,
which is the same order Tasty uses when building the StatusMap.
-}
buildTestMap :: OptionSet -> TestTree -> IO (IntMap TestInfo)
buildTestMap :: OptionSet -> TestTree -> IO (IntMap TestInfo)
buildTestMap OptionSet
opts TestTree
tree = do
  IORef Int
counterRef <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef (Int
0 :: Int)
  let Ap IO (IntMap TestInfo)
action = TreeFold (Ap IO (IntMap TestInfo))
-> OptionSet -> TestTree -> Ap IO (IntMap TestInfo)
forall b. Monoid b => TreeFold b -> OptionSet -> TestTree -> b
foldTestTree (IORef Int -> TreeFold (Ap IO (IntMap TestInfo))
mkFold IORef Int
counterRef) OptionSet
opts TestTree
tree
  IO (IntMap TestInfo)
action

mkFold :: IORef Int -> TreeFold (Ap IO (IntMap TestInfo))
mkFold :: IORef Int -> TreeFold (Ap IO (IntMap TestInfo))
mkFold IORef Int
counterRef =
  (TreeFold (Ap IO (IntMap TestInfo))
forall b. Monoid b => TreeFold b
trivialFold :: TreeFold (Ap IO (IntMap TestInfo)))
    { foldSingle = \OptionSet
_ TestName
name t
_ -> IO (IntMap TestInfo) -> Ap IO (IntMap TestInfo)
forall (f :: * -> *) a. f a -> Ap f a
Ap (IO (IntMap TestInfo) -> Ap IO (IntMap TestInfo))
-> IO (IntMap TestInfo) -> Ap IO (IntMap TestInfo)
forall a b. (a -> b) -> a -> b
$ do
        Int
idx <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
counterRef
        IORef Int -> (Int -> Int) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int
counterRef (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        let info :: TestInfo
info =
              TestInfo
                { tiId :: Int
tiId = Int
idx
                , tiName :: Text
tiName = TestName -> Text
Text.pack TestName
name
                , tiPath :: [Text]
tiPath = []
                }
        IntMap TestInfo -> IO (IntMap TestInfo)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IntMap TestInfo -> IO (IntMap TestInfo))
-> IntMap TestInfo -> IO (IntMap TestInfo)
forall a b. (a -> b) -> a -> b
$ Int -> TestInfo -> IntMap TestInfo
forall a. Int -> a -> IntMap a
IntMap.singleton Int
idx TestInfo
info
    , foldGroup = \OptionSet
_opts TestName
groupName [Ap IO (IntMap TestInfo)]
children -> IO (IntMap TestInfo) -> Ap IO (IntMap TestInfo)
forall (f :: * -> *) a. f a -> Ap f a
Ap (IO (IntMap TestInfo) -> Ap IO (IntMap TestInfo))
-> IO (IntMap TestInfo) -> Ap IO (IntMap TestInfo)
forall a b. (a -> b) -> a -> b
$ do
        let Ap IO (IntMap TestInfo)
childAction = [Ap IO (IntMap TestInfo)] -> Ap IO (IntMap TestInfo)
forall a. Monoid a => [a] -> a
mconcat [Ap IO (IntMap TestInfo)]
children
        IntMap TestInfo
childMap <- IO (IntMap TestInfo)
childAction
        let prependGroup :: TestInfo -> TestInfo
prependGroup TestInfo
ti = TestInfo
ti{tiPath = Text.pack groupName : tiPath ti}
        IntMap TestInfo -> IO (IntMap TestInfo)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IntMap TestInfo -> IO (IntMap TestInfo))
-> IntMap TestInfo -> IO (IntMap TestInfo)
forall a b. (a -> b) -> a -> b
$ (TestInfo -> TestInfo) -> IntMap TestInfo -> IntMap TestInfo
forall a b. (a -> b) -> IntMap a -> IntMap b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TestInfo -> TestInfo
prependGroup IntMap TestInfo
childMap
    , foldResource = \OptionSet
_ ResourceSpec a
_ IO a -> Ap IO (IntMap TestInfo)
k ->
        IO a -> Ap IO (IntMap TestInfo)
k (TestName -> IO a
forall a. HasCallStack => TestName -> a
error TestName
"Convex.Tasty.Streaming.TreeMap: resource not available during fold")
    }