Variable expansion in strings and regexes

3 files changed, 96 insertions, 28 deletions

diff --git a/src/Main.hs b/src/Main.hs
index ffd6292..dc1cc42 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -1,11 +1,13 @@
 module Main where
 
+import Control.Arrow
 import Control.Applicative
 import Control.Concurrent
 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
@@ -14,7 +16,7 @@ import Data.Text (Text)
 import qualified Data.Text as T
 import qualified Data.Text.IO as T
 
-import Text.Read
+import Text.Read (readMaybe)
 import Text.Regex.TDFA
 import Text.Regex.TDFA.Text
 
@@ -74,8 +76,12 @@ data TestEnv = TestEnv
     , teOptions :: Options
     }
 
-newtype TestRun a = TestRun { fromTestRun :: ReaderT TestEnv (ExceptT () IO) a }
-    deriving (Functor, Applicative, Monad, MonadReader TestEnv, MonadIO)
+data TestState = TestState
+    { tsVars :: [(VarName, Text)]
+    }
+
+newtype TestRun a = TestRun { fromTestRun :: ReaderT TestEnv (StateT TestState (ExceptT () IO)) a }
+    deriving (Functor, Applicative, Monad, MonadReader TestEnv, MonadState TestState, MonadIO)
 
 instance MonadFail TestRun where
     fail str = do
@@ -90,6 +96,8 @@ instance MonadError () TestRun where
 
     catchError (TestRun act) handler = TestRun $ catchError act $ fromTestRun . handler
 
+instance MonadEval TestRun where
+    lookupStringVar name = maybe (fail $ "variable not in scope: '" ++ unpackVarName name ++ "'") return =<< gets (lookup name . tsVars)
 
 instance MonadOutput TestRun where
     getOutput = asks teOutput
@@ -97,8 +105,9 @@ instance MonadOutput TestRun where
 forkTest :: TestRun () -> TestRun ()
 forkTest act = do
     tenv <- ask
+    tstate <- get
     void $ liftIO $ forkIO $ do
-        runExceptT (runReaderT (fromTestRun act) tenv) >>= \case
+        runExceptT (flip evalStateT tstate $ flip runReaderT tenv $ fromTestRun act) >>= \case
             Left () -> atomically $ writeTVar (teFailed tenv) True
             Right () -> return ()
 
@@ -283,7 +292,9 @@ runTest out opts test = do
         <$> pure out
         <*> newTVarIO False
         <*> pure opts
-    (fmap $ either (const False) id) $ runExceptT $ flip runReaderT tenv $ fromTestRun $ do
+    tstate <- TestState
+        <$> pure []
+    (fmap $ either (const False) id) $ runExceptT $ flip evalStateT tstate $ flip runReaderT tenv $ fromTestRun $ do
         net <- initNetwork
 
         let sigHandler SignalInfo { siginfoSpecific = chld } = do
@@ -305,12 +316,15 @@ runTest out opts test = do
                 void $ spawnOn (Right node) pname Nothing $
                     fromMaybe (optDefaultTool opts) (lookup pname $ optProcTools opts)
 
-            Send pname line -> do
+            Send pname expr -> do
                 p <- getProcess net pname
+                line <- evalStringExpr expr
                 send p line
 
-            Expect pname regex pat -> do
+            Expect pname expr@(RegexExpr ps) -> do
                 p <- getProcess net pname
+                regex <- evalRegexExpr expr
+                pat <- evalStringExpr (StringExpr $ map (left T.pack) ps)
                 expect p regex pat
 
             Wait -> do
