Remove remaining Expr usage in TestStep

4 files changed, 110 insertions, 55 deletions

diff --git a/src/Parser/Statement.hs b/src/Parser/Statement.hs
index a65227d..4bed1ef 100644
--- a/src/Parser/Statement.hs
+++ b/src/Parser/Statement.hs
@@ -12,6 +12,7 @@ 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
@@ -39,11 +40,10 @@ letStatement = do
         addVarName off tname
         void $ eol
         body <- testBlock indent
-        return $ Pure $ TestBlock [ Let line tname e body ]
+        return $ Let line tname e body
 
 forStatement :: TestParser (Expr TestBlock)
 forStatement = do
-    line <- getSourceLine
     ref <- L.indentLevel
     wsymbol "for"
     voff <- stateOffset <$> getParserState
@@ -63,7 +63,9 @@ forStatement = do
         let tname = TypedVarName name
         addVarName voff tname
         body <- testBlock indent
-        return $ Pure $ TestBlock [ For line tname (unpack <$> e) body ]
+        return $ (\xs f -> mconcat $ map f xs)
+            <$> (unpack <$> e)
+            <*> LambdaAbstraction tname body
 
 exprStatement :: TestParser (Expr TestBlock)
 exprStatement = do
@@ -102,6 +104,11 @@ 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
 
@@ -116,6 +123,8 @@ 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
@@ -129,6 +138,8 @@ 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
 
@@ -137,6 +148,8 @@ instance ParamType a => ParamType (Maybe a) where
     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
 
