1 files changed, 192 insertions, 86 deletions

diff --git a/src/Parser/Statement.hs b/src/Parser/Statement.hs
index c7cdf5a..474fa03 100644
--- a/src/Parser/Statement.hs
+++ b/src/Parser/Statement.hs
@@ -1,16 +1,19 @@
 module Parser.Statement (
     testStep,
+    testBlock,
 ) where
 
 import Control.Monad
 import Control.Monad.Identity
 import Control.Monad.State
 
+import Data.Bifunctor
 import Data.Kind
 import Data.Maybe
 import Data.Set qualified as S
 import Data.Text qualified as T
 import Data.Typeable
+import Data.Void
 
 import Text.Megaparsec hiding (State)
 import Text.Megaparsec.Char
@@ -19,11 +22,14 @@ import qualified Text.Megaparsec.Char.Lexer as L
 import Network (Network, Node)
 import Parser.Core
 import Parser.Expr
+import Parser.Shell
 import Process (Process)
+import Script.Expr
+import Script.Expr.Class
 import Test
 import Util
 
-letStatement :: TestParser [TestStep]
+letStatement :: TestParser (Expr (TestBlock ()))
 letStatement = do
     line <- getSourceLine
     indent <- L.indentLevel
@@ -38,11 +44,10 @@ 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 [TestStep]
+forStatement :: TestParser (Expr (TestBlock ()))
 forStatement = do
-    line <- getSourceLine
     ref <- L.indentLevel
     wsymbol "for"
     voff <- stateOffset <$> getParserState
@@ -62,26 +67,70 @@ forStatement = do
         let tname = TypedVarName name
         addVarName voff tname
         body <- testBlock indent
-        return [For line tname (unpack <$> e) body]
+        return $ (\xs f -> mconcat $ map f xs)
+            <$> (unpack <$> e)
+            <*> LambdaAbstraction tname (TestBlockStep EmptyTestBlock . Scope <$> body)
+
+shellStatement :: TestParser (Expr (TestBlock ()))
+shellStatement = do
+    ref <- L.indentLevel
+    wsymbol "shell"
+    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 [ TestStep ]
+exprStatement :: TestParser (Expr (TestBlock ()))
 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 ]
+        [ continuePartial off ref expr
+        , unifyExpr off Proxy expr
         ]
   where
-    continuePartial :: ExprType a => Int -> Pos -> Expr a -> TestParser [ TestStep ]
+    continuePartial :: ExprType a => Int -> Pos -> Expr a -> TestParser (Expr (TestBlock ()))
     continuePartial off ref expr = do
         symbol ":"
         void eol
-        (fun :: Expr (FunctionType TestBlock)) <- unifyExpr off Proxy expr
+        (fun :: Expr (FunctionType (TestBlock ()))) <- unifyExpr off Proxy expr
         scn
         indent <- L.indentGuard scn GT ref
         blockOf indent $ do
@@ -91,7 +140,7 @@ exprStatement = do
             let fun' = ArgsApp args fun
             choice
                 [ continuePartial coff indent fun'
-                , (: []) . ExprStatement <$> unifyExpr coff Proxy fun'
+                , unifyExpr coff Proxy fun'
                 ]
 
 class (Typeable a, Typeable (ParamRep a)) => ParamType a where
@@ -104,9 +153,18 @@ class (Typeable a, Typeable (ParamRep a)) => ParamType a where
     paramDefault :: proxy a -> TestParser (ParamRep a)
     paramDefault _ = mzero
 
+    paramNewVariables :: proxy a -> ParamRep a -> NewVariables
+    paramNewVariables _ _ = NoNewVariables
+    paramNewVariablesEmpty :: proxy a -> NewVariables
+    paramNewVariablesEmpty _ = NoNewVariables -- to keep type info for optional parameters
+
     paramFromSomeExpr :: proxy a -> SomeExpr -> Maybe (ParamRep a)
     paramFromSomeExpr _ (SomeExpr e) = cast e
 
+    paramExpr :: ParamRep a -> Expr a
+    default paramExpr :: ParamRep a ~ a => ParamRep a -> Expr a
+    paramExpr = Pure
+
 instance ParamType SourceLine where
     parseParam _ = mzero
     showParamType _ = "<source line>"
@@ -114,11 +172,13 @@ instance ParamType SourceLine where
 instance ExprType a => ParamType (TypedVarName a) where
     parseParam _ = newVarName
     showParamType _ = "<variable>"
