Scope of variables denoted in parse tree

3 files changed, 175 insertions, 114 deletions

diff --git a/src/Main.hs b/src/Main.hs
index 29c45bc..6306c17 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -6,7 +6,6 @@ import Control.Concurrent.STM
 import Control.Monad
 import Control.Monad.Except
 import Control.Monad.Reader
-import Control.Monad.State
 
 import Data.List
 import Data.Maybe
@@ -76,11 +75,12 @@ data TestEnv = TestEnv
     }
 
 data TestState = TestState
-    { tsVars :: [(VarName, SomeVarValue)]
+    { tsNetwork :: Network
+    , tsVars :: [(VarName, SomeVarValue)]
     }
 
-newtype TestRun a = TestRun { fromTestRun :: ReaderT TestEnv (StateT TestState (ExceptT () IO)) a }
-    deriving (Functor, Applicative, Monad, MonadReader TestEnv, MonadState TestState, MonadIO)
+newtype TestRun a = TestRun { fromTestRun :: ReaderT (TestEnv, TestState) (ExceptT () IO) a }
+    deriving (Functor, Applicative, Monad, MonadReader (TestEnv, TestState), MonadIO)
 
 instance MonadFail TestRun where
     fail str = do
@@ -89,30 +89,29 @@ instance MonadFail TestRun where
 
 instance MonadError () TestRun where
     throwError () = do
-        failedVar <- asks teFailed
+        failedVar <- asks $ teFailed . fst
         liftIO $ atomically $ writeTVar failedVar True
         TestRun $ throwError ()
 
     catchError (TestRun act) handler = TestRun $ catchError act $ fromTestRun . handler
 
 instance MonadEval TestRun where
-    lookupVar name = maybe (fail $ "variable not in scope: '" ++ unpackVarName name ++ "'") return =<< gets (lookup name . tsVars)
+    lookupVar name = maybe (fail $ "variable not in scope: '" ++ unpackVarName name ++ "'") return =<< asks (lookup name . tsVars . snd)
 
 instance MonadOutput TestRun where
-    getOutput = asks teOutput
+    getOutput = asks $ teOutput . fst
 
 forkTest :: TestRun () -> TestRun ()
 forkTest act = do
     tenv <- ask
-    tstate <- get
     void $ liftIO $ forkIO $ do
-        runExceptT (flip evalStateT tstate $ flip runReaderT tenv $ fromTestRun act) >>= \case
-            Left () -> atomically $ writeTVar (teFailed tenv) True
+        runExceptT (flip runReaderT tenv $ fromTestRun act) >>= \case
+            Left () -> atomically $ writeTVar (teFailed $ fst tenv) True
             Right () -> return ()
 
 atomicallyTest :: STM a -> TestRun a
 atomicallyTest act = do
-    failedVar <- asks teFailed
+    failedVar <- asks $ teFailed . fst
     res <- liftIO $ atomically $ do
         failed <- readTVar failedVar
         if failed then return $ Left ()
@@ -121,8 +120,8 @@ atomicallyTest act = do
         Left e -> throwError e
         Right x -> return x
 
-initNetwork :: TestRun Network
-initNetwork = do
+initNetwork :: (Network -> TestRun a) -> TestRun a
+initNetwork inner = do
     net <- liftIO $ do
         exists <- doesPathExist testDir
         when exists $ ioError $ userError $ testDir ++ " exists"
@@ -137,11 +136,11 @@ initNetwork = do
     void $ spawnOn (Left net) (ProcNameTcpdump) (Just softwareTermination) $
         "tcpdump -i br0 -w '" ++ testDir ++ "/br0.pcap' -U -Z root"
 
-    useGDB <- asks $ optGDB . teOptions
+    useGDB <- asks $ optGDB . teOptions . fst
     when useGDB $ do
         gdbInit =<< spawnOn (Left net) ProcNameGDB Nothing gdbCmd
 
-    return net
+    local (fmap $ \s -> s { tsNetwork = net }) $ inner net
 
 exitNetwork :: Network -> TestRun ()
 exitNetwork net = do
@@ -167,44 +166,51 @@ exitNetwork net = do
             ExitSuccess -> return ()
             ExitFailure code -> do
                 outLine OutputChildFail (Just $ procName p) $ T.pack $ "exit code: " ++ show code
-                liftIO . atomically . flip writeTVar False =<< asks teFailed
+                liftIO . atomically . flip writeTVar False =<< asks (teFailed . fst)
 
     liftIO $ do
         callCommand $ "ip -all netns del"
         callCommand $ "ip link del group 1"
 
-    failed <- liftIO . atomically . readTVar =<< asks teFailed
+    failed <- liftIO . atomically . readTVar =<< asks (teFailed . fst)
     liftIO $ if failed then exitFailure
                        else removeDirectoryRecursive $ netDir net
 
-getNode :: Network -> NodeName -> TestRun Node
-getNode net nname@(NodeName tnname) = (find ((nname==).nodeName) <$> liftIO (readMVar (netNodes net))) >>= \case
-    Just node -> return node
-    _ -> do
-        let name = T.unpack tnname
-            dir = netDir net </> ("erebos_" ++ name)
-            node = Node { nodeName = nname
-                        , nodeNetwork = net
-                        , nodeDir = dir
-                        }
-
-        ip <- liftIO $ do
-            exists <- doesPathExist dir
-            when exists $ ioError $ userError $ dir ++ " exists"
-            createDirectoryIfMissing True dir
-
-            modifyMVar (netNodes net) $ \nodes -> do
-                let ip = "192.168.0." ++ show (11 + length nodes)
-                callCommand $ "ip netns add \""++ name ++ "\""
-                callCommand $ "ip link add \"veth_" ++ name ++ ".0\" group 1 type veth peer name \"veth_" ++ name ++ ".1\" netns \"" ++ name ++ "\""
-                callCommand $ "ip link set dev \"veth_" ++ name ++ ".0\" master br0 up"
-                callOn node $ "ip addr add " ++ ip ++ "/24 broadcast 192.168.0.255 dev \"veth_" ++ name ++ ".1\""
-                callOn node $ "ip link set dev \"veth_" ++ name++ ".1\" up"
-                callOn node $ "ip link set dev lo up"
-                return $ (node : nodes, ip)
-
-        modify $ \s -> s { tsVars = (VarName [tnname, T.pack "ip"], SomeVarValue (T.pack ip)) : tsVars s }
-        return node
+getNode :: NodeName -> (Node -> TestRun a) -> TestRun a
+getNode nname inner = do
+    net <- asks $ tsNetwork . snd
+    nodes <- liftIO (readMVar (netNodes net))
+    case find ((nname==).nodeName) nodes of
+        Just node -> inner node
+        _ -> createNode nname inner
+
+createNode :: NodeName -> (Node -> TestRun a) -> TestRun a
+createNode nname@(NodeName tnname) inner = do
+    net <- asks $ tsNetwork . snd
+    let name = T.unpack tnname
+        dir = netDir net </> ("erebos_" ++ name)
+        node = Node { nodeName = nname
+                    , nodeNetwork = net
+                    , nodeDir = dir
+                    }
+
+    ip <- liftIO $ do
+        exists <- doesPathExist dir
+        when exists $ ioError $ userError $ dir ++ " exists"
+        createDirectoryIfMissing True dir
+
+        modifyMVar (netNodes net) $ \nodes -> do
+            let ip = "192.168.0." ++ show (11 + length nodes)
+            callCommand $ "ip netns add \""++ name ++ "\""
+            callCommand $ "ip link add \"veth_" ++ name ++ ".0\" group 1 type veth peer name \"veth_" ++ name ++ ".1\" netns \"" ++ name ++ "\""
+            callCommand $ "ip link set dev \"veth_" ++ name ++ ".0\" master br0 up"
+            callOn node $ "ip addr add " ++ ip ++ "/24 broadcast 192.168.0.255 dev \"veth_" ++ name ++ ".1\""
+            callOn node $ "ip link set dev \"veth_" ++ name++ ".1\" up"
+            callOn node $ "ip link set dev lo up"
+            return $ (node : nodes, ip)
+
+    local (fmap $ \s -> s { tsVars = (VarName [tnname, T.pack "ip"], SomeVarValue (T.pack ip)) : tsVars s }) $ do
+        inner node
 
 callOn :: Node -> String -> IO ()
 callOn node cmd = callCommand $ "ip netns exec \"" ++ unpackNodeName (nodeName node) ++ "\" " ++ cmd
@@ -257,9 +263,10 @@ spawnOn target pname killWith cmd = do
     liftIO $ modifyMVar_ (netProcesses net) $ return . (process:)
     return process
 
-getProcess :: Network -> ProcName -> TestRun Process
-getProcess net pname = liftIO $ do
-    Just p <- find ((pname==).procName) <$> readMVar (netProcesses net)
+getProcess :: ProcName -> TestRun Process
+getProcess pname = do
+    net <- asks $ tsNetwork . snd
+    Just p <- find ((pname==).procName) <$> liftIO (readMVar (netProcesses net))
     return p
 
 tryMatch :: Regex -> [Text] -> Maybe ((Text, [Text]), [Text])
@@ -275,10 +282,10 @@ exprFailed desc (SourceLine sline) pname expr = do
         outLine OutputMatchFail pname $ T.concat [T.pack "  ", textVarName name, T.pack " = ", textSomeVarValue value]
     throwError ()
 
-expect :: SourceLine -> Process -> Expr Regex -> [VarName] -> TestRun ()
-expect (SourceLine sline) p expr vars = do
+expect :: SourceLine -> Process -> Expr Regex -> [VarName] -> TestRun () -> TestRun ()
+expect (SourceLine sline) p expr vars inner = do
     re <- eval expr
-    timeout <- asks $ optTimeout . teOptions
+    timeout <- asks $ optTimeout . teOptions . fst
     delay <- liftIO $ registerDelay $ ceiling $ 1000000 * timeout
     mbmatch <- atomicallyTest $ (Nothing <$ (check =<< readTVar delay)) <|> do
         line <- readTVar (procOutput p)
@@ -294,13 +301,14 @@ expect (SourceLine sline) p expr vars = do
                  throwError ()
 
              forM_ vars $ \name -> do
-                 cur <- gets (lookup name . tsVars)
+                 cur <- asks (lookup name . tsVars . snd)
                  when (isJust cur) $ do
                      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 (map SomeVarValue capture) ++ tsVars s }
              outLine OutputMatch (Just $ procName p) line