@@ -161,6 +174,23 @@ instance ExprType a => ParamType (Traced a) where
 
 data SomeParam f = forall a. ParamType a => SomeParam (Proxy a) (f (ParamRep 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
@@ -197,21 +227,29 @@ paramOrContext name = fromParamOrContext <$> param name
 cmdLine :: CommandDef SourceLine
 cmdLine = param ""
 
-newtype InnerBlock = InnerBlock { fromInnerBlock :: TestBlock }
+newtype InnerBlock a = InnerBlock { fromInnerBlock :: [ a ] -> TestBlock }
 
-instance ParamType InnerBlock where
-    type ParamRep InnerBlock = Expr TestBlock
+instance ExprType a => ParamType (InnerBlock a) where
+    type ParamRep (InnerBlock a) = ( [ TypedVarName a ], Expr TestBlock )
     parseParam _ = mzero
     showParamType _ = "<code block>"
-    paramExpr = fmap InnerBlock
+    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
+
+        combine f (x : xs) = f x xs
+        combine _ [] = error "inner block parameter count mismatch"
 
 innerBlock :: CommandDef TestBlock
-innerBlock = fromInnerBlock <$> param ""
+innerBlock = ($ ([] :: [ Void ])) <$> innerBlockFun
+
+innerBlockFun :: ExprType a => CommandDef (a -> TestBlock)
+innerBlockFun = (\f x -> f [ x ]) <$> innerBlockFunList
 
-innerBlockExpr :: CommandDef (Expr TestBlock)
-innerBlockExpr =
-    let CommandDef args fun = param ""
-     in CommandDef args (Pure . fmap fromInnerBlock . fun)
+innerBlockFunList :: ExprType a => CommandDef ([ a ] -> TestBlock)
+innerBlockFunList = fromInnerBlock <$> param ""
 
 newtype ExprParam a = ExprParam { fromExprParam :: a }
     deriving (Functor, Foldable, Traversable)
@@ -236,10 +274,15 @@ command name (CommandDef types ctor) = do
     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
@@ -321,19 +364,19 @@ testSubnet :: TestParser (Expr TestBlock)
 testSubnet = command "subnet" $ Subnet
     <$> param ""
     <*> (fromExprParam <$> paramOrContext "of")
-    <*> innerBlockExpr
+    <*> innerBlockFun
 
 testNode :: TestParser (Expr TestBlock)
 testNode = command "node" $ DeclNode
     <$> param ""
     <*> (fromExprParam <$> paramOrContext "on")
-    <*> innerBlockExpr
+    <*> innerBlockFun
 
 testSpawn :: TestParser (Expr TestBlock)
 testSpawn = command "spawn" $ Spawn
     <$> param "as"
     <*> (bimap fromExprParam fromExprParam <$> paramOrContext "on")
-    <*> innerBlockExpr
+    <*> innerBlockFun
 
 testExpect :: TestParser (Expr TestBlock)
 testExpect = command "expect" $ Expect
@@ -341,7 +384,7 @@ testExpect = command "expect" $ Expect
     <*> (fromExprParam <$> paramOrContext "from")
     <*> param ""
     <*> param "capture"
-    <*> innerBlockExpr
+    <*> innerBlockFunList
 
 testDisconnectNode :: TestParser (Expr TestBlock)
 testDisconnectNode = command "disconnect_node" $ DisconnectNode
diff --git a/src/Run.hs b/src/Run.hs
index f6eba39..b623f52 100644
--- a/src/Run.hs
+++ b/src/Run.hs
@@ -110,32 +110,13 @@ runTest out opts test variables = do
 
 evalBlock :: TestBlock -> TestRun ()
 evalBlock (TestBlock steps) = forM_ steps $ \case
-    Let sline (TypedVarName name) expr inner -> do
-        cur <- asks (lookup name . tsVars . snd)
-        when (isJust cur) $ do
-            outLine OutputError Nothing $ T.pack "variable '" `T.append` textVarName name `T.append` T.pack "' already exists on " `T.append` textSourceLine sline
-            throwError Failed
-        value <- eval expr
-        withVar name value $ evalBlock =<< eval inner
-
-    For sline (TypedVarName name) expr inner -> do
-        cur <- asks (lookup name . tsVars . snd)
-        when (isJust cur) $ do
-            outLine OutputError Nothing $ T.pack "variable '" `T.append` textVarName name `T.append` T.pack "' already exists on " `T.append` textSourceLine sline
-            throwError Failed
-        value <- eval expr
-        forM_ value $ \i -> do
-            withVar name i $ evalBlock =<< eval inner
-
-    Subnet name@(TypedVarName vname) parent inner -> do
-        withSubnet parent (Just name) $ \net -> do
-            withVar vname net $ evalBlock =<< eval inner
-
-    DeclNode name@(TypedVarName vname) net inner -> do
-        withNode net (Left name) $ \node -> do
-            withVar vname node $ evalBlock =<< eval inner
-
-    Spawn tvname@(TypedVarName vname@(VarName tname)) target inner -> do
+    Subnet name parent inner -> do
+        withSubnet parent (Just name) $ evalBlock . inner
+
+    DeclNode name net inner -> do
+        withNode net (Left name) $ evalBlock . inner
+
+    Spawn tvname@(TypedVarName (VarName tname)) target inner -> do
         case target of
             Left net -> withNode net (Right tvname) go
             Right node -> go node
@@ -144,15 +125,14 @@ evalBlock (TestBlock steps) = forM_ steps $ \case
             opts <- asks $ teOptions . fst
             let pname = ProcName tname
                 tool = fromMaybe (optDefaultTool opts) (lookup pname $ optProcTools opts)
-            withProcess (Right node) pname Nothing tool $ \p -> do
-                withVar vname p $ evalBlock =<< eval inner
+            withProcess (Right node) pname Nothing tool $ evalBlock . inner
 
     Send p line -> do
         outProc OutputChildStdin p line
         send p line
 
     Expect line p expr captures inner -> do
-        expect line p expr captures $ evalBlock =<< eval inner
+        expect line p expr captures $ evalBlock . inner
 
     Flush p regex -> do
         flush p regex
@@ -273,7 +253,7 @@ exprFailed desc sline pname exprVars = do
             ]
     throwError Failed
 
-expect :: SourceLine -> Process -> Traced Regex -> [TypedVarName Text] -> TestRun () -> TestRun ()
+expect :: SourceLine -> Process -> Traced Regex -> [TypedVarName Text] -> ([ Text ] -> TestRun ()) -> TestRun ()
 expect sline p (Traced trace re) tvars inner = do
     timeout <- asks $ optTimeout . teOptions . fst
     delay <- liftIO $ registerDelay $ ceiling $ 1000000 * timeout
@@ -299,7 +279,7 @@ expect sline p (Traced trace re) tvars inner = do
                      throwError Failed
 
              outProc OutputMatch p line
-             local (fmap $ \s -> s { tsVars = zip vars (map someConstValue capture) ++ tsVars s }) inner
+             inner capture
 
          Nothing -> exprFailed (T.pack "expect") sline (Just $ procName p) trace
 
diff --git a/src/Run/Monad.hs b/src/Run/Monad.hs
index 2882197..f605dfb 100644
--- a/src/Run/Monad.hs
+++ b/src/Run/Monad.hs
@@ -91,7 +91,7 @@ instance MonadError Failed TestRun where
     catchError (TestRun act) handler = TestRun $ catchError act $ fromTestRun . handler
 
 instance MonadEval TestRun where
-    lookupVar name = maybe (fail $ "variable not in scope: '" ++ unpackVarName name ++ "'") return =<< asks (lookup name . tsVars . snd)
+    askDictionary = asks (tsVars . snd)
     withVar name value = local (fmap $ \s -> s { tsVars = ( name, someConstValue value ) : tsVars s })
 
 instance MonadOutput TestRun where
diff --git a/src/Test.hs b/src/Test.hs
index 53e0f03..c69d5e1 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -31,6 +31,7 @@ module Test (
 ) where
 
 import Control.Monad
+import Control.Monad.Reader
 
 import Data.Char
 import Data.List
@@ -41,6 +42,7 @@ import Data.String
 import Data.Text (Text)
 import Data.Text qualified as T
 import Data.Typeable
+import Data.Void
 
 import Text.Regex.TDFA qualified as RE
 import Text.Regex.TDFA.Text qualified as RE
@@ -64,13 +66,11 @@ 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)
+    = Subnet (TypedVarName Network) Network (Network -> TestBlock)
+    | DeclNode (TypedVarName Node) Network (Node -> TestBlock)
+    | Spawn (TypedVarName Process) (Either Network Node) (Process -> TestBlock)
     | Send Process Text