+    paramNewVariables _ var = SomeNewVariables [ var ]
+    paramNewVariablesEmpty _ = SomeNewVariables @a []
 
 instance ExprType a => ParamType (Expr a) where
     parseParam _ = do
         off <- stateOffset <$> getParserState
-        SomeExpr e <- literal <|> variable <|> between (symbol "(") (symbol ")") someExpr
+        SomeExpr e <- literal <|> between (symbol "(") (symbol ")") someExpr
         unifyExpr off Proxy e
     showParamType _ = "<" ++ T.unpack (textExprType @a Proxy) ++ ">"
 
@@ -127,14 +187,20 @@ instance ParamType a => ParamType [a] where
     parseParam _ = listOf (parseParam @a Proxy)
     showParamType _ = showParamType @a Proxy ++ " [, " ++ showParamType @a Proxy ++ " ...]"
     paramDefault _ = return []
+    paramNewVariables _ = foldr (<>) (paramNewVariablesEmpty @a Proxy) . fmap (paramNewVariables @a Proxy)
+    paramNewVariablesEmpty _ = paramNewVariablesEmpty @a Proxy
     paramFromSomeExpr _ se@(SomeExpr e) = cast e <|> ((:[]) <$> paramFromSomeExpr @a Proxy se)
+    paramExpr = sequenceA . fmap paramExpr
 
 instance ParamType a => ParamType (Maybe a) where
     type ParamRep (Maybe a) = Maybe (ParamRep a)
     parseParam _ = Just <$> parseParam @a Proxy
     showParamType _ = showParamType @a Proxy
     paramDefault _ = return Nothing
+    paramNewVariables _ = foldr (<>) (paramNewVariablesEmpty @a Proxy) . fmap (paramNewVariables @a Proxy)
+    paramNewVariablesEmpty _ = paramNewVariablesEmpty @a Proxy
     paramFromSomeExpr _ se = Just <$> paramFromSomeExpr @a Proxy se
+    paramExpr = sequenceA . fmap paramExpr
 
 instance (ParamType a, ParamType b) => ParamType (Either a b) where
     type ParamRep (Either a b) = Either (ParamRep a) (ParamRep b)
@@ -147,62 +213,106 @@ instance (ParamType a, ParamType b) => ParamType (Either a b) where
                 (_ : _) -> fail ""
     showParamType _ = showParamType @a Proxy ++ " or " ++ showParamType @b Proxy
     paramFromSomeExpr _ se = (Left <$> paramFromSomeExpr @a Proxy se) <|> (Right <$> paramFromSomeExpr @b Proxy se)
+    paramExpr = either (fmap Left . paramExpr) (fmap Right . paramExpr)
+
+instance ExprType a => ParamType (Traced a) where
+    type ParamRep (Traced a) = Expr a
+    parseParam _ = parseParam (Proxy @(Expr a))
+    showParamType _ = showParamType (Proxy @(Expr a))
+    paramExpr = Trace
 
 data SomeParam f = forall a. ParamType a => SomeParam (Proxy a) (f (ParamRep a))
 
-data CommandDef a = CommandDef [(String, SomeParam Proxy)] ([SomeParam Identity] -> a)
+data NewVariables
+    = NoNewVariables
+    | forall a. ExprType a => SomeNewVariables [ TypedVarName a ]
+
+instance Semigroup NewVariables where
+    NoNewVariables <> x = x
+    x <> NoNewVariables = x
+    SomeNewVariables (xs :: [ TypedVarName a ]) <> SomeNewVariables (ys :: [ TypedVarName b ])
+        | Just (Refl :: a :~: b) <- eqT = SomeNewVariables (xs <> ys)
+        | otherwise = error "new variables with different types"
+
+instance Monoid NewVariables where
+    mempty = NoNewVariables
+
+someParamVars :: Foldable f => SomeParam f -> NewVariables
+someParamVars (SomeParam proxy rep) = foldr (\x nvs -> paramNewVariables proxy x <> nvs) (paramNewVariablesEmpty proxy) rep
+
+data CommandDef a = CommandDef [(String, SomeParam Proxy)] ([SomeParam Identity] -> Expr a)
 
 instance Functor CommandDef where
-  fmap f (CommandDef types ctor) = CommandDef types (f . ctor)
+    fmap f (CommandDef types ctor) = CommandDef types (fmap f . ctor)
 
 instance Applicative CommandDef where
