1 files changed, 51 insertions, 19 deletions

diff --git a/src/Parser/Statement.hs b/src/Parser/Statement.hs
index 7c2977d..474fa03 100644
--- a/src/Parser/Statement.hs
+++ b/src/Parser/Statement.hs
@@ -1,5 +1,6 @@
 module Parser.Statement (
     testStep,
+    testBlock,
 ) where
 
 import Control.Monad
@@ -43,7 +44,7 @@ letStatement = do
         addVarName off tname
         void $ eol
         body <- testBlock indent
-        return $ Let line tname e body
+        return $ Let line tname e (TestBlockStep EmptyTestBlock . Scope <$> body)
 
 forStatement :: TestParser (Expr (TestBlock ()))
 forStatement = do
@@ -68,23 +69,52 @@ forStatement = do
         body <- testBlock indent
         return $ (\xs f -> mconcat $ map f xs)
             <$> (unpack <$> e)
-            <*> LambdaAbstraction tname body
+            <*> LambdaAbstraction tname (TestBlockStep EmptyTestBlock . Scope <$> body)
 
 shellStatement :: TestParser (Expr (TestBlock ()))
 shellStatement = do
     ref <- L.indentLevel
     wsymbol "shell"
-    wsymbol "as"
-    pname <- newVarName
-    wsymbol "on"
-    node <- typedExpr
-    symbol ":"
-    void eol
-    void $ L.indentGuard scn GT ref
-    script <- shellScript
-    cont <- testBlock ref
-    return $ TestBlockStep EmptyTestBlock <$>
-        (SpawnShell pname <$> node <*> script <*> LambdaAbstraction pname cont)
+    parseParams ref Nothing Nothing
+
+  where
+    parseParamKeyword kw prev = do
+        off <- stateOffset <$> getParserState
+        wsymbol kw
+        when (isJust prev) $ do
+            registerParseError $ FancyError off $ S.singleton $ ErrorFail $
+                "unexpected parameter with keyword ‘" <> kw <> "’"
+
+    parseParams ref mbpname mbnode = choice
+        [ do
+            parseParamKeyword "as" mbpname
+            pname <- newVarName
+            parseParams ref (Just pname) mbnode
+
+        , do
+            parseParamKeyword "on" mbnode
+            node <- typedExpr
+            parseParams ref mbpname (Just node)
+
+        , do
+            off <- stateOffset <$> getParserState
+            symbol ":"
+            node <- case mbnode of
+                Just node -> return node
+                Nothing -> do
+                    registerParseError $ FancyError off $ S.singleton $ ErrorFail $
+                        "missing parameter with keyword ‘on’"
+                    return $ Undefined ""
+
+            void eol
+            void $ L.indentGuard scn GT ref
+            script <- shellScript
+            cont <- fmap Scope <$> testBlock ref
+            let expr | Just pname <- mbpname = LambdaAbstraction pname cont
+                     | otherwise = const <$> cont
+            return $ TestBlockStep EmptyTestBlock <$>
+                (SpawnShell mbpname <$> node <*> script <*> expr)
+        ]
 
 exprStatement :: TestParser (Expr (TestBlock ()))
 exprStatement = do
@@ -261,14 +291,14 @@ instance ExprType a => ParamType (InnerBlock a) where
         combine f (x : xs) = f x xs
         combine _ [] = error "inner block parameter count mismatch"
 
-innerBlock :: CommandDef (TestBlock ())
+innerBlock :: CommandDef (TestStep ())
 innerBlock = ($ ([] :: [ Void ])) <$> innerBlockFun
 
-innerBlockFun :: ExprType a => CommandDef (a -> TestBlock ())
+innerBlockFun :: ExprType a => CommandDef (a -> TestStep ())
 innerBlockFun = (\f x -> f [ x ]) <$> innerBlockFunList
 
-innerBlockFunList :: ExprType a => CommandDef ([ a ] -> TestBlock ())
-innerBlockFunList = fromInnerBlock <$> param ""
+innerBlockFunList :: ExprType a => CommandDef ([ a ] -> TestStep ())
+innerBlockFunList = (\ib -> Scope . fromInnerBlock ib) <$> param ""
 
 newtype ExprParam a = ExprParam { fromExprParam :: a }
     deriving (Functor, Foldable, Traversable)
@@ -351,7 +381,8 @@ testLocal = do
     void $ eol
 
     indent <- L.indentGuard scn GT ref
-    localState $ testBlock indent
+    localState $ do
+        fmap (TestBlockStep EmptyTestBlock . Scope) <$> testBlock indent
 
 testWith :: TestParser (Expr (TestBlock ()))
 testWith = do
@@ -377,7 +408,7 @@ testWith = do
     indent <- L.indentGuard scn GT ref
     localState $ do
         modify $ \s -> s { testContext = ctx }
-        testBlock indent
+        fmap (TestBlockStep EmptyTestBlock . Scope) <$> testBlock indent
 
 testSubnet :: TestParser (Expr (TestBlock ()))
 testSubnet = command "subnet" $ Subnet
@@ -395,6 +426,7 @@ testSpawn :: TestParser (Expr (TestBlock ()))
 testSpawn = command "spawn" $ Spawn
     <$> param "as"
     <*> (bimap fromExprParam fromExprParam <$> paramOrContext "on")
+    <*> (maybe [] fromExprParam <$> param "args")
     <*> innerBlockFun
 
 testExpect :: TestParser (Expr (TestBlock ()))