diff --git a/src/Parser.hs b/src/Parser.hs
index 4b98dde..ddbdfd6 100644
--- a/src/Parser.hs
+++ b/src/Parser.hs
@@ -1,5 +1,7 @@
 {-# LANGUAGE TemplateHaskell #-}
 
+{-# OPTIONS_GHC -Wno-orphans #-}
+
 module Parser (
     parseTestFile,
 ) where
@@ -10,7 +12,7 @@ import Control.Monad.State
 import Data.Char
 import Data.Set (Set)
 import qualified Data.Set as S
-import Data.Text (Text)
+import Data.Text qualified as T
 import qualified Data.Text.Lazy as TL
 import qualified Data.Text.Lazy.IO as TL
 import Data.Void
@@ -20,8 +22,6 @@ import Generics.Deriving.Base as G
 import Text.Megaparsec hiding (State)
 import Text.Megaparsec.Char
 import qualified Text.Megaparsec.Char.Lexer as L
-import Text.Regex.TDFA (defaultCompOpt, defaultExecOpt)
-import Text.Regex.TDFA.String
 
 import System.Exit
 
@@ -31,6 +31,9 @@ type TestParser = ParsecT Void TestStream (State (Set ProcName))
 
 type TestStream = TL.Text
 
+instance MonadEval TestParser where
+    lookupStringVar _ = return T.empty
+
 skipLineComment :: TestParser ()
 skipLineComment = L.skipLineComment $ TL.pack "#"
 
@@ -77,12 +80,23 @@ procName = label "process name" $ lexeme $ do
     cs <- takeWhileP Nothing (\x -> isAlphaNum x || x == '_' || x == '-')
     return $ ProcName $ TL.toStrict (c `TL.cons` cs)
 
-quotedString :: TestParser Text
+varExpansion :: TestParser VarName
+varExpansion = do
+    void $ char '$'
+    choice
+        [ VarName . (:[]) . TL.toStrict <$> takeWhile1P Nothing (\x -> isAlphaNum x || x == '_')
+        ,do void $ char '{'
+            name <- takeWhile1P Nothing (/='}')
+            void $ char '}'
+            return $ VarName $ T.splitOn (T.singleton '.') (TL.toStrict name)
+        ]
+
+quotedString :: TestParser StringExpr
 quotedString = label "string" $ lexeme $ do
     symbol "\""
     let inner = choice
             [ char '"' >> return []
-            , takeWhile1P Nothing (`notElem` "\"\\") >>= \s -> (s:) <$> inner
+            , takeWhile1P Nothing (`notElem` "\"\\$") >>= \s -> (Left (TL.toStrict s):) <$> inner
             ,do void $ char '\\'
                 c <- choice
                     [ char '\\' >> return '\\'
@@ -92,27 +106,30 @@ quotedString = label "string" $ lexeme $ do
                     , char 'r' >> return '\r'
                     , char 't' >> return '\t'
                     ]
-                (TL.singleton c:) <$> inner
+                (Left (T.singleton c) :) <$> inner
+            ,do name <- varExpansion
+                (Right name :) <$> inner
             ]
-    TL.toStrict . TL.concat <$> inner
+    StringExpr <$> inner
 
-regex :: TestParser (Regex, Text)
+regex :: TestParser RegexExpr
 regex = label "regular expression" $ lexeme $ do
     symbol "/"
     let inner = choice
             [ char '/' >> return []
-            , takeWhile1P Nothing (`notElem` "/\\") >>= \s -> (s:) <$> inner
+            , takeWhile1P Nothing (`notElem` "/\\$") >>= \s -> (Left (TL.unpack s) :) <$> inner
             ,do void $ char '\\'
                 s <- choice
-                    [ char '/' >> return (TL.singleton '/')
-                    , anySingle >>= \c -> return (TL.pack ['\\', c])
+                    [ char '/' >> return (Left $ "/")
+                    , anySingle >>= \c -> return (Left ['\\', c])
                     ]
                 (s:) <$> inner
+            ,do name <- varExpansion
+                (Right name :) <$> inner
             ]
-    pat <- TL.concat <$> inner
-    case compile defaultCompOpt defaultExecOpt ("^" ++ TL.unpack pat ++ "$") of
-         Left err -> fail err
-         Right re -> return (re, TL.toStrict pat)
+    expr <- RegexExpr <$> inner
+    _ <- evalRegexExpr expr -- test regex parsing with empty variables
+    return expr
 
 
 class GInit f where ginit :: f x
@@ -171,7 +188,7 @@ testSpawn = command "spawn"
 
 data SendBuilder = SendBuilder
     { _sendBuilderProc :: Maybe ProcName
-    , _sendBuilderLine :: Maybe Text
+    , _sendBuilderLine :: Maybe StringExpr
     }
     deriving (Generic)
 
@@ -188,7 +205,7 @@ testSend = command "send"
 
 data ExpectBuilder = ExpectBuilder
     { _expectBuilderProc :: Maybe ProcName
-    , _expectBuilderRegex :: Maybe (Regex, Text)
+    , _expectBuilderRegex :: Maybe RegexExpr
     }
     deriving (Generic)
 
@@ -200,8 +217,7 @@ testExpect = command "expect"
     , Param "" expectBuilderRegex regex
     ] $ \b -> Expect
         <$> (maybe (fail "missing 'from' <proc>") return $ b ^. expectBuilderProc)
-        <*> (maybe (fail "missing regex to match") (return . fst) $ b ^. expectBuilderRegex)
-        <*> (maybe (fail "missing regex to match") (return . snd) $ b ^. expectBuilderRegex)
+        <*> (maybe (fail "missing regex to match") return $ b ^. expectBuilderRegex)
 
 
 testWait :: TestParser [TestStep]
diff --git a/src/Test.hs b/src/Test.hs
index 465b424..d652f9b 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -4,12 +4,21 @@ module Test (
 
     ProcName(..), textProcName, unpackProcName,
     NodeName(..), textNodeName, unpackNodeName,
+
+    MonadEval(..),
+    VarName(..), unpackVarName,
+    StringExpr(..), evalStringExpr,
+    RegexExpr(..), evalRegexExpr,
 ) where
 
+import Control.Monad
+
+import Data.List
 import Data.Text (Text)
 import qualified Data.Text as T
 
 import Text.Regex.TDFA
+import Text.Regex.TDFA.String
 
 import Process
 
@@ -19,8 +28,8 @@ data Test = Test
     }
 
 data TestStep = Spawn ProcName NodeName
-              | Send ProcName Text
-              | Expect ProcName Regex Text
+              | Send ProcName StringExpr
+              | Expect ProcName RegexExpr
               | Wait
 
 newtype NodeName = NodeName Text
@@ -31,3 +40,32 @@ textNodeName (NodeName name) = name
 
 unpackNodeName :: NodeName -> String
 unpackNodeName (NodeName tname) = T.unpack tname
+
+
+class Monad m => MonadEval m where
+  lookupStringVar :: VarName -> m Text
+
+
+data VarName = VarName [Text]
+    deriving (Eq, Ord)
+
+unpackVarName :: VarName -> String
+unpackVarName (VarName name) = concat $ intersperse "." $ map T.unpack name
+
+data StringExpr = StringExpr [Either Text VarName]
+
+evalStringExpr :: MonadEval m => StringExpr -> m Text
+evalStringExpr (StringExpr xs) = fmap T.concat $ forM xs $ \case
+    Left text -> return text
+    Right var -> lookupStringVar var
+
+data RegexExpr = RegexExpr [Either String VarName]
+
+evalRegexExpr :: (MonadFail m, MonadEval m) => RegexExpr -> m Regex
+evalRegexExpr (RegexExpr xs) = do
+    parts <- forM xs $ \case
+        Left str -> return str
+        Right var -> concatMap (\c -> ['\\',c]) . T.unpack <$> lookupStringVar var
+    case compile defaultCompOpt defaultExecOpt $ concat $ concat [["^"], parts, ["$"]] of
+        Left err -> fail err
+        Right re -> return re