-  pure x = CommandDef [] (\case [] -> x; _ -> error "command arguments mismatch")
-  CommandDef types1 ctor1 <*> CommandDef types2 ctor2 =
-      CommandDef (types1 ++ types2) $ \params ->
-          let (params1, params2) = splitAt (length types1) params
-           in ctor1 params1 $ ctor2 params2
+    pure x = CommandDef [] (\case [] -> Pure x; _ -> error "command arguments mismatch")
+    CommandDef types1 ctor1 <*> CommandDef types2 ctor2 =
+        CommandDef (types1 ++ types2) $ \params ->
+            let (params1, params2) = splitAt (length types1) params
+             in ctor1 params1 <*> ctor2 params2
 
 param :: forall a. ParamType a => String -> CommandDef a
 param name = CommandDef [(name, SomeParam (Proxy @a) Proxy)] $ \case
-    [SomeParam Proxy (Identity x)] -> fromJust $ cast x
+    [SomeParam Proxy (Identity x)] -> paramExpr $ fromJust $ cast x
     _ -> error "command arguments mismatch"
 
-data ParamOrContext a
+newtype ParamOrContext a = ParamOrContext { fromParamOrContext :: a }
+    deriving (Functor, Foldable, Traversable)
 
 instance ParamType a => ParamType (ParamOrContext a) where
-    type ParamRep (ParamOrContext a) = ParamRep a
-    parseParam _ = parseParam @a Proxy
+    type ParamRep (ParamOrContext a) = ParamOrContext (ParamRep a)
+    parseParam _ = ParamOrContext <$> parseParam @a Proxy
     showParamType _ = showParamType @a Proxy
     paramDefault _ = gets testContext >>= \case
         se@(SomeExpr ctx)
-            | Just e <- paramFromSomeExpr @a Proxy se -> return e
+            | Just e <- paramFromSomeExpr @a Proxy se -> return (ParamOrContext e)
             | otherwise -> fail $ showParamType @a Proxy <> " not available from context type '" <> T.unpack (textExprType ctx) <> "'"
+    paramExpr = sequenceA . fmap paramExpr
 
 paramOrContext :: forall a. ParamType a => String -> CommandDef a
-paramOrContext name = CommandDef [(name, SomeParam (Proxy @(ParamOrContext a)) Proxy)] $ \case
-    [SomeParam Proxy (Identity x)] -> fromJust $ cast x
-    _ -> error "command arguments mismatch"
+paramOrContext name = fromParamOrContext <$> param name
 
 cmdLine :: CommandDef SourceLine
 cmdLine = param ""
 
-data InnerBlock
+newtype InnerBlock a = InnerBlock { fromInnerBlock :: [ a ] -> TestBlock () }
 
-instance ParamType InnerBlock where
-    type ParamRep InnerBlock = [TestStep]
+instance ExprType a => ParamType (InnerBlock a) where
+    type ParamRep (InnerBlock a) = ( [ TypedVarName a ], Expr (TestBlock ()) )
     parseParam _ = mzero
     showParamType _ = "<code block>"
+    paramExpr ( vars, expr ) = fmap InnerBlock $ helper vars $ const <$> expr
+      where
+        helper :: ExprType a => [ TypedVarName a ] -> Expr ([ a ] -> b) -> Expr ([ a ] -> b)
+        helper ( v : vs ) = fmap combine . LambdaAbstraction v . helper vs
+        helper [] = id
 
-instance ParamType TestStep where
-    parseParam _ = mzero
-    showParamType _ = "<code line>"
+        combine f (x : xs) = f x xs
+        combine _ [] = error "inner block parameter count mismatch"
 
