1 files changed, 43 insertions, 41 deletions

diff --git a/src/Parser/Statement.hs b/src/Parser/Statement.hs
index b2f3cd6..c7cdf5a 100644
--- a/src/Parser/Statement.hs
+++ b/src/Parser/Statement.hs
@@ -8,9 +8,8 @@ import Control.Monad.State
 
 import Data.Kind
 import Data.Maybe
-import qualified Data.Set as S
+import Data.Set qualified as S
 import Data.Text qualified as T
-import qualified Data.Text.Lazy as TL
 import Data.Typeable
 
 import Text.Megaparsec hiding (State)
@@ -24,16 +23,6 @@ import Process (Process)
 import Test
 import Util
 
-getSourceLine :: TestParser SourceLine
-getSourceLine = do
-    pstate <- statePosState <$> getParserState
-    return $ SourceLine $ T.concat
-        [ T.pack $ sourcePosPretty $ pstateSourcePos pstate
-        , T.pack ": "
-        , TL.toStrict $ TL.takeWhile (/='\n') $ pstateInput pstate
-        ]
-
-
 letStatement :: TestParser [TestStep]
 letStatement = do
     line <- getSourceLine
@@ -76,9 +65,34 @@ forStatement = do
         return [For line tname (unpack <$> e) body]
 
 exprStatement :: TestParser [ TestStep ]
-exprStatement  = do
-    expr <- typedExpr
-    return [ ExprStatement expr ]
+exprStatement = do
+    ref <- L.indentLevel
+    off <- stateOffset <$> getParserState
+    SomeExpr expr <- someExpr
+    choice
+        [ do
+            continuePartial off ref expr
+        , do
+            stmt <- unifyExpr off Proxy expr
+            return [ ExprStatement stmt ]
+        ]
+  where
+    continuePartial :: ExprType a => Int -> Pos -> Expr a -> TestParser [ TestStep ]
+    continuePartial off ref expr = do
+        symbol ":"
+        void eol
+        (fun :: Expr (FunctionType TestBlock)) <- unifyExpr off Proxy expr
+        scn
+        indent <- L.indentGuard scn GT ref
+        blockOf indent $ do
+            coff <- stateOffset <$> getParserState
+            sline <- getSourceLine
+            args <- functionArguments (checkFunctionArguments (exprArgs fun)) someExpr literal (\poff -> lookupVarExpr poff sline . VarName)
+            let fun' = ArgsApp args fun
+            choice
+                [ continuePartial coff indent fun'
+                , (: []) . ExprStatement <$> unifyExpr coff Proxy fun'
+                ]
 
 class (Typeable a, Typeable (ParamRep a)) => ParamType a where
     type ParamRep a :: Type
@@ -102,7 +116,10 @@ instance ExprType a => ParamType (TypedVarName a) where
     showParamType _ = "<variable>"
 
 instance ExprType a => ParamType (Expr a) where
-    parseParam _ = typedExpr
+    parseParam _ = do
+        off <- stateOffset <$> getParserState
+        SomeExpr e <- literal <|> variable <|> between (symbol "(") (symbol ")") someExpr
+        unifyExpr off Proxy e
     showParamType _ = "<" ++ T.unpack (textExprType @a Proxy) ++ ">"
 
 instance ParamType a => ParamType [a] where
@@ -259,12 +276,12 @@ testWith = do
     off <- stateOffset <$> getParserState
     ctx@(SomeExpr (_ :: Expr ctxe)) <- someExpr
     let expected =
-            [ SomeExprType @Network Proxy
-            , SomeExprType @Node Proxy
-            , SomeExprType @Process Proxy
+            [ ExprTypePrim @Network Proxy
+            , ExprTypePrim @Node Proxy
+            , ExprTypePrim @Process Proxy
             ]
     notAllowed <- flip allM expected $ \case
-        SomeExprType (Proxy :: Proxy a) | Just (Refl :: ctxe :~: a) <- eqT -> return False
+        ExprTypePrim (Proxy :: Proxy a) | Just (Refl :: ctxe :~: a) <- eqT -> return False
         _ -> return True
     when notAllowed $ registerParseError $ FancyError off $ S.singleton $ ErrorFail $ T.unpack $
         "expected " <> T.intercalate ", " (map (("'"<>) . (<>"'") . textSomeExprType) expected) <> ", expression has type '" <> textExprType @ctxe Proxy <> "'"
@@ -295,11 +312,6 @@ testSpawn = command "spawn" $ Spawn
     <*> paramOrContext "on"
     <*> innerBlock
 
-testSend :: TestParser [TestStep]
-testSend = command "send" $ Send
-    <$> paramOrContext "to"
-    <*> param ""
-
 testExpect :: TestParser [TestStep]
 testExpect = command "expect" $ Expect
     <$> cmdLine
@@ -308,16 +320,6 @@ testExpect = command "expect" $ Expect
     <*> param "capture"
     <*> innerBlock
 
-testFlush :: TestParser [TestStep]
-testFlush = command "flush" $ Flush
-    <$> paramOrContext "from"
-    <*> param ""
-
-testGuard :: TestParser [TestStep]
-testGuard = command "guard" $ Guard
-    <$> cmdLine
-    <*> param ""
-
 testDisconnectNode :: TestParser [TestStep]
 testDisconnectNode = command "disconnect_node" $ DisconnectNode
     <$> paramOrContext ""
@@ -340,8 +342,11 @@ testPacketLoss = command "packet_loss" $ PacketLoss
     <*> innerBlock
 
 
-testBlock :: Pos -> TestParser [TestStep]
-testBlock indent = concat <$> go
+testBlock :: Pos -> TestParser [ TestStep ]
+testBlock indent = blockOf indent testStep
+
+blockOf :: Pos -> TestParser [ a ] -> TestParser [ a ]
+blockOf indent step = concat <$> go
   where
     go = do
         scn
@@ -349,7 +354,7 @@ testBlock indent = concat <$> go
         optional eof >>= \case
             Just _ -> return []
             _ | pos <  indent -> return []
-              | pos == indent -> (:) <$> testStep <*> go
+              | pos == indent -> (:) <$> step <*> go
               | otherwise     -> L.incorrectIndent EQ indent pos
 
 testStep :: TestParser [TestStep]
@@ -361,10 +366,7 @@ testStep = choice
     , testSubnet
     , testNode
     , testSpawn
-    , testSend
     , testExpect
-    , testFlush
-    , testGuard
     , testDisconnectNode
     , testDisconnectNodes
     , testDisconnectUpstream