-    | Expect SourceLine Process (Traced Regex) [ TypedVarName Text ] (Expr TestBlock)
+    | Expect SourceLine Process (Traced Regex) [ TypedVarName Text ] ([ Text ] -> TestBlock)
     | Flush Process (Maybe Regex)
     | Guard SourceLine EvalTrace Bool
     | DisconnectNode Node TestBlock
@@ -89,9 +89,12 @@ textSourceLine SourceLineBuiltin = "<builtin>"
 
 
 class MonadFail m => MonadEval m where
+    askDictionary :: m VariableDictionary
     lookupVar :: VarName -> m SomeVarValue
+    lookupVar name = maybe (fail $ "variable not in scope: '" ++ unpackVarName name ++ "'") return . lookup name =<< askDictionary
     withVar :: ExprType e => VarName -> e -> m a -> m a
 
+type VariableDictionary = [ ( VarName, SomeVarValue ) ]
 
 newtype VarName = VarName Text
     deriving (Eq, Ord, Show)
@@ -150,6 +153,10 @@ instance ExprType Regex where
     textExprType _ = T.pack "regex"
     textExprValue _ = T.pack "<regex>"
 
+instance ExprType Void where
+    textExprType _ = T.pack "void"
+    textExprValue _ = T.pack "<void>"
+
 instance ExprType a => ExprType [a] where
     textExprType _ = "[" <> textExprType @a Proxy <> "]"
     textExprValue x = "[" <> T.intercalate ", " (map textExprValue x) <> "]"