-innerBlock :: CommandDef [TestStep]
-innerBlock = CommandDef [("", SomeParam (Proxy @InnerBlock) Proxy)] $ \case
-    [SomeParam Proxy (Identity x)] -> fromJust $ cast x
-    _ -> error "command arguments mismatch"
+innerBlock :: CommandDef (TestStep ())
+innerBlock = ($ ([] :: [ Void ])) <$> innerBlockFun
+
+innerBlockFun :: ExprType a => CommandDef (a -> TestStep ())
+innerBlockFun = (\f x -> f [ x ]) <$> innerBlockFunList
+
+innerBlockFunList :: ExprType a => CommandDef ([ a ] -> TestStep ())
+innerBlockFunList = (\ib -> Scope . fromInnerBlock ib) <$> param ""
+
+newtype ExprParam a = ExprParam { fromExprParam :: a }
+    deriving (Functor, Foldable, Traversable)
+
+instance ExprType a => ParamType (ExprParam a) where
+    type ParamRep (ExprParam a) = Expr a
+    parseParam _ = do
+        off <- stateOffset <$> getParserState
+        SomeExpr e <- literal <|> variable <|> between (symbol "(") (symbol ")") someExpr
+        unifyExpr off Proxy e
+    showParamType _ = "<" ++ T.unpack (textExprType @a Proxy) ++ ">"
+    paramExpr = fmap ExprParam
 
-command :: String -> CommandDef TestStep -> TestParser [TestStep]
+command :: String -> CommandDef (TestStep ()) -> TestParser (Expr (TestBlock ()))
 command name (CommandDef types ctor) = do
     indent <- L.indentLevel
     line <- getSourceLine
@@ -210,19 +320,24 @@ command name (CommandDef types ctor) = do
     localState $ do
         restOfLine indent [] line $ map (fmap $ \(SomeParam p@(_ :: Proxy p) Proxy) -> SomeParam p $ Nothing @(ParamRep p)) types
   where
-    restOfLine :: Pos -> [(Pos, [(String, SomeParam Maybe)])] -> SourceLine -> [(String, SomeParam Maybe)] -> TestParser [TestStep]
+    restOfLine :: Pos -> [(Pos, [(String, SomeParam Maybe)])] -> SourceLine -> [(String, SomeParam Maybe)] -> TestParser (Expr (TestBlock ()))
     restOfLine cmdi partials line params = choice
         [do void $ lookAhead eol
+            let definedVariables = mconcat $ map (someParamVars . snd) params
             iparams <- forM params $ \case
                 (_, SomeParam (p :: Proxy p) Nothing)
                     | Just (Refl :: p :~: SourceLine) <- eqT -> return $ SomeParam p $ Identity line
-                    | Just (Refl :: p :~: InnerBlock) <- eqT -> SomeParam p . Identity <$> restOfParts cmdi partials
+
+                    | SomeNewVariables (vars :: [ TypedVarName a ]) <- definedVariables
+                    , Just (Refl :: p :~: InnerBlock a) <- eqT
+                    -> SomeParam p . Identity . ( vars, ) <$> restOfParts cmdi partials
+
                 (sym, SomeParam p Nothing) -> choice
                     [ SomeParam p . Identity <$> paramDefault p
                     , fail $ "missing " ++ (if null sym then "" else "'" ++ sym ++ "' ") ++ showParamType p
                     ]
                 (_, SomeParam (p :: Proxy p) (Just x)) -> return $ SomeParam p $ Identity x
-            return [ctor iparams]
+            return $ (TestBlockStep EmptyTestBlock) <$> ctor iparams
 
         ,do symbol ":"
             scn
@@ -232,16 +347,16 @@ command name (CommandDef types ctor) = do
         ,do tryParams cmdi partials line [] params
         ]
 
-    restOfParts :: Pos -> [(Pos, [(String, SomeParam Maybe)])] -> TestParser [TestStep]
+    restOfParts :: Pos -> [(Pos, [(String, SomeParam Maybe)])] -> TestParser (Expr (TestBlock ()))
     restOfParts cmdi [] = testBlock cmdi
     restOfParts cmdi partials@((partIndent, params) : rest) = do
         scn
         pos <- L.indentLevel
         line <- getSourceLine
         optional eof >>= \case
-            Just _ -> return []
+            Just _ -> return $ Pure mempty
             _ | pos <  partIndent -> restOfParts cmdi rest
-              | pos == partIndent -> (++) <$> restOfLine cmdi partials line params <*> restOfParts cmdi partials
+              | pos == partIndent -> mappend <$> restOfLine cmdi partials line params <*> restOfParts cmdi partials
               | otherwise         -> L.incorrectIndent EQ partIndent pos
 
     tryParam sym (SomeParam (p :: Proxy p) cur) = do
@@ -258,7 +373,7 @@ command name (CommandDef types ctor) = do
         ]
     tryParams _ _ _ _ [] = mzero
 
