More compact test output

3 files changed, 163 insertions, 16 deletions

diff --git a/chad-fast.cabal b/chad-fast.cabal
index 7a1c641..aa4dfcc 100644
--- a/chad-fast.cabal
+++ b/chad-fast.cabal
@@ -73,17 +73,27 @@ test-suite example
   default-language: Haskell2010
   ghc-options: -Wall -threaded
 
+library test-framework
+  exposed-modules: Test.Framework
+  build-depends:
+    base,
+    hedgehog,
+    time,
+    transformers
+  hs-source-dirs: test-framework
+  default-language: Haskell2010
+  ghc-options: -Wall -threaded
+
 test-suite test
   type: exitcode-stdio-1.0
   main-is: Main.hs
   build-depends:
     chad-fast,
+    test-framework,
     base,
     containers,
     dependent-map,
     hedgehog,
-    tasty,
-    tasty-hedgehog,
     text,
     transformers,
   hs-source-dirs: test
diff --git a/test-framework/Test/Framework.hs b/test-framework/Test/Framework.hs
new file mode 100644
index 0000000..e8c1295
--- /dev/null
+++ b/test-framework/Test/Framework.hs
@@ -0,0 +1,145 @@
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE ExistentialQuantification #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE ImplicitParams #-}
+module Test.Framework (
+  TestTree,
+  testGroup,
+  testProperty,
+  withResource,
+  withResource',
+  runTests,
+  defaultMain,
+
+  -- * Compatibility
+  TestName,
+) where
+
+import Control.Monad (forM_)
+import Control.Monad.Trans.State.Strict
+import Control.Monad.IO.Class
+import Data.String (fromString)
+import Data.Time.Clock
+import System.Exit
+import System.IO (hFlush, stdout)
+
+import qualified Hedgehog as H
+import qualified Hedgehog.Internal.Config as H
+import qualified Hedgehog.Internal.Property as H
+import qualified Hedgehog.Internal.Report as H
+import qualified Hedgehog.Internal.Runner as H
+import qualified Hedgehog.Internal.Seed as H.Seed
+
+
+data TestTree
+  = Group String [TestTree]
+  | forall a. Resource (IO a) (a -> IO ()) (a -> TestTree)
+  | HP String H.Property
+
+type TestName = String
+
+testGroup :: String -> [TestTree] -> TestTree
+testGroup = Group
+
+-- | The @a -> TestTree@ function must use the @a@ only inside properties: when
+-- not actually running properties, it will be passed 'undefined'.
+withResource :: IO a -> (a -> IO ()) -> (a -> TestTree) -> TestTree
+withResource = Resource
+
+-- | Same caveats as 'withResource'.
+withResource' :: IO a -> (a -> TestTree) -> TestTree
+withResource' make fun = withResource make (\_ -> return ()) fun
+
+testProperty :: String -> H.Property -> TestTree
+testProperty = HP
+
+computeMaxLen :: TestTree -> Int
+computeMaxLen = go 0
+  where
+    go :: Int -> TestTree -> Int
+    -- go indent (Group name trees) = maximum (2*indent + length name : map (go (indent+1)) trees)
+    go indent (Group _ trees) = maximum (0 : map (go (indent+1)) trees)
+    go indent (Resource _ _ fun) = go indent (fun undefined)
+    go indent (HP name _) = 2 * indent + length name
+
+data Stats = Stats
+  { statsOK :: Int
+  , statsTotal :: Int }
+  deriving (Show)
+
+initStats :: Stats
+initStats = Stats 0 0
+
+newtype M a = M (StateT Stats IO a)
+  deriving newtype (Functor, Applicative, Monad, MonadIO)
+
+modifyStats :: (Stats -> Stats) -> M ()
+modifyStats f = M (modify f)
+
+runTests :: TestTree -> IO ExitCode
+runTests = \tree -> do
+  let M m = let ?maxlen = computeMaxLen tree in go 0 tree
+  starttm <- getCurrentTime
+  stats <- execStateT m initStats
+  endtm <- getCurrentTime
+  printStats stats (diffUTCTime endtm starttm)
+  return (if statsOK stats == statsTotal stats
+            then ExitSuccess else ExitFailure 1)
+  where
+    go :: (?maxlen :: Int) => Int -> TestTree -> M ()
+    go indent (Group name trees) = do
+      liftIO $ putStrLn (replicate (2 * indent) ' ' ++ name)
+      forM_ trees $ go (indent + 1)
+    go indent (Resource make cleanup fun) = do
+      value <- liftIO make
+      go indent (fun value)
+      liftIO $ cleanup value
+      return ()
+    go indent (HP name (H.Property config test)) = do
+      let thislen = 2*indent + length name
+      liftIO $ putStr (replicate (2*indent) ' ' ++ name ++ ": " ++ replicate (?maxlen - thislen) ' ')
+      liftIO $ hFlush stdout
+
+      seed <- H.Seed.random
+
+      starttm <- liftIO getCurrentTime
+      report <- liftIO $ H.checkReport config 0 seed test (outputProgress (?maxlen + 2) name)
+      endtm <- liftIO getCurrentTime
+
+      liftIO $ printResult name report (diffUTCTime endtm starttm)
+
+      modifyStats $ \stats -> stats { statsOK = fromEnum (H.reportStatus report == H.OK) + statsOK stats
+                                    , statsTotal = 1 + statsTotal stats }
+
+outputProgress :: Int -> String -> H.Report H.Progress -> IO ()
+outputProgress indent name report = do
+  str <- H.renderProgress H.EnableColor (Just (fromString name)) report
+  putStr (str ++ "\x1B[" ++ show (indent+1) ++ "G")
+  hFlush stdout
+
+printResult :: String -> H.Report H.Result -> NominalDiffTime -> IO ()
+printResult name report timeTaken = do
+  str <- H.renderResult H.EnableColor (Just (fromString name)) report
+  if timeTaken >= 0.01
+    then putStrLn ("\x1B[K" ++ str ++ " (" ++ prettySeconds (realToFrac timeTaken) ++ ")")
+    else putStrLn ("\x1B[K" ++ str)
+
+printStats :: Stats -> NominalDiffTime -> IO ()
+printStats stats timeTaken
+  | statsOK stats == statsTotal stats = do
+      putStrLn $ "\x1B[32mAll " ++ show (statsTotal stats) ++
+                 " tests passed (" ++ prettySeconds (realToFrac timeTaken) ++ ")\x1B[0m"
+  | otherwise =
+      let nfailed = statsTotal stats - statsOK stats
+      in putStrLn $ "\x1B[31mFailed " ++ show nfailed ++ " out of " ++ show (statsTotal stats) ++
+                    " tests (" ++ prettySeconds (realToFrac timeTaken) ++ ")\x1B[0m"
+
+prettySeconds :: Double -> String
+prettySeconds x =
+  let str = show (round (x * 100) :: Int)
+      str' = replicate (3 - length str) '0' ++ str
+      (pre, post) = splitAt (length str' - 2) str'
+  in pre ++ "." ++ post ++ "s"
+
+defaultMain :: TestTree -> IO ()
+defaultMain tree = runTests tree >>= exitWith
diff --git a/test/Main.hs b/test/Main.hs
index 83271fc..1ad7f75 100644
--- a/test/Main.hs
+++ b/test/Main.hs
@@ -14,7 +14,6 @@ import Control.Monad.Trans.Class (lift)
 import Control.Monad.IO.Class (liftIO)
 import Control.Monad.Trans.State
 import Data.Bifunctor
--- import qualified Data.Functor.Product as Product
 import Data.Int (Int64)
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
@@ -22,8 +21,7 @@ import qualified Data.Text as T
 import Hedgehog
 import qualified Hedgehog.Gen as Gen
 import qualified Hedgehog.Range as Range
-import Test.Tasty
-import Test.Tasty.Hedgehog
+import Test.Framework
 
 import Array
 import AST
@@ -209,12 +207,10 @@ adTestGen name expr envGenerator =
   let env = knownEnv @env
       exprS = simplifyFix expr
   in
-  withCompiled env expr $ \getprimalfun ->
-  withCompiled env (simplifyFix expr) $ \getprimalSfun ->
+  withCompiled env expr $ \primalfun ->
+  withCompiled env (simplifyFix expr) $ \primalSfun ->
   testGroup name
     [testProperty "compile primal" $ property $ do
-       primalfun <- liftIO getprimalfun
-       primalSfun <- liftIO getprimalSfun
        input <- forAllWith (showEnv env) envGenerator
 
        let outPrimalI = interpretOpen False input expr
@@ -225,9 +221,8 @@ adTestGen name expr envGenerator =
        outPrimalSC <- liftIO $ primalSfun input
        diff outPrimalSI (closeIsh' 1e-8) outPrimalSC
 
-    ,withCompiled (dne env) (dfwdDN exprS) $ \getdnfun ->
+    ,withCompiled (dne env) (dfwdDN exprS) $ \dnfun ->
      testProperty "compile fwdAD" $ property $ do
-       dnfun <- liftIO getdnfun
        input <- forAllWith (showEnv env) envGenerator
        dinput <- forAllWith (showEnv (dne env)) $ extendDNE env input
        let (outDNI1, outDNI2) = interpretOpen False dinput (dfwdDN expr)
@@ -235,11 +230,8 @@ adTestGen name expr envGenerator =
        diff outDNI1 (closeIsh' 1e-8) outDNC1
        diff outDNI2 (closeIsh' 1e-8) outDNC2
 
-    ,withResource (makeFwdADArtifactCompile env exprS) (\_ -> pure ()) $ \getfwdartifactC ->
+    ,withResource (makeFwdADArtifactCompile env exprS) (\_ -> pure ()) $ \fwdartifactC ->
      testProperty "chad" $ property $ do
-       primalSfun <- liftIO getprimalSfun
-       fwdartifactC <- liftIO getfwdartifactC
-
        annotate (concat (unSList (\t -> ppSTy 0 t ++ " -> ") env) ++ ppSTy 0 (typeOf expr))
 
        let dtermChad0 = ELet ext (EConst ext STF64 1.0) $ chad' defaultConfig env expr
@@ -288,7 +280,7 @@ adTestGen name expr envGenerator =
     envScalars SNil SNil = []
     envScalars (t `SCons` ts) (Value x `SCons` xs) = tanScalars t x ++ envScalars ts xs
 
-withCompiled :: SList STy env -> Ex env t -> (IO (SList Value env -> IO (Rep t)) -> TestTree) -> TestTree
+withCompiled :: SList STy env -> Ex env t -> ((SList Value env -> IO (Rep t)) -> TestTree) -> TestTree
 withCompiled env expr = withResource (compile env expr) (\_ -> pure ())
 
 term_build1_sum :: Ex '[TArr N1 (TScal TF64)] (TScal TF64)