Support multiple node variables with same name

4 files changed, 56 insertions, 52 deletions

diff --git a/src/Main.hs b/src/Main.hs
index b6c952f..221bfb4 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -152,42 +152,36 @@ exitNetwork net = do
     liftIO $ if failed then exitFailure
                        else removeDirectoryRecursive $ netDir net
 
-getNode :: NodeName -> (Node -> TestRun a) -> TestRun a
-getNode nname inner = do
+createNode :: TypedVarName Node -> (Node -> TestRun a) -> TestRun a
+createNode (TypedVarName vname) 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 <- 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)
+        node <- modifyMVar (netNodes net) $ \nodes -> do
+            let nname = nextNodeName vname $ map nodeName nodes
+                ip = "192.168.0." ++ show (11 + length nodes)
                 node = Node { nodeName = nname
                             , nodeIp = T.pack ip
                             , nodeNetwork = net
-                            , nodeDir = dir
+                            , nodeDir = netDir net </> ("erebos_" ++ unpackNodeName nname)
                             }
-
-            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, node)
 
-    local (fmap $ \s -> s { tsVars = (VarName tnname, SomeVarValue node) : tsVars s }) $ do
+        let name = unpackNodeName $ nodeName node
+            dir = nodeDir node
+
+        exists <- doesPathExist dir
+        when exists $ ioError $ userError $ dir ++ " exists"
+        createDirectoryIfMissing True dir
+
+        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 " ++ T.unpack (nodeIp node) ++ "/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
+
+    local (fmap $ \s -> s { tsVars = (vname, SomeVarValue node) : tsVars s }) $ do
         inner node
 
 callOn :: Node -> String -> IO ()
@@ -261,8 +255,8 @@ exprFailed desc (SourceLine sline) pname expr = do
         outLine OutputMatchFail prompt $ T.concat [T.pack "  ", textVarName name, T.pack " = ", textSomeVarValue value]
     throwError ()
 
-expect :: SourceLine -> Process -> Expr Regex -> [VarName] -> TestRun () -> TestRun ()
-expect (SourceLine sline) p expr vars inner = do
+expect :: SourceLine -> Process -> Expr Regex -> [TypedVarName Text] -> TestRun () -> TestRun ()
+expect (SourceLine sline) p expr tvars inner = do
     re <- eval expr
     timeout <- asks $ optTimeout . teOptions . fst
     delay <- liftIO $ registerDelay $ ceiling $ 1000000 * timeout
@@ -275,6 +269,8 @@ expect (SourceLine sline) p expr vars inner = do
                  return $ Just m
     case mbmatch of
          Just (line, capture) -> do
+             let vars = map (\(TypedVarName n) -> n) tvars
+
              when (length vars /= length capture) $ do
                  outProc OutputMatchFail p $ T.pack "mismatched number of capture variables on " `T.append` sline
                  throwError ()
@@ -317,7 +313,7 @@ evalSteps = mapM_ $ \case
             evalSteps inner
 
     Spawn pname nname inner -> do
-        either getNode ((>>=) . eval) nname $ \node -> do
+        either createNode ((>>=) . eval) nname $ \node -> do
             opts <- asks $ teOptions . fst
             p <- spawnOn (Right node) pname Nothing $
                 fromMaybe (optDefaultTool opts) (lookup pname $ optProcTools opts)
diff --git a/src/Network.hs b/src/Network.hs
index 5872d13..a3c7120 100644
--- a/src/Network.hs
+++ b/src/Network.hs
@@ -2,6 +2,7 @@ module Network (
     Network(..),
     Node(..),
     NodeName(..), textNodeName, unpackNodeName,
+    nextNodeName,
 ) where
 
 import Control.Arrow
@@ -26,20 +27,28 @@ data Node = Node
     , nodeDir :: FilePath
     }
 
-newtype NodeName = NodeName Text
+data NodeName = NodeName Text Word
     deriving (Eq, Ord)
 
 textNodeName :: NodeName -> Text
-textNodeName (NodeName name) = name
+textNodeName (NodeName name 0) = name
+textNodeName (NodeName name num) = name <> T.pack "~" <> T.pack (show num)
 
 unpackNodeName :: NodeName -> String
