Remove embedded Expr for some inner TestBlock parameters

3 files changed, 34 insertions, 30 deletions

diff --git a/src/Parser/Statement.hs b/src/Parser/Statement.hs
index 6dc3c56..21b24a6 100644
--- a/src/Parser/Statement.hs
+++ b/src/Parser/Statement.hs
@@ -191,19 +191,22 @@ paramOrContext name = fromParamOrContext <$> param name
 cmdLine :: CommandDef SourceLine
 cmdLine = param ""
 
-newtype InnerBlock = InnerBlock { fromInnerBlock :: Expr TestBlock }
+newtype InnerBlock = InnerBlock { fromInnerBlock :: TestBlock }
 
 instance ParamType InnerBlock where
+    type ParamRep InnerBlock = Expr TestBlock
     parseParam _ = mzero
     showParamType _ = "<code block>"
+    paramExpr = fmap InnerBlock
 
-instance ParamType TestStep where
-    parseParam _ = mzero
-    showParamType _ = "<code line>"
-
-innerBlock :: CommandDef (Expr TestBlock)
+innerBlock :: CommandDef TestBlock
 innerBlock = fromInnerBlock <$> param ""
 
+innerBlockExpr :: CommandDef (Expr TestBlock)
+innerBlockExpr =
+    let CommandDef args fun = param ""
+     in CommandDef args (Pure . fmap fromInnerBlock . fun)
+
 newtype ExprParam a = ExprParam { fromExprParam :: a }
     deriving (Functor, Foldable, Traversable)
 
@@ -230,7 +233,7 @@ command name (CommandDef types ctor) = do
             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 . InnerBlock <$> restOfParts cmdi partials
+                    | Just (Refl :: p :~: InnerBlock) <- eqT -> SomeParam p . Identity <$> restOfParts cmdi partials
                 (sym, SomeParam p Nothing) -> choice
                     [ SomeParam p . Identity <$> paramDefault p
                     , fail $ "missing " ++ (if null sym then "" else "'" ++ sym ++ "' ") ++ showParamType p
@@ -312,19 +315,19 @@ testSubnet :: TestParser (Expr TestBlock)
 testSubnet = command "subnet" $ Subnet
     <$> param ""
     <*> (fromExprParam <$> paramOrContext "of")
-    <*> innerBlock
+    <*> innerBlockExpr
 
 testNode :: TestParser (Expr TestBlock)
 testNode = command "node" $ DeclNode
     <$> param ""
     <*> (fromExprParam <$> paramOrContext "on")
-    <*> innerBlock
+    <*> innerBlockExpr
 
 testSpawn :: TestParser (Expr TestBlock)
 testSpawn = command "spawn" $ Spawn
     <$> param "as"
     <*> (bimap fromExprParam fromExprParam <$> paramOrContext "on")
-    <*> innerBlock
+    <*> innerBlockExpr
 
 testExpect :: TestParser (Expr TestBlock)
 testExpect = command "expect" $ Expect
@@ -332,7 +335,7 @@ testExpect = command "expect" $ Expect
     <*> (fromExprParam <$> paramOrContext "from")
     <*> param ""
     <*> param "capture"
-    <*> innerBlock
+    <*> innerBlockExpr
 
 testDisconnectNode :: TestParser (Expr TestBlock)
 testDisconnectNode = command "disconnect_node" $ DisconnectNode
diff --git a/src/Run.hs b/src/Run.hs
index 76545e4..31a3c9e 100644
--- a/src/Run.hs
+++ b/src/Run.hs
@@ -162,18 +162,18 @@ evalBlock (TestBlock steps) = forM_ steps $ \case
         testStepGuard line vars expr
 
     DisconnectNode node inner -> do
-        withDisconnectedUp (nodeUpstream node) $ evalBlock =<< eval inner
+        withDisconnectedUp (nodeUpstream node) $ evalBlock inner
 
     DisconnectNodes net inner -> do
-        withDisconnectedBridge (netBridge net) $ evalBlock =<< eval inner
+        withDisconnectedBridge (netBridge net) $ evalBlock inner
 
     DisconnectUpstream net inner -> do
         case netUpstream net of
-            Just link -> withDisconnectedUp link $ evalBlock =<< eval inner
-            Nothing -> evalBlock =<< eval inner
+            Just link -> withDisconnectedUp link $ evalBlock inner
+            Nothing -> evalBlock inner
 
     PacketLoss loss node inner -> do
-        withNodePacketLoss node loss $ evalBlock =<< eval inner
+        withNodePacketLoss node loss $ evalBlock inner
 
     Wait -> do
         void $ outPromptGetLine "Waiting..."
diff --git a/src/Test.hs b/src/Test.hs
index 772ac28..da4c82d 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -58,20 +58,21 @@ data Test = Test
 newtype TestBlock = TestBlock [ TestStep ]
     deriving (Semigroup, Monoid)
 
-data TestStep = forall a. ExprType a => Let SourceLine (TypedVarName a) (Expr a) (Expr TestBlock)
-              | forall a. ExprType a => For SourceLine (TypedVarName a) (Expr [ a ]) (Expr TestBlock)
-              | Subnet (TypedVarName Network) Network (Expr TestBlock)
-              | DeclNode (TypedVarName Node) Network (Expr TestBlock)
-              | Spawn (TypedVarName Process) (Either Network Node) (Expr TestBlock)
-              | Send Process Text
-              | Expect SourceLine Process (Expr Regex) [ TypedVarName Text ] (Expr TestBlock)
-              | Flush Process (Maybe Regex)
-              | Guard SourceLine EvalTrace Bool
-              | DisconnectNode Node (Expr TestBlock)
-              | DisconnectNodes Network (Expr TestBlock)
-              | DisconnectUpstream Network (Expr TestBlock)
-              | PacketLoss Scientific Node (Expr TestBlock)
-              | Wait
+data TestStep
+    = forall a. ExprType a => Let SourceLine (TypedVarName a) (Expr a) (Expr TestBlock)
+    | forall a. ExprType a => For SourceLine (TypedVarName a) (Expr [ a ]) (Expr TestBlock)
+    | Subnet (TypedVarName Network) Network (Expr TestBlock)
+    | DeclNode (TypedVarName Node) Network (Expr TestBlock)
+    | Spawn (TypedVarName Process) (Either Network Node) (Expr TestBlock)
+    | Send Process Text
+    | Expect SourceLine Process (Expr Regex) [ TypedVarName Text ] (Expr TestBlock)
+    | Flush Process (Maybe Regex)
+    | Guard SourceLine EvalTrace Bool
+    | DisconnectNode Node TestBlock
+    | DisconnectNodes Network TestBlock
+    | DisconnectUpstream Network TestBlock
+    | PacketLoss Scientific Node TestBlock
+    | Wait
 
 newtype SourceLine = SourceLine Text