Trace regex variables in expect using new type

3 files changed, 16 insertions, 6 deletions

diff --git a/src/Parser/Statement.hs b/src/Parser/Statement.hs
index 21b24a6..a65227d 100644
--- a/src/Parser/Statement.hs
+++ b/src/Parser/Statement.hs
@@ -153,6 +153,12 @@ instance (ParamType a, ParamType b) => ParamType (Either a b) where
     paramFromSomeExpr _ se = (Left <$> paramFromSomeExpr @a Proxy se) <|> (Right <$> paramFromSomeExpr @b Proxy se)
     paramExpr = either (fmap Left . paramExpr) (fmap Right . paramExpr)
 
+instance ExprType a => ParamType (Traced a) where
+    type ParamRep (Traced a) = Expr a
+    parseParam _ = parseParam (Proxy @(Expr a))
+    showParamType _ = showParamType (Proxy @(Expr a))
+    paramExpr = Trace
+
 data SomeParam f = forall a. ParamType a => SomeParam (Proxy a) (f (ParamRep a))
 
 data CommandDef a = CommandDef [(String, SomeParam Proxy)] ([SomeParam Identity] -> Expr a)
diff --git a/src/Run.hs b/src/Run.hs
index 54fdba6..f6eba39 100644
--- a/src/Run.hs
+++ b/src/Run.hs
@@ -273,9 +273,8 @@ exprFailed desc sline pname exprVars = do
             ]
     throwError Failed
 
-expect :: SourceLine -> Process -> Expr Regex -> [TypedVarName Text] -> TestRun () -> TestRun ()
-expect sline p expr tvars inner = do
-    re <- eval expr
+expect :: SourceLine -> Process -> Traced Regex -> [TypedVarName Text] -> TestRun () -> TestRun ()
+expect sline p (Traced trace re) tvars inner = do
     timeout <- asks $ optTimeout . teOptions . fst
     delay <- liftIO $ registerDelay $ ceiling $ 1000000 * timeout
     mbmatch <- atomicallyTest $ (Nothing <$ (check =<< readTVar delay)) <|> do
@@ -302,7 +301,7 @@ expect sline p expr tvars inner = do
              outProc OutputMatch p line
              local (fmap $ \s -> s { tsVars = zip vars (map someConstValue capture) ++ tsVars s }) inner
 
-         Nothing -> exprFailed (T.pack "expect") sline (Just $ procName p) =<< gatherVars expr
+         Nothing -> exprFailed (T.pack "expect") sline (Just $ procName p) trace
 
 flush :: Process -> Maybe Regex -> TestRun ()
 flush p mbre = do
diff --git a/src/Test.hs b/src/Test.hs
index effd00a..53e0f03 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -20,7 +20,7 @@ module Test (
     ExprListUnpacker(..),
     ExprEnumerator(..),
     Expr(..), varExpr, eval, evalSome,
-    EvalTrace, VarNameSelectors, gatherVars,
+    Traced(..), EvalTrace, VarNameSelectors, gatherVars,
     AppAnnotation(..),
 
     ArgumentKeyword(..), FunctionArguments(..),
@@ -70,7 +70,7 @@ data TestStep
     | DeclNode (TypedVarName Node) Network (Expr TestBlock)
     | Spawn (TypedVarName Process) (Either Network Node) (Expr TestBlock)
     | Send Process Text
-    | Expect SourceLine Process (Expr Regex) [ TypedVarName Text ] (Expr TestBlock)
+    | Expect SourceLine Process (Traced Regex) [ TypedVarName Text ] (Expr TestBlock)
     | Flush Process (Maybe Regex)
     | Guard SourceLine EvalTrace Bool
     | DisconnectNode Node TestBlock
@@ -261,6 +261,7 @@ data Expr a where
     Concat :: [Expr Text] -> Expr Text
     Regex :: [Expr Regex] -> Expr Regex
     Undefined :: String -> Expr a
+    Trace :: Expr a -> Expr (Traced a)
 
 data AppAnnotation b = AnnNone
                      | ExprType b => AnnRecord Text
@@ -302,6 +303,7 @@ eval = \case
             Left err -> fail err
             Right re -> return re
     Undefined err -> fail err
+    Trace expr -> Traced <$> gatherVars expr <*> eval expr
 
 evalSome :: MonadEval m => SomeExpr -> m SomeVarValue
 evalSome (SomeExpr expr) = fmap SomeVarValue $ VarValue
@@ -309,6 +311,8 @@ evalSome (SomeExpr expr) = fmap SomeVarValue $ VarValue
     <*> pure mempty
     <*> (const . const <$> eval expr)
 
+data Traced a = Traced EvalTrace a
+
 type VarNameSelectors = ( VarName, [ Text ] )
 type EvalTrace = [ ( VarNameSelectors, SomeVarValue ) ]
 
@@ -339,6 +343,7 @@ gatherVars = fmap (uniqOn fst . sortOn fst) . helper
         Concat es -> concat <$> mapM helper es
         Regex es -> concat <$> mapM helper es
         Undefined {} -> return []
+        Trace expr -> helper expr
 
     gatherSelectors :: forall b. Expr b -> Maybe (VarName, [Text])
     gatherSelectors = \case