-unpackNodeName (NodeName tname) = T.unpack tname
+unpackNodeName = T.unpack . textNodeName
+
+nextNodeName :: VarName -> [NodeName] -> NodeName
+nextNodeName (VarName tname) = go 0
+  where
+    go n [] = NodeName tname n
+    go n (NodeName tname' m : ns) | tname == tname' = go (max n m + 1) ns
+                                  | otherwise       = go n ns
 
 
 instance ExprType Node where
     textExprType _ = T.pack "node"
     textExprValue n = T.pack "n:" <> textNodeName (nodeName n)
-    emptyVarValue = Node (NodeName T.empty) T.empty undefined undefined
+    emptyVarValue = Node (NodeName T.empty 0) T.empty undefined undefined
 
     recordMembers = map (first T.pack)
         [ ("ip", RecordSelector $ nodeIp)
diff --git a/src/Parser.hs b/src/Parser.hs
index 74a5ade..a38d0c9 100644
--- a/src/Parser.hs
+++ b/src/Parser.hs
@@ -26,7 +26,7 @@ import qualified Text.Megaparsec.Char.Lexer as L
 
 import System.Exit
 
-import Network (Node, NodeName(..))
+import Network ()
 import Process (ProcName(..))
 import Test
 
@@ -107,15 +107,15 @@ identifier = do
 varName :: TestParser VarName
 varName = VarName <$> identifier
 
-newVarName :: forall a proxy. ExprType a => proxy a -> TestParser VarName
-newVarName proxy = do
+newVarName :: forall a. ExprType a => TestParser (TypedVarName a)
+newVarName = do
     off <- stateOffset <$> getParserState
-    name <- varName
-    addVarName off proxy name
+    name <- TypedVarName <$> varName
+    addVarName off name
     return name
 
-addVarName :: forall a proxy. ExprType a => Int -> proxy a -> VarName -> TestParser ()
-addVarName off _ name = do
+addVarName :: forall a. ExprType a => Int -> TypedVarName a -> TestParser ()
+addVarName off (TypedVarName name) = do
     gets (lookup name . testVars) >>= \case
         Just _ -> parseError $ FancyError off $ S.singleton $ ErrorFail $ T.unpack $
             T.pack "variable '" <> textVarName name <> T.pack "' already exists"
@@ -312,7 +312,7 @@ letStatement = do
     SomeExpr (e :: Expr a) <- someExpr
 
     localState $ do
-        addVarName @a off Proxy name
+        addVarName off $ TypedVarName @a name
         void $ eol
         body <- testBlock indent
         return [Let line name e body]
@@ -328,16 +328,12 @@ instance ParamType SourceLine where
     parseParam = mzero
     showParamType _ = "<source line>"
 
-instance ParamType NodeName where
-    parseParam = NodeName . textVarName <$> newVarName @Node Proxy
-    showParamType _ = "<node>"
-
 instance ParamType ProcName where
     parseParam = procName
     showParamType _ = "<proc>"
 
-instance ParamType VarName where
-    parseParam = newVarName @Text Proxy
+instance ExprType a => ParamType (TypedVarName a) where
+    parseParam = newVarName
     showParamType _ = "<variable>"
 
 instance ExprType a => ParamType (Expr a) where
diff --git a/src/Test.hs b/src/Test.hs
index a90035b..659107f 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -4,7 +4,7 @@ module Test (
     SourceLine(..),
 
     MonadEval(..),
-    VarName(..), textVarName, unpackVarName,
+    VarName(..), TypedVarName(..), textVarName, unpackVarName,
     ExprType(..),
     SomeVarValue(..), fromSomeVarValue, textSomeVarValue,
     RecordSelector(..),
@@ -33,9 +33,9 @@ data Test = Test
     }
 
 data TestStep = forall a. ExprType a => Let SourceLine VarName (Expr a) [TestStep]
-              | Spawn ProcName (Either NodeName (Expr Node)) [TestStep]
+              | Spawn ProcName (Either (TypedVarName Node) (Expr Node)) [TestStep]
               | Send ProcName (Expr Text)
-              | Expect SourceLine ProcName (Expr Regex) [VarName] [TestStep]
+              | Expect SourceLine ProcName (Expr Regex) [TypedVarName Text] [TestStep]
               | Guard SourceLine (Expr Bool)
               | Wait
 
@@ -49,6 +49,9 @@ class MonadFail m => MonadEval m where
 newtype VarName = VarName Text
     deriving (Eq, Ord)
 
+newtype TypedVarName a = TypedVarName VarName
+    deriving (Eq, Ord)
+
 textVarName :: VarName -> Text
 textVarName (VarName name ) = name