Parser: collect toplevel definitions using Writer

2 files changed, 21 insertions, 8 deletions

diff --git a/src/Parser.hs b/src/Parser.hs
index 405622e..5c79390 100644
--- a/src/Parser.hs
+++ b/src/Parser.hs
@@ -5,7 +5,9 @@ module Parser (
 ) where
 
 import Control.Monad.State
+import Control.Monad.Writer
 
+import Data.Maybe
 import Data.Set qualified as S
 import Data.Text qualified as T
 import Data.Text.Lazy qualified as TL
@@ -24,8 +26,8 @@ import Parser.Statement
 import Test
 import Test.Builtins
 
-parseTestDefinition :: TestParser Test
-parseTestDefinition = label "test definition" $ toplevel $ do
+parseTestDefinition :: TestParser ()
+parseTestDefinition = label "test definition" $ toplevel ToplevelTest $ do
     block (\name steps -> return $ Test name $ concat steps) header testStep
     where header = do
               wsymbol "test"
@@ -47,7 +49,10 @@ parseTestModule absPath = do
         , do
             return $ [ T.pack $ takeBaseName absPath ]
         ]
-    moduleTests <- many parseTestDefinition
+    (_, toplevels) <- listen $ many $ choice
+        [ parseTestDefinition
+        ]
+    let moduleTests = catMaybes $ map (\case ToplevelTest x -> Just x; {- _ -> Nothing -}) toplevels
     eof
     return Module { .. }
 
@@ -56,9 +61,13 @@ parseTestFile path = do
     content <- TL.readFile path
     absPath <- makeAbsolute path
     let initState = TestParserState
-            { testVars = map (fmap someVarValueType) builtins
+            { testVars = concat
+                [ map (fmap someVarValueType) builtins
+                ]
             , testContext = SomeExpr RootNetwork
             }
-    case evalState (runParserT (parseTestModule absPath) path content) initState of
+        (res, _) = flip evalState initState $ runWriterT $ runParserT (parseTestModule absPath) path content
+
+    case res of
          Left err -> putStr (errorBundlePretty err) >> exitFailure
          Right testModule -> return testModule
diff --git a/src/Parser/Core.hs b/src/Parser/Core.hs
index b932523..341d9ca 100644
--- a/src/Parser/Core.hs
+++ b/src/Parser/Core.hs
@@ -2,6 +2,7 @@ module Parser.Core where
 
 import Control.Monad
 import Control.Monad.State
+import Control.Monad.Writer
 
 import Data.Text (Text)
 import qualified Data.Text.Lazy as TL
@@ -15,10 +16,13 @@ import qualified Text.Megaparsec.Char.Lexer as L
 import Network ()
 import Test
 
-type TestParser = ParsecT Void TestStream (State TestParserState)
+type TestParser = ParsecT Void TestStream (WriterT [ Toplevel ] (State TestParserState))
 
 type TestStream = TL.Text
 
+data Toplevel
+    = ToplevelTest Test
+
 data TestParserState = TestParserState
     { testVars :: [(VarName, SomeExprType)]
     , testContext :: SomeExpr
@@ -65,8 +69,8 @@ localState inner = do
     put s
     return x
 
-toplevel :: TestParser a -> TestParser a
-toplevel = L.nonIndented scn
+toplevel :: (a -> Toplevel) -> TestParser a -> TestParser ()
+toplevel f = tell . (: []) . f <=< L.nonIndented scn
 
 block :: (a -> [b] -> TestParser c) -> TestParser a -> TestParser b -> TestParser c
 block merge header item = L.indentBlock scn $ do