@@ -251,11 +258,13 @@ data ExprEnumerator a = ExprEnumerator (a -> a -> [a]) (a -> a -> a -> [a])
 
 
 data Expr a where
+    Let :: forall a b. ExprType b => SourceLine -> TypedVarName b -> Expr b -> Expr a -> Expr a
     Variable :: ExprType a => SourceLine -> VarName -> Expr a
     DynVariable :: TypeVar -> SourceLine -> VarName -> Expr DynamicType
     FunVariable :: ExprType a => FunctionArguments SomeArgumentType -> SourceLine -> VarName -> Expr (FunctionType a)
     ArgsApp :: FunctionArguments SomeExpr -> Expr (FunctionType a) -> Expr (FunctionType a)
     FunctionEval :: Expr (FunctionType a) -> Expr a
+    LambdaAbstraction :: ExprType a => TypedVarName a -> Expr b -> Expr (a -> b)
     Pure :: a -> Expr a
     App :: AppAnnotation b -> Expr (a -> b) -> Expr a -> Expr b
     Concat :: [Expr Text] -> Expr Text
@@ -282,8 +291,26 @@ instance Monoid a => Monoid (Expr a) where
 varExpr :: ExprType a => SourceLine -> TypedVarName a -> Expr a
 varExpr sline (TypedVarName name) = Variable sline name
 
+
+newtype SimpleEval a = SimpleEval (Reader VariableDictionary a)
+    deriving (Functor, Applicative, Monad)
+
+runSimpleEval :: SimpleEval a -> VariableDictionary -> a
+runSimpleEval (SimpleEval x) = runReader x
+
+instance MonadFail SimpleEval where
+    fail = error . ("eval failed: " <>)
+
+instance MonadEval SimpleEval where
+    askDictionary = SimpleEval ask
+    withVar name value (SimpleEval inner) = SimpleEval $ local (( name, someConstValue value ) : ) inner
+
+
 eval :: forall m a. MonadEval m => Expr a -> m a
 eval = \case
+    Let _ (TypedVarName name) valExpr expr -> do
+        val <- eval valExpr
+        withVar name val $ eval expr
     Variable sline name -> fromSomeVarValue sline name =<< lookupVar name
     DynVariable _ _ _ -> fail "ambiguous type"
     FunVariable _ sline name -> funFromSomeVarValue sline name =<< lookupVar name
@@ -294,6 +321,9 @@ eval = \case
     FunctionEval efun -> do
         FunctionType fun <- eval efun
         return $ fun mempty
+    LambdaAbstraction (TypedVarName name) expr -> do
+        dict <- askDictionary
+        return $ \x -> runSimpleEval (eval expr) (( name, someConstValue x ) : dict)
     Pure value -> return value
     App _ f x -> eval f <*> eval x
     Concat xs -> T.concat <$> mapM eval xs
@@ -321,6 +351,7 @@ gatherVars = fmap (uniqOn fst . sortOn fst) . helper
   where
     helper :: forall b. Expr b -> m EvalTrace
     helper = \case
+        Let _ (TypedVarName var) _ expr -> filter ((var /=) . fst . fst) <$> helper expr
         Variable _ var
             | isInternalVar var -> return []
             | otherwise -> (: []) . (( var, [] ), ) <$> lookupVar var
@@ -331,6 +362,7 @@ gatherVars = fmap (uniqOn fst . sortOn fst) . helper
             vs <- mapM (\(SomeExpr e) -> helper e) $ M.elems args
             return $ concat (v : vs)
         FunctionEval efun -> helper efun
+        LambdaAbstraction (TypedVarName var) expr -> filter ((var /=) . fst . fst) <$> helper expr
         Pure _ -> return []
         e@(App (AnnRecord sel) _ x)
             | Just (var, sels) <- gatherSelectors x