+             local (fmap $ \s -> s { tsVars = zip vars (map SomeVarValue capture) ++ tsVars s }) inner
+
          Nothing -> exprFailed (T.pack "expect") (SourceLine sline) (Just $ procName p) expr
 
 testStepGuard :: SourceLine -> Expr Bool -> TestRun ()
@@ -312,6 +320,41 @@ allM :: Monad m => [a] -> (a -> m Bool) -> m Bool
 allM (x:xs) p = p x >>= \case True -> allM xs p; False -> return False
 allM [] _ = return True
 
+evalSteps :: [TestStep] -> TestRun ()
+evalSteps = mapM_ $ \case
+    Let (SourceLine sline) 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` sline
+            throwError ()
+        value <- eval expr
+        local (fmap $ \s -> s { tsVars = (name, SomeVarValue value) : tsVars s }) $ do
+            evalSteps inner
+
+    Spawn pname nname inner -> do
+        getNode nname $ \node -> do
+            opts <- asks $ teOptions . fst
+            void $ spawnOn (Right node) pname Nothing $
+                fromMaybe (optDefaultTool opts) (lookup pname $ optProcTools opts)
+            evalSteps inner
+
+    Send pname expr -> do
+        p <- getProcess pname
+        line <- eval expr
+        send p line
+
+    Expect line pname expr captures inner -> do
+        p <- getProcess pname
+        expect line p expr captures $ evalSteps inner
+
+    Guard line expr -> do
+        testStepGuard line expr
+
+    Wait -> do
+        outPrompt $ T.pack "Waiting..."
+        void $ liftIO $ getLine
+        outClearPrompt
+
 runTest :: Output -> Options -> Test -> IO Bool
 runTest out opts test = do
     tenv <- TestEnv
@@ -319,10 +362,9 @@ runTest out opts test = do
         <*> newTVarIO False
         <*> pure opts
     tstate <- TestState
-        <$> pure []
-    (fmap $ either (const False) id) $ runExceptT $ flip evalStateT tstate $ flip runReaderT tenv $ fromTestRun $ do
-        net <- initNetwork
-
+        <$> pure (error "network not initialized")
+        <*> pure []
+    (fmap $ either (const False) id) $ runExceptT $ flip runReaderT (tenv, tstate) $ fromTestRun $ initNetwork $ \net -> do
         let sigHandler SignalInfo { siginfoSpecific = chld } = do
                 processes <- readMVar (netProcesses net)
                 forM_ processes $ \p -> do
@@ -336,36 +378,7 @@ runTest out opts test = do
                              Stopped sig -> err $ T.pack $ "child stopped with signal " ++ show sig
         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, SomeVarValue value) : tsVars s }
-
-            Spawn pname nname -> do
-                node <- getNode net nname
-                void $ spawnOn (Right node) pname Nothing $
-                    fromMaybe (optDefaultTool opts) (lookup pname $ optProcTools opts)
-
-            Send pname expr -> do
-                p <- getProcess net pname
-                line <- eval expr
-                send p line
-
-            Expect line pname expr captures -> do
-                p <- getProcess net pname
-                expect line p expr captures
-
-            Guard line expr -> do
-                testStepGuard line expr
-
-            Wait -> do
-                outPrompt $ T.pack "Waiting..."
-                void $ liftIO $ getLine
-                outClearPrompt
+        flip catchError (const $ return ()) $ evalSteps $ testSteps test
 
         _ <- liftIO $ installHandler processStatusChanged oldHandler Nothing
         exitNetwork net
diff --git a/src/Parser.hs b/src/Parser.hs
index 2ab64ef..5b8f003 100644
--- a/src/Parser.hs
+++ b/src/Parser.hs
@@ -75,7 +75,7 @@ block :: (a -> [b] -> TestParser c) -> TestParser a -> TestParser b -> TestParse
 block merge header item = L.indentBlock scn $ do
     h <- header
     choice
-        [ do try $ void $ lexeme $ char ':'
+        [ do symbol ":"
              return $ L.IndentSome Nothing (merge h) item
         , L.IndentNone <$> merge h []
         ]
@@ -285,15 +285,19 @@ getSourceLine = do
 letStatement :: TestParser [TestStep]
 letStatement = do
     line <- getSourceLine
+    indent <- L.indentLevel
     wsymbol "let"
     name <- VarName . (:[]) <$> identifier
     symbol "="
     SomeExpr (e :: Expr a) <- someExpr
     void $ eol
 
+    s <- get
     addVarName @a Proxy name
-    return [Let line name e]
+    body <- testBlock indent
+    put s
 
+    return [Let line name e body]
 
 class Typeable a => ParamType a where
     parseParam :: TestParser a
@@ -347,51 +351,79 @@ param name = CommandDef [(name, SomeParam (Proxy @a))] (\[SomeParam (Identity x)
 cmdLine :: CommandDef SourceLine
 cmdLine = param ""
 
-command :: String -> CommandDef a -> TestParser [a]
+data InnerBlock
+
+instance ParamType InnerBlock where
+    parseParam = mzero
+    showParamType _ = "<code block>"
+
+instance ParamType TestStep where
+    parseParam = mzero
+    showParamType _ = "<code line>"
+
+innerBlock :: CommandDef [TestStep]
+innerBlock = CommandDef [("", SomeParam (Proxy @InnerBlock))] (\[SomeParam (Identity x)] -> fromJust $ cast x)
+
+command :: String -> CommandDef TestStep -> TestParser [TestStep]
 command name (CommandDef types ctor) = do
+    indent <- L.indentLevel
     line <- getSourceLine
-    L.indentBlock scn $ do
-        wsymbol name
-        helper line $ map (fmap $ \(SomeParam (_ :: Proxy p)) -> SomeParam $ Nothing @p) types
+    wsymbol name
+    restOfLine indent [] line $ map (fmap $ \(SomeParam (_ :: Proxy p)) -> SomeParam $ Nothing @p) types
   where
-    helper line params = choice
+    restOfLine :: Pos -> [(Pos, [(String, SomeParam Maybe)])] -> SourceLine -> [(String, SomeParam Maybe)] -> TestParser [TestStep]
+    restOfLine cmdi partials line params = choice
         [do void $ lookAhead eol
             iparams <- forM params $ \case
                 (_, SomeParam (Nothing :: Maybe p))
                     | Just (Refl :: p :~: SourceLine) <- eqT -> return $ SomeParam $ Identity line
+                    | Just (Refl :: p :~: InnerBlock) <- eqT -> SomeParam . Identity <$> restOfParts cmdi partials
                 (sym, SomeParam (Nothing :: Maybe p)) -> choice
                     [ SomeParam . Identity <$> paramDefault @p
                     , fail $ "missing " ++ (if null sym then "" else "'" ++ sym ++ "' ") ++ showParamType @p Proxy
                     ]
                 (_, SomeParam (Just x)) -> return $ SomeParam $ Identity x
-            return $ L.IndentNone [ctor iparams]
+            return [ctor iparams]
 
         ,do symbol ":"
-            return $ L.IndentSome Nothing (return . concat) $ do
-                line' <- getSourceLine
-                L.indentBlock scn $ helper line' params
+            scn
+            indent <- L.indentLevel
+            restOfParts cmdi ((indent, params) : partials)
 
-        ,do tryParams line [] params
+        ,do tryParams cmdi partials line [] params
         ]
 
+    restOfParts :: Pos -> [(Pos, [(String, SomeParam Maybe)])] -> TestParser [TestStep]
+    restOfParts cmdi [] = testBlock cmdi
+    restOfParts cmdi partials@((partIndent, params) : rest) = do
+        scn
+        pos <- L.indentLevel
+        line <- getSourceLine
+        optional eof >>= \case
+            Just _ -> return []
+            _ | pos <  partIndent -> restOfParts cmdi rest
+              | pos == partIndent -> (++) <$> restOfLine cmdi partials line params <*> restOfParts cmdi partials
+              | otherwise         -> L.incorrectIndent EQ partIndent pos
+
     tryParam sym (SomeParam (cur :: Maybe p)) = do
         when (not $ null sym) $ wsymbol sym
         when (isJust cur) $ do
             fail $ "multiple " ++ (if null sym then "unnamed" else "'" ++ sym ++ "'") ++ " parameters"
         SomeParam . Just <$> parseParam @p
 
-    tryParams line prev ((sym, p) : ps) = choice $
+    tryParams cmdi partIndent line prev ((sym, p) : ps) = choice $
         (if null sym then reverse else id) {- try unnamed parameter as last option -} $
         [do p' <- tryParam sym p
-            helper line $ concat [reverse prev, [(sym, p')], ps]
-        ,do tryParams line ((sym, p) : prev) ps
+            restOfLine cmdi partIndent line $ concat [reverse prev, [(sym, p')], ps]
+        ,do tryParams cmdi partIndent line ((sym, p) : prev) ps
         ]
-    tryParams _ _ [] = mzero
+    tryParams _ _ _ _ [] = mzero
 
 testSpawn :: TestParser [TestStep]
 testSpawn = command "spawn" $ Spawn
     <$> param "as"
     <*> param "on"
+    <*> innerBlock
 
 testSend :: TestParser [TestStep]
 testSend = command "send" $ Send
@@ -404,6 +436,7 @@ testExpect = command "expect" $ Expect
     <*> param "from"
     <*> param ""
     <*> param "capture"
+    <*> innerBlock
 
 testGuard :: TestParser [TestStep]
 testGuard = command "guard" $ Guard
@@ -416,16 +449,31 @@ testWait = do
     wsymbol "wait"
     return [Wait]
 
+testBlock :: Pos -> TestParser [TestStep]
+testBlock indent = concat <$> go
+  where
+    go = do
+        scn
+        pos <- L.indentLevel
+        optional eof >>= \case
+            Just _ -> return []
+            _ | pos <  indent -> return []
+              | pos == indent -> (:) <$> testStep <*> go
+              | otherwise     -> L.incorrectIndent EQ indent pos
+
+testStep :: TestParser [TestStep]
+testStep = choice
+    [ letStatement
+    , testSpawn
+    , testSend
+    , testExpect
+    , testGuard
+    , testWait
+    ]
+
 parseTestDefinition :: TestParser Test
 parseTestDefinition = label "test definition" $ toplevel $ do
-    block (\name steps -> return $ Test name $ concat steps) header $ choice
-        [ letStatement
-        , testSpawn
-        , testSend
-        , testExpect
-        , testGuard
-        , testWait
-        ]
+    block (\name steps -> return $ Test name $ concat steps) header testStep
     where header = do
               wsymbol "test"
               lexeme $ TL.toStrict <$> takeWhileP (Just "test name") (/=':')
diff --git a/src/Test.hs b/src/Test.hs
index a2419a7..8db7552 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -33,10 +33,10 @@ data Test = Test
     , testSteps :: [TestStep]
     }
 
-data TestStep = forall a. ExprType a => Let SourceLine VarName (Expr a)
-              | Spawn ProcName NodeName
+data TestStep = forall a. ExprType a => Let SourceLine VarName (Expr a) [TestStep]
+              | Spawn ProcName NodeName [TestStep]
               | Send ProcName (Expr Text)
-              | Expect SourceLine ProcName (Expr Regex) [VarName]
+              | Expect SourceLine ProcName (Expr Regex) [VarName] [TestStep]
               | Guard SourceLine (Expr Bool)
               | Wait