Let statement

3 files changed, 25 insertions, 3 deletions

diff --git a/src/Main.hs b/src/Main.hs
index dc5ffd4..7c6d587 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -288,7 +288,7 @@ expect (SourceLine sline) p expr vars = do
              forM_ vars $ \name -> do
                  cur <- gets (lookup name . tsVars)
                  when (isJust cur) $ do
-                     outLine OutputMatchFail (Just $ procName p) $ T.pack "variable '" `T.append` textVarName name `T.append` T.pack "' already exists on " `T.append` sline
+                     outLine OutputError (Just $ procName p) $ T.pack "variable '" `T.append` textVarName name `T.append` T.pack "' already exists on " `T.append` sline
                      throwError ()
 
              modify $ \s -> s { tsVars = zip vars capture ++ tsVars s }
@@ -336,6 +336,14 @@ runTest out opts test = do
         oldHandler <- liftIO $ installHandler processStatusChanged (CatchInfo sigHandler) Nothing
 
         flip catchError (const $ return ()) $ forM_ (testSteps test) $ \case
+            Let (SourceLine sline) name expr -> do
+                 cur <- gets (lookup name . tsVars)
+                 when (isJust cur) $ do
+                     outLine OutputError Nothing $ T.pack "variable '" `T.append` textVarName name `T.append` T.pack "' already exists on " `T.append` sline
+                     throwError ()
+                 value <- eval expr
+                 modify $ \s -> s { tsVars = (name, value) : tsVars s }
+
             Spawn pname nname -> do
                 node <- getNode net nname
                 void $ spawnOn (Right node) pname Nothing $
diff --git a/src/Parser.hs b/src/Parser.hs
index fa85f8c..bce5a02 100644
--- a/src/Parser.hs
+++ b/src/Parser.hs
@@ -184,6 +184,18 @@ getSourceLine = do
         , TL.toStrict $ TL.takeWhile (/='\n') $ pstateInput pstate
         ]
 
+
+letStatement :: TestParser [TestStep]
+letStatement = do
+    line <- getSourceLine
+    wsymbol "let"
+    name <- VarName . (:[]) <$> identifier
+    sc
+    symbol "="
+    value <- stringExpr
+    return [Let line name value]
+
+
 command :: (Generic b, GInit (Rep b)) => String -> [Param b] -> (SourceLine -> b -> TestParser a) -> TestParser [a]
 command name params fin = do
     origline <- getSourceLine
@@ -292,7 +304,8 @@ testWait = do
 parseTestDefinition :: TestParser Test
 parseTestDefinition = label "test definition" $ toplevel $ do
     block (\name steps -> return $ Test name $ concat steps) header $ choice
-        [ testSpawn
+        [ letStatement
+        , testSpawn
         , testSend
         , testExpect
         , testGuard
diff --git a/src/Test.hs b/src/Test.hs
index e7e1255..80ee966 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -30,7 +30,8 @@ data Test = Test
     , testSteps :: [TestStep]
     }
 
-data TestStep = Spawn ProcName NodeName
+data TestStep = Let SourceLine VarName (Expr Text)
+              | Spawn ProcName NodeName
               | Send ProcName (Expr Text)
               | Expect SourceLine ProcName (Expr Regex) [VarName]
               | Guard SourceLine (Expr Bool)