test: Proper intermixing of GCC warnings with test output

5 files changed, 91 insertions, 44 deletions

diff --git a/bench/Main.hs b/bench/Main.hs
index ec9264b..6db77b5 100644
--- a/bench/Main.hs
+++ b/bench/Main.hs
@@ -36,7 +36,7 @@ import Simplify
 
 gradCHAD :: KnownEnv env => CHADConfig -> Ex env (TScal TF64) -> IO (SList Value env -> IO (Double, Rep (Tup (D2E env))))
 gradCHAD config term =
-  compile knownEnv $
+  compileStderr knownEnv $
     simplifyFix $ pruneExpr knownEnv $
     simplifyFix $ unMonoid $
     simplifyFix $
diff --git a/src/Compile.hs b/src/Compile.hs
index f2063ee..bf7817a 100644
--- a/src/Compile.hs
+++ b/src/Compile.hs
@@ -8,7 +8,7 @@
 {-# LANGUAGE PolyKinds #-}
 {-# LANGUAGE TupleSections #-}
 {-# LANGUAGE TypeApplications #-}
-module Compile (compile) where
+module Compile (compile, compileStderr) where
 
 import Control.Applicative (empty)
 import Control.Monad (forM_, when, replicateM)
@@ -71,28 +71,30 @@ debugAllocs :: Bool; debugAllocs = toEnum 0
 -- | Emit extra C code that checks stuff
 emitChecks :: Bool; emitChecks = toEnum 0
 
+-- | Returns compiled function plus compilation output (warnings)
 compile :: SList STy env -> Ex env t
-        -> IO (SList Value env -> IO (Rep t))
+        -> IO (SList Value env -> IO (Rep t), String)
 compile = \env expr -> do
   codeID <- atomicModifyIORef' uniqueIdGenRef (\i -> (i + 1, i))
 
   let (source, offsets) = compileToString codeID env expr
   when debugPrintAST $ hPutStrLn stderr $ "Compiled AST: <<<\n" ++ ppExpr env expr ++ "\n>>>"
   when debugCSource $ hPutStrLn stderr $ "Generated C source: <<<\n\x1B[2m" ++ lineNumbers source ++ "\x1B[0m>>>"
-  lib <- buildKernel source "kernel"
+  (lib, compileOutput) <- buildKernel source "kernel"
 
   let result_type = typeOf expr
       result_size = sizeofSTy result_type
 
-  return $ \val -> do
-    allocaBytes (koResultOffset offsets + result_size) $ \ptr -> do
-      let args = zip (reverse (unSList Some (slistZip env val))) (koArgOffsets offsets)
-      serialiseArguments args ptr $ do
-        callKernelFun lib ptr
-        ok <- peekByteOff @Word8 ptr (koOkResOffset offsets)
-        when (ok /= 1) $
-          ioError (mkIOError userErrorType "fatal error detected during chad kernel execution (memory has been leaked)" Nothing Nothing)
-        deserialise result_type ptr (koResultOffset offsets)
+  let function val = do
+        allocaBytes (koResultOffset offsets + result_size) $ \ptr -> do
+          let args = zip (reverse (unSList Some (slistZip env val))) (koArgOffsets offsets)
+          serialiseArguments args ptr $ do
+            callKernelFun lib ptr
+            ok <- peekByteOff @Word8 ptr (koOkResOffset offsets)
+            when (ok /= 1) $
+              ioError (mkIOError userErrorType "fatal error detected during chad kernel execution (memory has been leaked)" Nothing Nothing)
+            deserialise result_type ptr (koResultOffset offsets)
+  return (function, compileOutput)
   where
     serialiseArguments :: [(Some (Product STy Value), Int)] -> Ptr () -> IO r -> IO r
     serialiseArguments ((Some (Product.Pair t (Value arg)), off) : args) ptr k =
@@ -100,6 +102,15 @@ compile = \env expr -> do
         serialiseArguments args ptr k
     serialiseArguments _ _ k = k
 
+-- | 'compile', but writes any produced C compiler output to stderr.
+compileStderr :: SList STy env -> Ex env t
+              -> IO (SList Value env -> IO (Rep t))
+compileStderr env expr = do
+  (fun, output) <- compile env expr
+  when (not (null output)) $
+    hPutStrLn stderr $ "[chad] Kernel compilation GCC output: <<<\n" ++ output ++ ">>>"
+  return fun
+
 
 data StructDecl = StructDecl
     String  -- ^ name
diff --git a/src/Compile/Exec.hs b/src/Compile/Exec.hs
index cc6d5fa..bbccf1c 100644
--- a/src/Compile/Exec.hs
+++ b/src/Compile/Exec.hs
@@ -30,7 +30,7 @@ debug = False
 -- The IORef wrapper is required for the finalizer to attach properly (see the 'Weak' docs)
 data KernelLib = KernelLib !(IORef (FunPtr (Ptr () -> IO ())))
 
-buildKernel :: String -> String -> IO KernelLib
+buildKernel :: String -> String -> IO (KernelLib, String)
 buildKernel csource funname = do
   template <- (++ "/tmp.chad.") <$> getTempDir
   path <- mkdtemp template
@@ -51,11 +51,6 @@ buildKernel csource funname = do
     ExitSuccess -> return ()
     ExitFailure{} -> hPutStrLn stderr $ "[chad] Kernel compilation failed! Source: <<<\n" ++ lineNumbers csource ++ ">>>"
 
-  when (not (null gccStdout)) $
-    hPutStrLn stderr $ "[chad] Kernel compilation: GCC stdout: <<<\n" ++ gccStdout ++ ">>>"
-  when (not (null gccStderr)) $
-    hPutStrLn stderr $ "[chad] Kernel compilation: GCC stderr: <<<\n" ++ gccStderr ++ ">>>"
-
   case ec of
     ExitSuccess -> return ()
     ExitFailure{} -> do
@@ -72,7 +67,7 @@ buildKernel csource funname = do
   _ <- mkWeakIORef ref (do numLeft <- atomicModifyIORef' numLoadedCounter (\n -> (n-1, n-1))
                            when debug $ hPutStrLn stderr $ "[chad] unloading kernel " ++ path ++ " (" ++ show numLeft ++ " left)"
                            dlclose dl)
-  return (KernelLib ref)
+  return (KernelLib ref, gccStdout ++ (if null gccStdout then "" else "\n") ++ gccStderr)
 
 foreign import ccall "dynamic"
   wrapKernelFun :: FunPtr (Ptr () -> IO ()) -> Ptr () -> IO ()
diff --git a/test-framework/Test/Framework.hs b/test-framework/Test/Framework.hs
index 80711b2..b7d0dc2 100644
--- a/test-framework/Test/Framework.hs
+++ b/test-framework/Test/Framework.hs
@@ -1,9 +1,12 @@
+{-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE DerivingVia #-}
 {-# LANGUAGE ExistentialQuantification #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE ImplicitParams #-}
 {-# LANGUAGE ImportQualifiedPost #-}
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TupleSections #-}
 module Test.Framework (
@@ -11,12 +14,16 @@ module Test.Framework (
   testGroup,
   groupSetCollapse,
   testProperty,
-  withResource,
-  withResource',
   runTests,
   defaultMain,
   Options(..),
 
+  -- * Resources
+  withResource,
+  withResource',
+  TestCtx,
+  outputWarningText,
+
   -- * Compatibility
   TestName,
 ) where
@@ -29,11 +36,13 @@ import Control.Monad (forM, when, forM_)
 import Control.Monad.IO.Class
 import Data.IORef
 import Data.List (isInfixOf, intercalate)
-import Data.Maybe (isJust, mapMaybe, fromJust)
+import Data.Maybe (mapMaybe, fromJust)
+import Data.Monoid (All(..), Any(..), Sum(..))
 import Data.PQueue.Prio.Min qualified as PQ
 import Data.String (fromString)
 import Data.Time.Clock
 import GHC.Conc (getNumProcessors)
+import GHC.Generics (Generic, Generically(..))
 import System.Console.ANSI qualified as ANSI
 import System.Console.Concurrent (outputConcurrent)
 import System.Console.Regions
@@ -57,10 +66,16 @@ type TestName = String
 
 data TestTree
   = Group GroupOpts String [TestTree]
-  | forall a. Resource String (IO a) (a -> IO ()) (a -> TestTree)
+  | forall a. Resource String ((?testCtx :: TestCtx) => IO a) ((?testCtx :: TestCtx) => a -> IO ()) (a -> TestTree)
       -- ^ Name is not specified by user, but inherited from the tree below
   | HP String H.Property
 
+data TestCtx = TestCtx
+  { tctxOutput :: String -> IO () }
+
+outputWarningText :: (?testCtx :: TestCtx) => String -> IO ()
+outputWarningText = tctxOutput ?testCtx
+
 -- Not exported because a Resource is not supposed to have a name in the first place
 treeName :: TestTree -> String
 treeName (Group _ name _) = name
@@ -82,13 +97,13 @@ groupSetCollapse (Group opts name trees) = Group opts { goCollapse = True } name
 groupSetCollapse _ = error "groupSetCollapse: not called on a Group"
 
 -- | The @a -> TestTree@ function must use the @a@ only inside properties: the
--- functoin will be passed 'undefined' when exploring the test tree (without
+-- function will be passed 'undefined' when exploring the test tree (without
 -- running properties).
-withResource :: IO a -> (a -> IO ()) -> (a -> TestTree) -> TestTree
+withResource :: ((?testCtx :: TestCtx) => IO a) -> ((?testCtx :: TestCtx) => a -> IO ()) -> (a -> TestTree) -> TestTree
 withResource make cleanup fun = Resource (treeName (fun undefined)) make cleanup fun
 
 -- | Same caveats as 'withResource'.
-withResource' :: IO a -> (a -> TestTree) -> TestTree
+withResource' :: ((?testCtx :: TestCtx) => IO a) -> (a -> TestTree) -> TestTree
 withResource' make fun = withResource make (\_ -> return ()) fun
 
 testProperty :: String -> H.Property -> TestTree
@@ -226,7 +241,7 @@ runTests options = \tree' ->
                          successVar <- newEmptyMVar
                          runTreePar Nothing [] [] tree successVar
                          readMVar successVar
-             else isJust <$> runTreeSeq 0 [] tree
+             else getAll . seqresAllSuccess <$> runTreeSeq 0 [] tree
       stats <- readIORef statsRef
       endtm <- getCurrentTime
       let ?istty = isterm in printStats (treeNumTests tree) stats (diffUTCTime endtm starttm)
@@ -284,6 +299,9 @@ runTreePar topmparregion revidxlist revpath toptree@Resource{} topoutvar = runRe
       let pathitem = '[' : show depth ++ "](" ++ inhname ++ ")"
           path = intercalate "/" (reverse (pathitem : revpath))
           idxlist = reverse revidxlist
+      let ?testCtx = TestCtx (\str ->
+                       outputConcurrent (ansiYellow ++ "## Warning for " ++ path ++ ":" ++ ansiReset ++
+                                         "\n" ++ str ++ "\n"))
       submitOrRunIn mparregion idxlist Nothing $ \makeRegion -> do
         setConsoleRegion makeRegion ('|' : path ++ " [R] making...")
 
@@ -337,37 +355,51 @@ submitOrRunIn (Just reg) _idxlist outvar fun = do
   result <- fun reg
   forM_ outvar $ \mvar -> putMVar mvar result
 
+data SeqRes = SeqRes
+  { seqresHaveWarnings :: Any
+  , seqresAllSuccess :: All
+  , seqresNumLines :: Sum Int }
+  deriving (Generic)
+  deriving (Semigroup, Monoid) via Generically SeqRes
+
 -- | If all tests are successful, returns the number of output lines produced
 runTreeSeq :: (?options :: Options, ?stats :: IORef Stats, ?maxlen :: Int, ?istty :: Bool)
-           => Int -> [String] -> TestTree -> IO (Maybe Int)
+           => Int -> [String] -> TestTree -> IO SeqRes
 runTreeSeq indent revpath (Group opts name trees) = do
   putStrLn (replicate (2 * indent) ' ' ++ name) >> hFlush stdout
   starttm <- getCurrentTime
-  mlns <- fmap (fmap sum . sequence) . forM trees $
-            runTreeSeq (indent + 1) (name : revpath)
+  res <- fmap mconcat . forM trees $
+           runTreeSeq (indent + 1) (name : revpath)
   endtm <- getCurrentTime
-  case mlns of
-    Just lns | goCollapse opts, ?istty -> do
+  if not (getAny (seqresHaveWarnings res)) && getAll (seqresAllSuccess res) && goCollapse opts && ?istty
+    then do
       let thislen = 2*indent + length name
+      let Sum lns = seqresNumLines res
       putStrLn $ concat (replicate (lns+1) (ANSI.cursorUpCode 1 ++ ANSI.clearLineCode)) ++
                  ANSI.setCursorColumnCode 0 ++
                  replicate (2 * indent) ' ' ++ name ++ ": " ++ replicate (?maxlen - thislen) ' ' ++
                  ansiGreen ++ "OK" ++ ansiReset ++
                  prettyDuration False (realToFrac (diffUTCTime endtm starttm))
-      return (Just 1)
-    _ -> return ((+1) <$> mlns)
+      return (mempty { seqresNumLines = 1 })
+    else return (res <> (mempty { seqresNumLines = 1 }))
 runTreeSeq indent path (Resource _ make cleanup fun) = do
+  outputted <- newIORef False
+  let ?testCtx = TestCtx (\str -> do
+                   atomicModifyIORef' outputted (\_ -> (True, ()))
+                   putStrLn (ansiYellow ++ "## Warning for " ++ (intercalate "/" (reverse path)) ++
+                             ":" ++ ansiReset ++ "\n" ++ str))
   value <- make
-  success <- runTreeSeq indent path (fun value)
+  res <- runTreeSeq indent path (fun value)
   cleanup value
-  return success
+  warnings <- readIORef outputted
+  return (res <> (mempty { seqresHaveWarnings = Any warnings }))
 runTreeSeq indent path (HP name prop) = do
   let thislen = 2*indent + length name
   let prefix = replicate (2*indent) ' ' ++ name ++ ": " ++ replicate (?maxlen - thislen) ' '
   when ?istty $ putStr prefix >> hFlush stdout
   (ok, rendered) <- runHP (outputProgress (?maxlen + 2)) path name prop
   putStrLn ((if ?istty then ANSI.clearFromCursorToLineEndCode else prefix) ++ rendered) >> hFlush stdout
-  return (if ok then Just 1 else Nothing)
+  return (mempty { seqresAllSuccess = All ok, seqresNumLines = 1 })
 
 runHP :: (?options :: Options, ?stats :: IORef Stats, ?maxlen :: Int)
       => (H.Report H.Progress -> IO ())
@@ -489,10 +521,11 @@ ansi :: (?istty :: Bool) => String -> String
 ansi | ?istty = id
      | otherwise = const ""
 
-ansiRed, ansiGreen, ansiReset :: (?istty :: Bool) => String
-ansiRed   = ansi (ANSI.setSGRCode [ANSI.SetColor ANSI.Foreground ANSI.Dull ANSI.Red])
-ansiGreen = ansi (ANSI.setSGRCode [ANSI.SetColor ANSI.Foreground ANSI.Dull ANSI.Green])
-ansiReset = ansi (ANSI.setSGRCode [ANSI.Reset])
+ansiRed, ansiYellow, ansiGreen, ansiReset :: (?istty :: Bool) => String
+ansiRed    = ansi (ANSI.setSGRCode [ANSI.SetColor ANSI.Foreground ANSI.Dull ANSI.Red])
+ansiYellow = ansi (ANSI.setSGRCode [ANSI.SetColor ANSI.Foreground ANSI.Vivid ANSI.Yellow])
+ansiGreen  = ansi (ANSI.setSGRCode [ANSI.SetColor ANSI.Foreground ANSI.Dull ANSI.Green])
+ansiReset  = ansi (ANSI.setSGRCode [ANSI.Reset])
 
 -- getTermIsDark :: IO Bool
 -- getTermIsDark = do
diff --git a/test/Main.hs b/test/Main.hs
index 4bc9082..d586973 100644
--- a/test/Main.hs
+++ b/test/Main.hs
@@ -11,6 +11,7 @@
 {-# LANGUAGE UndecidableInstances #-}
 module Main where
 
+import Control.Monad (when)
 import Control.Monad.Trans.Class (lift)
 import Control.Monad.Trans.State
 import Data.Bifunctor
@@ -352,7 +353,10 @@ adTestGenChad testname config env envGenerator expr exprS primalSfun | Dict <- e
       dtermSChadSUS = simplifyFix $ unMonoid dtermSChadS
       dtermSChadSUSP = pruneExpr env dtermSChadSUS
   in
-  withResource (makeFwdADArtifactCompile env exprS) (\_ -> pure ()) $ \fwdartifactC ->
+  withResource' (do (fun, output) <- makeFwdADArtifactCompile env exprS
+                    when (not (null output)) $
+                      outputWarningText $ "Forward AD compile GCC output: <<<\n" ++ output ++ ">>>"
+                    return fun) $ \fwdartifactC ->
   withCompiled env dtermSChadSUSP $ \dcompSChadSUSP ->
     testProperty testname $ property $ do
       annotate (concat (unSList (\t -> ppSTy 0 t ++ " -> ") env) ++ ppSTy 0 (typeOf expr))
@@ -416,7 +420,11 @@ adTestGenChad testname config env envGenerator expr exprS primalSfun | Dict <- e
       diff tansCompSChadSUSP closeIshE' tansFwd
 
 withCompiled :: SList STy env -> Ex env t -> ((SList Value env -> IO (Rep t)) -> TestTree) -> TestTree
-withCompiled env expr = withResource (compile env expr) (\_ -> pure ())
+withCompiled env expr = withResource' $ do
+  (fun, output) <- compile env expr
+  when (not (null output)) $
+    outputWarningText $ "Kernel compilation GCC output: <<<\n" ++ output ++ ">>>"
+  return fun
 
 gen_gmm :: Gen (SList Value [R, R, R, I64, TMat R, TMat R, TMat R, TMat R, TVec R, I64, I64, I64])
 gen_gmm = do