-testLocal :: TestParser [TestStep]
+testLocal :: TestParser (Expr (TestBlock ()))
 testLocal = do
     ref <- L.indentLevel
     wsymbol "local"
@@ -266,9 +381,10 @@ testLocal = do
     void $ eol
 
     indent <- L.indentGuard scn GT ref
-    localState $ testBlock indent
+    localState $ do
+        fmap (TestBlockStep EmptyTestBlock . Scope) <$> testBlock indent
 
-testWith :: TestParser [TestStep]
+testWith :: TestParser (Expr (TestBlock ()))
 testWith = do
     ref <- L.indentLevel
     wsymbol "with"
@@ -292,75 +408,65 @@ 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 [TestStep]
+testSubnet :: TestParser (Expr (TestBlock ()))
 testSubnet = command "subnet" $ Subnet
     <$> param ""
-    <*> paramOrContext "of"
-    <*> innerBlock
+    <*> (fromExprParam <$> paramOrContext "of")
+    <*> innerBlockFun
 
-testNode :: TestParser [TestStep]
+testNode :: TestParser (Expr (TestBlock ()))
 testNode = command "node" $ DeclNode
     <$> param ""
-    <*> paramOrContext "on"
-    <*> innerBlock
+    <*> (fromExprParam <$> paramOrContext "on")
+    <*> innerBlockFun
 
-testSpawn :: TestParser [TestStep]
+testSpawn :: TestParser (Expr (TestBlock ()))
 testSpawn = command "spawn" $ Spawn
     <$> param "as"
-    <*> paramOrContext "on"
-    <*> innerBlock
+    <*> (bimap fromExprParam fromExprParam <$> paramOrContext "on")
+    <*> (maybe [] fromExprParam <$> param "args")
+    <*> innerBlockFun
 
-testExpect :: TestParser [TestStep]
+testExpect :: TestParser (Expr (TestBlock ()))
 testExpect = command "expect" $ Expect
     <$> cmdLine
-    <*> paramOrContext "from"
+    <*> (fromExprParam <$> paramOrContext "from")
     <*> param ""
     <*> param "capture"
-    <*> innerBlock
+    <*> innerBlockFunList
 
-testDisconnectNode :: TestParser [TestStep]
+testDisconnectNode :: TestParser (Expr (TestBlock ()))
 testDisconnectNode = command "disconnect_node" $ DisconnectNode
-    <$> paramOrContext ""
+    <$> (fromExprParam <$> paramOrContext "")
     <*> innerBlock
 
-testDisconnectNodes :: TestParser [TestStep]
+testDisconnectNodes :: TestParser (Expr (TestBlock ()))
 testDisconnectNodes = command "disconnect_nodes" $ DisconnectNodes
-    <$> paramOrContext ""
+    <$> (fromExprParam <$> paramOrContext "")
     <*> innerBlock
 
-testDisconnectUpstream :: TestParser [TestStep]
+testDisconnectUpstream :: TestParser (Expr (TestBlock ()))
 testDisconnectUpstream = command "disconnect_upstream" $ DisconnectUpstream
-    <$> paramOrContext ""
+    <$> (fromExprParam <$> paramOrContext "")
     <*> innerBlock
 
-testPacketLoss :: TestParser [TestStep]
+testPacketLoss :: TestParser (Expr (TestBlock ()))
 testPacketLoss = command "packet_loss" $ PacketLoss
-    <$> param ""
-    <*> paramOrContext "on"
+    <$> (fromExprParam <$> paramOrContext "")
+    <*> (fromExprParam <$> paramOrContext "on")
     <*> innerBlock
 
 
-testBlock :: Pos -> TestParser [ TestStep ]
+testBlock :: Pos -> TestParser (Expr (TestBlock ()))
 testBlock indent = blockOf indent testStep
 
-blockOf :: Pos -> TestParser [ a ] -> TestParser [ a ]
-blockOf indent step = concat <$> go
-  where
-    go = do
-        scn
-        pos <- L.indentLevel
-        optional eof >>= \case
-            Just _ -> return []
-            _ | pos <  indent -> return []
-              | pos == indent -> (:) <$> step <*> go
-              | otherwise     -> L.incorrectIndent EQ indent pos
-
-testStep :: TestParser [TestStep]
+testStep :: TestParser (Expr (TestBlock ()))
 testStep = choice
     [ letStatement
     , forStatement
+    , shellStatement
     , testLocal
     , testWith
     , testSubnet