Network expression type and context

5 files changed, 44 insertions, 21 deletions

diff --git a/src/Main.hs b/src/Main.hs
index 7671aa3..09810a3 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -1,3 +1,5 @@
+{-# LANGUAGE OverloadedStrings #-}
+
 module Main (main) where
 
 import Control.Arrow
@@ -96,6 +98,7 @@ instance MonadError Failed TestRun where
 
 instance MonadEval TestRun where
     lookupVar name = maybe (fail $ "variable not in scope: '" ++ unpackVarName name ++ "'") return =<< asks (lookup name . tsVars . snd)
+    rootNetwork = asks $ tsNetwork . snd
 
 instance MonadOutput TestRun where
     getOutput = asks $ teOutput . fst
@@ -168,12 +171,13 @@ withNetwork inner = do
 
     return res
 
-createNode :: TypedVarName Node -> (Node -> TestRun a) -> TestRun a
-createNode (TypedVarName vname) inner = do
-    net <- asks $ tsNetwork . snd
+createNode :: Expr Network -> Maybe (TypedVarName Node) -> (Node -> TestRun a) -> TestRun a
+createNode netexpr tvname inner = do
+    let vname = fromTypedVarName <$> tvname
+    net <- eval netexpr
     node <- liftIO $ do
         node <- modifyMVar (netNodes net) $ \nodes -> do
-            let nname = nextNodeName vname $ map nodeName nodes
+            let nname = nextNodeName (fromMaybe (VarName "node") vname) $ map nodeName nodes
                 ip = "192.168.0." ++ show (11 + length nodes)
                 node = Node { nodeName = nname
                             , nodeIp = T.pack ip
@@ -197,7 +201,7 @@ createNode (TypedVarName vname) inner = do
         callOn node $ "ip link set dev lo up"
         return node
 
-    withVar vname node $ inner node
+    maybe id (flip withVar node) vname $ inner node
 
 callOn :: Node -> String -> IO ()
 callOn node cmd = callCommand $ "ip netns exec \"" ++ unpackNodeName (nodeName node) ++ "\" " ++ cmd
@@ -302,12 +306,17 @@ evalSteps = mapM_ $ \case
         value <- eval expr
         withVar name value $ evalSteps inner
 
-    DeclNode name@(TypedVarName vname) inner -> do
-        createNode name $ \node -> do
+    DeclNode name@(TypedVarName vname) net inner -> do
+        createNode net (Just name) $ \node -> do
             withVar vname node $ evalSteps inner
 
-    Spawn (TypedVarName vname@(VarName tname)) nname inner -> do
-        either createNode ((>>=) . eval) nname $ \node -> do
+    Spawn (TypedVarName vname@(VarName tname)) target inner -> do
+        case target of
+            Left nname -> createNode RootNetwork (Just nname) go
+            Right (Left net) -> createNode net Nothing go
+            Right (Right node) -> go =<< eval node
+      where
+        go node = do
             let pname = ProcName tname
             opts <- asks $ teOptions . fst
             p <- spawnOn (Right node) pname Nothing $
diff --git a/src/Network.hs b/src/Network.hs
index d1d00bc..8048c72 100644
--- a/src/Network.hs
+++ b/src/Network.hs
@@ -45,6 +45,11 @@ nextNodeName (VarName tname) = go 0
                                   | otherwise       = go n ns
 
 
+instance ExprType Network where
+    textExprType _ = T.pack "network"
+    textExprValue _ = T.pack "s:0"
+    emptyVarValue = Network undefined undefined undefined
+
 instance ExprType Node where
     textExprType _ = T.pack "node"
     textExprValue n = T.pack "n:" <> textNodeName (nodeName n)
diff --git a/src/Network.hs-boot b/src/Network.hs-boot
index 820fdaf..1b5e9c4 100644
--- a/src/Network.hs-boot
+++ b/src/Network.hs-boot
@@ -1,4 +1,5 @@
 module Network where
 
+data Network
 data Node
 data NodeName
diff --git a/src/Parser.hs b/src/Parser.hs
index 903ad54..f8c5b0e 100644
--- a/src/Parser.hs
+++ b/src/Parser.hs
@@ -29,7 +29,7 @@ import qualified Text.Megaparsec.Char.Lexer as L
 
 import System.Exit
 
-import Network (Node)
+import Network (Network, Node)
 import Process (Process, ProcName(..))
 import Test
 import Util
@@ -40,7 +40,7 @@ type TestStream = TL.Text
 
 data TestParserState = TestParserState
     { testVars :: [(VarName, SomeExprType)]
-    , testContext :: Maybe SomeExpr
+    , testContext :: SomeExpr
     }
 
 data SomeExprType = forall a. ExprType a => SomeExprType (Proxy a)
@@ -53,6 +53,7 @@ textSomeExprType (SomeExprType p) = textExprType p
 
 instance MonadEval TestParser where
     lookupVar name = maybe (fail $ "variable not in scope: '" ++ unpackVarName name ++ "'") (return . someEmptyVar) =<< gets (lookup name . testVars)
+    rootNetwork = return emptyVarValue
 
 skipLineComment :: TestParser ()
 skipLineComment = L.skipLineComment $ TL.pack "#"
@@ -386,13 +387,13 @@ instance ParamType a => ParamType [a] where
     parseParam _ = listOf (parseParam @a Proxy)
     showParamType _ = showParamType @a Proxy ++ " [, " ++ showParamType @a Proxy ++ " ...]"
     paramDefault _ = return []
-    paramFromSomeExpr _ (SomeExpr e) = cast e <|> ((:[]) <$> cast e)
+    paramFromSomeExpr _ se@(SomeExpr e) = cast e <|> ((:[]) <$> paramFromSomeExpr @a Proxy se)
 
 instance (ParamType a, ParamType b) => ParamType (Either a b) where
     type ParamRep (Either a b) = Either (ParamRep a) (ParamRep b)
     parseParam _ = try (Left <$> parseParam @a Proxy) <|> (Right <$> parseParam @b Proxy)
     showParamType _ = showParamType @a Proxy ++ " or " ++ showParamType @b Proxy
-    paramFromSomeExpr _ (SomeExpr e) = (Left <$> cast e) <|> (Right <$> cast e)
+    paramFromSomeExpr _ se = (Left <$> paramFromSomeExpr @a Proxy se) <|> (Right <$> paramFromSomeExpr @b Proxy se)
 
 data SomeParam f = forall a. ParamType a => SomeParam (Proxy a) (f (ParamRep a))
 
@@ -418,8 +419,9 @@ instance ParamType a => ParamType (ParamOrContext a) where
     parseParam _ = parseParam @a Proxy
     showParamType _ = showParamType @a Proxy
     paramDefault _ = gets testContext >>= \case
-        Just se | Just e <- paramFromSomeExpr @a Proxy se -> return e
-        _ -> fail $ showParamType @a Proxy <> " not available from context"
+        se@(SomeExpr ctx)
+            | Just e <- paramFromSomeExpr @a Proxy se -> return e
+            | otherwise -> fail $ showParamType @a Proxy <> " not available from context type '" <> T.unpack (textExprType ctx) <> "'"
 
 paramOrContext :: forall a. ParamType a => String -> CommandDef a
 paramOrContext name = CommandDef [(name, SomeParam (Proxy @(ParamOrContext a)) Proxy)] (\[SomeParam Proxy (Identity x)] -> fromJust $ cast x)
@@ -515,7 +517,8 @@ testWith = do
     off <- stateOffset <$> getParserState
     ctx@(SomeExpr (_ :: Expr ctxe)) <- someExpr
     let expected =
-            [ SomeExprType @Node Proxy
+            [ SomeExprType @Network Proxy
+            , SomeExprType @Node Proxy
             , SomeExprType @Process Proxy
             ]
     notAllowed <- flip allM expected $ \case
@@ -529,12 +532,13 @@ testWith = do
 
     indent <- L.indentGuard scn GT ref
     localState $ do
-        modify $ \s -> s { testContext = Just ctx }
+        modify $ \s -> s { testContext = ctx }
         testBlock indent
 
 testNode :: TestParser [TestStep]
 testNode = command "node" $ DeclNode
     <$> param ""
+    <*> paramOrContext "on"
     <*> innerBlock
 
 testSpawn :: TestParser [TestStep]
@@ -617,7 +621,7 @@ parseTestFile path = do
     content <- TL.readFile path
     let initState = TestParserState
             { testVars = []
-            , testContext = Nothing
+            , testContext = SomeExpr RootNetwork
             }
     case evalState (runParserT parseTestDefinitions path content) initState of
          Left err -> putStr (errorBundlePretty err) >> exitFailure
diff --git a/src/Test.hs b/src/Test.hs
index 836992c..16936bb 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -34,8 +34,8 @@ data Test = Test
     }
 
 data TestStep = forall a. ExprType a => Let SourceLine VarName (Expr a) [TestStep]
-              | DeclNode (TypedVarName Node) [TestStep]
-              | Spawn (TypedVarName Process) (Either (TypedVarName Node) (Expr Node)) [TestStep]
+              | DeclNode (TypedVarName Node) (Expr Network) [TestStep]
+              | Spawn (TypedVarName Process) (Either (TypedVarName Node) (Either (Expr Network) (Expr Node))) [TestStep]
               | Send (Expr Process) (Expr Text)
               | Expect SourceLine (Expr Process) (Expr Regex) [TypedVarName Text] [TestStep]
               | Guard SourceLine (Expr Bool)
@@ -50,12 +50,13 @@ newtype SourceLine = SourceLine Text
 
 class MonadFail m => MonadEval m where
   lookupVar :: VarName -> m SomeVarValue
+  rootNetwork :: m Network
 
 
 newtype VarName = VarName Text
     deriving (Eq, Ord)
 
-newtype TypedVarName a = TypedVarName VarName
+newtype TypedVarName a = TypedVarName { fromTypedVarName :: VarName }
     deriving (Eq, Ord)
 
 textVarName :: VarName -> Text
@@ -118,6 +119,7 @@ data Expr a where
     Regex :: [Expr Regex] -> Expr Regex
     UnOp :: (b -> a) -> Expr b -> Expr a
     BinOp :: (b -> c -> a) -> Expr b -> Expr c -> Expr a
+    RootNetwork :: Expr Network
 
 eval :: MonadEval m => Expr a -> m a
 eval (Variable name) = fromSomeVarValue name =<< lookupVar name
@@ -130,6 +132,7 @@ eval (Regex xs) = mapM eval xs >>= \case
         Right re -> return re
 eval (UnOp f x) = f <$> eval x
 eval (BinOp f x y) = f <$> eval x <*> eval y
+eval (RootNetwork) = rootNetwork
 
 gatherVars :: forall a m. MonadEval m => Expr a -> m [(VarName, SomeVarValue)]
 gatherVars = fmap (uniqOn fst . sortOn fst) . helper
@@ -141,6 +144,7 @@ gatherVars = fmap (uniqOn fst . sortOn fst) . helper
     helper (Regex es) = concat <$> mapM helper es
     helper (UnOp _ e) = helper e
     helper (BinOp _ e f) = (++) <$> helper e <*> helper f
+    helper (RootNetwork) = return []
 
 
 data Regex = RegexCompiled Text RE.Regex